【论文】训练万亿参数模型的内存优化（ZeRO）

论文地址

标题：ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

链接：https://ar5iv.labs.arxiv.org/html/1910.02054

作者：Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He (Microsoft)

发表年份：2019 (arXiv:1910.02054)

代码仓库：DeepSpeed (https://github.com/microsoft/DeepSpeed)

主要结论与创新点

ZeRO (Zero Redundancy Optimizer) 是一个革命性的内存优化系统，通过在数据并行训练中消除内存冗余，实现了三个关键突破：

核心创新：首次提出将模型状态（优化器状态、梯度、参数）进行分区的数据并行训练方法，在保持数据并行通信效率的同时，获得了接近模型并行的内存效率。这使得训练万亿参数模型成为可能，仅需1024个GPU即可训练1万亿参数的模型。

性能突破：在400个GPU上训练超过100B参数的模型，实现了超线性加速，达到了15 Petaflops的吞吐量。相比现有最优方案，模型大小增加了8倍，性能提升了10倍。

易用性提升：可以训练高达13B参数的模型（如比Megatron GPT 8.3B和T5 11B更大）而无需模型并行，降低了科学家使用大模型的难度。

论文背景的核心问题

大模型训练的内存挑战

深度学习模型规模的快速增长带来了前所未有的挑战。在自然语言处理领域，模型规模从BERT-large的0.3B参数迅速增长到GPT-2的1.5B、Megatron-LM的8.3B，再到T5的11B参数。要将模型规模继续扩展到万亿参数级别，我们面临着根本性的内存限制问题。

单个设备（如GPU或TPU）的内存容量是有限的。以当前主流的32GB GPU为例，即使是最基础的数据并行（Data Parallelism, DP）方法，也无法在单卡上训练超过1.4B参数的模型，因为数据并行需要在整个训练过程中在每张卡上保存完整的模型副本。

现有方案的局限性

为了应对大模型训练的内存问题，研究者们提出了多种并行训练策略，但每种方案都有其固有的局限性：

数据并行（DP）：具有良好的计算和通信效率，但内存效率极差。它在所有数据并行进程中复制完整的模型状态，导致内存浪费。对于大模型，这种冗余使得训练无法进行。

模型并行（MP）：将模型垂直分割，在不同设备间分布计算和参数，可以获得较高的内存效率。但这种方法将计算切分得过于细粒度，每层都需要进行设备间通信，导致通信开销巨大。在单节点内，由于GPU间通信带宽高，MP尚可工作；但跨节点时，网络带宽成为瓶颈，效率急剧下降。论文测试了40B参数模型在跨两个DGX-2节点的情况，每张V100 GPU仅能达到约5 Tflops（不到硬件峰值的5%）。

其他方案：Pipeline Parallelism（PP）、CPU-Offloading等方法都在功能、易用性、内存效率和计算/通信效率之间做出了不同程度的权衡，但都无法满足训练万亿参数模型的需求。

内存消耗的构成分析

重要说明：本文中提到的"内存"主要指GPU显存（VRAM），这是深度学习训练的主要瓶颈。ZeRO-DP主要优化GPU显存，而ZeRO-R中的CPU Offload功能会涉及CPU内存（RAM）来进一步扩展容量。

为了解决这个问题，论文首先对训练过程中的内存消耗进行了深入分析，将内存使用分为两个主要部分：

1. 模型状态（Model States）：对于大模型，这是GPU显存消耗的主要部分，包括：

• 优化器状态（Optimizer States）：如Adam优化器中的动量（momentum）和方差（variance），存储在GPU显存中

• 梯度（Gradients）：反向传播过程中计算出的梯度，存储在GPU显存中

• 参数（Parameters）：模型本身的权重参数，存储在GPU显存中

2. 残差状态（Residual States）：主要占用GPU显存，但在ZeRO-R中可以通过CPU Offload转移到CPU内存，包括激活值（activations）、临时缓冲区（temporary buffers）和内存碎片（fragmented memory）。

以一个使用混合精度训练和Adam优化器的模型为例，优化器状态占用的内存通常是参数大小的2-3倍（float32格式），再加上梯度（通常与参数同大小，float32或float16），模型状态的总内存占用可能是参数大小的4-8倍甚至更多。

问题的解决思路

ZeRO的核心思想

基于对内存消耗构成的分析，ZeRO提出了一个根本性的洞察：现有并行训练方法的内存冗余问题不在于计算本身，而在于模型状态的存储方式。数据并行在每张卡上存储完整的模型状态副本，这是极大的浪费；模型并行虽然通过分区避免了冗余，但代价是通信效率的严重下降。

ZeRO的关键创新在于：能否在保持数据并行的计算和通信效率的同时，获得模型并行的内存效率？

动态分区策略

ZeRO的答案是肯定的。它通过动态分区（Dynamic Partitioning）策略，在训练过程中动态地将模型状态分布到不同的数据并行进程中，而不是静态地复制到每个进程。这样既消除了内存冗余，又保持了数据并行原有的计算粒度和通信模式。

具体来说，ZeRO引入了**通信调度（Communication Schedule）**机制：在需要某个参数分片进行计算或更新时，通过all-gather操作收集必要的参数分片；使用完毕后，再将其释放。这种按需分配的方式，使得每张卡只需要存储模型状态的\frac{1}{N_d}（N_d为数据并行度），而不是完整的副本。

多阶段优化策略

ZeRO-DP设计了三个渐进的优化阶段，可以独立启用或组合使用：

阶段1：优化器状态分区（P_{os}）：将优化器状态分布到不同的数据并行进程中，实现4倍的内存减少。

阶段2：梯度分区（P_{os+g}）：在阶段1的基础上，进一步分区梯度，实现8倍的内存减少。

阶段3：参数分区（P_{os+g+p}）：在前两个阶段的基础上，再分区参数本身，内存减少与数据并行度N_d成线性关系。例如，64个GPU可以实现64倍的内存减少。

每个阶段的启用都会带来内存收益，同时保持或略微增加通信开销，但总体而言，通信开销的增加是可控的。

实现方案

ZeRO与模型拆分的关系

重要概念澄清：ZeRO-DP不拆分模型结构本身，它只拆分模型状态（参数、梯度、优化器状态）。这是ZeRO与模型并行（Model Parallelism）的关键区别。

ZeRO-DP的工作原理：

• 每个GPU仍然维护完整的模型结构（包括所有层的定义）

• 但每个GPU只存储部分模型状态（参数分片、梯度分片、优化器状态分片）

• 在前向/反向传播时，通过all-gather临时收集完整参数进行计算

• 计算完成后立即释放，只保留自己负责的分片

为什么不需要拆分模型结构？

关键在于：神经网络是逐层计算的，不需要同时将所有层的参数都加载到内存中。ZeRO-DP利用这一点：

• 持久存储：每个GPU只存储\frac{1}{N_d}的参数分片（例如，64个GPU时，每个GPU只存储\frac{1}{64}的参数）

• 临时收集：计算某一层时，临时收集该层的完整参数（可能是几GB）

• 立即释放：该层计算完成后，立即释放临时收集的参数

• 逐层处理：逐层重复这个"收集-计算-释放"的过程

因此，虽然完整模型可能远超单卡容量（如100GB），但通过增加数据并行度（如64个GPU），每个GPU的持久内存占用只有\frac{100GB}{64} = 1.56GB，再加上临时计算时的单层参数（约2GB），远小于完整模型的内存需求。

ZeRO可以结合模型并行使用：

虽然ZeRO-DP本身不需要拆分模型，但在超大模型训练中，ZeRO经常与模型并行结合使用。模型并行有两种主要方式：

方式1：Pipeline Parallelism（流水线并行）

• 拆分方式：按模型的层（Layer）拆分，不同GPU负责不同的层

• 前向传播：数据流水线式地流经不同GPU上的不同层

GPU0: Layer 1-5 → GPU1: Layer 6-10 → GPU2: Layer 11-15 → ...

输入 → 激活值 → 激活值 → 激活值 → 输出

• 反向传播：梯度反向流水线传播

• 优点：实现相对简单，对模型结构的侵入性较小

• 缺点：需要流水线处理，每个时刻总有GPU在等待，GPU利用率可能不高（可以通过流水线调度优化）

方式2：Tensor Parallelism（张量并行/模型并行）

• 拆分方式：按层内的维度拆分，通常是矩阵乘法的列或行

• 前向传播：矩阵乘法按列或行拆分，计算后立即汇总

例如：Y = XW（X: [B, H], W: [H, D]）

GPU0: Y₁ = X @ W[:, :D/2] } → All-gather → 完整Y

