fix(muon): split the fully-gathered QKV tensor before orthogonalizing under GTP

[fanshiqing]

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

What does this PR do?

Muon runs a different update rule on linear_qkv depending on whether GTP is
enabled, silently. This makes the decision layout-invariant.

Problem

The split-QKV decision tests the local shard row count. Under GTP the weight is
row-sharded, so the same weight takes different paths:

layout	param.shape[0]	% 6528	path
TP1	6528 (whole)	0 ✓	split → 3 Newton–Schulz
GTP64	102 (shard)	✗	1 whole-matrix NS

Three facts make it matter:

• The rules are far apart — 26 % different overall, 95 % on k and v (they are
  2.9 % of the rows, so joint orthogonalization lets q dominate them).
• It does not wash out — NS amplifies the ~1e-7 fp32 reduction-order difference
  between layouts by ~10⁴x, and that floor is flat across condition number 1e1–1e6.
• It is invisible — the skip is a logging.DEBUG record that production never sees.

Measured at 256 GPUs (identical init on both arms), GTP64 vs matched TP1:
training loss +0.173 % — a near-constant 0.09–0.12 B token lag, 4.3x outside
the run-to-run noise floor. GTP64 + Adam on the identical topology shows none, as
expected: Adam is elementwise, amplification ~1x.

Fix

1. Decide from the global shape — qkv_global_rows():
   param.shape[0] * gtp_size - pad_length.
2. Split after the all-gather — at 6528 rows over 64 shards the q/k/v boundaries
   fall mid-shard, so no rank can split correctly before the gather.
3. take the GTP degree from the optimizer's own group, gated on
   is_gtp_weight_remat -- the same signal the split path uses

Result

	GTP64 vs its own TP1 baseline	vs noise floor (−0.038 … +0.040 %)
before	+0.1731 %	4.3x outside
after	−0.0342 %	✅ inside

(4 h from scratch, matched window 4.0–16.8 B, n=31. The residual is negative — GTP
slightly ahead — i.e. an overshoot past zero, not a leftover deficit.)

• Bitwise: TP1 and GTP now produce identical qkv updates (max |diff| 0.000e+00).
• Throughput: 2,527 → 2,536 ms/it at GTP64 (+0.36 %, inside cross-job noise). The
  split partitions the same NS FLOPs rather than adding any, on 8 attention layers.
• Control: --muon-no-split-qkv on both arms also lands inside the floor
  (+0.008 %), confirming the asymmetry was the mechanism. Not the right fix — turning
  the split off costs +0.216 % lm loss at TP1, so GTP should gain it, not TP1 lose it.

Supporting measurements

Split vs whole-matrix NS, production settings (16 steps, polar_express, spectral,
extra_scale_factor 0.2), same 6528×2688 gradient:

block	rows	relative difference
q	6144	0.116
k / v	192	0.947
whole	6528	0.255

NS amplification of a 1e-7 relative input perturbation:

cond	1e1	1e3	1e6
qkv 6528×2688	10,525x	9,720x	20,374x

Also eliminated with data before landing here: expert routing, grad-norm inflation,
data assignment, loss/MTP reduction scope, router BF16 quantization, vocab-sized
weights, initialization.

---

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run tools/autoformat.sh on my PR

tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_qkv.py — 14 cases, all
passing at world 2 and 4. Nothing existing could have caught this:
test_muon_optimizer_qkv_split sets param.is_qkv by hand on one rank, and
test_gtp_muon.py's parity cases carry no qkv weight.

• test_local_shard_would_have_disabled — the regression test proper; pins the shape
  that used to fail, 102 × 64 = 6528.
• test_split_after_gather_matches_tp1 — GTP result == TP1 result, per shard.
• test_gtp_split_differs_from_whole_matrix — negative control, so a regression to
  the old behaviour cannot pass trivially.
• test_row_sharded_modes_raise, test_padding_is_excluded,
  test_decision_is_layout_invariant, test_non_qkv_shape_is_rejected.

test_gtp_muon.py passes identically before and after (6 passed / 8 skipped both
sides; the 8 need an 8-GPU host).

Reviewer notes

• Behaviour change: blockwise/distributed with muon_split_qkv=True now
  raises. Those modes keep the weight row-sharded so no rank holds a q/k/v boundary —
  they were silently getting a different update rule. --muon-no-split-qkv is the
  opt-out.
• qkv_global_rows reads is_qkv/qkv_split_shapes (tensor-parallel attrs) and
  group/pad_length (GTP attrs) off the same param. Both are copied at the same
  sites today; a path copying only one set would compute global_rows wrongly.
• Same pattern, not fixed here: mixer.in_proj [12352, 2688] is also fused
  ([z,x,B,C,dt]) and GTP-sharded, escaping only via use_muon=False.
• MLA still unhandled (pre-existing TODO(deyuf)).

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

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[fanshiqing]

/claude review

[fanshiqing]

/ok to test b142e6a

[fanshiqing]

/ok to test 886b565

[fanshiqing]

/claude review

[fanshiqing]

/ok to test a6eb95e

[fanshiqing]

/ok to test 1c72d57

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34833388075