📚 课程目录第4章 ZeRO优化器详解

4、ZeRO优化器详解：ZeRO-1、ZeRO-2、ZeRO-3的分片策略、显存节省原理与适用场景

在数据并行训练中，每张 GPU 都保存一份完整的模型参数、梯度和优化器状态。以 Adam 为例，混合精度训练下每个参数大约需要 16 字节：2 字节 fp16 参数、2 字节 fp16 梯度、4 字节 fp32 主参数、4 字节 fp32 动量、4 字节 fp32 方差。当模型达到百亿参数时，仅优化器状态就会占用数十 GB 显存，单卡根本无法容纳。ZeRO（Zero Redundancy Optimizer）的核心思想就是：把这些冗余的副本切分到各张 GPU 上，需要时再通过通信临时聚合。ZeRO-1、ZeRO-2、ZeRO-3 分别对应优化器状态、梯度、参数三个层级的切分，显存节省逐级递增，通信开销也逐级上升。

一、核心原理：冗余副本的逐级消除

理解 ZeRO 的关键是区分三类显存占用：参数（Parameters）、梯度（Gradients）、优化器状态（Optimizer States）。在标准数据并行中，这三者在每张卡上都是完整副本，属于纯粹的冗余。ZeRO 的分片策略就是按顺序消除这些冗余：

ZeRO-1：只切分优化器状态。每张卡仍持有完整参数和梯度，但只负责更新自己那部分参数对应的优化器状态。更新完成后通过 All-Gather 把参数同步回所有卡。
ZeRO-2：在 ZeRO-1 基础上再切分梯度。反向传播时梯度通过 Reduce-Scatter 直接归约到对应卡上，每张卡只保留自己负责的那份梯度，不再需要完整梯度副本。
ZeRO-3：进一步切分参数本身。每张卡只保存 1/N 的参数，前向和反向计算需要某层参数时，通过 All-Gather 临时收集完整参数，计算完立即释放。

显存节省公式（以 N 张卡、混合精度 Adam 为例）：

基线每卡显存 ≈ 2P + 2P + 4P + 4P + 4P = 16P 字节

ZeRO-1 每卡 ≈ 2P + 2P + (4P + 4P + 4P)/N = 4P + 12P/N

ZeRO-2 每卡 ≈ 2P + (2P + 12P)/N = 2P + 14P/N

ZeRO-3 每卡 ≈ 16P/N

二、方法流程：一次训练迭代中的通信时序

以 ZeRO-2 为例，一次完整的前向 + 反向 + 更新流程如下：

前向传播：每张卡使用完整参数计算自己 micro-batch 的损失，无需通信。
反向传播：计算梯度时，对每层梯度执行 Reduce-Scatter，将归约后的梯度分片分发到对应卡上。
优化器更新：每张卡用自己持有的梯度分片和优化器状态分片，更新对应的参数分片。
参数同步：通过 All-Gather 将各卡更新后的参数分片聚合，使每张卡重新获得完整参数，进入下一轮迭代。

ZeRO-3 的流程更复杂：前向传播每进入一层，先对该层参数做 All-Gather，计算完立即释放；反向传播同样需要重新 All-Gather 参数，并在计算完梯度后做 Reduce-Scatter。因此 ZeRO-3 的通信量约为 ZeRO-2 的 1.5 倍，但显存节省也最彻底。

三、配置实操：DeepSpeed 中的 ZeRO 阶段设置

在 DeepSpeed 配置文件中，通过 zero_optimization.stage 指定阶段：

{
  "train_batch_size": 64,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true
  }
}

若使用 ZeRO-3，还需增加 stage3_prefetch_bucket_size、stage3_param_persistence_threshold 等参数，并可开启 offload_optimizer 或 offload_param 将部分状态卸载到 CPU。启动命令通常为：

deepspeed --num_gpus=8 train.py --deepspeed ds_config.json

四、适用场景与选型对照

阶段	分片对象	显存节省	通信开销	典型适用场景
ZeRO-1	优化器状态	约 4 倍	低	优化器状态占主导，通信带宽有限
ZeRO-2	优化器状态 + 梯度	约 8 倍	中	梯度显存压力大，追求性价比
ZeRO-3	优化器状态 + 梯度 + 参数	接近 N 倍	高	单卡无法容纳参数的超大模型

选型建议：优先从 ZeRO-2 开始尝试。如果单卡仍 OOM，再升级到 ZeRO-3 并配合 CPU Offload。ZeRO-1 通常只在通信极度受限或模型较小时使用。注意 ZeRO-3 对模型代码有一定侵入性，部分自定义层需要额外适配。

五、常见误区与复盘

误区一：阶段越高越好。ZeRO-3 虽然显存最省，但通信量最大，训练速度可能明显下降。在显存足够时，ZeRO-2 往往比 ZeRO-3 更快。

误区二：忽略 bucket size 调优。allgather_bucket_size 和 reduce_bucket_size 过小会导致通信碎片化，过大则增加显存峰值。通常从 5e8 开始调优。

误区三：ZeRO-3 与模型并行混用不当。ZeRO-3 本身已切分参数，若再叠加张量并行，需仔细规划切分维度，避免通信冲突。

复盘时建议记录三个指标：每卡显存峰值、单步迭代时间、通信占比。通过对比不同 ZeRO 阶段的实测数据，才能找到当前硬件与模型规模下的最优配置。