1. 为什么大模型训练绕不开 ZeRO-3 和 MoE如果你最近在折腾大模型训练大概率会在各种技术群里看到两个词反复出现DeepSpeed ZeRO-3和MoE。前者解决的是“显存不够、模型塞不下”的问题后者解决的是“参数规模上去了但计算量不想同步爆炸”的问题。这两个东西单独拎出来都不算新但组合在一起基本就是当前百亿到千亿参数模型训练的主流姿势。我自己第一次接触 ZeRO-3 是在训练一个 13B 级别的模型时单卡 80G 显存连模型权重都放不下更别提优化器状态和梯度了。当时试过梯度累积、梯度检查点能省的地方都省了结果还是 OOM。后来切到 ZeRO-3配合 CPU Offload才勉强把训练跑起来。而 MoE 则是在另一个场景下进入视野的——当你想把模型参数量从 13B 拉到 100B 以上但预算只允许你维持和 13B 差不多的计算量时MoE 几乎是唯一解。这篇文章面向的是已经有一定 PyTorch 训练基础、准备上手大模型训练或者正在被显存问题折磨的工程师。我会把 ZeRO-3 的分片逻辑、MoE 的路由机制、两者结合时的坑以及实际配置参数怎么调全部拆开讲清楚。不堆公式尽量用工程视角说人话。2. ZeRO-3 到底在分什么从显存账单说起2.1 训练显存的四大开销要理解 ZeRO-3先得知道训练时显存到底被谁吃了。以一个参数量为 Ψ 的模型为例混合精度训练fp16/bf16下显存开销主要来自四块模型权重fp16 下占 2Ψ 字节梯度fp16 下占 2Ψ 字节优化器状态如果用 Adamfp32 的动量加方差加主权重副本占 12Ψ 字节激活值跟 batch size、序列长度、模型深度强相关波动最大加起来一个 10B 参数的模型光权重梯度优化器状态就要 221216 字节/参数也就是 160GB。这还没算激活值。单卡 80G 根本不可能8 卡 80G 也就 640G勉强够但激活值一上来就崩。注意很多人只盯着模型权重算显存忽略了优化器状态才是大头。Adam 的 12Ψ 是 fp16 权重的 6 倍这是新手最容易误判的地方。2.2 ZeRO 三个阶段的本质区别ZeRO 的核心思想就一句话把原本每张卡都存一份的东西切成 N 份分散到 N 张卡上用的时候再通信拿回来。三个阶段切的东西不同阶段切分对象每卡显存近似通信量ZeRO-1优化器状态2Ψ 2Ψ 12Ψ/N低ZeRO-2优化器状态 梯度2Ψ (2Ψ12Ψ)/N中ZeRO-3优化器状态 梯度 权重(2Ψ2Ψ12Ψ)/N高ZeRO-1 只切优化器状态省得有限ZeRO-2 再切梯度省得更多ZeRO-3 连权重都切了所以又叫参数分片。这也是为什么 ZeRO-3 能让单卡装下远超显存容量的模型代价是前向和反向时每层都要 all-gather 权重通信量显著上升。我实测过一个 7B 模型8 卡 A100 80GZeRO-2 下每卡显存占用约 42G能跑 batch size 8ZeRO-3 下每卡显存占用约 18G能跑 batch size 16但每步耗时增加约 35%这个 trade-off 很典型ZeRO-3 用通信换显存。如果你的卡间带宽够NVLink/NVSwitch这个代价可以接受如果是普通 PCIe通信瓶颈会非常明显。2.3 ZeRO-3 的通信模式all-gather 与 reduce-scatterZeRO-3 在每一层的前向传播前需要把该层的权重从各卡 all-gather 回来算完再释放反向传播时同样要 all-gather 权重算梯度算完梯度用 reduce-scatter 分散回各卡。这个“用完即弃”的机制是 ZeRO-3 省显存的关键。但这里有个细节all-gather 是阻塞的。如果模型层数多、每层参数少通信次数会非常频繁延迟累积起来很吓人。DeepSpeed 为此做了通信重叠优化把下一层的 all-gather 和当前层的计算 overlap 起来但效果取决于模型结构和带宽。实操心得ZeRO-3 下如果发现 GPU 利用率上不去比如只有 30%大概率是通信瓶颈。这时候可以试试调大reduce_bucket_size和prefetch_bucket_size让每次通信搬更多数据减少通信次数。3. MoE 架构用稀疏激活撬动参数量3.1 MoE 的核心思想与路由机制MoEMixture of Experts的思路很朴素与其让所有参数都参与每次计算不如把参数分成多个“专家”每次只激活其中一小部分。这样总参数量可以做得很大但实际计算量只跟激活的专家数相关。一个典型的 MoE 层包含N 个专家网络每个专家通常是一个 FFN前馈网络一个门控网络Router/Gate决定每个 token 送给哪些专家组合逻辑把被选中专家的输出按门控权重加权求和以 Top-2 路由为例每个 token 会被送到得分最高的 2 个专家输出是这两个专家输出的加权和。这样总参数量是 N 倍但计算量只有 2 倍假设每个专家和原 FFN 同规模。3.2 稀疏激活带来的显存与计算解耦MoE 最迷人的地方在于参数量和计算量的解耦。传统稠密模型参数量翻倍计算量也翻倍MoE 可以把参数量翻 8 倍但计算量只增加 2 倍Top-2。这意味着你可以用相对固定的算力预算撑起一个参数量大得多的模型。但这里有个常见误解MoE 并不会减少显存占用。所有专家的参数都得存在显存里只是计算时不全用而已。所以“MoE 架构要全部参数进显存吗”这个问题的答案是要全部都要进显存。这也是为什么 MoE 几乎必须和 ZeRO-3 搭配使用——专家参数太多单卡根本放不下。3.3 负载均衡MoE 训练最头疼的问题MoE 训练最怕什么路由塌缩。也就是门控网络学歪了把所有 token 都送给同几个专家其他专家饿死。这时候模型退化成稠密模型MoE 的优势全没了而且被压榨的那几个专家梯度爆炸训练直接崩。解决负载均衡通常有两类手段辅助损失auxiliary loss在总 loss 里加一项惩罚专家负载不均容量因子capacity factor限制每个专家最多处理多少 token超出的丢弃或走残差DeepSpeed 的 MoE 实现里aux_loss和capacity_factor都是可配的。我一般会把aux_loss系数设在 0.01 左右capacity_factor设 1.25 到 1.5 之间。太小了均衡不住太大了浪费算力。踩过的坑有一次aux_loss设成 0.001前 1000 步看着挺正常到 5000 步突然路由全挤到 2 个专家上loss 直接起飞。后来调到 0.01 才稳住。这个参数不能省。4. ZeRO-3 与 MoE 结合11 大于 2 还是互相拖累4.1 为什么 MoE 几乎必须配 ZeRO-3前面说了MoE 的专家参数全都要进显存。假设你有 64 个专家每个专家是一个 100M 参数的 FFN光专家参数就 6.4B。再加上注意力层、嵌入层总参数量轻松上 10B。如果不用 ZeRO-3单卡根本放不下。而且 MoE 的专家是稀疏激活的这意味着如果用数据并行每张卡上的专家参数大部分时间都在“闲置”但显存照样占着。ZeRO-3 的参数分片正好解决这个问题每个专家参数被切到所有卡上需要哪个专家就 all-gather 哪个专家的参数不用的时候不占额外显存。4.2 通信放大效应与 bucket 调优但 ZeRO-3 MoE 不是简单叠加通信量会放大。原因在于 MoE 的路由是动态的每个 token 去的专家不一样导致 all-gather 的模式不规整。稠密模型里每层的 all-gather 是固定的可以很好地 overlapMoE 里专家选择是运行时决定的通信模式更复杂。DeepSpeed 对此的优化是把专家参数按 bucket 分组尽量让同一 bucket 里的专家被一起 all-gather。这时候reduce_bucket_size和prefetch_bucket_size的调优就很关键。我一般会从默认值开始然后按 2 的幂次往上试{ zero_optimization: { stage: 3, reduce_bucket_size: 5e7, prefetch_bucket_size: 5e7, stage3_prefetch_bucket_size: 5e7, stage3_param_persistence_threshold: 1e5, stage3_max_live_parameters: 1e9, stage3_max_reuse_distance: 1e9 } }stage3_param_persistence_threshold这个参数值得说一下小于这个阈值的参数不会被分片而是每张卡都留一份。对于 MoE 里的小参数比如门控网络设一个合适的阈值可以避免频繁通信。4.3 专家并行与数据并行的混合策略当专家数量很多时单纯靠 ZeRO-3 分片可能还不够这时候会引入专家并行Expert Parallelism。思路是把不同专家放到不同卡上每张卡负责一部分专家。token 根据路由结果被发送到对应卡上计算算完再发回来。DeepSpeed 的 MoE 支持ep_size参数来指定专家并行度。比如 8 卡ep_size4表示每 4 张卡组成一个专家并行组组内分专家组间做数据并行。这个配置需要根据专家数量和卡数仔细算ep_size 总卡数 / 数据并行度 每个专家并行组负责的专家数 总专家数 / ep_size我一般会让每个专家并行组至少负责 8 个专家太少了解析开销大太多了显存又不够。5. 完整实操从零配置一个 ZeRO-3 MoE 训练5.1 环境准备与 DeepSpeed 安装避坑“安装 deepspeed 包一直报错”是热词里高频出现的问题。DeepSpeed 的安装确实容易翻车主要是因为它依赖 PyTorch 版本、CUDA 版本、编译器版本三者严格匹配。我的建议是不要直接 pip install deepspeed而是从源码装并且先装好匹配的 PyTorch# 先确认 PyTorch 和 CUDA 版本 python -c import torch; print(torch.__version__, torch.version.cuda) # 安装匹配的 deepspeed pip install deepspeed --no-build-isolation # 如果报错试试指定 CUDA 架构 DS_BUILD_OPS1 pip install deepspeed --no-build-isolation常见报错及对策报错信息原因解决nvcc not found没装 CUDA toolkit装对应版本 CUDA toolkitundefined symbolPyTorch 和 DeepSpeed 编译架构不匹配重装匹配版本OOM during build编译时内存不够加MAX_JOBS4限制并行编译fused_adam not found自定义算子没编译设DS_BUILD_OPS1重装实操心得如果你只是要用 ZeRO-3 和 MoE不需要编译所有自定义算子。可以只编译需要的DS_BUILD_CPU_ADAM1 DS_BUILD_FUSED_ADAM1 pip install deepspeed。这样编译快很多也不容易出错。5.2 模型定义MoE 层的实现要点用 DeepSpeed 的 MoE模型侧需要把普通的 FFN 替换成 MoE 层。DeepSpeed 提供了MoE类但实际用的时候我更喜欢自己写一个轻量封装方便控制路由逻辑import torch import torch.nn as nn from deepspeed.moe.layer import MoE class MoEFFN(nn.Module): def __init__(self, hidden_size, ffn_size, num_experts8, top_k2): super().__init__() self.moe MoE( hidden_sizehidden_size, expertnn.Sequential( nn.Linear(hidden_size, ffn_size), nn.GELU(), nn.Linear(ffn_size, hidden_size) ), num_expertsnum_experts, ep_size1, # 专家并行度根据卡数调整 ktop_k, capacity_factor1.25, eval_capacity_factor2.0, min_capacity4, use_residualFalse ) def forward(self, x): output, aux_loss, _ self.moe(x) return output, aux_loss这里几个参数需要重点解释k2Top-2 路由每个 token 走 2 个专家capacity_factor1.25每个专家容量是平均负载的 1.25 倍min_capacity4最小容量防止专家太少时容量为 0use_residualFalse是否加残差连接一般 MoE 层本身在 Transformer 里已经有残差了5.3 DeepSpeed 配置文件的参数计算配置文件是 ZeRO-3 MoE 的核心。下面这份是我在 8 卡 A100 80G 上训练 13B MoE 模型8 专家Top-2的配置可以直接参考{ train_batch_size: 64, train_micro_batch_size_per_gpu: 2, gradient_accumulation_steps: 4, bf16: { enabled: true }, zero_optimization: { stage: 3, offload_optimizer: { device: cpu, pin_memory: true }, offload_param: { device: cpu, pin_memory: true }, overlap_comm: true, contiguous_gradients: true, reduce_bucket_size: 5e7, stage3_prefetch_bucket_size: 5e7, stage3_param_persistence_threshold: 1e5, stage3_max_live_parameters: 1e9, stage3_max_reuse_distance: 1e9, stage3_gather_16bit_weights_on_model_save: true }, moe: { enabled: true, ep_size: 1, use_tutel: false }, gradient_clipping: 1.0, steps_per_print: 100, optimizer: { type: Adam, params: { lr: 1e-4, betas: [0.9, 0.95], eps: 1e-8, weight_decay: 0.1 } }, scheduler: { type: WarmupDecayLR, params: { warmup_min_lr: 0, warmup_max_lr: 1e-4, warmup_num_steps: 2000, total_num_steps: 100000 } } }几个关键参数的计算逻辑train_batch_size与train_micro_batch_size_per_gpu的关系train_batch_size micro_batch_size × gradient_accumulation_steps × data_parallel_size 64 2 × 4 × 8这里 data_parallel_size 是 8因为 ep_size1所有卡都做数据并行。reduce_bucket_size的选择 这个值决定了每次 all-gather/reduce-scatter 搬多少数据。太小了通信次数多太大了显存峰值高。经验公式是reduce_bucket_size ≈ 单卡显存 × 0.1 / 280G × 0.1 / 2 4G但实际不会设这么大一般 5e750M 元素到 1e8 之间比较稳。stage3_max_live_parameters 这个参数控制同时保留多少参数在显存里。设太小会导致频繁重新 all-gather设太大又浪费显存。一般设成模型总参数量的 1/10 左右。5.4 启动训练与日志观察启动脚本deepspeed --num_gpus8 train.py \ --deepspeed ds_config.json \ --model_name_or_path ./model \ --output_dir ./output \ --num_train_epochs 3 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 4 \ --learning_rate 1e-4 \ --bf16 True \ --logging_steps 10 \ --save_steps 1000启动后重点观察几个指标loss 曲线MoE 的 loss 前期波动比稠密模型大正常aux_loss辅助损失应该稳定在一个小值如果持续上升说明负载不均专家负载分布DeepSpeed 会打印每个专家的 token 数理想情况是均匀分布GPU 利用率ZeRO-3 下如果低于 50%考虑调 bucket size显存占用用nvidia-smi观察应该稳定不涨实操心得MoE 训练前 1000 步建议把aux_loss系数调大一点比如 0.05先把路由均衡住之后再降回 0.01。这就像新员工入职先培训再上岗一上来就放任自流容易跑偏。6. 常见问题与排查速查表6.1 训练不稳定与 loss 尖刺MoE ZeRO-3 训练最常见的现象就是 loss 突然尖刺。原因通常有三个路由塌缩前面说过aux_loss太小导致专家负载不均。排查方法是打印专家负载直方图如果某个专家吃了超过 50% 的 token基本就是这个问题。对策是调大aux_loss或者临时降低学习率。梯度爆炸MoE 的梯度比稠密模型更容易爆炸因为专家是稀疏激活的单个专家的梯度可能很大。对策是开梯度裁剪gradient_clipping设 1.0 是标配不稳定的话降到 0.5。bf16 精度问题bf16 的动态范围比 fp16 大但精度低。MoE 的门控网络对精度敏感有时候 bf16 下路由会抖。对策是门控网络用 fp32 计算或者整体切回 fp16但 fp16 需要 loss scaling。6.2 显存溢出OOM的定位思路ZeRO-3 下 OOM 的排查和普通训练不太一样因为显存是动态的。定位思路确认是哪个阶段 OOM前向、反向还是优化器更新日志里会有提示检查stage3_max_live_parameters设太大导致 all-gather 的参数太多检查prefetch_bucket_size预取太多也会爆显存检查激活值ZeRO-3 不管激活值激活值爆了要开 gradient checkpointing检查 CPU offload如果 offload 了还 OOM说明是激活值或临时 buffer 的问题我整理了一个速查表现象可能原因对策前向 OOMall-gather 参数太多降stage3_max_live_parameters反向 OOM梯度 buffer 太大降reduce_bucket_size优化器更新 OOM优化器状态没 offload开offload_optimizer随机 OOM激活值峰值开 gradient checkpointing保存时 OOM权重聚合开stage3_gather_16bit_weights_on_model_save6.3 通信瓶颈与吞吐优化如果 GPU 利用率低、训练慢大概率是通信瓶颈。优化手段按优先级第一开overlap_comm让通信和计算重叠这是最有效的。但注意它需要额外的显存OOM 的话要权衡。第二调大 bucket size减少通信次数。但 bucket 太大会增加显存峰值需要试。第三用 NVLink如果卡间是 PCIe通信带宽是硬伤换 NVLink 机器能提升几倍。第四减少 all-gather 频率stage3_param_persistence_threshold调大让小参数不分片。第五考虑 ZeRO-2如果显存够ZeRO-2 的通信量比 ZeRO-3 小很多。MoE 的专家参数如果能在 ZeRO-2 下放下优先用 ZeRO-2。6.4 检查点保存与恢复的坑ZeRO-3 下保存检查点是个大坑。因为权重是分片的保存时需要 all-gather 成完整权重这个过程可能 OOM。DeepSpeed 提供了stage3_gather_16bit_weights_on_model_save选项但即使开了保存大模型时也可能爆。我的做法是保存时用zero_to_fp32.py脚本离线转换避免训练时 OOM或者只保存分片权重恢复时再加载分片定期保存别等到训练快结束才存万一 OOM 就白跑了# 离线转换分片检查点为完整权重 python zero_to_fp32.py ./checkpoint ./output_model.bin踩过的坑有一次训练了 3 天保存时 OOM 导致检查点损坏只能从头再来。后来改成每 500 步保存一次分片权重虽然占空间但安全。7. 一些实战中的参数调优经验7.1 学习率与 warmup 的设置MoE 模型的学习率通常比稠密模型小因为专家多、参数多大学习率容易导致路由震荡。我的经验值稠密 13Blr 2e-4 到 3e-4MoE 13B8 专家lr 1e-4 到 1.5e-4MoE 100Blr 5e-5 到 1e-4warmup 步数也要相应增加一般占总步数的 2% 到 5%。MoE 的路由需要时间稳定warmup 太短容易早期就塌缩。7.2 专家数量与容量的权衡专家数量不是越多越好。专家越多参数量越大但路由的决策空间也越大训练难度上升。我的经验小模型 1B4 到 8 个专家中模型1B 到 10B8 到 16 个专家大模型 10B16 到 64 个专家容量因子capacity_factor一般设 1.25 到 1.5。太小了 token 被丢弃多训练信号损失太大了算力浪费。可以观察“丢弃率”指标如果超过 10% 就调大容量。7.3 监控指标与早停策略MoE 训练需要监控的指标比稠密模型多主 loss正常下降aux loss稳定小值不持续上升专家负载熵越高越均匀理想是接近 log(N)token 丢弃率低于 10%梯度范数稳定不爆炸如果 aux loss 持续上升、负载熵持续下降说明路由在塌缩需要干预。早停策略可以设成连续 500 步 aux loss 上升就降低学习率或增大 aux loss 系数。8. 我个人在实际操作中的体会ZeRO-3 和 MoE 这套组合说白了就是“用通信换显存用稀疏换参数”。听起来很美但实际调起来坑不少。我最大的体会是不要一上来就追求大参数量先把小规模跑通。我见过太多人直接上 100B MoE结果路由塌缩、通信瓶颈、OOM 全来了调了两周没进展。正确的路径是先用 1B 左右的稠密模型把 ZeRO-3 跑通理解分片和通信的行为然后换成 8 专家的 MoE在小规模上把路由和负载均衡调稳最后再放大到目标规模。每一步都验证 loss 曲线、专家负载、显存占用、GPU 利用率这四个指标。另外DeepSpeed 的文档虽然全但很多参数的实际效果跟版本强相关。我建议锁定一个版本别频繁升级。我目前用的是 0.14.xMoE 相关的 API 比较稳定。升级前一定要在小的 checkpoint 上验证别拿正式训练冒险。最后分享一个小技巧MoE 训练时把门控网络的学习率单独设大一点比如主学习率的 2 倍能让路由更快收敛。这个在 DeepSpeed 里可以通过参数组实现虽然官方文档没提但实测有效。
大模型训练实战：ZeRO-3与MoE结合的原理、配置与调优
1. 为什么大模型训练绕不开 ZeRO-3 和 MoE如果你最近在折腾大模型训练大概率会在各种技术群里看到两个词反复出现DeepSpeed ZeRO-3和MoE。前者解决的是“显存不够、模型塞不下”的问题后者解决的是“参数规模上去了但计算量不想同步爆炸”的问题。这两个东西单独拎出来都不算新但组合在一起基本就是当前百亿到千亿参数模型训练的主流姿势。我自己第一次接触 ZeRO-3 是在训练一个 13B 级别的模型时单卡 80G 显存连模型权重都放不下更别提优化器状态和梯度了。当时试过梯度累积、梯度检查点能省的地方都省了结果还是 OOM。后来切到 ZeRO-3配合 CPU Offload才勉强把训练跑起来。而 MoE 则是在另一个场景下进入视野的——当你想把模型参数量从 13B 拉到 100B 以上但预算只允许你维持和 13B 差不多的计算量时MoE 几乎是唯一解。这篇文章面向的是已经有一定 PyTorch 训练基础、准备上手大模型训练或者正在被显存问题折磨的工程师。我会把 ZeRO-3 的分片逻辑、MoE 的路由机制、两者结合时的坑以及实际配置参数怎么调全部拆开讲清楚。不堆公式尽量用工程视角说人话。2. ZeRO-3 到底在分什么从显存账单说起2.1 训练显存的四大开销要理解 ZeRO-3先得知道训练时显存到底被谁吃了。以一个参数量为 Ψ 的模型为例混合精度训练fp16/bf16下显存开销主要来自四块模型权重fp16 下占 2Ψ 字节梯度fp16 下占 2Ψ 字节优化器状态如果用 Adamfp32 的动量加方差加主权重副本占 12Ψ 字节激活值跟 batch size、序列长度、模型深度强相关波动最大加起来一个 10B 参数的模型光权重梯度优化器状态就要 221216 字节/参数也就是 160GB。这还没算激活值。单卡 80G 根本不可能8 卡 80G 也就 640G勉强够但激活值一上来就崩。注意很多人只盯着模型权重算显存忽略了优化器状态才是大头。Adam 的 12Ψ 是 fp16 权重的 6 倍这是新手最容易误判的地方。2.2 ZeRO 三个阶段的本质区别ZeRO 的核心思想就一句话把原本每张卡都存一份的东西切成 N 份分散到 N 张卡上用的时候再通信拿回来。三个阶段切的东西不同阶段切分对象每卡显存近似通信量ZeRO-1优化器状态2Ψ 2Ψ 12Ψ/N低ZeRO-2优化器状态 梯度2Ψ (2Ψ12Ψ)/N中ZeRO-3优化器状态 梯度 权重(2Ψ2Ψ12Ψ)/N高ZeRO-1 只切优化器状态省得有限ZeRO-2 再切梯度省得更多ZeRO-3 连权重都切了所以又叫参数分片。这也是为什么 ZeRO-3 能让单卡装下远超显存容量的模型代价是前向和反向时每层都要 all-gather 权重通信量显著上升。我实测过一个 7B 模型8 卡 A100 80GZeRO-2 下每卡显存占用约 42G能跑 batch size 8ZeRO-3 下每卡显存占用约 18G能跑 batch size 16但每步耗时增加约 35%这个 trade-off 很典型ZeRO-3 用通信换显存。如果你的卡间带宽够NVLink/NVSwitch这个代价可以接受如果是普通 PCIe通信瓶颈会非常明显。2.3 ZeRO-3 的通信模式all-gather 与 reduce-scatterZeRO-3 在每一层的前向传播前需要把该层的权重从各卡 all-gather 回来算完再释放反向传播时同样要 all-gather 权重算梯度算完梯度用 reduce-scatter 分散回各卡。这个“用完即弃”的机制是 ZeRO-3 省显存的关键。但这里有个细节all-gather 是阻塞的。如果模型层数多、每层参数少通信次数会非常频繁延迟累积起来很吓人。DeepSpeed 为此做了通信重叠优化把下一层的 all-gather 和当前层的计算 overlap 起来但效果取决于模型结构和带宽。实操心得ZeRO-3 下如果发现 GPU 利用率上不去比如只有 30%大概率是通信瓶颈。这时候可以试试调大reduce_bucket_size和prefetch_bucket_size让每次通信搬更多数据减少通信次数。3. MoE 架构用稀疏激活撬动参数量3.1 MoE 的核心思想与路由机制MoEMixture of Experts的思路很朴素与其让所有参数都参与每次计算不如把参数分成多个“专家”每次只激活其中一小部分。这样总参数量可以做得很大但实际计算量只跟激活的专家数相关。一个典型的 MoE 层包含N 个专家网络每个专家通常是一个 FFN前馈网络一个门控网络Router/Gate决定每个 token 送给哪些专家组合逻辑把被选中专家的输出按门控权重加权求和以 Top-2 路由为例每个 token 会被送到得分最高的 2 个专家输出是这两个专家输出的加权和。这样总参数量是 N 倍但计算量只有 2 倍假设每个专家和原 FFN 同规模。3.2 稀疏激活带来的显存与计算解耦MoE 最迷人的地方在于参数量和计算量的解耦。传统稠密模型参数量翻倍计算量也翻倍MoE 可以把参数量翻 8 倍但计算量只增加 2 倍Top-2。这意味着你可以用相对固定的算力预算撑起一个参数量大得多的模型。但这里有个常见误解MoE 并不会减少显存占用。所有专家的参数都得存在显存里只是计算时不全用而已。所以“MoE 架构要全部参数进显存吗”这个问题的答案是要全部都要进显存。这也是为什么 MoE 几乎必须和 ZeRO-3 搭配使用——专家参数太多单卡根本放不下。3.3 负载均衡MoE 训练最头疼的问题MoE 训练最怕什么路由塌缩。也就是门控网络学歪了把所有 token 都送给同几个专家其他专家饿死。这时候模型退化成稠密模型MoE 的优势全没了而且被压榨的那几个专家梯度爆炸训练直接崩。解决负载均衡通常有两类手段辅助损失auxiliary loss在总 loss 里加一项惩罚专家负载不均容量因子capacity factor限制每个专家最多处理多少 token超出的丢弃或走残差DeepSpeed 的 MoE 实现里aux_loss和capacity_factor都是可配的。我一般会把aux_loss系数设在 0.01 左右capacity_factor设 1.25 到 1.5 之间。太小了均衡不住太大了浪费算力。踩过的坑有一次aux_loss设成 0.001前 1000 步看着挺正常到 5000 步突然路由全挤到 2 个专家上loss 直接起飞。后来调到 0.01 才稳住。这个参数不能省。4. ZeRO-3 与 MoE 结合11 大于 2 还是互相拖累4.1 为什么 MoE 几乎必须配 ZeRO-3前面说了MoE 的专家参数全都要进显存。假设你有 64 个专家每个专家是一个 100M 参数的 FFN光专家参数就 6.4B。再加上注意力层、嵌入层总参数量轻松上 10B。如果不用 ZeRO-3单卡根本放不下。而且 MoE 的专家是稀疏激活的这意味着如果用数据并行每张卡上的专家参数大部分时间都在“闲置”但显存照样占着。ZeRO-3 的参数分片正好解决这个问题每个专家参数被切到所有卡上需要哪个专家就 all-gather 哪个专家的参数不用的时候不占额外显存。4.2 通信放大效应与 bucket 调优但 ZeRO-3 MoE 不是简单叠加通信量会放大。原因在于 MoE 的路由是动态的每个 token 去的专家不一样导致 all-gather 的模式不规整。稠密模型里每层的 all-gather 是固定的可以很好地 overlapMoE 里专家选择是运行时决定的通信模式更复杂。DeepSpeed 对此的优化是把专家参数按 bucket 分组尽量让同一 bucket 里的专家被一起 all-gather。这时候reduce_bucket_size和prefetch_bucket_size的调优就很关键。我一般会从默认值开始然后按 2 的幂次往上试{ zero_optimization: { stage: 3, reduce_bucket_size: 5e7, prefetch_bucket_size: 5e7, stage3_prefetch_bucket_size: 5e7, stage3_param_persistence_threshold: 1e5, stage3_max_live_parameters: 1e9, stage3_max_reuse_distance: 1e9 } }stage3_param_persistence_threshold这个参数值得说一下小于这个阈值的参数不会被分片而是每张卡都留一份。对于 MoE 里的小参数比如门控网络设一个合适的阈值可以避免频繁通信。4.3 专家并行与数据并行的混合策略当专家数量很多时单纯靠 ZeRO-3 分片可能还不够这时候会引入专家并行Expert Parallelism。思路是把不同专家放到不同卡上每张卡负责一部分专家。token 根据路由结果被发送到对应卡上计算算完再发回来。DeepSpeed 的 MoE 支持ep_size参数来指定专家并行度。比如 8 卡ep_size4表示每 4 张卡组成一个专家并行组组内分专家组间做数据并行。这个配置需要根据专家数量和卡数仔细算ep_size 总卡数 / 数据并行度 每个专家并行组负责的专家数 总专家数 / ep_size我一般会让每个专家并行组至少负责 8 个专家太少了解析开销大太多了显存又不够。5. 完整实操从零配置一个 ZeRO-3 MoE 训练5.1 环境准备与 DeepSpeed 安装避坑“安装 deepspeed 包一直报错”是热词里高频出现的问题。DeepSpeed 的安装确实容易翻车主要是因为它依赖 PyTorch 版本、CUDA 版本、编译器版本三者严格匹配。我的建议是不要直接 pip install deepspeed而是从源码装并且先装好匹配的 PyTorch# 先确认 PyTorch 和 CUDA 版本 python -c import torch; print(torch.__version__, torch.version.cuda) # 安装匹配的 deepspeed pip install deepspeed --no-build-isolation # 如果报错试试指定 CUDA 架构 DS_BUILD_OPS1 pip install deepspeed --no-build-isolation常见报错及对策报错信息原因解决nvcc not found没装 CUDA toolkit装对应版本 CUDA toolkitundefined symbolPyTorch 和 DeepSpeed 编译架构不匹配重装匹配版本OOM during build编译时内存不够加MAX_JOBS4限制并行编译fused_adam not found自定义算子没编译设DS_BUILD_OPS1重装实操心得如果你只是要用 ZeRO-3 和 MoE不需要编译所有自定义算子。可以只编译需要的DS_BUILD_CPU_ADAM1 DS_BUILD_FUSED_ADAM1 pip install deepspeed。这样编译快很多也不容易出错。5.2 模型定义MoE 层的实现要点用 DeepSpeed 的 MoE模型侧需要把普通的 FFN 替换成 MoE 层。DeepSpeed 提供了MoE类但实际用的时候我更喜欢自己写一个轻量封装方便控制路由逻辑import torch import torch.nn as nn from deepspeed.moe.layer import MoE class MoEFFN(nn.Module): def __init__(self, hidden_size, ffn_size, num_experts8, top_k2): super().__init__() self.moe MoE( hidden_sizehidden_size, expertnn.Sequential( nn.Linear(hidden_size, ffn_size), nn.GELU(), nn.Linear(ffn_size, hidden_size) ), num_expertsnum_experts, ep_size1, # 专家并行度根据卡数调整 ktop_k, capacity_factor1.25, eval_capacity_factor2.0, min_capacity4, use_residualFalse ) def forward(self, x): output, aux_loss, _ self.moe(x) return output, aux_loss这里几个参数需要重点解释k2Top-2 路由每个 token 走 2 个专家capacity_factor1.25每个专家容量是平均负载的 1.25 倍min_capacity4最小容量防止专家太少时容量为 0use_residualFalse是否加残差连接一般 MoE 层本身在 Transformer 里已经有残差了5.3 DeepSpeed 配置文件的参数计算配置文件是 ZeRO-3 MoE 的核心。下面这份是我在 8 卡 A100 80G 上训练 13B MoE 模型8 专家Top-2的配置可以直接参考{ train_batch_size: 64, train_micro_batch_size_per_gpu: 2, gradient_accumulation_steps: 4, bf16: { enabled: true }, zero_optimization: { stage: 3, offload_optimizer: { device: cpu, pin_memory: true }, offload_param: { device: cpu, pin_memory: true }, overlap_comm: true, contiguous_gradients: true, reduce_bucket_size: 5e7, stage3_prefetch_bucket_size: 5e7, stage3_param_persistence_threshold: 1e5, stage3_max_live_parameters: 1e9, stage3_max_reuse_distance: 1e9, stage3_gather_16bit_weights_on_model_save: true }, moe: { enabled: true, ep_size: 1, use_tutel: false }, gradient_clipping: 1.0, steps_per_print: 100, optimizer: { type: Adam, params: { lr: 1e-4, betas: [0.9, 0.95], eps: 1e-8, weight_decay: 0.1 } }, scheduler: { type: WarmupDecayLR, params: { warmup_min_lr: 0, warmup_max_lr: 1e-4, warmup_num_steps: 2000, total_num_steps: 100000 } } }几个关键参数的计算逻辑train_batch_size与train_micro_batch_size_per_gpu的关系train_batch_size micro_batch_size × gradient_accumulation_steps × data_parallel_size 64 2 × 4 × 8这里 data_parallel_size 是 8因为 ep_size1所有卡都做数据并行。reduce_bucket_size的选择 这个值决定了每次 all-gather/reduce-scatter 搬多少数据。太小了通信次数多太大了显存峰值高。经验公式是reduce_bucket_size ≈ 单卡显存 × 0.1 / 280G × 0.1 / 2 4G但实际不会设这么大一般 5e750M 元素到 1e8 之间比较稳。stage3_max_live_parameters 这个参数控制同时保留多少参数在显存里。设太小会导致频繁重新 all-gather设太大又浪费显存。一般设成模型总参数量的 1/10 左右。5.4 启动训练与日志观察启动脚本deepspeed --num_gpus8 train.py \ --deepspeed ds_config.json \ --model_name_or_path ./model \ --output_dir ./output \ --num_train_epochs 3 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 4 \ --learning_rate 1e-4 \ --bf16 True \ --logging_steps 10 \ --save_steps 1000启动后重点观察几个指标loss 曲线MoE 的 loss 前期波动比稠密模型大正常aux_loss辅助损失应该稳定在一个小值如果持续上升说明负载不均专家负载分布DeepSpeed 会打印每个专家的 token 数理想情况是均匀分布GPU 利用率ZeRO-3 下如果低于 50%考虑调 bucket size显存占用用nvidia-smi观察应该稳定不涨实操心得MoE 训练前 1000 步建议把aux_loss系数调大一点比如 0.05先把路由均衡住之后再降回 0.01。这就像新员工入职先培训再上岗一上来就放任自流容易跑偏。6. 常见问题与排查速查表6.1 训练不稳定与 loss 尖刺MoE ZeRO-3 训练最常见的现象就是 loss 突然尖刺。原因通常有三个路由塌缩前面说过aux_loss太小导致专家负载不均。排查方法是打印专家负载直方图如果某个专家吃了超过 50% 的 token基本就是这个问题。对策是调大aux_loss或者临时降低学习率。梯度爆炸MoE 的梯度比稠密模型更容易爆炸因为专家是稀疏激活的单个专家的梯度可能很大。对策是开梯度裁剪gradient_clipping设 1.0 是标配不稳定的话降到 0.5。bf16 精度问题bf16 的动态范围比 fp16 大但精度低。MoE 的门控网络对精度敏感有时候 bf16 下路由会抖。对策是门控网络用 fp32 计算或者整体切回 fp16但 fp16 需要 loss scaling。6.2 显存溢出OOM的定位思路ZeRO-3 下 OOM 的排查和普通训练不太一样因为显存是动态的。定位思路确认是哪个阶段 OOM前向、反向还是优化器更新日志里会有提示检查stage3_max_live_parameters设太大导致 all-gather 的参数太多检查prefetch_bucket_size预取太多也会爆显存检查激活值ZeRO-3 不管激活值激活值爆了要开 gradient checkpointing检查 CPU offload如果 offload 了还 OOM说明是激活值或临时 buffer 的问题我整理了一个速查表现象可能原因对策前向 OOMall-gather 参数太多降stage3_max_live_parameters反向 OOM梯度 buffer 太大降reduce_bucket_size优化器更新 OOM优化器状态没 offload开offload_optimizer随机 OOM激活值峰值开 gradient checkpointing保存时 OOM权重聚合开stage3_gather_16bit_weights_on_model_save6.3 通信瓶颈与吞吐优化如果 GPU 利用率低、训练慢大概率是通信瓶颈。优化手段按优先级第一开overlap_comm让通信和计算重叠这是最有效的。但注意它需要额外的显存OOM 的话要权衡。第二调大 bucket size减少通信次数。但 bucket 太大会增加显存峰值需要试。第三用 NVLink如果卡间是 PCIe通信带宽是硬伤换 NVLink 机器能提升几倍。第四减少 all-gather 频率stage3_param_persistence_threshold调大让小参数不分片。第五考虑 ZeRO-2如果显存够ZeRO-2 的通信量比 ZeRO-3 小很多。MoE 的专家参数如果能在 ZeRO-2 下放下优先用 ZeRO-2。6.4 检查点保存与恢复的坑ZeRO-3 下保存检查点是个大坑。因为权重是分片的保存时需要 all-gather 成完整权重这个过程可能 OOM。DeepSpeed 提供了stage3_gather_16bit_weights_on_model_save选项但即使开了保存大模型时也可能爆。我的做法是保存时用zero_to_fp32.py脚本离线转换避免训练时 OOM或者只保存分片权重恢复时再加载分片定期保存别等到训练快结束才存万一 OOM 就白跑了# 离线转换分片检查点为完整权重 python zero_to_fp32.py ./checkpoint ./output_model.bin踩过的坑有一次训练了 3 天保存时 OOM 导致检查点损坏只能从头再来。后来改成每 500 步保存一次分片权重虽然占空间但安全。7. 一些实战中的参数调优经验7.1 学习率与 warmup 的设置MoE 模型的学习率通常比稠密模型小因为专家多、参数多大学习率容易导致路由震荡。我的经验值稠密 13Blr 2e-4 到 3e-4MoE 13B8 专家lr 1e-4 到 1.5e-4MoE 100Blr 5e-5 到 1e-4warmup 步数也要相应增加一般占总步数的 2% 到 5%。MoE 的路由需要时间稳定warmup 太短容易早期就塌缩。7.2 专家数量与容量的权衡专家数量不是越多越好。专家越多参数量越大但路由的决策空间也越大训练难度上升。我的经验小模型 1B4 到 8 个专家中模型1B 到 10B8 到 16 个专家大模型 10B16 到 64 个专家容量因子capacity_factor一般设 1.25 到 1.5。太小了 token 被丢弃多训练信号损失太大了算力浪费。可以观察“丢弃率”指标如果超过 10% 就调大容量。7.3 监控指标与早停策略MoE 训练需要监控的指标比稠密模型多主 loss正常下降aux loss稳定小值不持续上升专家负载熵越高越均匀理想是接近 log(N)token 丢弃率低于 10%梯度范数稳定不爆炸如果 aux loss 持续上升、负载熵持续下降说明路由在塌缩需要干预。早停策略可以设成连续 500 步 aux loss 上升就降低学习率或增大 aux loss 系数。8. 我个人在实际操作中的体会ZeRO-3 和 MoE 这套组合说白了就是“用通信换显存用稀疏换参数”。听起来很美但实际调起来坑不少。我最大的体会是不要一上来就追求大参数量先把小规模跑通。我见过太多人直接上 100B MoE结果路由塌缩、通信瓶颈、OOM 全来了调了两周没进展。正确的路径是先用 1B 左右的稠密模型把 ZeRO-3 跑通理解分片和通信的行为然后换成 8 专家的 MoE在小规模上把路由和负载均衡调稳最后再放大到目标规模。每一步都验证 loss 曲线、专家负载、显存占用、GPU 利用率这四个指标。另外DeepSpeed 的文档虽然全但很多参数的实际效果跟版本强相关。我建议锁定一个版本别频繁升级。我目前用的是 0.14.xMoE 相关的 API 比较稳定。升级前一定要在小的 checkpoint 上验证别拿正式训练冒险。最后分享一个小技巧MoE 训练时把门控网络的学习率单独设大一点比如主学习率的 2 倍能让路由更快收敛。这个在 DeepSpeed 里可以通过参数组实现虽然官方文档没提但实测有效。
