Scaling up Muon for Large-Scale Language Model Training

文章作者/机构: Jingyuan Liu1, Jianlin Su1, Xingcheng Yao2, Zhejun Jiang1, Guokun Lai1, Yulun Du1, Yidao Qin1, Weixin Xu1, Enzhe Lu1, Junjie Yan1, Yanru Chen1, Huabin Zheng1, Yibo Liu1, Shaowei Liu1, Bohong Yin1, Weiran He1, Han Zhu1, Yuzhi Wang1, Jianzhou Wang1, Mengnan Dong1, Zheng Zhang1, Yongsheng Kang1, Hao Zhang1, Xinran Xu1, Yutao Zhang1, Yuxin Wu1, Xinyu Zhou1 ∗, Zhilin Yang1
1 Moonshot AI, 2 UCLA

A1 主要贡献

近期，基于矩阵正交化的Muon优化器【16, Muon: An optimizer for hidden layers in neural networks. 2024. URL: https://kellerjordan.github.io/posts/muon/】在小规模语言模型训练中展现了强大的性能，但其在大模型上的可扩展性尚未得到验证。本报告旨在解决将Muon优化器扩展应用于大规模语言模型训练（数十亿参数，数万亿tokens）所面临的挑战 。

核心问题与挑战：
1. 如何有效地将基于矩阵正交化的优化器扩展到大规模模型训练中。
2. 如何在分布式环境中计算近似正交化。
3. 此类优化器是否能泛化应用于包括预训练和监督微调（SFT）在内的不同训练阶段。

技术贡献：
本研究在Muon的基础上，系统地识别并解决了其在大规模训练场景下的局限性，主要技术贡献包括：
* Muon的有效扩展分析：通过广泛分析，我们发现权重衰减（weight decay）对Mu-on的可扩展性至关重要。此外，我们提出了针对Muon参数级更新规则的尺度调整策略。这些调整使得Muon无需进行超参数调优即可直接应用于大规模训练，并显著提高了训练稳定性。
* 高效的分布式实现：我们开发了Muon的分布式版本，采用了类似ZeRO-1【31, ZeRO: Memory optimizations Toward Training Trillion Parameter Models. 2020. SC20】的优化方式，在保持算法数学特性的同时，实现了最优的内存效率并降低了通信开销。
* 缩放定律验证：我们进行了缩放定律研究，将Muon与强大的AdamW基线进行了比较，并证明了Muon的卓越性能。根据缩放定律结果，在计算最优训练（compute-optimal training）的设置下，Muon达到与AdamW相当的性能仅需约52%的训练FLOPs。

基于这些改进，我们发布了使用Muon训练的16B参数MoE模型Moonlight，该模型推进了性能与训练FLOPs之间的帕累托前沿。我们同时开源了内存最优且通信高效的分布式Muon实现、预训练模型、指令微调模型以及中间检查点，以支持未来对可扩展优化技术的研究。

A3 背景知识

2.1 背景

Muon优化器 Muon【16, Muon: An optimizer for hidden layers in neural networks. 2024. URL: https://kellerjordan.github.io/posts/muon/】是最近提出的一种用于优化可表示为矩阵的神经网络权重的优化器。在第t次迭代时，给定当前权重$W_{t-1}$、动量$\mu$、学习率$\eta_t$和目标函数$L_t$，Muon优化器的更新规则如下 ：

$$\begin{aligned} \begin{aligned} \mathbf{M}_{t} & =\mu \mathbf{M}_{t-1}+\nabla \mathcal{L}_{t}\left(\mathbf{W}_{t-1}\right) \\ \mathbf{O}_{t} & =\text { Newton-Schulz }\left(\mathbf{M}_{t}\right)^{1} \\ \mathbf{W}_{t} & =\mathbf{W}_{t-1}-\eta_{t} \mathbf{O}_{t} \end{aligned} \end{aligned}$$

此处，$M_t$是第t次迭代时的梯度动量，当t=0时设为零矩阵。在公式1中，采用了牛顿-舒尔茨迭代过程（Newton-Schulz iteration process）【3, Old Optimizer, New Norm: An Anthology. 2024. arXiv: 2409.20325 [cs.LG]】来近似求解$(M_t M_t^T)^{-1/2} M_t$。设$M_t$的奇异值分解（SVD）为$U\Sigma V^T$，则有$(M_t M_t^T)^{-1/2} M_t = UV^T$，这将$M_t$正交化。直观上，正交化可以确保更新矩阵是同构的，从而防止权重沿着少数几个主导方向学习【16, Muon: An optimizer for hidden layers in neural networks. 2024. URL: https://kellerjordan.github.io/posts/muon/】 。

