ZeRO-3遇上MoE：多卡训练显存爆炸与通信瓶颈的破解之道

前几天有个朋友跟我抱怨，手里8张A100，想训一个20B的稠密模型，结果OOM。batch size从32一路降到1，还是OOM。他说了一句让我印象深刻的话：不是有8张卡吗？几十张卡还不够？这个问题其实非常典型——很多人低估了大模型训练里优化器状态那部分开销，也低估了数据并行下显存冗余的严重程度。更别说MoE这种参数爆炸式增长的架构，如果你把专家层当成普通参数去套ZeRO-3，通信开销能直接把训练拖到怀疑人生。

这篇文章我打算把DeepSpeed ZeRO-3和MoE训练这两件事彻底讲透。适合的人群很明确：已经在用多卡训练、知道什么是数据并行和模型并行，但对"为什么ZeRO-3能省显存"只有模糊印象的人；以及想尝试MoE架构、但被路由机制和负载均衡折磨过的人。我会从显存账开始算，讲到ZeRO-3的通信本质，再拆MoE为什么不能简单套ZeRO-3，最后给一份能直接跑通参考的DeepSpeed配置和训练过程中的避坑经验。

1. 显存问题不是"模型太大"，而是"每一张卡都在复制同一个模型"

1.1 数据并行的显存账：三层冗余

很多人对显存的第一反应是"模型权重有多大，显存就要多大"。这个直觉在小模型时代成立，但到了大模型训练阶段，完全不够用。

你用数据并行在8张卡上训一个模型，最朴素的理解是"每张卡各算一批数据"，这没问题。但问题是，8张卡里每一张都存了完整的模型权重、完整的梯度，以及完整的优化器状态。也就是说，同样的东西被复制了8份。这里有个很关键的点——优化器状态往往比模型本身还大。

以稠密Transformer为例，假设模型有10B参数，跑在FP16下，光权重就是20GB。如果你用Adam优化器，它需要维护三份FP32状态：master weights（参数主副本）、momentum（一阶动量）、variance（二阶动量）。每份40GB，三份就是120GB，加上FP16的权重20GB和梯度20GB，整体训练状态下，参数量10B的模型实际占用高达160GB。注意，这只是一个GPU在数据并行下需要的显存。8张卡就是8份，1.28TB。

你这时候再回头看"20B模型800G显存应该够了吧"这种直觉，就明白为什么完全不靠谱了。显存不够的根本原因不是模型大，而是你把优化器状态和梯度都摆到了每一张卡上。

1.2 从ZeRO-1到ZeRO-3：先分优化器，再分梯度，最后分参数

ZeRO这个名字的意思是Zero Redundancy Optimizer，它的核心就是消灭这种冗余。它把训练状态拆成三样东西：optimizer state（优化器状态）、gradient（梯度）、parameter（参数）。这三个阶段是逐步递进的：

ZeRO-1和ZeRO-2相对容易理解：优化器状态本来就是训练中没有直接参与计算的"附属品"，谁持有哪一块，在需要更新参数时再聚合起来就行。ZeRO-3则是把模型参数本身也分掉了，这就引出了一个新的核心矛盾——参数都分掉了，前向计算的时候算子要从哪里取权重？

这也是我在往下讲之前，想先给读者建立的一个基本框架：ZeRO-3最本质的特点不是"显存少"，而是训练过程中每一次拿到完整参数的行为，都是一个临时状态。它是把"常驻完整参数"变成了"瞬态完整参数"。

2. ZeRO-3的瞬时全员汇聚：参数是租来的，不是常住的

2.1 一次前向计算里发生了什么

假设你有一个24层Transformer，用ZeRO-3把参数均匀分到32张卡上，那么每张卡实际常住的只是"所有层参数的1/32"。这听起来很美好，但问题来了：计算第一层的时候，第一层的权重W1并不是整块放在你这张卡上的，它在32张卡上各有一片。

这时候系统会做一次all-gather通信：把32张卡上的W1碎片全部汇聚起来，让每张卡都拿到完整的W1，然后才能做矩阵乘法。做完第一层，立刻把W1的完整副本丢掉。到了第二层，再对W2做同样的事。

