Add per-head Muon QKV orthogonalization and fix TP>num_query_groups with Muon. Generalize to MLA and GQA

[mkhona-nvidia]

Summary

• add an opt-in --muon-split-qkv-per-head mode that applies Muon's Newton-Schulz orthogonalization independently to each physical Q, gate, K, and V head
• derive per-head layouts for MHA/GQA and gated attention while preserving the existing projection-level QKV split by default
• keep QKV parameters on the Muon split path when tensor parallelism fragments query-group blocks, including TP > num_query_groups
• reconstruct fragmented layouts across GTP-remat and tensor-parallel row shards, then restore each rank's local update
• exclude each TP partition's GTP alignment padding before orthogonalization and restore zero padding before taking the local GTP shard
• expose the behavior through both Muon and AdaptiveMuon and add focused unit coverage

Root cause

The projection-split path previously required each rank-local QKV row shard to contain an integral number of query-group blocks. When TP split a block across ranks, the parameter was not tagged for QKV splitting. The optimizer now records the global fused layout and gathers fragmented blocks before applying Q/K/V or per-head orthogonalization.

GTP alignment padding is applied independently to every TP-local weight slice. Treating physical GTP shard rows as logical QKV rows either caused a layout mismatch during optimizer construction or allowed padding to affect head and projection boundaries. The gather path now strips padding after reconstructing each TP slice, before the TP gather, and restores inert zero padding in reverse order.

Validation

• git diff --check
• Python syntax compilation for all changed files
• Black 26.3.0
• isort 5.13.2
• Ruff 0.9.x
• Pylint 3.2.6 (10.00/10)
• added multi-rank coverage for fragmented gated and ungated projection splits, plus fragmented per-head splits
• added a TP=2, GTP=2 regression covering padded per-head splits and padded TP > num_query_groups projection splits

The distributed pytest coverage was not run locally because this macOS workspace does not provide the four-GPU CUDA/distributed test environment required by Megatron's unit-test runner.

[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.

[mkhona-nvidia]

/claude review

[claude]

Light review. Three things worth addressing before this leaves draft.

1) The graceful skip was replaced by a hard failure (megatron/core/optimizer/__init__.py:793-847)

Previously a QKV param whose row count did not match sum(qkv_split_shapes) was logged at DEBUG and left with is_qkv=False, so Muon fell back to whole-matrix orthogonalization. Now param.is_qkv = True is set unconditionally and a mismatch raises RuntimeError.

qkv_split_shapes is computed once per model chunk from model_chunk.config and reused for every linear_qkv.weight. For heterogeneous models (heterogeneous_block_specs, nemotron-nas), HeterogeneousTransformerConfig.get_config_for_layer overrides num_query_groups per block (megatron/core/transformer/heterogeneous/heterogeneous_config.py:246-250), so layers whose num_query_groups differs from the top-level config will now abort optimizer construction instead of degrading gracefully.

Suggest keeping the fallback for the shapes-do-not-match case and reserving RuntimeError for layouts that should have matched:

                if expected_global_rows != sum(global_split_shapes):
                    log_single_rank(
                        logger,
                        logging.DEBUG,
                        f"Emerging optimizer QKV split skipped for {name}: "
                        f"global_rows={sum(global_split_shapes)}, local_rows={param.shape[0]}, "
                        f"tp_size={tp_size}, gtp_remat_size={gtp_size}",
                    )
                    param.is_qkv = False
                    continue

Alternatively, compute qkv_split_shapes per parameter rather than caching one value for the whole chunk.

2) The uniform per-head path cannot run against the pinned dependency

pyproject.toml:239 and uv.lock:1136 pin emerging_optimizers to v0.2.0, but _require_batched_newton_schulz_support() requires >= 0.3.0. Uniform head sizes — i.e. every real model, since all splits are kv_channels — hit len(set(split_shapes)) == 1 at emerging_optimizers.py:411, so --muon-split-qkv-per-head raises RuntimeError in the CI container and in any env built from the lockfile.

