[Feature Request] Enable Muon Optimizer for Large-Scale Pre-Training

[gali-leilei]

Background

Muon is a momentum-based optimizer that applies Newton-Schulz orthogonalization to the gradient before the update step. It has been validated at scale:

• Kimi Moonlight: 3B/16B-parameter MoE model trained on 5.7T tokens — validated Muon with weight decay, distributed Muon, and hyperparameter matching with AdamW
• Kimi K2: 1040B-parameter MoE model (32B active) trained on 15.5T tokens — added QK clipping for attention stability

The core idea: for matrix-shaped parameters (attention projections, MLP weights), replace the Adam update with a Nesterov-SGD step where the gradient is first passed through Newton-Schulz iteration to approximate msign(G) = U @ diag(sign(σᵢ)) @ Vᵀ. Non-matrix parameters (embeddings, output logits, norms, biases) continue to use AdamW.

What Needs to Change

1. Config Schema (torchtitan/components/optimizer.py)

The current OptimizersContainer.Config has a single lr and AdamW-specific fields. Muon requires a mixed-optimizer config with separate hyperparameters for each group:

@dataclass(kw_only=True, slots=True)
class Config(Configurable.Config):
    name: str = "AdamW"  # or "MuonWithAdamW"
    lr: float = 8e-4              # AdamW lr (for embeddings, norms, etc.)
    lr_muon: float = 0.02         # Muon lr (for matrix params)
    beta1: float = 0.9
    beta2: float = 0.95
    eps: float = 1e-8
    weight_decay: float = 0.1
    muon_weight_decay: float = 0.0
    muon_momentum: float = 0.95
    ns_steps: int = 5             # Newton-Schulz iteration count
    ns_a: float = 3.4445
    ns_b: float = -4.7750
    ns_c: float = 2.0315

2. Parameter Group Splitting

Currently the optimizer receives a flat parameter list (optimizer.py:119). Muon requires splitting parameters into two groups:

def _split_param_groups(model: nn.Module) -> tuple[list, list]:
    """Returns (muon_params, adam_params)."""
    muon_params, adam_params = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if (
            param.ndim == 2
            and not any(skip in name for skip in ["tok_embeddings", "output", "norm"])
        ):
            muon_params.append(param)
        else:
            adam_params.append(param)
    return muon_params, adam_params

Per Keller Jordan's empirical findings, applying Muon to wq, wk, wv separately (rather than as a fused QKV matrix) yields better results. torchtitan's Llama model already stores them as separate nn.Linear modules (attention.py:428-482), so this is naturally satisfied.

3. Newton-Schulz Orthogonalization

def newtonschulz5(G: Tensor, steps: int = 5, eps: float = 1e-7,
                  a: float = 3.4445, b: float = -4.7750, c: float = 2.0315) -> Tensor:
    """Approximately compute msign(G) = U @ diag(sign(σ)) @ Vᵀ via quintic iteration."""
    assert G.ndim == 2
    X = G.bfloat16()
    X = X / (X.norm() + eps)
    if G.size(0) > G.size(1):
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X
    if G.size(0) > G.size(1):
        X = X.T
    return X.to(G.dtype)

The polynomial (a, b, c) = (3.4445, -4.7750, 2.0315) was chosen by Keller Jordan for fast convergence; Kimi Moonlight uses these same values.

4. Distributed Muon (FSDP2 Compatibility)

This is the critical engineering challenge. Under FSDP2, parameters are DTensor objects sharded across the data-parallel mesh. The Newton-Schulz step operates on the full gradient matrix, so it must run after gradient all-reduce.

Two approaches to evaluate:

Option A: All-gather before NS

• All-gather sharded gradient into the full matrix on each rank
• Apply Newton-Schulz locally on the full matrix
• Scatter the update back
• Similar to torch.optim.Muon in PyTorch 2.9 (pytorch#148819)

Option B: Per-rank NS + all-reduce

• Apply NS to the local gradient shard
• All-reduce the NS-transformed gradient
• Kimi's Moonlight paper approach — avoids the all-gather overhead

Note: As of PyTorch 2.9, torch.optim.Muon does not yet support distributed training (confirmed in pytorch#148819). A custom distributed implementation is required.

5. Checkpoint Compatibility

torchtitan uses get_optimizer_state_dict / set_optimizer_state_dict with flatten_optimizer_state_dict=True (optimizer.py:140-157). For a mixed-optimizer setup:

• Muon state per matrix param: momentum_buffer (1st-order only, no 2nd moment)
• AdamW state per non-matrix param: exp_avg, exp_avg_sq, step

The flattened format should still work since state is keyed by parameter FQN, but this requires testing.

6. FP8 Interaction

Do not apply FP8 GEMM to the Newton-Schulz computation. Kimi K2 explicitly disables FP8 for Muon steps (confirmed officially). The post_optimizer_hook registered for Float8 should be excluded from Muon parameter groups.

Recommended Hyperparameters

Based on Kimi Moonlight and Keller Jordan's experiments:

Hyperparameter	Recommended Value	Notes
Muon LR	0.02	~25× larger than AdamW LR of 8e-4
Muon momentum	0.95	Nesterov-style
Muon weight decay	0.0	Kimi Moonlight; AdamW handles it for non-Muon params
NS steps	5	Good accuracy/cost trade-off
(a, b, c)	(3.4445, -4.7750, 2.0315)	Keller Jordan values; Kimi Moonlight uses same

Proposed Work Items

• Design MuonWithAdamW optimizer class that wraps two internal optimizers with parameter group routing
• Implement Newton-Schulz orthogonalization kernel (bfloat16, numerically stable)
• Implement distributed Muon: evaluate Option A (all-gather) vs Option B (per-rank NS + all-reduce) for FSDP2
• Add MuonWithAdamW to _resolve_optimizer_cls() and extend config schema
• Test checkpoint save/load with mixed optimizer state under FSDP2 + DCP
• Validate exclusion of FP8 post-hook from Muon parameters
• Benchmark convergence vs AdamW baseline on Llama3 8B / 70B
• Document hyperparameter matching between Muon and AdamW (learning rate scaling)

References

• Keller Jordan's reference implementation and blog: https://kellerjordan.github.io/posts/muon/
• Kimi Moonlight paper (Muon at 16B scale): https://arxiv.org/pdf/2502.16982
• Kimi K2 technical report (Muon at 1040B scale): https://github.com/MoonshotAI/Kimi-K2/blob/main/tech_report.pdf
• PyTorch 2.9 Muon (single-node): https://docs.pytorch.org/docs/2.9/generated/torch.optim.Muon.html
• PyTorch distributed Muon tracking issue: addition of muon optimizer to torch.optim pytorch#148819

[gali-leilei]

duplicate of issue #148819