文章标题：ZeRO：面向万亿参数模型训练的内存优化
作者/机构：Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He (Microsoft)

A1 主要贡献

本文旨在解决训练参数量从数十亿到万亿级别的大型深度学习模型所面临的内存限制和效率挑战。现有并行化方案（如数据并行、模型并行）在内存效率、计算/通信效率和易用性之间存在权衡，难以进一步扩展模型规模。

为应对此挑战，本文提出了 ZeRO（Zero Redundancy Optimizer，零冗余优化器），一个旨在优化内存使用、提升训练速度和模型规模的解决方案。ZeRO的核心创新点如下：
1. 提出ZeRO-DP，一种增强的数据并行方案：ZeRO-DP通过划分模型状态（优化器状态、梯度和参数）而非复制，消除了数据并行训练中的内存冗余。它通过动态通信调度，在实现与模型并行相当的内存效率的同时，保持了数据并行的高计算粒度和低通信量。ZeRO-DP分为三个优化阶段：
* $P_{os}$ (Optimizer State Partitioning)：优化器状态划分，可减少4倍内存占用。
* $P_g$ (Gradient Partitioning)：增加梯度划分，可减少8倍内存占用。
* $P_p$ (Parameter Partitioning)：再增加参数划分，内存减少量与数据并行度 $N_d$ 成线性关系，可达 $N_d$ 倍。
2. 提出ZeRO-R，用于优化残差内存：在ZeRO-DP解决了模型状态内存瓶颈后，ZeRO-R进一步优化激活值、临时缓冲区和内存碎片这三类残差状态的内存占用。
* 通过划分激活值并按需重构，解决了模型并行中激活值复制导致的冗余。
* 通过管理临时缓冲区的大小，平衡内存与计算效率。
* 通过主动进行内存碎片整理，提高了内存利用率。
3. 将ZeRO与模型并行（MP）结合：ZeRO可以与MP结合使用，进一步降低内存占用。理论上，数据并行度为 $N_d$、模型并行度为 $N_m$ 时，每台设备的内存占用可减少 $N_d \times N_m$ 倍，使得在有限的硬件上以合理的批量大小高效训练超大规模模型成为可能。
4. 实现并验证了ZeRO-100B：本文实现了ZeRO的一个子集（$P_{os+g}$ + ZeRO-R），命名为ZeRO-100B。实验证明，该系统能在400个GPU上高效训练超过1000亿参数的模型，性能达到15 Petaflops，模型规模比现有技术（SOTA）提升8倍，性能提升10倍。同时，它支持在不使用模型并行的情况下训练高达130亿参数的模型，极大地降低了大规模模型训练的门槛。基于ZeRO，研究人员已成功训练出当时全球最大的语言模型Turing-NLG（170亿参数）。

A3 背景知识/关键Observation/设计原则

3. 内存都去哪儿了？

引言
以一个15亿参数的GPT-2模型为例，其16位精度的权重（或参数）仅需3GB内存，然而在拥有32GB内存的单个GPU上，使用Tensorflow或PyTorch却无法进行训练。这引发了一个问题：内存究竟被什么占用了？在模型训练期间，大部分内存被模型状态（即优化器状态、梯度和参数组成的张量）消耗。除此之外，剩余的内存被残差状态（激活值、临时缓冲区和内存碎片）占用。我们将从这两个方面详细分析内存消耗。

3.1 模型状态：优化器状态、梯度和参数

模型状态是内存消耗的主体
在训练过程中，大部分设备内存被模型状态所消耗。以Adam【索引6，Adam: A method for stochastic optimization，Diederik P. Kingma and Jimmy Ba，2015，ICLR 2015】为例，它是深度学习训练中最流行的优化器之一。Adam需要存储两个优化器状态：梯度的时间平均动量和方差，用于计算更新。因此，使用Adam训练模型必须有足够的内存来存放动量和方差的副本。此外，还需要内存来存储梯度和权重本身。在这三种与参数相关的张量中，优化器状态通常消耗内存最多，尤其是在应用混合精度训练时。

混合精度训练的内存分析
在当前一代NVIDIA GPU上训练大型模型，最先进的方法是混合精度（fp16/32）训练【索引23，Mixed precision training，Paulius Micikevicius et al.，2017】，其中参数和激活值存储为fp16，从而能够利用这些GPU上的高吞吐量Tensor Core单元【索引24，NVIDIA Tesla V100 GPU architecture，2017】。在混合精度训练中，前向和后向传播都使用fp16的权重和激活值。然而，为了在后向传播结束时有效计算和应用更新，混合精度优化器会保留一份fp32的参数副本以及所有其他优化器状态的fp32副本。

