ZeRO：迈向训练万亿参数模型的内存优化技术

Samyam Rajbhandari*, Jeff Rasley*, Olatunji Ruwase, Yuxiong He {samyamr, jerasley, olruwase, yuxhe}@microsoft.com

摘要

大型深度学习模型带来了显著的精度提升，但训练数十亿到万亿参数的模型极具挑战性。现有的数据并行和模型并行等解决方案在将这些模型适配到有限的设备内存中时，同时保持计算、通信和开发效率方面存在根本性的局限。我们开发了一种名为“零冗余优化器”（ZeRO）的新型解决方案，旨在优化内存，在大幅提高可高效训练的模型规模的同时，显著提升训练速度。ZeRO 消除了数据并行和模型并行训练中的内存冗余，同时保持了低通信量和高计算粒度，使我们能够以持续的高效率将模型规模扩展到与设备数量成正比的水平。我们对内存需求和通信量的分析表明：ZeRO 有潜力利用当今的硬件将模型扩展到超过 1 万亿个参数。

我们实现并评估了 ZeRO：它在 400 个 GPU 上以超线性加速比训练了超过 100B 参数的大型模型，实现了 15 Petaflops 的吞吐量。这代表了模型规模提升了 8 倍，性能较现有最先进技术提升了 10 倍。在易用性方面，ZeRO 可以训练高达 13B 参数的大型模型（例如，比 Megatron GPT 8.3B 和 T5 11B 更大），而无需科学家难以应用的模型并行技术。最后，研究人员利用 ZeRO 的系统突破，创建了世界上最大（17B 参数）且具有破纪录精度的语言模型。

1 扩展引言

深度学习（DL）模型正变得越来越大，模型规模的增加带来了显著的精度提升。在自然语言处理（NLP）领域，Transformer 架构为诸如 Bert-large (0.3B) [1]、GPT-2 (1.5B) [2]、Megatron-LM (8.3B) [3]、T5 (11B) [4] 等大型模型铺平了道路。为了实现模型规模从数百亿到万亿参数的持续增长，我们面临着训练它们的挑战——它们显然无法装入单个设备（如 GPU 或 TPU）的内存中，而简单地增加设备数量并不能帮助扩展训练。

基础的数据并行（DP）并不能减少每个设备的内存占用，在当前一代具有 32GB 内存的 GPU 上，对于超过 1.4B 参数的模型，内存就会耗尽。其他现有的解决方案，如流水线并行（PP）、模型并行（MP）、CPU 卸载等，在功能性、易用性以及内存和计算/通信效率之间进行了权衡，但所有这些对于实现速度和规模的训练都至关重要。

在训练大型模型的现有不同解决方案中，MP 可能是最有前途的。当前文献中最大的模型，即 11B T5 模型 [4] 和 Megatron-LM 8.3B [3]，都是由模型并行驱动的，分别在 Mesh-Tensorflow [5] 和 Megatron-LM [3] 中实现。然而，MP 无法在这些模型规模之外进一步扩展。MP 在垂直方向上拆分模型，将每一层的计算和参数划分到多个设备上，需要在每一层之间进行大量的通信。结果是，它们在单节点内（节点间 GPU 通信带宽较高）工作良好，但效率在跨节点后迅速下降 [3]。我们使用 Megatron-LM 在两个 DGX-2 节点上测试了一个 40B 参数的模型，观察到每个 V100 GPU 的性能约为 5 Tflops（不到硬件峰值的 5%）。

那么，我们如何克服现有解决方案的局限性，更高效地训练大型模型呢？为了回答这个问题，我们首先分析了现有系统在模型训练中的全谱内存消耗，并将其分为两部分：1）对于大型模型，大部分内存被模型状态（model states）占用，其中包括优化器状态（如 Adam [6] 中的动量和方差）、梯度和参数。2）剩余内存被激活值、临时缓冲区和不可用的碎片内存消耗，我们统称为残余状态（residual states）。我们开发了 ZeRO（零冗余优化器）来优化这两部分的内存效率，同时获得高计算和通信效率。由于这两部分面临不同的挑战，我们相应地开发并讨论了它们的解决方案。

优化模型状态内存：模型状态在训练过程中通常消耗最多的内存，但现有的 DP 和 MP 方法并没有提供令人满意的解决方案。DP 具有良好的计算/通信效率，但内存效率较差；而 MP 可能具有较差的计算/通信效率。更具体地说，DP 在所有数据并行进程中复制整个模型状态，导致冗余的内存消耗；而 MP 虽然通过划分这些状态获得了高内存效率，但往往导致过于细粒度的计算和昂贵的通信，从而降低了扩展效率。此外，所有这些方法都在整个训练过程中静态地维护所有必需的模型状态，即使并非所有模型状态在训练期间都一直需要。基于这些观察，我们开发了 ZeRO-DP（ZeRO 驱动的数据并行），它实现了 DP 的计算/通信效率，同时达到了 MP 的内存效率。ZeRO-DP 通过划分模型状态而不是复制它们，消除了跨数据并行进程的内存状态冗余，并通过在训练期间使用动态通信调度，保留了 DP 的计算粒度和通信量，从而保持了计算/通信效率。

