diff --git a/tests/optim/test_muon.py b/tests/optim/test_muon.py index de770480e..22162d2ea 100644 --- a/tests/optim/test_muon.py +++ b/tests/optim/test_muon.py @@ -23,9 +23,10 @@ from xtuner._testing.testcase import DeterministicDDPTestCase from xtuner.v1.config import FSDPConfig from xtuner.v1.config.optim import MuonConfig +from xtuner.v1.engine.train_engine import TrainEngine from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseMLP -from xtuner.v1.optim.muon import zeropower_via_newtonschulz5 +from xtuner.v1.optim.muon import _muonsplit_newton_schulz, zeropower_via_newtonschulz5 # ─── Test: Newton-Schulz functions ─────────────────────────────────────────── @@ -58,6 +59,34 @@ def test_triton_vs_pytorch(self): assert not torch.isnan(result2).any() torch.testing.assert_close(result1, result2, atol=3e-2, rtol=3e-2) + def test_muonsplit_matches_independent_blocks(self): + self.create_pg("cuda") + split_sizes = (6, 2, 6, 2) + logical_rows = sum(split_sizes) + G = torch.randn(logical_rows + 2, 4, device="cuda", dtype=torch.float32) + + result = _muonsplit_newton_schulz( + G, + epsilon=1e-7, + num_experts=1, + newton_schulz_func=zeropower_via_newtonschulz5, + split_sizes=split_sizes, + adjust_lr="rms_norm", + ) + + expected_chunks = [] + offset = 0 + for size in split_sizes: + chunk = G.narrow(0, offset, size) + chunk = zeropower_via_newtonschulz5(chunk, epsilon=1e-7) + chunk.mul_(0.2 * math.sqrt(max(size, G.size(1)))) + expected_chunks.append(chunk) + offset += size + expected_chunks.append(torch.zeros_like(G[logical_rows:])) + expected = torch.cat(expected_chunks) + + torch.testing.assert_close(result, expected) + # ─── Model for end-to-end tests ────────────────────────────────────────────── @@ -362,6 +391,44 @@ def test_muon_single_gpu_matches_reference(self): msg=f"mismatch on '{name}': max_abs={abs_diff.max().item():.2e}, max_rel={rel_diff.max().item():.2e}", ) + def test_clip_grad_norm_only_clips_adamw(self): + self.create_pg("cuda") + device = "cuda" + model = ToyMoEModelConfig(compile_cfg=False).build().to(device) + model.fully_shard( + FSDPConfig( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + torch_compile=False, + ) + ) + + optim_config = MuonConfig(max_grad_norm=1.0) + optimizer = optim_config.build(model) + total_squared = 0 + adamw_squared = 0 + for group in optimizer.param_groups: + is_muon = group["algorithm"] == "muon" + grad_value = 2.0 if is_muon else 1.0 + for param in group["params"]: + param.grad = torch.full_like(param, grad_value) + total_squared += param.numel() * grad_value**2 + if not is_muon: + adamw_squared += param.numel() + + engine = object.__new__(TrainEngine) + engine.model = model + engine.optimizer = optimizer + engine.optim_cfg = optim_config + grad_norm = engine.clip_grad_norm() + + torch.testing.assert_close(grad_norm, torch.tensor(math.sqrt(total_squared), device=device)) + clip_coef = min(1.0, optim_config.max_grad_norm / (math.sqrt(adamw_squared) + 1e-6)) + for group in optimizer.param_groups: + expected = 2.0 if group["algorithm"] == "muon" else clip_coef + for param in group["params"]: + torch.testing.assert_close(param.grad.to_local(), torch.full_like(param.grad.to_local(), expected)) + # ─── Test: end-to-end FSDP ─────────────────────────────────────────────────── @@ -430,7 +497,9 @@ def test_muon_fsdp_matches_reference(self, enable_all2all: bool): fsdp_loss.backward() # ── Production Muon optimizer step ──────────────────────────────────── - muon_config = MuonConfig(lr=LR, momentum=MU, weight_decay=WD, eps=EPSILON, betas=BETAS, enable_all2all=enable_all2all) + muon_config = MuonConfig( + lr=LR, momentum=MU, weight_decay=WD, eps=EPSILON, betas=BETAS, enable_all2all=enable_all2all + ) optim = muon_config.build(model) optim.step() diff --git a/xtuner/v1/config/optim.py b/xtuner/v1/config/optim.py index 6f28cb7e6..16801f640 100644 --- a/xtuner/v1/config/optim.py +++ b/xtuner/v1/config/optim.py @@ -83,6 +83,12 @@ def build(self, model): trainable_parameters_names = model.trainable_parameters() trainable_names = {name for name, _ in trainable_parameters_names} + muon_split_sizes = {} + for module in model.modules(): + get_muon_split_sizes = getattr(module, "get_muon_split_sizes", None) + if callable(get_muon_split_sizes): + muon_split_sizes.update(get_muon_split_sizes()) + untrainable_names = [] num_total = 0 num_total_requires_grad = 0 @@ -201,6 +207,7 @@ def build(self, model): use_triton=False, epsilon=self.eps, enable_all2all=self.enable_all2all, + muon_split_sizes=muon_split_sizes, ) return optimizer diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 61a1f58fe..e6e23cf36 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -257,11 +257,28 @@ def init_model_weights(self): def clip_grad_norm(self, do_clip: bool = True, dtype=torch.float32): ProberList.before_clip_grad_norm(self.model) self.model.scale_and_reduce_grad() + params = self.model.trainable_parameters() grads = [p.grad for _, p in params if p.grad is not None] grad_norm, grouped_grads = cal_grad_norm(grads, dtype=dtype) + if do_clip: - clip_coef = self.optim_cfg.max_grad_norm / (grad_norm + 1e-6) + if any(group.get("algorithm") == "muon" for group in self.optimizer.param_groups): + clip_params = ( + p + for group in self.optimizer.param_groups + if group.get("algorithm") != "muon" + for p in group["params"] + ) + clip_grads = [p.grad for p in clip_params if p.grad is not None] + if clip_grads: + clip_grad_norm, grouped_grads = cal_grad_norm(clip_grads, dtype=dtype) + else: + clip_grad_norm = torch.zeros((), device=DEVICE, dtype=dtype) + grouped_grads = {} + else: + clip_grad_norm = grad_norm + clip_coef = self.optim_cfg.max_grad_norm / (clip_grad_norm + 1e-6) clip_coef_clamped = torch.clamp(clip_coef, max=1.0) for grads in grouped_grads.values(): device = grads[0].device diff --git a/xtuner/v1/module/attention/dsa_mla.py b/xtuner/v1/module/attention/dsa_mla.py index 063e740b7..6f9bd1020 100644 --- a/xtuner/v1/module/attention/dsa_mla.py +++ b/xtuner/v1/module/attention/dsa_mla.py @@ -276,6 +276,14 @@ def __init__( indexer_backend=self.sparse_mla_backend, ) + def get_muon_split_sizes(self) -> dict[nn.Parameter, tuple[int, ...]]: + """Return the logical row blocks used by GLM MuonSplit.""" + return { + self.q_b_proj.weight: (self.qk_nope_head_dim, self.qk_rope_head_dim) * self.num_attention_heads, + self.kv_a_proj_with_mqa.weight: (self.kv_lora_rank, self.qk_rope_head_dim), + self.kv_b_proj.weight: (self.qk_nope_head_dim, self.v_head_dim) * self.num_attention_heads, + } + def forward( self, hidden_states: torch.Tensor, diff --git a/xtuner/v1/optim/muon.py b/xtuner/v1/optim/muon.py index 0b24ccc94..2af57e85b 100644 --- a/xtuner/v1/optim/muon.py +++ b/xtuner/v1/optim/muon.py @@ -24,6 +24,7 @@ import math from collections import defaultdict +from functools import partial from itertools import chain, product from typing import Callable, Generator, Iterator, Literal, Sequence, cast, overload @@ -43,7 +44,11 @@ def maybe_to_local(tensor: list[Tensor]) -> list[Tensor]: return [t.to_local() if isinstance(t, DTensor) else t for t in tensor] -def create_param_batches(params: Sequence[Tensor], batch_size: int) -> Generator[list[Tensor], None, None]: +def create_param_batches( + params: Sequence[Tensor], + batch_size: int, + extra_group_key: Callable[[Tensor], object] | None = None, +) -> Generator[list[Tensor], None, None]: """Batch parameters into groups of size `batch_size`. Tensors in each batch will have identical shape, sharding, and dtype. @@ -52,7 +57,8 @@ def create_param_batches(params: Sequence[Tensor], batch_size: int) -> Generator groups = defaultdict(list) for p in params: sharding = p.placements if isinstance(p, DTensor) else None - groups[(p.shape, sharding, p.dtype)].append(p) + extra_key = extra_group_key(p) if extra_group_key is not None else None + groups[(p.shape, sharding, p.dtype, extra_key)].append(p) # Create batches from grouped parameters for group in groups.values(): @@ -61,6 +67,18 @@ def create_param_batches(params: Sequence[Tensor], batch_size: int) -> Generator yield batch +def _get_muon_lr_ratio( + fan_out: int, + fan_in: int, + adjust_lr: Literal["rms_norm", "spectral_norm", "none"], +) -> float: + if adjust_lr == "none": + return 1.0 + if adjust_lr == "spectral_norm": + return math.sqrt(fan_out / fan_in) + return 0.2 * math.sqrt(max(fan_out, fan_in)) + + def pad_batch(batch: list[Tensor], batch_size: int) -> list[Tensor]: """Insert dummy tensors so the batch has exactly `batch_size` elements.""" assert len(batch) > 0 @@ -275,6 +293,8 @@ class Muon(Optimizer): remainder_strategy (str): Communication strategy for parameter batches smaller than world size. ``"agrs"`` uses all-gather + reduce-scatter without batch padding. ``"pad_all2all"`` restores the original FSDP2 Muon behavior by zero-padding the batch to world size and using all-to-all. + muon_split_sizes (dict[Tensor, tuple[int, ...]] | None): Logical row blocks that Muon should + orthogonalize and scale independently. Used by GLM MuonSplit attention projections. Muon optimizer algorithm by Keller Jordan: https://kellerjordan.github.io/posts/muon/ FSDP2 Muon uses all-to-all communications: https://www.essential.ai/blog/infra @@ -295,6 +315,7 @@ def __init__( newton_schulz_func: Callable | None = None, enable_all2all: bool = True, remainder_strategy: Literal["agrs", "pad_all2all"] = "agrs", + muon_split_sizes: dict[Tensor, tuple[int, ...]] | None = None, ): # Check hyperparameters if lr < 0.0: @@ -328,6 +349,7 @@ def __init__( super().__init__(params, defaults) self._enable_all2all = enable_all2all self._remainder_strategy = remainder_strategy + self._muon_split_sizes = muon_split_sizes or {} # Pre-compute lr adjustment ratios for each Muon parameter based on global shape. # This must happen at init time because DTensor.shape here is guaranteed to be @@ -340,16 +362,18 @@ def __init__( ne = group.get("num_experts", 1) for p in group["params"]: state = self.state[p] - if adj == "none": + split_sizes = self._muon_split_sizes.get(p) + if split_sizes is not None: + if p.ndim != 2 or ne != 1: + raise ValueError("MuonSplit only supports regular 2D Muon parameters.") + if not split_sizes or any(size <= 0 for size in split_sizes): + raise ValueError(f"Invalid MuonSplit sizes: {split_sizes}") + if sum(split_sizes) > p.shape[-2]: + raise ValueError(f"MuonSplit sizes {split_sizes} exceed parameter shape {tuple(p.shape)}") + # Each logical block applies its own ratio inside the orthogonalization callback. state["lr_ratio"] = 1.0 - elif adj == "spectral_norm": - fan_out = p.shape[-2] // ne - fan_in = p.shape[-1] - state["lr_ratio"] = math.sqrt(fan_out / fan_in) - elif adj == "rms_norm": - A = p.shape[-2] // ne - B = p.shape[-1] - state["lr_ratio"] = 0.2 * math.sqrt(max(A, B)) + else: + state["lr_ratio"] = _get_muon_lr_ratio(p.shape[-2] // ne, p.shape[-1], adj) # Newton-Schulz configuration if newton_schulz_func is not None: @@ -664,16 +688,31 @@ def _create_muon_tasks( group_world_size = group_process_group.size() if group_process_group is not None else 1 # Create batches within this mesh group - for params in create_param_batches(mesh_params, batch_size=group_world_size): + for params in create_param_batches( + mesh_params, + batch_size=group_world_size, + extra_group_key=self._muon_split_sizes.get, + ): gradients: list[Tensor] = [g for p in params if (g := p.grad) is not None] assert len(gradients) == len(params), "Some gradients became None after filtering" states = [self._get_or_initialize_state(p, algo_name) for p in params] momentums = [s["momentum"] for s in states] - lr_ratios = [s["lr_ratio"] for s in states] + lr_ratios = [1.0 if p in self._muon_split_sizes else s["lr_ratio"] for p, s in zip(params, states)] assert len(set(lr_ratios)) == 1, f"Found different lr_ratios: {set(lr_ratios)}" + split_sizes = self._muon_split_sizes.get(params[0]) + assert all(self._muon_split_sizes.get(p) == split_sizes for p in params) + newton_schulz_func = self._newton_schulz_func + if split_sizes is not None: + newton_schulz_func = partial( + _muonsplit_newton_schulz, + newton_schulz_func=self._newton_schulz_func, + split_sizes=split_sizes, + adjust_lr=group["adjust_lr"], + ) + is_remainder = len(params) < group_world_size # When all-to-all is disabled, every sharded batch uses AGRS. Otherwise remainder # batches follow the configured strategy: current AGRS behavior or the original @@ -698,7 +737,7 @@ def _create_muon_tasks( epsilon=epsilon, nesterov=nesterov, flatten=flatten, - newton_schulz_func=self._newton_schulz_func, + newton_schulz_func=newton_schulz_func, comm_strategy="agrs", shard_dim=sharded_tensor_dim, process_group=group_process_group, @@ -735,7 +774,7 @@ def _create_muon_tasks( epsilon=epsilon, nesterov=nesterov, flatten=flatten, - newton_schulz_func=self._newton_schulz_func, + newton_schulz_func=newton_schulz_func, comm_strategy=comm_strategy, shard_dim=sharded_tensor_dim, process_group=comm_pg, @@ -1425,6 +1464,45 @@ def muon_update_newton_schulz( return newton_schulz_func(X, epsilon=epsilon, num_experts=num_experts).reshape(original_shape) +def _muonsplit_newton_schulz( + X: Tensor, + epsilon: float | Tensor, + num_experts: int, + *, + newton_schulz_func: Callable, + split_sizes: tuple[int, ...], + adjust_lr: Literal["rms_norm", "spectral_norm", "none"], +) -> Tensor: + """Orthogonalize and scale unequal logical row blocks independently.""" + if X.ndim != 2 or num_experts != 1: + raise ValueError(f"MuonSplit expects one 2D matrix, got shape={tuple(X.shape)}, num_experts={num_experts}") + + logical_rows = sum(split_sizes) + if logical_rows > X.size(-2): + raise ValueError(f"MuonSplit sizes {split_sizes} exceed input shape {tuple(X.shape)}") + + chunks = X.narrow(-2, 0, logical_rows).split(split_sizes, dim=-2) + chunks_by_size: dict[int, list[tuple[int, Tensor]]] = defaultdict(list) + for index, (size, chunk) in enumerate(zip(split_sizes, chunks)): + chunks_by_size[size].append((index, chunk)) + + results: dict[int, Tensor] = {} + for size, indexed_chunks in chunks_by_size.items(): + batch = torch.cat([chunk for _, chunk in indexed_chunks], dim=-2) + batch = newton_schulz_func(batch, epsilon=epsilon, num_experts=len(indexed_chunks)) + lr_ratio = _get_muon_lr_ratio(size, X.size(-1), adjust_lr) + if lr_ratio != 1.0: + batch.mul_(lr_ratio) + for (index, _), result in zip(indexed_chunks, batch.split(size, dim=-2)): + results[index] = result + + output = torch.cat([results[index] for index in range(len(split_sizes))], dim=-2) + if logical_rows < X.size(-2): + padding = torch.zeros_like(X.narrow(-2, logical_rows, X.size(-2) - logical_rows)) + output = torch.cat([output, padding], dim=-2) + return output + + def zeropower_via_newtonschulz5(G: Tensor, epsilon: float = 1e-7, num_experts: int = 1): """Newton-Schulz iteration to approximate the orthogonalization of X.