Add sequence and output dimension subbatch (#10947)

* add subbatch

* remove backward_dw

* deal 0 size

* fix subbatch bug

* fix cuda 700

* polish code

* fix nan

Author: liuruyan

paddlenlp/transformers/deepseek_v2/configuration.py

@@ -188,6 +188,9 @@ def __init__(
         adaptive_remained_O1_recompute_ratio=0,
         offline_quant_expert_weight=True,
         clear_origin_weight_when_offline_quant=True,
+        mlp_bwd_subbatch_rows=0,
+        mlp_fwd_subbatch_rows=0,
+        output_subbatch_rows=0,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -245,6 +248,9 @@ def __init__(
         self.adaptive_remained_O1_recompute_ratio = adaptive_remained_O1_recompute_ratio
         self.offline_quant_expert_weight = offline_quant_expert_weight
         self.clear_origin_weight_when_offline_quant = clear_origin_weight_when_offline_quant
+        self.mlp_bwd_subbatch_rows = mlp_bwd_subbatch_rows
+        self.mlp_fwd_subbatch_rows = mlp_fwd_subbatch_rows
+        self.output_subbatch_rows = output_subbatch_rows
 
         super().__init__(
             pad_token_id=pad_token_id,

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -1641,15 +1641,18 @@ def build_schedule_node(self):
                 if DSV3_USE_FP8_GEMM:
                     attn_and_gate_node = ScheduleNode(self.attn_compute_for_fusion, name="attn_and_gate_node")
 
+                    # recompute_fwd_gate_up_ may be 1, 0 or -1, 1 means recompute, 0 means disable recompute, -1 means adaptive recompute.
                     recompute_fwd_gate_up_ = 1 if self.layer_idx in self.config.recompute_fwd_gate_up_list else 0
-                    recompute_fwd_gate_up_ = (
-                        -1 if self.config.adaptive_remained_O1_recompute_ratio else recompute_fwd_gate_up_
-                    )
+                    if recompute_fwd_gate_up_ == 0 and self.config.adaptive_remained_O1_recompute_ratio:
+                        recompute_fwd_gate_up_ = -1
 
                     fp8_fusion_moe_node = FusionMoeNode(
                         self.mlp,
                         recompute_fwd_gate_up=recompute_fwd_gate_up_,
                         is_split_group_gemm=self.config.is_split_group_gemm,
+                        mlp_fwd_subbatch_rows=self.config.mlp_fwd_subbatch_rows,
+                        mlp_bwd_subbatch_rows=self.config.mlp_bwd_subbatch_rows,
+                        output_subbatch_rows=self.config.output_subbatch_rows,
                         name="fp8_fusion_moe_node",
                     )
                     post_process_node = PostProcessNode(