Muon is Scalable for LLM Training 和 RMS-to-RMS

在近期大模型（LLM/ViT）优化器的研究（尤其是 Muon 、Scalable Muon 及 AdaMuon 等衍生方案）中，RMS-to-RMS（均方根对齐 / Rescaling） 是一个关键的更新幅值缩放策略。

它的核心作用是：把二阶/正交化优化器（如 Muon）算出来的更新矩阵幅值，强制对齐（Normalize/Align）到传统 AdamW 优化器的 RMS（Root-Mean-Square，均方根）水平上。

1. 为什么需要 RMS-to-RMS？（背景与痛点）

在前面讨论的 Muon 优化器中，我们通过 Newton-Schulz 迭代把梯度矩阵转换为了正交矩阵 Ot=UVTO_t = UV^TOt=UVT：

正交矩阵 UVTUV^TUVT 的每一个奇异值都被强行拉平为了 111。
对于一个维度为 m×nm \times nm×n 的矩阵，正交化后的 Frobenius 范数（矩阵元素的平方和再开方）被固定为 min⁡(m,n)\sqrt{\min(m, n)}min(m,n) 。

这就带来了一个极大的工程问题：它的更新幅值数量级（RMS）与 AdamW 并不在同一量级上！

AdamW 的 RMS 幅值： AdamW 会除以二阶矩的均方根（vt\sqrt{v_t}vt ），其更新矩阵每个元素的均方根 RMS 通常稳定在接近 1.01.01.0 的水平。
原始 Muon 的 RMS 幅值： 矩阵的 RMS 变成了 min⁡(m,n)m×n\sqrt{\frac{\min(m, n)}{m \times n}}m×nmin(m,n) ，不同形状的层（如 4096×40964096 \times 40964096×4096 vs 4096×143364096 \times 143364096×14336）其幅值差异极大。

如果直接混合使用 AdamW 和 Muon，你需要为不同层单独扫 LR（学习率），导致 AdamW 经过验证的现成 LR Schedule（学习率调度策略）在 Muon 上无法直接复用。

2. RMS-to-RMS 是怎么计算的？

为了解决这个问题，研究者（如 Moonshot AI 在 Muon is Scalable for LLM Training 中）提出了 RMS-to-RMS Scaling 因子 γ\gammaγ：

设参数矩阵维度为 m×nm \times nm×n，正交化后的更新矩阵为 Ot=Newton-Schulz(Mt)O_t = \text{Newton-Schulz}(M_t)Ot=Newton-Schulz(Mt)。

RMS-to-RMS 策略引入一个缩放因子 γ\gammaγ：

γ=0.2⋅max⁡(m,n)\gamma = 0.2 \cdot \sqrt{\max(m, n)}γ=0.2⋅max(m,n)

或者更通用的形式：将更新矩阵 OtO_tOt 的 RMS 强制缩放到与 AdamW 对齐的常数尺度：

Updatet=γ⋅Ot\text{Update}_t = \gamma \cdot O_tUpdatet=γ⋅Ot

通过加上这个 RMS-to-RMS 的系数 γ\gammaγ：

统一幅值： 消除矩阵长宽比（Aspect Ratio）对更新步长的影响；
完美兼容 AdamW： 使得 Muon 优化器可以无缝复用 AdamW 的学习率（Learning Rate）和权重衰减（Weight Decay）超参数。

3. RMS-to-RMS 带来的工程收益

维度	原始 Muon (无 RMS-to-RMS)	带有 RMS-to-RMS 的 Scalable Muon / AdaMuon
学习率调优	需要重新扫描全新的 LR 范围	直接复用 AdamW 已有的 LR 与 Cosine Schedule
不同层幅值	方阵与长方形权重矩阵的更新幅度不一致	所有形状矩阵的均方根更新步长保持对齐
数值稳定性	容易在大矩阵（如 Hidden Size 14336）中导致权重过快膨胀	避免精度溢出（尤其在 bf16/fp8 训练时更稳定）

一句话概括：RMS-to-RMS 是一种对齐机制。它通过匹配 AdamW 的 RMS 均方根更新量，给正交化后的矩阵加上一个"量纲换算系数"，让先进的 Muon/二阶优化器能够像插拔模块一样，直接套用现有的 AdamW 训练生态。

RMS-to-RMS（均方根对齐/缩放）确实是 Scalable Muon （即 Moonshot AI 团队发表的论文 "Muon is Scalable for LLM Training"）中最核心的理论与工程贡献之一，但不是唯一的贡献。

