AIBAT：教师主导的语境化语言评估工具，如何让AI从打分者变成行为分析师
2026/10/17 13:23:57
近十年大模型训练里的那只"显存怪兽"，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 把"切分"推进到了不同程度，对应三个 stage：
| Stage | 切分范围 | 每张卡模型状态占用（K 为参数量字节数） | 7B 模型、8 卡时每卡理论占用 |
|---|---|---|---|
| 基线（普通数据并行） | 不切分 | 2K（参数）+ 2K（梯度）+ 12K（优化器状态）= 16K | 112GB |
| Stage 1 | 仅优化器状态 | 2K + 2K + 12K/N | 约 38.5GB |
| Stage 2 | 优化器状态 + 梯度 | 2K + 2K/N + 12K/N | 约 26.25GB |
| Stage 3 | 参数 + 梯度 + 优化器状态 | 16K/N | 14GB |
注意这个表里的 N 是 GPU 数量。以 8 卡跑 7B 为例：普通数据并行每卡 112GB，直接 OOM；Stage 1 降到 38.5GB，Stage 2 降到 26.25GB，Stage 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 显存（HBM） | 1 ~ 2 TB/s | 纳秒级 | 数十 GB |
| CPU 内存（DRAM） | 50 ~ 100 GB/s | 百纳秒级 | 数百 GB ~ 数 TB |
| PCIe 总线 | 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 实测结果：显存与吞吐的此消彼长
| 配置 | 每卡模型状态 | 每卡实测显存 | 相对吞吐 | 备注 |
|---|---|---|---|---|
| 普通数据并行 | 112GB | OOM | 无法运行 | 单卡装不下 |
| Stage 1 | 38.5GB | 约 47GB | 1.00x | 显存压力大幅缓解 |
| Stage 2 | 26.25GB | 约 33GB | 0.97x | 梯度切分引入少量通信 |
| Stage 3 | 14GB | 约 22GB | 0.92x | 参数 all-gather 带来额外开销 |
| Stage 3 + CPU Offload | 更低 | 约 14GB | 0.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 的吞吐和显存数据都测一遍，整理成表再决策。磨刀不误砍柴工，这套流程帮我省掉的返工时间，比我花在研究配置上的时间多得多。