ZeRO-3 的关键不是"参数被切了"这句话,而是参数在需要时会被临时 gather 到 GPU 上,用完再释放。 每个 data parallel rank 长期只保存参数、梯度、优化器状态的一部分。 forward 到某层前,DeepSpeed 把这一层需要的参数 all-gather 出来; 计算结束后释放完整参数,只保留 shard。 backward 时同理,梯度通过 reduce-scatter 回到各自 shard。

ZeRO-3 与 Offload 的收益/代价

最常见、也最稳的 offload 入口是先把 optimizer state 放到 CPU。 Adam 的 optimizer state 通常包含 fp32 master weight、momentum、variance, 比模型参数本体还要重。把这部分搬出 GPU,能显著降低显存占用。 但 optimizer step 会吃 CPU 和 PCIe,吞吐会下降。

{
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 8,
  "steps_per_print": 50,

  "bf16": {
    "enabled": true
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5,
      "betas": [0.9, 0.95],
      "eps": 1e-8,
      "weight_decay": 0.1
    }
  },

  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 500000000,
    "stage3_prefetch_bucket_size": 50000000,
    "stage3_param_persistence_threshold": 100000,
    "stage3_max_live_parameters": 1000000000,
    "stage3_max_reuse_distance": 1000000000,

    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  },

  "gradient_clipping": 1.0,
  "wall_clock_breakdown": true
}

关键字段解释

NVMe Offload 是 ZeRO-Infinity 路线里的强力武器: 当 CPU 内存都不够,或想训练远超 GPU/CPU 常驻能力的模型时, 可以把参数或 optimizer state 放到本地 NVMe。 但这不是免费的魔法。SSD 带宽、IOPS、队列深度、文件系统、NUMA、临时目录空间都会进入训练性能模型。 所以它适合"为了放下更大模型"的场景,不适合作为默认高吞吐方案。

{
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 16,

  "bf16": { "enabled": true },

  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,

    "offload_param": {
      "device": "nvme",
      "nvme_path": "/local_nvme/deepspeed_param",
      "pin_memory": true,
      "buffer_count": 5,
      "buffer_size": 100000000,
      "max_in_cpu": 1000000000
    },

    "offload_optimizer": {
      "device": "nvme",
      "nvme_path": "/local_nvme/deepspeed_optimizer",
      "pin_memory": true,
      "buffer_count": 4,
      "fast_init": true
    },

    "stage3_prefetch_bucket_size": 50000000,
    "stage3_param_persistence_threshold": 10000,
    "stage3_max_live_parameters": 500000000,
    "stage3_max_reuse_distance": 500000000
  },

  "wall_clock_breakdown": true
}

NVMe 字段解释

DeepSpeed 集成有两条常见路径: 一条是用 DeepSpeed launcher 和原生训练脚本; 另一条是 HuggingFace Trainer / Accelerate 读取 ds_config.json。 不管哪条路径,第一轮都不要追求最大模型。先用小模型、小数据、十几个 step 跑通, 看日志确认 ZeRO stage、offload device、显存曲线和 step time。

DeepSpeed launcher

# 单机 4 卡示例
deepspeed --num_gpus=4 train.py \
  --deepspeed \
  --deepspeed_config ds_zero3_cpu.json \
  --model_name_or_path your-model \
  --per_device_train_batch_size 2 \
  --gradient_accumulation_steps 8

# 多机时使用 hostfile 或调度系统注入的环境变量,并确保每台机器都能访问本地 nvme_path
deepspeed --hostfile hostfile train.py --deepspeed --deepspeed_config ds_zero3_nvme.json

HuggingFace Trainer / Accelerate

# Trainer 常见用法
torchrun --nproc_per_node=4 train_hf.py \
  --deepspeed ds_zero3_cpu.json \
  --bf16 true \
  --per_device_train_batch_size 2 \
  --gradient_accumulation_steps 8

# Accelerate 常见用法
accelerate launch --config_file accelerate_config.yaml train_hf.py \
  --deepspeed ds_zero3_cpu.json

第一次跑只看这几件事

ZeRO-3 + Offload 的问题通常不是单点 bug,而是显存、通信、CPU 内存、磁盘 IO 的多方拉扯。 调参时要保留 baseline,一次只动一类参数。 最坏的方式是看到 OOM 就把所有 bucket 调小、看到慢就全部调大,最后不知道哪一项有效。

推荐调参顺序

# 0. 无 offload baseline:先确认模型、数据、loss 都正常
ds_zero3_no_offload.json

# 1. 只开 optimizer CPU offload:通常收益/风险比最好
offload_optimizer.device = "cpu"
offload_param = # 不配置

# 2. GPU 仍 OOM,再尝试 param CPU offload
offload_param.device = "cpu"
stage3_prefetch_bucket_size ↓
stage3_param_persistence_threshold ↓

# 3. CPU 内存也不够,最后考虑 NVMe
offload_param.device = "nvme"
offload_optimizer.device = "nvme"
nvme_path = "/local_nvme/..."

# 4. 每一步都记录
peak_gpu_mem, cpu_mem, nvme_util, step_time, tokens_per_sec

今天请输出一份可以复用的 DeepSpeed 实验笔记。 它不需要很长,但必须有三组对照:ZeRO-3 无 offload、ZeRO-3 + optimizer CPU offload、 ZeRO-3 + param/optimizer offload。每组记录显存、吞吐和瓶颈判断。

实验记录模板

检查清单

1. ZeRO-3 中参数 shard 在 forward 某层前后分别处于什么状态?
2. offload_optimizer 与 offload_param 分别搬走了什么?
3. CPU Offload 为什么通常比 NVMe Offload 更适合作为第一选择?
4. stage3_prefetch_bucket_size 调大和调小分别影响什么?
5. 为什么 pin_memory 能加速 CPU/GPU 传输,又可能带来内存压力?
6. 如何判断训练慢是 GPU compute 慢、NCCL 通信慢,还是 NVMe IO 慢?
7. 保存 ZeRO-3 checkpoint 时为什么要区分 sharded checkpoint 和完整 16-bit model?