ZeRO-DP 有三个主要的优化阶段（如图 1 所示），分别对应于优化器状态、梯度和参数的划分。当累积启用时：

优化器状态划分 ( $P_{os}$ )：4 倍内存减少，与 DP 相同的通信量；
增加梯度划分 ( $P_{os+g}$ )：8 倍内存减少，与 DP 相同的通信量；
增加参数划分 ( $P_{os+g+p}$ )：内存减少与 DP 程度 $N_d$ 成线性关系。例如，在 64 个 GPU 上拆分 ( $N_d = 64$ ) 将产生 64 倍的内存减少。通信量有适度的 50% 增加。

ZeRO-DP 消除了内存冗余，并使集群的全部聚合内存容量可用。在启用所有三个阶段的情况下，ZeRO 可以在仅 1024 个 NVIDIA GPU 上训练万亿参数模型。

$图 1：比较三种 ZeRO-DP 优化阶段下，每设备模型状态的内存消耗。\Psi 表示模型大小（参数数量），K 表示优化器状态的内存乘数，N_d 表示 DP 程度。在示例中，我们假设模型大小为 \Psi = 7.5B，DP 程度为 N_d = 64，基于混合精度训练和 Adam 优化器，K = 12。$

使用 Adam [6] 等优化器在 16 位精度下训练万亿参数模型，大约需要 16 TB 的内存来容纳优化器状态、梯度和参数。16 TB 除以 1024 等于 16 GB，这完全在 GPU 的合理范围内（例如，具有 32 GB 的片上内存）。

优化残余状态内存：在 ZeRO-DP 提升了模型状态的内存效率后，由激活值、临时缓冲区和不可用内存碎片消耗的剩余内存可能成为第二个内存瓶颈。我们开发了 ZeRO-R 来分别优化这三个因素消耗的残余内存。

对于激活值（为了执行反向传播而从前向传播中存储），我们注意到检查点技术 [7] 有所帮助，但对于大型模型来说还不够。因此，ZeRO-R 通过识别并消除现有 MP 方法中的激活值复制（通过激活值划分）来优化激活内存。它还在适当的时候将激活值卸载到 CPU。
ZeRO-R 定义了临时缓冲区的适当大小，以在内存和计算效率之间取得平衡。
我们观察到，由于不同张量的生命周期变化，训练期间会出现内存碎片。由于碎片化导致缺乏连续内存，即使有足够的空闲内存，也可能导致内存分配失败。ZeRO-R 根据张量的不同生命周期主动管理内存，防止内存碎片化。

ZeRO-DP 和 ZeRO-R 结合在一起，形成了一套强大的 DL 训练内存优化系统，我们统称为 ZeRO。

ZeRO 与 MP：由于 ZeRO 消除了 DP 中的内存低效问题，自然会问：我们还需要 MP 吗？什么时候需要？ZeRO 如何与 MP 配合使用？有了 ZeRO，MP 对于仅为了适配大型模型而言，变成了一个吸引力较低的选项。ZeRO-DP 在减少每设备内存占用方面至少与 MP 一样有效，或者在 MP 无法均匀划分模型时有时更有效。它还具有相当或更好的扩展效率。此外，数据并行非常易于使用，因此在不同的工作负载中广泛适用，而 MP 方法通常需要模型开发者修改模型，系统开发者开发分布式算子，且现有的工作（如 Megatron-LM）仅支持有限的算子和模型。

话虽如此，在某些情况下我们仍然希望利用 MP：i) 当与 ZeRO-R 一起使用时，MP 可以减少超大型模型的激活内存占用。ii) 对于激活内存不是问题的较小模型，当仅使用 DP 的聚合批次大小太大而无法获得良好的收敛性时，MP 也可以带来好处。在这种情况下，可以将 ZeRO 与 MP 结合使用，以在可接受的聚合批次大小下适配模型。

我们展示了 ZeRO 可以与 MP 结合，在每个设备上实现 $N_d \times N_m$ 倍的最大理论内存减少（其中 DP 程度为 $N_d$ ，MP 程度为 $N_m$ ）。这使我们能够在 1024 个 GPU 上适配万亿参数模型，并使用 16 路模型并行（在每个 DGX2 节点内）和 64 路跨节点数据并行，并使用适度的批次大小高效运行它！

实现与评估：ZeRO 中的全套优化技术可以让我们在当今的高端硬件集群（例如，具有 1K V100 GPU）上运行万亿参数模型，然而，硬件计算能力仍然太有限，训练时间可能长得不切实际（>1 年）。因此，我们本次实现的重点是高效支持比现有最先进技术（SOTA）参数多 10 倍（约 100B 参数）的模型，同时仍在当前硬件的计算能力范围内。我们实现并评估了 ZeRO 中的一个优化子集，称为 ZeRO-100B（ZeRO-DP 的 $P_{os+g}$ 加上 ZeRO-R），使我们能够实现这一目标。结果表明：

