2. 分布式训练基础
在进入 DeepSpeed ZeRO 的三个阶段之前，必须先建立分布式训练的“坐标系”。因为 ZeRO 并不是凭空出现的优化技巧，它本质上是对数据并行中冗余状态的一次系统性重构。如果不理解数据并行、模型并行、流水线并行各自切分的是什么，以及显存到底被谁吃掉，后面 ZeRO 的 Stage 1/2/3 就会变成死记硬背。
一、三种并行范式的切分对象
分布式训练的核心问题只有一个：单张 GPU 放不下模型、数据或计算量，于是把不同维度的负担分摊到多张卡上。三种主流范式切分的对象完全不同：
- 数据并行（Data Parallel, DP）：每张卡持有完整模型副本，把全局 batch 切成 N 份，各卡独立前向、反向，再通过 AllReduce 同步梯度。切分的是数据，模型不切。
- 模型并行（Model Parallel / Tensor Parallel, TP）：把单层内的权重矩阵按行或按列切开，分到不同卡上，层内计算需要频繁通信。切分的是单层参数。
- 流水线并行（Pipeline Parallel, PP）：把模型按层切成若干 stage，每个 stage 放在不同卡上，micro-batch 像流水线一样依次流过。切分的是模型深度。
实际的大模型训练通常是三者组合，即所谓 3D 并行：TP 负责层内、PP 负责层间、DP 负责数据副本。理解这一点，才能理解 ZeRO 的定位——它优化的是 DP 这一维度。
二、显存瓶颈：四类状态吃掉你的卡
为什么数据并行在单卡放得下模型时也会 OOM？因为训练时的显存占用远不止“模型参数”这一项。以混合精度 Adam 训练为例，单卡显存主要由四部分构成：
| 状态类别 | 典型精度 | 每参数字节数 | 说明 |
|---|---|---|---|
| 模型参数（FP16/BF16） | 2 字节 | 2 | 前向反向使用的权重 |
| 梯度（FP16/BF16） | 2 字节 | 2 | 反向传播产生 |
| 优化器状态（Adam） | FP32 | 12 | FP32 主权重 4 + 一阶动量 4 + 二阶动量 4 |
| 激活值 / 临时缓冲 | FP16 | 随 batch 变化 | 与序列长度、batch 强相关 |
可以看到，仅参数、梯度、优化器状态三项就达到 16 字节/参数。一个 7B 模型，这部分就是 7×10⁹ × 16 ≈ 112 GB，远超单张 80GB 的 A100。而这 16 字节中，真正“必须每卡一份”的只有参数和梯度，优化器状态在数据并行下是完全冗余的——这正是 ZeRO 的切入点。
三、一个具体的显存估算案例
假设用 8 张 A100（80GB）做纯数据并行训练一个 13B 模型，混合精度 Adam：
参数量 P = 13e9
每卡模型状态 = P × 16 字节
= 13e9 × 16 = 208 GB
单卡 A100 显存 = 80 GB
→ 208 GB > 80 GB，直接 OOM，连激活值都还没算
若改用 ZeRO Stage 1（只切优化器状态）：
优化器状态 12 字节切到 8 卡 → 12/8 = 1.5 字节/参数
每卡 = 13e9 × (2 + 2 + 1.5) = 71.5 GB
→ 勉强可跑，但激活值空间紧张
若改用 ZeRO Stage 3（参数+梯度+优化器全切）：
每卡 = 13e9 × 16 / 8 = 26 GB
→ 留出充足空间给激活值和 KV Cache这个估算说明：显存瓶颈不是单一因素，而是参数、梯度、优化器状态、激活值四者叠加。ZeRO 逐阶段地切分前三项，Stage 3 之后才真正把“模型状态”这一块压到最低。
四、常见误区与复盘
理解了这三种并行范式和显存构成，就能明白 ZeRO 为什么按“优化器状态 → 梯度 → 参数”的顺序分三个阶段切分：它切分的正是数据并行中那份最冗余、最占空间的模型状态，而把计算和通信模式尽量保持在数据并行的框架内。