企业官网建设流程全解析

DeepSpeed 集成 Muon 优化器实战：从正交化原理到 ZeRO 微调配置

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

导读

本文讲解 DeepSpeed 对 Muon（MomentUm Orthogonalized by Newton-schulz）优化器的原生支持。Muon 是一种针对神经网络隐藏 2D 权重矩阵设计的优化器，它用牛顿-舒尔茨迭代对动量矩阵做正交化处理，比 Adam 少维护一个动量缓冲区。读完本文，你将掌握：Muon 的收敛与显存优势来源、DeepSpeed 内部如何为它改造 ZeRO Stage 1/2 的参数分区更新流程、如何编写type: "muon"的 DeepSpeed 配置并配合muon_lr/adam_lr双学习率做 LLM 微调，以及它当前的支持边界与演进路线。

TL;DR

Muon 优化器正在被越来越多前沿 AI 实验室采用并用于大规模预训练，例如 Moonshot AI 使用 Muon 训练其 Kimi-K2-Thinking 基础大模型，智谱的 GLM-5（744B 参数）确认 GLM-4.5 与 GLM-5 预训练使用了 Muon 及按注意力头独立正交化的 "Muon Split" 技术，DeepSeek-V4（1.6T 参数）也使用 Muon 以获得更快收敛与更强训练稳定性。DeepSpeed 现已支持 Muon 优化器，社区围绕该特性的讨论与合并进度以 blogs/muon-optimizer/README.md 为准。

什么是 Muon 优化器？

Muon 是一种针对神经网络隐藏 2D 权重矩阵设计的优化器。它的核心思路是：取权重的梯度后先计算动量，再通过牛顿-舒尔茨迭代（Newton-Schulz iterations）对动量矩阵做正交化，最后用这个正交化后的矩阵去更新权重。由于 Muon 每个参数只维护一份动量缓冲区（而 Adam 需要一阶矩和二阶矩两份），因此优化器状态的显存占用比 Adam 更少。

正交化是收敛优势的关键

正交化步骤是 Muon 在预训练中获得收敛优势的关键。实际中，Transformer 中 2D 权重的梯度更新往往具有非常高的条件数——它们近似低秩，能量集中在少数几个大的奇异方向。通过对动量矩阵做正交化，Muon 将所有奇异值均衡化，相当于放大了那些会被大方向掩盖的、罕见但重要的更新方向，从而带来更好的样本效率。

以社区公开基准为例：在 NanoGPT speedrunning 基准中，Muon 相比 AdamW 将训练速度提升了约 35%；在 1.5B 参数规模上，达到 GPT-2 XL 水平的性能比 AdamW 快约 25%。Muon 还广泛用于 Keller Jordan 的 modded-NanoGPT、Andrej Karpathy 的 nanochat，其变体 MuonClip 则被生产级 LLM Kimi-K2 使用。

与 Adam 的内存对比

与 Adam 每个参数需要两份动量缓冲区（一阶矩 m 与二阶矩 v）不同，Muon 每个参数只需要一份动量缓冲区。这意味着对使用 Muon 的参数，优化器状态内存可以较 Adam 明显节省。

DeepSpeed 对 Muon 的支持方式与源码剖析

将 Muon 接入 DeepSpeed 并非简单替换优化器类，它需要在 ZeRO 分区更新流程的关键位置做改造。下面从源码角度拆解它如何在当前仓库中落地。

挑战：梯度在 ZeRO 中被展平了

将 Muon 应用于 DeepSpeed 的一个挑战在于：此前的优化器（SGD、Adam）把梯度看作展平的缓冲区，因此在它们所在的位置很难直接换入 Muon——因为那里的梯度缓冲区已经是展平的，无法还原成用于正交化的 2D 矩阵。DeepSpeed 的解决方案是把 Muon 更新移入 Stage 1 与 Stage 2 的DeepSpeedZeroOptimizer的get_flat_partition函数：在该函数内部，每个参数的梯度仍处于未展平状态，可以方便地施加 Muon 更新。对应实现位于 deepspeed/runtime/zero/stage_1_and_2.py，其中通过getattr(param, "use_muon", False)判断参数是否走 Muon 路径。

use_muon参数打标：谁是 2D 隐藏权重？

Muon 只作用于 2D 权重矩阵（attention 与 MLP 权重），对动量矩阵施加牛顿-舒尔茨正交化要求权重必须是 2D。非 2D 参数（embeddings、layer norm、bias、lm_head）回退到 AdamW。

