MoE大模型训练显存优化：DeepSpeed ZeRO-3与专家并行深度拆解

先说一个反直觉的结论当一个模型从稠密结构换成 MoE 结构时总参数量可能翻好几倍但如果你用对了并行策略单张 GPU 上的显存压力不仅不会等比上升反而可能因为专家参数被切得更碎而显著下降。这也是 MoE 能在不改变单卡存储上限的前提下把模型总规模推向万亿级别的核心原因。我经常看到初学大模型训练的朋友问一个问题MoE 架构要全部参数进显存吗这个问题问得不算错但答案高度取决于部署方式。如果单卡推理那没办法全部专家都得放进一张卡可如果是多卡训练或推理那完全不需要。而要把这个不需要变成现实靠的就是 DeepSpeed 的分布式训练能力尤其是 ZeRO-3 和专家并行Expert Parallelism的组合。这篇文章就把这套机制从头到尾拆开讲清楚包括 ZeRO-3 到底切了什么、MoE 的专家参数怎么分布、负载均衡损失为什么是隐形命门以及我实测中踩过的几个坑。1. MoE 的大和显存的大不是一回事1.1 单卡部署和多卡训练答案完全不同先说清楚最核心的认知偏差。MoEMixture of Experts专家混合模型的总参数量通常由两部分组成一部分是共享的稠密层比如 Embedding、Attention、LayerNorm 这些另一部分是大量的专家 FFN 层。假设你有 64 个专家每个专家 1.5 亿参数光专家部分就是 96 亿参数再加上共享层总参数量轻松上几十亿甚至上千亿。这时候你问MoE 架构要全部参数进显存吗在单卡推理场景下答案是要。因为路由器可能把任意 token 路由到任意专家你必须保证每个专家都在本卡可访问否则推理就断了。总参数有多大单卡显存就得有多大这也是 MoE 模型上线时通常也要多卡切分的原因。在多卡训练或多卡推理场景下答案是不要。通过专家并行EP64 个专家可以分布到 8 张卡上每张卡只持有 8 个专家通过 ZeRO-3共享层和优化器状态还能进一步切分到所有卡上。此时单卡显存压力取决于切分粒度而不是总参数量。这个区别一定要先立住后面讲的一切都建立在这个前提上。1.2 区分计算量和存储量两个指标另一个容易混淆的点是MoE 的参数量大不代表计算量大。在稠密 Transformer 里每个 token 要经过全部 FFN 层参数量和计算量成正比。但在 MoE 里路由器只为每个 token 挑选 top-K 个专家常见 K1 或 2也就是说每个 token 实际上只经过 1 个或 2 个专家的 FFN。如果总共有 64 个专家每个 token 只计算其中 2 个那计算量大约只相当于 2 个 FFN而不是 64 个。剩下的 62 个专家存在显存里但没有参与这次前向计算。这就引出一个非常实用的结论优化 MoE 训练存储量靠并行策略解决计算量靠 top-K 路由解决两个问题是独立的不能用同一个公式估算。我在给团队做显存规划时从来不看总参数 × 2 字节这种粗算法必须拆开参数分布来算。1.3 纯数据并行为什么是灾难这里要特意批评一种常见误配置拿到一个 MoE 模型不假思索地开 PyTorch DDP 或 DeepSpeed 纯数据并行。如果不开专家并行也不开 ZeRO-3纯数据并行意味着每个 GPU 上要复制一份完整的模型副本包括全部专家参数。假设你一共有 64 个专家每个 1.5 亿参数那每个 GPU 都得存下这 96 亿参数。32 张卡就是 32 份副本这不是分担显存而是把显存需求和总参数量绑死。更尴尬的是纯数据并行下 MoE 的专家路由是局部的每张卡只会在自己的本地专家里做选择除非你做全局负载均衡否则模型训练效果会明显劣化专家利用率也不均匀。所以结论很直接MoE 训练必须引入参数不复制的并行手段要么是专家并行EP要么是 ZeRO-3。这也是这篇文章标题把 DeepSpeed、ZeRO-3、MoE 三者放在一起的原因——它们本来就是配套出现的。2. ZeRO-3 的切分逻辑从 16 字节/参数说起2.1 一份参数在训练中到底占多少字节要理解 ZeRO-3先算一笔账。假设我们用混合精度训练fp16 存储 forward/backward 的模型参数fp32 维护优化器状态单个模型参数在训练过程中会对应这几种数据fp16 模型参数2 字节fp16 梯度2 字节fp32 优化器状态主权重副本 4 字节 动量momentum4 字节 方差variance4 字节共 12 字节总计每个参数约 16 字节一个 70B 稠密模型训练时的参数相关显存就是 70B × 16 字节 ≈ 1.12 TB。这块开销和 batch size、序列长度无关纯粹是模型规模和优化方式决定的。你一张 A100 80GB 就算掏空也塞不下 1.12TB更别提还有激活值和中间变量。ZeRO 的思路不是压缩这 16 字节而是把它们拆开让每一张卡只负责其中一小部分需要的时候再拼起来。这就是零冗余优化器Zero Redundancy Optimizer这个名字的由来——它消除的是不同并行副本之间的数据冗余。2.2 ZeRO-1/2/3 各自切掉了哪一块ZeRO 系列分三个阶段每个阶段切的东西不同ZeRO 阶段每卡仍持有的完整副本每卡被切片的对象70B 模型在 64 卡上的每卡存储量无 ZeRO纯 DP参数梯度优化器无约 1.12 TBZeRO-1参数梯度优化器状态约 293 GBZeRO-2参数优化器状态梯度约 155 GBZeRO-3无参数梯度优化器状态约 17.5 GB这个表我建议保存下来。做显存规划的时候先按这个表算出参数相关存储的下限再往上叠加激活值、通信缓冲区和框架开销就不会出现训练跑一半 OOM 的尴尬。ZeRO-3 之所以能彻底破解显存墙是因为它连模型参数本身也做了切分。64 张卡并行时每张卡只持有 70B 参数的 1/64也就是约 2.2GB 的参数分片。梯度是在反向传播过程中通过 reduce-scatter 操作汇总到各卡的优化器状态本来就按分片维护不需要在全集群范围内同步副本。2.3 参数重建all-gather 的代价与优化ZeRO-3 不是把参数切完就完事了。Transformer 的前向计算是按层进行的计算某一层时该层的完整权重必须存在否则没法做矩阵乘法。所以 ZeRO-3 在计算每一层之前会先通过 all-gather 把所有卡上的该层参数分片拼成完整权重这一层算完之后这份完整权重就被丢弃因为内存里不应该再保存冗余副本。打一个比方ZeRO-3 就像图书馆的临时借书需要读某一章时才从全集群把这一章的所有书页收集起来读完立刻还回去。好处是任何时刻每张卡上都没有完整的书坏处是每次借书都要触发一次跨卡通信。在实际训练里这种逐层 gather 如果完全是串行等待通信延迟会直接吃掉计算收益。因此 DeepSpeed 做了参数预取prefetch在当前层计算的同时提前把下一层需要的参数分片 gather 到本卡让通信和计算重叠起来。控制这个行为的主要是两个配置项stage3_prefetch_bucket_size预取缓冲区的粒度太大容易浪费显存太小通信重叠效果差。stage3_max_live_parameters允许驻留在显存中的参数规模上限显存紧张时可以调小。我在训 13B 稠密模型时试过调这两个参数最明显的体感是stage3_max_live_parameters设大了显存峰值高但速度快设小了峰值低但因为预取不足GPU 经常处于等通信的状态。折中方案一般取单层参数量的 2~4 倍具体值要结合网络带宽试。3. DeepSpeed 混合并行共享层走 ZeRO-3专家层走专家并行3.1 DeepSpeed MoE 的两个并行维度很多人以为 ZeRO-3 既然能把参数切得这么碎那直接用 ZeRO-3 切 MoE 不就行了理论上可以实践上很亏。原因在于专家参数在 MoE 层里是按专家为粒度组织的一个 token 只访问其中 top-K 个专家。如果用 ZeRO-3 把所有专家参数也切成 64 份那么每次要访问某个专家时都得对这个专家做一次 all-gather把它的完整权重从 64 张卡上拼回来。同一层里可能有几百个 token 被路由到不同的专家每次都需要跨卡 gather 不同专家的参数通信量会爆炸而且是碎片化的通信带宽利用率极低。所以 DeepSpeed 的做法是共享层Attention、Embedding 等走 ZeRO-3 切分MoE 层专家走专家并行EP切分。专家参数在 EP 维度上按专家编号静态分布每一张卡持有固定的一组专家分片token 不是去收集参数而是通过 all-to-all 通信把自己送到目标专家所在的卡上。具体到集群布局GPU 被组织成二维网格 EP × DP。假设 32 张 GPU设置 EP8、DP4每个专家并行组里有 8 张卡64 个专家被切到 8 张卡上每卡持有 8 个专家。4 个数据并行组共享同样的模型结构处理不同的 micro-batch。共享层参数除了被 EP 组复制外还可以由 ZeRO-3 继续在全部 32 张卡上切分。一句话总结MoE 层解决的是专家太多放不下EP 解决非专家层解决的是每卡都要有一份完整稠密层ZeRO-3 解决。这两个机制互补而不是二选一。3.2 一份可用的 zero_optimization moe 配置下面是我在实际训练中用过的一个 DeepSpeed 配置骨架关键项我都加了说明。注意 JSON 本身不能带注释这里展示的是完整配置含义在代码块后面解释。{ train_batch_size: 64, train_micro_batch_size_per_gpu: 2, zero_optimization: { stage: 3, contiguous_gradients: true, stage3_max_live_parameters: 1000000000, stage3_prefetch_bucket_size: 50000000, reduce_scatter: true }, moe: { num_experts: 64, top_k: 2, capacity_factor: 1.5, eval_capacity_factor: 2.0, min_capacity: 4, drop_tokens: true, aux_loss_coef: 0.01 } }逐项说明num_experts每个 MoE 层的专家总数。DeepSpeed 也支持给不同层配置不同专家数用列表传入即可。top_k每个 token 激活的专家数量训练阶段最常见的是 2推理阶段为了速度常用 1。capacity_factor训练时专家容量上限的松弛因子。专家的理想容量是每个专家平均接收的 token 数实际路由不可能完全均匀所以留出余量。1.0 意味着严格按平均容量收 token超出就丢1.2~2.0 是常见区间。eval_capacity_factor验证/推理阶段的容量因子通常要比训练大目的是少丢 token保证评估指标真实。min_capacity每个专家最少处理的 token 数。太小会让梯度稀疏太大又浪费计算我一般设 4~8。drop_tokens超出容量的 token 是否丢弃。设为false时溢出 token 会走框架里的保留专家路径取决于has_reserved_experts配置代价是增加通信量设为true则直接丢弃吞吐更高但可能影响质量。aux_loss_coef负载均衡辅助损失load balancing loss的权重这个在第 4 部分详细展开。3.3 capacity_factor 与 token 丢弃策略这里值得单独说一下 capacity_factor因为它是 MoE 训练里最容易被忽略、又最容易影响训练质量的参数。如果路由完全均匀64 个专家、batch 65536 个 token每个专家正好收到 1024 个 tokencapacity_factor1.0 就够了。但真实训练中 token 分布的指数长尾很常见热门专家可能收到 2000 个 token冷门专家只有 300 个。如果 capacity_factor 设得太紧热门专家的超额 token 会被疯狂丢弃模型损失里会莫名其妙出现一个不该有的精度缺口。设得太大每个专家都会 pad 空位置计算资源浪费在 padding token 上通信量也变大。我的做法是训练初期用 1.5等负载均衡损失降下来一点后再逐步收回到 1.2。不要一开始就追求 1.0 的高利用率除非你的 aux loss 已经收敛得很好。另外要留意drop_tokens对 loss 曲线的影响。如果发现训练 loss 突然出现周期性抖动先看看是不是 token 丢弃比例太高。DeepSpeed 的日志里会输出 MoE 层的统计信息包括每个专家接收的 token 数和丢弃比例这是一个非常重要的监控信号。4. 负载均衡MoE 训练的隐形命门附可运行的实现片段4.1 路由器坍缩为什么需要 load balancing lossMoE 训练里最经典的问题是路由器在训练过程中会坍缩少数几个专家被反复选中绝大多数专家长期接不到 token梯度稀疏甚至为零最终变成只有几个专家在干活整个 MoE 退化成一个小号稠密模型。这本质上是一个正反馈循环。某个专家前期因为随机初始化稍微强一点路由概率就高一点路由概率高这个专家拿到的梯度就多梯度多它就变得更强更强路由概率更高。于是赢者通吃其他专家被饿死。解决办法是在总损失里加入一个负载均衡辅助损失鼓励路由器把 token 尽量平均地分配给所有专家。这个技术最早在 GShard 和 Switch Transformer 里成熟起来DeepSpeed 的 MoE 层内部也内置了类似机制。4.2 一个最小可读的负载均衡损失实现有时我们需要在自定义模型里自己实现 MoE 层我给出一个简化但逻辑完整的负载均衡损失实现方便理解这个损失到底在算什么。import torch def load_balancing_loss(router_logits, top_k_indices, num_experts, top_k): # router_logits: [batch, seq, num_experts] # top_k_indices: [batch, seq, top_k]每个 token 被选中的 top_k 个专家编号 num_tokens router_logits.shape[0] * router_logits.shape[1] # 1. 统计每个专家在本 batch 中被选中的 token 次数 expert_counts torch.bincount(top_k_indices.flatten(), minlengthnum_experts).float() # 2. 每个 token 对被选中专家给出的归一化路由概率 probs torch.softmax(router_logits, dim-1) picked_probs torch.gather(probs, -1, top_k_indices) # [batch, seq, top_k] # 3. 按专家维度聚合平均路由概率 mean_probs torch.zeros(num_experts, devicerouter_logits.device) flattened_probs picked_probs.view(-1) flattened_experts top_k_indices.view(-1) mean_probs mean_probs.scatter_add_(0, flattened_experts, flattened_probs) mean_probs mean_probs / (num_tokens * top_k) # 4. Switch Transformer 风格的 aux loss aux_loss torch.sum(expert_counts * mean_probs) * num_experts return aux_loss这个实现的思路分三步用bincount统计每个专家实际接收了多少 token得到expert_counts这衡量的是负载的均匀程度。找到每个 token 的 top-K 专家后把这些 token 对相应专家的 softmax 概率收集起来按专家聚合得到每个专家的平均路由概率mean_probs。两个向量逐元素相乘并乘以专家数。如果所有专家接收 token 数完全相同且平均概率相同这个值会等于最小值某个专家独占所有 token 时值会变大从而在反向传播中惩罚这种坍缩。实际工程实现不会这么朴素框架会做 batch 化的并行归约但核心逻辑就是这样。aux_loss_coef就是给这个损失的乘数。DeepSpeed 的内置 MoE 层会自动把 aux loss 注入到总 loss 里如果你完全用自己的nn.Module实现 MoE则需要手动把它加到最终 loss 上否则你不会观察到任何负载均衡效果。4.3 aux_loss_coef 怎么调、负反馈怎么观察关于系数论文里最常见的默认值是 0.01但我在实践中建议从 0.001 起步。原因很简单aux loss 虽然能压制路由器坍缩但它本质上是在牺牲主任务的拟合能力去换取路由均匀性。系数太大模型参数会往平均分配 token这个目标偏移主任务 loss 不降反升系数太小又压不住坍缩。我试过在中小规模 MoE 上把系数从 0.01 调到 0.001训练 loss 明显更平滑专家利用率也稳定。怎么判断当前系数合不合适两个信号看主任务 loss 曲线是否出现平台期后突然跳降之类的异常说明路由器可能在某个阶段发生了坍缩后被 aux loss 硬拉回来。看专家利用率方差。DeepSpeed 日志里会输出每个专家平均接收的 token 数方差如果一直在高位且不下降就加大系数如果主任务 loss 明显劣化就减小系数。这里分享一个我常用的偏门技巧把eval_capacity_factor调大、训练时把drop_tokens打开然后定期观察验证集上的专家利用率。因为验证集没有那么多干扰梯度更能反映路由器是否已经学到了一个合理的分配策略。5. 实战踩坑显存估算、all-to-all 通信和数值稳定性5.1 显存估算够不够不能光看参数量很多人在开训前问我的模型能不能塞进 80G 显存这类问题的标准回答是看你怎么切。我提供一个粗估公式参数相关显存 ≈ 总参数 × 16 字节 / 参与切分的卡数如果用了 ZeRO-3 且共享层和专家层都在集群维度切分卡数就是全部 GPU 数量如果 MoE 层走 EP 而共享层走 ZeRO-3则专家部分按照 EP 并行度切分共享部分按全集群切分。举一个具体例子。假设模型共享层 20 亿参数64 个专家每个 1 亿参数总参数 84 亿。在 32 张 A100 上EP8、DP4共享层 20 亿参数按 ZeRO-3 被 32 卡切分每卡 0.625 亿参数约 10GB 的参数相关存储训练时含优化器状态。专家层 64 亿参数按 EP8 切分每卡 8 亿参数约 12.8GB。两项合计约每卡 22.8GB 的参数量相关显存再叠加激活值和通信 buffer80GB 完全可以扛住。用这个公式规划显存比猜测靠谱得多。值得注意的是激活值activation的显存和 batch size、序列长度、层数线性相关不能用参数公式去套。MoE 的激活值计算量虽然只和 top-K 相关但显存中保存的中间结果还是要按层数累积的。5.2 all-to-all 是 MoE 通信的隐形瓶颈ZeRO-3 的通信模式是 all-gather / reduce-scatter通信量和模型参数规模相关。MoE 专家并行则完全不同它是把 token 从当前卡搬运到专家所在的卡上通信量取决于 token 数量和 hidden size与专家参数大小无关。这个 all-to-all 通信是 MoE 训练里比 ZeRO-3 更隐蔽的性能瓶颈。原因在于每个 token 都要带上 hidden state以及可能的多层中间状态发给目标卡数据量 token 数 × hidden size × 字节数。micro-batch 越大、序列越长all-to-all 的通信量越大。跨节点情况下NVLink 变成了 IB 网络带宽下降一个数量级all-to-all 会直接把训练拖垮。我在跨节点环境下吃过一次大亏EP16 把专家切到了两个 node 上每个 step 里光等 all-to-all 就是普通稠密模型的好几倍。后来把 EP 控制在单节点内比如 A100 8 卡节点EP8让专家并行组尽量落在同一台机器的 NVLink 域内训练吞吐立刻回升。另外一个实用技巧是micro-batch size 不能太小。all-to-all 的固定通信开销比较高batch 太小会让通信占比变大。如果发现小 batch 下 GPU 利用率不高优先考虑增大 micro-batch而不是先去调 ZeRO-3 的预取参数。5.3 稳定性问题数值、学习率与梯度裁剪MoE 训练的 loss 曲线比稠密模型更容易出幺蛾子我遇到的典型问题包括数值溢出。路由器的 softmax 对 fp16 精度很敏感logits 一旦偏大softmax 后的梯度很容易溢出表现为 loss 突然变成 NaN。如果代码环境支持 bf16我建议直接切 bf16如果只能 fp16需要把初始 logits 做缩放或者加一个 logits 归一化层。学习率偏高。同一套数据稠密模型用 1e-4 没问题MoE 模型用同样的 LR 却可能震荡。我在多个规模上都验证过一个规律MoE 的合适学习率往往只有稠密模型的 1/5 到 1/2。不要急着上大 LR先用小 LR 把路由器的初始分布训稳再考虑提速。warmup 不足。MoE 训练建议把 warmup 拉长到总步数的 10%~20%。前期的路由器状态非常敏感短时间内从随机初始化跳到高置信度路由很容易早熟导致后面专家利用率怎么调都调不回来。我一般在 DeepSpeed 配置里用 LR scheduler 的 warmup 参数直接设 0.1某些模型上设过 0.2。梯度裁剪。把gradient_clipping设成 1.0对 MoE 来说几乎是必须的。MoE 层因为 token 到专家的分散映射梯度方差比稠密模型大很多不裁剪的话一个异常 batch 就可能毁掉整个训练进程。如果你遇到 loss 震荡按这个顺序排查先看是不是数值溢出再看 aux loss 系数是不是太大然后再看 capacity_factor 和 drop_tokens 是否导致 token 被大量丢弃最后才考虑调学习率和 warmup。大部分震荡问题都能在这四个环节里找到根因。最后说一个我个人一直保留的习惯任何 MoE 模型新开训练前我都会先在 8 卡小规模上跑 500 步打开 DeepSpeed 日志盯着两个数——每个专家的 token 接收方差和 token 丢弃比例。这两个指标正常了再上大规模不迟指标异常就回头调 capacity_factor 和 aux_loss_coef。磨刀不误砍柴工这一步帮我避免过至少三次大规模训练的返工。