Fa3 recompute (#10994)

* fix

* add fa3 rc

---------

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

Author: liuruyan

paddlenlp/transformers/deepseek_v2/configuration.py

@@ -183,6 +183,7 @@ def __init__(
         send_mtp_embed=False,
         using_post_norm_recompute=False,
         recompute_fwd_gate_up=0,
+        recompute_fa3=0,
         is_split_group_gemm=False,
         fakse_gate_restrict_balance=False,
         adaptive_remained_O1_recompute_ratio=0,
@@ -243,6 +244,7 @@ def __init__(
         self.send_mtp_embed = send_mtp_embed
         self.using_post_norm_recompute = using_post_norm_recompute
         self.recompute_fwd_gate_up = recompute_fwd_gate_up
+        self.recompute_fa3 = recompute_fa3
         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.py

@@ -1334,6 +1334,7 @@ def forward(
         eps,
         kv_lora_rank,
         softmax_scale,
+        recompute_fa3=False,
     ):
 
         bsz = q_init.shape[0]
@@ -1439,26 +1440,50 @@ def forward(
                 softmax_scale,
             )
         elif FA_VERSION == 3:
-            ctx.save_for_backward(
-                q_init,
-                kv_init,
-                attn_out,
-                softmax_lse,
-                q_ln_weight,
-                kv_ln_weight,
-                q_up_weight,
-                kv_up_weight,
-                rotary_emb,
-                num_heads,
-                q_head_dim,
-                qk_nope_head_dim,
-                v_head_dim,
-                qk_rope_head_dim,
-                position_ids,
-                eps,
-                kv_lora_rank,
-                softmax_scale,
-            )
+            if recompute_fa3:
+                ctx.save_for_backward(
+                    q_init,
+                    kv_init,
+                    None,
+                    None,
+                    q_ln_weight,
+                    kv_ln_weight,
+                    q_up_weight,
+                    kv_up_weight,
+                    rotary_emb,
+                    num_heads,
+                    q_head_dim,
+                    qk_nope_head_dim,
+                    v_head_dim,
+                    qk_rope_head_dim,
+                    position_ids,
+                    eps,
+                    kv_lora_rank,
+                    softmax_scale,
+                    recompute_fa3,
+                )
+            else:
+                ctx.save_for_backward(
+                    q_init,
+                    kv_init,
+                    attn_out,
+                    softmax_lse,
+                    q_ln_weight,
+                    kv_ln_weight,
+                    q_up_weight,
+                    kv_up_weight,
+                    rotary_emb,
+                    num_heads,
+                    q_head_dim,
+                    qk_nope_head_dim,
+                    v_head_dim,
+                    qk_rope_head_dim,
+                    position_ids,
+                    eps,
+                    kv_lora_rank,
+                    softmax_scale,
+                    recompute_fa3,
+                )
         else:
             assert False, f"invalid {FA_VERSION=}"
 
@@ -1508,10 +1533,17 @@ def backward(ctx, dout):
                 eps,
                 kv_lora_rank,
                 softmax_scale,
+                recompute_fa3,
             ) = ctx.saved_tensor()
         else:
             assert False, f"invalid {FA_VERSION=}"
 
+        if FA_VERSION == 2:
+            assert not recompute_fa3
+            assert attn_out is not None and softmax_lse is not None
+        if FA_VERSION == 3 and not recompute_fa3:
+            assert attn_out is not None and softmax_lse is not None
+
         q_ln_t, q_ln_invar = fused_ln.fused_rms_norm(q_init, q_ln_weight, eps)
 
         q_ln_fp8, q_ln_scale, q_ln_trans_fp8, q_ln_trans_scale = paddle.incubate.nn.functional.fp8_quant_blockwise(
@@ -1591,6 +1623,27 @@ def backward(ctx, dout):
                 v_grad = v_grad[..., :v_head_dim]
                 q_grad = q_grad * softmax_scale
         elif FA_VERSION == 3:
+            # recompute fa3
+            if recompute_fa3:
+                logger.info("Enable fa3 recomputation")
+                attn_out, softmax_lse = _C_ops.flash_attn_v3(
+                    query_states,
+                    key_states,
+                    value_states,
+                    None,  # q_v_
+                    None,  # q_descale_
+                    None,  # k_descale_
+                    None,  # v_descale_
+                    softmax_scale,
+                    True,
+                    -1,  # window_size_left
+                    -1,  # window_size_right
+                    0.0,  # softcap
+                    1,  # num_splits
+                    False,  # manual_set_pack_gqa
+                    False,  # pack_gqa_
+                    0,  # sm_margin
+                )
             with paddle.no_grad():
                 q_grad, k_grad, v_grad = _C_ops.flash_attn_v3_grad(
                     query_states,
@@ -1728,6 +1781,7 @@ def __init__(
         eps,
         kv_lora_rank,
         softmax_scale,
+        recompute_fa3=False,
     ) -> None:
         super().__init__()
         self._dtype = self._helper.get_default_dtype()
@@ -1764,6 +1818,7 @@ def __init__(
             self.eps,
             self.kv_lora_rank,
             self.softmax_scale,
+            self.recompute_fa3,
         ) = (
             rotary_emb,
             num_heads,
@@ -1774,6 +1829,7 @@ def __init__(
             eps,
             kv_lora_rank,
             softmax_scale,
+            recompute_fa3,
         )
         set_parameter_color([self.q_up_weight, self.kv_up_weight], "memory_attn")
 
@@ -1805,6 +1861,7 @@ def forward(self, q_init, kv_init, position_ids):
             self.eps,
             self.kv_lora_rank,
             self.softmax_scale,
+            recompute_fa3=self.recompute_fa3,
         )
 
 
@@ -1962,7 +2019,7 @@ def forward(self, x):
 class DeepseekV2Attention(nn.Layer):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
-    def __init__(self, config: DeepseekV2Config, layerwise_recompute: bool = False):
+    def __init__(self, config: DeepseekV2Config, layerwise_recompute: bool = False, recompute_fa3: bool = False):
         super().__init__()
         self.config = config
         self.attention_dropout = config.attention_dropout
@@ -1987,6 +2044,8 @@ def __init__(self, config: DeepseekV2Config, layerwise_recompute: bool = False):
             self.seq_length = config.seq_length
         self.sequence_parallel = config.sequence_parallel
 
+        self.recompute_fa3 = recompute_fa3
+
         self.input_layernorm = DeepseekV2RMSNorm(config)
 
         # Note that we will actually perform a recompute only if both enable_recompute and layerwise_recompute are set to True
@@ -2038,7 +2097,7 @@ def linear_dtype_gaurd():
             if DSV3_USE_ATTEN_RECOMPUTE:
                 self.fused_rms_norm_linear = FusedRMSLinear(self.hidden_size, config.q_lora_rank, config.kv_lora_rank + config.qk_rope_head_dim, 1e-6)
                 kv_up_dim = self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim)
-                self.memory_recompute_att = MemroyRecomputeAttn(config.q_lora_rank, config.kv_lora_rank, config.q_lora_rank, self.num_heads * self.q_head_dim, config.kv_lora_rank, kv_up_dim, self.rotary_emb, self.num_heads, self.q_head_dim, self.qk_nope_head_dim, self.v_head_dim, self.qk_rope_head_dim, 1e-6, self.kv_lora_rank, self.softmax_scale)
+                self.memory_recompute_att = MemroyRecomputeAttn(config.q_lora_rank, config.kv_lora_rank, config.q_lora_rank, self.num_heads * self.q_head_dim, config.kv_lora_rank, kv_up_dim, self.rotary_emb, self.num_heads, self.q_head_dim, self.qk_nope_head_dim, self.v_head_dim, self.qk_rope_head_dim, 1e-6, self.kv_lora_rank, self.softmax_scale, recompute_fa3=self.recompute_fa3)
                 self.o_proj = FP8KeepXLinear(self.num_heads * self.v_head_dim, self.hidden_size, bias_attr=config.attention_bias)
             else:
 
@@ -2263,7 +2322,9 @@ def forward(
 
 
 class DeepseekV2DecoderLayer(nn.Layer):
-    def __init__(self, config: DeepseekV2Config, layer_idx: int, layerwise_recompute: bool = False):
+    def __init__(
+        self, config: DeepseekV2Config, layer_idx: int, layerwise_recompute: bool = False, recompute_fa3: bool = False
+    ):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
@@ -2274,7 +2335,9 @@ def __init__(self, config: DeepseekV2Config, layer_idx: int, layerwise_recompute
 
         self.hidden_size = config.hidden_size
 
-        self.self_attn = DeepseekV2Attention(config=config, layerwise_recompute=layerwise_recompute)
+        self.self_attn = DeepseekV2Attention(
+            config=config, layerwise_recompute=layerwise_recompute, recompute_fa3=recompute_fa3
+        )
 
         DeepseekV2MLPClass = FP8Mlp if DSV3_USE_FP8_GEMM else DeepseekV2MLP
 

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -2023,6 +2023,27 @@ def compute_recompute_fwd_gate_up_list(pp_nums, all_dl_nums, dense_dl_nums, reco
                     ret.append(recompute_fwd_gate_up_list[i] + k)
             return ret
 
+        def compute_recompute_fa3_list(pp_nums, all_dl_nums, recompute_fa3):
+            all_layers_nums = all_dl_nums + 4  # embedding, rms, lm_head, mtp
+            segment_size = all_layers_nums // pp_nums
+            recompute_fa3_list = [0]
+            for idx in range(segment_size - 1, all_dl_nums, segment_size):
+                recompute_fa3_list.append(idx)
+
+            # If `recompute_fa3` is a Boolean value and is True, means all O1 will be recomputed.
+            # Otherwise `recompute_fa3` should be an integer representing how many O1 are recomputed.
+            assert isinstance(recompute_fa3, (int, bool))
+            if type(recompute_fa3) is bool:
+                enable_k_o1_rc = segment_size if recompute_fa3 is True else 0
+            else:
+                enable_k_o1_rc = recompute_fa3
+
+            ret = []
+            for i in range(len(recompute_fa3_list)):
+                for k in range(min(segment_size, enable_k_o1_rc)):
+                    ret.append(recompute_fa3_list[i] + k)
+            return ret
+
         pp_nums = (
             self.config["pipeline_parallel_degree"] * 2
             if self.config.use_dualpipev
@@ -2034,7 +2055,11 @@ def compute_recompute_fwd_gate_up_list(pp_nums, all_dl_nums, dense_dl_nums, reco
             self.config.first_k_dense_replace,
             self.config.recompute_fwd_gate_up,
         )
+        recompute_fa3_list = compute_recompute_fa3_list(
+            pp_nums, self.config.num_hidden_layers, self.config.recompute_fa3
+        )
 
+        logger.info(f"recompute_fa3_list: {recompute_fa3_list}")
         logger.info(f"recompute_fwd_gate_up_list: {recompute_fwd_gate_up_list}")
         config.recompute_fwd_gate_up_list = recompute_fwd_gate_up_list
 
@@ -2045,6 +2070,7 @@ def compute_recompute_fwd_gate_up_list(pp_nums, all_dl_nums, dense_dl_nums, reco
                     config=config,
                     layer_idx=i,
                     layerwise_recompute=i not in self.no_recompute_layers,
+                    recompute_fa3=i in recompute_fa3_list,
                 ),
                 f"{self._base_model.base_model_prefix}.layers.{i}",
             )