📚 课程目录第3章 ZeRO概述

3、ZeRO概述：ZeRO的动机、与DDP/模型并行的区别、ZeRO的三个阶段（Optimizer State、Gradient、Parameter）的核心思想

随着Transformer模型规模从1亿参数增长到千亿甚至万亿级别，分布式训练面临的核心矛盾不再是"算力不够"，而是"显存放不下"。以 175B为例，仅参数就需要700GB显存（FP16），加上Adam优化器状态（FP32主权重+动量+方差）和梯度，总显存需求超过2.8TB，远超单卡A100 80GB的容量。传统的分布式训练方案——数据并行（DDP）和模型并行——在应对这一挑战时各有局限。DeepSpeed的ZeRO（Zero Redundancy Optimizer）通过系统性地消除显存冗余，在不牺牲计算效率的前提下，将训练超大规模模型的内存瓶颈大幅缓解。本章将深入剖析ZeRO的动机、与传统并行策略的本质区别，以及其三个核心阶段（Optimizer State、Gradient、Parameter）的递进式优化思想。

3.1 ZeRO的动机：显存墙与通信效率的双重挑战

在大规模模型训练中，显存消耗主要来自四个部分：模型参数（Parameters）、梯度（Gradients）、优化器状态（Optimizer States）以及中间激活值（Activations）。以Adam优化器为例，每个参数需要保存：FP16参数副本（2字节）、FP16梯度副本（2字节）、FP32主权重（4字节）、FP32动量（4字节）、FP32方差（4字节），合计每个参数需要16字节。对于一个10B参数的模型，仅参数和优化器状态就需要160GB显存，远超单卡容量。

传统DDP在每个GPU上复制完整的模型副本，虽然计算并行度高，但显存冗余严重。模型并行（包括张量并行和流水线并行）虽然切分了参数，但引入了大量的通信开销，且对层间依赖敏感。ZeRO的核心洞察在于：数据并行中每个GPU持有的是完全相同的参数和优化器状态，这种冗余是可以通过分片（Sharding）消除的。ZeRO的目标是在保持数据并行通信模式（All-Reduce）的基础上，将显存占用从"每卡全量"降为"每卡分片"，同时通过通信调度避免通信成为瓶颈。

  核心动机总结：ZeRO不是要替代数据并行或模型并行，而是通过消除数据并行中的显存冗余，使得数据并行能够扩展到更大的模型规模。它保留了数据并行的计算效率和易用性，同时获得了接近模型并行的显存效率。

3.2 与DDP/模型并行的本质区别

要理解ZeRO的价值，需要先厘清它与传统并行策略的差异。下表从多个维度对比了DDP、模型并行和ZeRO三种方案：

对比维度	DDP（数据并行）	模型并行（张量/流水线）	ZeRO（分片数据并行）
参数存储	每卡全量复制	每卡持有部分参数	每卡持有参数分片（Stage 3）
优化器状态	每卡全量复制	每卡持有对应参数的状态	每卡持有分片状态（Stage 1）
梯度通信	All-Reduce 全量梯度	前向/反向传递中多次通信	Reduce-Scatter + All-Gather
通信量	2 × 参数量（每次迭代）	取决于切分粒度，通常更高	Stage 1/2: 同DDP；Stage 3: 3 × 参数量
显存扩展性	随卡数线性复制，无扩展性	可扩展，但受通信拓扑限制	显存随卡数近似线性下降
实现复杂度	低（PyTorch原生支持）	高（需手动切分和调度）	中（DeepSpeed封装）
适用场景	模型可放入单卡	超大模型，单层超单卡	大模型，且需要高效扩展

从表中可以看出，DDP的瓶颈在于显存冗余，模型并行的瓶颈在于通信复杂度和实现难度。ZeRO巧妙地选择了中间路线：保持数据并行的通信模式，但通过分片消除冗余。在Stage 1和Stage 2中，通信量甚至与DDP完全相同，但显存占用大幅降低；Stage 3虽然增加了参数All-Gather的通信量，但换来了显存的线性扩展能力。

