Skip to content

Commit 457e217

Browse files
WaynezeeWaynezee
authored
fix dsv3_cli (#4706)
1 parent f74155f commit 457e217

1 file changed

Lines changed: 168 additions & 0 deletions

File tree

  • paddleformers/cli/train/deepseek_v3_pretrain

‎paddleformers/cli/train/deepseek_v3_pretrain/modeling.py‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,20 @@
5555

5656
from paddle import _C_ops
5757

58+
_flash_mask_available = False
59+
try:
60+
if paddle.cuda.is_available() and paddle.cuda.get_device_capability()[0] == 10:
61+
from paddlefleet_ops.flash_mask.cute.interface import (
62+
_flash_attn_bwd as _flash_attn_v4_bwd,
63+
)
64+
from paddlefleet_ops.flash_mask.cute.interface import (
65+
_flash_attn_fwd as _flash_attn_v4_fwd,
66+
)
67+
68+
_flash_mask_available = True
69+
except (ImportError, AttributeError):
70+
_flash_mask_available = False
71+
5872
try:
5973
from paddle.nn.functional.flash_attention import flash_attention
6074
except:
@@ -2418,6 +2432,17 @@ def forward(
24182432
False, # pack_gqa_
24192433
0, # sm_margin
24202434
)
2435+
elif fa_version == 4:
2436+
attn_out, softmax_lse = _flash_attn_v4_fwd(
2437+
query_states,
2438+
key_states,
2439+
value_states,
2440+
softmax_scale=softmax_scale,
2441+
causal=True,
2442+
return_lse=True,
2443+
startend_row_indices=None,
2444+
pack_gqa=False,
2445+
)
24212446
else:
24222447
assert False, f"invalid {fa_version=}"
24232448

@@ -2492,6 +2517,55 @@ def forward(
24922517
softmax_scale,
24932518
recompute_fa3,
24942519
)
2520+
elif fa_version == 4:
2521+
if recompute_fa3:
2522+
ctx.save_for_backward(
2523+
q_init,
2524+
kv_init,
2525+
None,
2526+
None,
2527+
q_ln_weight,
2528+
kv_ln_weight,
2529+
q_up_weight,
2530+
kv_up_weight,
2531+
rotary_emb,
2532+
num_heads,
2533+
q_head_dim,
2534+
qk_nope_head_dim,
2535+
v_head_dim,
2536+
qk_rope_head_dim,
2537+
position_ids,
2538+
eps,
2539+
custom_map.q_lens,
2540+
custom_map.out_proj_weight,
2541+
kv_lora_rank,
2542+
softmax_scale,
2543+
recompute_fa3,
2544+
)
2545+
else:
2546+
ctx.save_for_backward(
2547+
q_init,
2548+
kv_init,
2549+
attn_out,
2550+
softmax_lse,
2551+
q_ln_weight,
2552+
kv_ln_weight,
2553+
q_up_weight,
2554+
kv_up_weight,
2555+
rotary_emb,
2556+
num_heads,
2557+
q_head_dim,
2558+
qk_nope_head_dim,
2559+
v_head_dim,
2560+
qk_rope_head_dim,
2561+
position_ids,
2562+
eps,
2563+
custom_map.q_lens,
2564+
custom_map.out_proj_weight,
2565+
kv_lora_rank,
2566+
softmax_scale,
2567+
recompute_fa3,
2568+
)
24952569
else:
24962570
assert False, f"invalid {fa_version=}"
24972571

@@ -2561,6 +2635,30 @@ def backward(ctx, dout):
25612635
softmax_scale,
25622636
recompute_fa3,
25632637
) = ctx.saved_tensor()
2638+
elif fa_version == 4:
2639+
(
2640+
q_init,
2641+
kv_init,
2642+
attn_out,
2643+
softmax_lse,
2644+
q_ln_weight,
2645+
kv_ln_weight,
2646+
q_up_weight,
2647+
kv_up_weight,
2648+
rotary_emb,
2649+
num_heads,
2650+
q_head_dim,
2651+
qk_nope_head_dim,
2652+
v_head_dim,
2653+
qk_rope_head_dim,
2654+
position_ids,
2655+
eps,
2656+
q_lens,
2657+
out_proj_weight,
2658+
kv_lora_rank,
2659+
softmax_scale,
2660+
recompute_fa3,
2661+
) = ctx.saved_tensor()
25642662
else:
25652663
assert False, f"invalid {fa_version=}"
25662664

