DeepSpeed ZeRO 优化

<-->

461人浏览 · 2026-04-21 13:49:21

<--> · 2026-04-21 13:49:21 发布

一、为什么需要 ZeRO？（背景动机）

训练大模型时，GPU 显存主要被以下四部分占用：

类别	占用	说明
模型参数 (fp16)	$2Ψ2\Psi$ bytes	$Ψ\Psi$ = 参数量
梯度 (fp16)	$2Ψ2\Psi$ bytes	与参数同形
优化器状态 (fp32)	$12Ψ12\Psi$ bytes	Adam: fp32参数 + 一阶矩 + 二阶矩
合计	$16Ψ16\Psi$ bytes	GPT-3 ( $175 B$ ) ≈ 2.8 TB

核心矛盾：数据并行 (DDP) 每张卡都保存完整副本 → 显存严重冗余。

二、ZeRO 三个阶段核心思想

$Optimizer\text{ZeRO} = \text{Zero Redundancy Optimizer}$

阶段	分片内容	每卡显存（ $N$ 卡）
Stage 1	优化器状态	$4Ψ+12ΨN4\Psi + \frac{12\Psi}{N}$
Stage 2	优化器状态 + 梯度	$2Ψ+14ΨN2\Psi + \frac{14\Psi}{N}$
Stage 3	优化器状态 + 梯度 + 参数	$16ΨN\frac{16\Psi}{N}$

关键洞见：通信量与 DDP 相比不增加（Stage 1/2），Stage 3 会略有增加但换来了线性的显存扩展性。

三、Stage 1/2 实现细节（`stage_1_and_2.py`）

3.1 数据结构

bit16_groups              # fp16 参数，按 param_group 组织
single_partition_of_fp32_groups  # 本 rank 负责的 fp32 主权重分片
parallel_partitioned_bit16_groups # 所有 rank 的 fp16 分片

3.2 前向 → 反向 → 优化完整流程

Forward
  └── 正常前向传播（fp16 参数）

Backward
  └── 钩子触发 reduce_independent_p_g_buckets_and_remove_grads(param)
        ├── 攒桶 (IPG Bucket)
        ├── 桶满 → reduce_ipg_grads()
        │     ├── allreduce (Stage 1) 或 reduce_scatter (Stage 2)
        │     └── 每卡只保留自己分片的梯度
        └── 梯度写入 fp32 grad buffer

Step (optimizer.step)
  ├── scaled_global_norm()  # 梯度裁剪用的全局 L2 norm
  ├── _update_scale()       # 动态 loss scaling 更新
  ├── 循环每个 sub_group:
  │     ├── unscale_and_clip_grads()
  │     ├── _optimizer_step()    # 更新本 rank 的 fp32 分片
  │     └── 将更新后的 fp32 → fp16（copy_fp32_to_fp16）
  └── allgather → 恢复完整 fp16 参数

3.3 关键方法解析

get_flat_partition：把一组 tensor 的梯度按照分区偏移切出本 rank 所需的那一段，拼成连续 flat buffer，不足部分补零。

create_gradient_handling_hooks：为每个参数注册 register_grad_hook，梯度产生时立即触发桶聚合逻辑，实现 overlap_comm（计算与通信重叠）。

allreduce_no_retain vs reduce_scatter：

Stage 1：allreduce 后每卡仅保留自己分片
Stage 2：直接 reduce_scatter，通信更高效

四、Stage 3 实现细节（`stage3.py`）

4.1 参数生命周期（最重要！）

参数状态机:
  NOT_AVAILABLE  ──allgather──►  AVAILABLE
       ▲                              │
       └─────────partition────────────┘

前向前：allgather 拉取完整参数
前向后：立即 partition（释放非本 rank 部分）
反向前：再次 allgather（用于计算梯度）
反向后：reduce_scatter 梯度 + partition 参数

这由 DeepSpeedZeRoOffload / PartitionedParameterCoordinator 通过 forward/backward hook 自动管理。

4.2 prefetch 机制

prefetch_bucket_size  # 预取阈值
max_reuse_distance    # 多远的参数值得预取

在执行当前层时，异步预取下一层参数，隐藏 allgather 延迟。

4.3 CPU/NVMe Offload

offload_param:      参数卸载到 CPU/NVMe（前向时再拉回 GPU）
offload_optimizer:  优化器状态卸载到 CPU（用 DeepSpeedCPUAdam 在 CPU 上 step）

DeepSpeedCPUAdam 用多线程 AVX 指令在 CPU 上执行 Adam，比 torch.optim.Adam 快 5x。

4.4 IPGBucketZ3 数据结构