GPU1: Y₂ = X @ W[:, D/2:] }

• 反向传播：梯度反向传播时也需要相应的reduce操作

• 优点：同一时刻所有GPU都在工作，GPU利用率高

• 缺点：每层都需要通信（all-gather或all-reduce），跨节点时开销大

ZeRO与模型并行的组合：

在实际应用中，ZeRO-DP可以与模型并行结合：

• ZeRO-DP + Pipeline Parallelism：每个Pipeline阶段内部使用ZeRO-DP，减少每个阶段的内存占用

• ZeRO-DP + Tensor Parallelism：在Tensor Parallel的每个组内使用ZeRO-DP，进一步减少内存

例如，论文中训练170B模型时，使用了16路模型并行（MP=16）+ ZeRO数据并行，即：

• 模型在16个GPU间进行张量并行拆分（可能是按注意力头或FFN维度拆分）

• 在16个GPU组成的每个组之间，使用ZeRO进行数据并行

ZeRO-DP的三阶段内存优化

图1：ZeRO三阶段内存优化对比示意图

flowchart TB

subgraph Baseline["标准数据并行 (Baseline)"]

B1["GPU 0<br/>完整模型状态<br/>4Ψ bytes"]

B2["GPU 1<br/>完整模型状态<br/>4Ψ bytes"]

B3["GPU N<br/>完整模型状态<br/>4Ψ bytes"]

B1 -.->|"内存冗余"| B2

B2 -.->|"内存冗余"| B3

end

subgraph Stage1["阶段1: 优化器状态分区 (P_os)"]

S1_1["GPU 0<br/>参数: Ψ<br/>梯度: Ψ<br/>优化器状态: 2Ψ/Nd<br/>内存: 2Ψ + 2Ψ/Nd"]

S1_2["GPU 1<br/>参数: Ψ<br/>梯度: Ψ<br/>优化器状态: 2Ψ/Nd<br/>内存: 2Ψ + 2Ψ/Nd"]

S1_3["GPU N<br/>参数: Ψ<br/>梯度: Ψ<br/>优化器状态: 2Ψ/Nd<br/>内存: 2Ψ + 2Ψ/Nd"]

S1_1 -->|"约4x减少"| S1_2

end

subgraph Stage2["阶段2: 梯度分区 (P_os+g)"]

