Refine fp8 quant and combine backward overlap (#10830)

* refine

* refine

* refine

Author: zhangbo9674

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -396,9 +396,9 @@ def __init__(self, forward_nodes, backward_nodes, use_fuion=True):
         for f, b in zip(forward_nodes, backward_nodes):
             self.nodes.append(schedule_node_class(f, b, f"OverlapedNode_{len(self.nodes)}"))
 
-    def forward_backward(self, inputs, output_grad):
+    def forward_backward(self, inputs, output_grad, event_to_wait=None):
         for n in self.nodes:
-            inputs, output_grad = n.forward_backward(inputs, output_grad)
+            inputs, output_grad, event_to_wait = n.forward_backward(inputs, output_grad, event_to_wait)
         return inputs, output_grad
 
 
@@ -409,7 +409,7 @@ def __init__(self, forward_node, backward_node, name=""):
         self.backward_node = backward_node
         self.name = name
 
-    def forward_backward(self, inputs, output_grad):
+    def forward_backward(self, inputs, output_grad, event_to_wait=None):
         paddle.base.core.nvprof_nvtx_push("forward_backward")
         output_grad = self.backward_node.post_process_node.backward(output_grad)
 
@@ -594,7 +594,7 @@ def post_process_forward(self, inputs):
 
         return inputs
 
-    def post_process_backward(self, output_grad):
+    def post_process_backward(self, output_grad, event_to_wait=None):
         if self.send_mtp_embed:
             (
                 inputs_embeds_mtp_grad,
@@ -610,43 +610,51 @@ def post_process_backward(self, output_grad):
                 l_aux_grad,
                 final_hidden_states_grad,
             ) = self.post_process_node.backward(output_grad)
-        output_combine_grad = self.fp8_fusion_moe_node.combine_quant_node.backward(final_hidden_states_grad)
+        output_combine_grad, quant_event = self.fp8_fusion_moe_node.combine_quant_node.backward(
+            final_hidden_states_grad, event_to_wait
+        )
         if self.send_mtp_embed:
             return (
                 inputs_embeds_mtp_grad,
                 hidden_states_grad,
                 residual_grad,
                 l_aux_grad,
                 output_combine_grad,
+                quant_event,
             )
         else:
             return (
                 hidden_states_grad,
                 residual_grad,
                 l_aux_grad,
                 output_combine_grad,
+                quant_event,
             )
 
-    def combine_backward(self, output_grad, async_finish=False):
+    def combine_backward(self, output_grad, async_finish=False, allocate_on_comm_stream=False):
         if self.send_mtp_embed:
             (
                 inputs_embeds_mtp_grad,
                 hidden_states_grad,
                 residual_grad,
                 l_aux_grad,
                 output_combine_grad,
+                quant_event,
             ) = output_grad
         else:
             (
                 hidden_states_grad,
                 residual_grad,
                 l_aux_grad,
                 output_combine_grad,
+                quant_event,
             ) = output_grad
 
         hidden_states_out_grad = self.fp8_fusion_moe_node.combine_node.backward(
             output_combine_grad,
             async_finish=async_finish,
+            previous_event=quant_event,
+            allocate_on_comm_stream=allocate_on_comm_stream,
         )
 
         if self.send_mtp_embed:
@@ -811,35 +819,36 @@ def __init__(self, forward_node, backward_node, name=""):
         self.backward_node = backward_node
         self.name = name
 
-    def forward_backward(self, inputs, output_grad):
+    def forward_backward(self, inputs, output_grad, event_to_wait=None):
         paddle.base.core.nvprof_nvtx_push("forward_backward")
 
         paddle.base.core.nvprof_nvtx_push("post_process_backward")
-        output_grad = self.backward_node.post_process_backward(output_grad)
+        output_grad = self.backward_node.post_process_backward(output_grad, event_to_wait)
         paddle.base.core.nvprof_nvtx_pop()
 
         paddle.base.core.nvprof_nvtx_push("combine_backward")
-        output_grad = self.backward_node.combine_backward(output_grad, async_finish=True)
+        output_grad = self.backward_node.combine_backward(output_grad, async_finish=True, allocate_on_comm_stream=True)
         # get combine event
-        combine_backward_event = deep_ep.get_event_from_comm_stream( self.backward_node.moe_group.id)
+        combine_backward_event = deep_ep.get_event_from_comm_stream(self.backward_node.moe_group.id)
         paddle.base.core.nvprof_nvtx_pop()
 
         paddle.base.core.nvprof_nvtx_push("attn_forward")
         inputs = self.forward_node.attn_forward(inputs)
         paddle.base.core.nvprof_nvtx_pop()
-        attn_compute_event = deep_ep.get_event_from_calc_stream(self.forward_node.moe_group.id)
 
+        attn_compute_event = deep_ep.get_event_from_calc_stream(self.forward_node.moe_group.id)
 
-        combine_backward_event.calc_stream_wait( self.backward_node.moe_group.id )
+        combine_backward_event.calc_stream_wait(self.backward_node.moe_group.id)
         paddle.base.core.nvprof_nvtx_push("mlp_backward_dx")
         output_grad = self.backward_node.mlp_backward(output_grad)
         paddle.base.core.nvprof_nvtx_pop()
+
         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
         )
         paddle.base.core.nvprof_nvtx_pop()
