ZeRO 把哪些训练状态分片了？

问题从哪里来

普通数据并行在每张卡放一份模型和训练状态。Adam 的动量、方差及混合精度主权重会比前向权重本身占用更多空间。ZeRO 在保持数据并行思路的同时，把这些状态分散到多个设备。

关键机制

Stage 1 分片优化器状态；Stage 2 再分片梯度；Stage 3 进一步分片参数，计算某层时按需收集。用 P、G、O 表示参数、梯度、优化器字节数，理想每卡状态从 P+G+O 逐步变为 P+G+O/n、P+(G+O)/n、(P+G+O)/n。

ZeRO-3 ideal state/device = (P + G + O) / n

算例与证据

教学算例 P=2、G=2、O=12 字节／参数，n=4：复制状态为16，三个阶段理想值依次为7、5.5、4。不同精度与优化器的分项不同；这些不是峰值显存，未计激活、收集缓冲与通信工作区。

收益之外的新约束

更深分片会带来参数收集、通信和调度压力。有效性取决于网络、计算重叠和微批量；训练装得下之后还要解决吞吐，不可只按理论状态字节选配置。