半夜盯着nvidia-smi发呆的场景跑大模型的人多少都经历过一块 80G 的 A100加载 7B 参数做全量微调优化器状态一开直接 OOM换成 8 卡分布式训练模型反而跑起来了但通信开销把训练效率拖下去一大截。这正是 DeepSpeed ZeRO-3 与 MoE 训练要解决的两个核心问题。简单说ZeRO-3 负责把模型参数、梯度、优化器状态全部拆开摊到多张卡上MoE 则让模型在参数总量膨胀的同时不把计算量同步抬上天。把这两件事放在同一篇文章里讲是因为今天训练大规模稀疏模型时它们几乎总是成对出现。这篇文章会从显存账本讲起把 ZeRO-3 的切分机制、MoE 的路由原理、两者怎么配合落地以及我在实际训练中踩过的坑一次说清楚。适合正在做多卡分布式训练、或者准备上手稀疏大模型的同学参考。不涉及数学推导全部用算账和实操视角来讲。1. 显存容量矛盾为什么单卡永远装不下大模型1.1 一个 7B 模型实际要吃掉多少显存很多人对显存占用的理解还停留在“模型权重多大就要多大显存”实际上全量训练时的显存需求比这大得多。以常见的 7B 参数模型为例现在主流做法是混合精度训练模型权重复制一份 BF16 用作前向反向计算同时保留一份 FP32 的“主权重”优化器用 Adam还得维护一阶动量 m 和二阶动量 v。逐项算下来模型权重BF162 字节 × 7B 14GB梯度BF162 字节 × 7B 14GBFP32 主权重4 字节 × 7B 28GBAdam 的 m 和 v各 4 字节 × 7B合计 56GB四项加起来单卡全量训练 7B 模型大约需要 112GB 显存。这就是为什么单张 80G 的 A100 在真正全量微调时一定会 OOM。顺着这个公式继续算13B 模型需要超过 200GB68B 模型超过 1TB。显存墙不是靠换一张更大的卡能解决的——你买到 1000GB 显存的服务器之前必须先学会怎么把显存用高效。我见过不少同学最先想到的解法是梯度累积、混合精度、或者干脆用 LoRA 只训练低秩 adapter。这些方案都能解决问题但各有边界LoRA 只适合微调场景预训练和全量微调仍然要面对全参数优化梯度累积只是把一次大步拆成小步并没有减少单步所需的总显存。这时候就需要分布式训练里真正管用的手法——把状态切出去。1.2 从 ZeRO-1 到 ZeRO-3 的演进逻辑ZeRO 的核心思想其实特别朴素一个分布式训练系统里有三种主要状态——模型参数、梯度、优化器状态。传统 Data ParallelDDP模式下每张卡都保存一份完整副本数据并行只切数据不切模型显存浪费严重。ZeRO-1 只切优化器状态每张卡只负责一部分参数的 m、v 和主权重通过通信让每张卡轮流得到全部状态完成更新显存瞬间下降一半以上。ZeRO-2 在 ZeRO-1 基础上把梯度也切成多份反向传播后每张卡只保存自己负责的那一段梯度。ZeRO-3 更进一步把模型参数本身也切成碎片前向计算需要哪层参数就通过通信把那层参数临时“拼”出来用完立刻丢弃。从 ZeRO-1 到 ZeRO-3显存占用一路走低但通信量一直在涨。三者对比下来ZeRO 阶段切分内容单卡显存占用7B 模型通信开销变化DDP不切约 112GB基准ZeRO-1优化器状态约 70GB略增ZeRO-2梯度 优化器状态约 56GB明显增加ZeRO-3参数 梯度 优化器状态约 28GB大幅增加从数字上能看出ZeRO-3 对显存极度友好代价是通信量成为主要瓶颈。实际训练中通常用通信重叠、梯度分桶、参数预取这些技巧把这些代价“藏”起来而这些技巧的具体效果和你的模型结构、卡间网络都有直接关系。2. ZeRO-3 机制拆解参数分片、动态全收集、梯度规约2.1 切分对象与切分粒度ZeRO-3 切的最小组件不是整个模型层而是每个参数张量。一个 7B 模型有几百个参数张量每个张量会被按 rank 数量切成 N 份每个 rank 只持有原本那个张量的 1/N。这种细颗粒度的好处是负载均衡模型里不同的层参数量差异不大但 embedding 层和 classifier 头经常特别大按张量切分更容易保持各卡之间的显存压力均匀。这里有一个容易误解的点ZeRO-3 并不会把所有参数均匀塞到每个 rank 上之后就完事。它需要在 forward 过程中按需触发 all-gather 操作把当前计算层所需的参数重新组合成完整版本。某个参数张量如果位于第 12 层那么在第 12 层计算之前所有 rank 都要通信一次把自己手里的 1/N 碎片拿出来拼成一个完整参数供这一层的矩阵乘法使用。计算完之后完整的副本又会被抛弃各 rank 重新只保留自己的那 1/N 碎片。生活化理解这就像一组人合力读一本分成很多册的辞典谁先查某个词条就把包含这个词条的那一册借出来其他册还给书架。只借当下要用的那一册而不是把所有册都摊在桌上。代价是借还书通信频繁好处是桌面显存上永远只有一本书占地方。2.2 forward 和 backward 中的两类关键通信ZeRO-3 在 forward 阶段的主要通信模式是 all-gather。每个参数张量在被使用前都要做一次全收集让每个 rank 都拿到完整参数参与计算。到 backward 阶段每个 rank 对自己拿到的完整参数算出一份梯度但这部分梯度是“重复的、不完整的”——毕竟每个 rank 只负责该张量的一部分分片的更新。因此反向传播结束时要做一个 reduce-scatter把各 rank 手上重复的梯度按分片位置求和后切回自己的那一段最终每张卡只保存它负责的那一段梯度。这就产生了一个实操经验如果模型非常大层数很多all-gather 和 reduce-scatter 的次数会呈线性增长。默认情况下每次参数收集都是同步的也就是说 GPU 在等通信返回时经常处于空闲状态。解决办法是开通信重叠overlap_comm和参数预取prefetch_bucket_size在计算上一层的间隙提前发起下一层参数的收集请求让通信和计算并行。实测下来模型在 64 卡规模下开不开 overlap吞吐差距经常在 20% 以上。另外contiguous_gradients这个配置也很重要。它会把多个小梯度的 reduce-scatter 合并成一次大块通信避免零碎的通信请求把通信链路打满。如果库里默认是关闭状态我建议在 stage 3 下显式打开。2.3 一份可直接使用的 ZeRO-3 配置示例实际使用 DeepSpeed 时stage 3 的配置一般在 JSON 里完成。下面这份配置我在训练 6B 稠密模型时验证过可以直接抄{ train_batch_size: 32, gradient_accumulation_steps: 4, optimizer: { type: AdamW, params: { lr: 3e-4, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, zero_optimization: { stage: 3, contiguous_gradients: true, reduce_bucket_size: 5e7, reduce_scatter: true, overlap_comm: true, prefetch_bucket_size: 5e7, offload_optimizer: { device: cpu, pin_memory: true } }, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 }, steps_per_print: 100 }重点解释几个参数。reduce_bucket_size和prefetch_bucket_size控制每次通信携带的梯度/参数字节数单位是字节数还是元素数要看你用的 DeepSpeed 版本默认 5e7 左右是安全值。调太大单个 bucket 可能触及通信缓冲上限调太小通信次数太多反而慢。offload_optimizer把 Adam 状态放到 CPU 内存这能让单卡显存再降一截但 CPU 内存占用直接上涨且优化器步骤变慢。建议先在纯 GPU 模式下跑通确认性能瓶颈后再决定是否开 offload。注意如果开了offload_optimizer又把pin_memory设为 true主机会锁定一部分内存用于 DMA 传输这需要你在机器上预留足够物理内存否则系统会开始 swap速度断崖式下跌。启动命令也很简单deepspeed --num_gpus 8 train.py --deepspeed ds_config.json在train.py里只需要完成模型初始化后调用deepspeed.initialize框架会自动接管参数的切分、梯度通信和优化器步骤。第一次跑通 stage 3 的常见 bug 是模型用了自己实现的参数共享比如 word embedding 和 output layer 绑权值这类共享参数在一个 forward 里被多个层引用ZeRO-3 会对同一参数反复做 all-gather导致性能极差或者显存反而更高。遇到这种情况建议要么拆开共享要么用 DeepSpeed 提供的特殊处理接口而不是自己在 forward 里直接缓存param.data。3. MoE 稀疏激活参数翻倍算力不翻倍3.1 专家路由的直觉理解MoEMixture of Experts架构最近两年热度极高从 Switch Transformer、Mixtral 到 DeepSeek-MoE核心思路都一样把 Transformer 里的 Feed-Forward Network 换成一个 Router 加多个专家 FFN。每个 token 并不经过全部专家而是由 Router 给每个专家打分选出得分最高的 top-k 个专家做计算。这里的关键收益和代价需要分清楚。如果模型有 64 个专家但每个 token 只激活 2 个专家那么推理时单个 token 的计算量近似于一个 2 专家规模的 FFN而不是 64 个专家规模的 FFN。这就是“参数总量可以很大但计算量相对可控”的来历。不过要注意一个容易混淆的点训练和推理对显存的要求不一样。推理时模型权重已经被固定加载到显存之后只有被激活的专家会进入计算流程但训练时要对每个专家计算梯度、更新参数即使一个 token 只走两个专家所有 64 个专家的权重依然需要被存储和更新。所以训练 MoE 时模型参数“进显存”这件事是绕不开的差别只在于进一张卡的显存还是分散在多张卡的显存里。自然语言里“moe架构要全部参数进显存吗”这个问题完整回答是权重必须存在于显存体系里但通过 ZeRO-3 或专家并行参数可以分摊到多张卡上运行期按需收集不需要任何单卡装下全部参数。3.2 负载均衡MoE 训练的第一个拦路虎MoE 训练最容易出问题的地方是 Router 崩溃。所谓“崩溃”不是指程序报错而是 Router 快速偏好少数几个专家导致这几个专家被大量 token 挤爆其他专家几乎闲置。专家之间负载严重不均衡时模型退化得非常厉害相当于大部分计算和参数都没有被有效利用。解决方案在工程上已经比较成熟引入负载均衡辅助损失也就是 aux loss。基本思路是统计每个专家分配到的 token 比例 f_i以及 Router 对所有 token 分配给专家 i 的平均概率 P_i然后把 N × f_i × P_i 的和作为惩罚项加进总损失。当 token 分布均匀时这个值最小某个专家聚集大量 token 时惩罚变大训练会自动把 Router 往平均分配方向推。DeepSpeed 的 MoE 配置里对应的就是load_balance_loss_weight。这个权重的设置非常需要拿捏开大了 Router 会变得过于“和稀泥”每个专家都差不多稀疏激活退化成均匀计算能力上不去开小了负载均衡失效一会儿又出现专家过载。我自己的经验是从 0.01 起步观察训练日志里每个 step 的专家分配统计再微调。除了 aux loss实际训练里还要注意noisy_gate_policy。常见选项是Jitter给 Router 打分加上一点带噪声的扰动让 token 分布更随机一定程度上防止 Router 在训练初期就锁定在偏科状态。这个机制可以理解为给专家筛选加上一点“随机选择权”避免所有 token 都挤向同一个最初得分高的专家。3.3 DeepSpeed-MoE 的并行方式专家并行和 All-to-All 通信MoE 的专家层在 DeepSpeed 中通常采用专家并行Expert Parallelism。做法是把 N 个专家分配到不同 GPU 上每个 GPU 只保存一部分专家。某个 token 被 Router 选到某个专家时它所在的 GPU 需要把 token 的 hidden state 发给目标专家所在的 GPU计算完成后再把结果传回来。这个过程就是 all-to-all 通信。这也是 MoE 训练和纯稠密模型训练最大的差别稠密模型的通信主要是梯度同步和数据并行通信模式比较规律MoE 则会在每次前向和反向中穿插多次 all-to-all。这些通信请求是随 token 分配动态变化的很难静态优化因此网络带宽会成为 MoE 训练的主要瓶颈之一。如果你的集群是跨节点多机环境我的建议是尽量把专家并行范围控制在单个节点内。也就是说ep_size不要跨过多的机器因为节点间网络比如以太网的带宽通常低于节点内 NVLink而 all-to-all 对带宽极其敏感。跨节点之后训练吞吐的下降往往不是来自算力而是来自这些动态通信请求的排队时延。4. ZeRO-3 与 MoE 组合实战配置、通信与调参4.1 两种省显存机制如何叠加把 ZeRO-3 和 MoE 放在一起训练实际上是在同一份显存预算里做了两套分配策略对于模型中的稠密层使用 ZeRO-3 进行全参数切分对于 MoE 层使用专家并行把不同专家放在不同卡上同时利用 ZeRO-3 把专家层内部的参数再做分片处理。这样一张卡上既不会保存全部稠密层参数也不会保存全部专家参数任意时刻显存里主要保留的是当前计算层需要的那份完整参数。有人可能会问既然 MoE 层已经用专家并行分散了为什么还要 ZeRO-3 再切一遍原因是 MoE 层之外还有大量注意力层、层归一化、embedding 层这些稠密部分不会因为 MoE 而自动变小。如果只靠专家并行不动这些稠密层显存中占比很大的注意力权重依然每张卡各有一份仍然装不下很大的模型。所以实际 DeepSpeed-MoE 的推荐组合是稠密层走 ZeRO-3 分片专家层走 expert parallelism两者共用一套通信规划机制。关于训练时的显存占用可以借助nvidia-smi dmon或 DeepSpeed 自带的日志观察每张卡的显存曲线。如果发现某一张卡明显比其他卡多占用几个 GB很可能是专家并行分配时的ep_size设置不当或者某一层专家被 Router 频繁调用但物理分配不平衡导致的缓存压力。4.2 一份 DeepSpeed-MoE 的完整配置示例下面是 DeepSpeed-MoE 训练中比较典型的一份配置文件注释部分是我在实际项目里反复调出来的经验值{ train_batch_size: 32, gradient_accumulation_steps: 4, zero_optimization: { stage: 3, overlap_comm: true, contiguous_gradients: true, reduce_bucket_size: 5e7, prefetch_bucket_size: 5e7, offload_optimizer: { device: cpu, pin_memory: true } }, moe: { enable: true, num_experts: [8], top_k: 2, ep_size: 2, min_capacity: 4, drop_tokens: false, noisy_gate_policy: Jitter, load_balance_loss_weight: 0.01 }, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 } }逐个解释 MoE 相关参数num_experts: [8]这是一个数组表示每一层 MoE 层的专家数量。如果模型有 12 层要替换成 MoE且有 8 个专家的层和 16 个专家的层就写成[8, 8, 8, 8, 16, 16, ...]这样的序列。不能简单写一个数字DeepSpeed 会按数组长度依次对应每一层。top_k: 2每个 token 路由到的专家数量。top_k2 是常见选择效果比 1 稳计算开销也只多一个专家。ep_size: 2专家并行的卡数。这里设为 2 表示每个专家被复制并分配到 2 个 GPU 上其实是把一个专家按卡再分片适合检测连通性和通信开销。生产环境建议先按单节点卡数设置不要盲目跨节点。min_capacity: 4每个专家每个 batch 的最小容量。容量就是专家最多能处理的 token 数容量太小时很多 token 溢出太大会让每个专家计算大量本不该由它处理的 token。drop_tokens: false当 token 超出专家容量时的策略。false 表示不丢弃而是重新路由到其他专家效率更稳但可能造成重复计算。真正大规模训练中我建议保持 false避免数据丢失。load_balance_loss_weight: 0.01这个权重前面讲过是负载均衡的惩罚系数。4.3 训练流程中的关键实操建议我强烈建议第一次接触 DeepSpeed-MoE 的团队不要一上来就在大模型上完整跑。先构造一个“麻雀式”的小模型两层 Transformer、每层 4 个专家、top_k 1数据用随机生成的假数据。这样做的目的很简单——把通信链路、Router 初始化、all-to-all 的执行路径全部跑通确认没有在分布式环境里埋雷再扩大规模。小模型跑通之后有几个点值得专门记录第一训练刚开始 200 步内的 loss 曲线波动是非常正常的。MoE 的 Router 在初期相当于随机打分专家之间负载不均衡是常态。不要因为这一步的 loss 高就立刻判定模型废了先观察 200 到 500 步看 aux loss 是否进入下降趋势。第二drop_tokens和min_capacity的配合逻辑。如果drop_tokens设为 false那么当某个专家容量超限时 token 会被重新分配但这也意味着你实际上创建的 token-router 映射被打破了会有额外通信。训练中遇到吞吐异常下降时优先查这两项。第三不要忽略noisy_gate_policy的影响。在训练初期噪声可以让 Router 保持“探索欲”但训练后期如果仍然保持强噪声会让已收敛的 Router 不断被扰动。部分框架允许动态调整或衰减噪声如果平台不支持建议在固定迭代步数后手动切换策略。5. 高频问题与排查实录5.1 常见问题速查表我在多个训练任务里反复遇到下面这些问题整理成一张速查表问题现象根因分析排查与解决开了 stage 3 仍然 OOM通信缓冲区、梯度累积缓冲、优化器状态没有真正被分片覆盖检查 reduce_bucket_size 是否过大确认 optimizer 是否走 DeepSpeed 接管考虑开启 offload_optimizer训练吞吐大幅下降通信次数过多all-gather/ reduce-scatter 没有被隐藏打开 overlap_comm、contiguous_gradients确认 prefetch_bucket_size 设置合理检查模型是否在 forward 中重复使用同一参数部分专家显存占用异常高专家并行分配不均或 Router 严重偏科检查 ep_size 设置用训练日志观察专家 token 分配比例调大 load_balance_loss_weight训练loss震荡不收敛Router 对负载分配失去控制噪声过大降低学习率调小 noisy_gate_policy 的扰动幅度检查 aux loss 权重是否过高开 offload 后系统卡死/极慢CPU 内存不足、swap 被触发关闭 pin_memory降低 offload 比例升级物理内存改用 NVMe offloadall-to-all 通信消耗时间过长专家分布在跨节点网络带宽不足缩小 ep_size 到单节点内优先用 NVLink 连接专家通信减少 top_k 或专家层数量5.2 几类典型问题的现场记录以 stage 3 下仍然 OOM 举例。很多人以为开了 ZeRO-3 就一定不会显存爆炸实际上 ZeRO-3 只负责把参数、梯度和优化器状态切出去并不负责限制前向计算时的激活值。激活值显存在大 batch 下照样可以吃掉几十 GB。排查方式是把train_batch_size和gradient_accumulation_steps组合调整先设train_batch_size1观察显存占用如果显存依然打在 80G 边缘明显是激活值或者通信缓冲的问题如果降下来了就是批量大小相关的问题。另一个常见坑是模型 executor 或 loader 在 ZeRO-3 环境下被重复执行。比如某些库在model.generate时会对权重自动做缓存这个缓存可能完全绕开 DeepSpeed 的分片管理导致显存立刻上升。遇到这种情况可以检查模型对象的子模块是否在每次 forward 里通过.to(device)移动参数这种“手动搬运”会破坏 ZeRO-3 的动态参数收集逻辑。MoE 训练中比较隐蔽的问题是 token 丢包。min_capacity设置过小时大量 token 超出专家容量drop_tokensfalse的语义会把这些 token 重新分配但有些版本的框架在 re-route 时的随机性可能导致同一个 token 被多个专家重复计算。这类问题不会让程序报错但会让 loss 曲线出现一些奇怪的“跳跃”最终模型精度也偏低。建议在训练脚本里定期打印每个专家的 token 计数用统计数据而不是手感来判断是否出现丢包。还有一点是关于负载均衡损失的监控。不要只盯总 loss一定要单独把 aux loss 从训练日志里拉出来看曲线。如果 aux loss 从第一步起就一直非常低可能并不是因为路由均衡而是因为权重太低没有被学习到。反过来如果 aux loss 一直降不下来就要怀疑 Router 本身的结构或初始化有问题。5.3 组合场景下的实测心得把 ZeRO-3 和 MoE 一起上之后最直观的感受是“显存压力瞬间释放通信压力集中爆发”。我在一次实验里用 8 卡 H800 训练 16B MoE 模型开 ZeRO-3 之前连参数都加载不进去开完之后显存占用非常健康但训练吞吐一上来就比同规模稠密模型低了不少。用性能剖析工具打点之后发现绝大多数时间花在了 all-to-all 通信的等待上而不是计算本身。这说明显存问题确实被解决了但接下来要优化的是通信拓扑。结合以上经验再给一个比较实用的建议MoE 模型训练时的ep_size选择一定要同时考虑显存上限和通信拓扑。如果显存非常紧张可以适当增大ep_size让每个专家分到更多卡的显存但要注意跨节点通信成本如果显存还有余量优先降低ep_size减少通信频率你会看到吞吐立刻改善。最后再说说调参顺序。训练 MoE 的过程中变量太多一次只改一个参数是黄金法则。我通常的顺序是先固定num_experts和top_k然后调load_balance_loss_weight观察 aux loss 和专家负载曲线再调min_capacity和drop_tokens处理 token 溢出最后才碰noisy_gate_policy这类影响收敛动态的参数。每次都记录日志对比不要凭感觉。我对 ZeRO-3 和 MoE 最深的体会是这两个机制都极其依赖训练时的可观测性日志里多打印专家分布和通信耗时比任何理论分析都管用。