@dataclass
class IPGBucketZ3:
    buffer: Optional[torch.Tensor]       # 连续通信 buffer
    buffer_meta: Optional[torch.Tensor]  # 元信息
    params: List[torch.Tensor]           # 本桶中的参数
    elements: int                        # 已累积元素数

攒桶策略：元素数超过 reduce_bucket_size 时触发 reduce_scatter。

五、通信原语对照表

操作	含义	ZeRO 使用场景
`all_reduce`	所有卡求和，结果广播给所有卡	Stage 1 梯度聚合
`reduce_scatter`	所有卡求和，结果按分区分发	Stage 2/3 梯度聚合
`all_gather`	每卡贡献自己的分片，拼成完整张量	Stage 3 参数恢复
`reduce`	求和结果只发给指定 rank	特定 rank 聚合

六、Loss Scaling（混合精度训练）

问题：fp16 动态范围小（最小正数 ≈ $6×10−56\times10^{-5}$ ），梯度容易下溢为 0。

解法：
$scaled_loss=loss×S\text{scaled\_loss} = \text{loss} \times S$

反向传播得到 $\times \text{grad}$ ，更新前除以 $S$ 。

动态 Loss Scaling 策略（CreateLossScaler）：

连续 $N$ 步无溢出 → $\times 2$ （上调）
检测到 inf/nan → $S /2$ （下调），跳过本步更新

七、梯度裁剪（Gradient Clipping）

# 全局 L2 norm（跨所有 rank all_reduce）
scaled_global_grad_norm = self.scaled_global_norm()
global_grad_norm = scaled_global_grad_norm / loss_scale

# 裁剪
clip_coef = max_norm / (global_grad_norm + 1e-6)
if clip_coef < 1.0:
    grad *= clip_coef

ZeRO 难点：每张卡只有梯度的分片，必须先 all_reduce 各自的局部 norm 平方和，才能得到全局 norm。

八、ZeRO++ / HPZ（常见进阶问题）

特性	说明
`zero_hpz_partition_size`	在节点内先做参数分片（HPZ），减少跨节点通信
`zero_quantized_weights`	allgather 时用量化权重，降低通信量
`all_to_all_quant_reduce`	量化 reduce_scatter，进一步压缩梯度通信

九、常见面试题 & 标准答案

Q1：ZeRO Stage 2 和 DDP 的通信量有何区别？

DDP 用 all_reduce，通信量 $2\Psi$ （reduce + broadcast）。ZeRO-2 用 reduce_scatter + all_gather，梯度阶段 $Ψ\Psi$ （reduce_scatter），参数同步阶段 $Ψ\Psi$ （all_gather），总量相同，但 Stage 2 每卡不保存完整梯度，节省显存。

Q2：ZeRO Stage 3 前向传播时是否需要 allgather？

是的。Stage 3 每张卡只保存参数的 $1N\frac{1}{N}$ 分片。前向传播每层开始前必须 allgather 完整参数，层计算完后立即释放非本卡分片，通过 prefetch 隐藏延迟。

Q3：为什么 ZeRO-Offload 要用专门的 DeepSpeedCPUAdam？

标准 torch.optim.Adam 在 CPU 上是单线程的。DeepSpeedCPUAdam 用 C++ 多线程 + AVX-512 SIMD 指令并行更新参数，速度提升 5-7x，是 CPU offload 实用化的关键。

Q4：overlap_comm 是如何实现计算通信重叠的？

通过为每个参数注册梯度钩子，梯度一产生就立即触发桶聚合和 reduce，用独立的 CUDA stream（reduction_stream）执行通信，与下一层的反向计算在不同 stream 上并发执行。

Q5：ZeRO Stage 3 和模型并行（Tensor Parallel）有何区别？

ZeRO-3 沿数据并行维度分片参数，每张卡有完整的计算图，通过通信恢复完整参数；Tensor Parallel 沿参数内部维度切分（如矩阵按行/列切分），每张卡只计算矩阵乘法的一部分，需要在算子内部通信。两者可以组合使用（3D并行）。

Q6：如何判断是否发生梯度溢出？

通过 has_overflow_serial 检查 fp16 梯度中是否存在 inf 或 nan。若溢出则跳过本步参数更新（不调用 optimizer.step），同时将 loss scale 减半。

十、一句话总结

“ZeRO 的本质是将数据并行训练中优化器状态、梯度、参数三类冗余沿 rank 维度分片存储，通过 reduce_scatter/allgather 在需要时重建，将显存从 $O(16Ψ)O(16\Psi)$ 降至 $O(16Ψ/N)O(16\Psi/N)$ ，同时保持与 DDP 相当的通信量，是目前训练千亿级大模型最主流的显存优化方案。”