DeepSpeed 在模型引擎初始化器中做了一次参数解析来打标。该逻辑位于 deepspeed/init.py 的set_optimizer_flags：

def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None: if config_class.optimizer_name == MUON_OPTIMIZER: for name, p in model.named_parameters(): if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): setattr(p, "use_muon", True) else: setattr(p, "use_muon", False)

判定规则清晰可读：p.ndim >= 2且参数名中不含embed与lm_head时标记为use_muon=True，否则为False。当优化器为 Muon 时，凡是打上use_muon标记的参数都用 Muon 更新权重。如果你在代码中直接构造参数组使用 Muon，deepspeed/runtime/engine.py 会校验所有参数是否都配置了use_muon属性，缺失时提示用param.use_muon = True / False显式设置。

混合优化器的参数组构造：muon_lr与adam_lr

需要强调的是，Muon 是一种混合优化器：它只对 2D 隐藏权重使用 Muon 更新，其余参数（embeddings、layer norm、bias、lm_head）全部回退到 Adam。因此 DeepSpeed 配置里支持通过muon_lr（Muon 参数的学习率）与adam_lr（Adam 参数的学习率）分别设置两类参数的学习率。

在 deepspeed/runtime/engine.py 中，引擎读取优化器配置后会把模型参数按use_muon拆成两个参数组：

• Muon 参数组（use_muon=True，名为muon-params）接受lr、momentum、weight_decay、muon_lr、ns_method等键，其中muon_lr会覆盖lr；
• Adam 参数组（use_muon=False，名为adam-params）接受lr、betas、eps、weight_decay、adam_lr等键，其中adam_lr会覆盖lr。

随后构造MuonWithAuxAdam，把 Adam 参数组交给可选的辅助优化器处理。优化器类型常量MUON_OPTIMIZER = 'muon'定义在 deepspeed/runtime/config.py，默认的 AdamW 模式常量ADAM_W_MODE_DEFAULT = True也在同一文件。若模型含有 MoE 层，engine 还会调用split_params_into_different_moe_groups_for_optimizer将参数组进一步拆分，确保与专家并行（EP）参数分组兼容。

核心算子：牛顿-舒尔茨正交化内核

Muon 相关的 PyTorch 实现集中放在deepspeed/runtime/zero/muon/目录下，包含：

• deepspeed/runtime/zero/muon/original_muon.py：上游移植的 Muon 实现。定义了Muon、SingleDeviceMuon、MuonWithAuxAdam、SingleDeviceMuonWithAuxAdam四个优化器类；支持 bf16 下稳定运行的zeropower_via_newtonschulz5（五次牛顿-舒尔茨迭代，系数a=3.4445, b=-4.7750, c=2.0315，用于最大化零点斜率）、以及zeropower_via_gram_newtonschulz（在小型方阵 Gram 矩阵R = X @ X.T上迭代，对宽矩阵显著降低 FLOP，并使用 fp16 提高半精度数值精度）。
• deepspeed/runtime/zero/muon/muon_optimizer.py：DeepSpeed 对MuonWithAuxAdam的封装子类。核心是在step()中对use_muon=True的参数组只执行解耦权重衰减加动量更新：p.mul_(1 - lr * weight_decay)与p.add_(p.grad.reshape(p.shape), alpha=-lr)——因为此时参数是展平版本，不适合直接正交化，真正的正交化已前移到 ZeRO 的get_flat_partition阶段。而use_muon=False的参数组则交给外部传入的 Adam 优化器（如 FusedAdam）或内部联机 Adam 更新处理；若 FusedAdam 初始化失败且fallback_to_inline=True，会回退到 Muon 自带的联机 Adam 更新。

muon_update支持的两种牛顿-舒尔茨方法由NS_METHODS = {"standard", "gram"}限定，其中ns_method的默认值是"gram"（矩形矩阵上约 2 倍加速），"standard"则是原始迭代。这些算子都套用了 DeepSpeed 的@compiler.compile()装饰器，可借助编译后端加速。

ZeRO Stage 支持与reduce_scatter约束

Muon 在 DeepSpeed 中的使用与 ZeRO 状态划分方式强相关，官方文档明确说明了两条边界：