以Adam为例的具体计算
我们以Adam为例进行具体说明。使用Adam对一个包含 $\Psi$ 个参数的模型进行混合精度训练，需要足够的内存来存放一份fp16的参数和梯度，内存需求分别为 $2\Psi$ 和 $2\Psi$ 字节。此外，还需要存放优化器状态：一份fp32的参数副本、动量和方差，内存需求分别为 $4\Psi$、 $4\Psi$ 和 $4\Psi$ 字节。我们用K表示优化器状态的内存乘数，即存储它们所需的额外内存为 $K\Psi$ 字节。混合精度Adam的K=12。总共需要的内存为 $2\Psi + 2\Psi + K\Psi = 16\Psi$ 字节。对于像GPT-2这样拥有15亿参数的模型，这导致至少需要24GB的内存，远高于仅存放fp16参数所需的3GB内存。

3.2 残差内存消耗

激活值的内存占用
在训练过程中，激活值会占用大量内存【索引7，Training deep nets with sublinear memory cost，Tianqi Chen et al.，2016】。举个具体的例子，一个15亿参数的GPT-2模型，在序列长度为1K、批量大小为32的情况下进行训练，大约需要60GB的内存。激活检查点（或激活重计算）是减少激活内存的常用方法，它能将激活内存大约减少到总激活量的平方根，但代价是33%的重计算开销【索引7，Training deep nets with sublinear memory cost，Tianqi Chen et al.，2016】。这将使该模型的激活内存消耗减少到约8GB。

大模型激活内存的挑战
尽管激活检查点能显著减少内存，但对于更大的模型，即使使用了该技术，激活内存仍然可能变得非常大。例如，一个拥有1000亿参数的类GPT模型，即使使用激活检查点，在批量大小为32的情况下，也需要约60GB的内存。

临时缓冲区的内存占用
用于存储中间结果的临时缓冲区对于大型模型会消耗不可忽视的内存。诸如梯度all-reduce或梯度范数计算等操作，为了提高吞吐量，倾向于在应用操作前将所有梯度融合成一个扁平化的缓冲区。例如，跨设备的all-reduce带宽会随着消息尺寸的增大而提高。虽然梯度本身通常存储为fp16张量，但融合后的缓冲区可能是fp32张量，具体取决于操作类型。当模型规模很大时，这些临时缓冲区的大小不可小觑。例如，对于一个15亿参数的模型，一个扁平化的fp32缓冲区将需要6GB内存。

内存碎片问题
到目前为止，我们讨论了训练期间的实际内存消耗。此外，即使有大量可用内存，也可能因为内存碎片问题而耗尽可用内存。当一个内存请求无法找到足够大的连续内存块来满足时，即使总可用内存大于请求量，请求也会失败。我们在训练非常大的模型时观察到显著的内存碎片现象，在一些极端情况下，即使有超过30%的内存可用，也会导致内存不足的问题。

4. ZeRO：洞察与概览

ZeRO的优化组成
ZeRO包含两套优化：i) ZeRO-DP，旨在减少模型状态的内存占用；以及 ii) ZeRO-R，旨在减少残差内存消耗。我们在此概述这些优化及其背后的洞察，这些洞察使得ZeRO能够在减少内存占用的同时保持高效。请注意，效率是这里的关键：若无此约束，将所有参数状态移至CPU内存或任意增加模型并行度等简单方案也能减少内存占用。

4.1 ZeRO-DP的洞察与概览

ZeRO驱动的数据并行的三个核心洞察
ZeRO驱动的数据并行（DP）基于三个关键洞察：
a) DP比MP具有更好的扩展效率：因为MP在降低计算粒度的同时增加了通信开销。超过某个点后，更低的计算粒度会降低单个GPU的效率，而增加的通信开销则阻碍了跨GPU的扩展性，尤其是在跨越节点边界时。相反，DP具有更高的计算粒度和更低的通信量，从而实现更高的效率。
b) DP的内存效率低下：因为模型状态在所有数据并行进程中被冗余存储。相反，MP通过划分模型状态来获得内存效率。
c) DP和MP静态持有所有模型状态：这两种方法在整个训练过程中都保持所有需要的模型状态，但并非所有状态在任何时候都是必需的。例如，每一层对应的参数仅在该层的前向和后向传播期间才需要。

ZeRO-DP的设计思路
基于这些洞察，ZeRO-DP在实现MP的内存效率的同时，保留了DP的训练效率。ZeRO-DP通过划分模型状态而非复制它们（详见第5节），并采用一个动态通信调度，该调度利用了模型状态固有的时间局部性，同时最小化了通信量（详见第7节）。通过这样做，ZeRO-DP在增加DP度的同时，线性地减少了每台设备的模型内存占用，并将通信量维持在与默认DP相近的水平，从而保持了高效率。

