ZeRO 要解决的不是 activation 显存,而是模型训练状态的复制。 以 BF16/FP16 训练 + AdamW 为例,每个参数通常不只占 2 bytes: 还有梯度、FP32 master weight、Adam 的一阶/二阶动量。 DDP 会把这些状态在每张卡上都保存一份,所以卡数变多主要提升吞吐,并不自动扩大单模型容量。

一个 7B 模型的粗算

def gb(params_billion, bytes_per_param):
    return params_billion * 1e9 * bytes_per_param / 1024**3

params = 7
print("params bf16:", gb(params, 2), "GiB")
print("grad bf16:", gb(params, 2), "GiB")
print("master weight fp32:", gb(params, 4), "GiB")
print("adam m/v fp32:", gb(params, 8), "GiB")
print("training states total:", gb(params, 16), "GiB")

ZeRO 的全名是 Zero Redundancy Optimizer。名字里的 zero redundancy 很直白: 不要让每个 rank 都重复保存相同训练状态。 三个 stage 是逐级增加切分范围:先切 optimizer state,再切 gradient,最后切 parameter。

ZeRO 不是免费午餐。省显存的代价是通信和调度更复杂。 Day 15 学的 AllGather / ReduceScatter 在这里第一次变成主角: ZeRO-2 用 ReduceScatter 让每个 rank 只保留梯度分片; ZeRO-3 在每层计算前 AllGather 参数,计算完再释放。

ZeRO-3 一层的生命周期

# 概念伪代码:真实 DeepSpeed 内部会做 bucket、prefetch、overlap
for layer in model.layers:
    full_param = all_gather(param_shard[layer])   # 通信:取回本层完整参数
    activation = layer.forward(x, full_param)     # 计算:正常 forward
    release(full_param)                           # 显存:用完释放或缓存

for layer in reversed(model.layers):
    full_param = all_gather(param_shard[layer])   # backward 也需要参数
    grad = layer.backward(full_param)
    grad_shard = reduce_scatter(grad)             # 通信:只保留本 rank 梯度分片
    optimizer_update(local_state[layer], grad_shard)
    release(full_param)

RoadMap 指定今天读 ZeRO: Memory Optimization Toward Training Trillion Parameter Models。 不建议从公式一路硬啃。今天的读法是先抓主问题、主表格和主结论: 大模型训练的显存瓶颈来自模型状态复制,ZeRO 用分片状态把数据并行的冗余降下来, 再通过通信调度尽量保住吞吐。

论文笔记模板

# ZeRO 论文笔记

## 1. 问题
- DDP 中哪些状态在每个 rank 上重复?
- activation / model states / temporary buffer / fragmentation 的边界是什么?

## 2. 三阶段
- ZeRO-1: 切什么? 每卡剩什么? 通信多了什么?
- ZeRO-2: 切什么? 为什么 ReduceScatter 变重要?
- ZeRO-3: 切什么? 为什么需要参数 all-gather?

## 3. 代价
- 哪些场景 ZeRO-3 会慢?
- prefetch / bucket / overlap 在解决什么?

## 4. 我的结论
- 如果模型能用 DDP 放下,我会不会用 ZeRO?
- 什么时候选 ZeRO-2,什么时候必须 ZeRO-3?

Day 23 会专门做 DeepSpeed 实战。今天不追求完整训练大模型, 只要能读懂 zero_optimization 配置, 并用同一个小模型对比 stage 1/2/3 的显存日志即可。 如果本机没有多卡,先把配置和命令整理好,在有 GPU 的环境执行。

三个配置文件

// ds_zero1.json
{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 1,
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 1
  },
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 1e-4 }
  }
}

// ds_zero2.json
{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 1,
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "reduce_scatter": true,
    "allgather_partitions": true,
    "overlap_comm": true,
    "contiguous_gradients": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 1e-4 }
  }
}

// ds_zero3.json
{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 1,
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto"
  },
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 1e-4 }
  }
}

运行方式

# 环境检查
ds_report

# DeepSpeed 启动,等价于替你处理分布式 launcher 和 engine 初始化
deepspeed --num_gpus=4 train_tiny_transformer.py --deepspeed ds_zero1.json
deepspeed --num_gpus=4 train_tiny_transformer.py --deepspeed ds_zero2.json
deepspeed --num_gpus=4 train_tiny_transformer.py --deepspeed ds_zero3.json

# 如果用 HuggingFace Trainer,通常是把 ds_config 传给 TrainingArguments
# deepspeed="ds_zero2.json"

1. 以 BF16 + AdamW 为例,列出每个参数在训练中对应的参数、梯度、master weight、m/v 状态,估算 DDP 每卡模型状态显存。
2. 用一张表说明 ZeRO-1、ZeRO-2、ZeRO-3 分别切分 P/G/O 中的哪些部分,每卡状态显存如何随 world size 变化。
3. 解释为什么 ZeRO-2 更依赖 ReduceScatter,ZeRO-3 更依赖 AllGather。
4. 读 ZeRO 论文后,用自己的话总结 model states、activation、temporary buffer、fragmentation 四类显存的区别。
5. 设计一个实验比较 DDP、ZeRO-1、ZeRO-2、ZeRO-3:你会记录哪些指标,如何判断 stage 选择是否合理?
6. 列出 3 个 ZeRO-3 可能比 ZeRO-2 更慢的原因,并说明如何用 profiler 或 DeepSpeed 日志验证。