1. ZeRO Stage 1 和 Stage 2 支持 Muon，前提是关闭 optimizer offload，并搭配reduce_scatter使用。
2. 分区边界处理：如果某个 Muon 矩阵跨越了 ZeRO 分区边界，DeepSpeed 会在施加非线性 Muon 更新之前，将完整归约后的梯度提供给持有该矩阵分区的每一个 rank。从 deepspeed/runtime/zero/stage_1_and_2.py 的源码结构看，相关逻辑通过检查参数的use_muon属性以及分区数量（len(partition_ids) > 1）来决定是否为跨分区的 Muon 矩阵保留完整的梯度分区信息。
3. 暂不支持组合：Muon 搭配reduce_scatter同时又启用 optimizer offload 的组合目前不被支持。

ZeRO Stage 3 支持

在 deepspeed/runtime/zero/stage3.py 中同样存在完整的 Muon 分支：例如self.use_muon = isinstance(self.optimizer, MuonWithAuxAdam)，并在reduce_scatter、all2all_process_group、动量缓冲区的内存保存策略（save_muon_momentum_buffer_in_memory）、参数分区等环节针对 Muon 做专门处理。官方在 “What's Next” 中也确认 ZeRO Stage 2 与 Stage 3 支持均已合并进主干。

如何用 DeepSpeed 配置启用 Muon

在 DeepSpeed 中启用 Muon，只需把配置里的优化器type写为muon，并在params中给出各类参数。下面是一个与官方收敛实验配置（ZeRO Stage 2、bf16、4 GPU 微调 Moonlight-16B-A3B）相符的参考配置：

{ "train_batch_size": 16, "gradient_accumulation_steps": 2, "train_micro_batch_size_per_gpu": 2, "optimizer": { "type": "muon", "params": { "muon_lr": 1e-4, "adam_lr": 2e-6, "momentum": 0.95, "weight_decay": 0.1, "betas": [0.9, 0.95], "eps": 1e-8 } }, "zero_optimization": { "stage": 2, "reduce_scatter": true }, "bf16": { "enabled": true }, "gradient_clipping": 1.0 }

几个关键配置项说明：

• type: "muon"：对应常量MUON_OPTIMIZER = 'muon'，引擎据此走MuonWithAuxAdam构造路径。
• muon_lr：2D 隐藏权重（Muon 参数组）的学习率，会覆盖通用lr。官方微调实验使用1e-4。
• adam_lr：embeddings、layer norm、bias、lm_head（Adam 参数组）的学习率，同样覆盖lr。官方实验使用2e-6。
• momentum：Muon 的动量系数，默认值 0.95，通常直接可用。
• betas/eps：仅作用于回退 Adam 的参数组；若缺省，MuonWithAuxAdam会填入默认值（lr 3e-4、betas (0.9, 0.95)、eps 1e-10 等），Muon 参数组缺省时也有默认 lr 0.02、momentum 0.95、ns_method="gram"。
• ns_method：可选"gram"（默认，矩形矩阵上更快）或"standard"。
• zero_optimization.stage：Muon 支持的组合为 Stage 1/2（关闭 optimizer offload、使用reduce_scatter）与 Stage 3；Muon +reduce_scatter+ optimizer offload 的组合会直接报错。

测试用例 tests/unit/ops/muon/test_muon.py 对这些行为做了覆盖：例如test_ns_method_training验证不同ns_method在各 ZeRO stage 下可训练、test_ns_method_stage3验证 Stage 3 下的训练、test_muon_reduce_scatter_with_optimizer_offload_raises验证"Muon + reduce_scatter + optimizer offload 抛错"的约束、test_update_matches_full_gradient_reference验证 Muon 更新与全量梯度参考一致。tests/unit/v1/ops/muon/test_muon_partial_training.py则覆盖了部分参数训练场景。

使用 Muon 微调 LLM：实操演示

社区提供了一个可一键对比多种 DeepSpeed 训练特性的微调演示仓库deepspeed_finetune_demo，你可以在同一处测试不同 DeepSpeed 特性的性能，也可以直接用它验证 Muon 微调 LLM：

git clone https://github.com/delock/deepspeed_finetune_demo cd deepspeed_finetune_demo ./finetune.sh <NUM_GPUS> <MODEL_NAME> z2_muon.json

命令中的z2_muon.json即对应上文给出的 ZeRO Stage 2 + Muon 配置文件，<NUM_GPUS>为 GPU 数量、<MODEL_NAME>为目标模型名（例如下文实验中的 Moonlight-16B-A3B）。