模型规模：结合 MP，ZeRO-100B 可以高效运行 170B 参数的模型，而像仅使用 Megatron 的现有系统无法在超过 40B 参数的情况下高效扩展，如图 2 所示。这比 SOTA 的模型规模增加了 8 倍以上。
速度：改进的内存效率带来了更高的吞吐量和更快的训练速度。如图 2 所示，ZeRO 在 400 个 Nvidia V100 GPU 集群上运行 100B 参数模型，每个 GPU 的吞吐量超过 38 TFlops，聚合性能超过 15 Petaflops。这比相同模型规模的 SOTA 训练速度提高了 10 倍以上。

可扩展性：我们在 64-400 个 GPU 的范围内观察到超线性加速，当我们将 GPU 数量加倍时，性能提升超过两倍。这是 ZeRO-DP 的一个特性，它随着我们增加 DP 程度而减少了模型状态的内存占用，从而允许我们适配每个 GPU 更大的批次大小，从而获得更好的性能。我们预计随着 GPU 数量增加到 400 以上，这种行为将持续下去。
大型模型训练的民主化：ZeRO-100B 使数据科学家能够训练高达 13B 参数的模型，而无需任何需要模型重构的 MP 或 PP，其中 13B 比文献中最大的模型（T5，11B 参数）参数更多。因此，数据科学家可以自由地试验大型模型，而不必担心并行化。相比之下，现有系统（例如 PyTorch Distributed Data Parallel）在 1.4B 参数模型上就会耗尽内存。
新的 SOTA 模型：ZeRO 驱动了具有 17B 参数和破纪录精度的最大语言模型——Turing-NLG [9]。

我们将 ZeRO 作为我们开源 DL 训练优化库 DeepSpeed 的一部分进行共享。我们计划在 2020 年 5 月底之前发布本文中描述的所有实现，并通过启用 ZeRO-DP 第 3 阶段划分参数 ( $P_{os+g+p}$ ) 将其进一步扩展以支持 1 万亿参数。我们计划使 ZeRO 完全可供 DL 社区使用，以促进大规模模型训练的演进和民主化。

2 相关工作

2.1 数据、模型和流水线并行

并行化是大规模训练大型模型的核心策略。对于适合在设备内存中进行训练的模型，数据并行（DP）用于将训练扩展到多个设备。在 DP 中，模型参数在每个设备上被复制。在每一步中，小批量数据被均匀地划分到所有数据并行进程中，使得每个进程在不同的数据样本子集上执行前向和反向传播，并使用跨进程的平均梯度来本地更新模型。

当模型无法装入设备内存时，模型并行（MP）[5, 3] 和流水线并行（PP）[10, 11] 分别在垂直和水平方向上将模型拆分到进程中。第 1 节讨论了 ZeRO 如何与 DP 和 MP 相关。我们现在讨论 PP 以及它如何与减少内存消耗相关。

PP 在层之间水平拆分模型，在不同的设备上运行每个分区，并使用微批处理（micro-batching）来隐藏流水线气泡 [10, 11]。由于水平拆分和微批处理，诸如绑定权重（tied-weights）和批归一化（batch-normalization）之类的模型功能很难实现。流行的 PP 实现（如 G-pipe [10]）同时划分模型参数和总激活值，但需要与流水线分区数量成正比的批次大小来隐藏流水线气泡。大的批次大小会影响收敛速度，同时也需要大量的内存来存储激活值。PipeDream [12] 中的另一种 PP 实现保留了多个陈旧参数的副本以隐藏流水线气泡，而无需显著增加批次大小，使其内存效率较低。此外，该实现不等同于标准的 DL 训练，并对训练收敛有影响。相比之下，ZeRO 获得了与 PP 相同或更好的内存效率，而不会产生 PP 的功能、性能和收敛相关的限制。

2.2 基于非并行的方法来减少内存

除了 MP 和 PP，还有多条研究路线旨在减少 DL 训练的内存开销。

2.2.1 减少激活内存

多项工作专注于通过压缩 [13]、激活检查点 [7, 14] 或实时分析 [15] 来减少激活值的内存占用。这些努力是互补的，可以与 ZeRO 一起工作。事实上，ZeRO-R 中的激活内存减少与激活检查点并行工作。

2.2.2 CPU 卸载

[16, 17] 利用当今计算节点的异构特性，分别通过算法设计或虚拟内存将模型状态卸载到 CPU 内存。高达 50% 的训练时间可能花费在 GPU-CPU-GPU 传输上 [16]。ZeRO 的不同之处在于，它在不将模型状态存储到带宽因 PCI-E 而严重受限的 CPU 内存的情况下，显著减少了内存消耗。在极少数情况下，ZeRO-R 可能会为超大型模型卸载激活检查点以提高性能（详见 6.1 节）。

2.2.3 内存高效优化器

[18, 19] 专注于通过维护模型参数和梯度的更粗粒度统计信息来减少自适应优化方法的内存消耗，这可能会对模型收敛保证产生影响。ZeRO 与这些努力是正交的，其优化不会改变模型优化方法或影响模型收敛，而是有效地减少了每个设备的优化器状态和梯度的内存占用。

2.3 训练优化器