4.2 ZeRO-R的洞察与概览

4.2.1 减少激活内存

两个关键洞察
a) MP复制激活内存：MP虽然划分了模型状态，但通常需要复制激活内存。例如，如果我们将一个线性层的参数垂直切分并在两个GPU上并行计算，每个GPU都需要完整的激活值来计算其分区部分。
b) 大模型的算术强度高：对于像GPT-2或更大的模型，算术强度（每次迭代的计算量与每次迭代的激活检查点量的比率）非常大（≥10K），并且随隐藏维度线性增加。这使得即使在带宽较低的情况下，也能隐藏激活检查点的数据移动成本。

ZeRO-R的解决方案
ZeRO通过在GPU间划分激活检查点来消除MP中的内存冗余，并按需使用all-gather操作来重构它们。激活内存的占用与MP度成比例减少。对于非常大的模型，ZeRO甚至可以选择将激活分区移动到CPU内存，由于这些模型具有高算术强度，仍然可以实现良好的效率。

4.2.2 管理临时缓冲区

恒定大小缓冲区策略
ZeRO-R使用恒定大小的缓冲区，以避免临时缓冲区随着模型尺寸的增大而膨胀，同时使其足够大以保持效率。

4.2.3 管理内存碎片

内存碎片产生的原因与洞察
内存碎片是由于短生命周期和长生命周期内存对象的交错使用造成的。在前向传播中，激活检查点是长生命周期的，而重新计算的激活值是短生命周期的。类似地，在后向计算中，参数梯度是长生命周期的，而激活梯度是短生命周期的。基于这一洞察，ZeRO采取了相应措施。

ZeRO的内存碎片整理方案
ZeRO通过将激活检查点和梯度移动到预分配的连续内存缓冲区中，进行即时内存碎片整理。这不仅增加了内存可用性，还通过减少内存分配器寻找空闲连续内存的时间来提高效率。

A2 方法细节

5. 深入ZeRO-DP

ZeRO-DP的核心思想
现有的数据并行（DP）方法在每个设备上复制模型状态，导致了显著的内存开销。ZeRO-DP通过在数据并行进程间划分它们——包括优化器状态、梯度和参数——来消除这种内存冗余。图1量化并可视化了使用和不使用ZeRO-DP时的内存需求。该图展示了在累积地划分（1）优化器状态、（2）梯度和（3）参数冗余后的内存占用情况。我们将这三个步骤称为ZeRO-DP的三个优化阶段：$P_{os}$、$P_g$ 和 $P_p$，下面将详细阐述。

图1：比较模型状态在不同设备上的内存消耗，展示了ZeRO-DP三个优化阶段的效果。Ψ表示模型大小（参数数量），K表示优化器状态的内存乘数，Nd表示DP度。示例中，我们假设模型大小Ψ = 7.5B，DP度Nd = 64，并基于使用Adam优化器的混合精度训练，设K = 12。
图1：比较模型状态在不同设备上的内存消耗，展示了ZeRO-DP三个优化阶段的效果。Ψ表示模型大小（参数数量），K表示优化器状态的内存乘数，Nd表示DP度。示例中，我们假设模型大小Ψ = 7.5B，DP度Nd = 64，并基于使用Adam优化器的混合精度训练，设K = 12。

5.1 $P_{os}$ : 优化器状态划分

划分与更新机制
对于一个数据并行度为 $N_d$ 的情况，我们将优化器状态分为 $N_d$ 个相等的分区，使得第 $i$ 个数据并行进程只更新对应于第 $i$ 个分区的优化器状态。因此，每个数据并行进程只需存储和更新全部优化器状态的 $1/N_d$，并相应地只更新 $1/N_d$ 的参数。在每个训练步骤结束时，我们通过在数据并行进程间执行一次 all-gather 操作，来获取所有进程上完全更新后的参数。

内存节省效果
如图1所示，经过优化器状态划分后，内存消耗从 $4\Psi + K\Psi$ 减少到 $4\Psi + K\Psi/N_d$。在图1所描绘的具体例子中，一个7.5B参数的模型在使用$P_{os}$和64路数据并行（$N_d=64$）时需要31.4GB内存，而标准DP则需要120GB。此外，当 $N_d$ 很大时，模型状态的内存需求从 $4\Psi + 12\Psi = 16\Psi$ 字节减少到 $4\Psi + 12\Psi/N_d \approx 4\Psi$ 字节，实现了4倍的缩减。

5.2 $P_g$: 梯度划分

划分与归约机制
由于每个数据并行进程只更新其对应的参数分区，因此它也只需要相应参数的归约后梯度。因此，在后向传播过程中，当每一层的梯度可用时，我们只在负责更新相应参数的数据并行进程上对它们进行归约。归约之后，我们不再需要这些梯度，它们的内存可以被释放。这将持有梯度所需的内存占用从 $2\Psi$ 字节减少到 $2\Psi/N_d$。

