optimize dual pp overlap (PaddlePaddle#74527)

* optimize dual pp overlap

* polish code

* polish code

Author: phlrain

paddle/fluid/distributed/collective/deep_ep/include/event_pool.h

@@ -22,7 +22,7 @@ namespace deep_ep::detail {
 
 class EventPool {
  public:
-  EventPool() = default;
+  EventPool();
   EventPool(const EventPool&) = delete;
   EventPool(EventPool&&) = delete;
   ~EventPool();

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

@@ -22,6 +22,16 @@ EventPool &EventPool::Instance() {
   return pool;
 }
 
+EventPool::EventPool() {
+  for (size_t i = 0; i < 1000; ++i) {
+    cudaEvent_t new_event;
+    CUDA_CHECK(cudaEventCreate(&new_event));
+
+    cudaEventRecord(new_event, 0);
+    incomplished_events_.push(new_event);
+  }
+}
+
 EventPool::~EventPool() {
   const auto &DestroyEvent = [](cudaEvent_t event) {
     cudaError_t e = cudaEventDestroy(event);

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

@@ -37,7 +37,7 @@
     PipelineParallel,
 )
 from .pp_utils.batch_comm_helper import BatchCommHelper
-from .zero_bubble_utils import WeightGradStore
+from .zero_bubble_utils import EventStore, WeightGradStore
 
 __all__ = []
 
@@ -358,6 +358,10 @@ def _commit_and_wait_comm(
             else 0
         )
         if common_forward_ops_num == 0 and common_backward_ops_num == 0:
+            if EventStore.event is not None:
+                e_t = EventStore.event
+                EventStore.event = None
+                return e_t
             return deep_ep.get_event_from_custom_stream(
                 paddle.device.current_stream().stream_base
             )
@@ -387,13 +391,28 @@ def _commit_and_wait_comm(
                 pp_raw_stream
             )
 
+        backward_outer_event_wait = False
+        if EventStore.event is not None:
+            with paddle.device.stream_guard(
+                paddle.device.Stream(stream_base=pp_raw_stream)
+            ):
+                EventStore.event.current_stream_wait()
+
+            EventStore.set(None)
+            self.pp_group.process_group.set_outer_wait(True)
+
+            backward_outer_event_wait = True
+
         if common_backward_ops_num > 0:
             bwd_reqs = batch_isend_irecv(self.comm_backward_ops)
 
             if not use_stream_wait_event:
                 for req in bwd_reqs:
                     req.wait()
 
+        if backward_outer_event_wait:
+            self.pp_group.process_group.set_outer_wait(False)
+
         if use_stream_wait_event:
             forward_event_to_wait.current_stream_wait()
 

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

@@ -54,6 +54,15 @@ def clear(cls) -> None:
         cls.funcs_queue = queue.Queue()
 
 
+class EventStore:
+
+    event = None
+
+    @classmethod
+    def set(cls, event) -> None:
+        cls.event = event
+
+
 def fold_init_dims(tensor):
     # NOTE(zhangyuqin1998): Reshape a rank-3 tensor from P x M x N to (P * M) x N,
     # to keep weight_grad in a correct rank. See phi::FoldInitDims.