S2_1["GPU 0<br/>参数: Ψ<br/>梯度: Ψ/Nd<br/>优化器状态: 2Ψ/Nd<br/>内存: Ψ + 3Ψ/Nd"]

S2_2["GPU 1<br/>参数: Ψ<br/>梯度: Ψ/Nd<br/>优化器状态: 2Ψ/Nd<br/>内存: Ψ + 3Ψ/Nd"]

S2_3["GPU N<br/>参数: Ψ<br/>梯度: Ψ/Nd<br/>优化器状态: 2Ψ/Nd<br/>内存: Ψ + 3Ψ/Nd"]

S2_1 -->|"约8x减少"| S2_2

end

subgraph Stage3["阶段3: 参数分区 (P_os+g+p)"]

S3_1["GPU 0<br/>参数: Ψ/Nd<br/>梯度: Ψ/Nd<br/>优化器状态: 2Ψ/Nd<br/>内存: 4Ψ/Nd"]

S3_2["GPU 1<br/>参数: Ψ/Nd<br/>梯度: Ψ/Nd<br/>优化器状态: 2Ψ/Nd<br/>内存: 4Ψ/Nd"]

S3_3["GPU N<br/>参数: Ψ/Nd<br/>梯度: Ψ/Nd<br/>优化器状态: 2Ψ/Nd<br/>内存: 4Ψ/Nd"]

S3_1 -->|"Nd倍减少<br/>(线性扩展)"| S3_2

end

Baseline --> Stage1

Stage1 --> Stage2

Stage2 --> Stage3

上图中，\Psi表示模型参数数量，N_d表示数据并行度。每个阶段都在前一个阶段的基础上进一步减少内存占用，最终实现与并行度线性相关的内存减少。

ZeRO-DP通过三个阶段逐步消除模型状态的内存冗余，每个阶段对应不同的模型状态分片策略。

阶段1：优化器状态分区（P_{os}）

在标准的混合精度训练中，使用Adam优化器时，每个参数需要：

• 参数本身：\Psi bytes（参数数量）

• 梯度：\Psi bytes

• Adam动量（momentum）：\Psi bytes（float32）

• Adam方差（variance）：\Psi bytes（float32）

因此，优化器状态通常占用2\Psi bytes（以float32计算），而参数和梯度在混合精度下可能是\Psi bytes（float16）。总的内存占用约为4\Psi bytes（保守估计）。

在数据并行中，每个进程都需要存储这4\Psi bytes，因此N_d个进程的总内存为4\Psi \times N_d。

优化器状态分区策略：将优化器状态均匀分布到N_d个数据并行进程中，每个进程只存储\frac{1}{N_d}的优化器状态。在需要更新参数时，通过通信收集完整的梯度，然后每个进程只更新自己负责的那部分参数的优化器状态和参数值。

内存节省：

• 优化器状态：从2\Psi \times N_d减少到2\Psi（每个进程\frac{2\Psi}{N_d}）

• 总模型状态：从4\Psi减少到约\frac{4\Psi + 2\Psi(N_d-1)}{N_d}，近似为2\Psi + \frac{2\Psi}{N_d}（当N_d较大时）

实际测试中，阶段1通常能实现约4倍的内存减少。

阶段2：梯度分区（P_{os+g}）

在阶段1的基础上，进一步将梯度也进行分区。前向传播后，每个进程计算自己负责的数据的梯度，但这些梯度分布在所有参数上。通过all-reduce操作，每个进程都能获得完整的梯度，然后每个进程只保留自己负责的那部分参数的梯度。

内存节省：

• 梯度：从\Psi \times N_d减少到\Psi（每个进程\frac{\Psi}{N_d}）

• 结合阶段1，总模型状态内存减少到约\Psi + \frac{2\Psi + \Psi}{N_d}（当N_d较大时）

实际测试中，阶段2通常能实现约8倍的内存减少。

阶段3：参数分区（P_{os+g+p}）

这是最激进的优化阶段，将参数本身也进行分区。在标准的数据并行中，每个进程需要在内存中保存完整的参数副本用于前向和反向传播。ZeRO-DP通过动态all-gather策略，只在需要时收集参数分片。

核心机制：逐层计算，逐层收集，逐层释放

这是理解ZeRO-DP阶段3的关键。虽然计算时需要完整参数，但不是所有层的参数都同时需要。神经网络是逐层计算的，因此可以按层收集参数、计算、然后立即释放。

具体例子：

假设训练一个100GB参数的模型，使用64个GPU（N_d = 64），每层参数约2GB：

每个GPU的持久存储（始终占用）：

- 参数分片：100GB / 64 = 1.56GB（只存储自己负责的参数分片）

- 优化器状态分片：约3.12GB

- 梯度分片：约1.56GB

总计约6.24GB（持久占用）

  

前向传播 Layer 1 时（临时占用）：

- All-gather收集Layer 1的完整参数：约2GB（临时）

- 计算激活值：需要激活值内存（取决于batch size）

- 计算完成后：立即释放Layer 1的参数（2GB释放）

- 保存激活值用于反向传播

  

前向传播 Layer 2 时：

- All-gather收集Layer 2的完整参数：约2GB（临时）

- 使用Layer 1的激活值计算

- 计算完成后：立即释放Layer 2的参数

- 保存Layer 2的激活值

  

...以此类推逐层计算...

  

反向传播时：

- 对于每一层，重新All-gather收集参数

- 计算梯度后，立即释放参数

- 通过Reduce-scatter将梯度分发回对应的GPU

关键理解：

1. 持久内存：每个GPU只存储\frac{1}{N_d}的参数分片，这是持久占用的

