Add adaptived recomptue of o1 (#10891)

* add adaptived recomptue of O1

* add adaptived recomptue ratio

* add adaptived recomptue ratio

Author: liuruyan

paddlenlp/transformers/deepseek_v2/configuration.py

@@ -185,6 +185,7 @@ def __init__(
         recompute_fwd_gate_up=0,
         is_split_group_gemm=False,
         fakse_gate_restrict_balance=False,
+        adaptive_remained_O1_recompute_ratio=0,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -239,6 +240,7 @@ def __init__(
         self.recompute_fwd_gate_up = recompute_fwd_gate_up
         self.is_split_group_gemm = is_split_group_gemm
         self.fakse_gate_restrict_balance = fakse_gate_restrict_balance
+        self.adaptive_remained_O1_recompute_ratio = adaptive_remained_O1_recompute_ratio
 
         super().__init__(
             pad_token_id=pad_token_id,

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -1189,9 +1189,15 @@ def build_schedule_node(self):
             if self.mlp.using_flex_token:
                 if DSV3_USE_FP8_GEMM:
                     attn_and_gate_node = ScheduleNode(self.attn_compute_for_fusion, name="attn_and_gate_node")
+
+                    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_
+                    )
+
                     fp8_fusion_moe_node = FusionMoeNode(
                         self.mlp,
-                        recompute_fwd_gate_up=self.config.recompute_fwd_gate_up,
+                        recompute_fwd_gate_up=recompute_fwd_gate_up_,
                         is_split_group_gemm=self.config.is_split_group_gemm,
                         name="fp8_fusion_moe_node",
                     )

paddlenlp/transformers/moe_layer.py

@@ -704,11 +704,20 @@ def __init__(self, custom_map, max_topk, recompute_fwd_gate_up=False, is_split_g
             recompute_fwd_gate_up=recompute_fwd_gate_up,
             is_split_group_gemm=is_split_group_gemm,
         )
+
+        self.seq_length = custom_map.config.seq_length
+        self.num_experts_per_tok = custom_map.config.num_experts_per_tok
+        self.adaptive_remained_O1_recompute_ratio = custom_map.config.adaptive_remained_O1_recompute_ratio
+
+        self.recompute_fwd_gate_up = recompute_fwd_gate_up
         self.dispatched_indices = None
         self.dispatched_probs = None
         self.tokens_per_expert = None
         self.router_topk = max_topk
 
+    def set_recompute_fwd_gate_up(self, recompute_fwd_gate_up):
+        self.experts_group_gemm_node.recompute_fwd_gate_up = recompute_fwd_gate_up
+
     def reset_statue(self, with_dw=False):
         """
         重置所有状态变量。
@@ -771,6 +780,18 @@ def forward(self, hs_2d_dispatched, dispatched_indices, dispatched_probs):
             dispatched_indices._record_stream()
             dispatched_probs._record_stream()
 
+            # If adaptive O1 recompute is enabled, determine whether to enable recompute O1 based on the degree of imbalance
+            if self.recompute_fwd_gate_up == -1:
+                if (
+                    unzipped_tokens.shape[0]
+                    > self.seq_length * self.num_experts_per_tok * self.adaptive_remained_O1_recompute_ratio
+                ):
+                    # logger.debug(f"recompute_fwd_gate_up changed to True, Because the receives {unzipped_tokens.shape[0]} Tensors greater then {self.seq_length*self.num_experts_per_tok*self.adaptive_remained_O1_recompute_ratio}.")
+                    self.set_recompute_fwd_gate_up(True)
+                else:
+                    # logger.debug(f"recompute_fwd_gate_up changed to False, Because the receives {unzipped_tokens.shape[0]} Tensors less then {self.seq_length*self.num_experts_per_tok*self.adaptive_remained_O1_recompute_ratio}.")
+                    self.set_recompute_fwd_gate_up(False)
+
             # 2 experts
             padding_token_per_experts = [(x + 127) // 128 * 128 for x in self.tokens_per_expert]
             expert_out = self.experts_group_gemm_node.forward(
@@ -792,6 +813,16 @@ def forward(self, hs_2d_dispatched, dispatched_indices, dispatched_probs):
             dispatched_indices._record_stream()
             dispatched_probs._record_stream()
 
+            # If adaptive O1 recompute is enabled, determine whether to enable recompute O1 based on the degree of imbalance
+            if self.recompute_fwd_gate_up == -1:
+                if (
+                    unzipped_tokens.shape[0]
+                    > self.seq_length * self.num_experts_per_tok * self.adaptive_remained_O1_recompute_ratio
+                ):
+                    self.set_recompute_fwd_gate_up(True)
+                else:
+                    self.set_recompute_fwd_gate_up(False)
+
             # 2 experts
             padding_token_per_experts = [(x + 127) // 128 * 128 for x in self.tokens_per_expert]
             expert_out = self.experts_group_gemm_node.forward(