Validate num_attention_heads is divisible by num_query_groups in TransformerConfig

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

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