[muon] Per-head Newton-Schulz for attention projections

[alanhuangyoo]

Implements #8367 — per-head Newton–Schulz for attention projections, as proposed there.

Scope grew since I opened this: it started as the kernel only, but the metadata and config half
turned out not to depend on the two questions I left on the issue, so it is all here. Points 1–4
of your proposal, plus unit tests for 5; the convergence run is below under "what is not here".

1. Kernel

muon_update gains num_heads. With it set, an attention projection of shape
[num_heads * head_dim, in_features] is viewed as [num_heads, head_dim, in_features] and
Newton–Schulz runs on that batch, so each head is orthogonalized against itself instead of
sharing one update direction with every other head — the coupled-block behaviour Kimi K3
(arXiv:2607.24653 §2.5) and GLM-5 Muon Split (arXiv:2602.15763) both move away from. The existing
max(1, m/n)**0.5 scaling is applied per head block.

As you said, mostly a view: both kernels are already batch-safe, and muon_update already had a
batched branch with per-block scaling for expert groups. This reuses that path.

2. Metadata

set_optimizer_flags tags muon_num_heads alongside use_muon, so it follows the pattern
already there and does not require AutoTP to be on. Head structure comes off the model config:
q/o projections are blocked by num_attention_heads, k/v by num_key_value_heads — different
counts under GQA, and using the query count for k/v would silently split them wrong.

Two things it deliberately declines to guess at:

• Fused QKV stays on the full-matrix path. Its three sections split separately, and under GQA
  they do not share a head count, so treating the matrix as 3 * num_heads uniform blocks would
  be wrong. This is the question I raised on the issue; if you would rather it be handled, say
  which layout to assume and I will add it.
• Anything whose output dim does not divide by the head count is skipped with a warning
  rather than reshaped on a guess.

3. ZeRO integration

All six muon_update call sites pass the tag through. Each already operates on a whole parameter
rather than a flat shard — the ZeRO-1/2 path views the momentum back to tensor.size() and
asserts ndim > 1, the ZeRO-3 path takes param.grad directly, and the DDP paths index real
parameters out of params_pad — so no call site needed reshaping.

4. Config

Opt-in optimizer.params.per_head_muon: true, as suggested. Off by default; with it off,
muon_num_heads is None everywhere and every call site takes exactly the branch it took
before.

Tests

20 CPU-only cases across two files.

test_per_head_muon.py — the arithmetic:

• batched == per-head loop, over (4,8,32)/(2,16,32)/(8,4,64) and both NS methods
• num_heads=1 reproduces the full-matrix path
• per-head differs from full-matrix when heads are unbalanced: one head's gradient scaled
  100×, then asserting the other heads' updates differ from what full-matrix gives them and
  that the four head-update norms land within 1.5× of each other. This is the case that fails if
  num_heads is ignored, which is what makes the equivalence cases load-bearing.
• shape/divisibility rejection

test_per_head_muon_tagging.py — what gets tagged: query count for q/o, kv count for k/v under
GQA, fused QKV skipped, non-attention params untouched, non-divisible shapes skipped, opt-in
required, and use_muon tagging unchanged.

On tolerances: the equivalence cases compare at a bound derived from the kernel's own compute
dtype rather than a tuned epsilon. gram iterates in fp16 and newtonschulz5 in bf16, and NS
amplifies rounding, so batched and unbatched agree to a few ulps, not bitwise — measured
0.027–0.053 absolute against a bf16 eps of 0.0078, norm ratios 0.995–1.005. The tests assert
8 * finfo(dtype).eps elementwise plus a separate norm-ratio check, so scale is still pinned.

$ pytest tests/unit/runtime/zero/test_per_head_muon.py tests/unit/runtime/zero/test_per_head_muon_tagging.py
20 passed

$ pytest tests/unit/ops/muon/test_muon.py
2 failed, 144 skipped        # both also fail on upstream/master in this environment
                             # (TestMuonRejectsReduceScatter, unrelated to this change)

$ yapf==0.40.0 --diff  /  flake8
(clean)

What is not here

The convergence comparison. I have 8×H100 and can run full-matrix vs per-head on a small model
with matched seeds and steps and post loss curves — I would rather do that once you have looked
at the layout assumptions above, so I am not measuring the wrong thing.

[alanhuangyoo]

Pushed a correction and ran this end to end on 2×H100. Reporting both, including the part that
did not work.

A correction first

The previous commit's tagging had two mistakes, both of which produced a wrong update rather
than an error, so I want them on the record rather than quietly fixed:

1. o_proj / out_proj were tagged with the query head count. Their head structure is on
   the input dimension ([hidden, num_heads * head_dim]) while the split is on dim 0. With the
   usual hidden == num_heads * head_dim they still divide evenly, so the matrix was cut across
   the wrong axis silently. Q/K/V only now.