2. 临时内存：计算某一层时，需要临时收集该层的完整参数（可能是几GB），但计算完成后立即释放

3. 逐层处理：由于是逐层计算的，同一时刻只需要一层或几层的参数，不需要所有层的参数同时存在

4. 内存峰值：峰值 = 持久内存 + 单层参数临时内存 + 激活值内存，远小于完整模型内存

参数收集策略：

• 前向传播时：每个进程通过all-gather收集完成当前层计算所需的参数分片

• 计算完成后：立即释放这些参数分片（通过释放通信缓冲区）

• 反向传播时：重复同样的收集-计算-释放流程

内存节省：

• 持久参数：从\Psi \times N_d减少到\frac{\Psi}{N_d}（每个进程）

• 临时参数：需要时临时收集单层参数，计算完立即释放，不占用持久内存

• 结合前两个阶段，总模型状态内存减少到约\frac{4\Psi}{N_d}（每个进程）

当N_d = 64时，总内存减少约为64倍。

为什么这样可行？

即使单卡放不下完整模型（100GB），但通过64张卡分区，每张卡只需要持久存储约1.56GB的参数分片。计算时虽然需要临时收集完整参数（例如某层的2GB），但：

• 这是临时的，计算完立即释放

• 只是某一层的参数，不是所有层

• 逐层处理，同一时刻只需要当前层的参数

因此，虽然计算时需要完整参数，但通过按需收集和立即释放的策略，持久内存占用大大减少。

图6：ZeRO-DP阶段3的逐层内存使用示意图

gantt

title ZeRO-DP阶段3：单GPU的内存使用时间线（100GB模型，64 GPU）

dateFormat X

axisFormat %s

section 持久内存

参数分片 1.56GB (持久) :done, persistent, 0, 1000

section Layer 1计算

All-gather收集参数 :active, gather1, 0, 10

Layer 1计算 :active, compute1, 10, 50

释放Layer 1参数 :done, release1, 50, 55

保存激活值 :active, save_act1, 50, 100

section Layer 2计算

All-gather收集参数 :active, gather2, 100, 110

Layer 2计算 :active, compute2, 110, 150

释放Layer 2参数 :done, release2, 150, 155

保存激活值 :active, save_act2, 150, 200

section Layer 3计算

All-gather收集参数 :active, gather3, 200, 210

Layer 3计算 :active, compute3, 210, 250

释放Layer 3参数 :done, release3, 250, 255

关键观察：

• 蓝色（持久内存）：始终占用1.56GB（参数分片），这是持久的

• 绿色（临时内存）：每次计算一层时，临时收集约2GB的完整参数，计算完立即释放

• 时间轴：不同层的参数不会同时存在，它们是按顺序收集和释放的

• 峰值内存：峰值 ≈ 持久内存(1.56GB) + 单层参数(2GB) + 激活值，远小于完整模型(100GB)

这就是为什么即使完整模型放不下单卡，通过增加数据并行度，ZeRO-DP也能让训练成为可能。

通信开销分析

虽然ZeRO-DP通过分区显著减少了内存，但需要额外的通信操作。论文详细分析了各阶段的通信开销：

阶段1（P_{os}）：通信量与标准数据并行相同，因为仍然需要在每个进程中收集完整梯度进行all-reduce。

阶段2（P_{os+g}）：通信量与标准数据并行相同，因为梯度仍然需要all-reduce（虽然存储时分区了）。

阶段3（P_{os+g+p}）：需要在前向和反向传播的每一层进行参数all-gather操作，通信量约为标准数据并行的1.5倍（相比梯度all-reduce，参数all-gather的通信量稍大）。但考虑到内存的巨大节省，这个通信开销的增加是值得的。

ZeRO-R：残差状态优化

除了模型状态的优化，ZeRO还提出了ZeRO-R来优化残差状态（激活值、临时缓冲区和内存碎片）。ZeRO-R主要优化GPU显存，但可以通过CPU Offload扩展到CPU内存：

激活值分区（Activation Partitioning）：通过CPU offload或激活值重计算来减少激活值的GPU显存占用。CPU Offload将激活值从GPU显存转移到CPU内存，虽然会增加CPU-GPU传输开销，但可以释放大量GPU显存。

临时缓冲区分区（Temporary Buffer Partitioning）：将临时缓冲区也进行分区，避免GPU显存浪费。临时缓冲区也可以选择offload到CPU内存。

内存碎片整理（Memory Defragmentation）：通过智能的内存管理减少GPU显存碎片化，提高显存利用率。

ZeRO-R的内存策略：在GPU显存充足时，激活值和缓冲区保留在GPU上；在GPU显存不足时，可以选择offload到CPU内存，用容量换速度。

算法流程

ZeRO-DP的训练流程可以用以下算法描述：

算法：ZeRO-DP训练流程（阶段3完整版）

  

输入：模型M，数据并行度Nd，训练数据D

输出：训练好的模型参数

  

1. 初始化：

- 将参数Θ分区到Nd个进程：Θ = {Θ₁, Θ₂, ..., Θ_{Nd}}

- 将优化器状态O分区到Nd个进程：O = {O₁, O₂, ..., O_{Nd}}

  

2. 对于每个训练批次：

a. 数据并行：每个进程处理不同的数据子集

b. 前向传播（对于每一层L）：

- All-gather：收集完成层L所需的参数分片 Θ_L

- 计算：使用Θ_L计算激活值 A_L

- 释放：释放Θ_L（保留激活值A_L用于反向传播）

c. 反向传播（对于每一层L，从后往前）：

