diff --git a/src/megatron/bridge/training/utils/flop_utils.py b/src/megatron/bridge/training/utils/flop_utils.py index 00041e278e..64b85a769f 100644 --- a/src/megatron/bridge/training/utils/flop_utils.py +++ b/src/megatron/bridge/training/utils/flop_utils.py @@ -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 @@ -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) + 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 @@ -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: diff --git a/tests/unit_tests/training/utils/test_flop_utils.py b/tests/unit_tests/training/utils/test_flop_utils.py index 32804e79d3..047043f011 100644 --- a/tests/unit_tests/training/utils/test_flop_utils.py +++ b/tests/unit_tests/training/utils/test_flop_utils.py @@ -80,9 +80,17 @@ class MockModelConfig: qk_head_dim: int = 64 qk_pos_emb_head_dim: int = 0 v_head_dim: int = 64 + o_lora_rank: int = 0 + o_groups: int = 1 # Sliding window attention settings window_size: tuple | list | int | None = None window_attn_skip_freq: int | list | None = None + # DeepSeek-V4 hybrid attention settings + csa_compress_ratios: list[int] | None = None + csa_window_size: int = 128 + dsa_indexer_n_heads: int | None = None + dsa_indexer_head_dim: int | None = None + dsa_indexer_topk: int | None = None # GDN (Gated DeltaNet) settings experimental_attention_variant: str | None = None linear_attention_freq: int | list | None = None @@ -630,6 +638,232 @@ def test_gdn_more_gdn_layers_changes_flops(self): assert flops_freq4 != flops_freq8, "Different GDN ratios should produce different FLOPs" +class TestDeepSeekV4HybridFlops: + """Tests for DeepSeek-V4 hybrid attention FLOPs in the transformer path.""" + + def test_dsv4_hybrid_exact_flops(self): + """DSv4 hybrid FLOPs include sparse attention, compressor, and indexer terms.""" + batch_size = 1 + seq_len = 512 + hidden_size = 128 + num_layers = 3 + num_heads = 4 + v_head_dim = 32 + q_lora_rank = 16 + qk_head_dim = 24 + qk_pos_emb_head_dim = 8 + o_lora_rank = 16 + o_groups = 2 + window = 64 + idx_n_heads = 2 + idx_head_dim = 8 + idx_topk = 32 + ffn_hidden_size = 256 + vocab_size = 1024 + + 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=qk_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + o_lora_rank=o_lora_rank, + o_groups=o_groups, + csa_compress_ratios=[0, 4, 128], + 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 * (qk_head_dim + qk_pos_emb_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) + + sparse_attn_r0 = num_heads * window * v_head_dim * 2 + sparse_attn_r128 = num_heads * (window + (seq_len // 128) / 2) * 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 = num_heads * (window + avg_comp_4) * v_head_dim * 2 + main_compressor_term = hidden_size * (2 * v_head_dim) * 2 + hidden_size * v_head_dim * 2 + indexer_term = ( + hidden_size * (2 * idx_head_dim) * 2 + + q_lora_rank * idx_n_heads * idx_head_dim + + hidden_size * idx_n_heads + + idx_n_heads * idx_head_dim * (seq_len // 4) + ) + dsv4_extra_term = ( + 3 * 2 * (sparse_attn_r0 + sparse_attn_r4 + sparse_attn_r128 + main_compressor_term + indexer_term) + ) + self_attention_term = projection_term + dsv4_extra_term + + mlp_term = 3 * 2 * hidden_size * (ffn_hidden_size * 2 * num_layers) + logit_term = 3 * 2 * hidden_size * vocab_size + expected_flops = batch_size * seq_len * (mlp_term + self_attention_term + logit_term) + + actual_flops = num_floating_point_operations(cfg, batch_size=batch_size) + + assert actual_flops == expected_flops + + def test_dsv4_hybrid_validates_compress_ratio_length(self): + """CSA compress-ratio count must match decoder plus MTP layer count.""" + model_cfg = MockModelConfig( + num_layers=2, + mtp_num_layers=1, + multi_latent_attention=True, + experimental_attention_variant="dsv4_hybrid", + q_lora_rank=16, + o_lora_rank=16, + csa_compress_ratios=[0, 4], + dsa_indexer_n_heads=2, + dsa_indexer_head_dim=8, + dsa_indexer_topk=32, + ) + cfg = MockConfigContainer(model=model_cfg) + + with pytest.raises(ValueError, match=r"expected 3 \(num_layers=2, mtp_num_layers=1\)"): + num_floating_point_operations(cfg, batch_size=1) + + def test_dsv4_hybrid_validates_supported_compress_ratios(self): + """CSA compress ratios must be recognized before layer counts are used.""" + model_cfg = MockModelConfig( + num_layers=3, + multi_latent_attention=True, + experimental_attention_variant="dsv4_hybrid", + q_lora_rank=16, + o_lora_rank=16, + csa_compress_ratios=[0, 8, 128], + dsa_indexer_n_heads=2, + dsa_indexer_head_dim=8, + dsa_indexer_topk=32, + ) + cfg = MockConfigContainer(model=model_cfg) + + with pytest.raises(ValueError, match=r"unsupported values: \[8\]"): + num_floating_point_operations(cfg, batch_size=1) + + def test_dsv4_hybrid_requires_q_lora_rank(self): + """DSv4 hybrid FLOPs require q_lora_rank for projection accounting.""" + model_cfg = MockModelConfig( + multi_latent_attention=True, + experimental_attention_variant="dsv4_hybrid", + q_lora_rank=None, + csa_compress_ratios=[0] * 24, + ) + cfg = MockConfigContainer(model=model_cfg) + + with pytest.raises(ValueError, match="q_lora_rank must be set"): + num_floating_point_operations(cfg, batch_size=1) + + def test_dsv4_hybrid_requires_compress_ratios(self): + """DSv4 hybrid FLOPs require per-layer CSA compress ratios.""" + model_cfg = MockModelConfig( + multi_latent_attention=True, + experimental_attention_variant="dsv4_hybrid", + q_lora_rank=16, + csa_compress_ratios=None, + ) + cfg = MockConfigContainer(model=model_cfg) + + with pytest.raises(ValueError, match="csa_compress_ratios must be set"): + num_floating_point_operations(cfg, batch_size=1) + + def test_dsv4_hybrid_without_ratio4_layers_skips_indexer_terms(self): + """DSv4 hybrid FLOPs support CSA patterns without DSA-indexed ratio-4 layers.""" + batch_size = 1 + seq_len = 512 + hidden_size = 128 + num_layers = 2 + num_heads = 4 + v_head_dim = 32 + q_lora_rank = 16 + qk_head_dim = 24 + qk_pos_emb_head_dim = 8 + o_lora_rank = 16 + o_groups = 2 + window = 64 + ffn_hidden_size = 256 + vocab_size = 1024 + + 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=qk_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + o_lora_rank=o_lora_rank, + o_groups=o_groups, + csa_compress_ratios=[0, 128], + csa_window_size=window, + gated_linear_unit=False, + ) + cfg = MockConfigContainer(model=model_cfg) + + q_term = q_lora_rank * (hidden_size + num_heads * (qk_head_dim + qk_pos_emb_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) + + sparse_attn_r0 = num_heads * window * v_head_dim * 2 + sparse_attn_r128 = num_heads * (window + (seq_len // 128) / 2) * v_head_dim * 2 + main_compressor_term = hidden_size * v_head_dim * 2 + dsv4_extra_term = 3 * 2 * (sparse_attn_r0 + sparse_attn_r128 + main_compressor_term) + self_attention_term = projection_term + dsv4_extra_term + + mlp_term = 3 * 2 * hidden_size * (ffn_hidden_size * 2 * num_layers) + logit_term = 3 * 2 * hidden_size * vocab_size + expected_flops = batch_size * seq_len * (mlp_term + self_attention_term + logit_term) + + actual_flops = num_floating_point_operations(cfg, batch_size=batch_size) + + assert actual_flops == expected_flops + + @pytest.mark.parametrize( + ("missing_field", "error_message"), + [ + ("dsa_indexer_n_heads", "dsa_indexer_n_heads must be set"), + ("dsa_indexer_head_dim", "dsa_indexer_head_dim must be set"), + ("dsa_indexer_topk", "dsa_indexer_topk must be set"), + ], + ) + def test_dsv4_hybrid_ratio4_requires_indexer_config(self, missing_field, error_message): + """DSv4 hybrid ratio-4 layers require DSA indexer config for FLOPs accounting.""" + model_kwargs = { + "num_layers": 3, + "multi_latent_attention": True, + "experimental_attention_variant": "dsv4_hybrid", + "q_lora_rank": 16, + "o_lora_rank": 16, + "csa_compress_ratios": [0, 4, 128], + "dsa_indexer_n_heads": 2, + "dsa_indexer_head_dim": 8, + "dsa_indexer_topk": 32, + } + model_kwargs[missing_field] = None + model_cfg = MockModelConfig(**model_kwargs) + cfg = MockConfigContainer(model=model_cfg) + + with pytest.raises(ValueError, match=error_message): + num_floating_point_operations(cfg, batch_size=1) + + class TestHybridMtpPatternParsing: """Tests for hybrid/MTP pattern parsing in FLOPs accounting."""