support bw split (#10823)

Author: phlrain

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -665,6 +665,9 @@ def combine_backward(self, output_grad, async_finish=False):
                 hidden_states_out_grad,
             )
 
+    def mlp_backward_dw(self):
+        self.fp8_fusion_moe_node.mlp_node.backward_dw()
+
     def mlp_backward(self, output_grad):
         if self.send_mtp_embed:
             (
@@ -681,7 +684,9 @@ def mlp_backward(self, output_grad):
                 l_aux_grad,
                 hidden_states_out_grad,
             ) = output_grad
-        hs_dispatched_grad, dispatched_probs_grad = self.fp8_fusion_moe_node.mlp_node.backward(hidden_states_out_grad)
+        hs_dispatched_grad, dispatched_probs_grad = self.fp8_fusion_moe_node.mlp_node.backward(
+            hidden_states_out_grad, with_dw=False
+        )
 
         if self.send_mtp_embed:
             return (
@@ -790,7 +795,9 @@ def backward(self, output_grad=None, scaler=None):
         output_grad = self.post_process_backward(output_grad)
         output_grad = self.combine_backward(output_grad)
         output_grad = self.mlp_backward(output_grad)
+        # todo(phlrain): overlap here
         output_grad = self.dispatch_backward(output_grad)
+        self.mlp_backward_dw()
         output_grad = self.attn_backward(output_grad)
         return output_grad
 
@@ -820,19 +827,23 @@ def forward_backward(self, inputs, output_grad):
 
         calc_stream_wait(self.backward_node.moe_group.id)
         attn_compute_event = deep_ep.get_event_from_calc_stream(self.forward_node.moe_group.id)
-        paddle.base.core.nvprof_nvtx_push("mlp_backward")
+        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()
-
-        calc_stream_wait(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)
         paddle.base.core.nvprof_nvtx_pop()
+
+        paddle.base.core.nvprof_nvtx_push("dispatch_backward_dw")
+        self.backward_node.mlp_backward_dw()
+        paddle.base.core.nvprof_nvtx_pop()
+
+        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()

paddlenlp/transformers/fp8_utils.py

@@ -784,6 +784,7 @@ def bwd_dowm_input(self, expert_w2, unzipped_grad, o1, inplace_swiglu_prob=False
         # compute gemm
         if isinstance(unzipped_grad, tuple):
             (unzipped_grad_fp8, unzipped_grad_scale) = unzipped_grad
+            unzipped_grad_scale = unzipped_grad_scale.T.contiguous().T
         else:
             unzipped_grad_fp8, unzipped_grad_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
                 unzipped_grad, output_scale_transpose=True, quant_method="1x128", input_transpose=False
@@ -1015,3 +1016,56 @@ def backward(self, out_grad):
 
         self.reset_statue()
         return dx, probs_grad
+
+    @paddle.no_grad()
+    def backward_dx(self, out_grad):
+        # recompute expert_w2 and expert_w1
+        expert_w1 = [x.w1 for x in self.experts if x is not None]
+        expert_w2 = [x.w2 for x in self.experts if x is not None]
+
+        if self.recompute_fwd_gate_up:
+            o1 = self.fwd_gate_up(None, expert_w1, len(expert_w1), self.tokens_per_expert)
+        else:
+            o1 = self.o1
+
+        # do2
+        do1, o2_s, probs_grad = self.bwd_dowm_input(expert_w2, out_grad, o1, inplace_swiglu_prob=True)
+        del o1
+        self.o1 = None
+
+        self.do1 = do1
+        self.o2_s = o2_s
+
+        self.out_grad = out_grad
+
+        # dx
+        dx = self.bwd_gate_up_input(do1, expert_w1, dx=out_grad[0] if isinstance(out_grad, tuple) else out_grad)
+
+        return dx, probs_grad
+
+    @paddle.no_grad()
+    def backward_dw(self):
+        # recompute expert_w2 and expert_w1
+        expert_w1 = [x.w1 for x in self.experts if x is not None]
+        expert_w2 = [x.w2 for x in self.experts if x is not None]
+
+        # dw1
+        self.bwd_gate_up_weight(self.do1, None, expert_w1, clear_input=True)
+        self.input_fp8 = None
+        self.input_scale = None
+        self.input = None
+        self.do1 = None
+
+        # dw2
+        if isinstance(self.out_grad, tuple):
+            out_grad_dequant_fp16 = paddle.incubate.nn.functional.fused_act_dequant(self.out_grad[0], self.out_grad[1])
+            self.out_grad = None
+            self.bwd_down_weight(out_grad_dequant_fp16, self.o2_s, expert_w2)
+            del out_grad_dequant_fp16
+        else:
+            self.bwd_down_weight(self.out_grad, self.o2_s, expert_w2)
+
+        self.o2_s = None
+
+        self.reset_statue()
+        return

paddlenlp/transformers/moe_layer.py

@@ -747,7 +747,7 @@ def forward(self, hs_2d_dispatched, dispatched_indices, dispatched_probs):
         return expert_out_zipped
 
     @paddle.no_grad()
-    def backward(self, hidden_states_out_grad):
+    def backward(self, hidden_states_out_grad, with_dw=True):
         """
         反向传播函数。
 
@@ -772,7 +772,10 @@ def backward(self, hidden_states_out_grad):
         record_stream_for_multi_input(hidden_states_out_grad)
 
         # expert_grad
-        expert_out, probs_grad = self.experts_group_gemm_node.backward(unzipped_grad)
+        if with_dw:
+            expert_out, probs_grad = self.experts_group_gemm_node.backward(unzipped_grad)
+        else:
+            expert_out, probs_grad = self.experts_group_gemm_node.backward_dx(unzipped_grad)
 
         hs_dispatched_grad, dispatched_probs_grad = self.unzip_node.backward(
             expert_out,
@@ -781,9 +784,14 @@ def backward(self, hidden_states_out_grad):
             self.dispatched_indices,
             num_experts=len(self.tokens_per_expert),
         )
-        self.reset_statue()
+        if with_dw:
+            self.reset_statue()
         return hs_dispatched_grad, dispatched_probs_grad
 
+    @paddle.no_grad()
+    def backward_dw(self):
+        self.experts_group_gemm_node.backward_dw()
+
 
 class FusionMoeNode:
     def __init__(
@@ -836,11 +844,11 @@ def forward(self, hidden_states, probs, routing_map):
             return output
 
     @paddle.no_grad()
-    def backward(self, output_grad):
+    def backward(self, output_grad, with_dw=True):
         output_combine_grad = self.combine_quant_node.backward(output_grad)
         hidden_states_out_grad = self.combine_node.backward(output_combine_grad)
 
-        hs_dispatched_grad, dispatched_probs_grad = self.mlp_node.backward(hidden_states_out_grad)
+        hs_dispatched_grad, dispatched_probs_grad = self.mlp_node.backward(hidden_states_out_grad, with_dw=with_dw)
 
         if DSV3_USE_FP8_DISPATCH:
             hs_fp8_grad, token_probs_grad = self.dispatch_node.backward(hs_dispatched_grad, dispatched_probs_grad)
@@ -850,6 +858,10 @@ def backward(self, output_grad):
             hs_bf16_grad, token_probs_grad = self.dispatch_node.backward(hs_dispatched_grad, dispatched_probs_grad)
             return hs_bf16_grad, None, token_probs_grad
 
+    @paddle.no_grad()
+    def backward_dw(self):
+        self.mlp_node.backward_dw()
+
 
 class FusionMoe(paddle.autograd.PyLayer):
     @staticmethod