实现细节与效率
这实际上是一个Reduce-Scatter操作，其中对应不同参数的梯度被归约到不同的进程。为了在实践中更高效，我们采用了一种分桶（bucketization）策略，即我们将对应于特定分区的所有梯度分到同一个桶中，并一次性对整个桶执行归约操作。这在精神上类似于NVIDIA的AMP【索引25，Automatic mixed-precision，NVIDIA，2019】优化器如何将all-reduce梯度计算分桶以重叠通信和计算。在我们的案例中，我们在分区边界执行reduce操作而不是all-reduce，以减少内存占用并重叠计算与通信。

内存节省效果
通过同时移除梯度和优化器状态的冗余，我们将内存占用进一步降低到 $2\Psi + 14\Psi/N_d \approx 2\Psi$。如图1中的例子所示，一个7.5B参数的模型在使用$P_{os+g}$和64路DP（$N_d=64$）时仅需16.6GB内存，而标准DP则需要120GB。当$N_d$很大时，模型状态的内存需求从$2\Psi + 14\Psi = 16\Psi$字节减少到$2\Psi + 14\Psi/N_d \approx 2\Psi$字节，实现了8倍的缩减。

5.3 $P_p$: 参数划分

划分与通信机制
与优化器状态和梯度一样，每个进程只存储与其分区相对应的参数。当需要其分区之外的参数进行前向和后向传播时，这些参数会通过广播从相应的数据并行进程接收。虽然这乍一看似乎会产生显著的通信开销，但我们证明，这种方法仅将基线DP系统的总通信量增加了1.5倍，同时实现了与 $N_d$ 成比例的内存减少。

内存节省效果
通过参数划分，我们将一个 $\Psi$ 参数模型的内存消耗从 $16\Psi$ 减少到 $16\Psi/N_d$。如图1中的例子所示，一个7.5B参数的模型在使用$P_{os+p+g}$和64路DP（$N_d=64$）时，模型状态内存仅需1.9GB，而标准DP则需要120GB。这具有深远的意义：只要有足够的设备来共享模型状态，ZeRO就能使DP适应任意大小的模型。

5.4 对模型大小的影响

三个阶段的内存缩减效果
划分的三个阶段 $P_{os}$、$P_{os+g}$ 和 $P_{os+g+p}$ 分别将每个数据并行进程在模型状态上的内存消耗减少了最多4倍、8倍和 $N_d$ 倍。表1分析了几个示例模型在ZeRO-DP三个优化阶段下，随DP度变化的模型状态内存消耗。若不使用ZeRO，无论DP度如何，内存消耗都等于表中的第一行。请注意，当 $N_d=64$ 时，使用 $P_{os}$、$P_{os+g}$ 和 $P_{os+g+p}$ 的ZeRO分别可以训练高达7.5B、14B和128B参数的模型。当 $N_d=1024$ 时，启用所有优化的ZeRO（$P_{os+g+p}$）可以训练拥有1万亿参数的模型！甚至可能训练任意大小的模型！而没有ZeRO，仅凭DP能运行的最大模型参数不到15亿。


表1：ZeRO-DP中不同优化的每设备内存消耗与DP度的函数关系。粗体文本表示模型可以放入32GB V100 GPU集群的组合。

6. 深入ZeRO-R

6.1 $P_a$: 分区激活检查点

消除激活复制冗余
如4.2节所述，模型并行（MP）的设计 inherently 需要复制激活值，导致模型并行GPU之间存在冗余的激活副本。ZeRO通过划分激活值来消除这种冗余，并且只在计算使用前，逐个激活层地将它们以复制形式实例化。具体来说，一旦模型某层的前向传播计算完成，其输入激活值就会在所有模型并行进程间被划分，直到反向传播再次需要它时。此时，ZeRO使用一个all-gather操作来重新物化一个复制的激活副本。我们将此优化称为$P_a$。它与激活检查点【索引7，Training deep nets with sublinear memory cost，Tianqi Chen et al.，2016】协同工作，只存储分区的激活检查点而不是复制的副本。此外，在模型非常大且设备内存非常有限的情况下，这些分区的激活检查点还可以被卸载到CPU，以额外的通信成本将激活内存开销降至几乎为零，我们将在第7节讨论这一点。我们称之为$P_{a+cpu}$。

内存节省效果
通过分区激活检查点，ZeRO将激活内存占用减少了一个与模型并行（MP）度成比例的因子。考虑训练一个如表4所示的100B模型，批量大小为32，序列长度为1024，MP度为16。如果我们为每个transformer层检查点一个激活，仅存储激活检查点每GPU就需要约33GB内存。但使用ZeRO中的$P_a$，这个需求可以减少到每GPU约2GB。此外，这2GB还可以被卸载到CPU，将激活的内存占用降至几乎为零。