自适应优化方法 [20, 6, 21, 22] 对于实现大型模型的 SOTA 性能和精度至关重要。与 SGD 相比，它们以显著的内存占用为代价，为每个模型参数和梯度维护细粒度的一阶和二阶统计信息。ZeRO 可以将这些优化器的内存占用减少几个数量级，使得这些复杂的优化方法在具有适度设备内存的硬件上训练大型模型变得切实可行。它还使得开发和使用可能具有更好收敛性的更复杂和更耗内存的优化器成为可能。

3 内存都去哪了？

让我们退一步检查当前训练系统的内存消耗。例如，一个 1.5B 参数的 GPT-2 模型在 16 位精度下需要 3GB 的内存来存储其权重（或参数），然而，使用 Tensorflow 或 PyTorch，它无法在具有 32GB 内存的单个 GPU 上进行训练。人们可能想知道所有的内存都去哪了。在模型训练期间，大部分内存被模型状态消耗，即包含优化器状态、梯度和参数的张量。除了这些模型状态外，剩余的内存被激活值、临时缓冲区和碎片内存消耗，我们称之为残余状态。我们详细查看这两部分的内存消耗。

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

在训练期间，大部分设备内存被模型状态消耗。以 Adam [6] 为例，它是 DL 训练中最流行的优化器之一。Adam 需要存储两个优化器状态：i) 时间平均动量，ii) 梯度的方差，以计算更新。因此，要使用 Adam 训练模型，必须有足够的内存来容纳动量和梯度方差的副本。此外，还需要足够的内存来存储梯度和权重本身。在这三种与参数相关的张量中，优化器状态通常消耗最多的内存，特别是在应用混合精度训练时。

混合精度训练：在当前一代 NVIDIA GPU 上训练大型模型的最先进方法是通过混合精度（fp16/32）训练 [23]，其中参数和激活值存储为 fp16，从而能够使用这些 GPU 上的高吞吐量张量核心单元 [24]。在混合精度训练期间，前向和反向传播都使用 fp16 权重和激活值执行。然而，为了在反向传播结束时有效地计算和应用更新，混合精度优化器保留了参数的 fp32 副本以及所有其他优化器状态的 fp32 副本。

让我们以 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 亿参数的模型，这导致至少 24 GB 的内存需求，这显著高于仅容纳 fp16 参数所需的微薄 3 GB 内存。

3.2 残余内存消耗

在训练期间，激活值会占用大量内存 [7]。作为一个具体示例，使用序列长度为 1K 和批次大小为 32 训练的 1.5B 参数 GPT-2 模型需要大约 60 GB 的内存。激活检查点（或激活重计算）是一种常见的减少激活内存的方法，它以 33% 的重计算开销为代价，将激活内存减少约总激活值的平方根 [7]。这将使该模型的激活内存消耗减少到约 8 GB。

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

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

内存碎片：到目前为止，我们已经讨论了训练期间的实际内存消耗。此外，即使有大量可用内存，也可能耗尽可用内存。这可能是由于内存碎片造成的。如果缺乏足够的连续内存来满足请求，即使总可用内存大于请求的内存，内存请求也会失败。我们在训练超大型模型时观察到严重的内存碎片，在某些极端情况下，即使仍有超过 30% 的内存可用，也会导致内存不足（OOM）问题。

4 ZeRO：洞察与概述

ZeRO 有两组优化：i) 旨在减少模型状态内存占用的 ZeRO-DP，和 ii) 旨在减少残余内存消耗的 ZeRO-R。我们概述了这些优化及其背后的洞察，这使得 ZeRO 能够在保持高效的同时减少内存占用。请注意，效率是这里的关键：没有这个约束，像将所有参数状态移动到 CPU 内存，或任意增加 MP 程度这样的琐碎解决方案都可以减少内存占用。

4.1 洞察与概述：ZeRO-DP

ZeRO 驱动的 DP 基于三个关键洞察： a) DP 比 MP 具有更好的扩展效率，因为 MP 降低了计算粒度，同时也增加了通信开销。超过某一点，较低的计算粒度会降低每个 GPU 的效率，而增加的通信开销会阻碍跨 GPU 的可扩展性，特别是在跨越节点边界时。相反，DP 具有更高的计算粒度和更低的通信量，从而允许更高的效率。 b) DP 内存效率低下，因为模型状态在所有数据并行进程中冗余存储。相反，MP 划分模型状态以获得内存效率。 c) DP 和 MP 都保留了整个训练过程中所需的所有模型状态，但并非所有内容都一直需要。例如，对应于每一层的参数仅在该层的前向传播和反向传播期间需要。

基于这些洞察，ZeRO-DP 保留了 DP 的训练效率，同时达到了 MP 的内存效率。ZeRO-DP 划分模型状态而不是复制它们（第 5 节），并使用动态通信调度，利用模型状态内在的时间特性，同时最小化通信量（第 7 节）。通过这样做，ZeRO-DP 随着 DP 程度的增加，线性地减少了模型的每设备内存占用，同时保持通信量接近默认 DP，从而保持了效率。

4.2 洞察与概述：ZeRO-R

4.2.1 减少激活内存

