DeepSpeed ZeRO：训练显存优化的三阶段原理与实践

说明：本篇为综述/无第一手实验数据，采用公开文档与机制解析路线。

实验目标

大模型训练绕不开显存墙：70B 模型 BF16 权重 140GB，加上梯度、优化器状态（Adam 双动量），单卡 80GB 完全装不下。09 章的 FSDP 篇讲了一条路；本篇讲另一条主流路线 DeepSpeed ZeRO：它和 FSDP 谁是谁的"平替"？三阶段分别切了什么、省了多少？Offload 到 CPU/NVMe 的边界在哪？目标：建立 ZeRO 的完整认知——三阶段原理、与 FSDP 的对应关系、训练脚本接入方式，以及"该用哪一档"的选型判断。

代码逐段拆解

第一段：训练脚本接入 ZeRO

import deepspeed

# 原有 PyTorch 训练脚本改造：一行初始化 + engine 替换
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=params,
    config_params=ds_config,
)
for step, batch in enumerate(train_loader):
    loss = model_engine(batch)          # 前向
    model_engine.backward(loss)         # 反向 + ZeRO 内部通信
    model_engine.step()                 # 优化器更新

关键点：

1. 接入成本低：deepspeed.initialize 返回的 model_engine 替代原生 model，前向/反向/step 三调用不变——ZeRO 的通信编排（All-Gather/Reduce-Scatter）都封装在 engine 内部。
2. config_params 是全部魔法所在：零代码改动之外的调优全在配置文件里——优化器类型、ZeRO 档位、Offload 开关、通信后端。

第二段：ZeRO 配置的三档

# deepspeed_config.json：ZeRO 三档配置示例
{
  "train_batch_size": 256,
  "optimizer": { "type": "AdamW", "params": { "lr": 3e-5 } },
  "zero_optimization": {
    "stage": 2,                        # 1 / 2 / 3 三档
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "allgather_partitions": true,
    "reduce_scatter": true,
    "overlap_comm": true               # 通信与计算重叠
  },
  "communication_data_type": "bf16"
}

• stage 1：只切优化器状态（Adam 双动量占显存大头，约 2×权重）。
• stage 2：再切梯度（梯度 = 权重大小），每卡只留自己分片对应的梯度。
• stage 3：连权重也分片（类似 FSDP 的全分片），显存占用降到"单层可装"级别。
• offload_optimizer 到 CPU：优化器状态搬 CPU 内存，GPU 显存进一步省，代价是 CPU↔GPU 传输开销。
• overlap_comm：梯度通信与下一段计算重叠，抵消一部分分片通信延迟。

原理讲解

三阶段切什么、省多少

以 70B、BF16、AdamW 为例（估算口径）：

项	全量（每卡）	ZeRO-1	ZeRO-2	ZeRO-3
优化器状态（Adam）	2×权重	1/DP×	1/DP×	1/DP×
梯度	1×权重	1×	1/DP×	1/DP×
权重	1×权重	1×	1×	1/DP×
合计（DP=8）	4×权重	~3×	~1.5×	~0.5×+单层

（DP=数据并行度；每项按"每卡占用/权重倍数"计。）

理解要点：

• 优化器状态是大头：Adam 存 fp32 主权重 + 两动量 ≈ 2× 权重，还常是 fp32（16 字节/参数）。stage 1 切掉它，显存立刻减一半。
• stage 3 是"训练时只装一层"：权重分片后，单卡只需装当前计算层，理论上可训任意大模型——这与 FSDP 的 Full Shard 等价，两者机制几乎相同（FSDP 后来成为 PyTorch 原生实现，ZeRO-3 是其 DeepSpeed 实现）。
• 省显存的代价是通信：分片越多，All-Gather/Reduce-Scatter 通信量越大；ZeRO-2 比 ZeRO-3 通信少，能装下时优先低档位。

ZeRO vs FSDP：不是二选一

维度	ZeRO-3（DeepSpeed）	FSDP（PyTorch 原生）
机制	分片权重/梯度/优化器态	同左（Full Shard）
接入	deepspeed.initialize	torch.distributed.fsdp.FullyShardedDataParallel
生态	DeepSpeed 全家桶（Offload/梯度压缩）	PyTorch 原生、与 torch.compile 配合
适用	已有 DeepSpeed 依赖/需要 Offload	纯 PyTorch 项目优先

选型判断：纯 PyTorch 新项目优先 FSDP（原生、稳）；需要 CPU/NVMe Offload 或与 DeepSpeed 生态（MoE 训练等）配合时用 ZeRO。两者机制同源，迁移成本低。

Offload 的边界

offload_optimizer: device: cpu 把优化器状态搬 CPU 内存——代价是每步更新都要 CPU↔GPU 传一轮，训练吞吐可能掉 10%-30%。Offload 是"显存不够时的补救"，不是默认选项；NVMe Offload（ZeRO-Infinity）更进一步，但 IO 延迟更大，只适合超大模型+低速训练场景。判断标准：显存够（可用低档位或非 Offload 高档位）就绝不 Offload。

实验结果

（估算口径，非实测；DP=8、70B、BF16 单卡显存需求）

方案	单卡显存估算	通信开销	适用
DDP 全量	~4×权重（560GB+）	无	小模型
ZeRO-2	~1.5×权重（~200GB）	中	常见生产档
ZeRO-3	~0.5×权重+单层	高	超大模型
ZeRO-2 + CPU Offload	权重+少量	中+传输	显存紧张

工程启示

1. 先算"差多少"再选档位：用上面的倍数表算当前显存缺口，能 ZeRO-2 解决就不上 ZeRO-3——分片档位越高通信越重，吞吐收益递减。
2. Offload 是补救不是标配：显存够就别 Offload；不够时优先 CPU Offload 优化器态（影响最小），NVMe Offload 是最后手段。
3. overlap_comm 必开：通信与计算重叠对分片训练的吞吐影响很大，是 ZeRO/FSDP 配置里性价比最高的一个开关。
4. 与 Checkpoint/稳定性篇联动：分片训练下保存的是"分片快照"（每卡存自己的分片），恢复时要聚合——10 章稳定性篇的断点重启需按 ZeRO 分片格式适配。

下一篇预告：训练侧的显存问题解决后，推理侧的对应物已经讲过（12 章 KV Cache）——下一步回到应用章，拆解"科学智能体：从工具调用到工作流自动化"的落地工程。