所以你会看到，ZeRO-3的显存曲线是很"抖"的：层与层之间，显存占用会周期性上涨下落，而不是像数据并行那样，从头到尾顶着一个巨大的模型副本。如果你用nvidia-smi监控ZeRO-3训练过程，经常会看到显存数值上下跳动，那不是异常，那是它在"租参数"。理解了这一点，你对ZeRO-3省显存的机制才算真正入门。

2.2 为什么要配合gradient checkpointing和NVLink

因为参数是临时聚合的，所以通信次数很频繁。一层的forward需要一次all-gather，backward计算该层梯度之前还需要再all-gather一次同样的参数。也就是说，每个Transformer层至少触发两轮全员通信。

为了减少通信量，常见做法是打开gradient checkpointing。开启之后，前向过程不会保留每一层的激活值，而是在反向传播时重新计算一次前向。表面上这增加了计算量，但因为反向重新计算前向时，本来就要再一次all-gather参数，这跟反向计算梯度的参数获取是可以复用的。你省掉的，是那些常驻激活值占用的显存，以及与之相关的部分通信时间。

还有一个硬性要求：节点内的卡间互联最好有NVLink。因为all-gather的通信模式非常吃带宽，PCIe Gen4的带宽和NVLink差距在数倍以上。我自己测试过，在纯PCIe的机器上开ZeRO-3，一个中等规模模型训练速度可能比ZeRO-2还慢30%以上。道理很直接：你省下了显存，却没省下通信，反而添了新通信。

2.3 ZeRO-3不是无代价的：通信开销的账

ZeRO-3的另一个关键问题是越小的模型，越不值得用。为什么？因为显存省下来需要一个前提：模型大到单卡放不下。如果一个模型本来单卡就能放下，你强行开ZeRO-3，是在用频繁的全员通信去换取用不上的显存空余，增加了通信瓶颈，训练速度只会更慢。

这里我给一个很粗糙的经验值：当你的模型参数量在5B以下、且单卡显存足够容纳优化器状态时，用ZeRO-2可能就已经合适。ZeRO-3通常是在模型大到连"每卡一份参数副本"都放不下的时候，才真正体现价值。这也是为什么很多训练框架会建议，小规模微调用不上stage 3，stage 3是给预训练大模型准备的。

理解了这个"瞬时聚合"的机制，再看MoE，就会意识到一个严重问题：MoE的模型参数比同样计算量的稠密模型大得多，如果一个专家层的参数也走这种"全员分片、临时聚合"的路线，通信开销几乎是不可接受的。这就是下一章要展开的矛盾。

3. MoE的稀疏激活：参数可以很多，但每个token只走一小段

3.1 为什么MoE能被"放大"：计算与参数量解耦

MoE，全称Mixture of Experts，中文常叫专家混合。它把Transformer里的FFN层替换成一组并行的专家网络，每个专家是一个独立的FFN。路由网络（router）会根据当前token的特征，只挑选其中top-k个专家来计算。

MoE最反直觉的地方在于：它的参数量可以非常大，但每个token实际触发的计算量只相当于一个稠密小网络。比如Switch Transformer的经典设计，32个专家，每层有32个FFN，但每个token只激活其中1个专家。这样你的模型参数翻了32倍，前向计算的FLOPs几乎没变。

这就是"稀疏激活"的价值：参数规模与计算规模解耦。你可以在不增加计算预算的前提下，让模型记住更多知识、拥有更大的容量。这也是近两年很多大模型采用MoE架构的核心动机——同样的训练算力，容量可以大好几倍。

3.2 路由器是MoE的调度中枢

路由器本质上是一个很小的线性层，输入是token的向量表示，输出是"这个token去每一个专家的概率打分"。对打分做个softmax，选top-k，然后把这个token送到对应的专家里去算。

