近十年大模型训练里的那只显存怪兽ZeRO 是怎么一步步驯服的只要你想训练稍微像样一点的大模型早晚会撞上 ZeRO 这个词。我在过去几年里从单卡跑小模型开始一路做到在几十卡集群上训练千亿参数级别的模型ZeRO 的每个 stage 都亲手配过、也亲手踩过坑。ZeRO 这个名字正式出现已经是这个十年后半程的事了但它背后那条如何让更多设备协同训练更大模型的主线其实走了十年以上。这篇东西不打算复述论文公式而是用我自己的理解把这条显存不够怎么办的演进路线讲清楚为什么早期方案总在搬数据ZeRO 为什么选择分数据以及当你自己动手配置时到底该怎么选 stage、该在哪里设置 offload。文章会照顾从没用过分布式训练的新手也会给已经在用的人一些排错参考。1. 十年前的显存困局模型状态与残差状态是怎么把显存吃光的1.1 显存里其实躺着两笔完全不同的账很多人一开始以为训练时显存主要被模型权重占着等到真正跑起来才发现权重只是很小一部分。过去这些年里每一代显存扩容都会被大模型瞬间吃掉根本原因在于训练一个模型显存里同时要装下两笔账。第一笔叫模型状态包含三类东西模型参数、梯度、优化器状态。优化器状态在 Adam 时代尤其夸张——Adam 除了要保留一份高精度的 FP32 主权重副本还要为每个参数维护两个额外的状态一阶动量 m 和二阶动量 v这三个东西都是 4 字节。再加上参与前向和反向传播的 FP16 参数副本2 字节和 FP16 梯度2 字节一个参数平均要吃 16 字节。为了让你有体感我拿一个 70 亿参数的模型算一笔账70 亿 × 16 字节光模型状态就要 112GB。这个数已经超过单张 80GB 显存的数据中心显卡更别提普通 24GB 的消费级卡了。而这还只是第一笔账。第二笔叫残差状态包括前向传播过程中每一层需要保存的激活值、临时缓冲区以及碎片化导致的不可用显存。这笔账在长序列和深网络下同样凶悍——一个 seq_len 为 2048 的 7B 模型激活值很可能再吃掉十几个 GB。过去十年里显存不够的报错几乎都来自这两笔账的叠加而早期方案只盯着其中一部分做文章。1.2 数据并行撑了半程但显存永远不会被它摊薄2010 年代中后期大家普遍的做法是数据并行把数据切成多份分给多张卡每张卡各自算一份梯度然后再把梯度做一次全局同步all-reduce。这个方案解决的是算得太慢的问题不是装不下的问题。因为每张卡手里都握着一份完整的模型副本模型状态一行都不会减少。8 卡跑 7B 模型每张卡照样要装 112GB 模型状态该 OOM 还是 OOM。数据并行只能把整体吞吐提上去显存天花板纹丝不动。这就是当年很多人撞上的那堵墙想训练更大的模型但加卡根本解决不了内存容量问题。后来社区也尝试过把模型的层切开分到不同的卡上也就是模型并行。这种方式能实质性降低单卡显存但手工切分网络、处理跨卡通信非常痛苦每个层面的切分策略都要重新设计。它像是一种物理分家行得通但工程复杂度极高。所以当 ZeRO 提出与其搬数据不如分状态的时候很多人都觉得思路一下子被打开了。2. ZeRO 分而治之三个 stage 的显存数学账怎么算2.1 核心思路把每张卡存一份改成每张卡存一份的 N 分之一ZeRO 的最聪明之处是承认每张卡确实都需要看到完整的模型但不需要永久持有完整模型。它把模型状态切碎均匀分布到 N 张卡上训练过程中谁需要某块数据就临时找持有者要用完即还。这就像团队里有一本公用字典以前是每人复印一整本数据并行现在改为拆成 N 册分给大家保管查哪个词就去对应的人手里借阅。物理上没有任何数据被复制或丢失只是存放位置变了。因为这个特性ZeRO 在数学上和原来的训练完全等价收敛曲线不会因为切分而改变——这一点后面排错时非常重要。2.2 三个 stage 的显存账本对比ZeRO 把切分推进到了不同程度对应三个 stageStage切分范围每张卡模型状态占用K 为参数量字节数7B 模型、8 卡时每卡理论占用基线普通数据并行不切分2K参数 2K梯度 12K优化器状态 16K112GBStage 1仅优化器状态2K 2K 12K/N约 38.5GBStage 2优化器状态 梯度2K 2K/N 12K/N约 26.25GBStage 3参数 梯度 优化器状态16K/N14GB注意这个表里的 N 是 GPU 数量。以 8 卡跑 7B 为例普通数据并行每卡 112GB直接 OOMStage 1 降到 38.5GBStage 2 降到 26.25GBStage 3 降到 14GB。如果再加上激活值等残差状态Stage 3 是唯一能让 80GB 显存卡从容跑完训练甚至还能加大 batch的方案。2.3 通信模式变了reduce-scatter 与 all-gather 的代价切分不是免费的代价藏在通信里。普通数据并行反向传播结束后需要做一次 all-reduce把每张卡的梯度汇总后广播给所有人。ZeRO Stage 2 把这一步拆成了两个更轻量的操作先做 reduce-scatter把梯度按分片规则求和并分发到对应卡上这样每张卡拿到的是自己负责那一份的完整梯度更新完成后不需要再广播全量优化器状态因为状态本来就分在各卡手里。Stage 3 多了一层更重的通信每一层前向和反向之前都要先 all-gather 取回这一层的完整参数。也就是说参数每次都要在节点间流动一次。层数越多、模型越大这个流动量就越可观。所以 Stage 3 虽然显存最省通信压力也是最大的。高性能卡间互联比如单机 8 卡走高速互联总线还能压住这个开销跨机器走普通以太网时Stage 3 很可能比 Stage 2 慢一大截——这不是犹豫要不要省显存的问题而是要先搞清楚你的网络到底扛不扛得住。3. 显存还不够怎么办CPU Offload 与 NVMe Offload 的带宽账本3.1 Offload 的动机把 GPU 塞不下的东西挪到隔壁房间Stage 3 已经把模型状态分到了理论上最省的位置但如果模型大到连分的份额都塞不进单卡或者单卡数量不够分就得动用最后的手段把它挪出 GPU。Offload 的思路很简单GPU 显存是稀缺资源而 CPU 内存和 NVMe 磁盘容量大得多。把不常用的数据通常先是优化器状态再是参数和梯度放到 CPU 内存甚至磁盘上GPU 只在需要的瞬间把它取回来。听起来像显存不够内存来凑实际做起来要先算一笔带宽账因为跨设备搬数据的速度会直接卡住训练节奏。3.2 三层存储的带宽量级为什么顺序这么重要我整理了一张经常贴给组里新人的带宽对照表数值是量级概念不必精确到小数点存储层带宽量级延迟量级典型容量GPU 显存HBM1 ~ 2 TB/s纳秒级数十 GBCPU 内存DRAM50 ~ 100 GB/s百纳秒级数百 GB ~ 数 TBPCIe 总线16 ~ 32 GB/s微秒级与设备相关NVMe 磁盘2 ~ 4 GB/s数十微秒级数 TB训练的一个 step 里参数和梯度要被反复读写多次。如果一层的数据需要从 NVMe 取回来算完再送回去每步的 IO 等待很容易从毫秒级膨胀到秒级。所以社区给出的默认顺序是先 Offload 优化器状态到 CPU 内存再 Offload 梯度和参数最后才考虑 NVMe。一层层往下退每一层都是拿训练速度换模型规模。3.3 什么时候该用哪种 Offload根据我的实际经验决策依据可以这么概括只需要把显存压到刚好能跑优先开 Stage 2 CPU Offload 优化器状态。优化器状态是最大的一块12 字节/参数搬到 CPU 后每卡显存立刻降一大截而通信压力增加有限。模型再大一些Stage 3 也塞不下把参数也做 CPU Offload但这意味着每一层前向/反向前都要通过 PCIe 取参数吞吐会明显下降。建议同时调大 batch 和梯度累积步数用计算密度把搬运开销摊薄。已经到了单机 CPU 内存都紧张的程度才考虑 NVMe Offload。这种配置基本是在能跑和跑得快之间彻底倒向了前者适合做微调或推理不太适合长期大规模预训练。一句话Offload 是把数据放到更便宜的仓库仓库越便宜搬运费越贵。你在配置里写一行device: cpu很容易但一定要想清楚你的训练任务到底能不能承受这笔搬运费。4. 一次从单卡到 8 卡的迁移实测配置、参数与显存曲线4.1 我用一个 7B 模型做了逐步迁移测试为了让上面的账算落地我拿一个 70 亿参数的模型做了一次完整迁移测试8 张 80GB 显存的卡单机互联序列长度 2048全局 batch 保持一致。测试过程就是从普通数据并行一路切到 ZeRO Stage 3再叠加 CPU Offload记录每一步的显存和吞吐变化。第一步普通数据并行直接跑模型状态需要 112GB/卡还没等激活值加载就 OOM 了连一个 step 都跑不完。这一步验证了前面说的数据并行加再多卡也救不了单卡装不下。第二步启用 Stage 1。配置里只需要把优化器状态切成 8 份每卡只要维护 12×7B/8 ≈ 10.5GB 的优化器状态加上参数和梯度模型状态落到 38.5GB。加上激活值后实测约 47GB可以跑了。吞吐和基线理论上一致因为梯度同步逻辑基本没变。第三步启用 Stage 2。梯度也切成 8 份模型状态降到 26.25GB实测约 33GB 显存占用。这一步开始引入 reduce-scatter 替代 all-reduce实测吞吐下降很小不到 3%。第四步启用 Stage 3。参数、梯度、优化器状态全切分模型状态降到 14GB实测约 22GB。因为每层前向/反向都要 all-gather 参数吞吐比 Stage 2 又低了约 8%。但换来的是显存余量非常大甚至可以再翻倍 batch size 来挽回吞吐。4.2 一份可以直接抄的配置框架下面这份是某个主流分布式训练框架的 ZeRO 配置格式字段含义我用注释标出来。实际使用时根据框架版本略有差异但核心字段大同小异{ zero_optimization: { stage: 3, offload_optimizer: { device: cpu, pin_memory: true }, overlap_comm: true, reduce_bucket_size: 500m, allgather_bucket_size: 500m, reduce_scatter: true }, optimizer: { type: Adam, params: { lr: 3e-5, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.0 } }, scheduler: { type: WarmupLR, params: { warmup_min_lr: 0, warmup_max_lr: 3e-5, warmup_num_steps: 500 } } }几个关键字段的取舍理由offload_optimizer.device设为cpu后优化器状态会常驻 CPU 内存。pin_memory: true可以加快 CPU 到 GPU 的拷贝速度但如果 CPU 内存本身紧张这个开关会加重压力。overlap_comm: true让通信和计算重叠理论上能隐藏一部分 all-gather 开销。实测在单机高速互联下收益明显但在跨机器慢网络上提升有限。reduce_bucket_size和allgather_bucket_size控制单次通信的数据块大小。块越大单次通信耗时越长但次数越少块越小通信更细碎但更容易穿插到计算间隙里。500m 左右是个相对稳的起点慢网络可以往小调高速互联可以往大调。4.3 实测结果显存与吞吐的此消彼长配置每卡模型状态每卡实测显存相对吞吐备注普通数据并行112GBOOM无法运行单卡装不下Stage 138.5GB约 47GB1.00x显存压力大幅缓解Stage 226.25GB约 33GB0.97x梯度切分引入少量通信Stage 314GB约 22GB0.92x参数 all-gather 带来额外开销Stage 3 CPU Offload更低约 14GB0.70x 左右每步要跨 PCIe 搬运优化器状态这组数据是在单机高速互联环境下的结果。如果你换到多机普通千兆/万兆网络Stage 3 和 Offload 的吞吐下降会更明显。我在集群上实测下来最稳妥的做法是先跑一遍 Stage 2 作为基准确认网络能承受额外通信后再上 Stage 3。别一上来就追求最省显存那样往往把性能也一起省掉了。5. 后来者最容易踩的坑通信开销、收敛性和工程实现的取舍5.1 一张踩坑排查表这些年我见过太多人在 ZeRO 上翻车问题大多不在概念本身而在几个非常具体的位置。我把最常见的坑整理成一张排查表现象根因处理方式Stage 3 下显存还是溢出只切了模型状态激活值没有被处理打开激活检查点gradient checkpointing用一点计算换显存开了 Offload 后吞吐暴跌CPU 内存带宽不够搬运优化器状态成了瓶颈先 Offload 优化器状态别一股脑全搬再不行调大 batch多机训练时某几张卡明显慢网络拓扑不均匀慢卡拖累全局 all-gather换更快的机器间互联或者干脆在配速阶段先用 Stage 2保存的检查点无法在单卡加载Stage 3 下权重是分片保存的没有合并用框架提供的合并脚本把分片权重合并成完整权重再加载加载预训练权重时形状对不上权重按卡数切分加载逻辑没适配先统一切分数量或先合并再切分换了优化器之后流程报错某些优化器没有实现分区逻辑回退到基础 Adam 类优化器或用兼容模式5.2 收敛性分区不会改变数学但 Offload 会改变节奏前面说了ZeRO 的三种切分在数学上等价于普通数据并行只要全局 batch 不变收敛曲线不会因为切分而漂移。这个结论在实际验证过多次之后基本可以当作共识。真正会影响收敛的往往不是切分而是 Offload 导致的训练节奏变化。举个例子Offload 之后每步吞吐下降如果你为了追上原来的训练进度而调大学习率或减少 warmup就可能扰动收敛。我的建议是一次只动一个变量。先保持超参不变只看显存和吞吐确认稳定后再决定要不要调整 batch 或学习率策略。另外Stage 3 配合激活检查点几乎成了标配。激活检查点会在前向时丢弃中间激活值、反向时重新计算代价是额外的计算量但省下的显存非常可观。两者叠加后7B 模型在 22GB 左右显存就能跑起来这个组合我是强烈推荐的。5.3 我给新手的选型顺序从最保守的配置开始如果你正准备把一个大模型训练任务迁到多卡环境我给一个非常具体、可以照做的顺序先用单卡跑一个小 batch确认识别流程本身没问题。切到普通数据并行记录 OOM 时的 batch 上限——这会告诉你显存瓶颈在哪。开 Stage 1确认能跑通记录吞吐。开 Stage 2对比吞吐如果下降超过 5%先检查是不是网络问题再继续。只有 Stage 2 仍然 OOM 或无法放大 batch 时才上 Stage 3。最后再决定要不要 Offload。Offload 是最后一根稻草不是第一选择。这个过程看起来保守但它能让你在每一个节点都清楚瓶颈到底是谁。我见过太多人一步跳到 Stage 3 Offload结果吞吐惨不忍睹最后又不知道问题出在哪。按顺序一步步试排错成本最低。最后聊点我的个人操作习惯如果说这几年有什么值得分享的体会那就是 ZeRO 真正解决的远不只是显存不够这个问题它让一套统一的分片思路同时覆盖了参数、梯度和优化器状态也让它后来能被继续延伸到 CPU 内存和 NVMe 存储上。我自己每次配置新任务时都会先看一眼机器之间的网络到底处于什么水平再决定 Stage 和 Offload 的搭配。再分享一个小技巧在正式训练前先开一个只有几万步的探路训练把不同 stage 的吞吐和显存数据都测一遍整理成表再决策。磨刀不误砍柴工这套流程帮我省掉的返工时间比我花在研究配置上的时间多得多。