Fix slow convergence issue when enable overlap_p2p_comm (#10992)

Author: lshpku

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -1086,7 +1086,6 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
         paddle.base.core.nvprof_nvtx_push("mlp_forward")
         inputs = self.forward_node.mlp_forward(inputs)
         paddle.base.core.nvprof_nvtx_pop()
-        mlp_fwd_event = deep_ep.get_event_from_calc_stream(self.forward_node.moe_group.id)
 
         if pp_stream is not None:
             paddle.base.core.nvprof_nvtx_push("post_process_forward")
@@ -1098,7 +1097,7 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
 
         paddle.base.core.nvprof_nvtx_push("combine_forward")
         inputs = self.forward_node.combine_forward(
-            inputs, previous_event=mlp_fwd_event, async_finish=True, allocate_on_comm_stream=True
+            inputs, previous_event=final_out_event, async_finish=True, allocate_on_comm_stream=True
         )
         paddle.base.core.nvprof_nvtx_pop()
 
@@ -1109,9 +1108,6 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
         if pp_stream is not None:
             send_recv_stream = paddle.device.Stream(stream_base=pp_stream)
 
-            # combine_forward_event.custom_stream_wait( pp_stream)
-            # final_out_event.custom_stream_wait(pp_stream)
-
             paddle.base.core.nvprof_nvtx_push("pp stream add")
 
             with paddle.device.stream_guard(send_recv_stream):
@@ -1228,6 +1224,8 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
             paddle.base.core.nvprof_nvtx_pop()  # moe_attn
 
             paddle.base.core.nvprof_nvtx_push("dense_mlp_moe_dispatch")
+            if combine_bw_event_to_wait is not None:
+                combine_bw_event_to_wait.calc_stream_wait(self.forward_node.moe_group.id)
             output_grad = self.backward_node.mlp_node.backward(output_grad)
             inputs = self.forward_node.dispatch_forward(
                 inputs, previous_event=attn_fw_event, async_finish=True, allocate_on_comm_stream=True