[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 counts come from
AutoTPMeta.from_model_config, the repo's existing reader for them, and the per-head width from
the config:

leaf	heads	per-head width
q_proj, query, wq	num_attention_heads	head_dim
k_proj, v_proj, key, value, wk, wv	num_key_value_heads	head_dim
q_b_proj, and q_proj on an MLA config	num_attention_heads	qk_nope + qk_rope
kv_b_proj	num_attention_heads	qk_nope + v_head_dim

Different counts under GQA, and using the query count for k/v would silently split them wrong.
MLA does not use head_dim for either up-projection, so a num_heads * head_dim check rejects
both — on GLM-5.2 that is 16 × 64 = 1024 against a real 4096. Without a q_lora_rank there is no
q_a/q_b pair and the query up-projection is a plain q_proj at the same MLA width
(DeepSeek-V2-Lite).

Matching is on the leaf module name, not the path, so dense in attention.output.dense cannot
pull in intermediate.dense. Whatever the name suggests, the shape has to agree: dim 0 must
divide by the head count, and where a per-head width is known it must match exactly.

Three things it deliberately declines to guess at:

• The output projection stays on the full-matrix path everywhere. Its head structure is on
  the input dimension, and with the usual hidden == num_heads * head_dim splitting dim 0 still
  divides evenly — the shape check cannot catch that one, only the name exclusion can.

• 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

44 CPU-only cases across two files, plus 9 end-to-end.

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, kv count for k/v under
GQA, output projection excluded, MLP matrices named dense not mistaken for attention, fused QKV
skipped, non-attention params untouched, non-divisible shapes skipped, opt-in required, and
use_muon tagging unchanged. Real architectures rather than only synthetic configs: llama,
qwen2, mistral, gpt_neox, falcon, and MLA fixtures whose shapes were measured by instantiating
the attention module on the released GLM-5.2-0.8B-A0.8B and
DeepSeek-Coder-V2-Lite-Instruct configs.

tests/unit/ops/muon/test_per_head_muon_e2e.py — DistributedTest at world_size=2, ZeRO 1/2/3:
tags survive deepspeed.initialize into the ZeRO call sites, the flag is off by default, and
both paths train.

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
44 passed

$ pytest tests/unit/ops/muon/test_per_head_muon_e2e.py     # 2 x H20
9 passed

$ yapf --style .style.yapf --diff  /  flake8 --config .flake8
(clean)

End-to-end training

Requested by @pengdurice; run on 2 × H20 against both models @delock named, at
#8384 (comment). Short version: GLM-5.2 trains under ZeRO-2 and ZeRO-3 and Kimi-K3 under ZeRO-2, with and
without the flag, and the loss curves are indistinguishable. The reason is that both mini
checkpoints carry untrained weights, so their attention heads are interchangeable and the split
has nothing to separate. On
trained checkpoints the imbalance it targets is large — a 6–8x median head-norm spread in the
gradient, which whole-matrix Newton-Schulz only reduces to ~3x and per-head takes to ~1.1x.
Numbers, config, hardware and the two open layout questions are in that comment.

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

[pengdurice]

@alanhuangyoo, let's add end to end training results with a realistic model setup for a realistic training on GPUs with world size > 1 and report the loss etc

[pengdurice]

let's think about if possible not to all gather parameters for heads not useful for this rank.

[alanhuangyoo]

There are no heads that are not useful to this rank in this design, and I want to lay out why rather than just say no.

Every Muon path here partitions by whole parameter, round-robin over ranks, not by head. MuonWithAuxAdam.step takes params[base_i + rank], computes the whole update locally, and the all_gather replicates the updated parameter so each rank has it for the next forward. Stage 3 does the same thing with gradients (stage3.py:1699), and ZeRO-1/2 calls muon_update on the full-shape gradient before narrowing into the flat partition (stage_1_and_2.py:2172). Since each rank runs a full replica of the parameter, it needs all of the heads.

Per-head does not touch any of that. It changes what happens inside muon_update on the rank that already owns the parameter — a view of [out, in] as [heads, head_dim, in] before the same batched NS kernel — so the collective and its volume are identical with the flag on and off.

The thing your question does point at is real, though, and I think it is a separate change: the work partition, not the gather. When a group has fewer parameters than ranks, params_pad pads with torch.empty_like and those ranks compute on padding. Per-head makes a head-level partition natural, because the NS is already batched over a head dimension — heads of one parameter could be split across ranks and the orthogonalized blocks gathered, which would use the idle ranks and cut the per-rank NS cost. That changes the partition for all of Muon rather than only the per-head path, so it wants its own PR and its own measurements. Happy to open an issue for it if you think it is worth pursuing.

[alanhuangyoo]

Correcting the last paragraph of my reply — I measured the follow-up I suggested and it does not hold up, so please do not spend review time on it.

I said a head-level partition would use ranks that the round-robin leaves idle and cut the per-rank Newton-Schulz cost. Both parts are small or zero on real models.

Idle ranks. Only the final chunk can be short, so the wasted rank-slots are (ws - len(params) % ws) % ws out of ceil(len/ws) * ws:

Muon parameters	ws=4	ws=8	ws=16	ws=32
Llama-3-8B, 224	0%	0%	0%	0%
Qwen3-32B, 448	0%	0%	0%	0%
Kimi-K3-0.40B hybrid, 318	0.6%	0.6%	0.6%	0.6%

A transformer has 7 * num_layers Muon matrices, which is divisible by every world size people use. It only bites on a toy model — 14 parameters on 8 ranks is 12.5%.

Load balance. I expected the descending sort plus round-robin to hand rank 0 the largest matrix of every chunk. It does, but a transformer repeats each shape num_layers times, so after sorting each chunk of ws consecutive parameters holds one shape and every rank gets the same work. Modelling the Gram NS cost as the X @ X.mT that builds the Gram matrix, max/mean per-rank work is 1.00x for Llama-3-8B and Qwen3-32B at ws = 2, 4, 8 and 16.

So the redundancy your comment points at is real as a description of the design, but on the models this would run on there is nothing measurable to recover. Where it could still matter is a model whose Muon matrices are genuinely heterogeneous — mixed expert sizes, or a hybrid whose linear-attention projections are much smaller than its MLPs — and I have not measured one of those. If you would like me to, say so and I will; otherwise I would leave the partition alone.

The rest of my reply stands: the collective replicates a whole updated parameter, every rank needs every head of it, and per-head changes neither the collective nor its volume.

[pengdurice]

I think this one is more universal, please search for other places to see if there are any implementations for extracting num of attention / kv heads etc.

[alanhuangyoo]

Agreed, and it was there. Switched to AutoTPMeta.from_model_config (deepspeed/module_inject/tp_shard.py), which is the repo's existing reader: it descends into text_config and probes the spellings models actually use (num_attention_heads, n_head, attention_heads, ...) instead of assuming one attribute name.

That fixed a real gap rather than just deduplicating — my version read num_attention_heads only, so a config spelled n_head was silently untagged. There is a test for it (test_head_count_comes_from_the_shared_extractor).

[pengdurice]

is this compatible with MLA arch?

[alanhuangyoo]

It is now; it was not when you asked. MLA's up-projections are head-blocked but not at head_dim:

• q_b_proj is num_heads x (qk_nope_head_dim + qk_rope_head_dim)
• kv_b_proj is num_heads x (qk_nope_head_dim + v_head_dim)

and the down-projections q_a_proj / kv_a_proj_with_mqa mix latent and rope components with no head structure, so they stay on the full-matrix path. Checked against DeepSeek-V3, GLM-5.2 and Kimi-K3 shapes, including DeepSeek-V2-Lite, which has no q_lora_rank and therefore reaches the same MLA width through a plain q_proj.

That last case is why the tagger no longer picks a branch by which config fields exist: it collects every geometry a leaf name could plausibly have and lets the shape confirm one, which is what @delock asked for in his review.

[pengdurice]

we should do integration test where we do an e2e training loop, look for other examples using SimpleModel for e2e training

[alanhuangyoo]

Added: tests/unit/ops/muon/test_per_head_muon_e2e.py, a real deepspeed.initialize and training loop over ZeRO 1/2/3 — the tags reach the optimizer, the flag is off unless asked for, and training progresses either way.

Beyond the unit tests, the PR body has multi-GPU runs on real checkpoints with world size > 1, which is what you asked for in the thread comment. That writeup also reports a negative result and its cause: the two mini models @delock pointed at ship untrained weights, so their heads are interchangeable and per-head has nothing to act on. On trained checkpoints the gradient carries a 6-8x median head-norm imbalance that the whole-matrix path reduces to ~3x and per-head takes to ~1.1x.

[delock]

Hi @alanhuangyoo, I would suggest to test your implementation against these two mini-MLA models as well. These are the mini-version of the model that use per-head Muon. They are small enough to test your implementation and contains MLA. A correctness test should be enough and convergence would be optional.

You may want to check whether the modeling and setting is consistent with the original model itself in sense of validate per-head Muon implementation. And whether they could be trained without per-head muon with z2 or z3 (baseline)

1 inference-optimization/GLM-5.2-0.8B-A0.8B https://huggingface.co/inference-optimization/GLM-5.2-0.8B-A0.8B
2 inference-optimization/Kimi-K3-0.40B https://huggingface.co/inference-optimization/Kimi-K3-0.40B

[alanhuangyoo]

@pengdurice @delock thanks — all four of the review points are in, plus the mini-model check.
Summary of what changed and what I found.

MLA (@pengdurice's question, @delock's models)

It was not compatible, and the failure mode was a silent skip rather than a wrong split.
Fixed.

MLA does not use q/k/v projections at all; it blocks its two up-projections by head:

q_b_proj    [num_heads * (qk_nope_head_dim + qk_rope_head_dim), q_lora_rank]
kv_b_proj   [num_heads * (qk_nope_head_dim + v_head_dim),       kv_lora_rank]

Neither per-head width is head_dim, so the previous shape check rejected both. That is the
split GLM-5's Muon Split describes, so it is the one worth having.

Checked against the real tensor shapes of both models you named, read from their safetensors
headers rather than assumed:

inference-optimization/GLM-5.2-0.8B-A0.8B — num_attention_heads=16, qk_nope=192,
qk_rope=64, v_head_dim=128:

tensor	shape	tagged	why
q_b_proj	(4096, 512)	16	4096 = 16 × (192 + 64)
kv_b_proj	(5120, 128)	16	5120 = 16 × (192 + 128)
q_a_proj	(512, 2048)	—	down-projection, no head structure
kv_a_proj_with_mqa	(192, 2048)	—	latent + rope, does not split into heads
o_proj	(2048, 2048)	—	heads on the input axis
up_proj	(4096, 2048)	—	MLP

inference-optimization/Kimi-K3-0.40B — worth flagging: its text_config.model_type is
kimi_linear, not MLA. The layers are q_proj/k_proj/v_proj at [256, 1024] alongside
q_conv1d/k_conv1d/v_conv1d, i.e. linear attention with convolution; the q_lora_rank fields in
the config are not used by the modeling code. Those names do match the standard-attention list,
so what keeps it off the per-head path is the shape check: 8 heads × head_dim 74 = 592 ≠ 256.
Correct outcome, but by the guard rather than by recognition — happy to add explicit handling if
you would rather per-head applied there too.

Head extraction (@pengdurice)

You were right that this already exists. Now goes through
AutoTPMeta.from_model_config — the repo's stated "single source of truth for reading kv-head /
hidden / attention-head counts" — which descends into text_config and probes num_heads,
n_head, attention_heads and friends rather than the one hardcoded attribute I had.

End-to-end training (@pengdurice, @delock's baseline point)

tests/unit/ops/muon/test_per_head_muon_e2e.py, DistributedTest with world_size = 2, ZeRO
1/2/3, following the shape of the existing test_muon.py:

• tags survive deepspeed.initialize into the ZeRO call sites — q_proj→8, k_proj/v_proj→2
  under GQA, o_proj and MLP untagged
• per_head_muon off by default: every tag is None and each call site takes the branch it took
  before
• both paths train — that is the without-per-head baseline you asked for; loss decreases in
  each of the six configurations

9 passed        # 3 tests × ZeRO {1,2,3}, world_size=2, on 2×H100

(reduce_scatter: False in those configs, same as the existing Muon suite — Muon rejects
reduce-scatter.)

Not all-gathering unused heads (@pengdurice)

Looked at this and I do not think it belongs in this PR. On the ZeRO-3 path
_apply_distributed_muon_update calls _partitioned_buffers_all_gather to materialize whole
parameters, then each rank processes one whole parameter (params[base_i + rank]). The gather
granularity is the parameter, not the head, so per-head only changes how an already-gathered
matrix is sliced — it adds no communication. Gathering only the heads a rank needs would mean
changing ZeRO-3's parameter-gathering strategy itself, which is a larger change with its own
correctness surface. Happy to open it separately if you want it pursued.

Totals

38 unit tests plus the 9 end-to-end ones. yapf and flake8 clean.

[pengdurice]

@pengdurice @delock thanks — all four of the review points are in, plus the mini-model check. Summary of what changed and what I found.

MLA (@pengdurice's question, @delock's models)

It was not compatible, and the failure mode was a silent skip rather than a wrong split. Fixed.

MLA does not use q/k/v projections at all; it blocks its two up-projections by head:

q_b_proj    [num_heads * (qk_nope_head_dim + qk_rope_head_dim), q_lora_rank]
kv_b_proj   [num_heads * (qk_nope_head_dim + v_head_dim),       kv_lora_rank]

Neither per-head width is head_dim, so the previous shape check rejected both. That is the split GLM-5's Muon Split describes, so it is the one worth having.

Checked against the real tensor shapes of both models you named, read from their safetensors headers rather than assumed:

inference-optimization/GLM-5.2-0.8B-A0.8B — num_attention_heads=16, qk_nope=192, qk_rope=64, v_head_dim=128:

tensor shape tagged why
q_b_proj (4096, 512) 16 4096 = 16 × (192 + 64)
kv_b_proj (5120, 128) 16 5120 = 16 × (192 + 128)
q_a_proj (512, 2048) — down-projection, no head structure
kv_a_proj_with_mqa (192, 2048) — latent + rope, does not split into heads
o_proj (2048, 2048) — heads on the input axis
up_proj (4096, 2048) — MLP
inference-optimization/Kimi-K3-0.40B — worth flagging: its text_config.model_type is kimi_linear, not MLA. The layers are q_proj/k_proj/v_proj at [256, 1024] alongside q_conv1d/k_conv1d/v_conv1d, i.e. linear attention with convolution; the q_lora_rank fields in the config are not used by the modeling code. Those names do match the standard-attention list, so what keeps it off the per-head path is the shape check: 8 heads × head_dim 74 = 592 ≠ 256. Correct outcome, but by the guard rather than by recognition — happy to add explicit handling if you would rather per-head applied there too.

Head extraction (@pengdurice)

You were right that this already exists. Now goes through AutoTPMeta.from_model_config — the repo's stated "single source of truth for reading kv-head / hidden / attention-head counts" — which descends into text_config and probes num_heads, n_head, attention_heads and friends rather than the one hardcoded attribute I had.

End-to-end training (@pengdurice, @delock's baseline point)

tests/unit/ops/muon/test_per_head_muon_e2e.py, DistributedTest with world_size = 2, ZeRO 1/2/3, following the shape of the existing test_muon.py:

• tags survive deepspeed.initialize into the ZeRO call sites — q_proj→8, k_proj/v_proj→2
  under GQA, o_proj and MLP untagged
• per_head_muon off by default: every tag is None and each call site takes the branch it took
  before
• both paths train — that is the without-per-head baseline you asked for; loss decreases in
  each of the six configurations

9 passed        # 3 tests × ZeRO {1,2,3}, world_size=2, on 2×H100

(reduce_scatter: False in those configs, same as the existing Muon suite — Muon rejects reduce-scatter.)

Not all-gathering unused heads (@pengdurice)

Looked at this and I do not think it belongs in this PR. On the ZeRO-3 path _apply_distributed_muon_update calls _partitioned_buffers_all_gather to materialize whole parameters, then each rank processes one whole parameter (params[base_i + rank]). The gather granularity is the parameter, not the head, so per-head only changes how an already-gathered matrix is sliced — it adds no communication. Gathering only the heads a rank needs would mean changing ZeRO-3's parameter-gathering strategy itself, which is a larger change with its own correctness surface. Happy to open it separately if you want it pursued.

Totals

38 unit tests plus the 9 end-to-end ones. yapf and flake8 clean.

Thank you for the response, instead of unit test, would you please run the e2e training and report your config, hardware spec, loss curves etc. I think this change is bit enough to warrant this kind of test.

[alanhuangyoo]

@pengdurice @delock — end-to-end training runs, on the two models you named plus the ones they turned out not to cover.

Four things up front, so you can pick what is worth reading:

1. Both models train, both stages, flag on and off, and the loss curves are indistinguishable. I am not claiming a convergence win from them, and section 4 shows why the numbers do not support one.
2. The reason is the models, not the implementation. Both mini checkpoints ship untrained weights, so their attention heads are interchangeable and the split has nothing to separate. On checkpoints that have been trained the imbalance it targets is a 6–8x median in the gradient, which the whole-matrix path only reduces to ~3x — section 5.
3. I was wrong about Kimi-K3 in my last comment. It is not linear-attention-only; it has two MLA layers and they do exercise this path. Section 2.
4. The MLA coverage had a gap that neither mini model reaches, which I found and fixed while checking this. Section 3.

Setup

GPU	2 × NVIDIA H20-3e, 140 GB each, driver 580.105.08
Host	192 cores, 2 TB RAM
Software	torch 2.9.1+cu128, transformers 5.16.1, DeepSpeed at this branch
Models	inference-optimization/GLM-5.2-0.8B-A0.8B (0.85 B, 6 layers, MLA + MoE) and inference-optimization/Kimi-K3-0.40B (0.42 B language side, 8 layers, MLA + linear attention + MoE + vision)
Data	wikitext-2-raw-v1, packed into 512-token windows, 23 767 lines of real text
Batch	micro 2 × 2 ranks × 1 accumulation step = 4 sequences of 512 tokens
Optimiser	Muon, lr 2e-3, momentum 0.95, grad clip 1.0, bf16, reduce_scatter: False
Steps	600, identical seed / init / data order / schedule between the two arms
Weights	from_pretrained, i.e. the checkpoints as published

The only difference between an "off" and an "on" run is per_head_muon. Section 4 was also run
a second time with a fresh random init instead of the published weights, and the two matrices
agree; the repeat and seed study in that section is from the random-init set, where I have five
runs per arm.

One note in passing, since it cost me a cycle: Kimi-K3-0.40B produces NaN from step 1 under a
fresh from_config init in bf16 — the forward overflows before any optimiser step runs, with the
flag on or off. Loading the published weights it is fine. That is upstream of this PR, but worth
knowing if you run it yourself.

1. Tagging on the real checkpoints

GLM-5.2-0.8B-A0.8B, read off the live model after deepspeed.initialize: 69 Muon parameters, 12 tagged, all of them the MLA up-projections, all at 16 heads.

parameter	shape	result
q_b_proj	(4096, 512)	per-head, 16 × 256 = qk_nope 192 + qk_rope 64
kv_b_proj	(5120, 128)	per-head, 16 × 320 = qk_nope 192 + v_head 128
q_a_proj	(512, 2048)	whole-matrix — down-projection
kv_a_proj_with_mqa	(192, 2048)	whole-matrix — latent + rope
o_proj	(2048, 2048)	whole-matrix — heads on the input axis
MLP / MoE experts / gate / DSA indexer	—	whole-matrix

Worth stating plainly why this model is the right one to check against: its head_dim is 64, so the num_heads * head_dim a tagger would naturally reach for gives 16 × 64 = 1024 against a real 4096. Reading the MLA widths is not a refinement here, it is the difference between working and silently splitting on the wrong boundary.

2. Kimi-K3-0.40B — I got this wrong earlier

I told you it is kimi_linear, not MLA, and that only the shape check keeps it off the per-head path. That was wrong. It is a hybrid, and I had only read the config rather than instantiating it.

linear_attn_config gives the layout: full_attn_layers: [4, 8], kda_layers: [1, 2, 3, 5, 6, 7]. So 2 of the 8 language layers are MLA and 6 are linear attention with convolution. Instantiated, 318 Muon parameters, 4 tagged:

layer kind	parameters	result
MLA (layers 4, 8)	q_b_proj (768, 256), kv_b_proj (1024, 128)	per-head, 8 heads — 8 × 96 and 8 × 128
KDA (6 layers)	q_proj/k_proj/v_proj (256, 1024)	whole-matrix
vision tower	wqkv (4608, 256)	whole-matrix — fused QKV

Layers 3 and 7 zero-indexed, matching full_attn_layers: [4, 8]. So the model does exercise the MLA path, on both of its MLA layers, and the tagging is right there. The KDA rejection is discussed under open questions below — it is not the shape check catching a mismatch, it is a head geometry this code does not read.

It also trains under ZeRO-2, with the flag on and off, from 18.5:

step	ZeRO-2 off	ZeRO-2 on
1	18.0000	18.0000
50	8.9150	8.9075
100	8.1713	8.1475
200	7.6288	7.6300
300	7.4150	7.4088
600	7.0462	7.0438

Last-100 means 7.1075 and 7.1006, within-window sd 0.21 — the same picture as GLM-5.2, for the
same reason.

No ZeRO-3 row for Kimi: stage 3 on this model is far slower here than stage 2 and I did not get a
600-step pair out of it. A 20-step run capped at five minutes produced zero steps, with both GPUs
pinned at 100% rather than idle, so it is grinding rather than deadlocked; a longer run reached
step 1 and had not reached step 2 a minute later. Same with the flag on as with it off, so it is
not this change, and I left it there since ZeRO-2 answers the baseline question. Worth knowing if
you run Kimi-K3 under stage 3 yourself.

3. A gap the two mini models did not cover

Checking the MLA path against other checkpoints turned up one this branch was missing. An MLA block only builds the q_a/q_b pair when q_lora_rank is set; without it the query up-projection is a plain q_proj, and its per-head width is still qk_nope + qk_rope:

# transformers/models/deepseek_v2/modeling_deepseek_v2.py
self.q_proj = (nn.Linear(self.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
               if self.q_lora_rank is None else None)

deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct is in that shape. Instantiating DeepseekV2Attention on its released config under transformers 5.16.1:

heads=16  head_dim=64  qk_nope=128  qk_rope=64  q_lora_rank=None
  q_proj.weight              (3072, 2048)    16 * 192
  kv_b_proj.weight           (4096, 512)     16 * 256
  kv_a_proj_with_mqa.weight  (576, 2048)
  o_proj.weight              (2048, 2048)

transformers populates head_dim = 64 even though the checkpoint's config omits it, so the shape check compared 3072 against 16 × 64 = 1024 and put the model back on the full-matrix path. Nothing was wrong — that is the behaviour before this feature — but a whole class of MLA checkpoints never reached the split. Fixed in c6a7793, with the released config as the test fixture; the two new tests fail on the previous commit and pass on this one.

Note o_proj there: 2048 divides by 16 evenly, so only the name exclusion keeps it off the per-head path. The shape check cannot catch that one.

4. Baselines and loss curves

Both stages train, with and without the flag, which is the baseline question. 25-step trailing
means, since the per-step loss on a 4-sequence batch is very noisy:

step	ZeRO-2 off	ZeRO-2 on	ZeRO-3 off	ZeRO-3 on
1	12.3147	12.3147	12.3147	12.3147
50	8.3168	8.2723	8.2812	8.2891
100	8.0154	8.0026	8.0428	7.9893
200	7.5475	7.5411	7.5814	7.5580
300	7.2224	7.2212	7.2188	7.2001
400	7.1211	7.1222	7.1265	7.1301
500	6.9100	6.9016	6.9023	6.9014
600	6.7657	6.7639	6.7614	6.7594

Mean of the last 100 steps: 6.8483, 6.8478, 6.8477, 6.8418, against a standard deviation of
about 0.29 inside each of those windows.

I am not claiming a convergence benefit from these runs, and the numbers do not support one. The four curves differ by 0.001 to 0.045 at each checkpoint, against a standard deviation of 0.29 inside the window each of those points averages over — the widest gap anywhere is about a sixth of the noise it sits in, and the median gap is a fiftieth of it. To put a number on that rather than eyeball it I reran the same configuration — same seed, same data order, same flag — five times per arm, and swept three seeds.

Mean of the last 100 steps, ZeRO-2, seed 1234, five runs of each arm that differ only in what the GPU does non-deterministically:

	n	mean	sd	range
per-head off	5	6.8436	0.0037	[6.8386, 6.8478]
per-head on	5	6.8411	0.0032	[6.8381, 6.8456]

The difference is +0.0025 against a pooled standard deviation of 0.0035 — 0.71 sd, with the two ranges overlapping across most of their span. Across seeds it does not even hold its sign:

seed	off	on	on − off
1234	6.8452	6.8381	−0.0071
2025	6.8577	6.8506	−0.0072
777	6.8571	6.8686	+0.0115

Two seeds favour per-head, the third reverses by more than either win. On this model, at this scale, the flag is a no-op for convergence, and section 5 is why.

5. Why these models cannot show one, and where it does show

Both mini checkpoints carry untrained weights. Loaded with from_pretrained, GLM-5.2-0.8B-A0.8B scores 12.31 on real text against ln(vocab) = 11.95, and the head-norm spread of its own q_b_proj and kv_b_proj weights is 1.008x and 1.012x. Its heads are interchangeable by construction, so the split has nothing to separate.

Head imbalance does appear once it starts training, but not much of it. Over 400 steps through the real Muon path, measuring what the optimiser actually orthogonalises — the Nesterov blend, not the raw gradient — medians across the run:

	gradient	momentum	after whole-matrix	after per-head
q_b_proj	1.67x	2.15x	1.04x	1.11x
kv_b_proj	1.25x	1.23x	1.07x	1.03x

At that level both paths land in the same place, and per-head is not even consistently the tighter of the two — on q_b_proj it is the looser one, which at these ratios is Newton-Schulz approximation rather than anything about the split. That is the whole explanation for section 4: the mechanism is not being exercised.

It is exercised on checkpoints that have been trained. Same measurement, 8 steps on wikitext, median over all layers:

These are standard-attention checkpoints, not MLA — I could not find a trained MLA model small enough to run here — but the split is the same operation on the same kind of head-blocked matrix.

deepseek-ai/deepseek-coder-1.3b-base — 1.35 B, MHA, 16 heads, loss 3.58 → 3.38

	gradient	after whole-matrix	after per-head
k_proj	8.03x (max 39.08x)	3.10x (max 7.18x)	1.13x (max 2.04x)
q_proj	5.93x (max 23.33x)	2.80x (max 5.99x)	1.12x (max 1.62x)
v_proj	2.23x (max 3.92x)	3.35x (max 6.50x)	1.22x (max 2.04x)

Qwen/Qwen2.5-0.5B — GQA, 14 query heads and 2 KV heads, loss 3.11 → 2.82

	gradient	after whole-matrix	after per-head
q_proj	3.65x (max 11187x)	1.37x (max 418.30x)	1.05x (max 1.66x)
k_proj	1.19x (max 3.32x)	1.02x	1.02x
v_proj	1.25x (max 3.74x)	1.02x	1.02x

Two things there. A trained model's heads are not interchangeable — the gradient carries a 6–8x median imbalance on deepseek-coder, and one Qwen layer has a head so nearly dead that the ratio is four digits. And a fixed number of Newton-Schulz steps on the whole matrix removes only part of it: 8.03x becomes 3.10x, and the Qwen outlier is still 418x after orthogonalisation. Per-head is what takes it to ~1.1x.

Comparing medians of different quantities is not much of an argument, so the same runs paired per layer, 24 layers each:

	spread grows through whole-matrix	spread grows through per-head	per-head tighter than whole-matrix
k_proj	0/24	0/24	24/24
q_proj	0/24	0/24	24/24
v_proj	22/24	0/24	24/24

Per-head is tighter than whole-matrix in every layer of every projection. And v_proj is the sharp case: in 22 of its 24 layers the whole-matrix path comes out more unbalanced across heads than the momentum going in — 2.23x in, 3.35x out at the median.

The GQA row also confirms the KV head count is used where it should be: q_proj tags at 14, k_proj and v_proj at 2, o_proj not at all.

DeepSeek-V2-Lite at 16 B was the closest trained MLA checkpoint and I did not want to pull 30 GB onto this box for it. If either of you has a smaller one, I will run the same measurement against it.

6. Implementation audit

Two properties I checked because a reviewer would, both on the real shapes.

The split does not change how large the update is. If it did, the two arms would run at different effective learning rates and section 4 would be meaningless. Ratio of per-head to whole-matrix update norm:

	GLM q_b	GLM kv_b	Kimi q_b	Kimi kv_b	DSv2 q_proj	Llama q_proj	Llama k_proj
ratio	1.013x	1.077x	0.895x	1.000x	1.070x	0.929x	0.843x

The spread there is Newton-Schulz, not the scaling. This implementation runs a
fixed-coefficient iteration, so it does not drive the singular values to exactly 1 on
either path — both land in a band. Singular values after undoing the fixed scale, averaged
over every head block and over five draws, where 1.0 would be fully orthogonalised:

	whole-matrix	per-head	draws where a singular value falls below 0.1
GLM q_b_proj (4096, 512) → 16 × (256, 512)	0.9157	0.9258	whole 0/5, per-head 0/5
GLM kv_b_proj (5120, 128) → 16 × (320, 128)	0.7902	0.8502	whole 0/5, per-head 0/5
Llama k_proj (1024, 4096) → 8 × (128, 4096)	0.9532	0.8030	whole 0/5, per-head 0/5
Llama q_proj (4096, 4096) → 32 × (128, 4096)	0.8518	0.8031	whole 5/5 (min 0.0001), per-head 0/5

The last row is the one worth flagging. On the square projection the whole-matrix
iteration drops a singular value to ~1e-4 in every draw — the update effectively loses a
direction — while no per-head block goes below 0.68 in any draw. The k_proj row is the
other direction: per-head blocks are thinner, and the iteration leaves them at 0.80 where
the full matrix reaches 0.95, which is where that 0.843x norm ratio comes from.

What that means in practice: on a k_proj-shaped parameter the per-head path takes a step
about 16% smaller than the whole-matrix path would, so turning the flag on slightly
rebalances the effective step size between projections rather than only redistributing it
across heads. It is small and it is one-sided by shape, but it is not nothing, and it is
worth knowing before the flag is recommended to anyone.

I looked at whether raising ns_steps closes the k_proj gap and it does not, in the way
I first assumed. Because the coefficients are fixed rather than adaptive, the mean is not
monotone in the step count — on that shape the whole-matrix path reads 0.9528 at 5 steps
and 0.8292 at 6. So I am reporting the band rather than proposing a step-count change; if
you want the per-head blocks driven closer to 1 it is a coefficient question, not a
ns_steps question, and it is not something I would fold into this PR.

7. Tests

On the same 2 × H20:

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

$ pytest tests/unit/ops/muon/test_per_head_muon_e2e.py
9 passed          # 3 tests x ZeRO {1,2,3}, world_size 2

$ pytest tests/unit/ops/muon/test_muon.py -k "TestMuonConfigs and (adam-1-0.05-128-10-True-False or adam-2-0.01-32-5-True-False or adam-1-0.05-128-10-False-False)"
2 failed, 1 passed          # this branch
2 failed, 1 passed          # upstream master 05daf05, same two tests, same reason

The two TestMuonConfigs failures are the offload_optimizer=True cases, and they are
op_builder.builder.CUDAMismatchException — this box's system CUDA does not match the one torch
was built against, so the CPUAdam op will not build. They fail identically on upstream master, so
they are the environment rather than this change. test_muon.py never sets per_head_muon in
any case, so muon_num_heads is None throughout it and nothing in that file reaches the new
code. A full run of the directory did not finish inside the 15 minutes I gave it, at 55% through.

yapf --style .style.yapf --diff and flake8 --config .flake8 are clean on both changed files.

The tagging tests are no longer only synthetic configs: llama, qwen2, mistral, gpt_neox and
falcon go through the real transformers config classes, and the MLA fixtures use shapes
measured by instantiating the attention module on the released GLM-5.2-0.8B-A0.8B and
DeepSeek-Coder-V2-Lite-Instruct configs. The two tests added in c6a7793 fail on the commit
before it and pass on it.

8. Open questions

Both are real head-blocked matrices in the models you named that this code leaves on the full-matrix path. Neither is a bug — the tagger declines rather than guessing — but both are decisions I would rather you made than assume.

GLM-5.2's DSA indexer. indexer.wq_b is (512, 512), which is index_n_heads 8 × index_head_dim 64, blocked on dim 0:

# transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py
self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)

Its head count lives in index_n_heads, not num_attention_heads, so AutoTPMeta does not see it. indexer.wk is (64, 2048), a single shared K with no head axis, and weights_proj is (8, 2048), one row per head, so neither is a candidate. Only wq_b is. Does Muon Split apply to the sparse-attention indexer in GLM-5, or only to the attention proper?

Kimi-K3's KDA layers. Their q_proj/k_proj/v_proj are (256, 1024) and genuinely head-blocked — linear_attn_config gives num_heads: 8, head_dim: 32, and 8 × 32 = 256. They are rejected here because the geometry is in that nested dict rather than in num_attention_heads/head_dim, which give 8 × 74 = 592. Does per-head Muon apply to linear-attention heads in K3, or is the split specific to full attention?

If the answer to either is yes, the general fix is probably to read the head count off the owning module — most attention implementations carry num_heads/head_dim as attributes, including both of these — rather than adding config keys one architecture at a time. That is a larger change than what is in this PR and I would rather not fold it in without you asking for it.

Not all-gathering unused heads

Unchanged from my previous answer, restating it since it was one of the four points: ZeRO-3 gathers whole parameters (_partitioned_buffers_all_gather) and each rank then processes one whole parameter, so per-head only changes how an already-gathered matrix is sliced and adds no communication. Gathering per head would mean changing ZeRO-3's gathering strategy itself. Happy to open that separately if you want it pursued.

[alanhuangyoo]

Two things since that comment, one of them a mistake of mine.

The decorator

Adding _per_head_orthogonalize above muon_update left @compiler.compile() attached to the new function instead of to muon_update, so muon_update lost the torch.compile it has on master. Unintended, and it applied to everyone using Muon rather than only to the per-head path. Fixed in c3e9334 — the decorated set is back to zeropower_via_newtonschulz5, zeropower_via_gram_newtonschulz and muon_update, matching master, and the helper does not need its own since it is called from inside a compiled function.

The CI failure, which I cannot pin on it

TestGradientAllreduceOpTraining::test[muon-zero1] and [muon-zero2] went red on the previous run — an fp16 gradient comparison between two gradient_allreduce_op settings, deltas around 2e-4 against rtol=1e-3, atol=1e-4.

I expected the decorator to be the cause and it is not, at least not here. On 2 × H20, ten runs:

tree	runs	result
upstream master 05daf05	2	2 passed
this branch, decorator restored	4	2 passed
this branch, exactly as it was pushed when CI failed	3	2 passed
branch merged with master, decorator restored	3	2 passed
branch merged with master, without the fix	2	2 passed

So it reproduces in none of them, including the state that failed, and including the merge CI actually tests. Different accelerator and a different torch build than the runner, so I cannot rule out that it is real there and only shows on that hardware — but I have no evidence it is this change, and the change does not touch the allreduce path or run any new code with per_head_muon off (muon_num_heads is None on every parameter, and muon_update returns before the new branch).

CI is running again on c3e9334. If it comes back red on the same two tests I will keep digging rather than call it flaky.

Does this affect the numbers in the previous comment

No. Every training run in it used the same build for both arms, so the on/off comparison is unaffected, and the head-spread measurements are ratios of 1.1x to 8x where compilation moves the last few digits. But the runs were made without muon_update compiled, which is worth saying plainly rather than leaving implied.

[alanhuangyoo]

Closing the loop on the CI question, since the rerun did not answer it: that run was cancelled at the 1h30m job timeout, not failed. Both of my runs today were, and it is not specific to this branch — the last 30 modal-torch-latest runs are 14 cancelled / 5 failure / 11 success, including runs on master and on ci-modal-merge-queue-only. @delock's #8404 and #8412 are reworking exactly this, so I assume it is known.

Since PR-level GPU CI cannot give the signal right now, I ran the whole file it flagged on 2 × H20, on upstream master and on this branch merged with master:

master 05daf05                 1 failed, 97 passed, 1 skipped   (14:07)
this branch merged with master 1 failed, 97 passed, 1 skipped   (14:30)

Same count, and the same single failure on both: TestZeroOffloadStage1::test, op_builder.builder.CUDAMismatchException — this box's system CUDA does not match the one torch was built against, so CPUAdam will not build. Environment, present on master.

TestGradientAllreduceOpTraining::test[muon-zero1] and [muon-zero2], the two that were red, pass on both. Together with the ten targeted runs in my previous comment that is thirteen runs without a reproduction, including the exact tree CI tests, so I am leaving it there rather than calling it flaky on no evidence — if it reappears on a run that completes, I will pick it up again.

One thing I did verify after restoring the decorator, because it was a gap rather than a formality: muon_update is compiled again, and num_heads is a varying Python int inside a compiled function, which the per-head path had only ever run without. Re-ran both suites against the compiled path on the merged tree.

[alanhuangyoo]

Numbers for that last paragraph, since "re-ran them" is not a measurement.

On the merged tree, against the compiled muon_update: 44 passed for the two unit files, 9 passed for the 2-GPU end-to-end suite.

And the thing I actually wanted to know — num_heads is a Python int that varies, so does compiling muon_update cost a graph per head count? Counting torch._dynamo unique graphs:

  num_heads=None   new graphs: 1
  num_heads=16     new graphs: 1
  num_heads=16     new graphs: 0
  num_heads=8      new graphs: 1
  num_heads=8      new graphs: 0
  num_heads=None   new graphs: 0
  num_heads=32     new graphs: 1
  total unique graphs: 4

One graph per distinct head count, repeats free. A model carries a query count and a KV count, so with the flag on that is at most two graphs beyond the one the whole-matrix path already compiles.

[delock]

Hi @alanhuangyoo , I agree that Kimi architecture support better be done in a seperate PR so we could have better discussion. If you open a seperate issue I'll assign it to you. I'll review and answer the rest of your discussion next week.

[alanhuangyoo]

Opened as #8420. It covers both cases I flagged — Kimi-K3's KDA layers and GLM-5.2's DSA indexer — with the shapes, the lines in each model's own code that produce them, and why the current tagger declines rather than mis-splits them.

I framed the two model questions first, since whether the split is meant to apply there is yours to answer and the code change follows from it either way. Happy to take it if you assign it.

No rush on the rest of the discussion here.

[delock]

Review: per-head Newton-Schulz (thanks for the thorough e2e writeup)

The kernel and the ZeRO plumbing look right to me — per-head reuses the batched NS kernels
behind a view, all six call sites read the tag through the same attribute channel as
use_muon, and communication volume is unchanged. The asks below are all about the tagger
(_attention_head_count), which is where the review weight of this PR sits.

Change requested

1. The dispatch cascade is order-dependent — make it candidate-based

q_proj is now in both _MLA_Q_LEAVES and _QUERY_HEAD_LEAVES, and which branch wins
depends on which config fields happen to exist, not on the model. Kimi-K3's KDA q_proj
is the live example: top-level qk_nope/rope are MLA leftovers, so the MLA branch claims
it and it survives only because 8×96 ≠ 256. That is safety by accident of a width
mismatch, not by recognition.

Suggested shape (the outcome is only num_heads, so two candidates that agree on the head
count are equivalent — ambiguity only exists when counts differ):

candidates = [(heads, width, source), ...]        # mla / kv / q, all evaluated
exact      = [c for c in candidates if rows == heads*width]
one exact match, or several with equal head counts  -> tag
several exact matches with different head counts    -> skip + warn (real ambiguity)
none                                                -> skip

This kills the ordering dependency, turns the Kimi case into principled safety, and is
~15 lines. While you are in there, splitting the function into three named steps would
make each one testable on its own:

classify(leaf) -> kind
geometry_candidates(kind, cfg) -> [(heads, width, source)]
confirm(param, candidates) -> (num_heads | None, reason)

2. An explicit opt-in must not silently degrade — and TP should refuse, not fall back

per_head_muon: true is an explicit request. Today every skip funnels to return None
with no signal at all. The PR description itself says "skipped with a warning", so part of
this is aligning the code with your own stated behavior. Principle: configuration that is
explicitly requested either applies, or reports why it did not — and where it can never
apply, refuses. Concretely:

• Hard error at initialize when the flag is on and zero parameters were tagged.
  This one rule covers the most severe case outright: under tensor parallelism the config
  describes the whole model while each rank holds a shard, so every attention projection
  fails the width check and the feature is silently off model-wide — users believe they
  are training with per-head dynamics and are not. It also catches misconfigurations and
  unrecognized architectures. Fail before training starts; the escape hatch is unsetting
  the flag. The message should list the likely causes (tensor parallelism active,
  architecture not recognized, model with no attention projections). Muon already has
  this culture: it rejects reduce_scatter rather than silently degrading.
• Warning, aggregated per leaf, for non-systemic shape mismatches — a leaf that
  matched an attention table but no candidate confirmed (q_proj×6: width-mismatch).
  This is the unrecognized-layout class (Kimi KDA today, until #8420). Keep it a
  warning, not an error: hybrid models still get per-head on their recognized leaves,
  and erroring would block the working half.
• Info summary whenever the flag is on, aggregated per leaf name —
  per_head_muon: tagged 12/69 Muon params — q_b_proj×6→16, kv_b_proj×6→16 — including
  the declared exclusions (o_proj, fused QKV, MLA down-projections), which are by design
  and should not warn.

The per-leaf skip reasons fall out of the confirm() step in the refactor above, so the
two changes are one piece of work.

3. Docs

per_head_muon has no documentation anywhere. Repo rule: new features ship with docs.
One section in the optimizer docs (flag, what gets tagged, what deliberately does not —
fused QKV, o_proj, MLA down-projections) would also be the right home for the
"shape confirms the name" contract.

4. License header

The three new test files carry # Copyright (c) Microsoft Corporation. New files in this
repo use the two-line header only — # SPDX-License-Identifier: Apache-2.0 + # DeepSpeed Team. Please drop the Microsoft line; a wrong copyright line is a licensing compliance
issue, not a style nit.

Tests requested (with #1)

• Kimi-K3 regression fixture: a non-MLA attention with MLA leftover config fields must
  land on None — pins today's accident.
• Synthetic ambiguity case: two candidates exactly matching with different head counts
  must skip + warn.

Open questions

• Not all-gathering unused heads: agreed this should not gate the PR. For the record,
  the reasoning checks out on my read: ZeRO-3's gather granularity is the whole parameter
  (_partitioned_buffers_all_gather), so per-head changes only how an already-gathered
  matrix is sliced and adds no communication, while gathering just the heads a rank needs
  means reworking ZeRO-3's partitioning contract itself — a separate change with its own
  correctness surface. @pengdurice — you raised this one; does a separate PR work for you,
  or would you rather it be pursued here?

Minor (non-blocking)

• _ns_tolerance / _norm_rtl duplicate the kernels' compute-dtype selection (method →
  dtype, plus the accelerator fallback). If a kernel ever changes precision, the tolerances
  drift silently — either 8× too tight (spurious flakes) or orders of magnitude too loose
  (test goes vacuous). Hoisting a shared ns_compute_dtype(ns_method) into
  original_muon.py, used by both the kernels and the tests, makes the drift structurally
  impossible.
• AutoTPMeta.from_model_config is constructed inside _attention_head_count, i.e. once
  per parameter, though it is loop-invariant — same probe sequence and dataclass rebuilt
  thousands of times for one model. Harmless in practice, but hoisting it above the loop
  in set_optimizer_flags (or threading a prebuilt meta into geometry_candidates) is
  free and reads better; it slots naturally into the refactor above.

Decorator fix (c3e93341) and the non-repro investigation both look right — thanks for
chasing it to thirteen runs.

[alanhuangyoo]

All six in 37740f6. Thanks — #1 and #2 were the right things to push on, and #1 found something I had not noticed.

1. Candidate-based dispatch

You are right that the Kimi case was safety by accident. Worse than that: it was accidental in a way that would have looked correct in review, because the outcome depended on which config fields exist rather than on the model.

classify(leaf)                        -> kind
geometry_candidates(kind, meta, cfg)  -> [(heads, width, source)]
confirm(param, candidates)            -> (num_heads | None, reason)

Agreement is not ambiguity, as you said — the head count is the whole output, so two candidates that confirm with the same count give the same answer. Disagreement is, and the skip names both: ambiguous:head-dim=12/mla-q=8.

Going through it turned up a weaker rule hiding in the old code. With no head_dim, the tagger fell back to divisibility alone — and rows % heads == 0 holds for matrices with no head structure at all. That is exactly how o_proj slipped through in the first version, and it was still there for any config without head_dim. Every candidate now carries a width, derived as hidden_size // num_attention_heads where a config omits head_dim, and a leaf with no width available is declined rather than tagged on divisibility. test_a_config_without_a_per_head_width_is_declined pins it.

2. Report, or refuse

Implemented as you specified. The hard error fires at initialize and names the causes:

per_head_muon is enabled but no attention projection could be tagged. Per-head
Newton-Schulz is therefore inactive for every parameter. Likely causes: tensor
parallelism is active, so each rank holds a shard whose width no longer matches
the config; the architecture's attention layout is not recognized; or the model
has no attention projections. Unset per_head_muon to train without it.
Leaves examined: {'o_proj': 'not-head-blocked', 'q_proj': 'width-mismatch'}

Non-systemic misses are aggregated per leaf as a warning, so a hybrid model keeps per-head on its recognized layers, plus an info line for what was tagged. Your point about the TP case is the one that convinced me it has to be an error rather than a warning: there is no partial result to keep, and the user has no way to notice.

3 & 4. Docs, header

per_head_muon documented in config-json.md — the table of what is tagged and what deliberately is not, the shape-confirms-the-name contract, and the reporting behaviour. The three test files now carry the two-line header.

Tests

Both you asked for, plus four the refactor made worth pinning:

• test_linear_attention_on_a_config_with_mla_leftovers_is_declined — Kimi-K3 hybrid fixture, parametrized over q/k/v
• test_candidates_that_disagree_on_the_head_count_are_skipped — synthetic ambiguity
• test_two_candidates_agreeing_on_the_head_count_are_not_ambiguous
• test_a_config_without_a_per_head_width_is_declined
• test_head_dim_is_derived_when_the_config_omits_it
• test_the_flag_errors_rather_than_silently_doing_nothing

52 unit and 9 end-to-end on 2 × H20, yapf and flake8 clean.

Minor

Both taken. AutoTPMeta.from_model_config is built once per model in set_optimizer_flags; the model-taking _attention_head_count stays for single-parameter callers. ns_compute_dtype(ns_method) moved into original_muon.py and is now the single definition the kernels and the test tolerances both read — the drift you describe is no longer expressible.

Open question

Kimi and the DSA indexer are #8420, which you assigned to me. I will not start it until the two model questions there have an answer, since the code follows from them either way.

[alanhuangyoo]

Note on the red check here, since a red X reads as "this PR broke something" and that is not what happened.

The modal-torch-latest / DeepSpeedAI CI job ended with

RuntimeError: run pytest failed with exit code 137

137 is SIGKILL — the sandbox killed pytest mid-suite. There is not a single FAILED or ERROR line anywhere in the log; the run was killed at roughly 12% of collection, inside tests/unit/v1/..., which neither this PR nor its tests touch. collect tests and DCO pass on the same commit.

The same thing hit #8356 on the same day, and three other branches of mine that ran within the same hour (#8362, #8433, #8435) went green, so it is intermittent rather than a property of this tree. I cannot re-run it — that needs write access to the repo. Any maintainer re-running the failed job should be enough; happy to push an empty commit instead if that is easier.

[alanhuangyoo]

Found a correctness bug in this PR while looking at what happens under tensor parallelism, and it turned into the strongest argument for the feature. Fixed in 943fd03.

The bug

set_optimizer_flags runs at deepspeed/__init__.py:215, before DeepSpeedEngine.__init__ reaches _configure_tensor_parallel. AutoTP's _tp_partition then does params_list[idx].data = _partition — the same Parameter object, a smaller .data — so the tag made against the whole model rides onto a shard.

Nothing catches it, because the stale count usually still divides. Llama, 8 heads of 32, autotp_size: 2, per_head_muon: true:

model.layers.0.self_attn.q_proj.weight  shape=(128, 256)  muon_num_heads=8  per-head width=16 (true 32)
model.layers.0.self_attn.k_proj.weight  shape=(128, 256)  muon_num_heads=8  per-head width=16 (true 32)
model.layers.0.self_attn.v_proj.weight  shape=(128, 256)  muon_num_heads=8  per-head width=16 (true 32)

128 % 8 == 0, so the divisibility check in _per_head_orthogonalize passes and Newton-Schulz runs on 8 blocks of 16 — half of each head — silently. Measured against the whole-model per-head update, that is 29% off.

My description also said tensor parallelism was the likely reason nothing gets tagged. That was wrong for AutoTP in the opposite direction: everything gets tagged, incorrectly. Corrected in the error message and the docs.

The fix, and why it is more than a repair

A column-parallel split is on dim 0, which is the axis the heads are on and the axis per-head Newton-Schulz batches over. A rank therefore holds whole heads, and the per-head width is invariant under the split while the count is not. So: record the width, re-derive the count from the shard after partitioning, and drop any shard whose rows are not a multiple of the width.

The consequence is that per-head Newton-Schulz is exact under tensor parallelism. Not close — equal:

per-head, tp=2, gram      rel diff to the same rows of the whole-matrix result = 0.00e+00
per-head, tp=2, standard                                                        0.00e+00
per-head, tp=4, gram                                                            0.00e+00
per-head, tp=4, standard                                                        0.00e+00

test_per_head_on_a_shard_is_the_shard_of_per_head asserts torch.equal, for both kernels at tp=2 and tp=4.

The whole-matrix path has no such property, and that is what makes the comparison mean something:

whole matrix, tp=2, rank 0  rel diff to the same rows of the whole-matrix result = 0.5273
whole matrix, tp=2, rank 1                                                         0.5234

Where that leaves Muon under AutoTP generally

That second number is not about this PR — it is what Muon already does with autotp_size > 1, and it is a real divergence. Same model, same seed, same batch, one step, q_proj:

gradient   tp=1 vs tp=2   relative difference 7.0e-07     (fp32 all-reduce ordering)
update     tp=1 vs tp=2   relative difference 3.9e-01     cosine 0.936, norm 11% larger

Filed separately as #8437 so it is not tangled up with this PR. With per_head_muon: true the same run is limited by the half-precision floor of the Newton-Schulz iteration itself — a 1e-7 perturbation of the input already moves the gram (fp16) output by 1.2e-3 and the standard (bf16) output by 5.8e-3, flat across three orders of magnitude of input perturbation — rather than by the split.

Tests

tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py (7 cases) pins the resolution: the count follows the shard, an unsharded model is unchanged, a shard that splits a head is dropped while its neighbours keep theirs, untagged parameters are untouched, a model with no tags is a no-op, and the same "an opt-in must not silently do nothing" error as the tagging pass.

tests/unit/ops/muon/test_per_head_muon_under_sharding.py (8 cases) pins the arithmetic above plus an end-to-end deepspeed.initialize with autotp_size: 2 on 2 GPUs, asserting the tags describe the rank's heads and that a step keeps every parameter finite. Reverting only the engine hook fails that one and nothing else.

57 passed across the three per-head files; yapf and flake8 clean.

[alanhuangyoo]

One more from the same thread, and it is why #8438 exists.

Per-head tagging had the same shape problem as use_muon does on master: under deepspeed.zero.Init a partitioned parameter's data is a flat placeholder and the layer's shape is on ds_shape, so param.shape reads as 1-D for every parameter in the model. The width check confirmed nothing, and this PR's error fired on a model whose layout it can read perfectly well:

q_proj: shape=(0,) ndim=1 ds_shape=torch.Size([256, 256]) status=ZeroParamStatus.NOT_AVAILABLE
-> ValueError: per_head_muon is enabled but no attention projection could be tagged

5d16495 reads the layer's shape instead — the same thing _shape_before_zero3_partition already does for AutoTP. ZeRO-3 with zero.Init and per_head_muon: true now tags q/k/v at 8 heads and leaves o_proj and the MLP alone, as it does without zero.Init.

Chasing that turned up a bigger one that is not this PR's: on master, set_optimizer_flags tests p.ndim >= 2 for use_muon, so under zero.Init no parameter is tagged at all and ZeRO-3 finds no sub-group using Muon. Muon never runs; every parameter goes down the AdamW branch and nothing reports it.

	parameters with use_muon	muon_update calls, 6 steps
built normally	14	84
built under zero.Init	0	0

That is #8438, against master, separate from this. It is worth landing on its own regardless of what happens here — zero.Init is how a model that does not fit on one device gets built, which is the case ZeRO-3 is for.

58 passed across the per-head files here; yapf and flake8 clean.

[delock]

@alanhuangyoo thanks for extensive test of muon path. I can help review your subsequent PRs if you would like to go down this path. My bandwidth allows me to review one Muon related PR at a time, I guess that aligns with your plan.

Hi @pengdurice requested your re-review to see if all your comments had been addressed. Thanks!

[alanhuangyoo]

You put that politely, so let me not: seven Muon PRs in two days isn't a plan, it's me opening things as I found them without once thinking about who has to read them. The bugs were worth reporting, but the pacing was mine to get right and I didn't.

So you're not the one deciding what to pick up first — in the order I'd argue for, worst first:

	PR	why
1	#8442	{"optimizer": {"type": "Muon"}} with no zero_optimization block trains with SGD. Stage defaults to 0, Newton-Schulz only runs inside the ZeRO optimizers, so the plainest possible Muon config quietly gets a different optimizer. 20 lines.
2	#8438	Under zero.Init every parameter reads as 1-D, so use_muon is set on none of them and ZeRO-3 finds no Muon sub-group. Muon never runs at all.
3	#8433	Fixes #7746, open since December. Reproduces on master with the reporter's own script — it needs their four ranks, which is why two-rank attempts looked clean.
4	#8435	fp16: the first loss-scale overflow writes NaN into the momentum and the next step writes it back out. The run never recovers.
5	#8440	Muon flattens away the param groups it's handed, so the usual no-weight-decay-on-biases grouping is silently dropped. AdamW keeps it.

Those five are independent of this PR and of each other. #8436 is the only one that has to wait, since it stacks on this branch.

And I'm not adding to the pile — nothing new goes to Muon from me until some of these clear. #8437, #8439 and #8443 are issues with measurements attached rather than PRs, so they need no review time; leave them until the queue's empty.

If a different order suits you better, or you'd rather I close some and reopen them later, just say. I'd rather three of these land than seven sit.

[delock]

@alanhuangyoo Thanks for your ordering, it looks good to me. I don't mean that it is too much, if we find a bug, better record it today than never. Let's do it one at a time and keep the pace.

You put that politely, so let me not: seven Muon PRs in two days isn't a plan, it's me opening things as I found them without once thinking about who has to read them. The bugs were worth reporting, but the pacing was mine to get right and I didn't.

So you're not the one deciding what to pick up first — in the order I'd argue for, worst first:

PR why
1 #8442 {"optimizer": {"type": "Muon"}} with no zero_optimization block trains with SGD. Stage defaults to 0, Newton-Schulz only runs inside the ZeRO optimizers, so the plainest possible Muon config quietly gets a different optimizer. 20 lines.
2 #8438 Under zero.Init every parameter reads as 1-D, so use_muon is set on none of them and ZeRO-3 finds no Muon sub-group. Muon never runs at all.
3 #8433 Fixes #7746, open since December. Reproduces on master with the reporter's own script — it needs their four ranks, which is why two-rank attempts looked clean.
4 #8435 fp16: the first loss-scale overflow writes NaN into the momentum and the next step writes it back out. The run never recovers.
5 #8440 Muon flattens away the param groups it's handed, so the usual no-weight-decay-on-biases grouping is silently dropped. AdamW keeps it.
Those five are independent of this PR and of each other. #8436 is the only one that has to wait, since it stacks on this branch.

And I'm not adding to the pile — nothing new goes to Muon from me until some of these clear. #8437, #8439 and #8443 are issues with measurements attached rather than PRs, so they need no review time; leave them until the queue's empty.

If a different order suits you better, or you'd rather I close some and reopen them later, just say. I'd rather three of these land than seven sit.

[alanhuangyoo]

@pengdurice — a re-review request, not a new argument.

Your changes-requested from 2 Sep is the only thing left on this one (@delock approved on 7 Sep), and I think what you asked for is in. You wanted end-to-end training on a realistic model at world size > 1, with the loss reported. That is the 4 Sep comment above: both mini-MLA models @delock named, 2 GPUs, per-head vs whole-matrix on the same seed and batch, loss curves for each. I also reported the part that did not work — MLA's q_b_proj is not per-head-shaped, so per-head does not apply there, and I said so rather than quietly counting it as covered.

Since then the branch has picked up a correctness fix of its own (943fd03: per-head tagging was computing head count from the pre-shard width, so under tensor parallelism a shard got the wrong number of heads) and a merge with current master. On the merged tree I re-ran everything on 2×H20:

• my per-head suites: 60 + 8 + 9 = 77 passed
• upstream's tests/unit/ops/muon/test_muon.py: 66 failed / 80 passed — identical test names to a pristine master checkout in the same environment, all 66 from one CUDAMismatchException in the op builder (that box's nvcc does not match torch's CUDA). I diffed the two failure lists rather than comparing counts.

If anything from your original comment is still open I would rather hear it than assume it is closed — but if it reads as answered, a re-review would unblock this.

[pengdurice]

let's consolidate the tests a bit. 3 files for one feature is excessive.

[alanhuangyoo]

Done in 789a1e7 — five files down to two, split by what they need to run rather than one per concern:

file	runs on	absorbed
tests/unit/runtime/zero/test_per_head_muon.py	CPU	the arithmetic; which parameters are tagged and with how many heads; re-resolving the count against a tensor-parallel shard
tests/unit/ops/muon/test_per_head_muon_accelerator.py	GPU	a shard of the per-head result is the per-head result of the shard; end-to-end training

Nothing dropped or rewritten — 77 tests before, 77 after, same names. Verified on 2×H20 with both files in one run.

Two mechanical notes:

• The tagging module and the tensor-parallel module each defined a class called _Attn, and they were different classes. The tensor-parallel one is now _ShardedAttn.
• The accelerator file is not also called test_per_head_muon.py. These directories have no __init__.py, and two same-named modules break collection when both are selected in one run — I tried it first and got an import error, so the name carries the distinction instead.

@pengdurice — this was the only thing outstanding from your side as far as I can tell; the end-to-end training results you asked for on 2 Sep are in the 4 Sep comment above. If anything else is still open I would rather hear it than assume it is closed.