Implementation:Hiyouga LLaMA Factory Muon Optimizer

Knowledge Sources	Hiyouga_LLaMA_Factory
Domains	Optimization, Training
Last Updated	2026-02-06 19:00 GMT

Overview

Implements the Muon optimizer (MomentUm Orthogonalized by Newton-Schulz) for training large transformer models with orthogonalized weight updates.

Description

The Muon optimizer combines SGD with momentum and a Newton-Schulz orthogonalization post-processing step. For each 2D parameter update, the gradient is replaced with the nearest orthogonal matrix computed via the zeropower_via_newtonschulz5 function, a quintic Newton-Schulz iteration that runs stably in bfloat16 on GPU. The iteration uses optimized coefficients (a=3.4445, b=-4.7750, c=2.0315) that maximize the slope at zero for rapid convergence. Learning rates are scaled by 0.2 * sqrt(max(A, B)) based on the parameter matrix dimensions. Parameters that are not 2D, or are identified as embedding/head layers, are optimized using an internal AdamW implementation with standard bias correction. The optimizer is based on research from Moonshot AI's Moonlight project and Keller Jordan's original Muon implementation.

Usage

Use this optimizer as an alternative to AdamW when training large transformer models from scratch. It is particularly suited for training with large batch sizes. Note that it may not work well for fine-tuning pretrained models or with small batch sizes.

Code Reference

Source Location

• Repository: Hiyouga_LLaMA_Factory
• File: src/llamafactory/third_party/muon/muon.py
• Lines: 1-226

Signature

def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int) -> torch.Tensor

class Muon(torch.optim.Optimizer):
    def __init__(
        self,
        lr=1e-3,
        wd=0.1,
        muon_params=None,
        momentum=0.95,
        nesterov=True,
        ns_steps=5,
        adamw_params=None,
        adamw_betas=(0.9, 0.95),
        adamw_eps=1e-8,
    )

    def adjust_lr_for_muon(self, lr: float, param_shape: list[int]) -> float
    def step(self, closure=None) -> Optional[float]

Import

from llamafactory.third_party.muon.muon import Muon

I/O Contract

Inputs

Name	Type	Required	Description
muon_params	iterable	Yes	Parameters to optimize with Muon (must be 2D tensors)
lr	float	No	Learning rate; updates will have spectral norm of lr (default: 1e-3)
wd	float	No	Weight decay coefficient (default: 0.1)
momentum	float	No	Momentum for internal SGD (default: 0.95)
nesterov	bool	No	Whether to use Nesterov-style momentum (default: True)
ns_steps	int	No	Number of Newton-Schulz iteration steps (default: 5)
adamw_params	iterable	No	Parameters to optimize with the internal AdamW fallback
adamw_betas	tuple[float, float]	No	Beta coefficients for internal AdamW (default: (0.9, 0.95))
adamw_eps	float	No	Epsilon for internal AdamW numerical stability (default: 1e-8)

Outputs

Name	Type	Description
loss	Optional[float]	The loss value if a closure is provided, otherwise None

Usage Examples

from llamafactory.third_party.muon.muon import Muon

# Separate 2D weight parameters from others
muon_params = [p for name, p in model.named_parameters() if p.ndim == 2 and "embed" not in name and "head" not in name]
adamw_params = [p for name, p in model.named_parameters() if p.ndim != 2 or "embed" in name or "head" in name]

optimizer = Muon(
    muon_params=muon_params,
    adamw_params=adamw_params,
    lr=0.02,
    momentum=0.95,
    nesterov=True,
    ns_steps=5,
)

Related Pages

• Hiyouga_LLaMA_Factory_DPO_Workflow - Training workflows that can use custom optimizers