那不应该啊, 如果没有这个技术, 难道openai训练的时候一张内存卡非常大吗?

GitHub Copilot: 这是个很好的问题！在 ZeRO 出现之前，业界确实用了多种"土方法"，并不是靠单张超大显存卡硬撑的。

ZeRO 之前的解决方案

1. 模型并行（Model Parallelism）—— 最早的主流方案

按层切分（Pipeline Parallelism）：

GPU 0: Layer 1-8    → 前向传给 GPU 1
GPU 1: Layer 9-16   → 前向传给 GPU 2
GPU 2: Layer 17-24  → 输出 loss
         反向时逐级传回

GPT-2（2019）、GPT-3（2020）都大量使用了这种方式，OpenAI 内部用的是 Megatron-LM（NVIDIA 做的）。

2. Tensor Parallelism —— Megatron 的核心贡献

把单个矩阵乘法切开：

$\quad \Rightarrow \quad Y = X[W_1 | W_2] = [XW_1 | XW_2]$

每张卡只算一半的矩阵乘，再 all_reduce 拼结果。GPT-3 就是用 8路 Tensor Parallel + 多节点 训练的。

3. 梯度检查点（Gradient Checkpointing）

以时间换空间：前向传播时不保存中间激活值，反向时重新计算。显存减少 $L\sqrt{L}$ 倍（ $L$ = 层数），但训练时间增加约 33%。

4. 混合精度训练（fp16）

把参数和激活从 fp32 变成 fp16，显存直接减半，这是 2018 年就普及的基础技术。

GPT-3 实际是怎么训的？

GPT-3（1750亿参数，2020年）用了：

技术	作用
Tensor Parallel（8路）	单节点内切分参数
Pipeline Parallel（多节点）	跨节点切分层
fp16 混合精度	减少显存和通信量
硬件：V100 80GB * 数千张	暴力堆卡

ZeRO 论文是 2020年 3月 才发布的，和 GPT-3 几乎同期，所以 GPT-3 训练时 ZeRO 还没被广泛采用。

ZeRO 的真正革命性在哪？

对比维度	模型并行	ZeRO
代码改动	大，需要重写模型结构	几乎为零，对用户透明
通信效率	层间串行流水，有气泡	与 DDP 通信量相当
适用场景	超大模型必用	中大模型首选
工程复杂度	高	低（一行配置开启）

一句话：ZeRO 之前不是靠"一张超大卡"，而是靠模型并行 + 堆卡数量 + 大量工程改造，门槛极高，只有 OpenAI/Google/NVIDIA 这类顶级机构才能玩得起。ZeRO 的价值在于让普通研究者也能用数据并行的简单方式训练大模型。

nathon-lee: 相当于 megratron 按照模型的维度切, 也可以将这些参数分散到多个GPU上, 那为什么还诞生了deepspeed呢?

GitHub Copilot: 这个问题问到点子上了！让我们从根本上对比两者。

Megatron 切的是"参数"，但没解决"状态"问题

Megatron Tensor Parallel 确实把参数切到了多卡，但训练时显存不只是参数：

$每卡还是完整的！\text{总显存} = \underbrace{\text{参数}}_{\text{Megatron 切了}} + \underbrace{\text{梯度} + \text{优化器状态}}_{\text{❌ 每卡还是完整的！}}$

以 Adam 优化器为例，每个参数对应 3 份 fp32 数据：

fp32 主权重
一阶矩 $m$
二阶矩 $v$

Megatron 把参数切了，但每张卡上这 3 份优化器状态依然是针对自己那份参数的完整副本，并没有进一步节省。

更本质的区别：切法不同，目的不同

	Megatron-LM	DeepSpeed ZeRO
切分维度	模型结构内部（矩阵行/列）	数据并行维度（按 rank 分片）
主要解决	单层参数放不下	优化器状态/梯度冗余
代码侵入性	高：必须改模型代码，每个算子都要适配	低：对模型透明，一行配置
适用模型	必须是 Transformer 特定结构	任意模型
通信时机	算子内部强制同步（`all_reduce`）	梯度聚合时（可异步）
能否单独用	能	能

一个具体例子感受差异

假设训练一个 10B 参数模型，用 8 张 80GB A100：

只用 Megatron（8路 Tensor Parallel）：

每卡参数量   = 10B / 8 = 1.25B × 2字节(fp16) = 2.5 GB  ✅
每卡优化器   = 1.25B × 12字节(fp32 Adam)     = 15 GB   ❌ 还是很大
每卡梯度     = 1.25B × 2字节                  = 2.5 GB
合计 ≈ 20 GB（还没算激活值）

