feat(core): add the Muon optimizer with an AdamW fallback

@@ -445,6 +445,38 @@ Optimizers     name: rmsprop opt: <class 'torch.optim.rmsprop.RMSprop'>     name: rprop opt: <class 'torch.optim.rprop.Rprop'>     name: novograd opt: <class 'nemo.core.optim.novograd.Novograd'>+    name: muon opt: <class 'nemo.core.optim.muon.Muon'>++``muon`` (Muon with an AdamW fallback) decides per parameter by shape:+a matrix below ``adam_dim_threshold`` in both dimensions takes the Muon update, every other parameter AdamW.+The shape alone does not tell an embedding from a hidden matrix,+so exclude small-vocabulary embeddings and tied output projections explicitly with ``use_muon`` on their parameter group,+in Python or, for a model attribute, through ``optim_param_groups``:++.. code-block:: python++    from nemo.core.optim import Muon++    optimizer = Muon(+        [+            {"params": hidden_matrices},+            {"params": [model.embedding.weight], "use_muon": False, "weight_decay": 0.0},+        ],+        lr=0.02,+        weight_decay=0.01,+    )++.. code-block:: yaml++    model:+        optim:+            name: muon+            lr: 0.02+            weight_decay: 0.01+        optim_param_groups:+            embedding:+                use_muon: false+                weight_decay: 0.0  Optimizer Params ~~~~~~~~~~~~~~~~

@@ -21,6 +21,7 @@     AdamaxParams,     AdamParams,     AdamWParams,+    MuonParams,     NovogradParams,     OptimizerParams,     RMSpropParams,

