3、ZeRO技术概述：ZeRO三个阶段原理与对比
在数据并行训练中，每张 GPU 都保存一份完整的模型状态：参数、梯度和优化器状态。以 Adam 为例，每个参数需要保存 FP16 参数、FP16 梯度、FP32 主权重、FP32 一阶动量和 FP32 二阶动量，合计约 16 字节/参数。一个 10B 参数模型仅模型状态就需约 160GB 显存，远超单卡容量。ZeRO（Zero Redundancy Optimizer）的核心思想是：把这些冗余的模型状态切分到各张 GPU 上，需要时再通过通信临时聚合，从而在不增加单卡计算量的前提下大幅降低显存占用。
ZeRO 分为三个阶段，逐级递进地消除冗余：Stage 1 切分优化器状态，Stage 2 再切分梯度，Stage 3 进一步切分参数。每推进一个阶段，显存占用进一步下降，但通信量相应增加。理解这三个阶段的切分对象、通信时机与显存收益，是掌握 DeepSpeed 配置调优的基础。
核心原理：逐级消除冗余
ZeRO 的切分逻辑建立在数据并行的冗余之上。在标准 DDP 中，每张卡都持有完整的参数、梯度和优化器状态，N 张卡就有 N 份完全相同的副本。ZeRO 的做法是：把可切分的状态按 N 等份划分，每张卡只保留其中一份，需要完整状态时通过 AllGather 或 ReduceScatter 临时获取。
- Stage 1：只切分优化器状态。参数和梯度仍完整复制，前向反向计算无需额外通信。优化器更新时，每张卡只更新自己负责的那份参数，再通过 AllGather 把更新后的参数同步给所有卡。显存从 16 字节/参数降至约 4 字节/参数。
- Stage 2：在 Stage 1 基础上切分梯度。反向传播产生梯度后，通过 ReduceScatter 把梯度按分片归约到对应卡上，每张卡只保留自己负责的梯度分片。显存进一步降至约 2 字节/参数，通信量增加约 50%。
- Stage 3：在 Stage 2 基础上切分参数。每张卡只保存自己负责的参数分片，前向和反向计算需要参数时，通过 AllGather 临时聚合完整参数，计算完立即释放。显存降至约 0.25 字节/参数，但通信量翻倍。
关键结论：ZeRO 三阶段是“用通信换显存”的递进方案。Stage 1 几乎无额外通信代价，是性价比最高的起点；Stage 2 适合梯度显存成为瓶颈的场景；Stage 3 能训练超大模型，但通信开销最大，通常需要配合通信重叠和分片策略优化。
方法流程：一次训练迭代中三阶段的数据流
以 Stage 2 为例，一次完整的训练迭代包含以下步骤：
- 前向传播：每张卡使用完整参数计算自己 batch 的损失，参数此时是完整副本，无需通信。
- 反向传播：计算得到完整梯度，但梯度尚未归约。
- 梯度 ReduceScatter：将完整梯度按分片归约到各卡，每张卡得到自己负责的梯度分片，其余分片被释放。
- 优化器更新：每张卡用自己那份梯度分片和优化器状态分片，更新对应的参数分片。
- 参数 AllGather：将各卡更新后的参数分片聚合，恢复完整参数，供下一轮前向使用。
Stage 3 的区别在于：前向传播前需要先 AllGather 参数，反向传播前也需要再次 AllGather 参数（因为反向使用的参数在前向结束后已被释放），因此通信次数更多。DeepSpeed 通过通信与计算重叠（如预取下一层参数）来隐藏部分开销。
案例与场景：如何选择阶段
| 场景 | 推荐阶段 | 理由 |
|---|---|---|
| 单卡可放下模型状态，仅想提升 batch size | Stage 1 | 几乎无额外通信，显存收益明显 |
| 梯度显存占用大，优化器状态已切分仍不足 | Stage 2 | 梯度切分进一步释放显存，通信增加可接受 |
| 模型参数本身超过单卡容量（如 10B+） | Stage 3 | 参数切分是唯一能放下模型的方式 |
| 高速互联（NVLink/InfiniBand）环境 | Stage 3 + 通信重叠 | 高带宽可抵消通信开销，显存收益最大化 |
| 低速网络或跨节点通信瓶颈明显 | Stage 1 或 Stage 2 | 避免 Stage 3 的频繁 AllGather 拖慢训练 |
误区与复盘
误区一：阶段越高越好。Stage 3 显存最省，但通信量最大。在低速网络下，Stage 3 可能比 Stage 2 慢很多。应根据模型规模、网络带宽和显存瓶颈综合选择。
误区二：ZeRO 会减少计算量。ZeRO 不改变计算量，每张卡的计算量与数据并行一致。它只改变模型状态的存储方式和通信模式。
误区三：Stage 1 没有通信开销。Stage 1 在优化器更新后需要 AllGather 参数，只是通信量相对较小，并非零开销。
复盘时建议关注三个指标：单卡显存峰值、每步通信耗时占比、吞吐量（tokens/s）。如果显存已满足但吞吐下降明显，说明通信成为瓶颈，可尝试降低阶段或开启通信重叠；如果显存仍不足，则需提升阶段或结合 CPU Offload。DeepSpeed 配置中通过 zero_optimization.stage 指定阶段，实际调优时应从 Stage 1 开始逐步验证。
{
"zero_optimization": {
"stage": 2,
"allgather_partitions": true,
"reduce_scatter": true,
"overlap_comm": true,
"contiguous_gradients": true
},
"train_batch_size": 64,
"fp16": { "enabled": true }
}理解 ZeRO 三阶段的切分对象与通信代价，是后续配置 DeepSpeed 和排查显存/吞吐问题的前提。实际项目中，建议先用小规模模型验证阶段选择，再逐步放大到目标规模。