如果把原始的 Muon 比作"在小模型上证明了正交化优化器可行性"的算法原型，那么 Scalable Muon 的主要贡献可以概括为解决了 Muon 向千亿级 LLM 扩展时的"三大痛点" ，而 RMS-to-RMS 则是解决其中超参数移植与跨层不稳定性的核心钥匙。

Scalable Muon 的三大核心贡献

1. 理论与幅值对齐：RMS-to-RMS Scaling（最关键的机制创新）

解决的痛点： 原始 Muon 的正交化输出矩阵 Ot=UVTO_t = UV^TOt=UVT 的元素均方根（RMS）极大地依赖于矩阵的具体形状（长宽 m×nm \times nm×n）。这导致不同形状的权重层（如 Self-Attention 的 QKVQKVQKV 投影阵 vs MLP 的 Up/Down 投影阵）接收到的实际更新步长严重失衡；同时，原始 LR 无法直接沿用 AdamW 的经验值。
RMS-to-RMS 的贡献： 论文推导并提出了精确的 RMS-to-RMS 缩放系数（如乘以 max⁡(m,n)\sqrt{\max(m, n)}max(m,n) 相关的比例因子），把 Muon 的更新幅值从量纲上映射到与 AdamW 的 RMS 处在同一级别。
收益： 使得大模型训练时，用户可以100% 无缝复用 AdamW 扫好的 Learning Rate 和 Weight Decay 超参数，且各层更新幅度保持高度稳定。

2. 分布式工程高效落地：Distributed Newton-Schulz (ZeRO/FSDP 融合)

解决的痛点： 原始 Muon 需要拿到完整的二维梯度矩阵才能进行 Newton-Schulz 迭代。在大模型使用 FSDP/ZeRO-3（参数与梯度分片）分布式训练时，这会导致极高的跨节点通信开销（频繁 All-Gather），严重拖慢 Wall-clock 训练速度。
Scalable Muon 的贡献： 设计了分布式 Newton-Schulz 正交化算子，将矩阵正交化迭代过程巧妙地与 ZeRO/FSDP 的梯度通信与参数更新重叠（Overlap）。
收益： 使得 Muon 训练 7B~70B+ 大模型时的通信开销降低到与传统 AdamW 几乎完全一致。

3. 数学稳定性与低精度训练支持（BF16/FP8 优化）

解决的痛点： 极深网络在低精度（BF16/FP8）训练时，Newton-Schulz 迭代中 Xk(XkTXk)2X_k (X_k^T X_k)^2Xk(XkTXk)2 的高阶矩阵乘法极易发生数值下溢（Underflow）或奇异值爆炸。
Scalable Muon 的贡献： 提出了迭代过程中的动态梯度截断与谱半径范数归一化预处理。
收益： 保证了千卡集群长达数周的 LLM 预训练过程"零 Loss Spike（损失骤增）"，极大提升了训练鲁棒性。

和 AdamW 配合

在实际的 Vision Transformer (ViT) 和 LLM 训练中，Muon 通常是作为"混合优化器"的一部分，与 AdamW 配合使用的。

它并不是一个彻底替代所有参数更新的单体优化器，而是针对 Transformer 架构不同模块的特性进行的分工协作。

1. 它们是如何分工配合的？

在具体的网络架构中，Muon 和 AdamW 负责更新的参数类型有着明确的界限：

Muon 负责更新：内部的二维重权重矩阵（2D Heavy Weight Matrices）
包含模块： Self-Attention 中的 Wq,Wk,Wv,WoW_q, W_k, W_v, W_oWq,Wk,Wv,Wo 投影矩阵，以及 MLP / SwiGLU 层中的线性变化矩阵（Up/Down/Gate Projection）。
原因： 这些参数都是高维的二维矩阵（比如 4096×40964096 \times 40964096×4096），具有极强的空间方向性。Muon 擅长通过 Newton-Schulz 迭代对这些矩阵做正交化更新，能够极大发挥其极速收敛和解耦曲率的优势。
AdamW 负责更新：一维参数与非矩阵参数（1D & Non-Matrix Parameters）
包含模块： 嵌入层（Embedding）、LayerNorm / RMSNorm 的缩放与偏移参数、偏置项（Biases）、Positional Embeddings（位置编码）等。
原因： 这些参数要么是一维向量，要么逻辑上无法做二维极分解正交化。用传统的 AdamW 去自适应缩放它们，既稳定又高效。

