optimize mlp bw split (#10931)

Co-authored-by: phlrain <[email protected]>

Author: phlrain

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -27,9 +27,7 @@
     ScheduleNode,
     SharedLayerDesc,
 )
-from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import (
-    WeightGradStore
-)
+from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import WeightGradStore
 
 try:
     from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import EventStore
@@ -721,17 +719,12 @@ def combine_backward(self, output_grad, previous_event=None, async_finish=False,
         ret = (inputs_embeds_mtp_grad, *ret) if self.send_mtp_embed else ret
         return ret
 
-    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:
             inputs_embeds_mtp_grad, hidden_states_grad, residual_grad, l_aux_grad, hidden_states_out_grad = output_grad
         else:
             hidden_states_grad, residual_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, with_dw=False
-        )
+        hs_dispatched_grad, dispatched_probs_grad = self.fp8_fusion_moe_node.mlp_node.backward(hidden_states_out_grad)
 
         ret = (hidden_states_grad, residual_grad, l_aux_grad, hs_dispatched_grad, dispatched_probs_grad)
         ret = (inputs_embeds_mtp_grad, *ret) if self.send_mtp_embed else ret
@@ -809,7 +802,6 @@ def backward(self, output_grad=None, scaler=None):
         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
 
@@ -853,7 +845,11 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
 
         combine_backward_event.calc_stream_wait(self.backward_node.moe_group.id)
         paddle.base.core.nvprof_nvtx_push("mlp_backward_dx")
+        assert WeightGradStore.funcs_queue.empty()
+        WeightGradStore.enabled = True
         output_grad = self.backward_node.mlp_backward(output_grad)
+        WeightGradStore.enabled = False
+        WeightGradStore.flush()
         paddle.base.core.nvprof_nvtx_pop()
 
         output_grad_event = deep_ep.get_event_from_calc_stream(self.backward_node.moe_group.id)
@@ -889,7 +885,6 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
             final_out = self.forward_node.post_process_node.forward_without_residual(inputs)
             paddle.base.core.nvprof_nvtx_pop()
 
-
         final_out_event = deep_ep.get_event_from_calc_stream(self.forward_node.moe_group.id)
 
         paddle.base.core.nvprof_nvtx_push("combine_forward")
@@ -922,13 +917,13 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
             paddle.base.core.nvprof_nvtx_pop()
 
         dispatch_backward_event.calc_stream_wait(self.backward_node.moe_group.id)
-        
+
         paddle.base.core.nvprof_nvtx_push("attn_backward")
         assert WeightGradStore.funcs_queue.empty()
         WeightGradStore.enabled = True
         output_grad = self.backward_node.attn_backward(output_grad)
         event_to_wait = deep_ep.get_event_from_calc_stream(self.backward_node.moe_group.id)
-        
+
         if EventStore is not None:
             EventStore.set(event_to_wait)
 
@@ -937,16 +932,12 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
         WeightGradStore.pop()
         assert WeightGradStore.funcs_queue.empty()
 
-        WeightGradStore.enabled = False
-        WeightGradStore.flush()
-        WeightGradStore.pop()
-        assert WeightGradStore.funcs_queue.empty()
         paddle.base.core.nvprof_nvtx_pop()
 
         # residual add
         if pp_stream is None:
             combine_forward_event.calc_stream_wait(self.forward_node.moe_group.id)
-            
+
             final_out = self.forward_node.post_process_node.forward_without_residual(inputs)
             if final_out.shape[-1] != combine_fwd_out.shape[-1]:
                 final_out[:, :, : combine_fwd_out.shape[-1]] += combine_fwd_out  # 直接广播并相加