两个关键洞察是： a) MP 划分模型状态，但通常需要复制激活内存。例如，如果我们垂直拆分线性层的参数并在两个 GPU 上并行计算它们，每个 GPU 都需要整个激活值来计算其分区。 b) 对于 GPT-2 或更大的模型，算术强度（每次迭代的计算量与每次迭代的激活检查点量的比率）非常大（ $\ge 10K$ ），并且随着隐藏维度线性增加，使得即使在带宽较低时，也可以隐藏激活检查点的数据移动成本。

ZeRO 通过跨 GPU 划分激活检查点消除了 MP 中的内存冗余，并使用 allgather 按需重建它们。激活内存占用与 MP 程度成比例减少。对于超大型模型，ZeRO 甚至可以选择将激活分区移动到 CPU 内存，同时由于这些模型中的高算术强度，仍然实现良好的效率。

4.2.2 管理临时缓冲区

ZeRO-R 使用恒定大小的缓冲区，以避免随着模型规模增加，临时缓冲区变得过大，同时使它们足够大以保持高效。

4.2.3 管理碎片内存

内存碎片是短生命周期和长生命周期内存对象交错的结果。在前向传播期间，激活检查点是长生命周期的，但重计算的激活值是短生命周期的。类似地，在反向计算中，激活梯度是短生命周期的，而参数梯度是长生命周期的。基于这一洞察，ZeRO 通过将激活检查点和梯度移动到预分配的连续内存缓冲区来执行即时内存碎片整理。这不仅增加了内存可用性，还通过减少内存分配器查找空闲连续内存的时间来提高效率。

5 深入了解 ZeRO-DP

虽然现有的 DP 方法在每个设备上复制模型状态并引入了显著的内存开销，但 ZeRO-DP 通过跨数据并行进程划分它们（优化器状态、梯度和参数）消除了这种内存冗余。图 1 量化并可视化了有无 ZeRO-DP 的内存需求。该图显示了在累积划分 (1) 优化器状态、(2) 梯度和 (3) 参数冗余后的内存占用。我们将它们称为 ZeRO-DP 的三个优化阶段： $P_{os}$ 、 $P_g$ 和 $P_p$ ，我们在下面详细阐述。

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

对于 $N_d$ 的 DP 程度，我们将优化器状态分为 $N_d$ 个相等的分区，使得第 $i$ 个数据并行进程仅更新对应于第 $i$ 个分区的优化器状态。因此，每个数据并行进程只需要存储和更新总优化器状态的 $1/N_d$ ，然后仅更新参数的 $1/N_d$ 。我们在每个训练步骤结束时跨数据并行进程执行 all-gather，以获得跨所有数据并行进程的完全更新的参数。

内存节省：如图 1 所示，优化状态划分后的内存消耗从 $4\Psi + K\Psi$ 减少到 $4\Psi + \frac{K\Psi}{N_d}$ 。如图 1 所示的具体示例，一个 7.5B 参数的模型在使用 $P_{os}$ 和 64 路 DP ( $N_d = 64$ ) 时需要 31.4GB 的内存，而标准 DP 需要 120 GB。此外，当 $N_d$ 很大时，模型状态上的内存需求从 $4\Psi + 12\Psi = 16\Psi$ 字节减少到 $4\Psi + \frac{12\Psi}{N_d} \approx 4\Psi$ 字节，导致 4 倍的减少。

5.2 $P_g$ ：梯度划分

由于每个数据并行进程仅更新其对应的参数分区，它只需要对应参数的归约梯度。因此，随着每一层的每一个梯度在反向传播期间变得可用，我们仅在负责更新对应参数的数据并行进程上归约它们。归约后，我们不再需要梯度，它们的内存可以被释放。这减少了容纳梯度所需的内存占用，从 $2\Psi$ 字节减少到 $\frac{2\Psi}{N_d}$ 。

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

内存节省：通过消除梯度和优化器状态冗余，我们将内存占用进一步减少到 $2\Psi + \frac{14\Psi}{N_d} \approx 2\Psi$ 。如图 1 所示的示例，一个 7.5B 参数的模型在使用 $P_{os+g}$ 和 64 路 DP ( $N_d = 64$ ) 时仅需要 16.6 GB 的内存，而标准 DP 需要 120 GB。当 $N_d$ 很大时，模型状态的内存需求从 $2\Psi + 14\Psi = 16\Psi$ 字节减少到 $2\Psi + \frac{14\Psi}{N_d} \approx 2\Psi$ 字节，导致 8 倍的减少。

5.3 $P_p$ ：参数划分

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

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

5.4 对模型规模的影响

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

6 深入了解 ZeRO-R

6.1 $P_a$ ：划分激活检查点

正如 4.2 中所讨论的，MP 设计上需要复制激活值，导致跨模型并行 GPU 的激活值冗余副本。ZeRO 通过划分激活值消除了这种冗余，并且仅在计算前，一次一个激活层以复制形式实现它们。更具体地说，一旦计算了模型某一层的前向传播，输入激活值就会跨所有模型并行进程进行划分，直到在反向传播期间再次需要它。此时，ZeRO 使用 all-gather 操作来重新实现激活值的复制副本。我们将此优化称为 $P_a$ 。它与激活检查点 [7] 结合使用，仅存储划分的激活检查点而不是复制副本。此外，在超大型模型和非常有限的设备内存的情况下，这些划分的激活检查点也可以卸载到 CPU，从而以额外的通信成本将激活内存开销减少到几乎为零，我们将在第 7 节讨论。我们将此称为 $P_{a+cpu}$ 。

