diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index f8779d674e2..ed709cbe46a 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1303,9 +1303,21 @@ def __post_init__(self): if self.kv_channels is None: self.kv_channels = self.hidden_size // self.num_attention_heads - if self.num_query_groups is None: + # num_query_groups == 0 is treated the same as None (both mean "use num_attention_heads"). + # This keeps minimal configs valid (num_attention_heads itself defaults to 0) and ensures + # a real attention config never reaches attention initialization with a zero query-group + # count, which would otherwise divide by zero downstream. + if self.num_query_groups is None or self.num_query_groups == 0: self.num_query_groups = self.num_attention_heads + if self.num_attention_heads > 0 and ( + self.num_query_groups <= 0 or self.num_attention_heads % self.num_query_groups != 0 + ): + raise ValueError( + f"num_query_groups ({self.num_query_groups}) must be a positive divisor of " + f"num_attention_heads ({self.num_attention_heads})." + ) + if ( self.num_query_groups % self.tensor_model_parallel_size != 0 and self.tensor_model_parallel_size % self.num_query_groups != 0 diff --git a/tests/unit_tests/transformer/test_transformer_config.py b/tests/unit_tests/transformer/test_transformer_config.py new file mode 100644 index 00000000000..574d8715bcb --- /dev/null +++ b/tests/unit_tests/transformer/test_transformer_config.py @@ -0,0 +1,49 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest + +from megatron.core.transformer.transformer_config import TransformerConfig + + +class TestTransformerConfig: + def test_num_query_groups_divides_num_attention_heads(self): + config = TransformerConfig( + num_layers=2, hidden_size=128, num_attention_heads=32, num_query_groups=8 + ) + assert config.num_query_groups == 8 + + def test_num_query_groups_defaults_to_num_attention_heads(self): + config = TransformerConfig(num_layers=2, hidden_size=128, num_attention_heads=32) + assert config.num_query_groups == 32 + + def test_num_query_groups_not_dividing_num_attention_heads_raises(self): + with pytest.raises(ValueError, match="must be a positive divisor of num_attention_heads"): + TransformerConfig( + num_layers=2, hidden_size=128, num_attention_heads=32, num_query_groups=5 + ) + + def test_num_query_groups_larger_than_num_attention_heads_raises(self): + with pytest.raises(ValueError, match="must be a positive divisor of num_attention_heads"): + TransformerConfig( + num_layers=2, hidden_size=128, num_attention_heads=4, num_query_groups=8 + ) + + def test_negative_num_query_groups_raises(self): + with pytest.raises(ValueError, match="must be a positive divisor of num_attention_heads"): + TransformerConfig( + num_layers=2, hidden_size=128, num_attention_heads=4, num_query_groups=-1 + ) + + def test_zero_num_query_groups_normalized_to_num_attention_heads(self): + # num_query_groups == 0 is treated like None: normalized to num_attention_heads, so a + # real attention config never reaches attention init with a zero query-group count. + config = TransformerConfig( + num_layers=2, hidden_size=128, num_attention_heads=2, num_query_groups=0 + ) + assert config.num_query_groups == 2 + + def test_minimal_config_without_attention_heads_is_allowed(self): + # num_attention_heads defaults to 0 in minimal configs used by many non-attention tests; + # num_query_groups then normalizes to 0 and validation is skipped. + config = TransformerConfig(num_layers=1, kv_channels=1) + assert config.num_query_groups == 0