6.2 CB: 恒定大小缓冲区

平衡内存与计算效率
ZeRO仔细选择临时数据缓冲区的大小，以平衡内存和计算效率。在训练期间，某些操作的计算效率高度依赖于输入大小，较大的输入能实现更高的效率。例如，一个大的all-reduce操作比小的操作能获得高得多的带宽。因此，为了获得更好的效率，像NVIDIA Apex或Megatron这样的高性能库会在应用这些操作前，将所有参数融合成一个单一的缓冲区。然而，融合缓冲区的内存开销与模型大小成正比，可能会成为瓶颈。例如，对于一个3B参数的模型，一个32位的融合缓冲区将需要12GB内存。为了解决这个问题，当模型变得过大时，我们简单地使用一个性能高效的恒定大小融合缓冲区。这样做，缓冲区大小就不再依赖于模型大小，并且通过保持缓冲区足够大，我们仍然可以实现良好的效率。

6.3 MD: 内存碎片整理

内存碎片的成因
模型训练中的内存碎片是由于激活检查点和梯度计算过程中的内存使用模式造成的。在使用激活检查点的前向传播过程中，只有选定的激活被存储用于反向传播，而大多数激活被丢弃，因为它们可以在反向传播时重新计算。这造成了短生命周期内存（被丢弃的激活）和长生命周期内存（被检查点的激活）的交错，导致内存碎片。类似地，在反向传播过程中，参数梯度是长生命周期的，而激活梯度以及计算参数梯度所需的任何其他缓冲区是短生命周期的。这种短期和长期内存的交错再次导致内存碎片。

内存碎片带来的问题
有限的内存碎片在内存充裕时通常不是问题，但对于在有限内存下进行的大模型训练，内存碎片会导致两个问题：i) 即使有足够的可用内存，也因为缺少连续内存而导致内存不足（OOM），ii) 由于内存分配器花费大量时间搜索连续内存块以满足内存请求，导致效率低下。

ZeRO的解决方案
ZeRO通过为激活检查点和梯度预分配连续的内存块，并在它们产生时将其复制到预分配的内存中，从而实现即时内存碎片整理。MD不仅使ZeRO能够用更大的批量大小训练更大的模型，而且在内存有限的情况下进行训练时也能提高效率。

7. ZeRO-DP的通信分析

引言
由于ZeRO通过消除内存冗余来提升模型大小，一个自然的问题是，我们是否用通信量换取了内存效率。换言之，与基线数据并行（DP）方法相比，ZeRO驱动的DP方法的通信量是多少？答案分两部分：i) 使用$P_{os}$和$P_g$时，ZeRO-DP不产生额外通信，同时实现高达8倍的内存缩减；ii) 在$P_{os}$和$P_g$的基础上使用$P_p$时，ZeRO-DP最多产生1.5倍的通信量，同时将内存占用进一步减少$N_d$倍。本节将对此进行分析。我们首先简要概述标准DP的通信量。

7.1 数据并行通信量

标准DP的通信机制
在数据并行训练期间，所有数据并行进程的梯度在反向传播结束时进行平均，然后计算下一步的更新。这个平均过程通过一个all-reduce通信集合操作来完成。对于大模型，all-reduce通信完全受限于通信带宽，因此，我们的分析仅限于每个数据并行进程发送和接收的总通信量。

All-reduce的实现与通信量计算
最先进的all-reduce实现采用两步法：第一步是reduce-scatter操作，它将数据的不同部分归约到不同的进程上；第二步是all-gather操作，每个进程收集所有进程上的归约后数据。这两步的结果就是一个all-reduce。reduce-scatter和all-gather都使用流水线方法实现，对于一个包含$\Psi$个元素的数据，每一步都会产生$\Psi$个元素的总数据移动量。因此，标准DP在每个训练步骤中产生$2\Psi$的数据移动量。

7.2 ZeRO-DP通信量

7.2.1 使用$P_{os+g}$的通信量

通信操作与总量
使用梯度划分后，每个进程只存储更新其对应参数分区所需的梯度部分。因此，ZeRO不再需要all-reduce，而只需对梯度执行一次scatter-reduce操作，产生的通信量为$\Psi$。在每个进程更新了它负责的参数分区后，会执行一次all-gather操作，以从所有数据并行进程收集所有更新后的参数。这也产生$\Psi$的通信量。所以，每个训练步骤的总通信量为$\Psi + \Psi = 2\Psi$，与基线DP完全相同。

7.2.2 使用$P_{os+g+p}$的通信量

