迎接DeepSeek开源周[Kimi先开为敬]发布开源最新优化技术
月之暗面通过系统分析和改进，成功地将 Muon 应用于 3B/16B 参数的 MoE 模型训练，训练了 5.7 万亿个令牌。结果表明，Muon 可以有效地替代 AdamW 作为大规模 LLM 训练的标准优化器，在训练效率和模型性能方面具有显著优势。通过开源实现、Moonlight 模型和中间训练检查点，论文旨在促进可扩展优化技术的研究，并加速 LLMs 训练方法的发展。
代码 & 实现：
https://github.com/MoonshotAI/Moonlight
全系列模型 (预训练, 指令微调 & 中间检查点)：
https://huggingface.co/moonshotai
技术报告 Paper：
https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
具体来说，
- 添加权重衰减: 通过在Muon中引入标准的AdamW权重衰减机制，解决了模型参数和层输出RMS增长过大的问题。
- 调整参数更新尺度: 通过调整Muon的参数更新规则，确保在不同形状的矩阵之间保持一致的更新RMS，从而提高训练稳定性。
- 分布式实现: 开发了基于ZeRO-1风格的Muon分布式版本，实现内存优化和通信效率的提升。
实验设计
- 模型架构: 使用类似于Deepseek-V3-Small的模型架构，并对其进行了一些小的修改，以适应Moonlight模型的需求。
- 数据集: 使用Kimi团队提供的5.7万亿token的数据集进行预训练。
- 训练过程: 训练分为几个阶段，逐步增加学习率和批量大小，并在训练过程中使用不同的数据质量进行优化。
结果与分析
1.一致性更新 RMS：实验表明，调整后的学习率方法（Adjusted LR）比基线方法（Baseline）和仅保持与 AdamW 一致 RMS 的方法（Update Norm）表现更好。
2.扩展性验证：通过扩展法实验，Muon 在计算最优设置下仅需约 52% 的训练 FLOPs 即可达到与 AdamW 相当的性能
。[不得了啊！]
3.预训练性能：Moonlight 模型在 1.2T tokens 时的性能显著优于使用 AdamW 的 Moonlight-A 模型，证明了 Muon 的扩展有效性。
4.微调性能：在监督微调阶段，Muon 预训练和微调的模型表现优于仅使用 AdamW 预训练和微调的模型，但在微调阶段使用不同的优化器时，Muon 的优势不明显。