@@ -2569,6 +2667,8 @@ def backward(ctx, dout):
25692667
assert attn_out is not None and softmax_lse is not None
25702668
if fa_version == 3 and not recompute_fa3:
25712669
assert attn_out is not None and softmax_lse is not None
2670+
if fa_version == 4 and not recompute_fa3:
2671+
assert attn_out is not None and softmax_lse is not None
25722672

25732673
q_ln_t, q_ln_invar = paddle.incubate.nn.functional.fused_rms_norm_ext(q_init, q_ln_weight, eps)
25742674

@@ -2722,6 +2822,71 @@ def backward(ctx, dout):
27222822
0.0,
27232823
0,
27242824
)
2825+
elif fa_version == 4:
2826+
# recompute fa4
2827+
if recompute_fa3:
2828+
with paddle.no_grad():
2829+
attn_out, softmax_lse = _flash_attn_v4_fwd(
2830+
query_states,
2831+
key_states,
2832+
value_states,
2833+
softmax_scale=softmax_scale,
2834+
causal=True,
2835+
return_lse=True,
2836+
startend_row_indices=None,
2837+
pack_gqa=False,
2838+
)
2839+
2840+
# padding x and quant
2841+
bsz = value_states.shape[0]
2842+
attn_out_reshape = attn_out.reshape([bsz * q_lens, -1]).contiguous()
2843+
d_attn_out_reshape_shape = attn_out_reshape.shape
2844+
attn_out_reshape = FP8LinearFunctionBase.padding(attn_out_reshape, 0)
2845+
(
2846+
attn_out_reshape_t_fp8,
2847+
attn_out_reshape_t_scale,
2848+
) = paddle.incubate.nn.functional.fp8_quant_blockwise(
2849+
attn_out_reshape,
2850+
output_scale_transpose=True,
2851+
quant_method="1x128",
2852+
input_transpose=True,
2853+
return_transpose_only=True,
2854+
)
2855+
dout_2d = dout.reshape([-1, dout.shape[-1]])
2856+
2857+
# ===== dx = deep_gemm(dout_fp8, w_fp8)
2858+
d_attn_out, dout_t_fp8, dout_t_scale = FP8LinearFunctionBase.compute_fp8_linear(
2859+
dout_2d, out_proj_weight, weight_transpose=False, return_mode="with_input_transpose_quant"
2860+
)
2861+
d_attn_out = d_attn_out.reshape(d_attn_out_reshape_shape)
2862+
FP8LinearFunctionBase.compute_expert_w_grad(
2863+
attn_out_reshape_t_fp8,
2864+
attn_out_reshape_t_scale,
2865+
dout_t_fp8,
2866+
dout_t_scale,
2867+
True,
2868+
True,
2869+
out_proj_weight,
2870+
paddle.float32,
2871+
)
2872+
2873+
dout = d_attn_out
2874+
dout = dout.reshape(attn_out.shape)
2875+
2876+
with paddle.no_grad():
2877+
flashmask_info = None
2878+
q_grad, k_grad, v_grad = _flash_attn_v4_bwd(
2879+
query_states,
2880+
key_states,
2881+
value_states,
2882+
attn_out,
2883+
dout,
2884+
softmax_lse,
2885+
flashmask_info,
2886+
softmax_scale=softmax_scale,
2887+
causal=True,
2888+
deterministic=bool(paddle.get_flags(["FLAGS_cudnn_deterministic"])["FLAGS_cudnn_deterministic"]),
2889+
)
27252890
else:
27262891
assert False, f"invalid {fa_version=}"
27272892

@@ -2852,6 +3017,9 @@ def __init__(
28523017
super().__init__()
28533018
self._dtype = self._helper.get_default_dtype()
28543019

3020+
if paddle.cuda.is_available() and paddle.cuda.get_device_capability()[0] >= 10 and _flash_mask_available:
3021+
fa_version = 4
3022+
28553023
self.q_ln_weight = paddle.create_parameter(
28563024
shape=[q_norm_hidden_size],
28573025
dtype=self._dtype,

0 commit comments

Comments
 (0)