Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

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

Open
  • Is something wrong?
#8,367 5 comments 4 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

@0z5a is already working on this.

Since Sep 18, 2026.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
45/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python
Domain
distributed-systems, machine-learning, performance

Research direction

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.

Written by the indexing model from the issue text.

Description

enhancement training

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

Happy to help implement this.

Dominant language
Python
Stars
43.3k
Forks
5.1k
Avg merge
4d 21h
Merged PRs (30d)
66

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepspeedai/DeepSpeed

All issues in deepspeedai/DeepSpeed

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.