From a2e0e4fa9d0e61ba9da8bae853e2ffa83474e86e Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Wed, 12 Aug 2026 16:47:58 -0700 Subject: [PATCH] fix(training): align DSv4 packed FLOPs with MCore Signed-off-by: Chen Cui --- .../bridge/training/utils/flop_utils.py | 21 +++-- .../training/utils/test_flop_utils.py | 91 +++++++++++++++++++ 2 files changed, 105 insertions(+), 7 deletions(-) diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index 10ffe3a69f..116b95662a 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -668,8 +668,10 @@ def num_floating_point_operations( # the result matches the legacy constant-length estimate. if seqlen_squared_sum is not None and seqlen_sum > 0: core_attn_seq_factor = seqlen_squared_sum / seqlen_sum + effective_seqlen_squared_sum = seqlen_squared_sum else: core_attn_seq_factor = effective_seq_length + effective_seqlen_squared_sum = seqlen_sum * effective_seq_length # If the model provider has a custom TFLOPS calculation method, use it (non-LoRA only). if not is_lora and hasattr(cfg.model, "_get_num_floating_point_operations"): @@ -1049,6 +1051,7 @@ def transformer_flops(): ffn_expansion_factor = 3 if cfg.model.gated_linear_unit is True else 2 experimental_attention_variant = getattr(cfg.model, "experimental_attention_variant", None) + dsv4_hybrid_core_attn_term = 0 if cfg.model.multi_latent_attention: """ @@ -1117,10 +1120,9 @@ def transformer_flops(): 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 - ) + # Window work is token-linear; compressed-KV attention scales with sum_i(sequence_length_i^2). + sparse_attn_r128 = n_layers_r128 * cfg.model.num_attention_heads * window * v_head_dim * 2 + sparse_attn_r128_core = n_layers_r128 * cfg.model.num_attention_heads * v_head_dim / 128 main_compressor_term = ( n_layers_r4 * cfg.model.hidden_size * (2 * v_head_dim) * 2 @@ -1138,8 +1140,9 @@ def transformer_flops(): 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)) + # Match MCore's nominal ratio-4 selection estimate, which uses the configured sequence length. + effective_topk_4 = min(idx_topk, cfg.model.seq_length // 4) + avg_comp_4 = effective_topk_4 * (1 - effective_topk_4 * 4 / (2 * cfg.model.seq_length)) sparse_attn_r4 = ( n_layers_r4 * cfg.model.num_attention_heads * (window + avg_comp_4) * v_head_dim * 2 ) @@ -1147,14 +1150,17 @@ def transformer_flops(): 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) ) + # Dense indexer scoring is quadratic and therefore uses the runtime squared-length sum below. + indexer_scoring_core = n_layers_r4 * idx_n_heads * idx_head_dim / 4 else: sparse_attn_r4 = 0 indexer_term = 0 + indexer_scoring_core = 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) + dsv4_hybrid_core_attn_term = 3 * 2 * (sparse_attn_r128_core + indexer_scoring_core) elif experimental_attention_variant == "dsa": # DSA replaces dense MLA core attention with top-k attention while retaining a # dense lightning indexer. The attention/indexer geometry follows equations 1-2 @@ -1556,6 +1562,7 @@ def count_indexer_layers(layer_count: int) -> int: # Logit. + 3 * 2 * cfg.model.hidden_size * padded_vocab_size * (mtp_num_layers + 1) ) + total_floating_point_operations += effective_seqlen_squared_sum * dsv4_hybrid_core_attn_term return total_floating_point_operations + _compute_vit_flops() def _compute_vit_flops(): diff --git a/tests/unit_tests/training/utils/test_flop_utils.py b/tests/unit_tests/training/utils/test_flop_utils.py index a9ee3878b4..78f2d8f59d 100644 --- a/tests/unit_tests/training/utils/test_flop_utils.py +++ b/tests/unit_tests/training/utils/test_flop_utils.py @@ -1057,6 +1057,97 @@ def test_dsv4_hybrid_exact_flops(self): assert actual_flops == expected_flops + def test_dsv4_hybrid_packed_flops_match_mcore_split(self): + """Packed DSv4 FLOPs split token-linear work from quadratic sparse work.""" + batch_size = 2 + seq_len = 256 + hidden_size = 512 + num_layers = 4 + num_heads = 8 + v_head_dim = 64 + q_lora_rank = 128 + o_lora_rank = 64 + o_groups = 2 + window = 64 + idx_n_heads = 4 + idx_head_dim = 32 + idx_topk = 16 + ffn_hidden_size = 2048 + vocab_size = 1024 + compress_ratios = [0, 4, 128, 128] + packed_lengths = [64, 64, 128, 256] + seqlen_sum = sum(packed_lengths) + seqlen_squared_sum = sum(length**2 for length in packed_lengths) + + model_cfg = MockModelConfig( + num_layers=num_layers, + hidden_size=hidden_size, + seq_length=seq_len, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_heads, + vocab_size=vocab_size, + multi_latent_attention=True, + experimental_attention_variant="dsv4_hybrid", + q_lora_rank=q_lora_rank, + qk_head_dim=32, + qk_pos_emb_head_dim=32, + v_head_dim=v_head_dim, + o_lora_rank=o_lora_rank, + o_groups=o_groups, + csa_compress_ratios=compress_ratios, + csa_window_size=window, + dsa_indexer_n_heads=idx_n_heads, + dsa_indexer_head_dim=idx_head_dim, + dsa_indexer_topk=idx_topk, + gated_linear_unit=False, + ) + cfg = MockConfigContainer(model=model_cfg) + + q_term = q_lora_rank * (hidden_size + num_heads * v_head_dim + 1) + kv_term = hidden_size * v_head_dim + v_head_dim + o_term = num_heads * v_head_dim * o_lora_rank + o_groups * o_lora_rank * hidden_size + projection_term = 3 * 2 * num_layers * (q_term + kv_term + o_term) + + n_layers_r0 = compress_ratios.count(0) + n_layers_r4 = compress_ratios.count(4) + n_layers_r128 = compress_ratios.count(128) + sparse_attn_r0 = n_layers_r0 * num_heads * window * v_head_dim * 2 + sparse_attn_r128_window = n_layers_r128 * num_heads * window * v_head_dim * 2 + effective_topk = min(idx_topk, seq_len // 4) + avg_comp_4 = effective_topk * (1 - effective_topk * 4 / (2 * seq_len)) + sparse_attn_r4 = n_layers_r4 * num_heads * (window + avg_comp_4) * v_head_dim * 2 + main_compressor_term = ( + n_layers_r4 * hidden_size * (2 * v_head_dim) * 2 + n_layers_r128 * hidden_size * v_head_dim * 2 + ) + indexer_token_term = ( + n_layers_r4 * hidden_size * (2 * idx_head_dim) * 2 + + n_layers_r4 * q_lora_rank * idx_n_heads * idx_head_dim + + n_layers_r4 * hidden_size * idx_n_heads + ) + self_attn_token_term = projection_term + 3 * 2 * ( + sparse_attn_r0 + sparse_attn_r4 + sparse_attn_r128_window + main_compressor_term + indexer_token_term + ) + self_attn_core_term = ( + 3 * 2 * (n_layers_r128 * num_heads * v_head_dim / 128 + n_layers_r4 * idx_n_heads * idx_head_dim / 4) + ) + + mlp_term = 3 * 2 * hidden_size * (ffn_hidden_size * 2 * num_layers) + logit_term = 3 * 2 * hidden_size * vocab_size + expected_flops = ( + seqlen_sum * (mlp_term + self_attn_token_term + logit_term) + seqlen_squared_sum * self_attn_core_term + ) + + actual_flops = num_floating_point_operations( + cfg, + batch_size=batch_size, + seqlen_sum=seqlen_sum, + seqlen_squared_sum=seqlen_squared_sum, + ) + bshd_flops = num_floating_point_operations(cfg, batch_size=batch_size) + + assert actual_flops == expected_flops + assert actual_flops < bshd_flops + def test_dsv4_hybrid_validates_compress_ratio_length(self): """CSA compress-ratio count must match decoder plus MTP layer count.""" model_cfg = MockModelConfig(