dtensor-muon 0.1.0

DTensor-based Distributed Muon Optimizer

dtensor-muon

A distributed-ready implementation of the Muon optimizer built on PyTorch DTensor. It runs the orthogonalization step efficiently across sharded parameters (FSDP / tensor-parallel meshes) and falls back to Adam/AdamW for the parameters Muon doesn't apply to — all from a single optimizer instance.

Unlike the original distributed implementation, when used with FSDP2, there is zero communication overhead as the optimizer can run orthogonalization directly over the local sharded parameters.

What is Muon?

Muon updates 2D+ parameters by taking the momentum-smoothed gradient and replacing it with its nearest semi-orthogonal matrix (the "zero-power" or matrix-sign of the gradient) before applying it. The orthogonalization is computed iteratively so it stays cheap on the GPU. This implementation provides two iteration schemes:

• newton_schulz — the classic Newton-Schulz iteration, with a fused Triton Gram-matrix kernel.
• polar_express — the Polar Express scheme (arXiv:2505.16932), which uses precomputed coefficients for faster convergence.

Features

• DTensor / distributed first — orthogonalization is run across the device mesh, with a dedicated fast path for parameters sharded only along dim 0 (FSDP).
• Unified Muon + Adam — one optimizer handles both. Parameters Muon can't update (1D tensors, embeddings, the LM head, etc.) are routed to a fused Adam/AdamW path via per-group configuration.
• Three variants:
  • Muon — the reference per-parameter implementation.
  • MuonForeach — batched foreach operations for higher throughput.
  • MuonLP — Experimental quantized (4-bit / 8-bit / fp8) optimizer states via torchao for reduced memory.
• Cautious weight decay — an extension beyond the original Muon (which has no cautious variant): weight decay is applied only where the update and parameter share a sign (u * p > 0), following the Cautious Optimizers technique. On by default; set use_cautious_wd=False for plain (non-cautious) decoupled weight decay.
• Nesterov momentum, automatic shape-based learning-rate scaling, optional torch.compile, and Triton kernels.

Requirements

• Python ≥ 3.12
• PyTorch ≥ 2.12.1
• Triton (for the fused Newton-Schulz kernel)
• torchao (optional, only for MuonLP)

Installation

# Install from PyPI
uv pip install dtensor-muon

# Install it directly from the repository:
uv pip install git+https://github.com/tonyf/dtensor-muon.git

# include torchao for the low-precision optimizer (MuonLP)
uv pip install "dtensor-muon[lp] @ git+https://github.com/tonyf/dtensor-muon.git"

Development

Clone the repository and sync the environment:

git clone https://github.com/tonyf/dtensor-muon.git
cd dtensor-muon
uv venv
source .venv/bin/activate
uv sync                  # core install
uv sync --extra lp       # include torchao for the low-precision optimizer

Usage

Muon is applied to 2D+ weight matrices, while norms, biases, and embeddings are typically left to Adam. Configure this with param groups using the "algorithm" key:

import torch
from dtensor_muon import Muon

model = ...

muon_params = [p for n, p in model.named_parameters() if p.ndim >= 2 and "embed" not in n]
adam_params = [p for n, p in model.named_parameters() if p.ndim < 2 or "embed" in n]

optimizer = Muon(
    [
        {"params": muon_params},                      # "muon" is the default
        {"params": adam_params, "algorithm": "adamw"},
    ],
    lr=1e-3,
    wd=0.1,
)

# standard training loop
for batch in dataloader:
    loss = model(batch).loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Choosing a variant

from dtensor_muon import Muon, MuonForeach, MuonLP   # MuonLP requires the `lp` extra

MuonForeach and MuonLP are drop-in replacements with the same constructor. Use MuonForeach for throughput and MuonLP to shrink optimizer-state memory.

Selecting an orthogonalization strategy

optimizer = Muon(params, orthogonalization_strategy="polar_express")  # or "newton_schulz"

Key options

Option	Default	Description
lr	1e-3	Learning rate.
wd	0.1	Weight decay.
use_cautious_wd	True	Cautious weight decay — apply decay only where update and param share a sign (u * p > 0). An addition not present in the original Muon; set False for plain weight decay.
momentum	0.95	Muon momentum.
nesterov	True	Use Nesterov momentum.
ns_steps	5	Orthogonalization iteration steps.
orthogonalization_strategy	"polar_express"	"newton_schulz" or "polar_express".
adam_betas	(0.9, 0.95)	Betas for the Adam path.
is_adamw	True	Decoupled (AdamW) vs. coupled weight decay for the Adam path.
fused_adam	True	Use the fused Adam kernel.
compile	False	torch.compile the per-parameter step.

Most options can also be overridden per param group.

Benchmarks

A self-contained benchmark suite lives in benchmark/. Run the whole thing (kernels, single-device optimizers, and the distributed orthogonalization paths) with:

uv run python benchmark/run.py            # full run, writes benchmark/RESULTS.md
uv run python benchmark/run.py --quick    # small shapes, fast smoke

It degrades gracefully: CUDA-only sections are skipped on a CPU-only host, and the distributed section uses NCCL/CUDA when ≥2 GPUs are visible and falls back to gloo/CPU otherwise. Full numbers and methodology are in benchmark/RESULTS.md.

The snapshot below was measured on 2× NVIDIA RTX PRO 6000 Blackwell, PyTorch 2.12.1+cu130.

Triton Gram kernel vs x @ x.mT — a clear win, 1.5–1.9× across shapes/dtypes (e.g. (32, 2048, 1024) bf16: 1.13 ms → 0.61 ms, 1.85×).

Orthogonalization loops (5 steps, bf16) — the fused Triton loop is at parity with or slightly faster than the torch.compiled PyTorch loop (≈0.96–1.08×), reflecting the Gram-kernel win net of the rest of the iteration (the non-symmetric B @ X matmul is a plain matmul in both). The uncompiled eager loop is ~1.3× slower than either.

Single-device optimizer step (48 weight matrices, full step()), vs a naive reference Muon:

optimizer	step (ms)	vs naive
naive (Newton-Schulz)	158	1.00× (ref)
Muon	260–280	0.56–0.61×
MuonForeach	257–270	0.59–0.62×

On this Blackwell card the naive eager loop is still faster than the torch.compile + Triton paths at these shapes (native bf16 matmul is very fast, so per-step() Python and compile overhead dominate); MuonForeach's batching roughly matches per-parameter Muon. (MuonLP targets optimizer-state memory, not step time, so it isn't in this table.)

Distributed orthogonalization (2× GPU, NCCL, params sharded on dim 0):

path	per-call (ms)	vs single-device
single-device (replicated, no collectives)	42.6	1.00× (ref)
FSDP fast path (foreach_zeropower_3d_fsdp)	20.4	2.09×
general path (foreach_zeropower + redistribute)	33.2	1.28×

The FSDP fast path works on each rank's local shard with no collectives — ~2.1× the single-device baseline (batch split across ranks) and ~1.6× faster than the general redistribute path. This is the core payoff of the DTensor design.

Project layout

src/dtensor_muon/
├── optim/            # Muon, MuonForeach, MuonLP optimizers
├── orthogonalize/    # zeropower dispatch, Newton-Schulz & Polar Express, DTensor handling
├── kernels/          # Triton Gram-matrix kernel
└── utils/            # DTensor and foreach helpers

Testing and linting

uv run pytest        # run the test suite
uv run ruff check    # lint
uv run ty            # type check