Add stepped O1 recompute (#11010)

* add_stepped_rc

* polish code

* remove stepevent

* fix bug

Author: liuruyan

paddlenlp/transformers/deepseek_v2/configuration.py

@@ -182,6 +182,7 @@ def __init__(
         use_dualpipev=False,
         send_mtp_embed=False,
         using_post_norm_recompute=False,
+        stepped_recompute_fwd_gate_up=False,
         recompute_fwd_gate_up=0,
         recompute_fa3=0,
         is_split_group_gemm=False,
@@ -245,6 +246,7 @@ def __init__(
         self.using_post_norm_recompute = using_post_norm_recompute
         self.recompute_fwd_gate_up = recompute_fwd_gate_up
         self.recompute_fa3 = recompute_fa3
+        self.stepped_recompute_fwd_gate_up = stepped_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

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -33,6 +33,7 @@
     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
 
@@ -598,6 +599,7 @@ def __init__(
         mlp_layer,
         send_mtp_embed,
         using_post_norm_recompute=False,
+        stepped_recompute_fwd_gate_up=False,
         name="",
     ):
         self.attn_and_gate_node = attn_and_gate_node
@@ -606,6 +608,7 @@ def __init__(
         self.send_mtp_embed = send_mtp_embed
 
         self.using_post_norm_recompute = using_post_norm_recompute
+        self.stepped_recompute_fwd_gate_up = stepped_recompute_fwd_gate_up
         self.name = name
 
         self.moe_group = mlp_layer.moe_group
@@ -1058,6 +1061,8 @@ def backward_for_fusion(self, output_grad, combine_bw_event_to_wait=None, pp_str
         return output_grad, event_to_wait
 
     def forward(self, inputs):
+        if self.stepped_recompute_fwd_gate_up:
+            self.fp8_fusion_moe_node.mlp_node.set_recompute_fwd_gate_up(True)
         inputs = self.attn_forward(inputs)
         inputs = self.dispatch_forward(inputs)
         inputs = self.mlp_forward(inputs)
@@ -1820,6 +1825,7 @@ def build_schedule_node(self):
                         mlp_layer=self.mlp,
                         send_mtp_embed=self.config.send_mtp_embed,
                         using_post_norm_recompute=self.config.using_post_norm_recompute,
+                        stepped_recompute_fwd_gate_up=self.config.stepped_recompute_fwd_gate_up,
                         name="FusionFp8DecoderLayerNode",
                     )
                 else: