开源赛道太挤了！月之暗面开源新版Muon优化器

首页 » 开源赛道太挤了！月之暗面开源新版Muon优化器

文章来源于互联网:开源赛道太挤了！月之暗面开源新版Muon优化器

省一半算力跑出2倍效果，月之暗面开源优化器Muon，同预算下全面领先。

月之暗面和 DeepSeek 这次又「撞车」了。

上次是论文，两家几乎前后脚放出改进版的注意力机制，可参考《撞车 DeepSeek NSA，Kimi 杨植麟署名的新注意力架构 MoBA 发布，代码也公开》、《刚刚！DeepSeek 梁文锋亲自挂名，公开新注意力架构 NSA》。

这次是开源。

上周五，DeepSeek 刚刚官宣这周要连续开源 5 个代码库，却被月之暗面深夜截胡了。

昨天，月之暗面抢先一步开源了改进版 Muon 优化器，比 AdamW 优化器计算效率提升了 2 倍。

团队人员表示，原始 Muon 优化器在训练小型语言模型方面表现出色，但其在扩展到更大模型方面的可行性尚未得到证明。因此，团队人员确定了两种对扩展 Muon 至关重要的技术：

添加权重衰减：对扩展到更大模型至关重要。
一致的 RMS 更新：在模型更新上执行一致的均方根。

这些技术使得 Muon 能够在大规模训练中直接使用，而无需调整超参数。Scaling law 实验表明，与计算最优训练的 AdamW 相比，Muon 的计算效率提升了 2 倍。

基于这些改进，月之暗面推出了 Moonlight，这是一个 3B/16B 参数的 Mixture-of-Expert（MoE）模型，使用 Muon 进行了 5.7 万亿 tokens 的训练。该模型刷新了当前的「帕累托前沿」，换句话说，在相同的训练预算下，没有其他模型能在所有性能指标上同时超越它。

与之前的模型相比，Moonlight 也以更少的训练 FLOPs 获得了更好的性能。

如下图所示，该研究进行了 Scaling law 研究，将 Muon 与强大的 AdamW 基线进行了比较，结果展示了 Muon 的卓越性能。Muon 实现了与 AdamW 训练相当的性能，同时仅需要大约 52% 的训练 FLOP。

月之暗面不但开源了内存优化且通信高效的 Muon 实现代码，并且还发布了预训练、指令调优以及中间检查点，以支持未来的研究。

论文《 MUON IS SCALABLE FOR LLM TRAINING 》。

论文地址：https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
代码地址：https://github.com/MoonshotAI/Moonlight
模型地址：https://huggingface.co/moonshotai/Moonlight-16B-A3B

研究介绍

扩展 Muon

Muon 优化器为 OpenAI 研究者 Keller Jordan 等人在 2024 年提出的，他们的研究表明在小规模训练中 Muon 的表现显著优于 AdamW。

但月之暗面发现，当将其扩展到训练更大模型并使用更多 token 时，模型性能提升逐渐减弱。他们观察到，权重和层输出的 RMS 值持续增长，最终超出了 bf16 的高精度范围，这可能会损害模型的性能。

为了解决这个问题，月之暗面在 Muon 中引入了标准的 AdamW（Loshchilov 等人，2019）权重衰减机制。

为了探究这一机制，研究者对 Muon 进行了有无权重衰减的实验，他们训练了一个包含 800M 参数和 100B token（约为最优训练 token 量的 5 倍）的模型。图 2 展示了使用 AdamW、无权重衰减的原始 Muon 以及带权重衰减的 Muon 训练的模型的验证损失曲线。

结果显示，虽然原始 Muon 在初期收敛速度更快，但一些模型权重随时间的推移增长过大，可能会影响模型的长期性能。

加入权重衰减后解决了这一问题 —— 结果表明，带权重衰减的 Muon 优于原始 Muon 和 AdamW，获得了更低的验证损失。公式 3 为表达式，其中 λ 为权重衰减比率。

一致的 RMS 更新。研究者发现 Adam 和 AdamW 的一个重要特性是，它们将更新的 RMS 维持在 1 左右。然而，月之暗面发现 Muon 更新 RMS 会根据参数矩阵形状的变化而变化，具体如下引理 1 所示：

为了在不同形状矩阵之间保持一致的 RMS 更新，该研究通过

来扩展 Muon 矩阵更新，从而抵消引理 1 中提到的影响。

在实际应用中，研究者通常将 AdamW 与 Muon 结合使用，以处理非矩阵参数。本文希望优化器超参数（学习率 η、权重衰减 λ）能够在矩阵参数和非矩阵参数之间共享。

因此他们提出将 Muon 更新的 RMS 调整到与 AdamW 相似的范围。他们通过以下调整将 Muon 更新 RMS 缩放至这一范围：

分布式 Muon

月之暗面团队还提出了一种基于 ZeRO-1 的分布式解决方案，称为分布式 Muon（Distributed Muon）。分布式 Muon 遵循 ZeRO-1 在数据并行（DP）上对优化器状态进行划分，并与普通的 ZeRO-1 AdamW 优化器相比引入了两个额外的操作，算法 1 描述了分布式 Muon 的实现。

实验

RMS 的一致性

为了使所有矩阵参数更新的 RMS 值与 AdamW 的 RMS 保持一致，研究团队尝试了两种方法来控制参数更新的 RMS，并将其与只用了 AdamW 的基线的 RMS 进行了对比。

由于大规模训练模型时，会出现各种意料之外的情况，因此，研究团队测试了 Muon 对训练早期阶段的影响。当矩阵维度差异增大时，更新 RMS 不一致的问题会更加明显。该团队对模型架构进行了微调，用标准的 2 层 MLP 替换了 Swiglu MLP，并将其矩阵参数的形状从 [H, 2.6H] 改为 [H, 4H]。

团队评估了模型的损失，并监控了关键参数的 RMS，尤其是形状为 [H, H] 的注意力查询权重和形状为 [H, 4H] 的 MLP 权重。

实验结果表明（见表 1），Update Norm 和 Adjusted LR 均优于基线方法，且 Adjusted LR 的计算成本更低，因此被选用于后续实验。

Muon 的 Scaling Law

为了与 AdamW 公平比较，该团队在一系列基于 Llama 架构的模型上对 Muon 进行了拓展。

对于 Muon，由于其 RMS 与 AdamW 匹配，团队直接复用了 AdamW 的超参数。

实验结果显示，拟合的 Scaling Law 曲线表明，在计算最优设置下，Muon 仅需约 52% 的训练 FLOPs 即可达到与 AdamW 相当的性能。这进一步说明了 Muon 在大规模语言模型训练中的高效性。

使用 Muon 进行预训练

为了评估 Muon 在模型架构中的表现，该团队使用 DeepSeek-V3-Small 架构从头开始预训练了 Moonlight 模型。

Moonlight 模型总共进行了 5.7 万亿 tokens 的训练，但在训练到 1.2 万亿 tokens 的阶段，团队将其与 DeepSeek-V3-Small（使用 1.33T tokens 训练的 2.4B/16B 参数 MoE 模型）和 Moonlight-A（与 Moonlight 设置相同，但使用 AdamW 优化器）进行了比较。如表 4 所示，Moonlight 在语言、数学和编码等任务上都显著优于 Moonlight-A，证明了 Muon 的扩展优势。