- All-gather：收集完成层L所需的参数分片 Θ_L

- 计算：使用Θ_L和激活值A_L计算梯度 G_L

- Reduce-scatter：将梯度G_L分区并分发到对应进程

- 释放：释放Θ_L

d. 参数更新（每个进程独立）：

- 使用本进程负责的梯度分区更新优化器状态和参数

3. 返回训练后的模型参数

这个算法的关键在于：参数只在需要时通过all-gather收集，使用完毕后立即释放，从而大幅减少内存占用。

图2：ZeRO-DP阶段3的前向和反向传播流程

sequenceDiagram

participant GPU0 as GPU 0<br/>(持有参数分片 Θ₁)

participant GPU1 as GPU 1<br/>(持有参数分片 Θ₂)

participant GPUN as GPU N<br/>(持有参数分片 Θ_N)

Note over GPU0,GPUN: 前向传播 - 层L

GPU0->>GPU0: 计算需要参数 Θ_L

GPU1->>GPU1: 计算需要参数 Θ_L

GPUN->>GPUN: 计算需要参数 Θ_L

GPU0->>GPU1: All-gather: 收集 Θ_L 的所有分片

GPU0->>GPUN: All-gather: 收集 Θ_L 的所有分片

GPU1->>GPU0: All-gather: 发送分片 Θ₂

GPU1->>GPUN: All-gather: 收集 Θ_L 的所有分片

GPUN->>GPU0: All-gather: 发送分片 Θ_N

GPUN->>GPU1: All-gather: 发送分片 Θ_N

Note over GPU0,GPUN: 每个GPU现在都有完整的 Θ_L

GPU0->>GPU0: 使用 Θ_L 计算激活值 A_L

GPU1->>GPU1: 使用 Θ_L 计算激活值 A_L

GPUN->>GPUN: 使用 Θ_L 计算激活值 A_L

GPU0->>GPU0: 释放 Θ_L (保留 A_L)

GPU1->>GPU1: 释放 Θ_L (保留 A_L)

GPUN->>GPUN: 释放 Θ_L (保留 A_L)

Note over GPU0,GPUN: 反向传播 - 层L

GPU0->>GPU0: All-gather: 重新收集 Θ_L

GPU1->>GPU1: All-gather: 重新收集 Θ_L

GPUN->>GPUN: All-gather: 重新收集 Θ_L

GPU0->>GPU0: 使用 Θ_L 和 A_L 计算梯度 G_L

GPU1->>GPU1: 使用 Θ_L 和 A_L 计算梯度 G_L

GPUN->>GPUN: 使用 Θ_L 和 A_L 计算梯度 G_L

GPU0->>GPU1: Reduce-scatter: 将梯度分片分发

GPU0->>GPUN: Reduce-scatter: 将梯度分片分发

GPU1->>GPU0: Reduce-scatter: 接收梯度分片 G₁

GPU1->>GPUN: Reduce-scatter: 将梯度分片分发

GPUN->>GPU0: Reduce-scatter: 接收梯度分片 G_N

GPUN->>GPU1: Reduce-scatter: 接收梯度分片 G_N

Note over GPU0,GPUN: 每个GPU只保留自己负责的梯度分片

GPU0->>GPU0: 使用梯度分片更新参数 Θ₁

GPU1->>GPU1: 使用梯度分片更新参数 Θ₂

GPUN->>GPUN: 使用梯度分片更新参数 Θ_N

这个流程图展示了ZeRO-DP阶段3的关键操作：在前向传播时通过all-gather收集完整参数进行计算，在反向传播时通过reduce-scatter将梯度分区并分发到对应的GPU，最终每个GPU只保留并更新自己负责的参数分片。

重要说明：注意这里每个GPU仍然需要完整的模型结构（所有层的代码），只是参数存储时分区了。这与模型并行不同，模型并行是将模型结构本身拆分到不同GPU上。

图4：ZeRO-DP vs 模型并行的对比示意图

flowchart TB

subgraph ZERO["ZeRO-DP（不拆分模型结构）"]

Z1["GPU 0<br/>完整模型结构<br/>参数分片: Θ₁<br/>所有层代码"]

Z2["GPU 1<br/>完整模型结构<br/>参数分片: Θ₂<br/>所有层代码"]

Z3["GPU N<br/>完整模型结构<br/>参数分片: Θ_N<br/>所有层代码"]

Z1 -.->|"All-gather<br/>临时收集完整Θ"| Z2

Z2 -.->|"计算后释放<br/>只保留分片"| Z3

end

subgraph MP["模型并行（拆分模型结构）"]

subgraph MP1["Pipeline Parallelism"]

P1["GPU 0<br/>Layer 1-5<br/>完整参数"]

P2["GPU 1<br/>Layer 6-10<br/>完整参数"]

P3["GPU N<br/>Layer 11-15<br/>完整参数"]

P1 -->|"激活值传递"| P2

P2 -->|"激活值传递"| P3

end

subgraph MP2["Tensor Parallelism"]

T1["GPU 0<br/>所有层<br/>参数分片 W₁"]

T2["GPU 1<br/>所有层<br/>参数分片 W₂"]

T1 <-->|"每层计算后<br/>All-gather"| T2

end

end

ZERO -->|"可以组合使用"| MP

从图中可以看出，ZeRO-DP的每个GPU都有完整的模型结构，只是参数存储时分区；而模型并行是将模型结构本身拆分（按层或按维度）到不同GPU上。两者可以结合使用，例如：先用Tensor Parallelism将模型按维度拆分，然后在每个Tensor Parallel组内使用ZeRO-DP进一步优化内存。

