[Feature] Per-Head Muon: per-head Newton-Schulz orthogonalization for attention projections

[delock]

Motivation

DeepSpeed now supports the Muon optimizer across ZeRO stages 1/2/3. In 2026, two frontier labs independently converged on the same refinement for attention projections: orthogonalizing per attention head rather than over the full projection matrix.

• Kimi K3 (arXiv:2607.24653, §2.5, "Per-Head Muon"): partitions the momentum of Q/K/V projections along the head dimension and orthogonalizes each head's block separately. Their motivation: full-matrix orthogonalization treats all heads as a single coupled block, so heads with larger gradient/momentum scales dominate the shared update direction while smaller-scale heads receive insufficiently normalized updates. Per-head orthogonalization equalizes update scale across heads, improves stability at larger scales, and is cheaper (Newton–Schulz on tall per-head blocks costs less than on the full projection).
• Zhipu GLM-5 "Muon Split" (arXiv:2602.15763): the same idea applied to MLA up-projection matrices (split by attention head, orthogonalize each head independently), closing the performance gap between MLA and GQA under Muon.

Today in DeepSpeed, muon_update() (deepspeed/runtime/zero/muon/original_muon.py) applies Newton–Schulz to the full 2D parameter matrix — i.e., the "coupled block" behavior both papers describe.

Proposal

1. Optimizer kernel: add a per-head mode to muon_update — reshape the update/momentum to [..., num_heads, head_dim] and run the existing bmm-based Newton–Schulz on the batched per-head blocks (the NS implementations already support batched inputs, so this is mostly a view/reshape branch). Apply the existing max(1, m/n)**0.5 scaling per head block.
2. Parameter metadata pipeline: tag attention projection parameters with head structure (num_heads, head_dim, fused-vs-split QKV layout), similar to the existing use_muon tagging. The AutoTP model scan / presets already identify attention head structure, so this metadata can be sourced from there.
3. ZeRO integration: the engine call site (post-unflatten, per-parameter muon_update) stays unchanged — just pass the head shape so the update can be viewed per-head and reshaped back before being applied to the flat buffer.
4. Config surface: opt-in switch, e.g. optimizer.params.per_head_muon: true (or per param-group), following the precedent of separate muon_lr / adam_lr.
5. Tests: unit tests for per-head vs full-matrix orthogonalization equivalence on synthetic attention-shaped params, plus convergence sanity on a small model.

References

• Kimi K3 technical report §2.5: https://arxiv.org/abs/2607.24653
• Muon Split (GLM): https://arxiv.org/abs/2602.15763
• Current Muon implementation: deepspeed/runtime/zero/muon/original_muon.py, deepspeed/runtime/zero/muon/muon_optimizer.py, ZeRO call site in deepspeed/runtime/zero/stage_1_and_2.py

Happy to help implement this.

[alanhuangyoo]

I would like to take this one, if nobody is on it yet.

Read through the current implementation to check the shape of the work. Confirming your point 1:
both Newton–Schulz kernels in original_muon.py are already batch-safe — assert G.ndim >= 2,
.mT, size(-2)/size(-1), and norm(dim=(-2, -1), keepdim=True) — and muon_update already
has a batched branch with per-block scaling:

if is_expert_group:
    scale = max(1, update.size(-2) / update.size(-1))**0.5
    update = ns_fn(update, steps=ns_steps) * scale

which is the same shape the per-head path needs. So the kernel side really is a view/reshape
plus reusing that branch, and the substance is in 2–4: sourcing num_heads / head_dim and the
fused-vs-split QKV layout, threading it to the call site, and the config surface.

Two things I would want your steer on before I write much:

1. Where the head metadata comes from. You mention the AutoTP model scan. Would you rather it
   be read off AutoTPMeta (which now carries num_kv_heads / num_attention_heads after
   [AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta #8241), or tagged onto the parameter next to use_muon so non-AutoTP users can set it too?
   The second is more code but does not tie per-head Muon to AutoTP being enabled.

2. Fused QKV. For a fused [3 * num_heads * head_dim, hidden] projection, I assume the split
   should be per head within each of Q/K/V rather than treating the fused matrix as
   3 * num_heads uniform heads — and with GQA the K/V blocks have a different head count than
   Q. Is that the layout you have in mind, or would you rather the first cut only handle split
   QKV and skip fused?

On validation: I have 8×H100 available, so beyond the unit tests I can run a small-model
convergence comparison (full-matrix vs per-head, matched seeds/steps) and report loss curves
rather than just asserting shapes.

[alanhuangyoo]

@delock implementation is up in #8384, built to the five points in your proposal.

Kernel, metadata, ZeRO integration and the config switch are all there; per_head_muon is
opt-in and off by default. Verified rather than asserted:

• Arithmetic — batched NS equals orthogonalizing each head alone, both NS methods, tolerance
  derived from the kernel's own compute dtype rather than tuned to pass.
• Behaviour — with one head's gradient scaled 100×, per-head and full-matrix must differ and
  the four head update norms land within 1.5× of each other. This is the case that fails if
  num_heads is ignored.
• Recognition — 5 real architectures. llama/qwen2/mistral tag q_proj→8, k_proj/v_proj→2
  under GQA; gpt_neox/falcon tag nothing, which is where a naive name match would wrongly catch
  dense_h_to_4h.
• End to end — 2×H100, ZeRO-1, bf16, real steps: q=8 k=2 v=2 o=None mlp=None embed=None.
• Mechanism — as head gradient scales diverge, full-matrix update norms spread out
  (max/min 1.009 → 1.266) while per-head stays flat at 1.015.

Two questions from my earlier comment are still open and shape the remaining work, so I left
both on the conservative side for now: fused QKV stays on the full-matrix path, and o_proj is
excluded because its head structure is on the input axis. Happy to extend either once you say
which layout you want assumed.