大模型训练走到今天单卡单机的时代早就过去了。但凡参数规模上到百亿、千亿甚至只是想在有限显存里塞下一个稍微像样的模型你都会撞上同一堵墙显存不够。DeepSpeed 的 ZeRO 系列就是为解决这堵墙而生的而 ZeRO-3 是其中把显存优化做到最激进的一档。与此同时MoE混合专家架构又给训练带来了另一套完全不同的挑战——它不追求把每个参数都用上而是让不同的 token 走不同的专家路径。把 ZeRO-3 和 MoE 放在一起训练是很多团队在做的组合但这两者的脾气并不完全对付。这篇内容就围绕 DeepSpeed ZeRO-3 与 MoE 训练展开把显存到底省在哪、MoE 的参数到底要不要全进显存、两者结合时哪些地方容易翻车一条条拆开讲清楚。适合已经跑过单机训练、准备上分布式、或者正在被显存和通信问题折磨的从业者参考。1. 先把显存这笔账算明白才知道 ZeRO-3 省的是什么很多人一上来就背 ZeRO 的三个 stage但真到调参的时候还是懵根本原因是没搞清楚训练时显存到底被谁吃掉了。我们先把这笔账摊开。1.1 训练显存的四大块开销在标准的混合精度训练里显存主要被四部分占据模型参数FP16/BF16 权重本身占 2 字节每参数。梯度和参数同精度也是 2 字节每参数。优化器状态如果用 AdamFP32 的 master weight、一阶动量、二阶动量加起来是 12 字节每参数444。激活值前向传播中间结果跟 batch size、序列长度、模型深度强相关这块最不可控。拿一个 10B 参数的模型举例光参数梯度优化器状态就是 221216 字节每参数也就是 160GB。这还没算激活值。单张 80GB 的卡连静态开销都放不下更别提激活了。这就是为什么必须做切分。1.2 ZeRO 三个 stage 到底切了什么ZeRO 的核心思路是既然数据并行里每张卡都存了一份完整的模型状态那这份冗余就是浪费。它把模型状态切成若干份每张卡只存一份需要的时候再通信拿回来。Stage切分对象单卡显存相对基线通信量ZeRO-1优化器状态约 1/4低ZeRO-2优化器状态 梯度约 1/8中ZeRO-3优化器状态 梯度 参数约 1/NN 为卡数高ZeRO-1 只切优化器状态省得有限但通信开销小适合显存压力不大的场景。ZeRO-2 把梯度也切了性价比通常最高是很多团队的主力选择。ZeRO-3 把参数也切了每张卡只保留自己负责的那一片参数前向和反向时按需从其他卡 gather 过来用完就释放。注意ZeRO-3 省显存最狠但代价是通信量显著上升。参数量越大、卡越多这个通信开销越明显。如果你的瓶颈是通信而不是显存盲目上 ZeRO-3 反而会更慢。1.3 为什么 ZeRO-3 的通信开销最大ZeRO-3 在每一层的前向和反向都要做一次参数的 all-gather。前向时把完整参数拼出来算算完释放反向时再拼一次算梯度。这意味着参数量被反复搬运。相比之下 ZeRO-2 只在梯度归约时通信频率低得多。所以选 stage 的逻辑很清晰显存够用就别上 ZeRO-3。只有当参数本身大到单卡放不下、或者激活值已经把显存挤爆时ZeRO-3 才是必要的。我见过不少团队一上来就 ZeRO-3结果训练速度掉了一半其实 ZeRO-2 加梯度检查点就能解决。2. MoE 的参数到底要不要全部进显存这是被问得最多的问题之一也是 MoE 训练里最容易产生误解的地方。答案要分情况说不能一句话概括。2.1 MoE 的结构决定了它的显存特性MoE 层里有一组专家expert每个 token 经过路由router后只激活其中 top-k 个专家。比如 8 个专家、top-2 激活那每个 token 实际只用到 2 个专家的计算。关键点在于激活的专家少不代表参数少。整个 MoE 层的所有专家参数加起来可能非常庞大但每次前向只有一部分参与计算。这就带来一个矛盾——计算量按激活比例算但显存要按全部参数算。2.2 稠密训练 vs 稀疏训练的显存账假设一个 MoE 模型总参数 100B但每次激活只有 20B如果按稠密方式训练所有 100B 参数都要参与梯度更新优化器状态、梯度全都要存显存按 100B 算。如果按稀疏方式训练只有被激活的专家参与更新理论上显存可以按激活部分算。但现实是绝大多数框架在训练时仍然需要把全部专家参数加载到显存里因为路由是动态的你无法预知下一个 batch 会激活哪些专家。所以MoE 训练时参数通常还是要全部进显存除非你用了专家并行expert parallelism把不同专家放到不同卡上。2.3 专家并行怎么解决这个问题专家并行的思路很直接既然专家之间是独立的那就把不同专家分到不同设备上。每张卡只负责一部分专家token 通过 all-to-all 通信被送到对应专家所在的卡上计算算完再送回来。这样一来单卡显存只需要装下自己负责的那几个专家而不是全部。代价是引入了 all-to-all 通信这个通信在跨节点时开销很大是 MoE 训练的主要瓶颈之一。提示专家并行和 ZeRO-3 可以叠加使用。ZeRO-3 切分的是每张卡内部的参数、梯度、优化器状态专家并行切分的是专家在不同卡之间的分布。两者维度不同组合起来能进一步压低单卡显存。3. ZeRO-3 和 MoE 结合时那些配置项到底怎么设理论讲完落到实操。这一节把关键配置项和它们背后的逻辑讲透避免你照抄配置却不知道在改什么。3.1 DeepSpeed 配置文件的骨架一个典型的 ZeRO-3 MoE 配置大概长这样{ train_batch_size: 64, gradient_accumulation_steps: 4, fp16: { enabled: true }, zero_optimization: { stage: 3, offload_optimizer: { device: cpu, pin_memory: true }, offload_param: { device: cpu, pin_memory: true }, overlap_comm: true, contiguous_gradients: true, stage3_gather_16bit_weights_on_model_save: true }, moe: { enabled: true, ep_size: 8, use_tutel: false } }这里几个参数值得单独说。3.2 offload 到底该不该开offload_optimizer和offload_param把优化器状态和参数卸载到 CPU 内存。开了之后显存压力骤降但训练速度会明显变慢因为 CPU 和 GPU 之间的 PCIe 带宽远低于显存带宽。我的经验是显存实在不够再开 offload。如果开了 offload 之后训练慢到无法接受那说明你的并行策略有问题应该优先考虑加卡或者调整专家并行度而不是硬扛 offload。offload 是最后的救命稻草不是常规配置。3.3 overlap_comm 和 contiguous_gradients 的作用overlap_comm让通信和计算重叠进行能有效掩盖一部分 all-gather 的延迟。contiguous_gradients把梯度整理成连续内存减少内存碎片对 ZeRO-3 这种频繁分配释放的场景很有帮助。这两个参数在 ZeRO-3 下基本是默认要开的除非你遇到特定的稳定性问题。实测下来开了 overlap_comm 之后ZeRO-3 的吞吐能提升 10% 到 20%具体取决于模型结构和网络拓扑。3.4 ep_size 怎么定ep_size是专家并行度表示专家被切分到多少张卡上。它的取值要和总卡数、专家数量配合。假设你有 64 张卡、64 个专家那 ep_size 设成 8 意味着每 8 张卡一组每组负责 8 个专家。设成 64 就是每个专家独占一张卡。ep_size 越大单卡专家越少显存越省但 all-to-all 通信范围越大。注意ep_size 必须能整除总卡数也必须能整除专家总数否则会报错。配置前先算清楚这两个数。4. 训练跑起来之后那些让人抓狂的报错怎么排配置写对了不代表就能顺利跑。MoE ZeRO-3 的组合在实战里有一批高频报错这一节按排查链路来讲。4.1 安装 deepspeed 就报错先别怀疑代码很多人第一步就卡在pip install deepspeed。这个包编译依赖比较重常见报错集中在几个方向CUDA 版本不匹配deepspeed 编译时会检测本机 CUDA如果 PyTorch 编译用的 CUDA 版本和系统 CUDA 不一致就会报错。解决办法是先确认torch.version.cuda再装对应版本的 deepspeed。缺少编译工具需要 gcc、g 等。报错信息里如果有gcc: command not found装一下 build-essential 就行。JIT 编译超时deepspeed 有些算子会在首次运行时 JIT 编译如果环境变量没设好会卡住。可以设置DS_BUILD_OPS0先跳过自定义算子编译跑通流程再补。我一般建议先用pip install deepspeed --no-build-isolation试能省掉不少隔离环境带来的坑。4.2 专家负载不均衡导致的训练崩溃MoE 最经典的问题就是路由塌缩——所有 token 都涌向少数几个专家其他专家饿死。表现是 loss 突然飙升或者直接 NaN。排查思路是这样的先打印每个专家被选中的频率。如果某个专家占比超过 50%基本可以确认是负载不均衡。检查有没有加负载均衡损失load balancing loss。这是 MoE 训练的标准配置通常是一个辅助 loss鼓励 token 均匀分配到各专家。检查 router 的初始化。router 权重初始化太极端会导致早期就塌缩。负载均衡损失的代码逻辑大致是统计每个专家被选中的比例和理想均匀分布做对比算一个辅助损失加进总 loss。系数一般设在 0.01 量级太大影响主任务太小起不到作用。4.3 ZeRO-3 下的参数 gather 超时ZeRO-3 在保存 checkpoint 或者做某些操作时需要把所有分片的参数 gather 回来。如果模型很大、卡很多这个 gather 可能超时。stage3_gather_16bit_weights_on_model_save这个参数就是控制保存时是否 gather 的。如果保存经常超时可以设成 false保存分片权重后续再合并。另外通信超时时间可以通过NCCL_TIMEOUT环境变量调大。4.4 显存明明够却 OOM 的诡异情况有时候显存监控显示还有余量但就是 OOM。这种情况在 ZeRO-3 MoE 下通常是内存碎片导致的。contiguous_gradients能缓解一部分但更彻底的办法是调整round_robin_gradients让梯度分配更均匀。还有一种可能是激活值峰值。MoE 的 all-to-all 通信会产生额外的缓冲区这部分显存容易被忽略。可以开梯度检查点gradient checkpointing把激活值压下来用计算换显存。5. 让训练真正跑快的几个调优方向跑通只是第一步跑快才是本事。这一节讲几个实测有效的调优方向。5.1 通信和计算的重叠程度ZeRO-3 的性能瓶颈几乎全在通信上。overlap_comm能重叠一部分但重叠效果取决于模型结构。如果某一层的计算量太小通信还没藏完计算就结束了那这层就是瓶颈。一个实用的做法是调整 bucket size。DeepSpeed 会把参数按 bucket 分组做 all-gatherbucket 太小通信次数多太大则重叠粒度粗。默认值通常够用但在 MoE 场景下可以适当调大因为专家层的参数块比较大。5.2 专家并行的拓扑选择all-to-all 通信对网络拓扑很敏感。如果专家并行组跨了节点通信要走慢速网络开销会暴涨。所以尽量让专家并行组落在同一个节点内用 NVLink 或高速互联。假设一个节点 8 张卡那 ep_size 设成 8 就能保证专家并行不跨节点。如果非要设成 16就要接受跨节点通信的代价。这个取舍在配置阶段就要想清楚。5.3 batch size 和梯度累积的配合MoE 训练对 batch size 比较敏感因为 batch 太小会导致路由统计不稳定负载均衡损失波动大。但 batch 太大又吃显存。折中方案是用梯度累积。train_batch_size设成实际想要的大 batchgradient_accumulation_steps设成累积步数两者相除就是单步的 micro batch。这样既保证了路由统计的稳定性又控制了单步显存。提示梯度累积步数增加会拉长单次迭代时间但不会增加显存。如果显存是瓶颈而时间不是这是最划算的调法。5.4 混合精度的选择BF16 相比 FP16 动态范围更大不容易溢出在 MoE 训练里更稳。如果硬件支持 BF16优先用它。FP16 需要配合 loss scaling而 MoE 的辅助损失会让 loss 尺度变化更复杂loss scaling 调起来更麻烦。实测下来同样的配置换成 BF16训练稳定性提升明显尤其是训练后期 loss 已经很小的时候FP16 更容易出现梯度下溢。6. 几个容易被忽略的实战细节最后聊几个文档里不太写、但实际会踩的细节。6.1 checkpoint 的兼容性ZeRO-3 保存的 checkpoint 是分片的和普通 checkpoint 格式不一样。如果你中途想换并行策略比如从 ZeRO-3 换到 ZeRO-2checkpoint 需要先合并转换。这个转换过程对 MoE 模型更麻烦因为专家参数还要考虑并行度的映射。建议在训练早期就把并行策略定下来别中途大改。如果非要改预留出转换和验证的时间。6.2 学习率 warmup 要更保守MoE 模型因为路由的存在早期训练更不稳定。warmup 步数建议比同等规模的稠密模型更长一些让 router 有时间找到合理的分配。我一般会把 warmup 比例从 1% 提到 2% 到 3%。6.3 监控指标要盯紧专家利用率除了常规的 loss、学习率、吞吐MoE 训练一定要盯专家利用率。如果发现某些专家长期不被激活要么是负载均衡没做好要么是专家数量设多了。专家数量不是越多越好要和数据规模、任务复杂度匹配。6.4 别忽视数据质量对路由的影响MoE 的路由是根据 token 表示来决定的如果训练数据分布很偏路由也会偏。数据里如果某类样本特别多对应的专家就会被过度激活。所以数据配比在 MoE 训练里比稠密模型更关键预处理阶段就要把分布调匀。我个人在几次 MoE 项目里最大的体会是ZeRO-3 解决的是显存问题MoE 解决的是容量和计算效率问题但两者叠加之后真正的难点从能不能跑变成了跑得稳不稳、快不快。配置只是起点负载均衡、通信拓扑、数据分布这些软性的东西才是决定训练成败的关键。如果一开始就冲着省显存去堆配置很容易在稳定性上栽跟头。先把小规模跑通、把专家利用率盯住再逐步放大这条路走得最稳。