ZeRO

ZeRO 在数据并行 worker 之间切分原本重复保存的训练状态，以降低单卡显存。

阶段心智模型

  • Stage 1：切分 optimizer states。
  • Stage 2：再切分 gradients。
  • Stage 3：再切分 parameters，并在需要时通信聚合。

阶段越高，单卡内存通常越低，但通信、调度和实现复杂度可能增加。具体行为应以当前 DeepSpeed 配置与官方文档为准。

与其他层的关系

  • Data Parallelism 定义 worker 处理不同数据。
  • DeepSpeed 提供 ZeRO 等训练引擎能力。
  • Ray 侧重任务/角色编排，不是 ZeRO 的替代品。

三阶段的因果取舍

1
2
3
4
5
6
Stage 1: optimizer states 分片
→ 省下 Adam 状态
Stage 2: + gradients 分片
→ 再省梯度副本，但 backward 需要同步
Stage 3: + parameters 分片
→ 单卡可容纳更大模型，但 forward/backward 频繁 all-gather

阶段升级应由显存账本驱动，而不是盲目追求 Stage 3。若通信已经主导 step time，增加切分可能让吞吐下降；若 checkpoint、参数 offload 或故障恢复不兼容，训练可运行也不代表可运维。记录 stage、bucket、offload、通信后端和 checkpoint 格式，才能复现结果。

相关节点：Training Memory Accounting、DeepSpeed。参考：ZeRO 论文。

Stage 2 和 Stage 3 的运行时差别

Stage 2 仍让每个 rank 持有完整参数，但只保存自己负责的梯度和 optimizer state；参数在 forward/backward 中不需要频繁 gather。Stage 3 连参数也切开，某层即将执行时 all-gather 它的权重，算完后再释放或重新分片。于是 Stage 3 的收益是模型规模边界被推开，代价是参数通信进入每个 layer 的关键路径。

这也改变了 checkpoint 的语义。某个 rank 保存的不是完整模型，恢复时需要按相同或兼容的分片元数据重组；直接把普通 DDP state dict 塞进 Stage 3 配置，可能在加载时隐式 gather，瞬间触发 OOM。上线前用真实恢复流程做一次“保存—杀掉 worker—重新加载—继续训练”测试，不能只验证文件存在。

offload 不是免费的显存

把 optimizer 或参数 offload 到 CPU/NVMe 能降低 GPU 峰值，却引入 PCIe、主机内存和 I/O 带宽瓶颈。若每步都在等待权重搬运，GPU 利用率会下降；增加 stage 反而可能使 step time 变长。记录 all-gather、reduce-scatter、CPU copy、NVMe I/O 的时间占比，才能决定是降低 stage、调整 bucket，还是改变硬件拓扑。

ZeRO 只解决重复训练状态，不会自动解决 sequence activation 或 data pipeline。若 OOM 位于 backward，先看 Training Memory Accounting 中的 activation 项；若卡在参数 gather，则调 stage、bucket 和通信 overlap。