@@ -28,6 +28,7 @@     'AdamaxParams',     'AdagradParams',     'AdamWParams',+    'MuonParams',     'RMSpropParams',     'RpropParams', ]@@ -235,6 +236,37 @@ class AdafactorParams(OptimizerParams):     warmup_init: bool = False  +@dataclass+class MuonParams(OptimizerParams):+    """+    Configuration of the Muon optimizer with an AdamW fallback (``nemo.core.optim.muon.Muon``).++    Args:+        lr: learning rate of the Muon update+        momentum: momentum factor of the gradient buffer+        nesterov: apply the momentum buffer in Nesterov form+        ns_steps: number of Newton-Schulz iterations+        weight_decay: decoupled weight decay+        scale_mode: shape factor of the orthogonalized update, ``'fan_in'``, ``'fan_out'`` or ``'none'``+        adam_betas: coefficients of the AdamW moving averages+        adam_eps: term added to the AdamW denominator+        adam_lr_ratio: the AdamW learning rate is ``lr * adam_lr_ratio``+        adam_dim_threshold: matrices with a dimension of at least this size use AdamW+        use_muon: force or exclude the Muon update; None selects by shape+    """++    momentum: float = 0.95+    nesterov: bool = True+    ns_steps: int = 5+    weight_decay: float = 0+    scale_mode: str = 'fan_in'+    adam_betas: Tuple[float, float] = (0.9, 0.95)+    adam_eps: float = 1e-8+    adam_lr_ratio: float = 0.025+    adam_dim_threshold: int = 8192+    use_muon: Optional[bool] = None++ def register_optimizer_params(name: str, optimizer_params: OptimizerParams):     """     Checks if the optimizer param name exists in the registry, and if it doesnt, adds it.@@ -293,4 +325,5 @@ def get_optimizer_config(name: str, **kwargs: Optional[Dict[str, Any]]) -> Optim     'rmsprop_params': RMSpropParams,     'rprop_params': RpropParams,     'adafactor_params': AdafactorParams,+    'muon_params': MuonParams, }

@@ -34,6 +34,7 @@     WarmupPolicy,     prepare_lr_scheduler, )+from nemo.core.optim.muon import Muon from nemo.core.optim.novograd import Novograd from nemo.core.optim.optimizer_with_main_params import MainParamsOptimizerWrapper from nemo.core.optim.optimizers import get_optimizer, parse_optimizer_args, register_optimizer

@@ -0,0 +1,314 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.+# SPDX-License-Identifier: Apache-2.0+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++"""Muon optimizer (MomentUm Orthogonalized by Newton-Schulz) with an AdamW fallback.++Muon (Keller Jordan et al., 2024, https://kellerjordan.github.io/posts/muon/) updates a weight matrix+with an approximately semi-orthogonal matrix derived from the momentum-averaged gradient,+computed with a quintic Newton-Schulz iteration in bfloat16.+Parameters that are not matrices (biases, normalization scales, convolution filters)+and matrices with a dimension of at least ``adam_dim_threshold`` (large embeddings and output projections)+are updated with AdamW at a fixed fraction of the learning rate,+as the public reference trainings (modded-nanogpt, nanochat) do with their embeddings and heads.+An embedding below the threshold is a matrix like any other and takes the Muon update+unless its parameter group sets ``use_muon=False``, see :class:`Muon`.+One optimizer instance handles both kinds,+so one learning-rate schedule drives both with a fixed ratio.+"""++from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union++import torch+from torch.optim.optimizer import Optimizer++__all__ = ['Muon', 'update_scale', 'zeropower_via_newtonschulz5']++SCALE_MODES = ('fan_in', 'fan_out', 'none')+++class Muon(Optimizer):+    """Muon for weight matrices with an AdamW fallback for the other parameters, see the module docstring.++    Which parameters take the Muon update is decided per parameter:+    with the default ``use_muon=None`` a parameter takes it+    when it is a matrix whose larger dimension is below ``adam_dim_threshold``,+    so biases, normalization scales, convolution filters and wide embeddings or output projections fall back to AdamW.+    The shape alone does not tell an embedding from a hidden matrix:+    a small-vocabulary embedding (``1026 x 512``) is below the default threshold and takes the Muon update,+    so put embeddings and tied output projections that are to stay on AdamW+    into a parameter group with ``use_muon=False`` (and usually ``weight_decay=0.0``)::++        Muon([{'params': hidden}, {'params': [embedding.weight], 'use_muon': False, 'weight_decay': 0.0}], lr=0.02)++    A parameter group can force the choice either way with ``use_muon=True`` or ``use_muon=False``;+    a forced parameter with more than two dimensions (a convolution filter)+    is orthogonalized as the matrix ``(out_channels, -1)``.++    The orthogonalized update of a ``(rows, columns)`` matrix is multiplied by a shape-dependent factor.+    ``scale_mode='fan_in'`` (the default) uses ``sqrt(max(1, rows / columns))``,+    which for the ``(out, in)`` layout of ``torch.nn.Linear``+    gives every update the root mean square ``1 / sqrt(fan_in)`` (the modded-nanogpt convention).+    ``'fan_out'`` uses ``sqrt(max(1, columns / rows))``, the root mean square ``1 / sqrt(fan_out)``,+    which is the same rule applied to the transposed ``(in, out)`` layout of other frameworks.+    ``'none'`` applies no factor.+    The same matrix in the two layouts receives the same update, transposed, when the two modes are paired,+    up to the bfloat16 rounding of the Newton-Schulz iteration.++    Weight decay is decoupled:+    it multiplies the parameter by ``1 - lr * weight_decay`` for Muon parameters+    and by ``1 - lr * adam_lr_ratio * weight_decay`` for AdamW parameters,+    so the relative decay per step is the same fraction of the respective learning rate;+    set ``weight_decay: 0.0`` on a parameter group to exclude it.++    Args:+        params: iterable of parameters or of parameter-group dicts+        lr: learning rate of the Muon update+        momentum: momentum factor of the gradient buffer+        nesterov: apply the momentum buffer in Nesterov form+        ns_steps: number of Newton-Schulz iterations+        weight_decay: decoupled weight decay+        scale_mode: ``'fan_in'``, ``'fan_out'`` or ``'none'``, see above+        adam_betas: coefficients of the AdamW moving averages+        adam_eps: term added to the AdamW denominator+        adam_lr_ratio: the AdamW learning rate is ``lr * adam_lr_ratio``+        adam_dim_threshold: matrices with a dimension of at least this size fall back to AdamW+        use_muon: force (``True``) or exclude (``False``) the Muon update; ``None`` selects by shape+    """++    def __init__(+        self,+        params: Union[Iterable[torch.Tensor], Iterable[Dict[str, Any]]],+        lr: float = 2e-2,+        momentum: float = 0.95,+        nesterov: bool = True,+        ns_steps: int = 5,+        weight_decay: float = 0.0,+        scale_mode: str = 'fan_in',+        adam_betas: Tuple[float, float] = (0.9, 0.95),+        adam_eps: float = 1e-8,+        adam_lr_ratio: float = 0.025,+        adam_dim_threshold: int = 8192,+        use_muon: Optional[bool] = None,+    ) -> None:+        if lr < 0.0:+            raise ValueError(f"Invalid learning rate: {lr}")+        if not 0.0 <= momentum < 1.0:+            raise ValueError(f"Invalid momentum: {momentum}")+        if ns_steps < 1:+            raise ValueError(f"Invalid number of Newton-Schulz steps: {ns_steps}")+        if weight_decay < 0.0:+            raise ValueError(f"Invalid weight decay: {weight_decay}")+        if not (0.0 <= adam_betas[0] < 1.0 and 0.0 <= adam_betas[1] < 1.0):+            raise ValueError(f"Invalid AdamW betas: {adam_betas}")+        if adam_eps < 0.0:+            raise ValueError(f"Invalid AdamW epsilon: {adam_eps}")+        if adam_lr_ratio < 0.0:+            raise ValueError(f"Invalid AdamW learning rate ratio: {adam_lr_ratio}")+        if adam_dim_threshold < 1:+            raise ValueError(f"Invalid AdamW dimension threshold: {adam_dim_threshold}")+        if scale_mode not in SCALE_MODES:+            raise ValueError(f"Unknown scale_mode {scale_mode!r}, expected one of {SCALE_MODES}")+        defaults = dict(+            lr=lr,+            momentum=momentum,+            nesterov=nesterov,+            ns_steps=ns_steps,+            weight_decay=weight_decay,+            scale_mode=scale_mode,+            adam_betas=tuple(adam_betas),+            adam_eps=adam_eps,+            adam_lr_ratio=adam_lr_ratio,+            adam_dim_threshold=adam_dim_threshold,+            use_muon=use_muon,+        )+        super().__init__(params, defaults)+        for group in self.param_groups:+            if group['scale_mode'] not in SCALE_MODES:+                raise ValueError(f"Unknown scale_mode {group['scale_mode']!r}, expected one of {SCALE_MODES}")+            if group['use_muon']:+                for param in group['params']:+                    if param.ndim < 2:+                        raise ValueError(+                            f"use_muon=True needs at least two dimensions, got shape {tuple(param.shape)}"+                        )++    @staticmethod+    def selects_muon(param: torch.Tensor, *, use_muon: Optional[bool], adam_dim_threshold: int) -> bool:+        """Whether ``param`` takes the Muon update under the given group options.++        Args:+            param: the parameter+            use_muon: the group's ``use_muon`` option, see the class docstring+            adam_dim_threshold: the group's ``adam_dim_threshold`` option++        Returns:+            True for the Muon update, False for the AdamW fallback+        """+        if use_muon is not None:+            return bool(use_muon)+        return param.ndim == 2 and max(param.shape) < adam_dim_threshold++    @torch.no_grad()+    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:+        """Performs a single optimization step.++        Args:+            closure: re-evaluates the model and returns the loss++        Returns:+            the loss returned by ``closure``, or None+        """+        loss = None+        if closure is not None:+            with torch.enable_grad():+                loss = closure()++        for group in self.param_groups:+            for param in group['params']:+                if param.grad is None:+                    continue+                state = self.state[param]+                if self.selects_muon(+                    param, use_muon=group['use_muon'], adam_dim_threshold=group['adam_dim_threshold']+                ):+                    _muon_update(+                        param,+                        param.grad,+                        state,+                        lr=group['lr'],+                        momentum=group['momentum'],+                        nesterov=group['nesterov'],+                        ns_steps=group['ns_steps'],+                        weight_decay=group['weight_decay'],+                        scale_mode=group['scale_mode'],+                    )+                else:+                    _adamw_update(+                        param,+                        param.grad,+                        state,+                        lr=group['lr'] * group['adam_lr_ratio'],+                        betas=group['adam_betas'],+                        eps=group['adam_eps'],+                        weight_decay=group['weight_decay'],+                    )+        return loss+++def zeropower_via_newtonschulz5(grad: torch.Tensor, steps: int = 5) -> torch.Tensor:+    """Approximate the orthogonal factor of a matrix with a quintic Newton-Schulz iteration.++    For the singular value decomposition ``grad = U S V^T`` the result approximates ``U V^T``:+    the iteration pushes the singular values towards one, not exactly onto it,+    with the modded-nanogpt coefficients,+    which are tuned for a fast start of the iteration rather than for convergence to exactly orthogonal factors.+    The iteration runs in bfloat16 on the wide orientation of the matrix (rows <= columns),+    so transposing the input transposes the result;+    the result has the dtype of the input.++    Args:+        grad: a matrix of shape (rows, columns), any float dtype+        steps: number of Newton-Schulz iterations++    Returns:+        the approximately semi-orthogonal matrix, same shape and dtype as ``grad``+    """+    if grad.ndim != 2:+        raise ValueError(f"Expected a matrix, got a tensor of shape {tuple(grad.shape)}")+    a, b, c = 3.4445, -4.7750, 2.0315+    x = grad.bfloat16()+    transpose = grad.size(0) > grad.size(1)+    if transpose:+        x = x.T+    x = x / (x.norm() + 1e-7)+    for _ in range(steps):+        aa = x @ x.T+        bb = b * aa + c * (aa @ aa)+        x = a * x + bb @ x+    if transpose:+        x = x.T+    return x.to(grad.dtype)+++def update_scale(shape: Sequence[int], scale_mode: str) -> float:+    """Factor on the orthogonalized update of a ``(rows, columns)`` matrix.++    Args:+        shape: the two dimensions of the matrix+        scale_mode: ``'fan_in'``, ``'fan_out'`` or ``'none'``, see :class:`Muon`++    Returns:+        the factor+    """+    rows, columns = shape+    if scale_mode == 'fan_in':+        return max(1.0, rows / columns) ** 0.5+    if scale_mode == 'fan_out':+        return max(1.0, columns / rows) ** 0.5+    if scale_mode == 'none':+        return 1.0+    raise ValueError(f"Unknown scale_mode {scale_mode!r}, expected one of {SCALE_MODES}")+++def _muon_update(+    param: torch.Tensor,+    grad: torch.Tensor,+    state: Dict[str, Any],+    *,+    lr: float,+    momentum: float,+    nesterov: bool,+    ns_steps: int,+    weight_decay: float,+    scale_mode: str,+) -> None:+    if 'momentum_buffer' not in state:+        state['momentum_buffer'] = torch.zeros_like(param)+    buf = state['momentum_buffer']+    buf.mul_(momentum).add_(grad)+    update = grad.add(buf, alpha=momentum) if nesterov else buf+    matrix = update if update.ndim == 2 else update.reshape(update.size(0), -1)+    orthogonal = zeropower_via_newtonschulz5(matrix, steps=ns_steps)+    if weight_decay != 0:+        param.mul_(1 - lr * weight_decay)+    param.add_(orthogonal.reshape_as(param), alpha=-lr * update_scale(matrix.shape, scale_mode))+++def _adamw_update(+    param: torch.Tensor,+    grad: torch.Tensor,+    state: Dict[str, Any],+    *,+    lr: float,+    betas: Tuple[float, float],+    eps: float,+    weight_decay: float,+) -> None:+    if 'exp_avg' not in state:+        state['step'] = 0+        state['exp_avg'] = torch.zeros_like(param)+        state['exp_avg_sq'] = torch.zeros_like(param)+    state['step'] += 1+    step = state['step']+    beta1, beta2 = betas+    exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']+    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)+    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)+    bias_correction1 = 1 - beta1**step+    bias_correction2 = 1 - beta2**step+    denom = (exp_avg_sq.sqrt() / (bias_correction2**0.5)).add_(eps)+    if weight_decay != 0:+        param.mul_(1 - lr * weight_decay)+    param.addcdiv_(exp_avg, denom, value=-lr / bias_correction1)

@@ -27,6 +27,7 @@ from nemo.core.config.optimizers import OptimizerParams, get_optimizer_config, register_optimizer_params from nemo.core.optim.adafactor import Adafactor from nemo.core.optim.adan import Adan+from nemo.core.optim.muon import Muon from nemo.core.optim.novograd import Novograd  from nemo.utils.model_utils import maybe_update_config_version@@ -43,6 +44,7 @@     'novograd': Novograd,     'adafactor': Adafactor,     'adan': Adan,+    'muon': Muon, }  try:

