Skip to content
Bud's Wiki

Addition of Muon optimizer to torch.optim — PyTorch team welcomes PR (Issue #148819)

TL;DR

Section titled “TL;DR”

GitHub feature-request issue (filed March 2025) asking the PyTorch team to add Keller Jordan’s Muon optimizer to torch.optim. The pointed-at comment (3070108227, July 15 2025) is the PyTorch maintainers’ response: they’re welcoming a PR, explicitly citing Moonshot’s MuonClip / Kimi K2 deployment as evidence that Muon is now scale-validated and worth shipping in the core library. Subsequent activity led to torch.optim.Muon shipping in PyTorch 2.12 (no distributed support yet — that work continues in torchtitan #2494).

Key claims

Section titled “Key claims”
  • Issue requests adding Muon to torch.optim, citing faster, more stable convergence than Adam at small scale.
  • The July-15-2025 maintainer comment welcomes a PR specifically in light of Moonshot’s MuonClip variant scaling Muon to Kimi K2 — i.e. the social proof for inclusion is “MuonClip worked at frontier MoE scale,” not the original CIFAR-10 / NanoGPT speedrun result.
  • As of PyTorch 2.9 / 2.12, torch.optim.Muon exists but does not support distributed training; users are expected to wrap it themselves for FSDP/HSDP/DCP setups (per the followup torchtitan tracking issue #2494).
  • Kimi K2 explicitly disables FP8 GEMM for the Newton–Schulz orthogonalization step inside Muon, and uses a MuonWithAdamW-style hybrid (Muon for 2D matrix params, AdamW for scalars/vectors/embeddings/LM head).

Method

Section titled “Method”

The issue itself is a single-paragraph feature request linking to Jordan’s Muon write-up and his reference implementation. The comment that @amit posted is the maintainers’ acceptance reply, inviting a PR. The actual mechanics of Muon are unchanged from the original recipe — momentum SGD, then orthogonalize the resulting matrix update via 5 quintic Newton–Schulz iterations with the (3.4445, -4.7750, 2.0315) coefficients, applied only to 2D weight matrices.

What’s new and wiki-relevant is the scaling story implicit in the comment thread: between the issue being filed (March 2025) and the maintainer green-light (July 2025), Moonshot shipped Kimi K2 with MuonClip — Muon plus a qk_clip post-step that rescales Q/K projections to bound attention logits — and demonstrated zero loss spikes through 5.7T → ≈15T tokens. This is the empirical event that moved Muon from “interesting speedrun curiosity” to “ship it in core.”

Results

Section titled “Results”

No quantitative claims in the comment itself; this is an infrastructure-tracking artifact. The downstream facts that justify filing it:

  • Muon (without clip) hit modded-NanoGPT speedruns ≈2.6× over the previous record at 124M scale.
  • MuonClip scaled to Kimi K2 (1T-param MoE, 32B-active) with no loss spikes.
  • torch.optim.Muon lands in PyTorch ≈2.12 (single-device only at filing time).

Why it’s interesting

Section titled “Why it’s interesting”

This is the canonical “research-lab optimizer becomes framework default” moment for Muon, and it’s the event that licenses many of the follow-on Muon papers already filed on the wiki — Muon is Not That Special: Random or Inverted Spectra Work Just as Well (claiming the specific spectra don’t matter), The Newton-Muon Optimizer (Newton-Muon variant), Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon (hardware-aware NS5 for Muon), Aurora: A Leverage-Aware Optimizer for Rectangular Matrices (leverage-aware extension), and the SSO line in Controlled LLM Training on Spectral Sphere (arguing Muon is only “half-aligned” with µP because it controls update spectra but not weight spectra). It also closes the loop with Kimi K2: Open Agentic Intelligence — the K2 announcement that the comment cites.

The fact that distributed-Muon is still unsolved (per torchtitan #2494, which calls out that K2’s release explicitly disables FP8 for the NS step) is the open infra problem that several of the cited follow-ups try to chip away at.

See also

Section titled “See also”