support send recv overlap pre code (#10872)

Author: phlrain

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -801,7 +801,7 @@ def forward_backward(self, inputs, output_grad, event_to_wait=None):
         dispatch_backward_event.calc_stream_wait(self.backward_node.moe_group.id)
         paddle.base.core.nvprof_nvtx_push("attn_backward")
         output_grad = self.backward_node.attn_backward(output_grad)
-        event_to_wait = paddle.device.current_stream().record_event()
+        event_to_wait = deep_ep.get_event_from_calc_stream(self.backward_node.moe_group.id)
         paddle.base.core.nvprof_nvtx_pop()
 
         combine_forward_event.calc_stream_wait(self.forward_node.moe_group.id)
@@ -1519,7 +1519,8 @@ def overlapped_forward_backward(
         backward_loss_fn_node,
         backward_input_grads,
         scaler,
-        event_to_wait=None,
+        combine_bw_event_to_wait = None,
+        pp_stream=None
     ):
         if backward_loss_fn_node is not None:
             if scaler:
@@ -1537,7 +1538,7 @@ def overlapped_forward_backward(
         forward_inputs = forward_pre_node.forward(forward_inputs)
         backward_input_grads = backward_pre_node.backward(backward_input_grads)
         forward_inputs, backward_input_grads, _ = overlap_node.forward_backward(
-            forward_inputs, backward_input_grads, event_to_wait
+            forward_inputs, backward_input_grads, combine_bw_event_to_wait
         )
         forward_inputs = forward_post_node.forward(forward_inputs)
         backward_input_grads = backward_post_node.backward(backward_input_grads)

paddlenlp/transformers/moe_layer.py

@@ -648,9 +648,10 @@ def backward(self, output_combine_grad, previous_event=None, async_finish=False,
 
 
 class Fp8CombineQuantNode:
-    def __init__(self, token_dispatcher, name="fp8_combine_quant_node"):
+    def __init__(self, token_dispatcher, moe_group=None, name="fp8_combine_quant_node"):
         self.token_dispatcher = token_dispatcher
         self.name = name
+        self.moe_group = moe_group
 
     @paddle.no_grad()
     def forward(self, output_combine):
@@ -666,13 +667,14 @@ def backward(self, output_grad, event_to_wait=None):
         # post combine grad
         if DSV3_USE_FP8_DISPATCH:
             if event_to_wait is not None:
+                assert self.moe_group is not None
+                event_to_wait.comm_stream_wait( self.moe_group.id)
                 buffer = get_buffer(self.token_dispatcher._comm_manager.group, get_hidden_bytes(output_grad))
                 custom_stream = paddle.device.Stream(stream_base=buffer.runtime.get_comm_stream())
-                custom_stream.wait_event(event_to_wait)
             else:
                 custom_stream = paddle.device.current_stream()
             with paddle.device.stream_guard(custom_stream):
-                output_combine_grad = paddle.reshape(output_grad, self.output_combine_shape)
+                output_combine_grad = paddle.reshape(output_grad, [-1, output_grad.shape[-1]])
                 # output_combine_grad quant to fp8
                 output_combine_grad_fp8, output_combine_grad_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
                     output_combine_grad, output_scale_transpose=False, quant_method="1x128", input_transpose=False
@@ -683,7 +685,7 @@ def backward(self, output_grad, event_to_wait=None):
                     quant_event = deep_ep.get_event_from_custom_stream(custom_stream.stream_base)
             return (output_combine_grad_fp8, output_combine_grad_scale), quant_event
         else:
-            output_combine_grad = paddle.reshape(output_grad, self.output_combine_shape)
+            output_combine_grad = paddle.reshape(output_grad, [-1, output_grad.shape[-1]])
             return output_combine_grad, None
 
 
@@ -873,7 +875,7 @@ def __init__(
             is_split_group_gemm=is_split_group_gemm,
         )
         self.combine_node = Fp8CombineNode(self.token_dispatcher)
-        self.combine_quant_node = Fp8CombineQuantNode(self.token_dispatcher)
+        self.combine_quant_node = Fp8CombineQuantNode(self.token_dispatcher, custom_map.moe_group)
         self.name = name
 
     @paddle.no_grad()