Validate num_attention_heads is divisible by num_query_groups in TransformerConfig

[huthvincent]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Rejects invalid grouped-query-attention configurations at config-construction time instead of letting them crash with a cryptic view() shape error deep inside attention.py.

TransformerConfig.__post_init__ now raises a clear ValueError when:

• num_query_groups is not positive (previously num_query_groups=0 passed validation because 0 % tp_size == 0), or
• num_attention_heads is not divisible by num_query_groups (hard GQA requirement; the per-group head count is computed with floor division).

Without this check, e.g. num_attention_heads=32, num_query_groups=5 constructs fine and then fails on the first forward pass with:

RuntimeError: shape '[8, 1, 5, 16]' is invalid for input of size 672

Also updates one existing unit test (test_trtllm_single_device_converter.py) that used num_query_groups=0 as a placeholder; it now uses num_query_groups=num_attn_heads (plain MHA, matching the mocked weight layout).

All existing in-tree configs and tests satisfy the new check (verified by grepping every num_query_groups usage in megatron/ and tests/).

Issue tracking

Linked issue: Fixes #5755

Contribution process

Pre-checks

• I have added relevant unit tests (tests/unit_tests/transformer/test_transformer_config.py, 5 cases: valid GQA, default groups, non-divisible, groups > heads, non-positive)
• I have added relevant functional tests (not applicable — config validation)
• I have added proper typing to my code
• I have added relevant documentation (error message is self-documenting)
• I have run the autoformatter on the changed files (black 24.10, isort)

Local verification (single GPU):

$ torchrun --nproc-per-node 1 -m pytest -q tests/unit_tests/transformer/test_transformer_config.py
5 passed
$ torchrun --nproc-per-node 1 -m pytest -q tests/unit_tests/export/trtllm/test_trtllm_single_device_converter.py
5 passed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[cspades]

@huthvincent Want me to run CI/CD?

[huthvincent]

@cspades Yes, please, thank you! CI would be much appreciated. All three PRs (#5756, #5759, #5761) are small, self-contained bug fixes; happy to have CI run on them whenever convenient.

[cspades]

/ok to test f6f5f53

[huthvincent]

@cspades Sorry for the extra round. Right after your /ok to test, I pushed a small follow-up commit (dd8db013c) that adds the required NVIDIA copyright headers to the changed files so the copyright-check job passes (my new test file used the old header format, and the trtllm test file I touched had no header). Since that moved the PR head, the earlier approval is now tied to the previous commit and CI did not start on the new one. Could you re-run /ok to test dd8db013c when convenient? Thank you, and sorry again for the churn.

[cspades]

/ok to test dd8db01

[huthvincent]

@cspades Update on this one. CI caught a real regression in my earlier commit: the num_query_groups positivity check I added rejected num_query_groups == 0, which is a valid sentinel (equal to num_attention_heads) relied upon by many minimal configs, so it broke a number of distributed tests. I have pushed 87ebbf113, which enforces only the divisibility requirement and only when num_query_groups > 0, reverts the unrelated test change, and adds coverage for the zero/minimal case. Verified locally against the configs that were failing. Sorry for the churn again. Could you re-run /ok to test 87ebbf113 when you have a chance? Thank you for your patience.

[cspades]

/ok to test 87ebbf1

[guihong-nv]

Thanks for the fix. The divisibility check looks good. I think one blocking edge case remains.
The current guard only runs when num_query_groups > 0, so zero and negative values still pass validation. In particular, the new test allows num_attention_heads=2, num_query_groups=0, but zero is not actually normalized to num_attention_heads; only None is. Attention initialization later divides by num_query_groups, so that configuration still fails downstream rather than at config construction.

[huthvincent]

@guihong-nv Great catch, thank you. You are right that 0 was not normalized and would divide by zero during attention initialization. I pushed 15a230d43, which treats num_query_groups == 0 the same as None (both normalize to num_attention_heads), so num_attention_heads=2, num_query_groups=0 now becomes a valid MHA config instead of failing downstream, while minimal configs (where num_attention_heads defaults to 0) stay valid. Zero and negative values are now rejected at construction whenever attention heads are present. I verified locally that the TRT-LLM converter and inference paths are unaffected, since they resolve to the same effective value, and updated the tests accordingly.

[huthvincent]

@cspades Sorry to bother you again. I pushed 15a230d43 to address @guihong-nv's review feedback on the num_query_groups == 0 edge case. Could you re-run /ok to test 15a230d43 when you have a chance, so CI can validate it? Thank you.

[cspades]

/ok to test 15a230d

[cspades]

@huthvincent You can ping @NVIDIA/mcore-oncall and they'll help you!