数据并行与 ZeRO：三级分片各省多少
内容提要
本文介绍数据并行与ZeRO优化技术。ZeRO通过三级分片（优化器状态、梯度、权重）减少显存占用，前两级不增通信量，第三级需额外all-gather。ZeRO-1每参数4+12/N字节，ZeRO-2为2+14/N，ZeRO-3降至16/N。分片可放CPU，用双缓冲减少GPU占用，但需PCIe带宽支持。
延伸解读
ZeRO 前两级为何“白拿”
ZeRO-1 和 ZeRO-2 在不增加通信量的前提下减少显存，关键在于它们将普通 DP 的 all-reduce 拆解为 reduce-scatter 与 all-gather，并在中间插入优化器 step 或梯度分片。这种拆解使得通信总量不变，却能让每卡只持有部分状态，因此几乎无额外代价。这也是当前训练普遍至少开启 ZeRO-1 的原因。
ZeRO-3 与张量/流水线并行的取舍
ZeRO-3 将权重也分片，但代价是通信量从 2Ψ 增至 3Ψ，且每层前向/反向都需等待 all-gather，通信延迟更高。相比之下，张量并行和流水线并行同样能切分权重，但通信模式更可控，因此在大模型训练中常被优先采用，ZeRO-3 则较少单独使用。
分片因子 N 的正确理解
ZeRO 中的 N 指同一份权重的副本数，而非总卡数。当叠加流水线并行或 expert 并行时，权重副本仅存在于同一 stage 或同一 EP rank 的卡之间，因此分片因子应取 DP 度。估算显存时若误用集群总卡数，会导致结果偏差，影响资源配置。
CPU 分片与双缓冲的通用模式
将梯度分片放到 CPU 后，GPU 上只需保留两个缓冲，轮流用于接收当前 chunk 的梯度并执行 reduce，同时下一个 chunk 写入另一缓冲。这种“大块状态在别处、按块流过 GPU”的模式，普遍只需两个槽位，但依赖 PCIe 带宽且需计算与拷贝重叠。Kimi K3 在训练和 RL 中均应用了此模式。
Q&A
ZeRO优化技术是什么？它主要解决什么问题？
ZeRO是一种用于数据并行训练的内存优化技术，通过将模型状态（优化器状态、梯度、权重）分片到多个GPU上，减少每张卡的显存占用，从而解决数据并行中每卡显存不减少的问题。
ZeRO的三个级别分别切分什么？每参数显存占用如何变化？
ZeRO-1切分优化器状态，每参数显存从16字节降至4+12/N字节；ZeRO-2进一步切分梯度，每参数显存降至2+14/N字节；ZeRO-3连权重也切分，每参数显存降至16/N字节。其中N是数据并行度。
ZeRO-1和ZeRO-2为什么不会增加通信量？
普通数据并行中的all-reduce操作等价于一次reduce-scatter加一次all-gather。ZeRO-1和ZeRO-2只是将all-reduce拆分为reduce-scatter和all-gather，并在中间插入优化器步骤，因此通信量保持不变，仍为2Ψ（Ψ为参数量）。
ZeRO-3相比前两级增加了什么通信开销？为什么？
ZeRO-3需要在前向和反向传播时对每一层临时进行all-gather操作以获取完整权重，因此通信量从2Ψ增加到3Ψ，并且每层都需要等待通信完成。
在流水线并行或专家并行下，ZeRO的分片因子N应该如何确定？
ZeRO的分片因子N是同一份权重的副本数，即数据并行度，而不是总卡数。在流水线并行下，权重只在同一流水线阶段的卡之间复制；在专家并行下，权重只在同一专家并行位置的卡之间复制。因此，估算显存时应使用数据并行度。
将梯度分片放到CPU后，为什么GPU上只需要两个缓冲区？
使用双缓冲机制：反向传播时，一个缓冲区（buffer A）用于当前chunk的梯度，进行reduce并累加到CPU上的分片；同时下一个chunk的梯度写入另一个缓冲区（buffer B）。两个缓冲区轮流使用，使得GPU上的梯度占用从全部参数减少到两个chunk的大小。
