在 DeepSpeed 中使用 Muon 优化器

要点速览

DeepSpeed 现已支持 Muon 优化器！Muon 优化器凭借其卓越表现，已获得前沿 AI 实验室的广泛采用。其中之一便是月之暗面（Moonshot AI），他们已采用 Muon 优化器来训练其大型基础模型，如 Kimi-K2-Thinking。本文将深入探讨 Muon 优化器是什么，以及它在 DeepSpeed 上的表现。

什么是 Muon 优化器？

Muon 是一款专为神经网络隐藏层二维权重设计的优化器。它获取权重的梯度，计算其动量，并应用牛顿-舒尔茨（Newton-Schulz）迭代对动量矩阵进行正交化，然后使用该正交化后的矩阵来更新权重。由于 Muon 仅维护一个动量缓冲区（而 Adam 需要两个），因此它在优化器状态上占用的内存更少。

正交化步骤是 Muon 在预训练中实现收敛优势的关键。在实践中，Transformer 中二维权重的梯度更新往往具有极高的条件数——它们几乎是低秩的，被少数几个较大的奇异方向所主导。通过对动量矩阵进行正交化，Muon 使所有奇异值趋于相等，从而有效地放大了那些本会被掩盖的、罕见但重要的更新方向。这带来了更好的样本效率：在 NanoGPT 速度竞赛基准测试中，Muon 的训练速度比 AdamW 提高了 35%，并且在 15 亿参数规模下，其达到 GPT-2 XL 级别性能的速度比 AdamW 快约 25%。

与需要为每个参数设置两个动量缓冲区的 Adam 优化器不同，Muon 优化器仅需要一个动量缓冲区。这意味着对于使用 Muon 优化器的参数，我们只需要分配一个动量缓冲区，与 Adam 相比可以节省内存。

Muon 被用于 Keller Jordan 的 NanoGPT 修改版、Andrej Karpathy 的 nanochat 中，Muon 的一个变体（MuonClip）也被生产级模型 月之暗面 Kimi-K2 所采用。最近，智谱 AI 的 GLM-5（744B 参数）确认在 GLM-4.5 和 GLM-5 的预训练中使用了 Muon 优化器，并结合了“Muon Split”技术，该技术通过注意力头对 MLA 上投影矩阵进行分割并独立正交化每个头，解决了在使用 Muon 时 MLA 和 GQA 之间的性能差距。DeepSeek-V4（1.6T 参数）也采用了 Muon 优化器以实现更快的收敛和更高的训练稳定性。

DeepSpeed 中的 Muon 优化器支持

在 DeepSpeed 中应用 Muon 优化器面临的一个挑战是，之前的优化器（SGD、Adam）将梯度视为扁平化（flattened）的缓冲区。因此，由于梯度缓冲区已经是扁平化的，很难在同一位置替换 Muon 优化器。我们将 Muon 更新移至 stage 1 和 stage 2 DeepSpeedZeroOptimizer 的 get_flat_partition 函数中，此时每个参数的梯度仍处于未扁平化阶段，因此我们可以轻松应用 Muon 更新。

Muon 优化器适用于二维权重矩阵（注意力权重和 MLP 权重）。它对动量矩阵应用牛顿-舒尔茨正交化，这要求权重必须是二维的。非二维参数（嵌入层、层归一化、偏置、lm_head）将回退到 AdamW。我们在模型引擎初始化器中应用了解析逻辑，将模型参数标记为 use_muon（仅当模型参数是二维且属于隐藏层时）。当使用 Muon 优化器时，任何被标记为 use_muon 的参数都将使用 Muon 优化器来更新权重。
请注意，Muon 是一种混合优化器：它仅对二维隐藏权重使用 Muon 更新，而对所有其他参数（嵌入层、层归一化、偏置、lm_head）回退到 Adam。DeepSpeed 配置支持通过 muon_lr（用于 Muon 参数）和 adam_lr（用于 Adam 参数）设置独立的学习率。

使用 Muon 优化器运行 DeepSpeed 微调

Deepspeed 微调演示是一个用于在单一位置使用不同 DeepSpeed 训练功能并比较其性能的示例。您可以使用它来测试使用 Muon 优化器对 LLM 模型进行微调。