2. 为什么不能全用 Muon？

维度限制： Muon 依赖于矩阵的极分解（Polar Decomposition）和 Newton-Schulz 迭代（计算 XXTXX X^T XXXTX）。这个机制专为二维矩阵设计，对于向量（1D Vector）或零散的标量而言，正交化失去了代数意义。
训练稳定性： Norm 层和 Embedding 层的梯度分布特征与 MLP/Attention 层的权重矩阵截然不同，强行正交化容易导致数值不稳定性。
生态兼容： 配合 AdamW 更新 1D 参数，能够让混合优化器（Muon + AdamW）继承 AdamW 在微调和部分基础层上的鲁棒性。

3. 代码中的配合方式（PyTorch 伪代码）

在开源社区（如 Keller Jordan 的原版实现或 Moonshot 的 Scalable Muon 项目）中，通常在构建 Optimizers 时这样配置：

python 复制代码

# 1. 拆分参数：按维度与层类型分组
muon_params = [p for p in model.parameters() if p.ndim == 2 and "embed" not in p.name]
adamw_params = [p for p in model.parameters() if p.ndim < 2 or "embed" in p.name]

# 2. 实例化两个优化器（使用联合 Step 封装）
optimizer_muon = Muon(muon_params, lr=0.02, momentum=0.95)
optimizer_adamw = AdamW(adamw_params, lr=3e-4, weight_decay=0.01)

# 3. 训练循环中依次更新
optimizer_muon.step()
optimizer_adamw.step()

在现代 Transformer 训练生态中，"Muon 负责骨架（2D 矩阵），AdamW 负责细节（1D / Embedding / Norm）" 已经成为了使用 Muon 时的标准落地范式。

这两项技术是 Scalable Muon （如 Moonshot AI 提出的 "Muon is Scalable for LLM Training" 框架）为了将 Muon 从单卡/小模型扩展到万卡千亿参数 LLM 预训练时，在分布式系统工程 与数学算法计算量两个层面做出的核心优化。

下面为你分别详细拆解它们的原理与实现机制：

一、分布式 Newton-Schulz 正交化算子（System & Distributed Communication Optimization）

1. 为什么需要这个算子？（原始痛点）

在大模型分布式训练中（如使用 PyTorch FSDP 或 DeepSpeed ZeRO-3），模型参数 W\boldsymbol{W}W 和梯度 G\boldsymbol{G}G 被切割（Shard）分发在不同的 GPU 上。

传统做法的瓶颈： Newton-Schulz 迭代需要基于完整的二维梯度矩阵 计算 XkXkTXkX_k X_k^T X_kXkXkTXk。如果按常规流程，每张卡必须先通过跨节点通信（All-Gather）拼凑出完整的梯度矩阵，在本地算完 5 轮正交化迭代后，再重新切分权重更新。
这样会导致频繁的通信阻塞，GPU 算力大量空转，Wall-clock 耗时大幅增加。

2. 算子工作的核心原理

分布式 Newton-Schulz 正交化算子通过通信与计算重叠（Overlap）以及分块矩阵乘法（Block-wise Parallelism）重构了计算流程：

梯度通信重叠（Communication Overlap）： 算子将梯度正交化过程直接融合进 FSDP/ZeRO 的反向传播与参数更新 Hook 中。在前一层梯度计算的同时，后一层的 Newton-Schulz 迭代已经利用当前卡的显存切片并行启动，隐藏了通信延迟。
分块正交化与规约（Block-distributed Matrix Multiply）：
对于超大维度矩阵（例如 Hidden Size =14336= 14336=14336 的 MLP 层），算子利用分布式的矩阵乘法算子（类似 Megatron-LM 的 Tensor Parallel 逻辑），将 Xk(XkTXk)X_k (X_k^T X_k)Xk(XkTXk) 的矩阵乘法拆解为多卡并行计算：

Ak=XkXkT在列并行卡上计算并在行并行卡间 Reduce-Scatter\boldsymbol{A}_k = X_k X_k^T \quad \text{在列并行卡上计算并在行并行卡间 Reduce-Scatter}Ak=XkXkT在列并行卡上计算并在行并行卡间 Reduce-Scatter

每一轮迭代都尽量在本地分片显存中直接完成，无需将完整矩阵恢复到单卡显存中。

