月之暗面开源改进版Muon优化器，算力需求比AdamW锐减48%，DeepSeek也适用

克雷西
2025-02-23
17:31:33

来源：量子位

克雷西 发自 凹非寺
量子位 | 公众号 QbitAI

算力需求比AdamW直降48%，OpenAI技术人员提出的训练优化算法Muon，被月之暗面团队又推进了一步！

团队发现了Muon方法的Scaling Law，做出改进并证明了Muon对更大的模型同样适用。

在参数量最高1.5B的不同Llama架构模型上，改进后的Muon算力需求仅为AdamW的52%。

同时团队还基于DeepSeek架构训练出了一个16B的MoE模型，与改进后的优化算法一同开源。

Muon技术博客发布当时主要适用于较小的模型和数据集，作者留下了三个悬而未决的问题：

• Muon能否用于更大规模的训练？
• Muon能否在更大规模的GPU集群上使用？
• Muon是否同样适用于微调和强化学习？
  现在月暗团队用实验给出了回答——全部都是Yes。

消息一出，当时Muon的作者也都很激动，主要作者Keller Jordan表示这是Muon规模化的首个成功报告，为团队送上了祝贺。

另一名贡献者，当时负责Muon规模化实验的Hyperbolic Labs联创兼CTO Yuchen Jin也表示，月暗团队的这项成果，是Muon的一次胜利。

将AdamW特点引入Muon

在介绍月暗团队的工作之前，先来了解一下Muon是个什么样的技术。

这是一种神经网络隐藏层的2D参数优化器，主要作者是OpenAI深度学习团队的Keller Jordan。

这项成果发表于去年的12月8日，而Keller也是去年12月加入的OpenAI。

Muon的核心思想是通过正交化梯度更新矩阵，避免参数更新陷入局部极小，使模型能够学习到更加多样化的特征表示。

在94%的精度下，Muon把CIFAR-10在A100上的训练时间从3.3秒缩短至2.6秒。

不过当时Muon团队只证明了其在小型模型和数据集上的可行性，对于较大的模型能否适用则是个未知数。

现在经过月暗团队的改进之后，Muon被证明对于更大的模型和数据集同样适用。

针对模型本身，团队吸收了AdamW中的一些特点，移植到了Muon当中，具体包括两个方面。

一是引入了权重衰减机制，在权重更新公式中添加了一个带有衰减系数的项。

这样做的原因是作者发现直接将Muon应用到大规模训练时，模型权重和层输出的幅度会持续增长，最终超出bf16的高精度表示范围，损害模型性能。

在训练一个8亿参数模型至100B tokens（约5倍计算预算最优）的过程中，团队对比了AdamW、无权重衰减的Muon和带权重衰减的Muon。

结果显示，带权重衰减的Muon在过拟合阶段取得了最佳效果，验证了权重衰减的必要性。

第二项改进，是调整了Muon的参数更新尺度，使不同形状矩阵参数的更新幅度保持一致，并与AdamW的更新幅度匹配。

Muon的一个特点是，对于形状为[A,B]的矩阵参数，其理论更新幅度为sqrt(1/max(A,B))。

这导致不同形状矩阵参数的更新幅度差异很大，比如对于MLP这种宽矩阵，更新会过小，而将每个head看作独立矩阵时，更新又会过大。

此外，这个幅度也与AdamW不一致，给超参数的设置带来困难。

为了让不同矩阵参数的更新幅度匹配，并与AdamW保持一致，作者尝试了几种改进方案，最终选择直接基于形状调整每个参数的学习率。

其中0.2是通过实验确定的一个常数，用于将Muon的更新尺度与AdamW对齐。

除了对Muon本身的改进，要想将Muon用于更大规模的训练，还需要将其扩展到分布式训练环境中。

由于Muon需要完整的梯度矩阵来计算正交化的更新量，而现有的分布式训练框架（如ZeRO-1、Megatron-LM等）都假设优化器状态可以独立地按元素切分到不同设备上，所以它们无法直接支持Muon。

为了解决这个问题，论文作者提出了分布式Muon的并行化策略。

它在ZeRO-1的基础上引入了两个额外的操作:

• 一是在每个数据并行组内做梯度聚合通信，将分散的梯度切片合并成完整的矩阵；
• 二是基于聚合后的梯度矩阵并行计算正交化的更新量，然后只保留与本地参数对应的那一部分。

这种实现方式在最小化内存占用和通信开销的同时，最大限度地保留了原始Muon算法的数学性质。

证明Muon扩展可行性

基于上述Muon改进,作者取得了以下成果，作者在Llama架构的一系列稠密模型上，进行了Muon和AdamW的模型缩放对比实验。

结果表明,在计算预算最优的情况下，Muon的样本效率是AdamW的1.92倍，即训练FLOPS只需AdamW的52%，就能达到相当的性能。

这一发现证实了Muon在大规模训练中的效率优势。

在此基础之上，作者以DeepSeek-V3-Small架构作为基础，用改进的Muon训练了Moonlight模型。

Moonlight是一个MoE模型，具有15.29B的总参数和2.24B激活参数，训练token量为5.7T。

与相同规模和数据量的模型相比，Moonlight在英语理解与推理（MMLU、TriviaQA、BBH）、代码生成（HumanEval、MBPP）、数学推理（GSM8K、MATH、CMATH）、中文理解（C-Eval、CMMLU）等各类任务上都取得了明显更好的性能。

即使与使用更大数据集训练的稠密模型相比，Moonlight也展现了极强的竞争力。

与多个知名语言模型的对比表明，Moonlight在性能-训练预算平面上推进了帕累托前沿（Pareto Frontier）。

（注：帕累托前沿是一个经济学和管理学中的概念，描述的是在多目标决策问题中所有可能的最优解的集合，这些解在多个目标之间取得了最佳平衡。在帕累托前沿上的每一个点，都意味着一个目标的改善必然以牺牲另一个目标为代价，因此它代表了在多个目标之间实现的最佳权衡。）

为了进一步分析Muon更新矩阵参数的内在机制，作者对比了Muon和AdamW训练得到的模型在不同训练阶段的参数矩阵奇异值谱。

结果发现，Muon优化的矩阵在各层各类参数上，总是比AdamW有更高的奇异值熵。这从经验上验证了Muon通过正交化来学习更多样化表示的直觉。

最后，在Moonlight模型的基础上，作者还探索了Muon在指导微调阶段的效果，结果表明，在预训练和微调阶段均使用Muon的效果是最佳的。

技术报告：
https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
Code：
https://github.com/MoonshotAI/Moonlight
Moonlight模型：
https://huggingface.co/moonshotai/Moonlight-16B-A3B

Read More