Count DSv4 attention FLOPs in HybridModel training metrics

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())

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