Muon 收敛性实验与评测结果

实验设置

官方用 Muon 微调了 Moonlight-16B-A3B（MoE 模型，总参数量 16B、激活参数量 3B），并在代码生成（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）
• 批大小：batch size 16，梯度累积 2，共 4 块 GPU
• 训练轮数：1 epoch，梯度裁剪 1.0

评测结果对比

在这 4 项指标中，Muon 有 3 项超过 AdamW：MBPP+（0.548 vs 0.534，+1.4 个百分点）、MMLU（0.678 vs 0.660，+1.8 个百分点）和 GSM8K（0.810 vs 0.805，+0.5 个百分点）。在 MBPP 基础测试集上 AdamW 略胜（0.661 vs 0.646，-1.5 个百分点），但在包含更多测试用例、更为严格的 MBPP+ 上 Muon 得分更高（0.548 vs 0.534），说明 Muon 的泛化能力更好。读者在使用前也应注意到，这是单一模型、特定领域微调数据上的观察结果，不同模型与任务上的收益需要自行验证。

Muon 的显存节省实测

由于 Muon 每个参数只维护一份动量缓冲区（而非 Adam 的两份），其优化器状态显存低于 Adam。但 Muon 是混合优化器：2D 隐藏权重用 Muon（1 份缓冲区），其余参数（embeddings、layer norm、lm_head 等）仍用 Adam（2 份缓冲区），因此实际节省取决于 2D 隐藏权重占参数总量的比例。

每参数状态内存对比

对于典型 Transformer，约 90% 的参数是 2D 隐藏权重，因此优化器状态内存大约能减少 45%。不过由于 GPU 总显存还包含模型权重、梯度和激活值，端到端显存降幅会更小（见下方实测）。

Qwen2.5-3B 微调实测峰值显存

官方在 8 块 A100（40GB）上、用与上文一致的配置（batch size 32、ZeRO Stage 2、bf16）微调 Qwen2.5-3B（数据集 tatsu-lab/alpaca）并测量每 GPU 峰值显存：

Muon 每 GPU 显存比 AdamW 低约 3 GiB（约 9%）。节省完全来自优化器状态：Muon 参数存一份动量缓冲区（4 字节），而 Adam 存两份（8 字节）。因为优化器状态只是总显存的一部分（其余为模型权重、梯度与激活），端到端降幅相对温和。对更大模型或显存更紧张的场景，这 9% 的节省可能恰好决定工作负载能否放进设备显存、还是需要触发 CPU offload。

演进路线（What's Next）

随着社区采纳度快速提升以及 Kimi-K2（1T 参数）、GLM-5（744B 参数）等生产级应用的出现，Muon 正被视为大规模训练中取代 Adam 成为默认优化器的有力候选。DeepSpeed 正在持续推进 Muon 的完整支持，当前进度如下：

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

如果你关注 DeepSpeed 中 Muon 的最新进展、想参与讨论或提交 PR，可以围绕上述路线图在仓库中发起 issue 或贡献代码，让 Muon 在 DeepSpeed 中更稳定、更快速。

关键源码索引

• Muon 支持公告与评测：blogs/muon-optimizer/README.md
• 优化器实现（含 Newton-Schulz 内核与 Gram 变体）：deepspeed/runtime/zero/muon/original_muon.py
• DeepSpeed 封装与 Adam 回退：deepspeed/runtime/zero/muon/muon_optimizer.py
• 参数打标逻辑：deepspeed/init.py
• 优化器名称与 AdamW 默认值：deepspeed/runtime/config.py
• 双参数组构造与muon_lr/adam_lr覆盖规则：deepspeed/runtime/engine.py
• ZeRO Stage 1/2 中use_muon分区更新处理：deepspeed/runtime/zero/stage_1_and_2.py
• ZeRO Stage 3 中 Muon 分支：deepspeed/runtime/zero/stage3.py
• Muon 单元测试：tests/unit/ops/muon/test_muon.py、tests/unit/v1/ops/muon/test_muon_partial_training.py

致谢与贡献

本文涉及的 DeepSpeed Muon 支持工作由 Wang, Zhipeng (@PKUWZP)、Peng Du (@pengdurice)、Chi McIsaac (@qimcis) 与 Ma, Guokai (@delock) 共同贡献。

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed