[Fix] Unify Muon matrix layout for NS and lr_ratio

[YanhuiDua]

Motivation

Muon updates need one consistent answer to: what is the logical matrix being orthogonalized, and what (rows, cols) should lr_ratio use? Two problems hide in this question.

Problem 1 — NS orthogonalization shape and lr_ratio shape disagree

The two places that interpret parameter layout disagreed:

1. Newton–Schulz (muon_update_newton_schulz)
   • With flatten=True and ndim >= 3 (e.g. Conv3d / Conv1d weights), the tensor was flattened to one 2D matrix via flatten(start_dim=1), e.g. [1152, 3, 2, 16, 16] → (1152, 1536).
2. lr_ratio (precomputed from p.shape[-2] // num_experts, p.shape[-1])
   • Ignored the flatten contract and only looked at the last two dims, e.g. the same Conv3d used (16, 16) for RMS scaling.

So NS orthogonalized a large flattened matrix, while the effective step size was scaled as if the matrix were the tiny trailing slice — update direction and learning-rate scale described different shapes. Additionally, the NS callback reinterpreted MoE layout via a separate num_experts argument, making the optimizer ↔ callback contract easy to get wrong (MuonSplit / Triton / custom NS).

Problem 2 — non-matrix tensors are routed into Muon at all

Qwen3.5-VL/Qwen3.5 language layers start with a depthwise causal conv1d [8192, 1, 4] (and the vision tower has patch_embed Conv3d [1152, 3, 2, 16, 16]). These are not matrices:

1. No cross-channel spectral structure. flatten=True turned [8192, 1, 4] into one [8192, 4] matrix, so NS orthogonalized across channels, coupling and magnitude-equalizing 8192 independent 4-tap filters — the update direction is distorted by construction.
2. lr_ratio was computed on the wrong slice. The old init used (p.shape[-2], p.shape[-1]) = (1, 4) → rms ratio 0.4, while NS actually ran on [8192, 4].
3. Fixing Problem 1 alone exposes the routing bug. With both keyed to the flattened [8192, 4], the ratio jumps to 0.2*sqrt(8192) ≈ 18.10 (×45.25), which reproduced as a grad_norm spike at step 10 (module-grad hooks pinned it to layers.0.self_attn.conv1d).
4. Mainstream implementations agree. Megatron-LM's Muon only accepts 2D non-embedding weights (megatron/core/optimizer/muon.py), and NeMo Automodel's dion grouping sends ndim == 2 to Muon and everything else to fused AdamW (nemo_automodel/components/optim/dion.py).

Fix

1. _muon_matrix_shape — single source of truth for logical Muon layout

Case	Result
num_experts > 1 on 2D [E*rows, cols]	(E, rows, cols) — per-expert NS
flatten=True or already 2D	(rows, prod(rest)) — one matrix (incl. conv flatten)
flatten=False and ndim >= 3	(batch, rows, cols) — batch of last-two-dim matrices

Both _param_lr_ratio and muon_update_newton_schulz call this helper, so orthogonalization shape and lr_ratio always match. Default / Triton NS then only accept already-shaped 2D/3D inputs (no silent [E*M, N] + num_experts reinterpretation on the optimizer path).

2. Route only 2D matrices to Muon (p.ndim == 2 in MuonConfig)

MuonConfig now sends 3D weights (depthwise conv1d, Conv3d) and 1D tensors to the AdamW group, matching Megatron-LM and NeMo Automodel. MoE fused expert weights are unaffected: they are stored as 2D [E * rows, cols] and keep the num_experts batch path. Regression tests in tests/optim/test_muon_param_routing.py (2D → Muon, conv1d → AdamW, 1D → AdamW, fused experts keep num_experts; the conv1d case fails under the old ndim >= 2 routing).

Change summary

• Add _muon_matrix_shape; wire NS + lr_ratio through it.
• Align Triton / built-in NS with the no-num_experts callback contract; harden MuonSplit accordingly.
• Route only 2D matrices to Muon (p.ndim == 2 in MuonConfig); 3D weights go to the AdamW group.
• Extend unit tests (incl. conv layout + param routing).

Results: PR vs main (Muon SFT)

Same configs, all2all, no checkpoints. Experiment 1 is 3×PR vs 3×main repeats (mean + min–max band). Experiments 2–3 are fresh A/B runs at the final PR head (both fixes, 5405706b): deterministic setup — step-1 losses match exactly on both sides.

Experiment 1 — GLM-5.2 Muon SFT (35 steps; 3×PR vs 3×main repeats)

Experiment 2 — Qwen3.5-VL Muon SFT (20 steps; PR vs main, both fixes)

Experiment 3 — Qwen3.5 text-only Muon SFT (20 steps; PR vs main, both fixes)

[YanhuiDua]

@claude review

[github-actions]

Claude:

Summary

本 PR 引入 _muon_matrix_shape 作为 Muon 逻辑矩阵布局(2D 单矩阵 / 3D 批矩阵)的唯一来源,并据此统一驱动 _param_lr_ratio 与 muon_update_newton_schulz,同时移除了 Newton-Schulz 函数中基于 num_experts 的重解释。方向正确,修正了 flatten=True 时 conv 参数 lr_ratio 用错维度、以及 flatten=False 时 3D 参数被塌缩为单一大矩阵这两个真实缺陷,但新的 NS 调用契约漏改了 MuonSplit 回调,导致该路径必然在运行期抛 TypeError。

Main Flowchart after this PR

flowchart TD
  A["Muon.step"] --> B["_create_muon_tasks"]
  B --> C["_build_muon_task<br/>_param_lr_ratio (changed)"]
  C --> D{"muon_split_sizes?"}
  D -- "no" --> E["self._newton_schulz_func"]
  D -- "yes" --> F["partial(_muonsplit_newton_schulz, ...)"]
  E --> G["muon_update_newton_schulz<br/>_muon_matrix_shape -> 2D/3D (changed)"]
  F --> G
  G --> H["newton_schulz_func(matrices, epsilon=epsilon)"]
  H --> I["zeropower / triton<br/>accept 2D or 3D (changed)"]
  H --> J["_muonsplit_newton_schulz<br/>TypeError: missing num_experts"]
  style C fill:#ffe9a8
  style G fill:#ffe9a8
  style I fill:#ffe9a8
  style J fill:#ffb3b3

Loading

核心原理实现与单测

• 布局契约:_muon_matrix_shape 的三个分支(expert 优先、flatten/已 2D、按前导维度批处理)与 zeropower_via_newtonschulz5、newton_schulz_triton 的 2D/3D 入参实现一致;Triton 侧 unsqueeze(0)/squeeze(0) 正确保持了返回 rank。
• lr_ratio 时机:lr_ratio 从 __init__ 移到 _build_muon_task,已确认此处 p.shape 仍是 DTensor 全局 shape(参数本身不会被原地 reshape,maybe_to_local 返回新张量),原注释担心的“局部 shard shape”不会发生。
• 单测覆盖:2D/3D NS 等价性、conv flatten 布局、expert 形状规则、legacy checkpoint 的 lr_ratio 刷新均有覆盖;MuonSplit 的真实调用路径(经 Muon.step())完全没有覆盖,现有测试直接调用 _muonsplit_newton_schulz 并显式传入 num_experts=1,恰好掩盖了本次契约变更引入的 Critical 缺陷。

抽象与信息隐藏评估

• Warning xtuner/v1/optim/muon.py:119-121:被声明为 “single source of truth” 的 _muon_matrix_shape 在 num_experts > 1 分支忽略 flatten 并硬编码 shape[1],对 ndim > 2 的 expert 参数会产生丢元素的形状,既不报错也不正确。

单测建议

• Warning tests/optim/test_muon_conv_layout.py:134-158:MuonSplit 与 expert 布局只通过私有 _param_lr_ratio 断言且整类标记为 @pytest.mark.gpu,缺少经 public Muon.step() 的真实路径覆盖,建议补一个 CPU 单进程 muon_split_sizes step 测试。

其他 Issues

Critical

• xtuner/v1/optim/muon.py:1692:新的 NS 调用不再传 num_experts,而 _muonsplit_newton_schulz 仍将其声明为必填参数,所有启用 MuonSplit 的训练(qwen3_vl vision、glm52 dsa_mla、gated_deltanet)在第一次 step 即抛 TypeError。

Warning

• xtuner/v1/optim/muon.py:1734 / xtuner/v1/optim/newton_schulz_triton.py:292:公开的 newton_schulz_func 回调契约与两个 NS 函数签名移除了 num_experts,属于无 deprecation 的 breaking change,外部自定义回调会直接 TypeError。

Verdict

REQUEST_CHANGES

[YanhuiDua]

@claude review

[github-actions]

Claude:

Summary

本 PR 引入 _muon_matrix_shape 作为 Muon 逻辑矩阵布局(2D 单矩阵 / 3D 批矩阵)的唯一来源,并据此统一驱动 _param_lr_ratio 与 muon_update_newton_schulz,同时从默认与 Triton Newton-Schulz 中移除基于 num_experts 的重解释。修正了两个真实缺陷:flatten=True 时 conv 参数的 lr_ratio 用 shape[-2]/shape[-1] 而 NS 实际用 (shape[0], prod(shape[1:]));flatten=False 时 3D+ 参数被塌缩成单个大矩阵。前几轮指出的 Critical(_muonsplit_newton_schulz 漏改 NS 调用契约)、expert 分支产生无效形状、以及 MuonSplit 缺少真实路径覆盖,在当前 head 均已修复,本轮未发现新的 Warning 及以上问题。

核心原理实现与单测

• 布局契约一致:_muon_matrix_shape 三分支(num_experts > 1 优先 → flatten 或已 2D → 按前导维度批处理)与 zeropower_via_newtonschulz5(纯 2D/3D @)、newton_schulz_triton(unsqueeze(0)/squeeze(0),batch 在 transpose 前取、A 缓冲在 transpose 后按 X.size(-2) 分配)实现一致,返回 rank 与输入一致。
• MuonSplit 等价性:torch.stack + unbind(0) 与旧的 cat(dim=-2) + split 等价,因为旧实现在 NS 内部又以 view(num_experts, -1, N) 还原为同一 3D 布局;mul_(lr_ratio) 作用于 NS 新张量,无别名风险。
• lr_ratio 迁移安全:_build_muon_task 在构造 AsyncTask(generator 开始执行)之前计算 ratio,此处 p.shape 仍是 DTensor 全局 shape,maybe_to_local 不原地修改参数;state["lr_ratio"] 在 muon.py 中无读者(更新实际使用 lr_ratios[0]),属 checkpoint/可观测性字段。
• 生产可达性:xtuner/v1/config/optim.py 仅选取 p.ndim >= 2 且固定 flatten=True,故 conv 3D/4D 参数确实走被修复的 flatten 路径;expert 组参数恒为 2D,新增的 assert len(shape) == 2 在生产不会触发。
• 单测覆盖:_muon_matrix_shape 契约(含两条 assert 消息)、2D 与 3D NS 等价性、conv 与等价二维参数经 optimizer.step() 逐位一致(Shard/Replicate × 3D/4D/5D)、各策略下 NS 矩阵行列与 lr_ratio 行列一致(tracking callback)、MuonSplit 经 public Muon.step()、legacy 与 DCP checkpoint 恢复后 ratio 刷新,均有真实代码路径覆盖。

其他 Issues

Warning

• xtuner/v1/optim/muon.py#L1732 / xtuner/v1/optim/newton_schulz_triton.py#L292:公开 NS 函数与 newton_schulz_func 回调契约移除 num_experts 属无 deprecation 的 breaking change(已在先前轮次 inline 指出,本轮不重复发布);影响面限于外部自定义回调——这两个函数未在 xtuner/v1/optim/__init__.py 导出,docs/ 中亦无引用。

Verdict

APPROVE