这句"送到对应的专家里去算"听起来很简单，但在分布式训练里，它意味着token要跨卡搬运。专家不是分布在每张卡上的完整副本，而是每张卡只拥有部分专家的完整副本。当一张卡上的token被路由到了另一张卡上的专家时，它必须通过通信把token的隐藏状态送过去，计算完专家输出再送回来。

这个过程叫all-to-all通信，是MoE训练里最绕不开的成本来源。我在实际工程里观察到的现象是：当专家数量多、token路由分散时，all-to-all的通信量可以占到整个训练通信量的30%以上，有时候甚至超过反向传播的梯度通信。这也是为什么，很多MoE训练框架在单机多卡（节点内）会优先走NVLink，跨节点则必须优化路由分组，否则训练吞吐会很难看。

3.3 MoE训练最直观的显存直觉

上面讲了，专家在分布式环境下不是"全员复制"，而是"分布存放"。那么问题来了：MoE里有一些层是稠密的（比如attention层、embedding层），这些层如果也要在每张卡上放一份完整副本，显存照样爆炸。

所以MoE模型的显存规划，要分成两块看：一块是基础层（attention、embedding等非专家层），另一块是专家层。两者需要的并行策略完全不同。如果你用数据并行训MoE，基础层和专家层都全员复制，那MoE参数翻了好多倍的直接后果就是显存翻好多倍，这显然是不可持续的。DeepSpeed的解法，是把这两块分开对待——这就是下一章的内容。

4. ZeRO-3 + MoE为什么需要专门设计：专家不能全员分片

4.1 常规ZeRO-3在专家层为什么不好使

用常规ZeRO-3去训MoE，最朴素的做法是：把整个模型的参数都分片，包括所有专家层。这样做显存确实省下来了——但代价是什么？是专家参数在每次计算前都要all-gather一次。假设你有32个专家，分成32个GPU，每个GPU持有1/32的专家参数，前向计算某个token时，它的top-1专家权重散在所有卡上，你需要先从32张卡把专家权重聚合到当前卡，才能算这个token的专家输出。

这不是一次性的，而是每一层、每一个batch都要做。最严重的情况是，token被路由到了很多不同的专家，你要反复多次发起large-scale的all-gather。通信量会膨胀到完全不可接受。我甚至见过小型实验里，启用常规ZeRO-3训MoE，整个训练过程几乎60%时间都花在通信等待上，GPU计算单元在摸鱼。

4.2 DeepSpeed-MoE的混合并行设计

DeepSpeed对MoE的处理思路很有意思：它没有让专家层也走"全员分片"，而是走了另一条路——专家按GPU分布，每个GPU持有专家的完整副本，而不是分片。

具体来说，在DeepSpeed-MoE里，模型被拆成两部分看待：

• 基础层（attention等稠密层）：走ZeRO-3的分片逻辑，该省显存就省显存。
• 专家层：每个GPU只放一部分专家，且这部分专家参数是完整的、不被分片的。

这叫做专家并行（expert parallelism）。比如一个MoE层有32个专家，你用8卡训练，那么每张卡通常放4个专家。每个token算出来要去哪个专家之后，就通过通信把token序列直接发到对应的GPU上，让那块GPU用完整的专家参数完成计算。

为什么要这么设计？因为专家层的计算具有离散性和局部性。一个token只会用到少数几个专家，如果把专家参数分成碎片，所有计算都得等聚合完成，等于把本来可以并行的小计算强行串行化。而"每个专家整块放某张卡上"，计算只需要在持有该专家的那张卡上完成，即使需要跨卡搬运token，也比频繁搬运参数碎片高效得多。

4.3 all-to-all通信：路由后的token如何找专家

这里补一下前面提到的all-to-all。在DeepSpeed-MoE框架里的流程大致是这样：

1. 每张卡上的router先算出每个token要去的专家编号。
2. 系统统计本卡token要去哪些远程专家，并按目标整理成一个通信张量。
3. 通过all-to-all原语，把token的hidden state发给目标卡。
4. 目标卡用本地完整的专家参数完成专家的FFN计算。
5. 再通过一次all-to-all把结果送回原卡。

