做大规模模型训练这几年被问得最多的一个问题就是“MoE 架构是不是要把全部参数都塞进显存”这个问题的背后其实是 DeepSpeed ZeRO-3 和 MoE 训练两条知识线没有打通。先把结论放在前头训练时并不是每张卡都要放下全部参数参数可以分片存在多张 GPU 上但是 MoE 的路由机制、负载均衡和 all-to-all 通信又会给 ZeRO-3 这套分片体系添不少麻烦。这篇文章按我自己做实验的路径把 ZeRO-3 怎么拆显存、MoE 为什么难训、DeepSpeed-MoE 怎么把两者揉在一起以及我踩过的坑一次讲透。1. 为什么现在要谈 ZeRO-3 和 MoE一张卡的放不下的那些参数1.1 从稠密模型到稀疏模型显存压力为什么会“质变”传统 Transformer 是稠密模型每一层前馈网络FFN的所有参数对每个 token 都会计算一遍。以 7B 模型为例FP16 精度下权重就要占 14GBAdam 优化器状态按每参数 16 字节算约 28GB梯度再加 14GB光这三项就是 56GB激活值和通信缓冲区还没算。这在单张 80GB 的 A100 上确实是“勉强能训”但已经没什么余量了。切成 MoE 之后总参数量往往是成倍甚至数十倍上涨。Switch Transformer 的 126B 模型权重就有 252GB优化器状态更多到离谱。这时候还想着“一张卡放下全部参数”根本不用讨论可行性。很多人被“稀疏模型”这个说法误导以为 MoE 参数量大但计算量小所以显存压力应该不大。这个理解在推理阶段对了一半在训练阶段是错的错得很离谱。训练时稀疏模型并没有真正“只用一小部分参数”它只是用一小部分参数做前向和反向的计算但梯度更新、优化器状态、以及路由统计都需要把整个专家系统的状态管理起来。所以 MoE 训练的第一道坎不是“显存够不够”而是“如何把巨大的参数量拆到多张卡上同时还能正确完成一轮完整的梯度同步”。ZeRO-3 解决的就是这个“拆”的问题。1.2 MoE 训练的“悖论”推理很稀疏训练很稠密我用一个生活化的类比来说。食堂窗口很多吃饭的人来了每人只去两三个窗口打菜这是 MoE 的推理但后厨每天备菜要把所有窗口的食材都准备好、清点完、补齐货这是 MoE 的训练。你不可能因为每个人只吃两个菜就让后厨只备两个窗口的菜。具体到训练过程前向计算时每个 token 经过门控网络Gate选 top-k 个专家只有这些专家参与计算结果但反向传播时梯度要回传到“被选中的专家”的参数上并更新。看起来确实只有部分专家被更新问题在于下一批数据可能选中完全不同的专家。整个数据分布上没有一个专家是永远不被选中的加上路由本身不稳定所以每个训练 step 结束后所有专家参数都需要等待梯度更新哪怕这个 step 只激活了其中几个。更根本的问题是优化器状态。Adam 需要保存每个参数的一阶动量 m 和二阶动量 v不管这个参数本轮有没有参与计算只要它在模型里就要维护这两份状态。这就是为什么说过“MoE 总参数再多优化器状态是按总参数量走的不是按激活参数走的”。理解了这个悖论后面看 ZeRO-3 的分片策略就会清楚它每一步到底在干什么。2. ZeRO-3 的核心原理把显存不足变成通信问题2.1 三步分片优化器状态、梯度、参数分别怎么拆DeepSpeed 的 ZeRO 分三个阶段很多人只看结论不细究过程结果配置参数的时候只能抄不会改。我先把三个阶段和背后的动机理一遍ZeRO 阶段分片内容单卡显存效果代价ZeRO-1优化器状态显存大头从 Adam 状态来这步最划算几乎无额外通信只是优化器状态分散在各卡ZeRO-2优化器状态 梯度梯度不再每卡全量保存反向传播时需要 reduce-scatter 梯度ZeRO-3优化器状态 梯度 模型参数理论上单卡显存随卡数线性下降前向和反向都要 all-gather 参数ZeRO-3 最核心的“拆参数”我多说几句。它做的不是张量并行那种按矩阵维度切块而是按层切分模型状态。比如一个 24 层的 Transformer8 张卡时每张卡只保存其中若干层的全部参数而不是每层都拆成 1/8。这样设计的好处是通信模式相对简单某一层计算之前把包含这层参数的权重拉到所有参与计算的卡上算完就释放。注意ZeRO-3 说的是“每一时刻只有少数层的参数驻留在显存”所以单卡峰值显存不是总参数除以卡数那么简单它取决于通信桶大小和当前活跃的计算层数。这也是为什么配置文件里有stage3_max_live_parameters这种参数来控制显存的上限拉低它能压显存拉高它能减少频繁加载带来的开销。2.2 “一进一出”的数据流前向、反向各在哪一步做通信理解 ZeRO-3 的通信盯住两个集体通信原语就够了all-gather 和 reduce-scatter。前向过程假设模型有 24 层第 1 层参数只存在 rank 0 上第 2 层只存在 rank 1 上。当计算第 1 层时rank 0 把自己保存的这层参数广播给所有其他 rank大家拿到的是完整的第 1 层权重算出这一层的激活值这一层算完所有 rank 立刻把这段参数从显存中释放。下一层同理。这就是 all-gather。反向过程反向计算某一层梯度时同样先把该层参数 all-gather 回来算完本层的梯度之后要把梯度 reduce-scatter 到保存该层参数的那个 rank 上。每个 rank 最后只保留自己参数分片对应的梯度然后用这些梯度更新自己那部分优化器状态。所以每一层前向都 pull 一次反向也 pull 一次反向结束还有一次梯度归并。这绝不是没代价的通信量通常和参数量成正比。这也是很多人跑 ZeRO-3 觉得比 ZeRO-2 慢的原因参数广播多了好几轮。DeepSpeed 的应对手段主要是overlap_comm让通信和计算重叠前向算第 n 层时后台提前 all-gather 第 n1 层参数。从效果看只要代码里没有明显的同步阻塞这个开关值得一直开着。3. MoE 训练为什么难路由、负载均衡和并行3.1 路由和 all-to-alltoken 是怎么去找专家的MoE 的难点不完全在“一体积大”而在于它的数据流不是单纯的“层内计算—层间传递”而是多了一步 token 和专家的配对。一个非 MoE 的 Transformertoken 在每个 GPU 上从头到尾算自己的大家之间只在梯度同步时通信但 MoE 层里每个 token 可能被路由到其他 GPU 上的专家所以数据必须“物理地”搬到对应专家的卡上。这个过程用的是all-to-all 通信每个数据并行 rank 上有一部分 token它们被 Gate 网络分到了不同的专家于是这些 token 需要被打包发给对应的 expert-parallel rank。换句话说数据不是跟着原来的 batch 走而是跟着专家走。反向传播再走一次 all-to-all 把 token 还原回原 rank。这个设计直接影响显存和通信。专家参数可以因为 ZeRO-3 分片而不用整卡保存但 all-to-all 需要额外的通信缓冲区尤其是当某个专家收到的 token 特别多时缓冲区会突然膨胀。我见过不少训练任务不是死在计算上而是死在 all-to-all 缓冲区的峰值内存上。3.2 负载均衡Aux Loss 和 z-loss 为什么必须有MoE 训练最经典的毛病是“赢者通吃”门控网络很快学会总是把 token 分配给少数几个专家其他专家变成摆设。一旦出现这种情况被频繁调用的专家梯度更新快能力越来越强门控更不愿意选别人形成恶性循环。最后模型容量虽然在那儿摆着但有效参数量可能只有 20%。解决办法是给门控网络加一个负载均衡损失大家常说的 Aux Loss。它的核心思想是统计每个专家实际分到的 token 比例 f_i 和门控给出平均路由概率 P_i让这两个分布尽量接近。当某个专家分到的 token 过多时f_i 大于 P_i损失就会变大梯度会推着门控把概率摊平。我贴一段常见公式的写法方便你直观感受L_aux alpha * num_experts * sum_i(f_i * P_i)其中 alpha 就是 DeepSpeed 配置里的moe_gate_loss_wt。还有一个 z-loss作用是防止门控 logits 的绝对值变大因为如果所有专家都被均匀选到但 logits 本身数值很大下一步会把部分专家压到趋近于零出现“隐性坍塌”。z-loss 就是直接对 gate 的 logits 做 L2 惩罚让路由保持在比较平滑的状态。这两个损失权重都不能一开始就调很大。我在实验里通常先把moe_gate_loss_wt设成 0.001观察训练两三百步后各专家接收 token 的直方图如果分布太歪再逐步加到 0.01z-loss 一般 0.001 起作用很小但能防止极端情况。4. DeepSpeed 里 ZeRO-3 和 MoE 的协同DeepSpeed-MoE 的实现4.1 两层并行data parallel × expert parallel有了前面的铺垫现在可以谈 DeepSpeed-MoE 的模型到底是怎么部署的。实际上在 DeepSpeed 的框架里一个 MoE Transformer 被拆成了两种并行模式非 MoE 层attention、embedding、layer norm 等走数据并行并且这些层的参数继续用 ZeRO-3 做分片MoE 层里的专家走 expert parallel专家们被平均分配到ep_size张卡上同一小组内的卡组成一个“专家并行组”门控网络Gate的权重很小通常每个 expert-parallel 组内复制一份避免为了一个 gate 参数做跨组通信。这里有一个容易混淆的点ZeRO-3 和 expert parallel 并不是二选一。专家参数既可以被 expert parallel 分布到不同卡上也可以在专家并行组内部再用 ZeRO-3 把每个专家继续分片。DeepSpeed 把这套组合称为ZeRO-3 MoE官方文档里也叫 DeepSpeed-MoE。它允许你让模型的训练规模扩展到几百 B 参数而单卡显存依然可控。另一个值得注意的设计是PR-MoEPyramid Residual MoE和Residual MoE这是 DeepSpeed 论文里提出的变体。它把一个 MoE 层的总专家数按比例分成 k 个 MoE 子层每组专家负责处理前一组的“残差”路由结果。这么做的核心动机是减少 all-to-all 的瓶颈不是让所有 token 都去全世界找专家而是让 token 先在本地算一层再让部分 token 去下一组。实际训练时我用这个结构发现负载均衡明显更好调因为每组专家的压力天然被切分了。4.2 一个能跑的配置文件与模型改动下面给一个完整的 DeepSpeed 配置文件我有意把 ZeRO-3 和 MoE 相关的参数都标注出来{ train_batch_size: 256, gradient_accumulation_steps: 4, optimizer: { type: Adam, params: { lr: 3e-4, betas: [0.9, 0.95], eps: 1e-8, weight_decay: 0.01 } }, zero_optimization: { stage: 3, overlap_comm: true, contiguous_gradients: true, reduce_bucket_size: 5e8, allgather_bucket_size: 5e8, stage3_max_live_parameters: 1e9, stage3_max_reuse_distance: 1e9, stage3_gather_16bit_weights_on_model_save: true }, moe: { enabled: true, ep_size: 8, moe_expert_count: 16, moe_frequency: 1, moe_gate_loss_wt: 0.001, moe_gate_loss_wt_max: 0.01, moe_z_loss_wt: 0.001 }, communication_data_type: fp16, gradient_clipping: 1.0, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 } }配置参数的含义我得专门解释几个因为抄配置容易改配置难。moe_frequency表示每隔几层插入一个 MoE 层为 1 代表每一层 FFN 都替换成 MoE为 2 则隔一层替换一次。ep_size是每个专家组的卡数它不一定等于总卡数比如 64 卡时可以把 MoE 层分成 8 个 validators每个验证器 8 卡这样同一层专家分布在 8 卡上而数据并行度是 64/88。moe_gate_loss_wt_max是这个损失的封顶上限训练轮数越多权重会在 0.001 到 0.01 之间线性递增防止后期门控为了减少负载均衡损失而牺牲模型本身的效果。模型侧的改动主要是把 Transformer 的 FFN 层替换成 DeepSpeed 的 MoE 实现。我给一个简化版的片段import torch import torch.nn as nn import deepspeed from deepspeed.moe.layer import MoE class FeedForward(nn.Module): def __init__(self, d_model, d_ff): super().__init__() self.w1 nn.Linear(d_model, d_ff) self.w2 nn.Linear(d_ff, d_model) def forward(self, x): return self.w2(torch.relu(self.w1(x))) class TransformerLayerWithMoE(nn.Module): def __init__(self, d_model, d_ff, expert_count, ep_size): super().__init__() self.attention nn.MultiheadAttention(d_model, 8) self.moe_ffn MoE( hidden_sized_model, expertFeedForward(d_model, d_ff), num_expertsexpert_count, ep_sizeep_size, use_residualFalse, fuse_gateTrue, top_k1, min_capacity4 ) def forward(self, x): x x self.attention(x, x, x)[0] x x self.moe_ffn(x) return x这里MoE包装器是 DeepSpeed 提供给用户的不需要自己实现门控和 all-to-allep_size会自动决定它内部的通信方式。调用训练入口时用model, optimizer, _, lr_scheduler deepspeed.initialize(configconfig_path, modelmodel)就把分片和 MoE 全接上了。注意MoE 层不参与 ZeRO-3 的参数分片时必须在模型初始化之前把dist.init_process_group完成否则ep_size的分组信息拿不到。4.3 显存估算一张表给个参考我经常被问“这配置到底吃多少显存”这个问题没法给单一数字因为和隐藏维度、专家数、序列长度、batch size 都有关系。但可以给一个按经验推算的参考表前提是 16 卡 A100 80G、隐藏维度 2048、FFN 维度 8192、expert_count16非 MoE 层用 ZeRO-3 全分片模型规模参考专家并行数 ep_size单卡理论峰值显存能否训练稠密 7B 转 MoE总参数 24B4接近 60GB可以但通信压力大稠密 13B 转 MoE总参数 45B855GB 左右可以缓冲区和激活需要调总参数 100B 以上1660GB 以上必须开 CPU offload 或加大卡数这张表不是精确值但能给你一个直觉MoE 总参数越大只靠 ZeRO-3 分片并不够还要把reduce_bucket_size和allgather_bucket_size调小给 all-to-all 缓冲区留空间。我曾为了省事把 bucket 调到 1e9结果 MoE 层的 all-to-all 直接超限。如果想省显存比较靠谱的组合是“ZeRO-3 param offload MoE 层稳稳卡在显存里”把非 MoE 层的参数送去 CPU给专家留出空间。5. 实战中我认为最值得注意的地方通信、负载均衡和保存加载5.1 all-gather 和 all-to-all 都吃通信怎么配置ZeRO-3 的通信和 MoE 的通信在同一个训练 step 里叠加。我不是第一次看到有人只注意 ZeRO-3 的 all-gather 而忘了 MoE 的 all-to-all结果 GPU 利用率飙到 80% 但训练吞吐只有几 TFLOPs因为大家都在等网络而不是在算。这几条配置是我实测有效的overlap_comm: true让参数预取和当前计算重叠。如果不开前向反向每一步都卡在通信等待上allgather_bucket_size和reduce_bucket_size设置成接近单卡网络带宽能容纳的大小比如 8 卡时设 5e832 卡以上缩到 2e8防止单个 all-gather 调用长时间占住网卡对 MoE 的 all-to-all 做微调DeepSpeed 在MoE层有use_rtsRandom Token Selection参数它可以限制每个专家的 capacity超出的 token 会被随机丢弃等价于给通信缓冲区加了一个保险丝。容量设成min_capacity4起步后面看 token 分布再调大。通信瓶颈的判断方法有个土办法nvidia-smi 看 GPU 利用率如果接近 100% 但 dmesg 里全是 NCCL warning或者训练日志每个 step 时间突然陡增多半是 all-to-all 被某个热点专家拖住了这时首要任务不是加卡而是把负载均衡损失权重拉上去。5.2 负载均衡的“玄学”loss 权重怎么调Aux Loss 的权重是这个领域最“玄”的超参数之一。权重太小个别专家饿死训练后期想救都难权重太大强行把 token 平均分给所有专家每个专家的数据分布被人为打散模型效果和稠密模型差距拉大。我的经验是先跑一次小规模基线把每个专家平均接收 token 的数量打出来看分布。分布方差在 20% 之内时可以不动某个专家超过平均 token 数一倍以上时把moe_gate_loss_wt从 0.001 往上加每 500 步看一次分布如果训练中后期出现突然不均衡先检查是不是学习率太大导致门控 logits 失控加 z-loss 比继续加 aux loss 更稳。还要注意一个监控技巧DeepSpeed 会在日志里打印 MoE 相关的 loss 分项比如moe_gate_loss和moe_z_loss。这两个值如果出现震荡说明门控在反复横跳可以尝试将学习率下降或调高top_k值比如从 1 变成 2会显著平滑路由分布。5.3 checkpoint 保存ZeRO-3 的当前权重怎么存训练任务跑到一半保存权重时很多人会踩一个大坑ZeRO-3 的每个 rank 只保存了自己分片的那部分参数单独存出来的文件根本不是一个完整权重文件。要保存出能直接被模型加载的完整权重必须把各 rank 的 16bit 权重都 gather 出来。配置文件里stage3_gather_16bit_weights_on_model_save设为 true可以在保存过程中自动做一次全局参数收集代价是保存时会有几秒到几十秒的通信尖峰卡越多越明显。另一个更稳妥的办法是用 DeepSpeed 自带的工具脚本zero_to_fp32.py。训练时按原样保存 ZeRO-3 的 ckpt保存后执行一次python zero_to_fp32.py . pt_global_step它会读取所有分片并还原成一个 FP32 完整权重。我自己总是两个方案同时用训练中每个 ckpt 目录保留分片文件最终导出统一 FP32 给下游。这样即使某个 step 保存的 checkpoint 没设gather_16bit_weights也不会出现“模型缺一半”的情况。6. 常见问题速查我在 MoE ZeRO-3 训练中踩过的坑6.1 高频问题对照表这里把我在不同项目里遇到的典型问题整理成速查表每一行都是真实踩过或者看别人踩过、最后找到原因的症状大概率原因解决办法训练几百步后某几个专家 token 数接近 0Aux Loss 权重太小或不平衡调大moe_gate_loss_wt同时增大top_kall-to-all 通信阶段剧烈卡顿某个专家 capacity 超限缓冲区反馈阻塞打开use_rts调小min_capacity前向时报 all_gather 超时allgather_bucket_size太大单次通信占用过久从 5e8 降到 2e8 或 1e8模型能存但加载后精度明显变差保存时没做 16bit 全量 gather开启stage3_gather_16bit_weights_on_model_save或zero_to_fp32.py导一次启动训练直接报显存不足激活值和通信缓冲区没预留空间开激活 checkpointing调小 bucket 大小CPU offload 后训练像死机一样offload 连续读写频率太高关掉参数 offload只 offload 优化器状态或换 NVMe 路径6.2 我的经验先在小规模上跑通负载均衡曲线如果你准备把一套成熟的稠密模型改成 MoE我强烈建议不要一步到位上几百 B 超大模型。先用一个 1B 左右甚至几百 M 的模型配 8 卡 A100把 MoE 从 8 个专家、4 个专家并行用到 32 个专家、8 个专家并行。重点不是看吞吐而是观察三个指标每个专家接收 token 的方差、MoE 层的 all-to-all 耗时、以及 gate loss 的收敛曲线。这些指标在小规模上跑一块基本能预示大规模训练时的问题。我在一个 13B 稠密模型改 MoE 的项目里一开始直接上了 64 卡 128 专家结果前 500 步训练 loss 一直在波动。后来退回单机 8 卡复现同样波动的训练 loss通过打印每个专家的 token 分布才发现是因为 embedding 输出维度太大导致门控输入不稳定和显存容量一点关系都没有。小规模复现问题比大规模盲目调参省很多时间。最后再分享一个我自己坚持了很久的习惯每个 MoE 实验开始前把deepspeed.moe.layer.MoE里的顶层is_moe_param标记打印出来搞清楚模型里哪些参数走 expert parallel、哪些走 ZeRO-3 分片。这个看似不起眼的动作能避免很多“明明是 MoE 配置没对齐却一直在调优化器参数”的无效操作。希望你不用走我那么多弯路。