3.3 ZeRO的三个阶段：从优化器状态到参数的全量分片

ZeRO的设计遵循"由易到难、由显存大头到显存小头"的递进逻辑。三个阶段分别针对优化器状态、梯度和参数进行分片，每个阶段都对应不同的显存收益和通信代价。

Stage 1：优化器状态分片（Optimizer State Partitioning）

在Adam优化器中，每个参数需要维护FP32主权重、动量和方差，这三者合计12字节/参数，是显存消耗的最大头。Stage 1的核心思想是：将优化器状态按数据并行度（DP）均匀切分，每个GPU只负责更新自己分片对应的参数。具体流程如下：

前向传播：每个GPU持有完整的参数副本，正常计算前向。
反向传播：每个GPU计算完整梯度，此时梯度是全量的。
梯度同步：执行All-Reduce（或Reduce-Scatter）得到全局梯度。
局部更新：每个GPU只更新自己分片对应的参数，使用本地的优化器状态。
参数同步：更新后的参数分片通过All-Gather广播给所有GPU，保证下一轮前向时参数一致。

Stage 1的显存收益非常显著：优化器状态从每卡全量降为每卡1/N（N为DP度数）。以7.5B模型为例，优化器状态从90GB降为约1.4GB（N=64）。但需要注意，Stage 1仍然要求每卡能容纳完整的参数和梯度，因此它适合"参数+梯度可放入单卡，但加上优化器状态后放不下"的场景。

Stage 2：梯度分片（Gradient Partitioning）

Stage 2在Stage 1的基础上，进一步将梯度也进行分片。核心变化在于梯度同步方式：从All-Reduce改为Reduce-Scatter。Reduce-Scatter将梯度归约和分片合并为一步：每个GPU只负责归约自己分片对应的梯度，得到局部归约后的梯度分片，然后丢弃其他分片的梯度。这样，每个GPU只需要存储1/N的梯度，显存进一步下降。

Stage 2的显存收益：梯度从每卡全量降为每卡1/N。对于7.5B模型，梯度从15GB降为约0.23GB（N=64）。此时，每卡显存需求约为：参数（15GB）+ 梯度分片（0.23GB）+ 优化器状态分片（1.4GB）≈ 16.6GB，已经可以放入单卡A100。Stage 2适合"参数可放入单卡，但参数+梯度+优化器状态放不下"的场景。

工程经验：在实际使用中，Stage 2是性价比最高的配置。它几乎不增加通信量（Reduce-Scatter + All-Gather的总通信量与All-Reduce相同），但显存节省接近8倍。对于大多数10B-20B规模的模型，Stage 2配合激活重计算（Activation Checkpointing）已经足够。

Stage 3：参数分片（Parameter Partitioning）

Stage 3是ZeRO的终极形态，将模型参数也进行分片。此时，每个GPU只持有1/N的参数分片，前向和反向传播过程中需要动态地All-Gather参数。具体流程为：

前向传播：在计算某一层时，先All-Gather该层的参数分片，得到完整参数，计算完后丢弃非本卡分片的参数。
反向传播：类似地，在计算梯度时，先All-Gather参数，计算梯度后丢弃。
梯度同步：使用Reduce-Scatter得到本卡负责的梯度分片。
参数更新：使用本卡的优化器状态更新本卡的参数分片。

Stage 3的显存收益是革命性的：参数从每卡全量降为每卡1/N。对于7.5B模型，每卡仅需存储约0.23GB参数，总显存需求约7.5GB（N=64），这意味着显存需求随卡数线性下降。但代价是通信量增加：每次前向和反向都需要All-Gather参数，通信量约为3 × 参数量（前向一次、反向一次、梯度归约一次），比DDP多50%。

风险提示：Stage 3虽然显存最优，但通信开销显著增加。在低带宽网络（如10GbE）或跨节点场景下，通信可能成为瓶颈。建议在以下条件满足时使用Stage 3：① 单卡确实无法容纳模型参数；② 网络带宽充足（建议≥25Gb/s）；③ 使用DeepSpeed的通信优化（如梯度压缩、通信重叠）。