通信操作与总量
参数划分后，每个数据并行进程只存储它所更新的参数。因此，在前向传播期间，它需要接收所有其他分区的参数。然而，这个过程可以流水线化以避免内存开销。在计算模型中对应于特定分区的部分之前，负责该分区的数据并行进程可以将权重广播给所有数据并行进程。该分区的前向传播完成后，这些参数就可以被丢弃。总通信量因此为$\Psi \times (N_d-1)/N_d \approx \Psi$。换句话说，我们将参数的all-gather操作重新调度，将其分散到整个前向传播过程中，并在使用后丢弃参数。但请注意，这个all-gather在反向传播中需要以相反的顺序再次发生。

最终通信量分析
因此，总通信量是这些all-gather操作产生的通信量与梯度reduce-scatter产生的通信量之和。总通信量因此为$3\Psi$，是基线DP的1.5倍。梯度和参数划分都利用了“并非所有梯度和参数状态在任何时候都需要”这一洞察，通过明智地通信状态来优化内存。

8. ZeRO-R的通信分析

引言
我们比较了ZeRO-R中分区激活检查点（$P_a$）与基线模型并行（MP）的通信量，并表明$P_a$引入的通信量增量通常不到基线MP的十分之一。此外，我们分析了$P_a$相对于数据并行（DP）通信量的通信开销，以确定$P_a$在何种情景下能通过允许更大的批量大小和减少DP通信来提高效率。我们利用这种分析来决定是否以及何时应用$P_a$和$P_{a+cpu}$。

$P_a$通信开销分析
分区激活检查点的通信量权衡取决于模型大小、检查点策略和MP策略。为了分享具体的见解，我们在基于SOTA MP方法Megatron-LM实现的Transformer模型背景下进行分析。

与Megatron-LM的对比
在带有激活检查点的Megatron-LM中，每个transformer块在前向传播中执行两次大小为 batch × seq_length × hidden_dim 的all-reduce操作，在前向重计算中执行两次，在后向传播中再执行两次。由于一次all-reduce的通信量是 2 × message_size，每个块的总通信量是 12 × seq_length × hidden_dim。

ZeRO-R的增量通信
当ZeRO-R划分激活检查点时，它需要在反向传播的每次激活检查点的前向重计算之前，增加一次all-gather操作。通常，我们为每个transformer块检查点输入激活，因此每个transformer块需要一次all-gather。由于all-gather的通信量是 message_size，$P_a$的通信开销因此是 seq_length * hidden_dim。所以，$P_a$的总通信开销不到模型并行原始通信量的10%。

$P_a$对DP通信的影响
当MP与DP结合使用时，$P_a$可以用来将数据并行通信量减少一个数量级，代价是模型并行通信量增加10%，并且在数据并行通信成为性能瓶颈时显著提升效率。注意，$P_a$将激活内存消耗减少了MP度，允许批量大小成比例增加。对于大型模型，MP度可以高达16（DGX-2节点上的GPU数量），允许批量大小增加最多16倍。数据并行训练的通信量与批量大小成反比。因此，由于$P_a$带来的批量大小数量级增加，可能导致数据并行通信量减少一个数量级。

$P_{a+cpu}$的通信开销
最后，如果应用$P_{a+cpu}$，分区的激活检查点被卸载到CPU，将激活内存需求降至几乎为零，代价是与$P_a$相比，增加了2倍的往返CPU内存的数据移动。在极端情况下，即使使用了$P_a$，DP通信量由于批量大小过小仍是主要瓶颈时，只要CPU数据传输开销小于DP通信量开销，$P_{a+cpu}$就能通过增加批量大小来提高效率，这对于小批量大小通常是成立的。

决策依据
给定模型和硬件特性，我们利用上述分析来决定是否以及何时应用$P_a$和$P_{a+cpu}$。

9. 迈向万亿参数模型

现状与挑战
目前已发表的最大模型参数量在百亿级别，训练这些模型已极具挑战性。达到万亿参数，即增大三个数量级，是必然趋势，但这条路将充满障碍、惊喜和创新。虽然我们不声称知道或解决了所有问题，但ZeRO从系统角度解决了最根本的挑战之一：能够在当前硬件上容纳如此规模的模型，并以良好的系统可扩展性进行训练。

相较于SOTA的飞跃
当前最先进的框架Megatron能够以可接受的吞吐量训练的最大模型是160-200亿参数的模型，这需要在一个DGX-2系统内完成。通过跨多个DGX节点进行模型并行来进一步扩展，会因节点间带宽有限而导致效率显著下降。

