做大规模模型训练这几年DeepSpeed 和 MoE 这两个词基本是绕不开的。DeepSpeed 的 ZeRO-3 解决了“单卡放不下参数”的痛MoE 架构则用稀疏激活的思路把模型容量往上推了好几个量级。但很多人把两者混在一起用的时候其实并没有真正搞懂它们各自在干什么以及配合起来有哪些坑。这篇文章我把 ZeRO-3 和 MoE 训练的核心机制、配置方法、踩坑经验全部分享出来适合正在做大模型训练、想了解显存优化原理、或者准备上手 MoE 模型但还没理清思路的工程师。先说说我为什么想写这个题目。之前团队训练百亿级稠密模型时光靠 ZeRO-3 已经快被通信开销拖垮了后来切到 MoE 架构发现参数量的瓶颈被打破了但显存和负载均衡又成了新问题。如果你也在“稠密模型显存不够MoE 模型不会训ZeRO-3 配不明白”这三个状态里打转这篇文章就是给你准备的。1. 先想清楚ZeRO-3 和 MoE 各自在解决什么1.1 一个最直接的疑问MoE 架构要全部参数进显存吗这是很多人第一次接触 MoE 时最容易搞混的问题。常规理解里模型越大显存需求越大那 MoE 这种参数动辄几千亿的架构岂不是要几块 A100 才能装下答案是否定的。MoEMixture of Experts混合专家架构的关键在于“稀疏激活”。它的总参数量确实很大但每次前向计算时门控网络只会选择其中一部分专家参与计算。比如一个 64 个专家的 MoE 层每个 token 通常只激活 top-2 个专家剩下的 62 个专家虽然存在但这一轮计算里它们不需要“干活”也不需要对它们做梯度更新。所以你要区分两个概念参数量和激活参数量。MoE 的总参数量决定了模型的“知识容量”激活参数量才决定单次计算的实际开销。这就像一个大公司的员工花名册上可能躺着几万人但某一天实际出勤干活、需要发工资的只有那么几百人。那“全部参数进显存吗”我的答案是取决于你用什么方式训练。如果你用最朴素的方式把所有专家参数都放在一张卡上那当然要全部进显存而且得不偿失。但配合 ZeRO-3 的参数分片或者 MoE 原生支持的专家并行策略每个 GPU 只需要放一部分专家参数这样显存压力就会被彻底摊开。还有一点很关键MoE 的“被激活”参数才参与计算和梯度所以显存里的优化器状态、梯度也只需要跟激活参数挂钩这也是 MoE 能跑万亿参数的原因之一。1.2 ZeRO-3 到底砍掉了哪部分显存开销ZeRO 的全称是 Zero Redundancy Optimizer意思是“零冗余优化器”。它要干掉的是分布式训练里最常见的浪费每张卡都存一份完整的模型参数、梯度、优化器状态。我用一个例子帮你估一下显存构成。假设你有一个 10B 参数的模型用 Adam 优化器训练混合精度fp16下模型参数10B × 2 bytes 20 GB梯度同样 20 GBAdam 优化器状态每个参数要存 fp32 的 master weight、momentum、variance也就是 10B × 12 bytes 120 GB加在一起就是 160 GB这还只是存储开销不算中间激活值。一张 80GB 的 A100 显然是放不下的而数据并行下如果有 8 张卡每张卡都要复制这 160 GB浪费得吓人。ZeRO 三个阶段的思路很直接ZeRO-1把优化器状态分片到各个 GPU每张卡只存 1/N 的优化器状态。ZeRO-2优化器状态 梯度都分片。ZeRO-3优化器状态 梯度 模型参数全部分片。到了 ZeRO-3理论上每张卡只需要存 1/N 的完整模型状态量。换句话说10B 模型的 160 GB 存储需求在 32 张卡上摊下来只用 5 GB 左右。这个数字变化是有质变意义的因为它意味着你不再被“单卡显存上限”卡脖子。1.3 为什么单独用 ZeRO-3 或单独用 MoE 都不够如果你只靠 ZeRO-3 去训练超大稠密模型核心瓶颈很快就会暴露通信。因为 ZeRO-3 在每次前向和后向时都需要通过 all-gather 把分片的参数临时聚合成完整参数计算完再丢弃。模型越大all-gather 的数据量越大通信时间占比会高得吓人GPU 算力再强也只能等着数据从网线里过来。如果只靠 MoE 而不做显存优化那就等于把“员工花名册”放在一个办公室里虽然每次只有几个人干活但花名册本身很占地方。一旦专家数量多了光存储就够你喝一壶。而且 MoE 还有负载均衡问题如果门控网络学偏了所有 token 都往同一个专家上挤那这个专家的显存和算力就会被瞬间打爆其他专家空转整个训练直接失效。ZeRO-3 MoE 的组合本质上是“把静态存储摊开 把动态计算聚焦”这两件事同时做了。ZeRO-3 负责让参数、梯度、优化器状态不要再无脑复制MoE 负责让计算量不要再无脑膨胀。二者搭配才能在大规模场景里既装得下又跑得动。2. ZeRO-3 的原理拆解分片、通信与计算的重构2.1 从 ZeRO-1 到 ZeRO-3 的演进逻辑很多人容易把 ZeRO 的三个阶段背下来但不清楚演进逻辑。我建议你从“谁在被复制”这个角度去理解。数据并行 DDP 时代最浪费的是什么是优化器状态和梯度在所有卡上都是完整一份。每张卡算出的梯度虽然不同但要 all-reduce 求平均最后每张卡拿到的平均梯度是一样的然后各自用这个梯度更新出一样的参数副本。这导致 N 张卡就等于 N 份参数的重复存储而大家干的事完全相同。ZeRO-1 先把优化器状态切开每张卡只负责更新自己那一份参数。ZeRO-2 又把梯度切开因为梯度在 all-reduce 之后没必要完整保留。到了 ZeRO-3索性把参数也切开每张卡只保留 1/N 的参数分片。这个演进的本质是把“数据并行”逐步改造成“混合并行”通信模式、计算模式都变了存储压力则从 O(N×模型规模) 降到了 θ(模型规模)。2.2 ZeRO-3 的动态分片机制参数不是一直全的ZeRO-3 最反直觉的地方在于训练过程中绝大多数时间你的显存里并没有“完整模型”。参数以分片形式躺在各张卡的专属位置只有当某个算子需要完整的参数时才会通过网络把各分片聚合起来。具体流程是这样的前向计算开始前每个 GPU 持有 embedding、部分 transformer 层参数的分片。当计算某个 transformer 层时该层参数所在的各个 GPU 通过 all-gather 把这一层参数集合成完整版计算这一层的前向。前向结束后这一层完整参数立即被丢弃显存释放。后向过程中为了算梯度需要重新 all-gather 参数再做一次前向重计算或利用存储的激活值然后计算梯度。梯度算完后通过 reduce-scatter 同步把梯度分片到各卡。各卡用自己负责的优化器状态分片更新自己负责的参数分片。这个过程你可能已经发现了它有点像“用完即走”的临时租房模式需要完整参数时就租一块地用完立刻退掉。因为这临时聚合的参数不会被保留所以显存峰值远低于“完整模型常驻”的方案。2.3 通信开销的权衡all-gather 与 reduce-scatter 为什么需要重叠ZeRO-3 的代价是通信量显著增加。每次前向/后向都要做一次全参数级别的 all-gather这在网络带宽有限时会变成极大的性能瓶颈。以 10B 模型为例每轮 step 光参数通信就是几十 GB 的体量。所以 ZeRO-3 的进阶用法几乎都会开启通信重叠communication overlap。DeepSpeed 里对应的配置项是overlap_comm和reduce_bucket_size核心思路是把大块通信切碎让计算和通信并行。你可能觉得“并行”是理所应当的但真正做优化的时候需要把参数请求的粒度调得足够小小到 GPU 在计算其他层的时候网络已经在搬运下一层的参数。我在训练 7B 模型时做过对比不开 overlap单 step 耗时 2.1 秒开了 overlap 并调整 reduce_bucket_size 后降到 1.2 秒左右。如果你在训练中观察到 GPU 利用率很低但网络吞吐很高往往就是通信没有和计算重叠好。2.4 ZeRO-3 的 offload 策略显存不够用内存来凑ZeRO-3 还有一个关键能力是 offload也就是把优化器状态、甚至参数和梯度放到 CPU 内存或 NVMe 上。这在大模型训练里几乎是标配。offload_optimizer: 把 Adam 的动量状态放到 CPUGPU 只保留参数和梯度。offload_param: 把参数也放到 CPUGPU 计算前再拉取。我自己用过 NVMe offload 训过 100B 级别的模型速度当然比全 GPU 慢很多但在硬件有限的条件下“慢但能跑”远好于“跑不起来”。这里提醒一点offload 千万不能盲目开offload_param会增加大量 CPU-GPU 通信如果 CPU 内存带宽不够反而会让训练速度暴跌。3. MoE 训练的核心机制门控路由、稀疏激活与负载均衡3.1 MoE 层的结构一个路由器加上一堆专家网络理解 MoE 的最好方式是把一个 Transformer 层拆开看。常规的 FFN前馈网络就是一个全连接层输入经过它之后得到输出。MoE 则把这一个 FFN 替换成一组 FFN每个 FFN 叫做一个“专家”。同时新增一个“门控网络”Router它的作用是为每个 token 选择最合适的 top-k 个专家。比如一个 128 专家的 MoE 层门控会输出一个 128 维的概率分布然后取概率最高的 2 个专家把 token 的 hidden state 分别送入这两个专家计算最后按概率加权求和。核心寄存器是模型变大了但计算量没有等比例变大。稠密模型和 MoE 模型如果激活参数量相同则 FLOPs 几乎一致但 MoE 的总参数量可能是前者的几十倍。这也是为什么 PaLM、GPT-4 这类模型可以宣称千亿万亿参数但实际推理时单 token 的计算量并没有夸张到离谱。3.2 稀疏激活下的反向传播只更新被激活的专家很多初学者会问MoE 反向传播时没有参与前向的专家需要更新吗答案是不需要也不应该。因为前向时这些专家根本没有产生计算梯度链就不会走它们。门控网络通常会在这轮更新以学习“哪些专家适合处理哪些 token”的分配策略。这里有个容易忽视的细节top-k 的选择不可导。门控输出概率分布是连续的但“选出哪几个专家”这个操作是离散的梯度没法直接流回门控。实践中的做法是用概率加权softmax 输出乘以专家输出这样梯度可以通过权重回流到门控参数。也就是说门控学的不是“正确选择专家的 0/1 标签”而是“每个专家应该分到多少权重”。3.3 负载均衡不用 aux lossMoE 会自我崩溃MoE 训练里最经典的问题就是负载不均衡load imbalance。如果门控网络发现某个专家特别“好用”它就会把越来越多 token 路由到那个专家形成正反馈这个专家学得更快门控更喜欢它其他专家则慢慢饿死最后整个模型退化成只有少数专家在工作容量优势荡然无存。解决办法是在损失函数里加一项辅助负载均衡损失auxiliary load balancing loss。DeepSpeed 的配置里对应aux_loss_coef参数。经典的负载均衡 loss 公式大概是loss_aux alpha * N * sum(f_i * P_i)其中N是专家数量f_i是 token 实际被路由到专家 i 的比例P_i是门控分配给专家 i 的平均概率。当 token 分布和门控概率分布一致时这个 loss 会趋向于最小从而鼓励均匀路由。alpha一般设 0.01 左右太小起不到约束作用太大会干扰主任务 loss。顺带说一句最近也有人在做动态专家分配、基于强化学习的负载均衡但工业界主流还是这个 aux loss 方案简单、稳定、好调。3.4 专家并行让每个 GPU 只管少数专家MoE 的天然并行方式就是“专家并行”Expert Parallelism把不同专家放在不同 GPU 上token 通过门控路由后需要被发送到对应专家所在的设备。这和 ZeRO-3 有联系也有区别。ZeRO-3 是“所有参数都不分青红皂白地分片”而专家并行是“沿着专家维度做物理切分”。两种模式可以叠加先用专家并行把专家分配到各设备再用 ZeRO-3 把每个专家内部的参数继续分片。这种叠加是 DeepSpeed-MoE 的典型用法。实际训练中会有比较多的通信模式专家并行需要做 token 的 All-to-All 通信把 token 发给对应专家ZeRO-3 需要做参数的 all-gather 通信。这两种通信如果同时发生网络压力会很大。所以一般建议 MoE 场景下把 ZeRO-3 的层级粒度调大一些减少拆片通信次数。4. ZeRO-3 MoE 的实操配置从 YAML 到损失函数4.1 DeepSpeed 配置文件逐项解读直接给一份我常用的 ZeRO-3 MoE 训练配置ds_config.json逐项说明{ train_batch_size: 512, train_micro_batch_size_per_gpu: 4, gradient_accumulation_steps: 4, zero_optimization: { stage: 3, offload_optimizer: { device: cpu, pin_memory: true }, overlap_comm: true, contiguous_gradients: true, reduce_bucket_size: 5e8, stage3_max_live_parameters: 1e9, stage3_prefetch_bucket_size: 5e7 }, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 32 }, moe: { num_experts: 64, top_k: 2, aux_loss_coef: 0.01, ep_size: 8 } }几个关键配置的选取逻辑stage3_max_live_parameters: 控制 ZeRO-3 中“同时保持完整”的参数上限。设得太大显存峰值升高设得太小通信次数变多。一般设为单层模型参数量左右比较合适。stage3_prefetch_bucket_size: 参数预取粒度调大可以增加通信连续性但不是越大越好要结合网络带宽。ep_size: 专家并行度。如果ep_size8表示每 8 个 GPU 负责一组专家通常建议 ep_size 不超过节点内的 GPU 数避免跨节点通信。num_experts: 总专家数需要能被ep_size整除否则分配不均。top_k: 每个 token 激活的专家数常见是 1 或 2。top_k2 效果一般更好但通信和计算翻倍。4.2 损失函数的改造主任务 aux loss在代码层面MoE 的职责是接住 DeepSpeed 传过来的 aux loss。使用 DeepSpeed 的 MoE 层时它会在内部计算 aux loss并通过loss model_loss aux_loss的方式合并。伪代码大致是这样import deepspeed from deepspeed.moe.layer import MoE # 把模型中的 FFN 替换为 MoE model.transformer.ffn MoE( hidden_sizehidden_size, expertself_ffn, num_experts64, top_k2, use_residualTrue ) # 训练循环里合并损失 model_loss criterion(logits, labels) aux_loss model.get_aux_loss() loss model_loss aux_loss loss.backward()这里有个很重要的点get_aux_loss()返回的是所有 MoE 层 aux loss 的累加值。如果模型很深累加值会很大需要配合aux_loss_coef缩放。DeepSpeed 的 MoE 层内部已经乘了系数所以代码里不需要再手动乘一次。4.3 稠密模型到 MoE 模型的改造路线实际业务中很少有人从零训一个 MoE更常见的是把已经训好的稠密模型“稀疏化”成 MoE或者在训练中途切换。DeepSpeed 提供了DenseToMoE的转换工具大致思路是把每个 Transformer 层的 FFN 复制 N 份作为专家初始权重。初始化门控网络通常是线性层 softmax。保留已有训练权重只对新增的专家和门控做 warmup。这个方式的优势是可以利用稠密模型已经学到的知识MoE 层的专家起点较高收敛速度远快于从零随机初始化。业内不少开源模型就是这么干的稠密基座 MoE 扩展。4.4 训练稳定性的关键参数aux_loss_coef 和 top_k 的联动调节这里分享一个我踩过坑后的调参经验。aux_loss_coef和top_k不是独立参数它们会互相影响。当top_k1时每个 token 只走一个专家负载不均衡的风险最大因为门控一旦学偏就没有任何缓冲。此时aux_loss_coef需要设得稍高比如 0.05。当top_k2时token 至少有两个选择某一个专家过载的影响会被另一专家分摊所以aux_loss_coef可以降到 0.005~0.01。如果你观察到训练中某个专家的 token 计数持续偏高而 aux loss 很小说明 aux_loss_coef 太小门控没有收到足够的均衡压力。我自己常用的检查方法是看日志里的Expert counts分布若 max/min 超过 5 倍就要警惕了。5. 常见问题排查与性能调优实录5.1 显存 OOM先别急着开 offload训练 MoE 遇到 OOM 时很多人的第一反应是开offload_param或把train_micro_batch_size_per_gpu调小。但我建议你按这个顺序排查检查是否开了contiguous_gradients。这个配置会额外保留一块连续内存给梯度如果本身显存紧张可以关掉试试。检查stage3_max_live_parameters是否过大。如果设成所有参数总量那 ZeRO-3 几乎退化成无分片。检查 MoE 层的临时显存。专家计算时每个 token 的中间激活值可能非常惊人尤其是路由到同一专家的 token 数很多时。这时用top_k1可以大幅减少临时激活。最后才考虑 offload。因为 offload 优化器状态虽然省显存但会显著拖慢训练。5.2 训练发散loss 变成 NaN 的排查方向MoE 训练比稠密模型更容易发散主要原因有门控输出概率出现过小的对数导致指数运算不稳定。建议检查门控层是否初始化正确DeepSpeed 的 MoE 层通常会用小方差初始化。fp16 的 loss scale 不够大。MoE 的梯度幅值在不同专家间差异很大有的专家梯度极小有的专家梯度极大混合精度下容易被下溢或上溢。可以把initial_scale_power从 32 增到 48或者开启动态 loss scale。top-k 为 2 时两个专家输出了差异巨大的值加权求和后可能出现极端值。这时需要给门控的 softmax 温度加一个较大值比如 1.0 不变或略微调大平滑概率分布。5.3 吞吐上不去通信瓶颈的判断方法ZeRO-3 MoE 的通信复杂度相当高。判断瓶颈的方法很简单看训练日志中每个 step 里 GPU 空闲等待的时间占比。如果 GPU 利用率长期低于 70%八成是通信阻塞。我的调优思路如下同一节点内用 NVLink跨节点用 RDMA。如果网络只有 RoCE 或者千兆以太网ZeRO-3 会很痛苦。这个硬件条件没法突破。overlap_comm必须开启并且要配合stage3_prefetch_bucket_size调整。过小的 bucket 会导致通信碎片化过大的 bucket 又会延迟参数到达时间影响计算。专家并行ep_size尽量不超过单机 GPU 数。跨节点做专家路由的 All-to-All 通信代价很高能避免就避免。使用gradient_accumulation_steps增大 batch size让每次通信对应的计算量变大摊薄通信成本。5.4 显存测量与监控用实际数据说话不要靠猜。我建议在训练启动时每 50 步打印一次以下信息max_memory_allocated、memory_reserved、每轮 step 耗时、各专家 token 计数、当前 aux loss 值。这组数据能帮你快速定位问题。举个例子Powerful 显存分布异常时你会看到memory_reserved远大于max_memory_allocated这说明分配器缓存了很多内存但没用上。此时可以尝试设置PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:512来缓解碎片。6. 我个人的实操体会最后聊几点我的感受。ZeRO-3 和 MoE 的配合从理论上看是完美互补但工程上要花很多精力去调通信和显存之间的平衡。我见过不少团队把 MoE 专家数量堆得很大却忽略了ep_size和节点拓扑的匹配结果通信量爆炸训练速度还不如小模型。我真的建议大家先在小规模上把配置跑通再放大。具体来说先用 4~8 卡、小专家数跑通完整流程观察 aux loss 是否收敛、专家负载是否均匀、通信是否有瓶颈。确认没有任何问题后再按照节点数线性扩展。如果你直接把 1000 亿参数配置扔到 64 卡上出了问题排查成本极高。再分享一个小技巧启动训练时加上--report_every 10并且把 DeepSpeed 日志里的Micorsoft DeepSpeed Profiling打开能看到每一步中计算、通信、同步分别花了多少时间。这个数据比任何性能分析工具都直观是定位瓶颈的第一手资料。限于篇幅这里讲的主要是方法论和核心机制具体的分布式训练脚本、数据加载优化每个团队的代码结构差异很大没法给出通用模板。但你只要把 ZeRO-3 的参数分片逻辑、MoE 的稀疏激活和负载均衡机制搞明白再结合自己的实际场景跑几轮应该很快就能找到手感。