git clone https://github.com/delock/deepspeed_finetune_demo
cd deepspeed_finetune_demo

./finetune.sh z2_muon.json

Muon 优化器收敛实验结果

我们通过微调 Moonlight-16B-A3B（一个总参数 16B、活跃参数 3B 的混合专家模型）测试了 Muon 优化器，并在代码生成（MBPP/MBPP+）、通用知识（MMLU）和数学推理（GSM8K）基准上进行了评估。每个基准测试都使用其特定领域的训练集。

训练配置

• 模型：Moonlight-16B-A3B (MoE, 16B 总量 / 3B 活跃)
• 训练数据集：用于 MBPP/MBPP+ 的 sahil2801/CodeAlpaca-20k，用于 MMLU 的 cais/mmlu (auxiliary_train, ~95k 条示例)，用于 GSM8K 的 meta-math/MetaMathQA (sample_rate=0.1, ~39.5k 条示例)
• ZeRO Stage 2, bf16, 专家并行 (autoep_size=4)
• 批次大小：16，梯度累积：2，4 张 GPU
• 1 个 epoch，梯度裁剪：1.0

评估结果

Muon 在 4 个指标中的 3 个上优于 AdamW：MBPP+（0.548 对 0.534，+1.4pp）、MMLU（0.678 对 0.660，+1.8pp）和 GSM8K（0.810 对 0.805，+0.5pp）。在 MBPP 基础测试中，AdamW 以微弱优势领先 Muon（0.661 对 0.646，-1.5pp），但 Muon 在具有额外测试用例的更严格的 MBPP+ 上得分更高（0.548 对 0.534），这表明其泛化能力更强。

Muon 优化器内存节省

Muon 优化器在优化器状态上使用的内存比 Adam 少，因为它为每个参数维护一个动量缓冲区，而不是两个（一阶和二阶矩）。
内存使用对比

请注意，Muon 是一种混合优化器：二维隐藏权重使用 Muon（1 个缓冲区），而其余参数（嵌入层、层归一化、lm_head）仍使用 Adam（2 个缓冲区）。实际内存节省取决于作为二维隐藏权重的参数比例。对于典型的 Transformer 模型，约 90% 的参数是二维隐藏权重，因此优化器状态内存大约减少了 45%。然而，由于 GPU 总内存还包括模型权重、梯度和激活值，端到端的内存减少量较小（见下文测量结果）。

测量 GPU 内存：Qwen2.5-3B 微调

我们使用上述相同的 8x A100 (40GB) 配置（批次大小 32，ZeRO Stage 2，bf16）测量了在 tatsu-lab/alpaca 数据集上微调 Qwen2.5-3B 期间的峰值 GPU 内存。

与 AdamW 相比，Muon 将每张 GPU 的内存需求减少了约 3 GiB（9%）。节省完全来自优化器状态：Muon 参数存储一个动量缓冲区（4 字节）而不是 Adam 的两个（8 字节）。然而，由于优化器状态只是 GPU 总内存（包括模型权重、梯度和激活值）的一个组成部分，因此端到端的减少量比较适中。对于更大的模型或更严格的内存预算，这 9% 的节省可能决定了工作负载是可以在设备上运行，还是需要 CPU 卸载。

未来展望

Muon 正在社区中迅速获得关注，Kimi-K2（1T 参数）和 GLM-5（744B 参数）的生产级采用表明，它已成为替代 Adam 作为大规模训练默认优化器的强有力竞争者。我们正在积极构建 DeepSpeed 对 Muon 的全面支持，一系列改进措施已经在进行中：

• ZeRO Stage 2 支持 — 已合并
• ZeRO Stage 3 支持 — 已合并
• 基于 Gram-Schmidt 的牛顿-舒尔茨迭代 — 一种更快的正交化内核，正在审查中
• CPU 卸载 — 进行中
• MuonClip — Kimi-K2 所使用的变体，计划中

我们欢迎任何关于 DeepSpeed 中 Muon 优化器支持的想法、反馈和贡献——请发起 Issue 进行讨论，或向 DeepSpeed 提交 PR。让我们一起让 Muon 在 DeepSpeed 中变得稳健且闪电般快速！