ZeRO的贡献
ZeRO极大地增加了可高效运行的模型大小。它使当前一代硬件能够运行大得多的模型，而无需细粒度的模型并行跨越节点边界。如表1所示，ZeRO在启用所有优化（$P_{os+g+p}$）后，仅使用数据并行就能在1024个GPU上容纳超过1万亿参数。或者，当与模型并行结合时（如表2所示），ZeRO可以在1024个GPU上，通过16路模型并行（在每个DGX2节点内）和64路数据并行（跨节点）来容纳超过1万亿参数。高效运行万亿参数模型不再是不可能的！

算力差距
然而，在可接受的时间范围内端到端地训练一个万亿参数模型，可能仍需要巨大的算力，这是当今AI集群所缺乏的。

资源需求估算
为了理解资源需求，我们与Bert-Large做一个简要比较。Bert-Large可以在一个1024 GPU的DGX-2H集群上训练67分钟【索引26，未在文中列出详细信息】。一个1万亿参数的模型，对于一个数据样本的计算量，可以轻易地比Bert-Large模型多3000倍（1万亿 / 3.3亿）。即使我们假设训练模型所需的序列长度和总样本数相同，在相同的硬件和相似的计算效率下，训练一个1T模型也需要140天。实际上，数据样本和序列长度很可能随模型大小增加而增加，导致训练时间超过一年。这将需要一个百亿亿次（exa-flop）级别的系统才能在合理时间内训练一个1T参数模型。但当这样的计算能力可用时，我们希望ZeRO能提供高效运行1T模型的系统技术。

A4 实验环境

  • 模型架构：实验采用类似GPT-2的Transformer模型。通过改变隐藏层维度和层数来构建不同参数规模的模型。具体模型配置参数见于论文中的表格（Table 4及附录）。
  • 硬件配置：
    • GPU：在一个由400个NVIDIA V100 GPU组成的集群上进行实验。
    • 平台：该集群由25个DGX-2节点构成。
    • 网络：节点间通信带宽为800 Gbps。
  • 软件配置：
    • 实现：ZeRO-100B在PyTorch中实现，接口与torch.nn.module兼容。
    • 基线：
      • 对于不使用模型并行的实验，基线是PyTorch的分布式数据并行（DDP）。
      • 对于使用模型并行的实验，基线是当时最先进的Megatron-LM框架（2019年9月版本）。
    • 依赖：实验中，ZeRO驱动的数据并行与Megatron-LM的模型并行相结合进行评估。

A4 实验结果

本文实现的ZeRO子集（$P_{os+g}$ + ZeRO-R），命名为 ZeRO-100B，在实验中展现了卓越的性能。

  • 速度与模型规模

    • 内容：比较了ZeRO-100B（结合MP）与单独使用Megatron-LM（MP）在不同模型大小下的单GPU吞吐量。
    • 结果：ZeRO-100B能在400个GPU上高效运行高达1700亿参数的模型，而Megatron-LM在超过40B参数后性能急剧下降。对于8B到100B参数的模型，ZeRO-100B实现了平均超过30%峰值性能的15 Petaflops总吞吐量。相较于基线，ZeRO-100B的模型规模提升超过8倍，速度提升高达10倍。
    • 结论：ZeRO通过避免跨节点模型并行的高通信开销，显著提升了大模型的训练速度和可支持规模。（见图2）

    图2：ZeRO训练吞吐量和相对于SOTA基线的加速比，随模型大小变化。对于ZeRO，MP总是在一个节点内；对于基线，大于40B的模型需要跨节点MP。
    图2：ZeRO训练吞吐量和相对于SOTA基线的加速比，随模型大小变化。对于ZeRO，MP总是在一个节点内；对于基线，大于40B的模型需要跨节点MP。

  • 超线性可扩展性

    • 内容：评估了ZeRO-100B在训练一个60B参数模型时，从64个GPU扩展到400个GPU的性能变化。
    • 结果：在64到400个GPU的范围内观察到超线性加速，即GPU数量翻倍时，性能增长超过两倍。
    • 结论：这是因为随着数据并行度（DP degree）增加，ZeRO-DP减少了每个GPU的内存占用，从而可以容纳更大的批次大小，提升了计算强度和吞吐量。（见图3）

    图3：使用ZeRO-100B训练60B参数模型的超线性可扩展性和单GPU训练吞吐量。
    图3：使用ZeRO-100B训练60B参数模型的超线性可扩展性和单GPU训练吞吐量。

  • 普及大规模模型训练

    • 内容：展示了仅使用ZeRO驱动的数据并行（无模型并行）所能训练的模型规模和性能。
    • 结果：ZeRO-100B在128个GPU上，无需模型并行即可训练高达130亿参数的模型，平均单GPU吞吐量超过40 TFlops。相比之下，标准数据并行最大只能训练14亿参数的模型，吞吐量不足20 TFlops。
    • 结论：ZeRO极大地降低了训练大型模型的门槛，数据科学家无需进行复杂的模型重构即可训练超过现有SOTA（如T5-11B）规模的模型，且不依赖于昂贵的节点内高速互联（如NVLINK）。（见图4）

    图4：使用ZeRO-DP的最大模型吞吐量。
    图4：使用ZeRO-DP的最大模型吞吐量。

  • 不同优化项的内存与性能分析

    • 内容：分析了ZeRO不同优化配置（C1-C5，如表3）对最大模型尺寸、内存消耗和性能的影响。

    • 结果：

      • 最大模型尺寸：启用$P_a$（分区激活）使模型尺寸从40B增至60B；启用$P_{os+g}$（优化器与梯度划分）进一步跃升至140B；启用$P_{a+cpu}$（激活卸载至CPU）再增至150B。（见图6）
      • 内存消耗：各优化项按预期减少了PyTorch缓存的峰值内存。对于100B模型，$P_{a+cpu}$的内存节省效果比40B模型更明显，证明其对极大模型的价值。（见图7）
      • 性能：性能提升与内存消耗降低正相关，因为更低的内存占用允许更大的批次。唯一的例外是$P_{a+cpu}$，由于CPU数据传输开销，它通常会降低性能，但对于无法运行或批次极小（如170B模型）的超大模型，它却是必要且有益的。（见图8）



      图6：最大模型尺寸

      图7：最大缓存分配

      图8：每GPU吞吐量

  • Turing-NLG：170亿参数的SOTA语言模型

    • 内容：使用ZeRO-100B端到端训练了Turing-NLG模型。
    • 结果：Turing-NLG成为当时全球最大的语言模型（17B参数），并在Webtext-103困惑度上取得了10.21的新SOTA记录。训练过程中，ZeRO-100B为该模型实现了41.4 TFlops/GPU的持续吞吐量。
    • 结论：ZeRO的系统性突破成功赋能了创纪录规模和精度的语言模型训练。（见图5）

    图5：由ZeRO赋能的SOTA模型Turing-NLG。
    图5：由ZeRO赋能的SOTA模型Turing-NLG。