只用 ZeRO-3（8路数据并行）：

每卡参数量   = 10B / 8 × 2字节  = 2.5 GB   ✅
每卡优化器   = 10B / 8 × 12字节 = 15 GB   ✅  (分片了！)
每卡梯度     = 10B / 8 × 2字节  = 2.5 GB  ✅
合计 ≈ 20 GB（显存节省效果更均衡）

那为什么现在两个都在用？

因为它们解决的是互补的问题，可以叠加：

3D 并行 = Tensor Parallel (Megatron)
         + Pipeline Parallel
         + ZeRO (DeepSpeed)

Megatron 解决：单层矩阵太大，一张卡放不下
ZeRO 解决：优化器状态冗余，N 张卡存了 N 份一样的东西

一句话总结：Megatron 是"把大块肉切小"，ZeRO 是"把重复买的东西退掉"。一个模型可能同时需要两种优化，这就是 DeepSpeed 诞生并与 Megatron 共存的根本原因。

nathon-lee: 这些参数难道不是每张GPU上都需要的吗? 如果打散到各个GPU上了, 那某个GPU需要用到其他的参数, 怎么办呢?

GitHub Copilot: 这正是 ZeRO-3 最核心的设计！答案就是 “用的时候再去要，用完立刻还”。

直觉类比：图书馆借书模式

想象 8 个人（GPU）共同写一篇论文：

DDP 模式：每人买一套完整的参考书（8份完全一样）→ 极度浪费
ZeRO-3 模式：每人只保管自己那1/8的书，需要某章节时去借，用完还回去

具体机制：AllGather + 用完即弃

正常状态（平时）：
  GPU 0 只存  Layer 0-2 的参数分片
  GPU 1 只存  Layer 3-5 的参数分片
  ...

执行 Layer 3 前向时：
  ┌─────────────────────────────────────────┐
  │  AllGather：所有 GPU 广播自己的分片      │
  │  → 每张 GPU 临时拼出 Layer 3 完整参数   │
  │  → 执行计算                             │
  │  → 计算完立刻释放！（非本卡分片删掉）   │
  └─────────────────────────────────────────┘

用代码里的状态机来说（stage3.py 中）：

参数状态：
  PARTITIONED  ──── AllGather ────►  AVAILABLE（临时）
      ▲                                    │
      └──────── 计算完成，立刻 partition ───┘

完整训练一步的参数流转

前向传播（逐层执行）：
  Layer 1: AllGather → 计算 → 释放
  Layer 2: AllGather → 计算 → 释放
  Layer 3: AllGather → 计算 → 释放
  ...

反向传播（逐层反向）：
  Layer N: AllGather → 算梯度 → ReduceScatter梯度 → 释放参数
  ...
  Layer 1: AllGather → 算梯度 → ReduceScatter梯度 → 释放参数

优化器更新：
  每张卡只更新自己负责的那 1/N 参数分片（不需要通信）

那通信开销不会很大吗？

确实会增加，但 ZeRO 用两个手段来隐藏延迟：

1. Prefetch 预取

# stage3.py 中的参数
prefetch_bucket_size = 50_000_000  # 提前预取下一层参数

# 执行 Layer 3 计算时，同时异步 AllGather Layer 4 的参数
# 计算和通信并行 → 延迟被隐藏

2. 反向时复用

前向时已经 AllGather 过的参数，
如果 max_reuse_distance 内还会用到（反向传播），
就暂时不释放，避免重复通信

和 Megatron 的本质区别再次体现

	Megatron Tensor Parallel	ZeRO-3
何时需要其他卡数据	每个算子内部，强制同步	每层开始前，可异步预取
通信粒度	细（每个矩阵乘都要通信）	粗（每层一次 AllGather）
能否overlap	难（算子内强依赖）	容易（层间有时间窗口）

一句话：ZeRO-3 的参数确实被打散了，但通过 AllGather"借来用，用完还" 的机制保证正确性，再通过 Prefetch 预取把通信延迟藏在计算里，所以实际训练速度损失远比想象中小。

最常见场景：1 进程 = 1 GPU = 1 rank，三者等价 ✓

严格来说：
  rank  = 进程在某个通信组内的编号（逻辑概念）
  GPU   = 物理硬件
  ZeRO 分片用的 rank = dp_process_group 内的编号
  
一张 GPU 上也可以跑多个进程（CPU offload 极端情况），
多张 GPU 也可以属于同一个 rank 所在的通信组。