@@ -0,0 +1,281 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.+# SPDX-License-Identifier: Apache-2.0+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++import copy++import pytest+import torch++from nemo.core import config, optim+from nemo.core.optim.muon import Muon, update_scale, zeropower_via_newtonschulz5+from nemo.core.optim.optimizers import AVAILABLE_OPTIMIZERS++LR = 0.02+WEIGHT_DECAY = 0.01+++def _parameters(seed: int = 0) -> torch.nn.ParameterDict:+    """Parameters of every kind the optimizer distinguishes, in a fixed order."""+    generator = torch.Generator().manual_seed(seed)++    def param(*shape):+        return torch.nn.Parameter(torch.randn(*shape, generator=generator))++    # assigned one by one: ParameterDict sorts the keys of a plain dict+    parameters = torch.nn.ParameterDict()+    parameters['hidden'] = param(48, 32)  # Muon+    parameters['tall'] = param(64, 16)  # Muon, fan_in scale 2+    parameters['bias'] = param(48)  # AdamW+    parameters['conv'] = param(8, 4, 3, 3)  # AdamW by shape+    parameters['embedding'] = param(8200, 16)  # AdamW, wide+    return parameters+++def _set_gradients(parameters: torch.nn.ParameterDict, seed: int) -> None:+    generator = torch.Generator().manual_seed(seed)+    for param in parameters.values():+        param.grad = torch.randn(param.shape, generator=generator)+++def _run(optimizer: Muon, parameters: torch.nn.ParameterDict, steps: int, first_seed: int = 100) -> None:+    for step in range(steps):+        _set_gradients(parameters, first_seed + step)+        optimizer.step()+++class TestMuon:+    @pytest.mark.unit+    def test_registry(self):+        assert AVAILABLE_OPTIMIZERS['muon'] is Muon+        parameters = _parameters()+        optimizer = optim.get_optimizer('muon')(parameters.parameters(), lr=LR)+        assert isinstance(optimizer, Muon)++        parsed = optim.parse_optimizer_args('muon', {'name': 'auto', 'momentum': 0.9, 'weight_decay': WEIGHT_DECAY})+        assert parsed['momentum'] == 0.9+        assert parsed['weight_decay'] == WEIGHT_DECAY+        assert parsed['adam_lr_ratio'] == config.optimizers.MuonParams().adam_lr_ratio+        assert parsed['scale_mode'] == 'fan_in'+        assert parsed['use_muon'] is None++    @pytest.mark.unit+    def test_parameter_selection(self):+        parameters = _parameters()+        selected = {+            name: Muon.selects_muon(param, use_muon=None, adam_dim_threshold=8192)+            for name, param in parameters.items()+        }+        assert selected == {'hidden': True, 'tall': True, 'bias': False, 'conv': False, 'embedding': False}+        assert not Muon.selects_muon(parameters['hidden'], use_muon=None, adam_dim_threshold=48)+        assert not Muon.selects_muon(parameters['hidden'], use_muon=False, adam_dim_threshold=8192)+        assert Muon.selects_muon(parameters['conv'], use_muon=True, adam_dim_threshold=8192)+        assert Muon.selects_muon(parameters['embedding'], use_muon=True, adam_dim_threshold=8192)++        with pytest.raises(ValueError, match="at least two dimensions"):+            Muon([{'params': [parameters['bias']], 'use_muon': True}], lr=LR)+        with pytest.raises(ValueError, match="scale_mode"):+            Muon(parameters.parameters(), lr=LR, scale_mode='rows')++    @pytest.mark.unit+    def test_newton_schulz_orthogonalizes(self):+        generator = torch.Generator().manual_seed(1)+        grad = torch.randn(16, 40, generator=generator)+        orthogonal = zeropower_via_newtonschulz5(grad)+        assert orthogonal.shape == grad.shape+        assert orthogonal.dtype == grad.dtype+        singular_values = torch.linalg.svdvals(orthogonal.double())+        assert singular_values.min() > 0.5+        assert singular_values.max() < 1.5+        # the input is far from orthogonal+        normalized = torch.linalg.svdvals(grad.double() / grad.double().norm())+        assert normalized.min() < 0.5++        # the iteration runs on the wide orientation, so a transposed input gives exactly the transposed result+        assert torch.equal(zeropower_via_newtonschulz5(grad.T), orthogonal.T)+        with pytest.raises(ValueError, match="matrix"):+            zeropower_via_newtonschulz5(torch.zeros(3, 4, 5))++    @pytest.mark.unit+    def test_muon_step_matches_definition(self):+        parameters = _parameters()+        optimizer = Muon(parameters.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)+        initial = {name: param.detach().clone() for name, param in parameters.items()}+        _run(optimizer, parameters, steps=2)++        expected = initial['hidden'].clone()+        buf = torch.zeros_like(expected)+        for step in range(2):+            generator = torch.Generator().manual_seed(100 + step)+            grad = torch.randn(48, 32, generator=generator)+            buf = buf * 0.95 + grad+            nesterov = grad + 0.95 * buf+            orthogonal = zeropower_via_newtonschulz5(nesterov, steps=5)+            expected = expected * (1 - LR * WEIGHT_DECAY) - LR * max(1.0, 48 / 32) ** 0.5 * orthogonal+        assert torch.allclose(parameters['hidden'], expected, atol=1e-7, rtol=1e-6)+        assert torch.equal(optimizer.state[parameters['hidden']]['momentum_buffer'], buf)++    @pytest.mark.unit+    def test_adamw_fallback_matches_torch_adamw(self):+        parameters = _parameters()+        reference = copy.deepcopy(parameters)+        betas, eps = (0.9, 0.95), 1e-8+        optimizer = Muon(+            parameters.parameters(),+            lr=LR,+            weight_decay=WEIGHT_DECAY,+            adam_betas=betas,+            adam_eps=eps,+            adam_lr_ratio=1.0,+        )+        adamw_names = ['bias', 'conv', 'embedding']+        adamw = torch.optim.AdamW(+            [reference[name] for name in adamw_names], lr=LR, betas=betas, eps=eps, weight_decay=WEIGHT_DECAY+        )+        for step in range(3):+            _set_gradients(parameters, 200 + step)+            _set_gradients(reference, 200 + step)+            optimizer.step()+            adamw.step()+        for name in adamw_names:+            torch.testing.assert_close(parameters[name], reference[name], atol=1e-7, rtol=1e-6)+            assert optimizer.state[parameters[name]]['step'] == 3++    @pytest.mark.unit+    def test_weight_decay_uses_the_group_learning_rate(self):+        parameters = _parameters()+        initial = {name: param.detach().clone() for name, param in parameters.items()}+        optimizer = Muon(+            [+                {'params': [parameters['hidden'], parameters['bias']]},+                {'params': [parameters['tall'], parameters['conv']], 'weight_decay': 0.0},+            ],+            lr=LR,+            weight_decay=WEIGHT_DECAY,+            adam_lr_ratio=0.1,+        )+        for param in parameters.values():+            param.grad = torch.zeros_like(param)+        optimizer.step()+        # with zero gradients both updates vanish and only the decay acts+        torch.testing.assert_close(parameters['hidden'], initial['hidden'] * (1 - LR * WEIGHT_DECAY))+        torch.testing.assert_close(parameters['bias'], initial['bias'] * (1 - LR * 0.1 * WEIGHT_DECAY))+        assert torch.equal(parameters['tall'], initial['tall'])+        assert torch.equal(parameters['conv'], initial['conv'])++    @pytest.mark.unit+    def test_scale_modes(self):+        assert update_scale((64, 16), 'fan_in') == 2.0+        assert update_scale((64, 16), 'fan_out') == 1.0+        assert update_scale((16, 64), 'fan_in') == 1.0+        assert update_scale((16, 64), 'fan_out') == 2.0+        assert update_scale((64, 16), 'none') == 1.0++        updates = {}+        for mode in ('fan_in', 'fan_out', 'none'):+            parameters = _parameters()+            initial = parameters['tall'].detach().clone()+            optimizer = Muon([parameters['tall']], lr=LR, scale_mode=mode)+            _run(optimizer, parameters, steps=1)+            updates[mode] = initial - parameters['tall']+        torch.testing.assert_close(updates['fan_in'], 2 * updates['fan_out'])+        assert torch.equal(updates['fan_out'], updates['none'])++        # the same matrix stored transposed receives the transposed update when the modes are paired+        parameters = _parameters()+        transposed = torch.nn.Parameter(parameters['tall'].detach().T.contiguous())+        optimizer = Muon([parameters['tall']], lr=LR, scale_mode='fan_in')+        optimizer_t = Muon([transposed], lr=LR, scale_mode='fan_out')+        for step in range(3):+            _set_gradients(parameters, 300 + step)+            transposed.grad = parameters['tall'].grad.T.contiguous()+            optimizer.step()+            optimizer_t.step()+        assert torch.equal(transposed, parameters['tall'].T)++    @pytest.mark.unit+    def test_forced_muon_flattens_filters(self):+        parameters = _parameters()+        flat = torch.nn.Parameter(parameters['conv'].detach().reshape(8, -1).clone())+        optimizer = Muon([{'params': [parameters['conv']], 'use_muon': True}], lr=LR, weight_decay=WEIGHT_DECAY)+        optimizer_flat = Muon([flat], lr=LR, weight_decay=WEIGHT_DECAY)+        for step in range(2):+            _set_gradients(parameters, 400 + step)+            flat.grad = parameters['conv'].grad.reshape(8, -1).clone()+            optimizer.step()+            optimizer_flat.step()+        assert torch.equal(parameters['conv'].reshape(8, -1), flat)+        assert 'momentum_buffer' in optimizer.state[parameters['conv']]++    @pytest.mark.unit+    def test_state_dict_round_trip(self):+        parameters = _parameters()+        optimizer = Muon(parameters.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)+        _run(optimizer, parameters, steps=2)+        saved_parameters = copy.deepcopy(parameters)+        saved_state = copy.deepcopy(optimizer.state_dict())+        _run(optimizer, parameters, steps=1, first_seed=102)++        restored = Muon(saved_parameters.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)+        restored.load_state_dict(saved_state)+        _run(restored, saved_parameters, steps=1, first_seed=102)+        for name in parameters:+            assert torch.equal(saved_parameters[name], parameters[name]), name+        assert restored.state[saved_parameters['bias']]['step'] == 3++    @pytest.mark.unit+    def test_scheduler_drives_both_learning_rates(self):+        parameters = _parameters()+        optimizer = Muon(parameters.parameters(), lr=1.0, adam_lr_ratio=0.1)+        scheduler = optim.lr_scheduler.WarmupPolicy(optimizer, warmup_steps=4, max_steps=10)+        initial = {name: param.detach().clone() for name, param in parameters.items()}+        for param in parameters.values():+            param.grad = torch.zeros_like(param)+        optimizer.param_groups[0]['weight_decay'] = 0.5+        optimizer.step()+        scheduler.step()+        lr = scheduler.get_last_lr()[0]+        assert lr == pytest.approx(2 / 5)+        # the first step ran at the warmup learning rate 1/5+        torch.testing.assert_close(parameters['hidden'], initial['hidden'] * (1 - 0.2 * 0.5))+        torch.testing.assert_close(parameters['bias'], initial['bias'] * (1 - 0.02 * 0.5))++    @pytest.mark.unit+    def test_production_shape_update_rms(self):+        # the first feed-forward matrix of a width-1024 Conformer, in torch's (out, in) layout+        matrix = torch.nn.Parameter(torch.randn(4096, 1024, generator=torch.Generator().manual_seed(5)))+        assert Muon.selects_muon(matrix, use_muon=None, adam_dim_threshold=8192)+        initial = matrix.detach().clone()+        optimizer = Muon([matrix], lr=1.0)+        matrix.grad = torch.randn(matrix.shape, generator=torch.Generator().manual_seed(6))+        optimizer.step()+        update = initial - matrix.detach()+        # fan_in: the update of an (out, in) matrix has root mean square 1 / sqrt(fan_in),+        # up to the Newton-Schulz approximation of the singular values+        assert update.pow(2).mean().sqrt().item() == pytest.approx(1 / 1024**0.5, rel=0.15)+        assert update.shape == (4096, 1024)++    @pytest.mark.unit+    def test_small_vocabulary_embedding_needs_an_explicit_exclusion(self):+        # a 1026 x 512 embedding is below the default threshold, so the shape heuristic selects Muon for it+        embedding = torch.nn.Embedding(1026, 512)+        assert Muon.selects_muon(embedding.weight, use_muon=None, adam_dim_threshold=8192)+        optimizer = Muon([{'params': [embedding.weight], 'use_muon': False, 'weight_decay': 0.0}], lr=LR)+        embedding.weight.grad = torch.ones_like(embedding.weight)+        optimizer.step()+        state = optimizer.state[embedding.weight]+        assert 'exp_avg' in state+        assert 'momentum_buffer' not in state+        # a wide embedding falls back by shape alone+        assert not Muon.selects_muon(torch.empty(8192, 16), use_muon=None, adam_dim_threshold=8192)