3.4 案例拆解：从DDP到ZeRO-3的迁移实战

下面通过一个实际案例，展示如何将一个使用DDP训练的模型（1.5B参数）迁移到ZeRO-3，并对比显存和吞吐的变化。

初始环境与约束：团队使用4台8卡A100（80GB）服务器，共32卡，NVLink内部互联，跨节点使用100Gb/s RoCE网络。原始DDP训练1.5B模型时，每卡显存占用约72GB（参数3GB + 梯度3GB + Adam状态18GB + 激活值48GB），接近OOM风险。训练吞吐约1200样本/秒。

迁移步骤：

# 原始DDP配置
python train.py --model -1.5b --distributed ddp

# 迁移到ZeRO-3（DeepSpeed配置）
# ds_config.json
{
  "train_batch_size": 1024,
  "gradient_accumulation_steps": 4,
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 3e-5 }
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" },
    "overlap_comm": true,
    "contiguous_gradients": true
  },
  "fp16": { "enabled": true }
}

关键参数说明：

stage: 3：启用ZeRO-3参数分片。
offload_optimizer：将优化器状态卸载到CPU内存，进一步释放GPU显存。
overlap_comm：通信与计算重叠，隐藏All-Gather延迟。
contiguous_gradients：梯度连续存储，减少内存碎片。

结果对比：

指标	DDP（32卡）	ZeRO-3（32卡）	变化
每卡显存占用	72 GB	38 GB	↓ 47%
训练吞吐	1200 样本/秒	1050 样本/秒	↓ 12.5%
可训练最大模型	~2B	~10B	↑ 5倍
通信量/迭代	~6 GB	~9 GB	↑ 50%

分析与复盘：ZeRO-3将每卡显存从72GB降至38GB，虽然吞吐下降了12.5%，但换来了5倍的模型规模扩展能力。吞吐下降的主要原因是参数All-Gather带来的额外通信。通过开启overlap_comm和offload_optimizer，实际吞吐损失已从最初的20%优化到12.5%。如果网络升级到200Gb/s，吞吐损失可进一步控制在5%以内。

  复盘经验：ZeRO-3并非"免费午餐"，它用通信换显存。在决定使用Stage 3之前，建议先评估：① 模型是否真的无法用Stage 2训练？② 网络带宽是否足够？③ 是否可以通过激活重计算（Activation Checkpointing）减少激活值显存，从而避免使用Stage 3？通常，激活值占用的显存可以通过重计算大幅降低，优先优化激活值往往比直接上Stage 3更划算。

3.5 常见误区与总结

在实际使用ZeRO时，有几个常见误区需要澄清：

误区一：ZeRO-3一定比ZeRO-2好。实际上，Stage 3的通信开销更大，在小规模集群或低带宽环境下，Stage 2可能更快。选择阶段应基于显存需求和网络条件综合判断。
误区二：ZeRO可以完全替代模型并行。当单个模型层（如超大Embedding或超大FFN）超过单卡显存时，ZeRO无法处理，仍需结合模型并行（如张量并行）进行层内切分。
误区三：ZeRO只适用于Transformer。ZeRO是通用优化器，适用于任何基于梯度下降的模型，包括CNN、RNN等。
误区四：开启ZeRO后不需要关注激活值显存。ZeRO只优化参数、梯度和优化器状态，激活值仍然占用大量显存。建议配合激活重计算（activation_checkpointing）使用。

总结：ZeRO通过三阶段递进式分片，系统地解决了数据并行中的显存冗余问题。Stage 1消除优化器状态冗余，Stage 2消除梯度冗余，Stage 3消除参数冗余。每个阶段都在显存节省和通信开销之间做出权衡。在实际工程中，建议从Stage 2起步，根据显存瓶颈和网络条件逐步升级到Stage 3，同时配合激活重计算、CPU卸载等辅助技术，实现训练效率和模型规模的最佳平衡。