muon

muon is the Muon optimizer (Jordan et al., 2024) with a grouped Newton-Schulz kernel, loadable through kernels. Muon orthogonalizes the momentum of every 2D parameter with a Newton-Schulz iteration before each step, and the cost of the optimizer is that iteration: 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 sits idle between launches. This kernel runs the iteration for a whole shape-bucket of parameters in a single batched tensor-core call (bf16 operands, fp32 accumulate), and rounds the bf16 write-back stochastically so bf16 parameters train with no fp32 master copy.

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. bf16 parameters in a Muon group are written back with stochastic rounding, so no fp32 master copy is needed. 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] -> all orthogonalized in one call

API

How it works

For a matrix G, Muon computes NS(G), an approximate orthogonal factor, with the quintic iteration X <- a X + (b A + c A^2) X, A = X X^T, coefficients (3.4445, -4.7750, 2.0315), five steps, after scaling X to unit Frobenius norm and transposing so the contracted dimension is the larger one. The singular values are pulled toward one, loosely by design.

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) over the whole bucket rather than G separate launches, with the two bf16-elementwise combines fused. The scaled update and decoupled weight decay are a single kernel; for a bf16 parameter the write-back is rounded 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 (biases, norm gains) and any group marked use_muon=False (embeddings, the LM head) take a standard decoupled-weight-decay AdamW path.

Measured

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

The advantage grows with the GPU, because batching is what keeps a large device fed: on the H200 and the RTX PRO 6000 the grouped iteration is faster at every size measured, including the large matrices where per-matrix GEMMs are compute bound. torch.compile fuses the per-matrix iteration but does not batch across distinct parameters, so the grouped call is faster than the compiled per-matrix loop as well. The L4 has few SMs, so a matrix of 768 or larger already fills it and grouping is at parity there; the gain on that class of GPU is in the many-small-matrix regimes (adapters, small models, mixture-of-experts).

Correctness

Equivalence to reference Muon is established at three levels, measured on H200 and RTX PRO 6000 through get_kernel. The kernel's grouped and single-matrix paths are bitwise identical (torch.equal), so the shape-bucket is not an approximation of the per-matrix result.

• The grouped iteration is the reference iteration. Run in fp64, grouped and per-matrix Newton-Schulz agree to 5.8e-15, the fp64 rounding floor, so batching a shape-bucket does not change the computation.
• The kernel is at least as accurate as the reference. Against an fp64 ground-truth Newton-Schulz over 49 cases (seven shapes by seven conditionings: Gaussian, ill-conditioned to 1e-6, low-rank, heavy-tailed, scaled by 1e+-3, near-orthogonal), the kernel's bf16 error is 0.65x the reference bf16 error in the median and never above 0.80x; its singular values are closer to the truth than the reference's (0.17 vs 0.22 worst) and the output direction matches the reference to a cosine above 0.978. Any kernel-versus-reference difference is bf16 rounding, and it favors the kernel.
• It holds dynamically. A 20M-parameter six-layer transformer trained 600 steps with kernel-Muon and with a reference-Muon differing only in the Newton-Schulz implementation keeps a loss relative gap under 1.2e-3 at every step; the two descend together while the parameters decorrelate only at the bf16 rate expected of two low-precision runs of one optimizer.

The stochastic-rounding write-back is unbiased (a sub-ULP update across 200k bf16 elements averages to the true value, where round-to-nearest leaves it unchanged) 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, as in the standard recipe.
• Optimizer state for a Muon matrix parameter is one fp32 momentum buffer, and the parameter may stay bf16 with no fp32 master, so a bf16 matrix parameter costs 4 bytes per element less than the bf16-master recipe.
• base_index sets the global-index offset for sharded updates 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" and the modular-duality view of Muon as steepest descent under the spectral 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.