4、ZeRO技术原理：ZeRO-1、ZeRO-2、ZeRO-3的分片策略与显存优化机制

在数据并行训练中，每张 GPU 都保存一份完整的模型参数、梯度与优化器状态。当模型规模增大时，这三类数据的显存占用会迅速超过单卡容量。ZeRO（Zero Redundancy Optimizer）的核心思想是：把这些原本被所有数据并行进程重复保存的数据切分开，按需通信、按需重建，从而在不改变模型数学等价性的前提下大幅降低单卡显存。

ZeRO 的三个阶段并不是三套独立算法，而是对“参数 / 梯度 / 优化器状态”三类冗余的逐级消除：ZeRO-1 切优化器状态，ZeRO-2 再切梯度，ZeRO-3 连参数也切。

4.1 三类显存冗余与分片对象

以 Adam 为例，假设模型有 Ψ 个参数，混合精度训练中每张卡通常需要保存：fp16 参数 2Ψ 字节、fp16 梯度 2Ψ 字节、fp32 优化器状态（一阶动量、二阶动量、fp32 主权重）共 12Ψ 字节。合计约 16Ψ 字节，其中优化器状态占了 75%。

  • 优化器状态：Adam 的 m、v 和 fp32 master weight，是显存占用最大的部分。
  • 梯度：反向传播产生的 fp16 梯度，在归约后每卡都有一份完整副本。
  • 参数：前向与反向计算必须使用完整参数，因此天然需要全量驻留。

ZeRO 的分片策略正是按上述顺序逐层深入：先切最占显存、通信代价最低的优化器状态，再切梯度，最后切参数。

ZeRO 三阶段分片对象与显存下降示意 Baseline (DDP) 优化器状态 12Ψ 梯度 2Ψ 参数 2Ψ ZeRO-1 优化器状态 /N 12Ψ/N 梯度 2Ψ 参数 2Ψ ZeRO-2 优化器状态 /N 12Ψ/N 梯度 /N 2Ψ/N 参数 2Ψ ZeRO-3 优化器 /N 12Ψ/N 梯度 /N 2Ψ/N 参数 /N 2Ψ/N N 为数据并行度；色块面积表示单卡显存占用，逐级缩小

4.2 ZeRO-1：只切优化器状态

ZeRO-1 将优化器状态按数据并行 rank 均分为 N 份，每个 rank 只维护自己那 1/N 的 m、v 和 fp32 主权重。前向和反向仍使用完整参数与完整梯度，因此计算图不变。

训练一步的流程是：各 rank 独立完成前向与反向，得到完整梯度；执行一次 reduce-scatter，把梯度按分片归约到对应 rank；每个 rank 用本地优化器状态更新自己负责的那部分参数；最后执行 all-gather，把更新后的参数分片广播回所有 rank，恢复完整参数。

# DeepSpeed 配置：ZeRO-1
{
  "zero_optimization": {
    "stage": 1,
    "reduce_bucket_size": 5e8,
    "allgather_bucket_size": 5e8
  },
  "train_batch_size": 64,
  "fp16": { "enabled": true }
}

ZeRO-1 的显存收益约为 4 倍（16Ψ → 4Ψ + 12Ψ/N），通信量只增加一次 reduce-scatter 和一次 all-gather，通常比纯 DDP 仅多约 1.5 倍通信量，是性价比最高的一档。

4.3 ZeRO-2：再切梯度

ZeRO-2 在 ZeRO-1 基础上，把梯度也按 rank 分片。反向传播时不再对完整梯度做 all-reduce，而是直接 reduce-scatter，每个 rank 只保留自己负责归约的那一段梯度。由于优化器状态本来就按同一分片规则存放，梯度归约完成后可以立即在本地更新对应参数分片，无需额外搬运。

  • 反向过程中梯度一旦产生就参与归约，显存峰值显著低于 ZeRO-1。
  • 参数仍需 all-gather 回全量，因此前向、反向仍使用完整参数。
  • 显存收益约 8 倍（16Ψ → 2Ψ + 14Ψ/N）。
ZeRO-2 是大多数中等规模模型（数亿到数十亿参数）的推荐起点：显存下降明显，通信开销可控，对代码侵入极小。

4.4 ZeRO-3：参数也分片

ZeRO-3 把参数也切分，每个 rank 只长期持有 1/N 的参数分片。计算某一层时，通过 all-gather 临时收集该层完整参数，计算完成后立即释放。因此显存占用与数据并行度 N 近似成反比，单卡显存可降到约 16Ψ/N。

代价是通信频率大幅上升：每一层前向和反向都要做一次参数 all-gather，反向还要对梯度做 reduce-scatter。为降低开销，DeepSpeed 引入参数预取与通信重叠，在计算当前层时提前收集下一层参数。

# DeepSpeed 配置：ZeRO-3
{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" },
    "offload_param": { "device": "cpu" },
    "stage3_prefetch_bucket_size": 5e7,
    "stage3_param_persistence_threshold": 1e5,
    "stage3_max_live_parameters": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "fp16": { "enabled": true }
}

4.5 三阶段对比与选型

阶段 分片对象 单卡显存（近似） 额外通信 适用场景
ZeRO-1 优化器状态 4Ψ + 12Ψ/N reduce-scatter + all-gather 显存略紧、通信敏感
ZeRO-2 优化器状态 + 梯度 2Ψ + 14Ψ/N reduce-scatter + all-gather 中等模型，通用推荐
ZeRO-3 优化器状态 + 梯度 + 参数 16Ψ/N 逐层 all-gather + reduce-scatter 超大模型、单卡装不下

4.6 常见误区与复盘

误区一：认为 ZeRO 阶段越高越好。ZeRO-3 通信量最大，在模型本可放入单卡时反而拖慢训练。误区二：忽略 batch size 与分片数的关系，N 增大后每卡有效 batch 变小，需同步调整梯度累积。误区三：开启 offload 后未评估 PCIe 带宽，CPU 卸载可能成为瓶颈。

复盘时建议按以下顺序排查：先确认 stage 是否与模型规模匹配；再检查 reduce_bucket_size、allgather_bucket_size 是否过小导致通信碎片化；最后观察日志中的通信占比与显存峰值，必要时在 ZeRO-2 与 ZeRO-3 之间做 A/B 对比。

# 观察显存与通信占比
deepspeed --num_gpus=8 train.py \
  --deepspeed ds_config_zero2.json \
  --log_steps 10

# 日志中关注
# zero_optimization stage = 2
# max memory allocated / reserved
# communication time vs compute time

理解 ZeRO 的关键，是把它看作“冗余消除”的渐进过程：优化器状态最先被切，梯度其次，参数最后。每一级都在显存与通信之间做新的权衡，选型时应以模型能否放入单卡、通信带宽是否充裕为判断依据，而不是盲目追求最高阶段。

← 上一章 📖 返回目录 下一章 →