用于矩阵正交化的牛顿-舒尔茨迭代 公式1是通过一个迭代过程计算的。首先，我们设置$X_0 = M_t/\|M_t\|_F$。然后，在每次迭代k中，我们按如下方式从$X_{k-1}$更新$X_k$：

$$\mathbf{X}_{k}=a \mathbf{X}_{k-1}+b\left(\mathbf{X}_{k-1} \mathbf{X}_{k-1}^{\mathrm{T}}\right) \mathbf{X}_{k-1}+c\left(\mathbf{X}_{k-1} \mathbf{X}_{k-1}^{\mathrm{T}}\right)^{2} \mathbf{X}_{k-1}$$

其中$X_N$是经过N次迭代后的结果。这里的a、b、c是系数。为了确保公式2的正确收敛，我们需要调整这些系数，使得多项式$f(x) = ax + bx^3 + cx^5$在1附近有一个不动点。在【16, Muon: An optimizer for hidden layers in neural networks. 2024. URL: https://kellerjordan.github.io/posts/muon/】的原始设计中，系数被设置 为 a = 3.4445, b = -4.7750, c = 2.0315，以便对于较小的初始奇异值，迭代过程能更快收敛。本研究遵循相同的系数设置。

范数约束下的最速下降 Bernstein等人【3, Old Optimizer, New Norm: An Anthology. 2024. arXiv: 2409.20325 [cs.LG]】提出将深度学习中的优化过程视为范数约束下的最速下降。从这个角度看，Muon与Adam【19, Adam: A Method for Stochastic Optimization. 2015. ICLR 2015】, 【26, Decoupled Weight Decay Regularization. 2019. ICLR】之间的差异可以看作是范数约束的差异。Adam是在一个从Max-of-Max范数动态调整的a范数约束下的最速下降，而Muon提供了一个位于某个大p值的Schatten-p范数的静态范围内的范数约束【10, The Case for Muon. 2024. URL: https://x.com/leloykun/status/1846842887839125941】。当公式1被精确计算时，Muon提供的范数约束将是谱范数。神经网络的权重被用作输入空间或隐藏空间上的算子，这些空间通常是（局部）欧几里得的 【5, Deep Learning Optimizers as Steepest Descent in Normed Spaces. 2024. URL: http://leloykun.github.io/ponder/steepest-descent-opt/】，因此对权重的范数约束应该是一个诱导算子范数（对于权重矩阵而言是谱范数）。从这个意义上说，Muon提供的范数约束比Adam更合理 。

A2 方法细节

2.2 扩展Muon

权重衰减 尽管Muon在小规模上表现显著优于AdamW，如【16, Muon: An optimizer for hidden layers in neural networks. 2024. URL: https://kellerjordan.github.io/posts/muon/】所示，我们发现在扩展到用更多tokens训练更大模型时，性能增益会减小。我们观察到，权重和层输出的均方根（RMS）都持续增长到一个很大的规模，超过了bf16的高精度范围，这可能会损害模型的性能。为解决此问题，我们将标准的AdamW 【26, Decoupled Weight Decay Regularization. 2019. ICLR】权重衰减机制引入到Muon中。

$$\mathbf{W}_t = \mathbf{W}_{t-1} - \eta_t (\mathbf{O}_t + \lambda \mathbf{W}_{t-1})$$

权重衰减的有效性验证 我们对有和没有权重衰减的Muon进行了实验，以了解其对LLM训练动态的影响。基于我们在3.2节的缩放定律研究，我们用100B tokens（约5倍于最优训练tokens数）训练了一个800M参数的模型。图2展示了使用AdamW、原始Muon（无权重衰减）和带权重衰减的Muon训练的模型的验证损失曲线。虽然原始Muon初期收敛更快，但我们观察到一些模型权重随时间增长过大，可能限制了模型的长期性能。增加权重衰减解决了这个问题——结果表明，带权重衰减的Muon优于原始Muon和AdamW，在过训练（over-train）的情况下实现了更低的验证损失。因此，我们将更新规则调整为公式3，其中$\lambda$是权重衰减率。

