---

一、核心思想

ZeRO 的全称是 Zero Redundancy Optimizer，核心思想是：在数据并行的基础上，把原本每张卡都重复保存的模型状态切分开，从而消除冗余、降低单卡显存占用。

在标准数据并行中，每张卡都保存一份完整的：

• 模型参数
• 梯度
• 优化器状态

这三部分在所有卡上完全重复，冗余度极高。ZeRO 就是把这三部分依次切分到不同卡上，对应三个阶段：

• Stage 1：切分优化器状态
• Stage 2：切分优化器状态和梯度
• Stage 3：切分优化器状态、梯度和模型参数

切得越彻底，单卡显存占用越低，但通信开销也越大。

---

二、先看标准数据并行的显存冗余

假设模型有 $\Psi$ 个参数，使用 Adam 优化器，混合精度训练。

标准数据并行下，每张卡需要保存：

• FP16 模型参数：$2\Psi$ 字节
• FP16 梯度：$2\Psi$ 字节
• FP32 主权重：$4\Psi$ 字节
• FP32 一阶动量：$4\Psi$ 字节
• FP32 二阶动量：$4\Psi$ 字节

合计每张卡约 $16\Psi$ 字节。如果模型有 10B 参数，每张卡光模型状态就要 160GB，单卡根本放不下。

ZeRO 要解决的就是这种冗余。

---

三、ZeRO Stage 1：切分优化器状态

1、优化了什么

Stage 1 只切分优化器状态，包括：

• FP32 主权重
• 一阶动量
• 二阶动量

假设有 $N$ 张卡，那么每张卡只保存 $1/N$ 的优化器状态。

2、显存变化

每张卡显存从约 $16\Psi$ 降到：

• FP16 参数：$2\Psi$
• FP16 梯度：$2\Psi$
• 优化器状态：$12\Psi / N$

当 $N$ 很大时，优化器状态占用大幅下降，每张卡显存趋近于 $4\Psi$。

3、通信特点

Stage 1 的通信模式和标准数据并行基本一致：

• 反向传播后做梯度 All-Reduce；
• 更新参数时，每张卡只更新自己负责的那部分优化器状态对应的参数；
• 更新完后需要做一次 All-Gather，把参数同步回所有卡。

所以 Stage 1 通信量增加不多，主要是参数广播。

4、适用场景

• 显存瓶颈主要在优化器状态；
• 希望以较小通信代价换取显存下降；
• 是 ZeRO 中最容易落地的一级。

---

四、ZeRO Stage 2：切分优化器状态和梯度

1、优化了什么

Stage 2 在 Stage 1 基础上，进一步切分梯度。

每张卡只保存自己负责的那部分梯度，而不是完整梯度。

2、显存变化

每张卡显存从约 $16\Psi$ 降到：

• FP16 参数：$2\Psi$
• 梯度：$2\Psi / N$
• 优化器状态：$12\Psi / N$

当 $N$ 很大时，每张卡显存趋近于 $2\Psi$。

3、通信特点

Stage 2 使用 Reduce-Scatter 代替 All-Reduce：

• 标准 DP 是 All-Reduce，每张卡都拿到完整平均梯度；
• Stage 2 是 Reduce-Scatter，每张卡只拿到自己负责的那部分平均梯度。

Reduce-Scatter 的通信量和 All-Reduce 相当，但每张卡只需要存一部分梯度，所以显存下降，通信量没有明显增加。

4、适用场景

• 显存瓶颈在梯度；
• 希望在 Stage 1 基础上进一步省显存；
• 通信开销仍然可控。

---

五、ZeRO Stage 3：切分优化器状态、梯度和模型参数

1、优化了什么

Stage 3 在 Stage 2 基础上，进一步切分模型参数。

每张卡只保存自己负责的那部分参数，而不是完整模型。

2、显存变化

每张卡显存从约 $16\Psi$ 降到：

• 参数：$2\Psi / N$
• 梯度：$2\Psi / N$
• 优化器状态：$12\Psi / N$

