4、ZeRO系列技术演进：ZeRO-1/2/3原理、分片策略、通信开销分析、与数据并行/模型并行的关系
ZeRO（Zero Redundancy Optimizer）是微软 DeepSpeed 提出的显存优化技术族，其核心思想非常朴素：数据并行中每张卡都保存了完整的模型状态，而这些状态在训练过程中高度冗余。ZeRO 通过把优化器状态、梯度、参数按数据并行度切分到各卡，让每张卡只保留 1/N 的模型状态，从而在不改变计算语义的前提下大幅降低单卡显存占用。ZeRO-1/2/3 分别对应切分粒度的逐级加深。
一、核心原理：三块显存与三级分片
训练一个参数量为 Ψ 的模型，混合精度 Adam 训练时每卡显存主要由三部分构成：优化器状态（fp32 动量、方差、fp32 主权重，约 12Ψ 字节）、梯度（fp16，约 2Ψ 字节）、参数（fp16，约 2Ψ 字节）。ZeRO 的三级分片正是按这个顺序逐层切开：
- ZeRO-1（Pos）：只切分优化器状态，梯度与参数仍全量保留。显存从 16Ψ 降到约 4Ψ + 12Ψ/N。
- ZeRO-2（Pos+g）：在 ZeRO-1 基础上再切分梯度，显存降到约 2Ψ + 14Ψ/N。
- ZeRO-3（Pos+g+p）：进一步切分参数，显存降到约 16Ψ/N，接近线性扩展。
下图展示了三级分片在单卡上的显存构成变化，以及每级新增的通信环节：
二、分片策略与通信开销分析
分片越深，显存越省，但通信量也越大。理解通信开销的关键是区分通信量（Volume）与通信次数（Frequency）。ZeRO-1 只在优化器步做一次 All-Gather，通信量与基线 DP 相当；ZeRO-2 把梯度规约改为 Reduce-Scatter，通信量不变但省了显存；ZeRO-3 则在前向和反向中反复 All-Gather 参数，通信量约为 ZeRO-2 的 1.5 倍。
| 级别 | 切分对象 | 单卡显存 | 相对通信量 | 适用场景 |
|---|---|---|---|---|
| ZeRO-1 | 优化器状态 | 4Ψ + 12Ψ/N | 1.0× | 显存略紧，通信敏感 |
| ZeRO-2 | 优化器状态 + 梯度 | 2Ψ + 14Ψ/N | 1.0× | 最常用，性价比最高 |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | 16Ψ/N | 约 1.5× | 超大模型，显存极限 |
三、与数据并行、模型并行的关系
ZeRO 本质上是数据并行的显存优化增强，而不是替代数据并行。它保留了 DP 的计算图与梯度语义，只是把冗余状态切片。与模型并行（TP/PP）的关系则是互补：
- 与张量并行（TP）：TP 切分单层内的矩阵运算，通信频繁且要求高带宽（通常限节点内 NVLink）。ZeRO-3 切分的是跨层的参数副本，通信可跨节点。两者可叠加，形成 TP+ZeRO-3 的 3D 并行。
- 与流水并行（PP）：PP 按层切分到不同设备，天然降低单卡参数。ZeRO-1/2 可与 PP 无缝组合，因为 PP 不改变每层内的 DP 冗余结构。
- 与纯 DP 的区别：纯 DP 每卡显存随模型线性增长；ZeRO-3 下每卡显存随 DP 度 N 近似线性下降，使 DP 也能训练千亿参数模型。
四、实操配置与常见误区
在 DeepSpeed 中，ZeRO 级别通过 JSON 配置的 zero_optimization.stage 指定：
{
"train_batch_size": 64,
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"allgather_partitions": true,
"overlap_comm": true,
"reduce_scatter": true,
"contiguous_gradients": true
},
"fp16": { "enabled": true }
}pin_memory，导致 CPU-GPU 拷贝成为瓶颈；③ 把 ZeRO 与 TP 混用时未调整 tensor_parallel 的 world size，造成分片维度冲突。
复盘来看，ZeRO 的演进路线是一条清晰的“用通信换显存”的曲线：ZeRO-1 几乎零额外通信代价，ZeRO-2 是当前工业界默认选择，ZeRO-3 则是显存极限下的最后手段。理解每一级切了什么、多付了什么通信，才能在真实集群带宽与模型规模之间做出正确取舍。