一致的更新RMS Adam和AdamW【19, Adam: A Method for Stochastic Optimization. 2015. ICLR 2015】,【26, Decoupled Weight Decay Regularization. 2019. ICLR】的一个重要特性是它们维持一个理论上的更新RMS在1附近。然而，我们证明了Muon的更新RMS会根据参数的形状而变化，如下引理所述：

引理1. 对于一个形状为[A, B]的满秩矩阵参数，其理论上的Muon更新RMS是$1/\sqrt{max(A, B)}$。

证明见附录A。我们监控了训练过程中Muon的更新RMS，发现它通常接近上述理论值。我们注意到这种不一致性在扩大模型尺寸时可能会产生问题：

更新RMS不一致的问题
* 当max(A, B)过大时，例如稠密的MLP矩阵，更新会变得过小，从而限制了模型的表达能力，导致次优性能。
* 当max(A, B)过小时，例如将GQA【33, Fast Transformer Decoding: One Write-Head is All You Need. 2019. arXiv: 1911.02150 [http://cs.NE]】或MLA【9, DeepSeek-V3 Technical Report. 2024. arXiv: 2412.19437 [http://cs.CL]】中的每个KV头视为一个单独的参数，更新会变得过大，从而导致训练不稳定，同样导致次优性能。

维持一致更新RMS的策略 为了在不同形状的矩阵间保持一致的更新RMS，我们提议根据引理1，将每个矩阵的Muon更新按其$\sqrt{max(A, B)}$进行缩放以抵消该效应。第3.1节的实验表明，这一策略对优化是有益的。

与AdamW的更新RMS匹配 Muon旨在更新基于矩阵的参数。在实践中，AdamW与Muon结合使用，处理非矩阵参数，如RMSNorm、LM head和嵌入参数。我们希望优化器的超参数（学习率$\eta$，权重衰减$\lambda$）能在矩阵和非矩阵参数间共享。

匹配更新RMS的具体调整 我们建议将Muon的更新RMS匹配到与AdamW相似的水平。根据经验观察，AdamW的更新RMS通常在0.2到0.4之间。因此，我们通过以下调整将Muon的更新RMS缩放到这个范围：

$$\mathbf{W}_t = \mathbf{W}_{t-1} - \eta_t (0.2 \cdot \mathbf{O}_t \cdot \sqrt{\max(A, B)} + \lambda \mathbf{W}_{t-1})$$

我们通过经验结果验证了这一选择（详见附录A）。此外，我们强调，通过这种调整，Muon可以直接重用为AdamW调优的学习率和权重衰减。

其他超参数 Muon还包含另外两个可调超参数：牛顿-舒尔茨迭代步数N和动量$\mu$。我们凭经验观察到，当设置N为10时，迭代过程会比N=5产生更精确的正交化结果，但不会带来更好的性能。因此，为了效率，本研究中我们设置N=5。我们没有观察到调整动量能带来持续的性能提升，所以我们选择了0.95，与【16, Muon: An optimizer for hidden layers in neural networks. 2024. URL: https://kellerjordan.github.io/posts/muon/】相同 。

2.3 分布式Muon

ZeRO-1和Megatron-LM Rajbhandari等人【31, ZeRO: Memory optimizations Toward Training Trillion Parameter Models. 2020. SC20】引入了ZeRO-1技术，该技术将昂贵的优化器状态（如主权重、动量）划分到整个集群中。Megatron-LM【34, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2020. arXiv: 1909.08053 [http://cs.CL]】将ZeRO-1集成到其原生并行设计中。基于Megatron-LM复杂的并行策略，如张量并行（TP）、流水线并行（PP）、专家并行（EP）和数据并行（DP），ZeRO-1的通信负载可以从在整个分布式世界中收集数据减少到仅在数据并行组内收集。

方法 ZeRO-1对AdamW是高效的，因为它以逐元素的方式计算更新。然而，Muon需要完整的梯度矩阵来计算更新。因此，原始的ZeRO-1不直接适用于Muon。

算法1 分布式Muon

需要: 完整梯度G，DP分区的动量m，DP分区的参数p，动量μ。
1: // 在DP上对G进行Reduce-scatter以获得正确的梯度
2: g = reduce_scatter(G, dp_group)
3: // 使用局部分区的动量m对g应用动量
4: g′ = update_with_momentum(g, m, µ)
5: // DP Gather：跨DP将g′收集成一个完整的矩阵G
6: G = gather(g′, dp_group)
7: // 计算Muon更新
8: U = Newton-Schulz(G)
9: // 丢弃U的其余部分，只保留本地分区u，然后应用更新规则
10: p′ = apply_update(p, u)
11: // 将更新后的p′ All-gather成P
12: P = all_gather(p′, dp_group)
13: // 返回更新的RMS用于日志记录
14: return $\sqrt{u^2.mean()}$

分布式Muon的实现 我们提出了一种基于ZeRO-1的新的分布式解决方案，称为分布式Muon。分布式Muon遵循ZeRO-1在DP上划分优化器状态，并与原始的ZeRO-1 AdamW优化器相比引入了两个额外的操作：

1. DP Gather。对于一个本地DP分区的主权重（模型权重的1/DP大小），此操作是将相应的分区梯度收集成一个完整的梯度矩阵。
2. 计算完整更新。在上述收集之后，如2.1节所述，在完整的梯度矩阵上执行牛顿-舒尔茨迭代步骤。请注意，之后我们将丢弃部分完整更新矩阵，因为我们只需要与本地参数对应的分区来执行更新。

分布式Muon的实现如算法1所述。分布式Muon引入的额外操作用蓝色标出。

分析 我们在几个方面将分布式Muon与经典的基于ZeRO-1的分布式AdamW（为简洁起见，称为分布式AdamW）进行了比较：
* 内存使用。Muon只使用一个动量缓冲区，而AdamW使用两个。因此，Muon优化器使用的额外内存是分布式AdamW的一半。
* 通信开销。对于每个设备，额外的DP gather仅由本地DP分区的参数p需要。因此，通信成本小于G的reduce-scatter或P的all-gather。此外，Muon仅需在bf16中进行牛顿-舒尔茨迭代步骤，与fp32相比，进一步将通信开销减少了50%。总体而言，分布式Muon的通信负载是分布式AdamW的(1, 1.25]倍。上限的计算方式是：分布式Muon的通信量为4（fp32 G reduce-scatter）+ 2（bf16 Muon gather）+ 4（fp32 P all-gather），而分布式AdamW为4 + 4。在实践中，由于我们通常使用多个DP进行训练，经验上的额外成本通常更接近下限1.5。
* 延迟。分布式Muon的端到端延迟比分布式AdamW大，因为它引入了额外的通信并需要运行牛顿-舒尔茨迭代步骤。然而，这并不是一个显著问题，因为（a）只需大约5个牛顿-舒尔茨迭代步骤就能得到一个好的结果（在2.2节中讨论），并且（b）优化器引起的端到端延迟与模型的前向-后向传播时间相比是微不足道的（例如，通常为1%到3%）。此外，一些工程技术，如重叠gather和计算，以及重叠优化器reduce-scatter与参数gather，可以进一步减少延迟。

在我们的分布式集群中训练大规模模型时，与AdamW相比，分布式Muon没有明显的延迟开销。我们将很快发布一个为开源项目Megatron-LM【34, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2020. arXiv: 1909.08053 [http://cs.CL]】实现分布式Muon的拉取请求。

A4 实验环境

• 模型架构:

  • 缩放定律实验: 使用了一系列参数量从399M到1.5B的Llama架构【11, The Llama 3 Herd of Models. 2024. arXiv: 2407.21783 [http://cs.AI]】稠密模型。
  • 预训练: 使用了Moonlight模型，这是一个基于deepseek-v3-small【9, DeepSeek-V3 Technical Report. 2024. arXiv: 2412.19437 [http://cs.CL]】架构的3B激活/16B总参数的MoE模型（包含嵌入层）。具体修改见附录C。
  • SFT: 使用了Qwen2.5-7B基础模型【44, Qwen2.5 Technical Report. 2024. arXiv preprint arXiv:2412.15115】。

• 数据集:

  • 缩放定律实验: 使用了自有的专有数据集。
  • 预训练: 使用了K. Team 2025【40, Kimi k1.5: Scaling Reinforcement Learning with LLMs. 2025】中描述的数据集，最大上下文长度为8K，共训练5.7T tokens。
  • SFT: 使用了开源的tulu-3-sft-mixture数据集【20, Tülu 3: Pushing Frontiers in Open Language Model Post-Training. 2024】。

• 硬件配置:

  • 未明确说明具体的硬件型号，但提及在“分布式集群”中进行训练。

• 软件配置:

  • 代码实现基于Megatron-LM【34, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2020. arXiv: 1909.08053 [http://cs.CL]】项目，并集成了类似ZeRO-1【31, ZeRO: Memory optimizations Toward Training Trillion Parameter Models. 2020. SC20】的分布式优化策略。

A4 实验结果

3.1 一致的更新RMS

实验设置: 为了解决Muon更新RMS随参数矩阵形状变化的问题，我们比较了三种控制更新RMS的方法：

1. 基线(Baseline)：将更新矩阵乘以0.2 * H（H为模型隐藏层大小），以使其RMS与AdamW保持一致。
   
   $$\mathbf{W}_{t}=\mathbf{W}_{t-1}-\eta_{t}(0.2 \cdot \mathbf{O}_{t} \cdot \sqrt{H}+\lambda \mathbf{W}_{t-1})$$

2. 更新范数(Update Norm)：直接对通过牛顿-舒ль茨迭代计算出的更新进行归一化，使其RMS严格等于0.2。
   $$\mathbf{W}_t = \mathbf{W}_{t-1} - \eta_t(0.2 \cdot \mathbf{O}_t / \text{RMS}(\mathbf{O}_t) + \lambda \mathbf{W}_{t-1})$$
3. 调整学习率(Adjusted LR)：根据每个更新矩阵的形状，将其学习率乘以一个因子0.2 * sqrt(max(A, B))。
   $$\mathbf{W}_t = \mathbf{W}_{t-1} - \eta_t(0.2 \cdot \mathbf{O}_t \cdot \sqrt{\max(A, B)} + \lambda \mathbf{W}_{t-1})$$
   实验在800M参数的模型上进行，并特意将SwiGLU MLP替换为标准2层MLP，以凸显不同矩阵维度（[H, 4H] vs [H, H]）带来的影响。

实验结果与分析: 如表1所示，在训练了4B tokens后：
* 性能: “更新范数”和“调整学习率”方法均取得了比基线更好的训练和验证损失。
* 权重RMS:
* 对于MLP权重（形状[H, 4H]），这两种方法的权重RMS大约是基线的两倍，这符合预期，因为更新RMS大约是基线的两倍。
* 对于注意力查询权重（形状[H, H]），“调整学习率”方法的权重RMS与基线相似，而“更新范数”方法的权重RMS则较大，与其MLP权重类似。

结论: 基于这些发现，我们选择了成本更低的“调整学习率”方法用于后续实验。

表1：控制Muon在不同模型参数上的更新RMS

3.2 Muon的缩放定律

实验设置: 为了与AdamW进行公平比较，我们在Llama架构的系列稠密模型（参数量从399M到1.5B）上进行了缩放定律实验。我们首先通过网格搜索为AdamW建立了强大的基线（详见附录B和表2）。对于Muon，由于其更新RMS已与AdamW匹配，我们直接复用了AdamW的最优超参数。

实验结果与分析: 图3展示了拟合的缩放定律曲线，拟合方程见表3。从图1a可以看出，在计算最优设置下，Muon达到与AdamW相当的性能仅需约52%的训练FLOPs。

结论: 实验证明了Muon相较于AdamW在计算效率上的显著优势。

表2：缩放定律模型和超参数

*以8K上下文长度的样本数量计。

表3：缩放定律曲线的拟合参数

3.3 使用Muon进行预训练

模型与训练: 我们使用Muon从头预训练了一个名为“Moonlight”的模型。该模型采用deepseek-v3-small架构，拥有2.24B激活参数和15.29B总参数（不含嵌入层）。模型在5.7T tokens上进行了多阶段训练，包括学习率的预热、余弦衰减以及一个使用高质量数据（数学、代码、推理）的冷却（cooldown）阶段。

性能评估:

1. 中期性能 (1.2T tokens): 如表4所示，我们将Moonlight在训练1.2T tokens时的检查点与同样架构和相近训练量的Deepseek-v3-Small以及我们自己用AdamW训练的基线（Moonlight-A）进行了比较。结果显示，Moonlight显著优于Moonlight-A，证明了Muon的扩展有效性，尤其在数学和代码相关任务上表现突出。
2. 最终性能 (5.7T tokens): 如表5所示，完成全部5.7T tokens训练后，Moonlight与同等规模和更大规模的公开模型进行了比较。结果表明，Moonlight优于架构相似、训练tokens数相同的模型，并且即使与训练数据量远超自身的稠密模型相比，也保持了强大的竞争力。如图1b所示，Moonlight位于模型性能与训练预算的帕累托前沿。

结论: Moonlight的强大性能验证了Muon作为大规模LLM训练的有效优化器，能够显著提升训练效率和最终模型表现。

表4：不同模型在约1.2T tokens时的比较。

† 报告的参数数量不包括嵌入参数。

表5：不同模型在各项基准测试上的比较。

† 报告的参数数量不包括嵌入参数。‡ 我们使用TriviaQA的完整数据集测试了所有列出的模型。

3.4 奇异谱动态

实验设置: 为了验证Muon能使权重矩阵在更多样化的方向上进行优化的直觉，我们对使用Muon和AdamW训练的权重矩阵进行了谱分析。我们计算了权重矩阵的SVD熵，其定义如下：

$$H(\sigma)=-\frac{1}{\log n} \sum_{i=1}^{n} \frac{\sigma_{i}^{2}}{\sum_{j=1}^{n} \sigma_{j}^{2}} \log \frac{\sigma_{i}^{2}}{\sum_{j=1}^{n} \sigma_{j}^{2}}$$

实验结果与分析:
* 如图4所示，在1.2T tokens的预训练过程中，所有训练检查点和所有类别的权重矩阵中，Muon训练的权重SVD熵均高于AdamW。
* 这种差异在用于专家选择的路由器（router）权重中尤为显著，表明混合专家（MoE）模型能从Muon中获益更多。
* 附录F中的奇异值分布可视化进一步证实，超过90%的权重矩阵在使用Muon优化时具有更高的SVD熵。

结论: 实验结果为“Muon能为权重矩阵提供更多样化的更新谱”这一直觉提供了强有力的经验证据，这有助于模型探索更多样的优化方向。

3.5 使用Muon进行监督微调（SFT）

实验目的: 本节探究了Muon优化器在LLM训练的标准SFT阶段的表现。

实验1：预训练和SFT优化器的可互换性:
* 设置: 我们使用Muon和AdamW分别对[email protected]（Muon预训练）和[email protected]（AdamW预训练）进行了SFT。
* 结果: 如表6所示，（Muon预训练 + Muon微调）的组合性能最佳。当SFT的优化器与预训练优化器不同时，性能会下降。这表明Muon在SFT阶段的优势与预训练阶段的优化器选择密切相关。

实验2：在公开预训练模型上使用Muon进行SFT:
* 设置: 我们使用Muon对公开的Qwen2.5-7B基础模型（AdamW预训练）进行SFT。
* 结果: 如表7所示，Muon微调的模型与Adam微调的模型性能相当，没有显示出显著优势。

结论: Muon带来的好处在SFT阶段依然存在，尤其是当预训练和微调都使用Muon时。然而，预训练和微调阶段的优化器不匹配会影响性能。为了获得最佳性能，更有效的方式是在预训练阶段就应用Muon。

表6：探究预训练和SFT阶段优化器可互换性的影响。

表7：应用于Qwen2.5-7B预训练模型的SFT中，Adam和Muon优化器的比较。

A7 补充细节

4 讨论

未来的研究可以从以下几个方向进一步探索和扩展当前发现。

将所有参数纳入Muon框架
目前，Muon优化器与Adam优化器结合使用，某些参数仍由Adam优化。尽管这种混合方法可行，但仍有改进空间。将所有参数的优化完全整合到Muon框架内是一个具有重要研究价值的课题。

将Muon扩展到Schatten范数
Muon优化器可以解释为在谱范数下的最速下降法。考虑到Schatten范数的广泛适用性和通用性，将Muon扩展到涵盖一般的Schatten范数是一个有前景的方向。这种扩展可能会解锁额外的优化能力，并可能产生优于当前基于谱范数的实现的结果。

理解并解决预训练-微调不匹配问题
实践中观察到的一个显著现象是，使用AdamW预训练的模型在用Muon进行微调时性能次优，反之亦然。这种优化器不匹配为有效利用大量AdamW预训练的检查点带来了巨大障碍，因此需要进行严格的理论研究。精确理解其底层机制对于设计出稳健有效的解决方案至关重要。

A5 结论

在本技术报告中，我们对Muon在LLM训练中的可扩展性进行了全面研究。通过系统的分析和改进，我们成功地将Muon应用于一个在5.7万亿tokens上训练的3B/16B参数的MoE模型。我们的结果表明，Muon可以有效地替代AdamW成为大规模LLM训练的标准优化器，在训练效率和模型性能方面都提供了显著的优势。通过开源我们的实现、Moonlight模型以及中间训练检查点，我们旨在促进可扩展优化技术的进一步研究，并加速LLM训练方法的发展。

A6 附录

A 更新RMS

引理1的证明

证明 不失一般性，考虑正交矩阵$U \in R^{n \times n}$和$V \in R^{m \times m}$，其中$n \ge m \ge r$。我们将证明对于$X = U[:,:r]V[:r,:]$（Muon的更新具有相同的格式），其RMS值为$\sqrt{r/mn}$。根据矩阵乘法的定义：

$$X_{i,j}=\sum_{k=1}^{r}U_{i,k}V_{k,j}$$

RMS可以表示为：

$$\begin{aligned} \begin{aligned} \text{RMS}(X)^{2} & =\frac{1}{m n} \sum_{i=1}^{n} \sum_{j=1}^{m} \sum_{k=1}^{r} U_{i, k}^{2} V_{k, j}^{2} \\ & =\frac{1}{m n} \sum_{k=1}^{r}\left(\sum_{i=1}^{n} U_{i, k}^{2}\right)\left(\sum_{j=1}^{m} V_{k, j}^{2}\right) \\ & =\frac{1}{m n} \sum_{k=1}^{r} 1 \\ & =\frac{r}{m n} \end{aligned} \end{aligned}$$

因此，$RMS(X) = \sqrt{r/mn}$。对于矩阵是满秩的常见情况，$r=m$，得出$RMS(X) = \sqrt{1/n}$。□

Muon和AdamW之间一致的更新RMS 正如2.2节所讨论的，我们希望在Muon和AdamW优化器之间匹配更新RMS。这通过在小规模模型上的实验得到了验证。我们将Muon的更新RMS设置在[0.05, 0.1, 0.2, 0.4, 0.8]的范围内，并以AdamW为基准。我们在表8中报告了2k步（约2B tokens）时的损失和代表性权重矩阵的RMS。从结果中，我们发现0.2 RMS和0.4 RMS的表现相似，且远优于其他设置。这些发现与我们的经验观察一致，即AdamW的更新RMS在0.2∼0.4的范围内。我们选择将Muon的更新RMS控制在0.2。

表8：Muon更新RMS实验

*除第一列外，所有其他候选项均使用具有受控RMS的Muon。

B AdamW基线缩放定律

确定最优参数 为确保我们实验的公平性和准确性，我们在我们的专有数据集上进行了一系列实验，以推导出AdamW的最优缩放定律参数。这包括在受限的计算预算（FLOPs, C）下，确定最优的模型大小(N)、训练tokens数量(D)、学习率(η)和批量大小(B)。【18, Scaling Laws for Neural Language Models. 2020. arXiv: 2001.08361 [cs.LG]】, 【14, Training Compute-Optimal Large Language Models. 2022. arXiv: 2203.15556 [http://cs.CL]】, 【4, Deepseek llm: Scaling open-source language models with longtermism. 2024. arXiv preprint arXiv:2401.02954】表9展示了我们系统性参数搜索过程的结果。

表9：缩放定律参数与计算预算（FLOPs）之间的经验关系

超参数搜索 为了系统地确定AdamW基线中的最优缩放定律超参数，我们采用了多阶段搜索协议。首先，我们选择了多个计算预算（FLOPs水平），并根据先前研究的经验指南初始化了模型大小、学习率和批量大小。对于每个固定的FLOPs约束，我们改变模型大小N，同时反向调整训练token数D以保持C = 6ND，从而探索模型容量和数据效率之间的权衡。每个配置都被训练至收敛，并记录验证损失以确定N和D的帕累托最优组合。随后，在固定最优N-D对的情况下，我们通过网格搜索来优化学习率和批量大小，确保所有配置的稳定性和收敛性。为了减轻局部最小值的影响并增强鲁棒性，这个迭代过程重复了2-3次，逐步缩小超参数空间。

优化过程在图5中得到了进一步说明，该图描绘了在不同FLOPs预算下，损失景观作为训练tokens、学习率和批量大小的函数。每条碗状曲线代表特定FLOPs水平的损失曲面，其独特的全局最小值对应于最优的超参数配置。

C 模型架构

架构选择 Muon与模型架构无关，我们使用了一个与【9, DeepSeek-V3 Technical Report. 2024. arXiv: 2412.19437 [http://cs.CL]】中描述的Deepseek-V3-Small相似的模型，因为它是一个强大的开源权重模型，可作为基准。我们在Moonlight模型中做了一些小修改，并在此列出：

多令牌预测（MTP） 在我们的实验中，MTP并未对预训练显示出显著的好处。为简单起见，我们没有在Moonlight模型中引入MTP层。

Auxfree偏置更新 在【9, DeepSeek-V3 Technical Report. 2024. arXiv: 2412.19437 [http://cs.CL]】中，auxfree偏置的更新方式为：$b_i = b_i + u \times sign(e_i)$，其中u是更新率，$b_i$是第i个专家的偏置，$e_i$是该专家的违规率。我们稍微修改了更新规则为：$b_i = b_i + u \times (sign(e_i) - sign(e).mean())$，其中$sign(e).mean()$是所有专家违规率符号的平均值，以控制偏置的大小，同时不改变topk选择逻辑。

门控缩放因子 Deepseek-V2-Lite没有使用门控缩放因子，而Deepseek-V3使用了2.5的缩放因子。我们使用了2.446的缩放因子，以控制与稠密模型相似的输出rms。计算我们的门控缩放因子的代码可以在图6中找到。

D 训练稳定性

无损失或梯度范数尖峰 Moonlight的训练过程非常平滑，我们没有遇到任何损失尖峰或梯度范数尖峰。损失和梯度范数曲线可以在图7中看到（Moonlight用蓝色表示，由AdamW训练的Moonlight-A用红色表示）。

最大注意力Logit 在训练过程中，我们观察到，尽管训练损失和梯度范数在整个过程中保持稳定，但在初始训练阶段，特定层中的最大注意力logit（计算为全局批次中单个最大的logit值）呈现出明显的上升轨迹，超过了100的阈值。值得注意的是，与替代优化器相比，AdamW在控制这一指标方面表现出更健康的行为。

大注意力Logits比例分析 为了进一步研究这一现象的影响，我们引入了大注意力logits比例指标，定义为一个批次内超过100的注意力logits的比例。如图7所示，该比例始终保持在较低水平（约$10^{-4}$），表明极大的logit值是稀疏的。此外，随着训练的进行，最大logit值逐渐减小，表明优化动态变得更加健康。

RMSNorm Gamma权重衰减 值得注意的是，对RMSNorm的gamma参数应用权重衰减对于确保训练稳定性至关重要，因为它能有效防止每层出现过高的输出RMS值。

E 与更昂贵模型的比较

性能对比 表10展示了我们的Moonlight模型（使用Muon优化）与使用更多计算资源训练的公开可用模型之间的比较分析，包括LLama3.1-8B【11, The Llama 3 Herd of Models. 2024. arXiv: 2407.21783 [http://cs.AI]】、Gemma-9B【38, Gemma 2: Improving open language models at a practical size. 2024. arXiv preprint arXiv:2408.00118】和Qwen2.5-7B【44, Qwen2.5 Technical Report. 2024. arXiv preprint arXiv:2412.15115】。图8展示了Moonlight与领域内可比模型在GSM8k性能基准上的表现。

表10：不同模型在各项基准测试上的比较。

† 报告的参数数量不包括嵌入参数。‡ 我们使用TriviaQA的完整数据集测试了所有列出的模型。

F 权重矩阵的奇异值分布

可视化分析 我们通过为每个矩阵绘制其奇异值按降序排列的线图来可视化权重矩阵的奇异值分布，并按最大值进行归一化。如图9和图10所示，我们发现，对于大多数权重矩阵，由Muon优化的奇异值分布比AdamW的更平坦，这进一步证实了Muon可以提供更多样化的更新谱的假设。