合计约 $16\Psi / N$。

当 $N$ 很大时，单卡显存占用可以下降接近 $N$ 倍，这就是 ZeRO-3 能显著降低单卡显存占用的原因。

3、前向和反向需要 All-Gather

因为每张卡只保存部分参数，所以在计算某一层时：

• 需要先通过 All-Gather 把这一层的完整参数收集起来；
• 计算完后，再释放掉不属于自己的参数。

反向传播类似，也需要按层 All-Gather 参数。

4、为什么通信开销上升

Stage 3 的通信开销上升，主要有三个原因：

第一，参数 All-Gather 次数多。 每一层前向都要 All-Gather 一次参数，反向还要再来一次。层数越多，All-Gather 次数越多。

第二，通信量随参数规模线性增长。 每层都要传输参数，整体通信量和模型参数量、层数都相关，远高于 Stage 1 和 Stage 2。

第三，计算和通信难以完全重叠。 虽然可以通过预取、流水线等方式重叠，但实现复杂，实际中很难完全隐藏通信。

所以 Stage 3 用通信换显存，显存降得最狠，但通信开销也最大。

5、适用场景

• 模型极大，单卡连参数都放不下；
• 显存是硬瓶颈，通信可以靠高带宽网络缓解；
• 常与张量并行、流水线并行组合使用。

---

六、三级对比

可以用一个表来总结：

阶段	切分内容	每卡显存	通信特点	显存节省
Stage 1	优化器状态	约 $4\Psi + 12\Psi/N$	与标准 DP 接近	中等
Stage 2	优化器状态 + 梯度	约 $2\Psi + 14\Psi/N$	用 Reduce-Scatter 替代 All-Reduce	较大
Stage 3	优化器状态 + 梯度 + 参数	约 $16\Psi/N$	每层都要 All-Gather，通信量大	最大

---

七、为什么说 ZeRO-3 省显存但通信开销上升

省显存的原因：

• 参数、梯度、优化器状态全部切分；
• 每张卡只保存 $1/N$ 的模型状态；
• 单卡显存占用大幅下降。

通信开销上升的原因：

• 前向每层都要 All-Gather 参数；
• 反向每层也要 All-Gather 参数；
• 通信次数多、通信量大；
• 计算和通信重叠难度高。

一句话总结：ZeRO-3 是用通信换显存，适合模型极大、显存极度紧张的场景。

---

八、简单评价这道面试题

1、这道题的价值

• 这是大模型分布式训练中非常核心的一道题；
• 一道题能考察 ZeRO 三个阶段的递进关系、显存计算和通信代价；
• 区分度好，初级可能只知道 ZeRO 能省显存，中级能说清三个阶段，高级能讲清通信开销和实际权衡；
• 贴近实战，DeepSpeed、PyTorch FSDP 等都与 ZeRO 思想密切相关。

2、这道题的不足

• 题目偏概念，如果没有实际训练经验，容易背答案；
• 没有涉及具体数字计算，比如给定模型规模算显存；
• 没有涉及 ZeRO-Offload、ZeRO-Infinity 等延伸内容；
• 对通信开销的考察还可以更深入，比如带宽、延迟、重叠策略。

3、建议的追问方向

• 给定一个 10B 模型和 64 张卡，ZeRO-1、2、3 分别每卡占多少显存？
• ZeRO-3 和 PyTorch FSDP 是什么关系？
• ZeRO-Offload 和 ZeRO-Infinity 解决了什么问题？
• ZeRO-3 中如何做参数预取来隐藏通信？
• ZeRO-3 和张量并行、流水线并行如何组合？

4、总体评价

这是一道很好的大模型训练系统面试题，适合考察候选人对显存优化和分布式通信的理解。它既能考察基础概念，也能延伸到工程权衡。如果候选人能把三个阶段的切分对象、显存变化、通信特点以及 ZeRO-3 的通信代价都讲清楚，基本可以判断其具备大模型分布式训练的扎实基础。

---

---