做大规模模型训练这几年，被问得最多的一个问题就是：“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 精度下权重就要占 14GB，Adam 优化器状态按每参数 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 层的 Transformer，8 张卡时每张卡只保存其中若干层的全部参数，而不是每层都拆成 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 第 n+1 层参数。从效果看，只要代码里没有明显的同步阻塞，这个开关值得一直开着。
3. MoE 训练为什么难：路由、负载均衡和并行
3.1 路由和 all-to-all：token 是怎么去找专家的
MoE 的难点不完全在“一体积大”，而在于它的数据流不是单纯的“层内计算—层间传递”，而是多了一步 token 和专家的配对。一个非 MoE 的 Transformer，token 在每个 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.01；z-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-MoE（Pyramid 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/8=8。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_size=d_model, expert=FeedForward(d_model, d_ff), num_experts=expert_count, ep_size=ep_size, use_residual=False, fuse_gate=True, top_k=1, min_capacity=4 ) def forward(self, x): x = x + self.attention(x, x, x)[0] x = x + self.moe_ffn(x) return x这里MoE包装器是 DeepSpeed 提供给用户的，不需要自己实现门控和 all-to-all；ep_size会自动决定它内部的通信方式。调用训练入口时，用model, optimizer, _, lr_scheduler = deepspeed.initialize(config=config_path, model=model)就把分片和 MoE 全接上了。注意：MoE 层不参与 ZeRO-3 的参数分片时，必须在模型初始化之前把dist.init_process_group完成，否则ep_size的分组信息拿不到。
4.3 显存估算：一张表给个参考
我经常被问“这配置到底吃多少显存”，这个问题没法给单一数字，因为和隐藏维度、专家数、序列长度、batch size 都有关系。但可以给一个按经验推算的参考表，前提是 16 卡 A100 80G、隐藏维度 2048、FFN 维度 8192、expert_count=16，非 MoE 层用 ZeRO-3 全分片：
| 模型规模参考 | 专家并行数 ep_size | 单卡理论峰值显存 | 能否训练 |
|---|---|---|---|
| 稠密 7B 转 MoE，总参数 24B | 4 | 接近 60GB | 可以，但通信压力大 |
| 稠密 13B 转 MoE，总参数 45B | 8 | 55GB 左右 | 可以，缓冲区和激活需要调 |
| 总参数 100B 以上 | 16 | 60GB 以上 | 必须开 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 卡时设 5e8，32 卡以上缩到 2e8，防止单个 all-gather 调用长时间占住网卡；- 对 MoE 的 all-to-all 做微调：DeepSpeed 在
MoE层有use_rts（Random Token Selection）参数，它可以限制每个专家的 capacity，超出的 token 会被随机丢弃，等价于给通信缓冲区加了一个保险丝。容量设成min_capacity=4起步，后面看 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 数接近 0 | Aux Loss 权重太小或不平衡 | 调大moe_gate_loss_wt，同时增大top_k |
| all-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 配置没对齐，却一直在调优化器参数”的无效操作。希望你不用走我那么多弯路。