图5：ZeRO-DP结合Tensor Parallelism的前向传播详细流程

sequenceDiagram

participant TP0 as Tensor Parallel Group 0<br/>(GPU 0-3, 4路TP)

participant TP1 as Tensor Parallel Group 1<br/>(GPU 4-7, 4路TP)

Note over TP0,TP1: ZeRO-DP数据并行度=2, Tensor Parallel度=4

Note over TP0: 组内Tensor Parallel计算Layer L

TP0->>TP0: GPU0-3: All-gather收集ZeRO参数分片<br/>(从数据并行组中)

TP0->>TP0: GPU0: W₁[:, :D/4] × X

TP0->>TP0: GPU1: W₁[:, D/4:D/2] × X

TP0->>TP0: GPU2: W₁[:, D/2:3D/4] × X

TP0->>TP0: GPU3: W₁[:, 3D/4:] × X

TP0->>TP0: All-gather汇总输出<br/>(Tensor Parallel通信)

TP0->>TP0: 释放ZeRO参数分片<br/>(保存激活值)

Note over TP1: 组内Tensor Parallel计算Layer L

TP1->>TP1: GPU4-7: All-gather收集ZeRO参数分片<br/>(从数据并行组中)

TP1->>TP1: GPU4: W₂[:, :D/4] × X

TP1->>TP1: GPU5: W₂[:, D/4:D/2] × X

TP1->>TP1: GPU6: W₂[:, D/2:3D/4] × X

TP1->>TP1: GPU7: W₂[:, 3D/4:] × X

TP1->>TP1: All-gather汇总输出<br/>(Tensor Parallel通信)

TP1->>TP1: 释放ZeRO参数分片<br/>(保存激活值)

Note over TP0,TP1: 反向传播时类似，但顺序相反

关键理解：

1. Tensor Parallelism拆分：在组内，每个GPU负责矩阵乘法的一部分（按列或行拆分），每层计算后需要All-gather汇总

2. ZeRO-DP优化：在数据并行组间，每个GPU只存储参数的分片，需要时All-gather收集完整参数

3. 组合效果：先用TP减少单组内每个GPU的参数量，再用ZeRO-DP在组间进一步分区，实现双重内存优化

应用优势与研究意义

可扩展性突破

ZeRO的最大贡献在于使训练万亿参数模型成为可能。以使用Adam优化器进行16位混合精度训练为例，1万亿参数的模型需要约16TB内存来存储优化器状态、梯度和参数。通过1024个GPU进行阶段3的ZeRO优化，每个GPU只需要约16GB内存，这在32GB GPU上是完全可以实现的。

论文通过理论分析和实验验证，证明了ZeRO具有良好的可扩展性。在400个GPU上，成功训练了170B参数的模型，并展示了训练更大模型的可行性。

性能表现

论文在多个模型规模上进行了性能评估，结果显示：

• 超线性加速：在400个GPU上训练100B+参数模型时，ZeRO实现了超线性加速（super-linear speedup），这在并行计算中是罕见的，主要得益于更大的batch size和更好的内存利用率。

• 吞吐量提升：相比基线方法，ZeRO在保持相同通信开销的情况下，实现了更高的吞吐量。例如，在128个GPU上训练60B参数模型时，ZeRO的吞吐量显著高于使用模型并行的Megatron-LM。

• 模型规模提升：ZeRO可以训练比现有最优方案大8倍的模型，同时实现10倍的性能提升。

易用性改进

传统的模型并行需要复杂的模型切分策略和通信调度，对研究者来说实现难度较高。ZeRO通过简单的配置即可启用，无需手动进行模型切分，大大降低了使用门槛。这使得研究者可以在不掌握复杂并行技术的情况下，训练比Megatron GPT 8.3B和T5 11B更大的模型。

实际应用价值

ZeRO的提出直接推动了更大规模语言模型的发展。论文中提到，研究者已经使用ZeRO的系统性突破创建了当时世界上最大的语言模型（17B参数），并取得了突破性的准确率。

后续，DeepSpeed（微软基于ZeRO实现的系统）被广泛应用于大模型训练，包括GPT-3、GPT-NeoX、Bloom等超大规模模型，都或多或少地采用了ZeRO的优化策略。

核心结论与工程启示

核心结论

通过对ZeRO论文的深入分析，我们可以提炼出以下5个核心结论：

结论1：模型状态分区是解决大模型训练内存瓶颈的关键。传统的并行策略要么复制模型状态导致内存浪费（DP），要么分区导致通信效率下降（MP）。ZeRO通过动态分区策略，在保持DP通信效率的同时获得了MP的内存效率，这是其成功的核心。

结论2：分阶段优化策略提供了灵活性和渐进性。ZeRO的三个优化阶段可以独立启用，为不同的应用场景提供了灵活的配置选择。对于中小规模模型，可能只需要阶段1或阶段2；对于超大规模模型，才需要启用阶段3。这种渐进式的设计使得ZeRO具有广泛的适用性。

结论3：通信开销的增加是可控的，且值得。阶段3虽然增加了约50%的通信开销，但换来了与数据并行度线性相关的内存减少。在内存受限的场景下，这个权衡是明显值得的。

结论4：残差状态的优化同样重要。虽然模型状态占用大部分内存，但激活值和临时缓冲区在大模型中也可能达到数百GB。ZeRO-R的激活值分区和CPU offload策略，进一步释放了内存空间。

