Merge branch 'develop' of https://github.com/PaddlePaddle/PaddleFormers into merge_paddlefleet_to_paddleformers

Author: risemeup1

paddleformers/cli/train/deepseek_v3_pretrain/modeling.py

@@ -54,6 +54,20 @@
 
 from paddle import _C_ops
 
+_flash_mask_available = False
+try:
+    if paddle.cuda.is_available() and paddle.cuda.get_device_capability()[0] == 10:
+        from paddlefleet_ops.flash_mask.cute.interface import (
+            _flash_attn_bwd as _flash_attn_v4_bwd,
+        )
+        from paddlefleet_ops.flash_mask.cute.interface import (
+            _flash_attn_fwd as _flash_attn_v4_fwd,
+        )
+
+        _flash_mask_available = True
+except (ImportError, AttributeError):
+    _flash_mask_available = False
+
 try:
     from paddle.nn.functional.flash_attention import flash_attention
 except:
@@ -2990,6 +3004,17 @@ def forward(
                 False,  # pack_gqa_
                 0,  # sm_margin
             )
+        elif fa_version == 4:
+            attn_out, softmax_lse = _flash_attn_v4_fwd(
+                query_states,
+                key_states,
+                value_states,
+                softmax_scale=softmax_scale,
+                causal=True,
+                return_lse=True,
+                startend_row_indices=None,
+                pack_gqa=False,
+            )
         else:
             raise AssertionError(f"invalid {fa_version=}")
 
@@ -3064,6 +3089,55 @@ def forward(
                     softmax_scale,
                     recompute_fa3,
                 )
+        elif fa_version == 4:
+            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,
+                    custom_map.q_lens,
+                    custom_map.out_proj_weight,
+                    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,
+                    custom_map.q_lens,
+                    custom_map.out_proj_weight,
+                    kv_lora_rank,
+                    softmax_scale,
+                    recompute_fa3,
+                )
         else:
             raise AssertionError(f"invalid {fa_version=}")
 
@@ -3141,6 +3215,30 @@ def backward(ctx, dout):
                 softmax_scale,
                 recompute_fa3,
             ) = ctx.saved_tensor()
+        elif fa_version == 4:
+            (
+                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,
+                q_lens,
+                out_proj_weight,
+                kv_lora_rank,
+                softmax_scale,
+                recompute_fa3,
+            ) = ctx.saved_tensor()
         else:
             raise AssertionError(f"invalid {fa_version=}")
 
@@ -3149,6 +3247,8 @@ def backward(ctx, dout):
             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
+        if fa_version == 4 and not recompute_fa3:
+            assert attn_out is not None and softmax_lse is not None
 
         q_ln_t, q_ln_invar = paddle.incubate.nn.functional.fused_rms_norm_ext(
             q_init, q_ln_weight, eps
@@ -3328,6 +3428,71 @@ def backward(ctx, dout):
                     0.0,
                     0,
                 )
+        elif fa_version == 4:
+            # recompute fa4
+            if recompute_fa3:
+                with paddle.no_grad():
+                    attn_out, softmax_lse = _flash_attn_v4_fwd(
+                        query_states,
+                        key_states,
+                        value_states,
+                        softmax_scale=softmax_scale,
+                        causal=True,
+                        return_lse=True,
+                        startend_row_indices=None,
+                        pack_gqa=False,
+                    )
+
+                    # padding x and quant
+                    bsz = value_states.shape[0]
+                    attn_out_reshape = attn_out.reshape([bsz * q_lens, -1]).contiguous()
+                    d_attn_out_reshape_shape = attn_out_reshape.shape
+                    attn_out_reshape = FP8LinearFunctionBase.padding(attn_out_reshape, 0)
+                    (
+                        attn_out_reshape_t_fp8,
+                        attn_out_reshape_t_scale,
+                    ) = paddle.incubate.nn.functional.fp8_quant_blockwise(
+                        attn_out_reshape,
+                        output_scale_transpose=True,
+                        quant_method="1x128",
+                        input_transpose=True,
+                        return_transpose_only=True,
+                    )
+                    dout_2d = dout.reshape([-1, dout.shape[-1]])
+
+                    # ===== dx = deep_gemm(dout_fp8, w_fp8)
+                    d_attn_out, dout_t_fp8, dout_t_scale = FP8LinearFunctionBase.compute_fp8_linear(
+                        dout_2d, out_proj_weight, weight_transpose=False, return_mode="with_input_transpose_quant"
+                    )
+                    d_attn_out = d_attn_out.reshape(d_attn_out_reshape_shape)
+                    FP8LinearFunctionBase.compute_expert_w_grad(
+                        attn_out_reshape_t_fp8,
+                        attn_out_reshape_t_scale,
+                        dout_t_fp8,
+                        dout_t_scale,
+                        True,
+                        True,
+                        out_proj_weight,
+                        paddle.float32,
+                    )
+
+                    dout = d_attn_out
+                    dout = dout.reshape(attn_out.shape)
+
+            with paddle.no_grad():
+                flashmask_info = None
+                q_grad, k_grad, v_grad = _flash_attn_v4_bwd(
+                    query_states,
+                    key_states,
+                    value_states,
+                    attn_out,
+                    dout,
+                    softmax_lse,
+                    flashmask_info,
+                    softmax_scale=softmax_scale,
+                    causal=True,
+                    deterministic=bool(paddle.get_flags(["FLAGS_cudnn_deterministic"])["FLAGS_cudnn_deterministic"]),
+                )
         else:
             raise AssertionError(f"invalid {fa_version=}")
 
@@ -3511,6 +3676,9 @@ def __init__(
         super().__init__()
         self._dtype = self._helper.get_default_dtype()
 
+        if paddle.cuda.is_available() and paddle.cuda.get_device_capability()[0] >= 10 and _flash_mask_available:
+            fa_version = 4
+
         self.q_ln_weight = paddle.create_parameter(
             shape=[q_norm_hidden_size],
             dtype=self._dtype,

paddleformers/transformers/deepseek_v4/modeling.py

@@ -513,6 +513,8 @@ def _gen_aoa_config(cls, config: DeepseekV4Config):
             "embed.weight -> model.embedding.embed_tokens.weight",
             "norm.weight -> model.norm.weight",
         ]
+        if mtp_num_layers > 0:
+            stmts.append("embed.weight -> model.mtp_embedding.embed_tokens.weight")
         if config.tie_word_embeddings:
             stmts += ["embed.weight -> model.lm_head.weight"]
         else:
@@ -835,6 +837,8 @@ def _gen_inv_aoa_config(cls, config: DeepseekV4Config):
             "model.embedding.embed_tokens.weight -> embed.weight",
             "model.norm.weight -> norm.weight",
         ]
+        if mtp_num_layers > 0:
+            stmts.append("model.mtp_embedding.embed_tokens.weight -> embed.weight")
         if config.tie_word_embeddings:
             stmts += ["model.lm_head.weight -> _"]
         else: