From 47340007a42312ca9dd6eed3357a88af25a89e3a Mon Sep 17 00:00:00 2001 From: Deyu Fu Date: Thu, 8 Oct 2026 17:34:50 +0800 Subject: [PATCH] Count DSv4 attention FLOPs in HybridModel training metrics Account for C/H/W attention using the actual hybrid pattern and real-token statistics, including compressors, indexer work, grouped output projection, and MTP overhead. Preserve the existing non-DSv4 calculations. Related to #5795. Adapted from NVIDIA/Megatron-LM#5142, #5358, and #5762. Signed-off-by: Deyu Fu --- megatron/training/training.py | 83 +++++++++++++++++- .../test_num_floating_point_operations.py | 86 +++++++++++++++++++ 2 files changed, 165 insertions(+), 4 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index 4199efb81aa..63ecd614b6d 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -188,13 +188,13 @@ get_args, get_energy_monitor, get_one_logger, + get_run_config, get_signal_handler, get_telemetry, get_tensorboard_writer, get_timers, get_train_state, get_wandb_writer, - get_run_config, ) from .theoretical_memory_usage import report_theoretical_memory from .utils import ( @@ -838,6 +838,65 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float return total_real_tokens / dedup, seqlen_squared_sum / dedup +def _dsv4_attention_flops(args, layer_counts, total_tokens, seqlen_squared_sum): + """Estimate DSv4 W/C/H attention FLOPs, including projections and compressors. + + Uses the model-FLOPs convention: multiply/add counts as two + operations and forward plus backward as three passes. Window and capped + Top-K work are token-linear; HCA and indexer scoring use the real sequence + length squared sum. Packed Top-K saturation uses the token-weighted mean + length, since these batch statistics do not retain individual document lengths. + """ + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + + if total_tokens == 0: + return 0 + window_layers = layer_counts[Symbols.WINDOW] + csa_layers = layer_counts[Symbols.CSA] + hca_layers = layer_counts[Symbols.HCA] + attention_layers = window_layers + csa_layers + hca_layers + hidden = args.hidden_size + heads = args.num_attention_heads + head_dim = args.v_head_dim + q_rank = args.q_lora_rank + output_rank = args.output_projection_lora_rank + output_groups = args.output_projection_groups + window = args.csa_window_size + sequence_length = seqlen_squared_sum / total_tokens + + # Read v_head_dim directly: DSv4 derives qk_head_dim on the layer config, + # so the raw argument namespace may still carry the ordinary MLA default. + query_projection = q_rank * (hidden + heads * head_dim + 1) + kv_projection = hidden * head_dim + head_dim + output_projection = heads * head_dim * output_rank + output_groups * output_rank * hidden + token_term = attention_layers * (query_projection + kv_projection + output_projection) + token_term += attention_layers * heads * window * head_dim * 2 + + # Ratio-4 compression overlaps two windows; ratio-128 uses one. + token_term += csa_layers * hidden * (2 * head_dim) * 2 + token_term += hca_layers * hidden * head_dim * 2 + core_term = hca_layers * heads * head_dim / 128 + + if csa_layers: + if getattr(args, "csa_dense_mode", False): + # Dense CSA attends to all compressed keys and has no indexer. + core_term += csa_layers * heads * head_dim / 4 + else: + indexer_heads = args.dsa_indexer_n_heads + indexer_dim = args.dsa_indexer_head_dim + topk = min(args.dsa_indexer_topk, sequence_length // 4) + average_compressed_keys = topk * (1 - topk * 4 / (2 * sequence_length)) + token_term += csa_layers * heads * average_compressed_keys * head_dim * 2 + token_term += csa_layers * ( + hidden * (2 * indexer_dim) * 2 + + q_rank * indexer_heads * indexer_dim + + hidden * indexer_heads + ) + core_term += csa_layers * indexer_heads * indexer_dim / 4 + + return 6 * (token_term * total_tokens + core_term * seqlen_squared_sum) + + def num_floating_point_operations( args, batch_size, @@ -1384,18 +1443,25 @@ def _split_spec_part(part): from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, get_hybrid_layer_counts, + parse_hybrid_pattern, ) + layer_counts = get_hybrid_layer_counts(args.hybrid_layer_pattern) num_mamba_layers, num_gdn_layers, num_attn_layers, num_mlp_layers, num_moe_layers = ( itemgetter(Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE)( - get_hybrid_layer_counts(args.hybrid_layer_pattern) + layer_counts ) ) mtp_num_layers = args.mtp_num_layers if mtp_num_layers is None: mtp_num_layers = 0 + has_dsv4_attention = any(layer_counts[symbol] for symbol in Symbols.DSV4_COMPRESS_RATIO_MAP) + if has_dsv4_attention: + # The static C/H/W specs select DSv4 even when the CLI variant is + # inferred only on TransformerConfig. The pattern also owns MTP depth. + mtp_num_layers = parse_hybrid_pattern(args.hybrid_layer_pattern).mtp_num_depths # Compute hybrid model FLOPs. - return int(hybrid_flops( + flops = hybrid_flops( total_tokens=total_real_tokens_in_batch, seqlen_squared_sum=seqlen_squared_sum_in_batch, hidden_size=args.hidden_size, @@ -1430,7 +1496,16 @@ def _split_spec_part(part): gdn_use_gdn2=(args.experimental_attention_variant == "gdn2"), vocab_size=args.padded_vocab_size, mtp_num_layers=mtp_num_layers, - )) + ) + if has_dsv4_attention: + flops += _dsv4_attention_flops( + args, layer_counts, total_real_tokens_in_batch, seqlen_squared_sum_in_batch + ) + # MTP normalization and the concatenated embedding/hidden projection. + flops += 6 * mtp_num_layers * ( + 3 * args.hidden_size + 2 * args.hidden_size ** 2 + ) * total_real_tokens_in_batch + return int(flops) else: # Compute standard Transformer model FLOPs. return int(transformer_flops()) diff --git a/tests/unit_tests/training/test_num_floating_point_operations.py b/tests/unit_tests/training/test_num_floating_point_operations.py index 6210b5ed2f0..aa1652857d3 100644 --- a/tests/unit_tests/training/test_num_floating_point_operations.py +++ b/tests/unit_tests/training/test_num_floating_point_operations.py @@ -140,3 +140,89 @@ def test_plain_attention_and_dense_mlp_flops_are_unchanged(layout, mtp_depth): actual = num_floating_point_operations(args, batch_size=3, **stats) assert actual == _expected_flops(tokens, squared_lengths, (0, 3 + depth, 3 + depth, 0), depth) + + +def _make_dsv4_args(pattern, **overrides): + """A small HybridModel whose FLOPs have independently calculated coefficients.""" + args = _make_args( + hybrid_layer_pattern=pattern, + experimental_attention_variant=None, + num_layers=len(pattern.split('/')[0].replace('|', '')), + seq_length=256, + multi_latent_attention=True, + group_query_attention=False, + q_lora_rank=2, + # Config derives this dimension from v_head_dim; args may retain this stale default. + qk_head_dim=128, + qk_pos_emb_head_dim=2, + v_head_dim=8, + output_projection_groups=2, + output_projection_lora_rank=3, + csa_window_size=2, + csa_dense_mode=False, + dsa_indexer_n_heads=2, + dsa_indexer_head_dim=4, + dsa_indexer_topk=2, + mamba_state_dim=4, + mamba_head_dim=2, + mamba_num_groups=1, + mamba_num_heads=2, + gdp_num_householder=1, + ) + for name, value in overrides.items(): + setattr(args, name, value) + return args + + +@pytest.mark.parametrize('symbol', ['W', 'C', 'H']) +@pytest.mark.parametrize('packed', [False, True]) +def test_dsv4_attention_flops_use_pattern_and_real_token_statistics(symbol, packed): + args = _make_dsv4_args(symbol) + # BSHD: two L=256 sequences. Packed: L=128 and L=384 in a padded allocation. + tokens = 512 + squared_lengths = 128**2 + 384**2 if packed else 2 * 256**2 + stats = ( + dict(total_real_tokens_in_batch=tokens, seqlen_squared_sum_in_batch=squared_lengths) + if packed + else {} + ) + if symbol != 'C': + args.dsa_indexer_n_heads = args.dsa_indexer_head_dim = args.dsa_indexer_topk = None + # Per-layer projection coefficient is 218 MACs; window is 64. C adds + # compressor/indexer coefficients 256/160 and a capped sparse core. H adds + # compressor coefficient 128 and a 1/128 compressed-attention core. + coefficients = {'W': (1692, 0), 'C': (4572, 12), 'H': (2460, 0.75)} + token_coefficient, core_coefficient = coefficients[symbol] + expected = (token_coefficient + 1536) * tokens + core_coefficient * squared_lengths + if symbol == 'C': + expected -= 1536 * tokens**2 / squared_lengths + before = vars(args).copy() + assert num_floating_point_operations(args, batch_size=2, **stats) == int(expected) + assert vars(args) == before + if packed: + args.seq_length = 2048 + assert num_floating_point_operations(args, batch_size=32, **stats) == int(expected) + + +def test_dsv4_flops_count_repeated_mtp_attention_and_projection(): + args = _make_dsv4_args('W-E/C-/C-', mtp_num_layers=None) + tokens, squared_lengths = 512, 2 * 256**2 + # Pattern: W + 2*C + 3*dense MLP + 1*MoE; two MTP norms/projections, + # three logits projections. The actual repeated pattern owns MTP depth. + expected = 27636 * tokens + 24 * squared_lengths - 3072 * tokens**2 / squared_lengths + assert num_floating_point_operations(args, batch_size=2) == int(expected) + assert args.mtp_num_layers is None + + +def test_dsv4_dense_csa_flops_do_not_require_an_indexer(): + args = _make_dsv4_args( + 'C', + csa_dense_mode=True, + dsa_indexer_n_heads=None, + dsa_indexer_head_dim=None, + dsa_indexer_topk=None, + ) + # Dense ratio-4 CSA: projections/window/compressor=538 MACs per token, + # plus the full compressed attention core; logits add1536 FLOPs per token. + expected = (6 * 538 + 1536) * 512 + 24 * (2 * 256**2) + assert num_floating_point_operations(args, batch_size=2) == expected