2. dense was matched anywhere in the parameter path, which also names MLP matrices —
   intermediate.dense, output.dense, dense_h_to_4h, dense_4h_to_h. Those have no head
   structure at all, and 4 * hidden divides by the head count just fine, so they were being
   split too. Matching is now on the leaf module name against explicit Q/K/V names, and the
   shape has to confirm the layout (dim 0 divisible by the head count, and equal to
   num_heads * head_dim wherever the config states head_dim).

Both have regression tests. Against the previous logic all five names those tests pin come back
tagged, so they are load-bearing rather than decorative.

Tagging, verified on GPU

Small GQA model (q_heads=8, kv_heads=2, split QKV), ZeRO-1 on 2 ranks, bf16, real training
steps — not a unit test:

per_head_muon: false    q=None k=None v=None  o=None  mlp=None  embed=None
per_head_muon: true     q=8    k=2    v=2     o=None  mlp=None  embed=None

GQA splits correctly (q by 8, k/v by 2), o_proj stays off the per-head path, and MLP and
embedding are untouched.

What the change actually does, measured

The papers' claim is that full-matrix orthogonalization lets heads with larger momentum dominate
the shared update direction. Measured directly on muon_update's output — head update norms,
averaged over 3 seeds, q_heads=8, head_dim=32, hidden=256:

gradient scale spread across heads	full-matrix max/min	per-head max/min	full-matrix CV	per-head CV
uniform (1×)	1.009	1.015	0.0030	0.0053
moderate (10×)	1.120	1.015	0.0408	0.0053
extreme (100×)	1.266	1.015	0.0867	0.0053

With head scales uniform the two agree, so per-head does not distort the balanced case. As head
scales diverge the full-matrix update norms spread out — max/min 1.009 → 1.266, CV up 29× —
while per-head stays flat regardless of the input spread. That is the mechanism the change is
for, and it does not depend on a task being hard enough to show it.

What I could not show

A convergence win. I ran full-matrix vs per-head on the same small model, matched seeds and
data, 400 steps × 2 seeds — but the task I used (next token from a fixed permutation) is learned
in about 40 steps and both runs go to ~0.0 with near-identical curves (0.0471 vs 0.0471,
0.0109 vs 0.0110). That says nothing about either optimizer, so I am not presenting it as
evidence. An earlier attempt with random tokens was worse still — random sequences have no
learnable structure, so both runs just sat at the entropy floor ln(512) = 6.24.

The papers' claim is about stability at scale, which a toy model is the wrong instrument for. If
you have a configuration you would consider a fair test, I have the GPUs to run it.

[alanhuangyoo]

Added coverage against real model architectures. The other tagging cases use a stand-in module
whose leaf names I picked myself, which is circular for a change whose job is recognizing the
names real models use.

Built from actual HF configs, transformers 5.16.1:

architecture	QKV layout	tagged	left on the full-matrix path
llama	split, GQA	q_proj→8, k_proj→2, v_proj→2	o_proj, gate_proj, up_proj, down_proj
qwen2	split, GQA	q_proj→8, k_proj→2, v_proj→2	same
mistral	split, GQA	q_proj→8, k_proj→2, v_proj→2	same
gpt_neox	fused	(nothing)	query_key_value, dense, dense_h_to_4h, dense_4h_to_h
falcon	fused	(nothing)	same

Two things this pins that the synthetic cases could not:

• GQA on real configs. q_proj gets the query head count and k_proj/v_proj the kv count,
  from the model's own config rather than from names I invented.
• The fused-QKV architectures are exactly where the old matching went wrong. gpt_neox and
  falcon name their output projection dense and their MLP matrices dense_h_to_4h /
  dense_4h_to_h. Substring matching on the path tagged all three — an MLP matrix split by a
  head count it has no relationship to. Now none of them are tagged, and the test asserts that
  by listing whatever is tagged when it fails, so a regression names the offender.

30 tests across the two files, all CPU-only.

[alanhuangyoo]

Note on the red modal-torch-latest / DeepSpeedAI CI here, so it does not read as this PR
breaking something: the run was cancelled, not failed. Every test that got to run passed —
the log reaches 61% with no failures and then ends on The operation was canceled. after
1h15m, which looks like the job's own time limit rather than anything in the diff.

It is not specific to this PR either. The last eight runs of that workflow:

running    3597bd0f8  Run multi-rank CPU unit tests in CI via LOCAL_SIZE
cancelled  e9584d3f7  [muon] Per-head Newton-Schulz ...        <- this PR
success    3de47dbad  Fix zero_to_fp32 --safe_serialization ...
cancelled  11876c644  [muon] Per-head Newton-Schulz ...        <- this PR
failure    576954404  Describe universal checkpoint shards ...
success    f009942cb  Filter --include against the real slots ...
cancelled  3bdabae08  fix: Only bind device id when needed ...
cancelled  2b9a606c9  Fix AutoTP + deep compile collectives ...

Four cancellations across three different authors' PRs.

For what it is worth, the tests this PR adds are CPU-only and take about 5 seconds for all 30,
so they are not what is pushing the job over its limit. Happy to re-trigger if you want a clean
run before reviewing.