Every new test that reaches the batched branch monkeypatches EMERGING_OPTIMIZERS_VERSION to Version("0.3.0"), so orthogonalize_fn(grad.view(len(split_shapes), head_rows, -1)) is only ever exercised against a stubbed center_rows, never against the real newton_schulz_tp. The 3D-input contract with upstream is therefore unvalidated. Either bump the pin in this PR, or state the version requirement in the --muon-split-qkv-per-head help text so users see the constraint before the traceback.

3) No coverage for the parameter-tagging logic

The 76 new lines in _get_megatron_emerging_optimizer are the highest-risk part of this change — TP/GTP offset math, local_start, qkv_gtp_pad_length derivation, and two new RuntimeError paths — and nothing tests them. All new tests hand-set qkv_split_shapes_global / qkv_split_heads_are_complete / qkv_split_groups_are_complete on bare torch.nn.Parameter objects, so the code that derives those attributes never runs. Notably, the headline fix (keeping QKV params on the Muon split path when TP fragments query-group blocks, including TP > num_query_groups) has no test that goes through the tagging path at all.

A test that builds a small GPT model with tensor_model_parallel_size > num_query_groups, calls get_megatron_optimizer with muon_split_qkv=True, and asserts on the resulting linear_qkv.weight attributes would cover the regression this PR fixes. Per the mcore-testing skill that needs torch.distributed.run --nproc-per-node N; tests/unit_tests/test_emerging_optimizers.py already initializes model parallel via Utils.initialize_model_parallel, so it fits there.

[mkhona-nvidia]

/claude review

[skyw]

Avoid ambiguous naming like "global", hard to interpret what it means.
If it is the convention of the code base, comment what it means.

[skyw]

Shouldn't need this path, get_pg_size and get_pg_rank returns 1 and 0 with None.

[skyw]

Better to name it from_transfomer_config.

[skyw]

looks like from_splits? repeated is on output not input, i.e. not from.

[FDecaYed]

I think we can keep the behavior simpler here: preserve the old TP Muon path only for the case it already supported correctly — projection-level QKV splitting where each TP rank has complete local query groups and GTP is not involved.

Per-head splitting and GTP should always use the new path. The same applies when projection-level QKV cannot be split locally. If muon_tp_mode=distributed is requested in those cases, please fall back to non-TP/duplicated NS and print a warning instead of silently ignoring the setting.

In short: keep the old behavior only for the old valid case; route all new cases（perhead/gtp/tp>group) through the new implementation. and print warning

[FDecaYed]

Now that #6381 is in main, there is one compatibility case to handle: uniform per-head splits use 3D batched NS, but EO’s SYRK path only supports 2D input. muon_split_qkv_per_head + muon_use_syrk would therefore fail.

Could we use the batched path only when SYRK is disabled, and fall back to one 2D call per head when it is enabled? It would also be good to add a test for this flag combination.

Longer term, I would prefer MCore to pin and support a single EO version instead of adding more version-specific branches here, but that can be a follow-up.

[buptzyb]

With this layout, enabling muon_split_qkv_per_head appears to run separate Newton–Schulz updates for every head’s content-Q (W^UQ) and RoPE-Q (W^QR) slices. Is per-head Muon for W^QR required by a specific model or recipe? If not, do we need mixed-granularity layout metadata so that W^UQ can be split per head while W^QR retains its intended whole-projection treatment?

The same question applies to the corresponding FusedMLA and AbsorbedMLA layouts below.

[mkhona-nvidia]

We think per-head muon should apply to each attention head individually. So I would be biased towards splitting every type of head regardless of RoPE (note that latet SOTA hybrid models do not have RoPE) for both GQA and MLA.

GLM-5 is where it was introduced: https://arxiv.org/html/2602.15763v2

and it has since been adopted by Kimi-K3. We expect it to be default for all models going forward