support bw split (#10915)

* support bw split

* fix bug

* polish code

Author: phlrain

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -80,6 +80,8 @@
 
 FA_VERSION = int(os.getenv("FA_VERSION", 2))
 
+from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import WeightGradStore
+
 from ..fp8_utils import FP8KeepXLinear, FP8Linear, FP8Mlp
 from .fp8_linear import Linear
 
@@ -1084,7 +1086,8 @@ def qkv_pre_process(
     target_key_value_shape = [0, 0, num_heads, qk_nope_head_dim + v_head_dim]
 
     q = q.reshape(shape=target_query_shape)
-    q_nope, q_pe = paddle.split(q, [qk_nope_head_dim, qk_rope_head_dim], axis=-1)
+    q_nope = q[..., :qk_nope_head_dim]
+    q_pe = q[..., qk_nope_head_dim:]
 
     # DeepSeekV2 kv_lora_rank+qk_rope_head_dim=512+64
 
@@ -1094,8 +1097,9 @@ def qkv_pre_process(
 
     # self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim = 128+64
     # self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim) = config.qk_nope_head_dim + self.v_head_dim = 128+128
+    k_nope = kv[..., :qk_nope_head_dim]
+    value_states = kv[..., qk_nope_head_dim:]
 
-    k_nope, value_states = paddle.split(kv, [qk_nope_head_dim, v_head_dim], axis=-1)
     kv_seq_len = value_states.shape[1]
 
     cos, sin = rotary_emb(value_states, seq_len=kv_seq_len)
@@ -1434,15 +1438,53 @@ def backward(ctx, dout):
         paddle.base.core._set_has_grad(False)
 
         # call up proj
-        d_kv_ln_t, d_kv_up_weight = _C_ops.matmul_grad(kv_ln_t, kv_up_weight, d_kv, False, False)
+        if hasattr(kv_up_weight, "main_grad"):
+            d_kv_ln_t = paddle.matmul(d_kv, kv_up_weight, transpose_y=True)
+
+            def kv_up_weight_grad(kv_ln_t, d_kv, kv_up_weight):
+               
+                with paddle.no_grad():
+                    w_grad_t = paddle.matmul( kv_ln_t.reshape([-1, kv_ln_t.shape[-1]]), d_kv.reshape([-1, d_kv.shape[-1]]), transpose_x=True)
+
+                    kv_up_weight.main_grad.add_( w_grad_t )
+
+            if WeightGradStore.enabled:
+            
+                WeightGradStore.put(partial(kv_up_weight_grad, kv_ln_t, d_kv, kv_up_weight))
+            else:
+                kv_up_weight_grad(kv_ln_t, d_kv, kv_up_weight)
+
+            d_kv_up_weight = None
+
+        else:
+            d_kv_ln_t, d_kv_up_weight = _C_ops.matmul_grad(kv_ln_t, kv_up_weight, d_kv, False, False)
+
 
         d_compressed_kv, d_kv_ln_weight = fused_ln.fused_rms_norm_grad_func(
             compressed_kv, kv_ln_weight, kv_ln_invar, d_kv_ln_t, eps
         )
 
         d_kv_init = paddle.concat([d_compressed_kv, d_k_pe], axis=-1)
 
-        d_q_ln_t, d_q_up_weight = _C_ops.matmul_grad(q_ln_t, q_up_weight, d_q, False, False)
+        if hasattr(q_up_weight, "main_grad"):
+            d_q_ln_t = paddle.matmul(d_q, q_up_weight, transpose_y=True)
+
+            def q_up_weight_grad(q_ln_t, d_q, q_up_weight):                
+
+                with paddle.no_grad():
+                    w_grad_t = paddle.matmul( q_ln_t.reshape([-1, q_ln_t.shape[-1]]), d_q.reshape([-1, d_q.shape[-1]]), transpose_x=True)
+                    q_up_weight.main_grad.add_( w_grad_t )
+
+            if WeightGradStore.enabled:            
+                WeightGradStore.put(partial(q_up_weight_grad, q_ln_t, d_q, q_up_weight))
+            else:
+                q_up_weight_grad(q_ln_t, d_q, q_up_weight)
+
+            d_q_up_weight = None
+
+        else:
+            d_q_ln_t, d_q_up_weight = _C_ops.matmul_grad(q_ln_t, q_up_weight, d_q, False, False)
+
         d_q_init, d_q_ln_weight = fused_ln.fused_rms_norm_grad_func(q_init, q_ln_weight, q_ln_invar, d_q_ln_t, eps)
 
         return d_q_init, d_kv_init, d_q_ln_weight, d_kv_ln_weight, d_q_up_weight, d_kv_up_weight

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -27,6 +27,14 @@
     ScheduleNode,
     SharedLayerDesc,
 )
+from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import (
+    WeightGradStore
+)
+
+try:
+    from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import EventStore
+except ImportError:
+    EventStore = None
 from paddle.distributed.fleet.recompute.recompute import recompute
 from paddle.distributed.fleet.utils.sequence_parallel_utils import ScatterOp
 
@@ -714,7 +722,10 @@ def combine_backward(self, output_grad, previous_event=None, async_finish=False,
         return ret
 
     def mlp_backward_dw(self):
-        self.fp8_fusion_moe_node.mlp_node.backward_dw()
+        if WeightGradStore.enabled:
+            WeightGradStore.put(self.fp8_fusion_moe_node.mlp_node.backward_dw)
+        else:
+            self.fp8_fusion_moe_node.mlp_node.backward_dw()
 
     def mlp_backward(self, output_grad):
         if self.send_mtp_embed:
@@ -914,8 +925,18 @@ def forward_backward(self, inputs, output_grad, combine_bw_event_to_wait=None, p
 
         paddle.base.core.nvprof_nvtx_pop()
         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)
+
+        WeightGradStore.enabled = False
+        WeightGradStore.flush()
+        WeightGradStore.pop()
+        assert WeightGradStore.funcs_queue.empty()
 
         paddle.base.core.nvprof_nvtx_pop()
 

paddlenlp/transformers/fp8_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from functools import partial
 
 import numpy
 import paddle
@@ -27,6 +28,8 @@ def swiglu(x, y=None):
             x, y = paddle.chunk(x, chunks=2, axis=-1)
         return F.silu(x) * y
 
+from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import WeightGradStore
+
 
 USE_DS_GEMM = os.getenv("USE_DS_GEMM", "False").lower() == "true"
 
@@ -239,16 +242,31 @@ def compute_expert_w_grad(
         if hasattr(weight, "main_grad"):
             if weight.main_grad is None:
                 weight.main_grad = paddle.zeros(shape=weight.shape, dtype=paddle.float32)
-            result = FP8LinearFunctionBase.kitchen_gemm(
-                input_t,
-                input_t_scale,
-                dout_t,
-                dout_t_scale,
-                is_a_1d_scaled,
-                is_b_1d_scaled,
-                weight.main_grad,
-                rtn_dtype,
-            )
+            if WeightGradStore.enabled:
+                WeightGradStore.put(
+                    partial(FP8LinearFunctionBase.kitchen_gemm,
+                    input_t,
+                    input_t_scale,
+                    dout_t,
+                    dout_t_scale,
+                    is_a_1d_scaled,
+                    is_b_1d_scaled,
+                    weight.main_grad,
+                    rtn_dtype, )
+                )
+                result = None
+
+            else:
+                result = FP8LinearFunctionBase.kitchen_gemm(
+                    input_t,
+                    input_t_scale,
+                    dout_t,
+                    dout_t_scale,
+                    is_a_1d_scaled,
+                    is_b_1d_scaled,
+                    weight.main_grad,
+                    rtn_dtype,
+                )
         else:
             if weight.grad is None:
                 weight.grad = paddle.zeros(shape=weight.shape, dtype=paddle.float32)
@@ -288,9 +306,17 @@ def common_fp8_mlp_bwd(do3, x_fp8, x_scale, x_t_fp8, x_t_scale, w1, w2, apply_ba
             o2, output_scale_transpose=True, quant_method="1x128", input_transpose=True, return_transpose_only=True
         )
         if apply_backward_hook:
-            FP8LinearFunctionBase.compute_expert_w_grad(
-                o2_t_fp8, o2_t_scale, do3_t_fp8, do3_t_scale, True, True, w2, rtn_dtype=paddle.float32
-            )
+            if WeightGradStore.enabled:
+                WeightGradStore.put(
+                    partial(
+                        FP8LinearFunctionBase.compute_expert_w_grad,
+                    o2_t_fp8, o2_t_scale, do3_t_fp8, do3_t_scale, True, True, w2, rtn_dtype=paddle.float32 )
+                )
+            else:
+
+                FP8LinearFunctionBase.compute_expert_w_grad(
+                    o2_t_fp8, o2_t_scale, do3_t_fp8, do3_t_scale, True, True, w2, rtn_dtype=paddle.float32
+                )
         else:
             dw2 = FP8LinearFunctionBase.kitchen_gemm(
                 o2_t_fp8, o2_t_scale, do3_t_fp8, do3_t_scale, True, True, rtn_dtype=paddle.float32
@@ -306,9 +332,17 @@ def common_fp8_mlp_bwd(do3, x_fp8, x_scale, x_t_fp8, x_t_scale, w1, w2, apply_ba
 
         # ===== dw1 = deep_gemm(x_t_fp8, do1_t_fp8) =====
         if apply_backward_hook:
-            FP8LinearFunctionBase.compute_expert_w_grad(
-                x_t_fp8, x_t_scale, do1_t_fp8, do1_t_scale, True, True, w1, rtn_dtype=paddle.float32
-            )
+            if WeightGradStore.enabled:
+                WeightGradStore.put(
+                    partial(
+                        FP8LinearFunctionBase.compute_expert_w_grad,
+                    x_t_fp8, x_t_scale, do1_t_fp8, do1_t_scale, True, True, w1, rtn_dtype=paddle.float32)
+                )
+            
+            else:
+                FP8LinearFunctionBase.compute_expert_w_grad(
+                    x_t_fp8, x_t_scale, do1_t_fp8, do1_t_scale, True, True, w1, rtn_dtype=paddle.float32
+                )
         else:
             dw1 = FP8LinearFunctionBase.kitchen_gemm(
                 x_t_fp8, x_t_scale, do1_t_fp8, do1_t_scale, True, True, rtn_dtype=paddle.float32