1. 课程导入

当你尝试训练一个百亿甚至千亿参数的大模型时，最先撞上的往往不是算法问题，而是显存墙。一张 80GB 的 A100，在混合精度下训练 175B 参数的 GPT-3 级别模型，仅模型参数、梯度和优化器状态三份数据就需要约 2.4TB 显存——这远超单卡容量，甚至超过单机 8 卡的聚合显存。DeepSpeed 的 ZeRO（Zero Redundancy Optimizer，零冗余优化器）正是为击穿这堵墙而生的核心技术。

本课程将用 30 个章节，系统拆解 ZeRO 的三个优化阶段，从原理到代码、从配置到排障，帮助你真正掌握大模型分布式训练的显存优化方法。

一、为什么需要 ZeRO：显存墙的本质

在数据并行（DP）训练中，每张 GPU 都保存一份完整的模型副本。以 Adam 优化器、混合精度训练为例，每个参数对应的显存开销大致如下：

问题的关键在于“冗余”：数据并行中每张卡都持有完全相同的这份 2.8TB 数据，而真正被并行利用的只有计算能力，存储却被重复了 N 份。ZeRO 的核心思想就是——把冗余的模型状态切分到各张卡上，需要时再通过通信临时聚合。

二、ZeRO 三阶段的核心原理

ZeRO 将模型状态分为三块：优化器状态、梯度、参数。按切分范围从小到大，形成三个阶段：

• ZeRO-1（优化器状态切分）：只切分优化器状态（含 FP32 主权重），显存降至约 1/4，通信量与标准 DP 相当，几乎无额外开销，是最容易落地的阶段。
• ZeRO-2（梯度 + 优化器状态切分）：在 ZeRO-1 基础上再切分梯度，显存降至约 1/8，通信量仍接近 DP，是性价比最高的常用配置。
• ZeRO-3（参数 + 梯度 + 优化器状态切分）：三块全部切分，显存降至约 1/16，可训练超大模型，但参数在前向/反向时需要 All-Gather，通信量显著增加，需配合通信重叠优化。

下图展示了三个阶段中，模型状态在各 GPU 上的分布变化：

三、方法流程：ZeRO 如何工作

以 ZeRO-2 为例，一次训练迭代的流程可以概括为：

1. 前向传播：每张卡用完整参数计算自己负责的微批次数据，得到损失。
2. 反向传播：计算梯度，但每张卡只保留自己负责切分的那部分梯度，其余梯度通过 Reduce-Scatter 直接归约到对应卡上，避免全量梯度驻留。
3. 优化器更新：每张卡只更新自己持有的优化器状态和对应参数分片。
4. 参数同步：通过 All-Gather 将各卡更新后的参数分片聚合，恢复完整参数供下一轮前向使用。

在 DeepSpeed 中，只需在配置文件中指定阶段即可启用：

{
  "train_batch_size": 64,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

其中 overlap_comm 让通信与计算重叠，contiguous_gradients 减少显存碎片，是生产环境常用的两个开关。

四、典型场景与选型建议

不同规模模型对 ZeRO 阶段的需求差异明显：

• 1B～7B 模型：单卡或少量卡即可容纳，ZeRO-1 通常足够，通信开销最小。
• 7B～70B 模型：ZeRO-2 是主流选择，显存与通信平衡最佳。
• 70B 以上模型：需要 ZeRO-3，并常与张量并行、流水并行组合使用。

五、常见误区与复盘

初学者最容易踩的坑有三个：

• 误区一：阶段越高越好。 ZeRO-3 通信量最大，若模型本可被 ZeRO-2 容纳，强行上 ZeRO-3 反而降低吞吐。
• 误区二：忽略 batch size 与梯度累积。 ZeRO 切分的是模型状态，不改变全局 batch 语义，配置时需同步调整 train_batch_size 与 gradient_accumulation_steps。
• 误区三：只看显存不看通信。 在多机低带宽环境下，ZeRO-3 的 All-Gather 可能成为瓶颈，需开启通信重叠并评估网络拓扑。

本课程后续将逐阶段深入：先讲清 ZeRO-1 的优化器状态切分与通信原语，再进入 ZeRO-2 的梯度归约，最后攻克 ZeRO-3 的参数分片与显存-通信权衡，并配套可运行的配置与排障案例。掌握这三阶段，你就拥有了训练大模型的显存钥匙。