你看，通信搬运的是token本身，而不是专家参数。token的hidden state大小远比模型参数小得多，所以这种设计对带宽的要求是可控制的。当然，当token数量非常大、路由特别分散时，all-to-all的token总量会显著增长，这时候就需要在路由策略上做文章——这正是下一章要说的负载均衡。

5. 路由崩塌与负载均衡：MoE训练的隐形杀手

5.1 训练一开始就会出现的"少数专家垄断"

MoE训练有个非常容易踩的坑：路由崩塌（routing collapse）。现象是，随着训练推进，路由器会越来越"偷懒"，总是把大部分token分给少数几个专家，其他专家长期收不到token，梯度几乎为零，变成了"死专家"。

为什么router会偷懒？因为初始化后，有些专家可能恰好对某些token的打分略高一点，这些专家得到训练后，它们的表征对当前阶段的数据更适配，router就会更倾向于选择它们，形成"富者愈富、贫者愈贫"的正反馈。最后，模型的容量优势被浪费掉，实际训练效果甚至不如一个稠密小模型。

这种问题在训练初期特别容易发生。我在实验里见过最夸张的情况：32个专家里只有3个专家收到了超过90%的token，其余29个专家基本没干活，但训练显存和通信开销一分钱都没少。

5.2 负载均衡loss从哪来，怎么用

为了让专家负载尽量均匀，主流做法是在训练loss里加一项负载均衡的辅助loss。以Switch Transformer的经典设计为例，它会统计两样东西：

• f_i：被路由到专家i的token比例（实测值）
• P_i：路由给专家i的平均概率（router给出的期望值）

辅助loss的形状大致是：alpha * N * sum(f_i * P_i)。当所有专家收到token的比例和router打分的期望接近均匀分布时，这一项取到最小值。alpha是权重系数，通常在1e-3到1e-2这个量级。

注意，这个辅助loss的核心思想是让router的预期分布与实际分布保持一致，并让两者都尽量均匀。如果只有P_i均匀而f_i不均匀，说明router打分的分布虽然平均，但top-1竞争让部分专家依然被饿死。所以f_i的均匀性对训练效果至关重要。

5.3 动态负载均衡与expert capacity的实际取舍

除了辅助loss，还有两个工程层面的参数需要关注：expert capacity和capacity factor。

expert capacity指的是每个专家最多能处理多少token，超过的部分会被丢弃或者通过aux loss惩罚。这个设计是为了防止某几个专家因为瞬间流量过载而拖慢训练。capacity factor则决定了容量上限的余量，比如capacity_factor=1.25表示每个专家可以超载25%的token。

模块里还有一个常见技能：noisy top-k routing。它在路由打分上人为加一点噪声，让router无法轻易对同一批token重复选择同一个专家，相当于给路由过程加了随机扰动。这个噪声的方差在训练中逐渐衰减，初期帮助探索，后期趋于稳定。

从我实操的角度给个建议：如果你刚开始训练MoE，优先观察每个batch的专家负载统计，而不要只看总loss。一旦发现某几个专家负载长期为0或者接近0，先别着急调大alpha——把alpha调大虽然能解决负载不均，但会严重稀释主loss对模型质量的约束，导致模型本身学不好。正确顺序是：先确认top-k设置和capacity factor是否合理，再微调alpha，最后才考虑改路由噪声策略。

6. 实操：DeepSpeed配置到常见问题排查

6.1 一份能跑通的ZeRO-3 + MoE配置长什么样

理论讲再多，终究要落地到配置。下面是一份DeepSpeed训练MoE模型时相对常见的配置结构，字段含义我都注释了：