内存节省：通过划分激活检查点，ZeRO 将激活占用空间减少了与 MP 程度成比例的因子。考虑训练表 4 中所示的 100B 模型，批次大小为 32，序列长度为 1024，MP 程度为 16。如果我们为每个 Transformer 层检查点一个激活值，仅存储激活检查点就需要每个 GPU 大约 33 GB 的内存。但在 ZeRO 的 $P_a$ 中，它可以减少到每个 GPU 大约 2 GB。此外，这 2GB 可以卸载到 CPU，从而将激活的内存占用减少到几乎为零。

6.2 $C_B$ ：恒定大小缓冲区

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

6.3 $M_D$ ：内存碎片整理

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

$表 2：通过内存分析的最大模型规模（左）和使用 ZeRO-OS 运行时的测量模型规模（右）。使用 P_{os} 的测量模型规模与理论最大值相匹配，证明了我们的内存分析为模型规模提供了现实的上限。$

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

ZeRO 通过为激活检查点和梯度预分配连续内存块，并在它们产生时将它们复制到预分配的内存中，从而执行即时的内存碎片整理。 $M_D$ 不仅使 ZeRO 能够以更大的批次大小训练更大的模型，而且在有限内存下训练时也提高了效率。

7 ZeRO-DP 的通信分析

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

7.1 数据并行通信量

在数据并行训练期间，所有数据并行进程的梯度在反向传播结束时、计算下一步更新之前进行平均。平均操作使用 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}$ 的通信量

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

因此，总通信量是这些 all-gather 产生的通信量与梯度 reduce-scatter 产生的通信量之和。因此，总通信量为 $3\Psi$ ，与基线相比为 1.5 倍。梯度和参数划分都利用了这样一个洞察：并非所有梯度和参数状态在任何时候都是必需的，从而通过明智地通信状态来优化内存。

8 ZeRO-R 的通信分析

我们将 ZeRO-R 中划分激活检查点 ( $P_a$ ) 的通信量与基线 MP 进行比较，并表明 $P_a$ 产生的通信量增加通常不到基线 MP 的十分之一。此外，我们分析了 $P_a$ 相对于 DP 通信量的通信开销，以确定 $P_a$ 通过允许更大的批次大小和减少 DP 通信来提高效率的场景。我们利用这种分析来决定是否以及何时应用 $P_a$ 以及 $P_{a+cpu}$ 。

划分激活检查点的通信量权衡取决于模型规模、检查点策略和 MP 策略。为了分享具体的见解，我们在使用 SOTA MP 方法（Megatron-LM）实现的 Transformer 模型背景下进行分析。

在带有激活检查点的 Megatron-LM 中，每个 Transformer 块在前向传播中执行两次大小为 $batch \times seq\_length \times hidden\_dim$ 的 all-reduce 操作，在前向重计算中执行两次，在反向传播中再执行两次。每个块的总通信量为 $12 \times seq\_length \times hidden\_dim$ ，因为 all-reduce 的通信量是 $2 \times message\_size$ 。

当 ZeRO-R 划分激活检查点时，它需要在每个激活检查点上的反向传播的前向重计算之前进行额外的 all-gather 操作。通常，我们为每个 Transformer 块检查点输入激活值，每个 Transformer 块需要一次 all-gather。因此， $P_a$ 的通信开销为 $seq\_length \times hidden\_dim$ ，因为 all-gather 的通信量是 $message\_size$ 。因此， $P_a$ 的总通信开销小于模型并行原始通信量的 10%。

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

最后，如果应用 $P_{a+cpu}$ ，划分的激活检查点被卸载到 CPU，以 2 倍的额外 CPU 内存数据移动为代价，将激活内存需求减少到几乎为零。在 DP 通信量因小批次大小（即使使用 $P_a$ ）而成为主要瓶颈的极端情况下， $P_{a+cpu}$ 可以通过增加批次大小来提高效率，只要 CPU 数据传输开销小于 DP 通信量开销，这对于小批次大小通常是成立的。

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

9 迈向 1 万亿参数

当今最大的已发布模型在 100 亿参数范围内，这已经具有挑战性。达到万亿参数，即大 3 个数量级，将不可避免地发生，但道路将充满障碍、惊喜和创新。虽然我们不声称知道或解决了所有这些问题，但 ZeRO 从系统角度解决了最根本的挑战之一：在当前硬件上适配这种规模模型的能力，同时允许它以良好的系统可扩展性进行训练。

从 SOTA 的飞跃：SOTA 框架 Megatron 在 DGX-2 系统中可以以可接受的吞吐量训练的最大模型是 16-20B 参数的模型。通过跨多个 DGX 节点进行模型并行来进一步扩展，由于有限的节点间带宽，会导致效率显著下降。

