Bind get stream method to python for deep_ep

paddle/fluid/distributed/collective/deep_ep/event.hpp

@@ -45,6 +45,11 @@ struct EventHandle {
  event->record(stream);
  }
 
+ explicit EventHandle(const phi::CUDAStream& stream) {
+ event = std::make_shared<deep_ep::detail::Event>();
+ event->record(stream.raw_stream());
+ }
+
  EventHandle(const EventHandle& other) = default;
 
  void current_stream_wait() const {
@@ -57,6 +62,7 @@ struct EventHandle {
 
 EventHandle GetEventHandleFromCalcStream(int context_ring_id);
 EventHandle GetEventHandleFromCommStream(int context_ring_id);
+EventHandle GetEventHandleFromCustomStream(const phi::CUDAStream& stream);
 
 inline deep_ep::detail::Event create_event(const cudaStream_t& s) {
  auto event = deep_ep::detail::Event();

paddle/fluid/distributed/collective/deep_ep/src/event.cc

@@ -34,4 +34,8 @@ EventHandle GetEventHandleFromCalcStream(int context_ring_id) {
 EventHandle GetEventHandleFromCommStream(int context_ring_id) {
  return EventHandle(detail::GetCommStreamFromGroup(context_ring_id));
 }
+
+EventHandle GetEventHandleFromCustomStream(const phi::CUDAStream& stream) {
+ return EventHandle(stream);
+}
 } // namespace deep_ep

paddle/fluid/pybind/deep_ep_api.cc

@@ -53,6 +53,9 @@ void BindDeepEPApi(pybind11::module *m) {
  m->def("get_event_handle_from_comm_stream",
  &deep_ep::GetEventHandleFromCommStream);
 
+ m->def("get_event_handle_from_custom_stream",
+ &deep_ep::GetEventHandleFromCustomStream);
+
  pybind11::class_<deep_ep::Buffer>(*m, "Buffer")
  .def(pybind11::init<int, int, int64_t, int64_t, bool, int>())
  .def("is_available", &deep_ep::Buffer::is_available)

python/paddle/distributed/communication/deep_ep/__init__.py

@@ -19,6 +19,7 @@
  EventOverlap,
  get_event_from_calc_stream,
  get_event_from_comm_stream,
+ get_event_from_custom_stream,
 )
 
 __all__ = [
@@ -27,4 +28,5 @@
  "Config",
  "get_event_from_calc_stream",
  "get_event_from_comm_stream",
+ "get_event_from_custom_stream",
 ]

python/paddle/distributed/communication/deep_ep/utils.py

@@ -101,3 +101,9 @@ def get_event_from_comm_stream(group_id: int) -> EventOverlap:
  return EventOverlap(
  event=paddle.base.core.get_event_handle_from_comm_stream(group_id)
  )
+
+
+def get_event_from_custom_stream(stream) -> EventOverlap:
+ return EventOverlap(
+ event=paddle.base.core.get_event_handle_from_custom_stream(stream)
+ )

python/paddle/distributed/fleet/meta_parallel/dualpipev.py

@@ -296,7 +296,7 @@ def _forward_backward_compute(
  _, backward_grads = self._get_backward_inputs(
  backward_phase, backward_acc_id
  )
-
+ event_to_wait = paddle.device.current_stream().record_event()
  # forward & backward
  forward_chunk = self._layers.get_schedule_chunk(chunk_id=forward_phase)
  backward_chunk = self.schedule_chunks[backward_phase][backward_acc_id]
@@ -309,6 +309,7 @@ def _forward_backward_compute(
  backward_loss_fn_node,
  backward_grads,
  self.scaler,
+ event_to_wait,
  )
  )
  self.schedule_chunks[backward_phase][backward_acc_id] = None