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 对 MuonMomentUm Orthogonalized by Newton-schulz优化器的原生支持。Muon 是一种针对神经网络隐藏 2D 权重矩阵设计的优化器它用牛顿-舒尔茨迭代对动量矩阵做正交化处理比 Adam 少维护一个动量缓冲区。读完本文你将掌握Muon 的收敛与显存优势来源、DeepSpeed 内部如何为它改造 ZeRO Stage 1/2 的参数分区更新流程、如何编写type: muon的 DeepSpeed 配置并配合muon_lr/adam_lr双学习率做 LLM 微调以及它当前的支持边界与演进路线。TL;DRMuon 优化器正在被越来越多前沿 AI 实验室采用并用于大规模预训练例如 Moonshot AI 使用 Muon 训练其 Kimi-K2-Thinking 基础大模型智谱的 GLM-5744B 参数确认 GLM-4.5 与 GLM-5 预训练使用了 Muon 及按注意力头独立正交化的 Muon Split 技术DeepSeek-V41.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_flagsdef 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_muonTrue否则为False。当优化器为 Muon 时凡是打上use_muon标记的参数都用 Muon 更新权重。如果你在代码中直接构造参数组使用 Muondeepspeed/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_lrMuon 参数的学习率与adam_lrAdam 参数的学习率分别设置两类参数的学习率。在 deepspeed/runtime/engine.py 中引擎读取优化器配置后会把模型参数按use_muon拆成两个参数组Muon 参数组use_muonTrue名为muon-params接受lr、momentum、weight_decay、muon_lr、ns_method等键其中muon_lr会覆盖lrAdam 参数组use_muonFalse名为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五次牛顿-舒尔茨迭代系数a3.4445, b-4.7750, c2.0315用于最大化零点斜率、以及zeropower_via_gram_newtonschulz在小型方阵 Gram 矩阵R X X.T上迭代对宽矩阵显著降低 FLOP并使用 fp16 提高半精度数值精度。deepspeed/runtime/zero/muon/muon_optimizer.pyDeepSpeed 对MuonWithAuxAdam的封装子类。核心是在step()中对use_muonTrue的参数组只执行解耦权重衰减加动量更新p.mul_(1 - lr * weight_decay)与p.add_(p.grad.reshape(p.shape), alpha-lr)——因为此时参数是展平版本不适合直接正交化真正的正交化已前移到 ZeRO 的get_flat_partition阶段。而use_muonFalse的参数组则交给外部传入的 Adam 优化器如 FusedAdam或内部联机 Adam 更新处理若 FusedAdam 初始化失败且fallback_to_inlineTrue会回退到 Muon 自带的联机 Adam 更新。muon_update支持的两种牛顿-舒尔茨方法由NS_METHODS {standard, gram}限定其中ns_method的默认值是gram矩形矩阵上约 2 倍加速standard则是原始迭代。这些算子都套用了 DeepSpeed 的compiler.compile()装饰器可借助编译后端加速。ZeRO Stage 支持与reduce_scatter约束Muon 在 DeepSpeed 中的使用与 ZeRO 状态划分方式强相关官方文档明确说明了两条边界ZeRO Stage 1 和 Stage 2 支持 Muon前提是关闭 optimizer offload并搭配reduce_scatter使用。分区边界处理如果某个 Muon 矩阵跨越了 ZeRO 分区边界DeepSpeed 会在施加非线性 Muon 更新之前将完整归约后的梯度提供给持有该矩阵分区的每一个 rank。从 deepspeed/runtime/zero/stage_1_and_2.py 的源码结构看相关逻辑通过检查参数的use_muon属性以及分区数量len(partition_ids) 1来决定是否为跨分区的 Muon 矩阵保留完整的梯度分区信息。暂不支持组合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 做专门处理。官方在 “Whats 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_lr2D 隐藏权重Muon 参数组的学习率会覆盖通用lr。官方微调实验使用1e-4。adam_lrembeddings、layer norm、bias、lm_headAdam 参数组的学习率同样覆盖lr。官方实验使用2e-6。momentumMuon 的动量系数默认值 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_methodgram。ns_method可选gram默认矩形矩阵上更快或standard。zero_optimization.stageMuon 支持的组合为 Stage 1/2关闭 optimizer offload、使用reduce_scatter与 Stage 3Muon 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 微调 LLMgit 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-A3BMoE 模型总参数量 16B、激活参数量 3B并在代码生成MBPP / MBPP、通用知识MMLU与数学推理GSM8K三类基准上评测每个基准使用各自领域专属的训练集。训练配置要点如下模型Moonlight-16B-A3BMoE16B 总参数 / 3B 激活参数训练数据集MBPP/MBPP 使用sahil2801/CodeAlpaca-20kMMLU 使用cais/mmluauxiliary_train约 95k 条GSM8K 使用meta-math/MetaMathQAsample_rate0.1约 39.5k 条并行与精度ZeRO Stage 2、bf16、专家并行autoep_size4批大小batch size 16梯度累积 2共 4 块 GPU训练轮数1 epoch梯度裁剪 1.0评测结果对比OptimizerLearning Rateadam_lr (for Muon)MBPPMBPPMMLUGSM8Kbaseline (pre-finetune)——0.4950.4310.4010.526AdamW2e-6—0.6610.5340.6600.805Muon1e-42e-60.6460.5480.6780.810在这 4 项指标中Muon 有 3 项超过 AdamWMBPP0.548 vs 0.5341.4 个百分点、MMLU0.678 vs 0.6601.8 个百分点和 GSM8K0.810 vs 0.8050.5 个百分点。在 MBPP 基础测试集上 AdamW 略胜0.661 vs 0.646-1.5 个百分点但在包含更多测试用例、更为严格的 MBPP 上 Muon 得分更高0.548 vs 0.534说明 Muon 的泛化能力更好。读者在使用前也应注意到这是单一模型、特定领域微调数据上的观察结果不同模型与任务上的收益需要自行验证。Muon 的显存节省实测由于 Muon 每个参数只维护一份动量缓冲区而非 Adam 的两份其优化器状态显存低于 Adam。但 Muon 是混合优化器2D 隐藏权重用 Muon1 份缓冲区其余参数embeddings、layer norm、lm_head 等仍用 Adam2 份缓冲区因此实际节省取决于 2D 隐藏权重占参数总量的比例。每参数状态内存对比OptimizerState Buffers per ParamMemory per ParameterAdam2 (m, v)8 bytesMuon1 (momentum)4 bytes对于典型 Transformer约 90% 的参数是 2D 隐藏权重因此优化器状态内存大约能减少 45%。不过由于 GPU 总显存还包含模型权重、梯度和激活值端到端显存降幅会更小见下方实测。Qwen2.5-3B 微调实测峰值显存官方在 8 块 A10040GB上、用与上文一致的配置batch size 32、ZeRO Stage 2、bf16微调 Qwen2.5-3B数据集 tatsu-lab/alpaca并测量每 GPU 峰值显存OptimizerPeak Memory per GPUSavings vs AdamWAdamW34.5 GiB—Muon31.4 GiB9%Muon 每 GPU 显存比 AdamW 低约 3 GiB约 9%。节省完全来自优化器状态Muon 参数存一份动量缓冲区4 字节而 Adam 存两份8 字节。因为优化器状态只是总显存的一部分其余为模型权重、梯度与激活端到端降幅相对温和。对更大模型或显存更紧张的场景这 9% 的节省可能恰好决定工作负载能否放进设备显存、还是需要触发 CPU offload。演进路线Whats Next随着社区采纳度快速提升以及 Kimi-K21T 参数、GLM-5744B 参数等生产级应用的出现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.pyDeepSpeed 封装与 Adam 回退deepspeed/runtime/zero/muon/muon_optimizer.py参数打标逻辑deepspeed/init.py优化器名称与 AdamW 默认值deepspeed/runtime/config.py双参数组构造与muon_lr/adam_lr覆盖规则deepspeed/runtime/engine.pyZeRO Stage 1/2 中use_muon分区更新处理deepspeed/runtime/zero/stage_1_and_2.pyZeRO Stage 3 中 Muon 分支deepspeed/runtime/zero/stage3.pyMuon 单元测试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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考