-        dispatch_forward_event = deep_ep.get_event_from_comm_stream( self.forward_node.moe_group.id )
+        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)
@@ -851,28 +860,28 @@ def forward_backward(self, inputs, output_grad):
         self.backward_node.mlp_backward_dw()
         paddle.base.core.nvprof_nvtx_pop()
 
-        dispatch_forward_event.calc_stream_wait( self.forward_node.moe_group.id)
+        dispatch_forward_event.calc_stream_wait(self.forward_node.moe_group.id)
         paddle.base.core.nvprof_nvtx_push("mlp_forward")
         inputs = self.forward_node.mlp_forward(inputs)
         paddle.base.core.nvprof_nvtx_pop()
 
         paddle.base.core.nvprof_nvtx_push("combine_forward")
         inputs = self.forward_node.combine_forward(inputs, async_finish=True)
         paddle.base.core.nvprof_nvtx_pop()
-        combine_forward_event = deep_ep.get_event_from_comm_stream( self.forward_node.moe_group.id)
-
+        combine_forward_event = deep_ep.get_event_from_comm_stream(self.forward_node.moe_group.id)
 
         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()
         paddle.base.core.nvprof_nvtx_pop()
 
         combine_forward_event.calc_stream_wait(self.forward_node.moe_group.id)
         paddle.base.core.nvprof_nvtx_push("post_process_forward")
         inputs = self.forward_node.post_process_forward(inputs)
         paddle.base.core.nvprof_nvtx_pop()
         paddle.base.core.nvprof_nvtx_pop()
-        return inputs, output_grad
+        return inputs, output_grad, event_to_wait
 
 
 def build_overlapped_nodes(forward_chunk, backward_chunk):
@@ -1579,6 +1588,7 @@ def overlapped_forward_backward(
         backward_loss_fn_node,
         backward_input_grads,
         scaler,
+        event_to_wait=None,
     ):
         if backward_loss_fn_node is not None:
             if scaler:
@@ -1595,7 +1605,9 @@ def overlapped_forward_backward(
         ) = build_overlapped_nodes(forward_chunk, backward_chunk)
         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)
+        forward_inputs, backward_input_grads = overlap_node.forward_backward(
+            forward_inputs, backward_input_grads, event_to_wait
+        )
         forward_inputs = forward_post_node.forward(forward_inputs)
         backward_input_grads = backward_post_node.backward(backward_input_grads)
 

paddlenlp/transformers/fused_a2a.py

@@ -343,10 +343,15 @@ def forward(self, x, group, handle, previous_event=None, async_finish=False):
 
         return combined_x
 
-    def backward(self, grad_output, previous_event=None, async_finish=False):
+    def backward(self, grad_output, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
         """Backward pass of fused combine."""
         out = fused_combine_backward_func(
-            grad_output, self.group, self.handle, previous_event=previous_event, async_finish=async_finish
+            grad_output,
+            self.group,
+            self.handle,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
         )
         self.reset_statue()
         return out

paddlenlp/transformers/moe_layer.py

@@ -27,11 +27,16 @@
 
 from ..utils.log import logger
 from .fp8_utils import FP8GroupGemmMlpFunctionNode
-from .fused_a2a import CombineNode, DispatchNode
+from .fused_a2a import CombineNode, DispatchNode, get_buffer, get_hidden_bytes
 from .moe_gate import PretrainedMoEGate
 from .moe_utils import UnZipNode, ZipNode
 from .token_dispatcher import MoEFlexTokenDispatcher, PreDispatchNode
 
+try:
+    import paddle.distributed.communication.deep_ep as deep_ep
+except ImportError:
+    deep_ep = None
+
 DSV3_USE_FP8_GEMM = os.getenv("DSV3_USE_FP8_GEMM", "False").lower() == "true"
 
 DSV3_USE_FP8_GROUP_GEMM = os.getenv("DSV3_USE_FP8_GROUP_GEMM", "False").lower() == "true"
@@ -622,12 +627,13 @@ def forward(self, hidden_states_out, previous_event=None, async_finish=False):
         return output_combine
 
     @paddle.no_grad()
-    def backward(self, output_combine_grad, previous_event=None, async_finish=False):
+    def backward(self, output_combine_grad, previous_event=None, async_finish=False, allocate_on_comm_stream=False):
         # combine grad -> fp8
         hidden_states_out_grad = self.combine_node.backward(
             output_combine_grad,
             previous_event=previous_event,
             async_finish=async_finish,
+            allocate_on_comm_stream=allocate_on_comm_stream,
         )
         return hidden_states_out_grad
 
@@ -647,18 +653,29 @@ def forward(self, output_combine):
         return output
 
     @paddle.no_grad()
-    def backward(self, output_grad):
+    def backward(self, output_grad, event_to_wait=None):
         # post combine grad
-        output_combine_grad = paddle.reshape(output_grad, self.output_combine_shape)
-
         if DSV3_USE_FP8_DISPATCH:
-            # 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
-            )
-            return (output_combine_grad_fp8, output_combine_grad_scale)
+            if event_to_wait is not None:
+                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 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
+                )
+                output_grad._record_stream()
+                quant_event = None
+                if event_to_wait is not 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:
-            return output_combine_grad
+            output_combine_grad = paddle.reshape(output_grad, self.output_combine_shape)
+            return output_combine_grad, None
 
 
 class FusionMlpNode: