4. ZeRO Stage 1

在混合精度训练中，显存占用的大头往往不是模型参数本身，而是优化器状态。以 Adam 为例，每个可训练参数需要维护一阶动量与二阶动量两份 FP32 副本，再加上 FP32 主权重，合计约 12 字节/参数。当模型达到十亿甚至百亿参数时，优化器状态会迅速成为显存瓶颈。ZeRO Stage 1 的核心思路非常直接：把优化器状态按数据并行进程切分，每个 GPU 只保存 1/N 的优化器状态，从而把这一部分显存降到原来的 1/N。

4.1 核心原理：分片、收集与释放

ZeRO Stage 1 只对优化器状态做分片，模型参数和梯度仍然在每个数据并行进程上完整保留。训练一步的流程可以概括为“前向反向照常、优化器状态分片更新、参数广播同步”三个阶段：

• 前向与反向：每个 rank 用完整参数计算梯度，梯度通过 AllReduce 在数据并行组内求平均，得到一致的完整梯度。
• 优化器更新：每个 rank 只负责自己分片内的那部分参数，用本地保存的优化器状态执行 Adam 更新，得到更新后的参数分片。
• 参数收集：通过 AllGather 把各 rank 更新后的参数分片收集起来，拼成完整参数，供下一轮前向使用。

关键点在于：优化器状态是“常驻分片”，而完整参数只在更新后短暂地通过 AllGather 重建。这样既保证了数学等价性，又显著降低了优化器状态的显存占用。

4.2 配置参数：一行开关背后的细节

在 DeepSpeed 中启用 ZeRO Stage 1 只需要在配置文件中设置 zero_optimization.stage = 1。但要让训练真正稳定高效，还需要关注几个配套参数：

4.3 实战演练：从零跑通 Stage 1

下面以 HuggingFace Transformers 模型为例，展示一个最小可运行的 DeepSpeed Stage 1 配置与启动命令。

{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 1,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 1,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_scatter": false
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 2e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  }
}

启动命令如下，注意 --num_gpus 要与实际数据并行规模一致：

deepspeed --num_gpus=4 train.py \
  --deepspeed ds_zero_stage1.json \
  --model_name_or_path bert-large-uncased \
  --per_device_train_batch_size 8 \
  --output_dir ./output

训练启动后，日志中会出现类似下面的显存统计，可以直观看到优化器状态被分片：

[Rank 0] Zero Stage 1: optimizer state partitioned
[Rank 0] Before: optimizer state 12.0 GB per GPU
[Rank 0] After : optimizer state  3.0 GB per GPU (1/4)
[Rank 0] AllGather bucket size: 500000000

4.4 误区与复盘

• 误区一：Stage 1 能省参数显存。不能。Stage 1 只分片优化器状态，模型参数和梯度仍然完整复制，参数显存不变。
• 误区二：分片后不需要通信。仍然需要 AllReduce 梯度、AllGather 参数，通信量比纯数据并行略高，但换来显存下降。
• 误区三：Stage 1 一定比 Stage 2 快。不一定。Stage 2 进一步分片梯度，显存更省，但通信模式不同；应根据显存与带宽实测选择。
• 复盘要点：先确认瓶颈是优化器状态而非激活值；再检查 allgather_bucket_size 是否过大导致显存峰值；最后用 overlap_comm 观察吞吐变化。