[Fix] Add GLM-5.2 MuonSplit and AdamW-only gradient clipping

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()
 

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

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

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,

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.