{ "train_batch_size": 32, "gradient_accumulation_steps": 4, "optimizer": { "type": "Adam", "params": { "lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8 } }, "zero_optimization": { "stage": 3, "allgather_partitions": true, "reduce_scatter": true, "overlap_comm": true, "contiguous_gradients": true }, "moe": { "enabled": true, "model_parallelism_size": 1, "expert_parallel_size": 8, "num_experts": 32, "top_k": 1, "aux_loss_coef": 1e-2, "capacity_factor": 1.25, "all2all_backend": "nvlink" } }

几个字段我说一下我的理解。expert_parallel_size代表专家并行度，在这个例子里是8，意思是专家层被分布到8张卡上，每张卡持有4个专家的完整副本。model_parallelism_size是模型并行度，通常配合Megatron-LM使用，纯DeepSpeed环境设成1就行。aux_loss_coef负载均衡loss的权重，我建议初次实验从1e-2往下调，先跑200步看路由健康度，如果负担均衡但loss收敛变差，就降到1e-3。

启动训练时可以加一句：

deepspeed --num_gpus=8 train.py \ --deepspeed_config deepspeed_config.json

训练起来之后，确认是否进入ZeRO-3阶段，看日志里有没有类似stage=3的初始信息。DeepSpeed启动时会把zero stage、optimizer类型、offload状态等打印出来，别忽略这些启动日志。

6.2 安装DeepSpeed常见的报错思路

看你在这个标题下还搜过安装报错的问题，我随手把这个经验也写下。pip install deepspeed报错的原因五花八门，见得最多的几类：

• CUDA版本和PyTorch版本错配：DeepSpeed有些算子需要现场编译，编译时找不到CUDA_HOME，或者找到的CUDA版本和PyTorch编译时用的CUDA不一致。最常见的报错就是CUDA_HOME not defined或failed to find libcudart。我的建议是，先把nvcc --version和python -c "import torch; print(torch.version.cuda)"打出来对比，确保同一个版本。
• gcc版本太老：DeepSpeed有些用C++写的extension，比如CPU Adam、sparse attention，对编译器版本有要求。如果编译卡在某个C++语法上，大概率是gcc版本问题，换到gcc 8以上的环境会顺利很多。
• 安装阶段跳过不需要的算子：如果你不需要某些half-precision或者sparse算子，可以在安装时设置DS_BUILD_OPS=0来跳过强制编译。实测下来，很多报错都能靠跳过那些重算子解决；等后续真用到某个功能时，再单独编译对应的算子，按需处理。

如果你只是想在本地试ZeRO-3，不必纠结所有算子都能编译——跳过sparse attention这类重算子，对常规Transformer训练没有影响。

6.3 训练过程中怎么观察显存和负载是否健康

训练跑起来后，我一般会同时看三个指标：

第一，单卡显存曲线的波动特征。ZeRO-3模式下，显存应该是呈现周期性波动的，波动幅度和模型层深度有关。如果你发现显存一直平稳高位不落，说明你可能没真正跑在stage 3，而是退化成了stage 2或数据并行模式。

第二，GPU利用率。MoE + ZeRO-3训练时，GPU利用率低于60%不一定不正常，因为很可能是通信和计算没有充分重叠。但如果出现周期性的大段空闲，同时网络吞吐暴涨，说明all-to-all通信在排队，这时候可以考虑调整all2all_backend，或减少跨节点路由的分布度。

第三，专家负载统计。DeepSpeed在训练MoE时通常会在日志里统计每个专家的token接收情况。看到类似历史负载表之后，重点看有没有专家长期为0。一旦发现"死专家"，我建议先停下看配置，不要让它持续跑完整个训练周期，否则到后期大概率白烧算力。

结尾

最后说句真心话：ZeRO-3和MoE这两样东西，单独任何一个，理解起来都不算太难；难的是把它们放在一起时，你要时刻切换思维——基础层在忙着分片、聚合，专家层在忙着搬token、找专家。很多人训练MoE失败，不是架构搭不出来，而是没意识到"显存省下来了，通信问题才刚浮出水面"。

个人经验是，新手刚上手MoE训练时，先别急着把专家数调到64甚至128。我更推荐起步阶段用8个专家配合top_k=1，在几十步内把路由负载、显存波动和通信时间都观测一遍，形成手感，然后再逐步扩大。小配置暴露问题最快，也最便宜。等你能把8个专家的负载调得相对均衡了，再碰大规模MoE，心态会稳很多。