A5 结论

从高性能计算和系统角度看，ZeRO代表了大型模型训练领域的一场革命性变革。本文的实现ZeRO-100B，已经将模型规模扩大了8倍，吞吐量提升超过10倍，在现代GPU集群上实现了超线性加速，并成功训练出当时世界最大的模型。然而，这仅仅是冰山一角。完整的ZeRO有潜力将模型规模再提升一个数量级，为未来万亿参数模型的训练提供可能。

作者对ZeRO最感乐观的一点是，它没有给数据科学家带来额外负担。与模型并行（MP）和流水线并行（PP）等现有方法不同，ZeRO无需重构模型，使用起来和标准数据并行一样简单，这使其成为未来大规模模型训练研究的首选方案。通过开源和社区反馈，作者计划让ZeRO完全服务于深度学习社区，以催化大规模模型训练的演进和普及。

A6 附录

AE 附录

模型配置与实验细节
这些表格包含了论文中实验所使用的所有模型配置和批量大小。


表5：与图2相关的模型配置，涉及ZeRO与基线吞吐量对比。

表6：与图3相关的模型配置，涉及超线性可扩展性。

表7：与图4相关的模型配置，涉及不同ZeRO配置下的最大模型尺寸。

表8：与图5相关的模型配置，涉及不同ZeRO配置下的内存分配。

表9：与图6相关的模型配置，涉及不同ZeRO配置下的吞吐量。

表10：与图7相关的模型配置，涉及在仅使用数据并行时评估最大模型尺寸与吞吐量。

实验设置的补充说明
在图2中，请注意一些基线实验的GPU总数是384或256，而不是ZeRO的400。这是因为GPU总数必须是模型并行（MP）度的倍数，而我们总共只有400个GPU可用。模型配置值还存在一些额外的约束，例如隐藏层大小必须能被注意力头数整除，隐藏层大小必须能被MP度整除，以及注意力头数必须能被MP度整除。对于基线，我们使用了能够容纳模型的最小的2的幂次方的GPU数量。例如，对于170B参数模型，基线使用的是256个GPU。由于我们只有400个GPU，我们只能用256个GPU运行基线。

基线对比的公平性说明
我们想指出，这实际上给了基线一个优势，因为更少的GPU意味着基线有更好的通信吞吐量。例如，在170B参数模型的情况下，基线的数据并行（DP）度为1，因此实际上没有DP通信开销。尽管基线有此优势，本文呈现的结果依然优越。

性能指标的说明
此外，我们想指出，我们比较的是每个GPU的性能，而不是聚合性能，因此，尽管给基线带来了轻微优势，结果仍然是同类比较（apples-to-apples）。