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

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.

Start by reading deepspeed/runtime/zero/muon/original_muon.py and muon_optimizer.py to understand muon_update and its Newton–Schulz paths, then inspect the ZeRO call site in deepspeed/runtime/zero/stage_1_and_2.py. Trace how existing use_muon metadata and AutoTP attention structure are propagated. Done means an opt-in per-head mode works through ZeRO and is covered by synthetic orthogonalization and small-model convergence tests.