Muon

class tensorplay.optim.Muon(params, lr: float = 0.001, weight_decay: float = 0.1, momentum: float = 0.95, nesterov: bool = True, ns_coefficients: tuple[float, float, float] = (3.4445, -4.775, 2.0315), eps: float = 1e-07, ns_steps: int = 5, adjust_lr_fn: str | None = None)

Implements Muon algorithm.

\[ \begin{align}\begin{aligned}\begin{split}\begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \gamma \text{ (lr)},\ \lambda \text{ (weight decay)},\ \mu \text{ (momentum)},\ \textit{nesterov}\in\{True,False\},\\ &\hspace{13mm}(a,b,c)\ \text{ (NS coefficients)},\ \varepsilon \text{ (epsilon)},\ k \text{ (NS steps)},\ \theta_0 \text{ (params)},\ f(\theta) \text{ (objective)} \\ &\textbf{initialize} : B_0 \leftarrow 0 \text{ (momentum buffer)} \\[-1.ex] &\rule{110mm}{0.4pt} \\ &\textbf{for}\ t=1\ \textbf{to}\ \ldots\ \textbf{do} \\[0.25ex] &\hspace{5mm} g_t \leftarrow \nabla_{\theta} f_t(\theta_{t-1}) \\[0.25ex] &\hspace{5mm} B_t \leftarrow \mu B_{t-1} + g_t \\[0.25ex] &\hspace{5mm} \widetilde{B}_t \leftarrow \begin{cases} g_t + \mu B_t, & \text{if nesterov}=True \\ B_t, & \text{if nesterov}=False \end{cases} \\[1.0ex] &\hspace{5mm} O_t \leftarrow \mathrm{NS}^{(a,b,c)}_{k}\!\big(\widetilde{B}_t;\ \varepsilon\big) \\[0.5ex] &\hspace{5mm} \theta_t \leftarrow \theta_{t-1} - \gamma\,\lambda\,\theta_{t-1} \quad\text{(decoupled weight decay)} \\[0.25ex]\end{split}\\\begin{split} &\hspace{5mm} \gamma \leftarrow \mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big) \\[0.25ex] &\hspace{5mm} \theta_t \leftarrow \theta_t - \gamma\, O_t \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\mathbf{return}\ \theta_t \\[-1.ex] &\rule{110mm}{0.4pt}s \end{aligned}\end{split}\end{aligned}\end{align} \]

Here, \(\mathrm{NS}^{(a,b,c)}_{k}(\cdot;\varepsilon)\) denotes \(k\) iterations of the Newton–Schulz orthogonalization operator parameterized by coefficients \((a,b,c)\) with numerical stabilization \(\varepsilon\).

The purpose for \(\mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big)\) is to make the orthogonalized update scale consistently across rectangular matrices.

Keller’s original implementation scales the update by \(\sqrt{\max\!\left(1, \frac{A}{B}\right)}\), where \(A\) and \(B\) are dimensions of the matrix being optimized, which represent fan-out and fan-in for a Linear weight matrix.

Moonshot’s implementation focuses on matching \(RMS\) of AdamW. The adjustment is computed as: \(\gamma \leftarrow {0.2}\gamma\,\sqrt{\max\!\left({A}, {B}\right)}\) The method is adopted from Muon is Scalable for LLM Training. Research results show that with this adjustment Muon can directly reuse the learning rate and weight decay tuned for AdamW.

Jeremy Bernstein in Deriving Muon proposes a scaling condition on the spectral norm, which scales the update by \(\sqrt{\frac{A}{B}}\). This is similar to the Keller’s “original” implementation but removes clamping down to 1.

We provide these options for the learning rate adjustment: “original”, which follows Keller’s implementation, “match_rms_adamw”, which refers to Moonshot’s implementation, and “spectral_unclamped”, which matches Bernstein’s implementation. If adjust_lr_fn is not specified, the default is “original”.

For further details regarding the algorithm we refer to Muon: An optimizer for hidden layers in neural networks, Muon is Scalable for LLM Training, and Deriving Muon.

Parameters:

• Other (params. Note that Muon is an optimizer for 2D parameters of neural network hidden layers.) – parameters, such as bias, and embedding, should be optimized by a standard method such as AdamW.

• lr (float, Tensor, optional) – learning rate (default: 1e-3).

• weight_decay (float, optional) – weight decay (L2 penalty). (default: 0.1)

• momentum (float, optional) – momentum factor (default: 0.95)

• nesterov (bool, optional) – enables Nesterov momentum. Only applicable when momentum is non-zero

• ns_coefficients (tuple of three floats, optional) – coefficients (a,b,c) for the Newton–Schulz orthogonalization polynomial (default: (3.4445, -4.775, 2.0315))

• eps (float, optional) – term added to the denominator for numerical stability. (default: 1e-07)

• ns_steps (int, optional) – number of Newton–Schulz iteration steps. (default: 5)

• adjust_lr_fn (str, optional) – function to adjust learning rate. One of “original”, “match_rms_adamw”, and “spectral_unclamped”. If not specified, we will default to use “original”. (default: None)

Example

>>> # xdoctest: +SKIP
>>> # Muon only supports 2D params; use a standard optimizer
>>> # such as AdamW for biases, embeddings, and other non-2D
>>> # parameters.
>>> muon_params = [
...     p for p in model.parameters() if p.ndim == 2
... ]
>>> other_params = [
...     p for p in model.parameters() if p.ndim != 2
... ]
>>> optim_muon = torch.optim.Muon(
...     muon_params, lr=0.02, momentum=0.95
... )
>>> optim_adamw = torch.optim.AdamW(
...     other_params, lr=3e-4, weight_decay=0.01
... )
>>> optim_muon.zero_grad()
>>> optim_adamw.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optim_muon.step()
>>> optim_adamw.step()

step(closure=None)[source]

Performs a single optimization step.