ZeRO 大幅增加了可高效运行的模型规模。它使当前一代硬件能够运行显著更大的模型，而无需细粒度的模型并行来跨越节点边界。如表 1 所示，ZeRO 在启用所有优化 ( $P_{os+g+p}$ ) 的情况下，仅使用 DP 即可在 1024 个 GPU 上适配超过 1 万亿参数。或者，当与模型并行结合时（如表 2 所示），ZeRO 可以在 1024 个 GPU 上适配超过 1 万亿参数，并使用 16 路模型并行（在每个 DGX2 节点内）和 64 路跨节点数据并行。高效运行万亿参数模型不再是不可能的！

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

为了了解资源需求，我们提供了一个与 Bert-Large 的简要比较。Bert-Large 可以在 1024 个 GPU 的 DGX-2H 集群上在 67 分钟内完成训练 [26]。一个 1 万亿参数的模型可以轻松包含比 Bert-Large 模型多 3000 倍（1 万亿 / 3.3 亿）的计算量。即使我们假设相同的序列长度和训练模型所需的样本总数，假设相同的硬件和相似的计算效率，训练一个 1T 模型也需要 140 天。在实践中，随着模型规模的增加，数据样本和序列长度都可能增加，需要一年多的时间来训练。在合理的时间内训练 1T 参数模型需要一个 exa-flop 系统。但当这种计算能力可用时，我们希望 ZeRO 将提供高效运行 1T 模型的系统技术。

10 实现与评估

我们的实现重点是支持高效训练约 100B 参数的模型，这些模型比当今最大的已发布模型（例如 T5-11B [4]）大一个数量级，同时在当前硬件（例如 1K V100 GPU）上在合理的时间范围内可训练。我们实现并评估了 ZeRO 中的一个优化子集——ZeRO-DP 中的 $P_{os+g}$ 加上 ZeRO-R——使我们能够实现这一目标。我们将此实现称为 ZeRO-100B。我们的结果表明，ZeRO-100B 可以高效训练高达 170B 参数的模型，比 SOTA 大 8 倍，速度快 10 倍，并具有改进的易用性。ZeRO-100B 驱动了 Turing-NLG，这是世界上已发布的最大模型，具有新的 SOTA 精度。

10.1 实现与方法论

实现：我们在 PyTorch 中实现了 ZeRO-100B，包括 $P_{os+g}$ 和 ZeRO-R 中的全套优化。其接口与任何实现为 torch.nn.module 的模型兼容。用户只需使用此接口包装他们的模型，并像使用经典 DP 一样利用 ZeRO 驱动的 DP。用户不需要修改他们的模型。ZeRO 驱动的 DP 可以与任何形式的 MP（包括 Megatron-LM）结合使用。

硬件：我们在一个由 400 个 V100 GPU（25 个 DGX-2 节点）组成的集群上进行了实验，该集群具有 800 Gbps 的节点间通信带宽。

基线：对于没有 MP 的实验，我们使用 torch 的分布式数据并行（DDP）作为基线。对于有 MP 的实验，我们使用 Megatron-LM，因为据我们所知，它是 SOTA。我们使用了来自 NVIDIA 的开源版本 Megatron-LM，日期为 2019 年 9 月。最新的 Megatron-LM 结果报告了使用 32 个 DGX-2 节点（总共 512 个 32GB V100 GPU）扩展到 16B 参数模型的能力 [3]。

ZeRO：没有 MP 的实验使用 ZeRO-100B 中的 ZeRO 驱动的 DP 实现。有 MP 的实验将 ZeRO 驱动的 DP 与 Megatron-LM 的 MP 结合。

模型配置：本节中介绍的模型是类似 GPT-2 [2] 的 Transformer 模型。我们改变隐藏维度和层数以获得具有不同参数数量的模型。表 4 显示了我们在实验中使用的配置参数，AE 附录中有更多详细信息。

10.2 速度与模型规模

ZeRO-100B 在 400 个 GPU 上高效运行高达 170B 参数的模型，比 Megatron-LM 大 8 倍以上。图 2 显示了使用 ZeRO-100B 结合 MP 与仅使用 Megatron MP 相比，不同模型规模下的每个 GPU 吞吐量。ZeRO-100B 在 8B 到 100B 参数的模型上平均实现了 15 PetaFlops 的持续吞吐量（超过峰值的 30%）。相比之下，基线 MP 性能随着模型规模的增加而迅速下降：MP 在 GPU 之间产生高通信量，并且为了适配更大的模型而跨越单个节点会导致通信带宽从 300GB/秒/链路（NVSwitch）下降到 12.5 GB/秒/链路（Infiniband EDR），从而导致显著的性能下降。ZeRO-100B 实现了比基线高出 10 倍的加速，在大型模型上表现显著优于基线。

对于 ZeRO-100B，超过 100B 后的性能轻微下降是由于缺乏足够的内存来运行更大的批次大小。我们预计随着我们增加 GPU 数量，性能会提高，因为 ZeRO-100B 具有超线性加速，我们将在下面讨论。

10.3 超线性可扩展性

