allco on comm (#10859)

Author: zhangbo9674

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -569,13 +569,18 @@ def mlp_forward(self, inputs):
         ret = (inputs_embeds_mtp, *ret) if self.send_mtp_embed else ret
         return ret
 
-    def combine_forward(self, inputs, async_finish=False):
+    def combine_forward(self, inputs, async_finish=False, previous_event=None, allocate_on_comm_stream=False):
         if self.send_mtp_embed:
             (inputs_embeds_mtp, hidden_states, residual, l_aux, hidden_states_out) = inputs
         else:
             (hidden_states, residual, l_aux, hidden_states_out) = inputs
 
-        output_combine = self.fp8_fusion_moe_node.combine_node.forward(hidden_states_out, async_finish=async_finish)
+        output_combine = self.fp8_fusion_moe_node.combine_node.forward(
+            hidden_states_out,
+            async_finish=async_finish,
+            previous_event=previous_event,
+            allocate_on_comm_stream=allocate_on_comm_stream and previous_event is not None,
+        )
 
         ret = (hidden_states, residual, l_aux, output_combine)
 
@@ -652,7 +657,7 @@ def mlp_backward(self, output_grad):
         ret = (inputs_embeds_mtp_grad, *ret) if self.send_mtp_embed else ret
         return ret
 
-    def dispatch_backward(self, output_grad, async_finish=False):
+    def dispatch_backward(self, output_grad, async_finish=False, previous_event=None, allocate_on_comm_stream=False):
         if self.send_mtp_embed:
             (
                 inputs_embeds_mtp_grad,
@@ -666,7 +671,11 @@ def dispatch_backward(self, output_grad, async_finish=False):
             hidden_states_grad, residual_grad, l_aux_grad, hs_dispatched_grad, dispatched_probs_grad = output_grad
 
         hs_grad, token_probs_grad = self.fp8_fusion_moe_node.dispatch_node.backward(
-            hs_dispatched_grad, dispatched_probs_grad, async_finish=async_finish
+            hs_dispatched_grad,
+            dispatched_probs_grad,
+            async_finish=async_finish,
+            previous_event=previous_event,
+            allocate_on_comm_stream=allocate_on_comm_stream and previous_event is not None,
         )
 
         ret = (hidden_states_grad, residual_grad, l_aux_grad, hs_grad, token_probs_grad)
@@ -755,6 +764,8 @@ def forward_backward(self, inputs, output_grad, event_to_wait=None):
         output_grad = self.backward_node.mlp_backward(output_grad)
         paddle.base.core.nvprof_nvtx_pop()
 
+        output_grad_event = deep_ep.get_event_from_calc_stream(self.backward_node.moe_group.id)
+
         paddle.base.core.nvprof_nvtx_push("dispatch_forward")
         inputs = self.forward_node.dispatch_forward(
             inputs, previous_event=attn_compute_event, async_finish=True, allocate_on_comm_stream=True
@@ -763,7 +774,9 @@ def forward_backward(self, inputs, output_grad, event_to_wait=None):
         dispatch_forward_event = deep_ep.get_event_from_comm_stream(self.forward_node.moe_group.id)
 
         paddle.base.core.nvprof_nvtx_push("dispatch_backward")
-        output_grad = self.backward_node.dispatch_backward(output_grad, async_finish=True)
+        output_grad = self.backward_node.dispatch_backward(
+            output_grad, async_finish=True, previous_event=output_grad_event, allocate_on_comm_stream=True
+        )
         paddle.base.core.nvprof_nvtx_pop()
         # get dispatch backward event
         dispatch_backward_event = deep_ep.get_event_from_comm_stream(self.backward_node.moe_group.id)
@@ -777,8 +790,12 @@ def forward_backward(self, inputs, output_grad, event_to_wait=None):
         inputs = self.forward_node.mlp_forward(inputs)
         paddle.base.core.nvprof_nvtx_pop()
 
+        inputs_event = deep_ep.get_event_from_calc_stream(self.forward_node.moe_group.id)
+
         paddle.base.core.nvprof_nvtx_push("combine_forward")
-        inputs = self.forward_node.combine_forward(inputs, async_finish=True)
+        inputs = self.forward_node.combine_forward(
+            inputs, async_finish=True, previous_event=inputs_event, allocate_on_comm_stream=True
+        )
         paddle.base.core.nvprof_nvtx_pop()
         combine_forward_event = deep_ep.get_event_from_comm_stream(self.forward_node.moe_group.id)
 

paddlenlp/transformers/fused_a2a.py

@@ -308,7 +308,9 @@ def forward(
 
         return recv_x, recv_token_probs, states
 
-    def backward(self, grad_output, grad_token_probs, previous_event=None, async_finish=False):
+    def backward(
+        self, grad_output, grad_token_probs, previous_event=None, async_finish=False, allocate_on_comm_stream=False
+    ):
         """Backward pass of fused dispatch."""
         out = fused_dispatch_backward_func(
             grad_output,
@@ -317,6 +319,7 @@ def backward(self, grad_output, grad_token_probs, previous_event=None, async_fin
             self.handle,
             previous_event=previous_event,
             async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
         )
         self.reset_statue()
         return out
@@ -329,12 +332,17 @@ def __init__(self, name="combine"):
     def reset_statue(self):
         self.handle = None
 
-    def forward(self, x, group, handle, previous_event=None, async_finish=False):
+    def forward(self, x, group, handle, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
         """Forward pass of fused combine."""
         states = dict()
         states["handle"] = handle
         combined_x = fused_combine_forward_func(
-            x, group, states, previous_event=previous_event, async_finish=async_finish
+            x,
+            group,
+            states,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
         )
 
         self.handle = handle

paddlenlp/transformers/moe_layer.py

@@ -596,13 +596,21 @@ def forward(
         return hs_2d_dispatched, dispatched_indices, dispatched_probs
 
     @paddle.no_grad()
-    def backward(self, hs_dispatched_grad, dispatched_probs_grad, previous_event=None, async_finish=False):
+    def backward(
+        self,
+        hs_dispatched_grad,
+        dispatched_probs_grad,
+        previous_event=None,
+        async_finish=False,
+        allocate_on_comm_stream=False,
+    ):
         # dispatch grad
         hs_grad, _, token_probs_grad = self.dispatch_act_node.backward(
             hs_dispatched_grad,
             dispatched_probs_grad,
             previous_event=previous_event,
             async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
         )
         return hs_grad, token_probs_grad
 
@@ -614,14 +622,15 @@ def __init__(self, token_dispatcher, name="fp8_combine_node"):
         self.name = name
 
     @paddle.no_grad()
-    def forward(self, hidden_states_out, previous_event=None, async_finish=False):
+    def forward(self, hidden_states_out, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
         # combine
         output_combine = self.combine_node.forward(
             hidden_states_out,
             self.token_dispatcher._comm_manager.group,
             self.token_dispatcher._comm_manager.handle,
             previous_event=previous_event,
             async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
         )
         output_combine.stop_gradient = False
         return output_combine