Skip to content

[training] fix: Update DeepSeek-V4 FLOPs calculation #4128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yaoyu-33 merged 3 commits into from
Jun 3, 2026
Merged

[training] fix: Update DeepSeek-V4 FLOPs calculation #4128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8c396ac
[training] fix: Update DeepSeek-V4 FLOPs calculation
cuichenx Jun 3, 2026
de9a3c9
[training] test: Cover DeepSeek-V4 FLOPs edge cases
cuichenx Jun 3, 2026
9a252df
[training] test: Cover DSv4 indexer validation
cuichenx Jun 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
163 changes: 130 additions & 33 deletions src/megatron/bridge/training/utils/flop_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -442,6 +442,8 @@ def transformer_flops():
# GLU: h->2*ffn_h and ffn_h->h = 3 projections; non-GLU: h->ffn_h and ffn_h->h = 2 projections.
ffn_expansion_factor = 3 if cfg.model.gated_linear_unit is True else 2

experimental_attention_variant = getattr(cfg.model, "experimental_attention_variant", None)

if cfg.model.multi_latent_attention:
"""
Basic arithmetic
Expand All @@ -456,48 +458,144 @@ def transformer_flops():
https://arxiv.org/abs/2305.10403
https://arxiv.org/abs/2205.05198
"""
## MLA
if not hasattr(cfg.model, "q_lora_rank") or cfg.model.q_lora_rank is None:
q_term = (
if experimental_attention_variant == "dsv4_hybrid":
# DeepSeek-V4 hybrid MLA uses sparse attention instead of the full
# core-attention terms used by DeepSeek-V2/V3 MLA. Projection costs
# are accounted here; sparse attention, compressor, and indexer
# costs are added below.
q_lora_rank = getattr(cfg.model, "q_lora_rank", None)
if q_lora_rank is None:
raise ValueError("q_lora_rank must be set for dsv4_hybrid FLOPs calculation")

qk_head_dim = getattr(cfg.model, "qk_head_dim", 64)
qk_pos_emb_head_dim = getattr(cfg.model, "qk_pos_emb_head_dim", 0)
v_head_dim = getattr(cfg.model, "v_head_dim", 64)
o_lora_rank = getattr(cfg.model, "o_lora_rank", 0)
o_groups = getattr(cfg.model, "o_groups", 1)

q_term = q_lora_rank * (
cfg.model.hidden_size
* cfg.model.num_attention_heads
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "qk_pos_emb_head_dim", 0))
+ cfg.model.num_attention_heads * (qk_head_dim + qk_pos_emb_head_dim)
+ 1 # q norm
)
else:
q_term = cfg.model.q_lora_rank * (
cfg.model.hidden_size
+ cfg.model.num_attention_heads
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "qk_pos_emb_head_dim", 0))
+ 1
kv_term = cfg.model.hidden_size * v_head_dim + v_head_dim # kv projection + kv norm
o_term = (
cfg.model.num_attention_heads * v_head_dim * o_lora_rank
+ o_groups * o_lora_rank * cfg.model.hidden_size
)
self_attn_term = (
3
* 2 # fwd(1) + bwd(2) *FMA
* num_layers
* (
## q lora + rope + q norm
q_term
## kv lora + rope + kv norm
+ getattr(cfg.model, "kv_lora_rank", 0)
* (
self_attn_term = 3 * 2 * num_layers * (q_term + kv_term + o_term)

compress_ratios = getattr(cfg.model, "csa_compress_ratios", None)
if compress_ratios is None:
raise ValueError("csa_compress_ratios must be set for dsv4_hybrid FLOPs calculation")
if len(compress_ratios) != num_layers:
raise ValueError(
f"Invalid length of csa_compress_ratios: {len(compress_ratios)}, "
f"expected {num_layers} "
f"(num_layers={cfg.model.num_layers}, mtp_num_layers={mtp_num_layers})."
)

supported_compress_ratios = {0, 4, 128}
unsupported_compress_ratios = [
ratio for ratio in compress_ratios if ratio not in supported_compress_ratios
]
if unsupported_compress_ratios:
raise ValueError(
"csa_compress_ratios contains unsupported values: "
f"{unsupported_compress_ratios}. Only 0, 4, and 128 are supported."
)

n_layers_r0 = sum(1 for ratio in compress_ratios if ratio == 0)
n_layers_r4 = sum(1 for ratio in compress_ratios if ratio == 4)
n_layers_r128 = sum(1 for ratio in compress_ratios if ratio == 128)
Comment thread
cuichenx marked this conversation as resolved.
Comment thread
cuichenx marked this conversation as resolved.
window = getattr(cfg.model, "csa_window_size", 128)

sparse_attn_r0 = n_layers_r0 * cfg.model.num_attention_heads * window * v_head_dim * 2
avg_comp_128 = (core_attn_seq_factor // 128) / 2
sparse_attn_r128 = (
n_layers_r128 * cfg.model.num_attention_heads * (window + avg_comp_128) * v_head_dim * 2
)

main_compressor_term = (
n_layers_r4 * cfg.model.hidden_size * (2 * v_head_dim) * 2
+ n_layers_r128 * cfg.model.hidden_size * v_head_dim * 2
)

if n_layers_r4 > 0:
idx_n_heads = getattr(cfg.model, "dsa_indexer_n_heads", None)
idx_head_dim = getattr(cfg.model, "dsa_indexer_head_dim", None)
idx_topk = getattr(cfg.model, "dsa_indexer_topk", None)
if idx_n_heads is None:
raise ValueError("dsa_indexer_n_heads must be set for dsv4_hybrid ratio==4 layers")
if idx_head_dim is None:
raise ValueError("dsa_indexer_head_dim must be set for dsv4_hybrid ratio==4 layers")
if idx_topk is None:
raise ValueError("dsa_indexer_topk must be set for dsv4_hybrid ratio==4 layers")

effective_topk_4 = min(idx_topk, core_attn_seq_factor // 4)
avg_comp_4 = effective_topk_4 * (1 - effective_topk_4 * 4 / (2 * core_attn_seq_factor))
sparse_attn_r4 = (
n_layers_r4 * cfg.model.num_attention_heads * (window + avg_comp_4) * v_head_dim * 2
)
indexer_term = (
n_layers_r4 * cfg.model.hidden_size * (2 * idx_head_dim) * 2
+ n_layers_r4 * q_lora_rank * idx_n_heads * idx_head_dim
+ n_layers_r4 * cfg.model.hidden_size * idx_n_heads
+ n_layers_r4 * idx_n_heads * idx_head_dim * (core_attn_seq_factor // 4)
)
else:
sparse_attn_r4 = 0
indexer_term = 0

sparse_attn_term = sparse_attn_r0 + sparse_attn_r4 + sparse_attn_r128
self_attn_term += 3 * 2 * (sparse_attn_term + main_compressor_term + indexer_term)
else:
## MLA
if not hasattr(cfg.model, "q_lora_rank") or cfg.model.q_lora_rank is None:
q_term = (
cfg.model.hidden_size
* cfg.model.num_attention_heads
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "qk_pos_emb_head_dim", 0))
)
else:
q_term = cfg.model.q_lora_rank * (
cfg.model.hidden_size
+ cfg.model.num_attention_heads
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "v_head_dim", 64))
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "qk_pos_emb_head_dim", 0))
+ 1
)
+ cfg.model.hidden_size * getattr(cfg.model, "qk_pos_emb_head_dim", 0)
## o proj
+ (cfg.model.num_attention_heads * getattr(cfg.model, "v_head_dim", 64)) * cfg.model.hidden_size
## core attn
+ core_attn_seq_factor
self_attn_term = (
3
* 2 # fwd(1) + bwd(2) *FMA
* num_layers
* (
cfg.model.num_attention_heads
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "qk_pos_emb_head_dim", 0))
## q lora + rope + q norm
q_term
## kv lora + rope + kv norm
+ getattr(cfg.model, "kv_lora_rank", 0)
* (
cfg.model.hidden_size
+ cfg.model.num_attention_heads
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "v_head_dim", 64))
+ 1
)
+ cfg.model.hidden_size * getattr(cfg.model, "qk_pos_emb_head_dim", 0)
## o proj
+ (cfg.model.num_attention_heads * getattr(cfg.model, "v_head_dim", 64))
* cfg.model.hidden_size
## core attn
+ core_attn_seq_factor
* (
cfg.model.num_attention_heads
* (getattr(cfg.model, "qk_head_dim", 64) + getattr(cfg.model, "qk_pos_emb_head_dim", 0))
)
/ 2
+ core_attn_seq_factor
* cfg.model.num_attention_heads
* getattr(cfg.model, "v_head_dim", 64)
/ 2
)
/ 2
+ core_attn_seq_factor * cfg.model.num_attention_heads * getattr(cfg.model, "v_head_dim", 64) / 2
)
)

else:
## MHA or GQA
Expand Down Expand Up @@ -563,7 +661,6 @@ def transformer_flops():
# When experimental_attention_variant is "gated_delta_net", a fraction of the
# layers use GDN instead of standard attention. Override self_attn_term with a
# weighted sum of GDN and standard-attention per-layer costs.
experimental_attention_variant = getattr(cfg.model, "experimental_attention_variant", None)
if experimental_attention_variant == "gated_delta_net":
linear_attention_freq = cfg.model.linear_attention_freq
if linear_attention_freq is None:
Expand Down
Loading
Loading