ZeRO-100B 展示了针对超大型模型规模的超线性可扩展性。图 3 显示了 60B 参数模型从 64 个 GPU 扩展到 400 个 GPU 的可扩展性结果，我们预计这一趋势将随着更多 GPU 而持续。 $P_{os+g}$ 减少了 ZeRO-100B 的每 GPU 内存消耗，并随着 DP 程度的增加，允许 ZeRO-100B 适配更大的每 GPU 批次大小，这反过来又通过增加算术强度提高了吞吐量。

10.4 大型模型训练的民主化

使用 MP 和 PP 对许多数据科学家来说具有挑战性，这是训练大型模型的众所周知的障碍。ZeRO 不需要对模型本身进行任何更改，它可以像基线 DP 一样简单地使用，同时提供显著提升的模型规模和速度。图 4 显示，ZeRO-100B 可以在 128 个 GPU 上训练高达 13B 参数的模型而无需 MP，每个 GPU 平均实现超过 40 TFlops 的吞吐量。相比之下，没有 ZeRO，仅使用 DP 的最大可训练模型具有 1.4B 参数，每个 GPU 的吞吐量低于 20 TFlops。此外，在没有 MP 通信开销的情况下，这些模型可以在不需要非常快速的节点内互连（如 NVLINK 或 NVSwitch，这是实现 MP 高效率所必需的）的低端计算节点上进行训练。

10.5 内存与性能分析

我们研究了不同优化对最大模型规模、内存消耗和性能的好处和影响。这些优化在表 3 中被称为配置 1 到 5 (C1-C5)。

最大模型规模：图 6 显示了通过为固定批次大小和 16 的 MP 启用不同 ZeRO 优化而得到的最大可训练模型。当使用 C1 与 C2 训练时，模型规模从 40B 增加到 60B，这是由于使用 $P_a$ 减少了 16 倍（MP 程度）的激活内存，而使用 C4 跃升至 140B 是由于启用了 $P_{os+g}$ ，与 C2 中的 $P_{os}$ 相比，它将模型状态的内存需求减半。使用 C5 增加到 150B 完全是由于通过将划分的激活检查点卸载到 CPU 内存，进一步减少了激活内存。

最大缓存内存：图 7 显示了 PyTorch 在 40B 和 100B 参数模型的每个训练迭代期间缓存的最大内存。缓存内存大小的减少正如从 C1 到 C2 所预期的那样。C2 和 C3 之间的内存消耗差异取决于模型状态大小与激活内存的比较，当激活内存较大时，差异可能会增加，或者当模型状态较大时，差异会减少。值得注意的是，对于 40B 模型，缓存内存并没有从 C4 减少到 C5，但对于 100B 模型却减少了。这仅仅是因为 100B 的激活内存大得多，以至于减少量是显而易见的。这使得 $P_{a+cpu}$ 成为我们在处理超大型模型时适配更大批次大小的宝贵工具。在图 8 中，170B 模型需要 $P_{a+cpu}$ 才能在不耗尽内存的情况下执行。

最大可实现性能：图 8 显示了不同优化集下的最佳可实现性能。请注意，性能提升对应于优化之间内存消耗的减少。如前所述，更低的内存消耗允许更大的批次大小，从而提高性能。唯一的警告是 60B 参数模型在 C4 和 C5 之间的性能下降。尽管内存消耗更低，但 C5 产生了到 CPU 的激活移动，这在大多数情况下会导致更差的性能，除了少数模型太大以至于没有 C5 根本无法运行，或者没有 C5 可以运行的批次大小非常小（例如图 8 中 170B 参数的模型）的情况。在训练期间， $P_{a+cpu}$ 仅在有益时才开启。

10.6 Turing-NLG，具有 17B 参数的 SOTA 语言模型

截至 2020 年 5 月 12 日，Turing-NLG 是世界上具有超过 17B 参数的最大模型。它以 10.21 的 Webtext-103 困惑度实现了语言模型的新 SOTA。Turing-NLG 是使用 ZeRO-100B 端到端训练的，图 5 显示了与之前的 SOTA（Megatron-LM 8.3B 参数模型）相比，超过 300K 次迭代的验证困惑度。ZeRO-100B 为该模型实现了 41.4 TFlops/GPU 的持续吞吐量。

11 结束语

从 HPC 和系统的角度来看，我们相信 ZeRO 代表了大型模型训练领域的革命性变革。虽然我们的实现 ZeRO-100B 实现了模型规模 8 倍的增加、吞吐量 10 倍以上的提升，在现代 GPU 集群上实现了超线性加速，并训练了世界上最大的模型，但这仅仅是冰山一角。ZeRO 整体上有潜力将模型规模再增加一个数量级，从而实现未来万亿参数模型的训练。

也许我们对 ZeRO 最感到乐观的是，它对数据科学家没有任何障碍。与 MP 和 PP 等现有方法不同，不需要模型重构，并且它像标准 DP 一样易于使用，这使得 ZeRO 成为未来大规模模型训练研究的主要候选者。通过开源和社区反馈，我们计划使 ZeRO 完全可供 DL 社区使用，以促进大规模模型训练的演进和民主化。

精粹译文