显存与算子融合（CUDA Kernel Fusion）： 针对 3.4445Xk−4.7750XkXkTXk+2.0315Xk(XkTXk)23.4445 X_k - 4.7750 X_k X_k^T X_k + 2.0315 X_k (X_k^T X_k)^23.4445Xk−4.7750XkXkTXk+2.0315Xk(XkTXk)2 这个固定的奇多项式，开发专门的 CUDA/Triton 融合算子（Fused Kernel），避免中间结果写入显存（Global Memory），极大地降低了带宽消耗。

二、动态低秩近似与更平滑的迭代停止条件（Algorithm & Compute Reduction Optimization）

1. 为什么需要这项优化？（原始痛点）

原始 Muon 无论梯度矩阵的奇异值分布如何，盲目地对所有层、所有 Step 统一执行固定次数（如 5 轮）的 Newton-Schulz 迭代。

每次迭代需要 2 次高维矩阵乘法 (GEMM) ，5 轮迭代就是 10 次 GEMM！
在 LLM 中，大量层的梯度矩阵本质上是低秩（Low-Rank）或其奇异值在训练中后期已经非常接近 1，继续做满 5 轮迭代造成了极大的算力浪费。

2. 动态低秩近似（Dynamic Low-Rank Approximation）

动态低秩近似通过降低参与正交化的矩阵有效维度，来直接裁剪 GEMM 的计算复杂度：

降维正交化： 当梯度矩阵 G∈Rm×n\boldsymbol{G} \in \mathbb{R}^{m \times n}G∈Rm×n（假设 m≪nm \ll nm≪n）表现出明显的低秩特性时，不再直接对原始大矩阵做正交化，而是将其投影到低秩子空间：

G≈UrΣrVrT\boldsymbol{G} \approx \boldsymbol{U}_r \boldsymbol{\Sigma}_r \boldsymbol{V}_r^TG≈UrΣrVrT

通过快速的 Randomized SVD 或 Subspace Iteration 提取主成分 r≪min⁡(m,n)r \ll \min(m, n)r≪min(m,n)，仅对 r×rr \times rr×r 的核心小矩阵执行 Newton-Schulz 迭代。

计算量裁剪： 原本矩阵乘法的计算复杂度从 O(m⋅n2)\mathcal{O}(m \cdot n^2)O(m⋅n2) 直接下降到 O(m⋅r2)\mathcal{O}(m \cdot r^2)O(m⋅r2)。正交化完成后，再通过低秩基底重新投影回原空间，矩阵乘法开销下降了 50% 以上。

3. 更平滑的迭代停止条件（Smooth Early-Stopping Criteria）

不再硬编码固定迭代次数，而是根据梯度矩阵的正交化程度（谱条件数/奇异值偏离度）动态决定何时提前退出迭代：

收敛度残差监控（Residual Norm Metric）： 在第 kkk 轮迭代时，通过轻量级标量估算残差 δk=∥XkXkT−I∥F\delta_k = \Vert{}X_k X_k^T - I\Vert{}Fδk=∥XkXkT−I∥F（或估算矩阵元素变化量 ∥Xk+1−Xk∥F∥Xk∥F\frac{\Vert{}X{k+1} - X_k\Vert{}_F}{\Vert{}X_k\Vert{}_F}∥Xk∥F∥Xk+1−Xk∥F）。
动态自适应早停（Adaptive Early-Exit）：
训练初期 / 梯度剧烈波动时： 残差较大，执行完整 4~5 轮迭代，确保梯度彻底正交化；
训练中后期 / 梯度已相对平滑时： 奇异值已经集中在 1 附近，残差 δk<ϵ\delta_k < \epsilonδk<ϵ（门限阈值），在第 2 轮或第 3 轮迭代后立刻提前退出（Early Exit）。
平滑松弛（Smooth Relaxation）： 引入平滑衰减因子，避免在迭代停止的临界步数产生突变的更新步长，保证更新量的连续性。

三、两者结合的最终收益

这两项改进相辅相成，构成了 Scalable Muon 在千亿大模型训练中的计算效率基石：

评估维度	原始 Muon	Scalable Muon (包含上述两项优化)
单步 Newton-Schulz GEMM 次数	强行固定 10 次	动态降至 4~6 次（节省 40%~60% 正交化计算量）
分布式通信开销	高（频繁 All-Gather 汇聚大矩阵）	低（通信与梯度计算重叠，开销对齐 AdamW）
万卡集群 Scaling 效率	受限于通信与显存，扩展性较差	高（支持 70B+ 参数规模模型的极速预训练）