muon

muon is the Muon optimizer (Jordan et al., 2024) with a grouped Newton-Schulz kernel, loadable through kernels. The reference baseline is an fp64 Newton-Schulz iteration, against which the grouped and per-matrix paths agree bitwise with each other and sit at or below reference bf16 error.

Muon orthogonalizes the momentum of every 2D parameter with a Newton-Schulz iteration before each step, and that iteration is the optimizer's cost: five steps of three matrix products per matrix. Run one matrix at a time, those products are too small to fill a large GPU, so the device idles between launches. This kernel runs the iteration for a whole shape-bucket of parameters in one batched tensor-core call, and rounds the bf16 write-back stochastically so bf16 parameters train with no fp32 master copy.

256 matrices orthogonalized in one call: the batch's singular-value spectrum migrates to the sigma = 1 line over five Newton-Schulz steps (88.5% inside [0.7, 1.3] at step five, bf16), in 2 ms against 131 ms for the per-matrix loop.

Usage

import torch
from kernels import get_kernel

muon = get_kernel("phanerozoic/muon", version=1, trust_remote_code=True)

# hidden 2D weights use Muon; embeddings, the LM head, and 1D params use AdamW
hidden = [p for n, p in model.named_parameters() if p.ndim == 2 and "embed" not in n and "lm_head" not in n]
rest   = [p for n, p in model.named_parameters() if p not in set(hidden)]
opt = muon.Muon([
    {"params": hidden, "use_muon": True,  "lr": 0.02, "weight_decay": 0.01},
    {"params": rest,   "use_muon": False, "adamw_lr": 3e-4},
])

for batch in loader:
    opt.zero_grad(); loss(model(batch)).backward(); opt.step()

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. The orthogonalization is also exposed directly:

O = muon.orthogonalize(G, steps=5)          # single matrix [R, C] -> bf16
Ob = muon.orthogonalize(G_stack)            # [G, R, C] batched in one call

API

Method

For a matrix G, Muon computes an approximate orthogonal factor with the quintic iteration X <- aX + (bA + cA^2)X, A = X X^T, coefficients (3.4445, -4.7750, 2.0315), five steps, after scaling to unit Frobenius norm and transposing so the contracted dimension is the larger one. The optimizer buckets its 2D parameters by shape and stacks each bucket into one [G, R, C] tensor; every Newton-Schulz step is then three batched GEMMs (cublasGemmStridedBatchedEx, bf16 operands, fp32 accumulate) with the two elementwise combines fused. The scaled update and decoupled weight decay are a single kernel; bf16 write-back rounds with counter-based stochastic rounding keyed by the element's global index, so the expected value is exact and sub-ULP updates accumulate. 1D parameters and any group marked use_muon=False take a standard AdamW path.

Measured

Newton-Schulz over a set of parameter matrices, grouped kernel against the same iteration per matrix in eager PyTorch and under torch.compile (max-autotune), median of repeated executions. Each cell is speedup over eager / over compile.

The advantage grows with the GPU, since batching is what keeps a large device fed; torch.compile fuses the per-matrix iteration but does not batch across distinct parameters. On the L4, matrices of 768 and larger already fill the card and grouping is at parity; the gain there is in many-small-matrix regimes (adapters, small models, mixture-of-experts).

Correctness

Measured on H200 and RTX PRO 6000 through get_kernel; the grouped and single-matrix paths are bitwise identical (torch.equal).

• The grouped iteration is the reference iteration: in fp64, grouped and per-matrix Newton-Schulz agree to 5.8e-15.
• Against an fp64 ground truth over 49 cases (seven shapes by seven conditionings), the kernel's bf16 error is 0.65x the reference bf16 error in the median and never above 0.80x; singular values are closer to the truth than the reference's (0.17 vs 0.22 worst) and the output direction matches to a cosine above 0.978.
• A 20M-parameter transformer trained 600 steps with kernel-Muon and with a reference-Muon keeps a loss relative gap under 1.2e-3 at every step.
• The stochastic-rounding write-back is unbiased and bitwise reproducible for a fixed (seed, step). Results are deterministic.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ and torch with bf16 tensor cores.
• Muon is for the 2D hidden weights; keep embeddings, the LM head, and 1D parameters in an AdamW group.
• Optimizer state per Muon matrix is one fp32 momentum buffer; the parameter may stay bf16 with no fp32 master.
• base_index sets the global-index offset so a parameter updated whole or across shards rounds identically.

References

Jordan et al., "Muon: An optimizer for the hidden layers of neural networks" (2024); Bernstein and Newhouse, "Old Optimizer, New Norm" (2024); Liu et al., "Muon is Scalable for LLM Training" (Moonshot, 2025); Amsel et al., "The Polar Express: Optimal Matrix Sign Methods" (2025).

License

Apache-2.0.