结论5：动态调度机制是高效实现的关键。ZeRO不是简单地静态分区参数，而是通过按需all-gather的动态调度，在需要时收集、使用后释放。这种动态性使得内存使用更加高效，避免了静态分区的通信开销问题。

图3：ZeRO动态调度机制示意图

flowchart LR

subgraph GPU["单个GPU的内存状态"]

A["参数分片<br/>Θ_i (持久存储)"]

B["优化器状态分片<br/>O_i (持久存储)"]

C["临时缓冲区<br/>(动态分配)"]

end

subgraph Forward["前向传播阶段"]

F1["All-gather参数<br/>收集 Θ_L"]

F2["计算激活值<br/>使用 Θ_L"]

F3["释放参数<br/>Θ_L → 缓冲区"]

end

subgraph Backward["反向传播阶段"]

B1["All-gather参数<br/>重新收集 Θ_L"]

B2["计算梯度<br/>使用 Θ_L 和 A_L"]

B3["Reduce-scatter梯度<br/>分区并分发"]

B4["释放参数<br/>Θ_L → 缓冲区"]

end

subgraph Update["参数更新阶段"]

U1["更新优化器状态<br/>使用梯度分片 G_i"]

U2["更新参数<br/>Θ_i ← G_i"]

end

GPU --> Forward

Forward --> F1

F1 --> F2

F2 --> F3

F3 --> Backward

Backward --> B1

B1 --> B2

B2 --> B3

B3 --> B4

B4 --> Update

Update --> U1

U1 --> U2

U2 --> GPU

style A fill:#e1f5ff

style B fill:#e1f5ff

style C fill:#fff4e1

这个流程图展示了ZeRO的动态调度机制：参数分片和优化器状态分片是持久存储在每个GPU上的，但在前向和反向传播时，需要通过all-gather动态收集完整参数进行计算，计算完成后立即释放回缓冲区，从而最大程度地减少内存占用。

工程启示

启示1：在什么场景下应该采用ZeRO

适用场景：

1. 模型规模超过单卡内存容量：当模型无法在单卡上训练，且使用传统数据并行会因内存不足而失败时，ZeRO是最佳选择。

2. 需要训练超大模型（10B+参数）：对于超大规模模型，ZeRO几乎是唯一能在保持良好性能的同时训练的方法。传统的模型并行在跨节点时效率极低。

3. 追求训练吞吐量：ZeRO可以支持更大的batch size，从而获得更高的训练吞吐量。这对于需要快速迭代的实验或生产环境特别重要。

4. 硬件资源有限：在GPU内存受限的环境中（如消费级GPU或云GPU），ZeRO可以通过增加数据并行度来训练更大的模型。

不适用场景：

1. 小规模模型（<1B参数）：对于小模型，ZeRO的开销可能超过收益，直接使用标准数据并行即可。

2. 单卡可以训练：如果模型可以在单卡上训练，直接使用数据并行或混合精度训练即可，无需ZeRO。

3. 对通信延迟极度敏感：在某些对延迟极度敏感的场景下（虽然这种情况很少），阶段3的额外通信开销可能不可接受。

启示2：如何避免常见问题

问题1：参数收集的时机不当导致性能下降

在实际使用中，参数的all-gather操作必须与计算流水线化（pipeline），避免计算与通信的串行化。DeepSpeed通过在CUDA stream上异步执行通信操作来实现这一点。

工程建议：

• 使用异步通信API（如NCCL的异步all-gather）

• 确保通信与计算重叠（communication-computation overlap）

• 监控GPU利用率，如果发现GPU空闲时间过长，可能是通信阻塞导致的

问题2：激活值内存仍然占用过大

即使启用了ZeRO的所有三个阶段，激活值仍然可能占用大量内存，特别是在Transformer的每一层都需要保存激活值用于反向传播时。

工程建议：

• 启用梯度检查点（Gradient Checkpointing）：用计算换内存，只保存部分层的激活值，需要时重新计算

• 使用ZeRO-R的激活值CPU offload：将激活值offload到CPU内存，虽然会增加CPU-GPU传输开销，但可以释放GPU内存

• 合理设置batch size：在内存和训练效果之间找到平衡点

问题3：通信带宽成为瓶颈

在多节点训练中，如果节点间网络带宽不足，阶段3的额外通信开销可能导致训练速度下降。

工程建议：

• 使用高速网络（如InfiniBand）进行多节点训练

• 考虑使用混合策略：在节点内使用模型并行，节点间使用ZeRO数据并行

• 监控通信时间占比，如果超过30%，可能需要重新评估并行策略

启示3：实际工程中的关键细节

细节1：优化器状态的数据类型选择

虽然论文提到使用混合精度训练，但优化器状态通常仍使用float32。这是因为Adam等优化器需要高精度来维护动量和方差的准确性。

实践建议：

• 参数和梯度使用float16（或bfloat16）

• 优化器状态使用float32

• 如果内存仍然不足，可以考虑使用更激进的优化器（如8-bit Adam），但可能影响收敛性

细节2：梯度累积的实现

在内存受限的情况下，通常需要使用梯度累积来模拟更大的batch size。ZeRO的梯度分区策略需要与梯度累积正确配合。

实践建议：

• 在梯度累积时，每个进程只累积自己负责的那部分梯度

• 确保在累积完成后才进行reduce-scatter操作

• 监控梯度更新的正确性，确保分区后的梯度更新逻辑正确

细节3：Checkpoint和恢复训练

由于参数是分区的，保存和加载checkpoint需要特殊处理。

实践建议：

• 保存checkpoint时，需要收集所有分区的参数

• 可以使用rank 0进程收集完整参数并保存，或使用分布式checkpoint格式

• 确保checkpoint格式兼容DeepSpeed的ZeRO实现

启示4：与其他方法的对比和选择

ZeRO vs 模型并行（MP）：

| 维度 | ZeRO | MP |

|----------|--------------------|--------------------|

| 内存效率 | 高（线性扩展） | 高（但受节点限制） |

| 通信效率 | 高（数据并行模式） | 低（跨节点时） |

| 易用性 | 高（自动分区） | 低（需要手动切分） |

| 适用规模 | 10B-1T+ | <100B（跨节点时） |

选择建议：

• 优先选择ZeRO：对于100B+参数模型，或需要在多节点上训练时

• 考虑混合策略：在单节点内使用MP（节点内通信快），节点间使用ZeRO

ZeRO vs Pipeline并行（PP）

Pipeline并行通过时间维度的流水线来提高GPU利用率，但与ZeRO可以互补。

选择建议：

• 组合使用：ZeRO + Pipeline并行可以获得最佳效果

• DeepSpeed：已经实现了ZeRO + Pipeline并行的组合，这是当前训练超大模型的主流方案

ZeRO vs CPU Offloading

CPU Offloading将参数或优化器状态offload到CPU内存，但会增加CPU-GPU传输开销。

选择建议：

• 优先使用ZeRO：如果GPU间通信带宽足够（如InfiniBand），ZeRO通常比CPU Offloading更高效

• 组合使用：在GPU通信带宽不足时，可以组合使用ZeRO和CPU Offloading

启示5：性能优化的方向

优化方向1：通信优化

虽然ZeRO的通信开销是可控的，但仍有优化空间：

• 使用压缩通信：对梯度或参数进行压缩后再传输（如1-bit Adam、梯度压缩）

• 通信调度优化：更智能地调度all-gather和reduce-scatter的时机，最大化通信计算重叠

• 网络拓扑感知：在多节点训练中，考虑网络拓扑结构，优先在同一交换机下的GPU间通信

优化方向2：内存管理优化

• 内存池（Memory Pool）：预分配内存池，减少内存分配和释放的开销

• 内存预取：在需要参数分片之前就开始all-gather，进一步重叠通信和计算

• 激活值优化：更激进的激活值offload策略，或使用更高效的激活值重计算算法

优化方向3：系统级优化

• 混合精度策略：更细粒度的混合精度控制，对不同部分使用不同的精度

• 动态batch size：根据当前内存使用情况动态调整batch size

• 自适应并行策略：根据模型大小和硬件配置自动选择最优的并行策略组合

实践案例说明

案例1：训练GPT-3规模模型

假设我们要训练一个175B参数的GPT-3规模模型：

• 内存需求估算：使用Adam优化器和混合精度，模型状态约需175B \times 4 bytes \times 4 = 2.8TB（考虑优化器状态）

• 硬件配置：假设有256个32GB GPU（总内存8TB）

• ZeRO策略：启用阶段3，每个GPU只需存储\frac{2.8TB}{256} = 11GB的模型状态，还有约21GB用于激活值和其他开销

• 通信考虑：使用InfiniBand网络，确保节点间通信带宽足够（>200Gbps）

• 实际部署：使用DeepSpeed ZeRO-3配置，结合Pipeline并行（如果需要），batch size设置为每GPU 1-2

案例2：在有限资源上训练大模型

假设只有8个16GB GPU，想训练13B参数的模型：

• 内存需求：13B \times 4 bytes \times 4 = 208GB模型状态

• 标准DP不可行：8个GPU总内存128GB，无法容纳

• ZeRO策略：启用阶段3，每个GPU存储\frac{208GB}{8} = 26GB，但单个GPU只有16GB，仍然不足

• 组合策略：使用ZeRO-3 + CPU Offloading，将部分激活值offload到CPU，或使用梯度检查点减少激活值内存

• 实际配置：DeepSpeed ZeRO-3 + activation offload + gradient checkpointing，batch size设为1

案例3：多节点训练优化

假设有4个节点，每个节点8个GPU，共32个GPU，训练40B参数模型：

• 节点内通信：NVLink，带宽高（>300GB/s）

• 节点间通信：以太网，带宽较低（10-25GB/s）

• 策略选择：

• 方案A（纯ZeRO）：32路数据并行，节点间通信频繁，可能成为瓶颈

• 方案B（混合并行）：节点内使用模型并行（8路MP），节点间使用ZeRO数据并行（4路DP）

• 推荐方案B：减少跨节点通信次数，提高整体效率

通过这些案例可以看出，ZeRO不是万能的，需要根据具体的硬件配置、网络条件和模型规模来选择最优的策略组合。

总结

ZeRO通过创新的动态分区策略，成功解决了大模型训练的内存瓶颈问题，使训练万亿参数模型成为可能。其核心价值在于：在不牺牲通信效率的前提下，实现了与数据并行度线性相关的内存减少。

对于工程实践者而言，ZeRO最重要的启示是：内存优化需要系统性的思考，不能只关注单一维度。ZeRO通过模型状态分区、残差状态优化、动态调度等多种技术的组合，实现了整体最优。在实际应用中，需要根据具体的硬件资源、网络条件和模型规模，灵活选择和组合不同的优化策略。

随着大模型时代的到来，ZeRO及其后续发展（如ZeRO-Infinity）将继续在大模型训练中发挥关键作用。掌握ZeRO的核心思想和工程实践，对于参与大模型训练的工程师和研究者来说，都是必不可少的。