《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》--中文翻译

转载请注明出处：https://blog.csdn.net/nocml/article/details/134949832

摘要

  大型深度学习模型可显著提高准确性，但训练数十亿至数万亿参数具有挑战性。 现有的解决方案（例如数据和模型并行性）在将这些模型适应有限的设备内存，同时获得计算、通信和开发效率方面表现出根本的局限性。 我们开发了一种新颖的解决方案，零冗余优化器 (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) 领域，Transformers 为大型模型铺平了道路，例如 Bert-large (0.3B) 、GPT-2 (1.5B) 、Megatron-LM (8.3B) ，T5 (11B) 。 为了使模型大小能够从数十亿个参数持续增长到数万亿个参数，我们在训练它们时遇到了挑战 - 它们显然不适合单个设备（例如 GPU 或 TPU）的内存，并且简单地添加更多设备将 无助于扩大训练规模。
  基本数据并行性 (DP) 不会减少每个设备的内存占用，并且在具有 32GB 内存的当代 GPU 上，对于参数超过 1.4B 的模型，会耗尽内存。 其他现有的解决方案，例如管道并行（PP）、模型并行（MP）、CPU卸载等，在功能、可用性以及内存和计算/通信效率之间进行权衡，但所有这些对于训练的速度和规模都至关重要。
  在训练大型模型的不同现有解决方案中，MP 可能是最有前景的。 当前文献中最大的模型，11B T5 模型和 Megatron-LM 8.3B ，均由模型并行提供支持，分别在 Mesh-Tensorflow 和 Megatron-LM 中实现 。 然而，MP 无法扩展到超出这些模型尺寸的范围。 MP 垂直分割模型，将每层中的计算和参数划分到多个设备上，需要每层之间进行大量通信。 因此，它们在 GPU 间通信带宽较高的单个节点内运行良好，但超出单个节点时效率会迅速下降 。我们使用 Megatron-LM 在两个 DGX-2 节点上测试了 40B 参数模型，并观察到每个 V100 GPU 的性能约为 5Tflops（低于硬件峰值的 5%）。
  那么，我们如何克服现有解决方案的局限性，更有效地训练大型模型呢？ 为了回答这个问题，我们首先分析现有系统在模型训练上的内存消耗，并将其分为两部分：1）对于大型模型，大部分显存被模型状态占用，其中包括优化器状态（ 例如 Adam 中的动量和方差、梯度和参数。 2）剩余显存被激活、临时缓冲区和不可用的碎片显存消耗，我们统称为剩余状态。 我们开发了 ZeRO（零冗余优化器），以优化两者的显存效率，同时获得高计算和通信效率。 由于这两个部分面临不同的挑战，我们相应地开发和讨论他们的解决方案。
  优化模型状态显存 模型状态在训练过程中通常会消耗最大量的显存，但现有的方法（例如 DP 和 MP）并不能提供令人满意的解决方案。 DP 具有良好的计算/通信效率，但显存效率较差，而 MP 的计算/通信效率较差。 更具体地说，DP 在所有数据并行过程中复制整个模型状态，导致冗余显存消耗； 虽然 MP 对这些状态进行分区以获得高显存效率，但通常会导致过于细粒度的计算和昂贵的通信，从而降低扩展效率。 此外，所有这些方法都静态地维护整个训练过程中所需的所有模型状态，即使在训练期间并非始终需要所有模型状态。 基于这些观察，我们开发了ZeRO-DP，ZeRO驱动的数据并行，它实现了DP的计算/通信效率，同时实现了MP的内存效率。 ZeRO-DP 通过划分模型状态而不是复制模型状态来消除数据并行过程中的内存状态冗余，并且通过在训练期间使用动态通信调度保留 DP 的计算粒度和通信量来保留计算/通信效率。
  ZeRO-DP 具有三个主要优化阶段（如图 1 所示），分别对应于优化器状态、梯度和参数的划分。 累积启用时：
  1) 优化器状态分区（$P_{os}$）：内存减少 4 倍，通信量与 DP 相同；
  2) 添加梯度分区（$P_{os+g}$）：内存减少8倍，通信量与DP相同；
  3) 添加参数分区（$P_{os+g+p}$）：内存减少与DP并行数量$N_d$呈线性关系。
例如，划分为 64 个 GPU ($N_d$ = 64) 将导致内存减少 64 倍。 通信量增加 50%。

  ZeRO-DP 消除了显存冗余，并使集群的全部聚合显存容量可用。 启用所有三个阶段后，ZeRO 可以在 1024 个 NVIDIA GPU 上训练万亿参数模型。 具有 Adam 等 16 位精度优化器的万亿参数模型需要大约 16 TB 的显存来保存优化器状态、梯度和参数。 16TB 除以 1024 就是 16GB，这完全在 GPU 的合理范围内（例如，设备上有 32GB 显存）。

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

  优化剩余状态显存 ZeRO-DP 提高了模型状态的显存效率后，激活、临时缓冲区和不可用显存片段消耗的其余显存可能会成为次要显存瓶颈。 我们开发ZeRO-R来分别优化这三个因素消耗的剩余显存。

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

  ZeRO-DP 和 ZeRO-R 组合在一起形成了一个强大的深度学习显存优化系统，我们统称为 ZeRO。

  ZeRO 和 MP： 既然 ZeRO 消除了 DP 中内存效率低下的问题，我们很自然会问：我们还需要 MP 吗？什么时候需要？ ZeRO 如何与 MP 配合使用？ 有了ZeRO，MP对于单独拟合大型模型的目的来说就不再那么有吸引力了。 ZeRO-DP 在减少每个设备显存占用方面至少与 MP 一样有效，当 MP 无法均匀划分模型时ZeRO-DP会更有效。 它还具有相当或更好的扩展效率。 此外，数据并行性非常易于使用，可以广泛适用于不同的工作负载，如今的 MP 方法通常需要模型开发人员进行一些工作来修改其模型，需要系统开发人员来制定分布式算子，而像 Megatron-LM 这样的现有工作仅支持有限的算子和模型集。
  话虽这么说，但在某些情况下我们仍然希望利用 MP：
与 ZeRO-R 一起使用时，MP 可以减少超大模型的激活显存占用。

• ii) 对于激活显存不成问题的较小模型，当单独使用 DP 的聚合batch size太大而无法实现良好收敛时，MP 也可以带来好处。 在这种情况下，可以将 ZeRO 与 MP 结合起来，以适应具有可接受的聚合batch size的模型。

  ZeRO可以与MP相结合，在DP并行数为$N_d$、MP并行数为$N_m$的每个设备上，理论上显存最多减少$N_d\times N_m$次。这可以使我们在1024个GPU上，使用16路MP（在每个DGX2节点内）和64路跨节点DP，拟合万亿参数模型，并使用适度的batch size高效运行它！

  实施和评估 ZeRO 中的全套优化可以让我们在当今的高端硬件集群上运行具有万亿参数的模型（例如，使用 1K V100 GPU），然而，硬件计算能力仍然太有限，训练时间可能会很长（>1年），因此，我们对该实现的重点是有效地支持参数为SOTA 10倍（～100B参数）的模型，同时保证仍在当前硬件的计算能力范围内。我们在 ZeRO 中实现并评估了一个名为 ZeRO-100B 的优化子集（ZeRO-DP 的 Pos+g 加 ZeRO-R），这使我们能够实现这一目标。 结果显示：
  模型大小与MP相结合，ZeRO-100B可以有效地运行170B参数模型，而现有的系统（如单独使用Megatron）无法有效地扩展到40B参数以上，如图2所示。与SOTA相比，参数规模扩大了8倍多。

图2:ZeRO训练吞吐量和不同型号SOTA基线的加速率。对于ZeRO，MP始终适合一个节点，而对于基线，大于40B的模型需要跨节点的MP。

  速度提高了内存效率，提供了更高的吞吐量和更快的训练。如图2所示，ZeRO在400 个 Nvidia V100 GPU的集群上运行100B参数模型，每个GPU超过38 TFlops，聚合性能超过15 Petaflops。在模型规模相同的情况下，这比业界最好的结果快10倍。

  可扩展性我们在64-400个GPU的范围内观察到超线性加速，当我们将GPU数量增加一倍时，性能会增加一倍多。这是ZeRO DP的一个特性，它减少了模型状态的显存占用，因为我们增加了DP的并行数量，使每个GPU能够适应更大的batch size，从而获得更好的性能。随着我们将GPU数量增加到400个以上，我们预计这种行为将继续下去。

  大模型训练的民主化ZeRO-100B使数据科学家能够训练具有高达13B参数的模型，而不需要任何需要模型重构的MP或PP，其中13B的参数比文献中最大的模型（T5具有11B参数）更多。因此，数据科学家可以自由地对大型模型进行实验，而不用担心并行性。相比之下，现有系统（例如PyTorch Distributed Data Parallel）的显存不足，参数模型为1.4B。

  新的SOTA模型$ZeRO$为最大的语言模型 17B参数的 Turing-NLG , 提供了破纪录的准确性。

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

2. 相关工作

2.1 数据、模型和管道并行

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

  当模型不适合设备显存时，模型并行性（MP）和流水线并行性（PP）分别以垂直和水平方式在进程之间划分模型。第1节讨论了ZeRO与DP和MP的关系。我们现在讨论PP以及它与减少显存消耗的关系。

  PP 将模型水平地跨层分割，每个分区在不同的设备上运行，并使用微批处理来隐藏管道气泡 。 由于水平分割和微批处理，诸如捆绑权重和批量归一化之类的模型功能很难实现。 流行的 PP 实现（例如 G-pipe ）对模型参数和总激活进行分区，但需要与管道分区数量成比例的batch size来隐藏管道气泡。 大的batch size会影响收敛速度，同时还需要大量显存来存储激活。 PipeDream 中 PP 的不同实现保留了历史参数的多个副本以隐藏管道气泡，而不会显着增加batch size 大小 ，从而降低显存占用率。 此外，该实现并不等同于标准的深度学习训练，并且会对训练收敛产生影响。 相比之下，ZeRO 获得与 PP 相同或更好的内存效率的同时，不会产生 与PP 相同的功能、性能和收敛相关的限制。

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

  除了 MP 和 PP 之外，还有多种工作旨在减少 DL 训练的显存开销。

2.2.1 减少激活显存

  多项努力集中在通过压缩、激活检查点或实时分析 [15] 来减少激活的显存占用。 这些努力是互补的，并且可以与 ZeRO 一起工作。 事实上，ZeRO-R 中的激活显存减少与激活检查点并行工作。

2.2.2 CPU卸载

  利用当今计算节点的异构性质，分别通过算法设计或虚拟化内存将模型状态卸载到 CPU 内存。 高达 50% 的训练时间会花费在 GPU-CPU-GPU 传输上 。 ZeRO 的不同之处在于，它显着降低了显存消耗，而无需将模型状态存储到 CPU 内存中，而 CPU 内存的带宽由于 PCI-E 而受到严重限制。 在极少数情况下，ZeRO-R 可能只卸载非常大模型的激活检查点以提高性能（详细信息请参阅第 6.1 节）。

2.2.3 显存效率优化器

  专注于通过维护模型参数和梯度的粗粒度统计来减少自适应优化方法的显存消耗，这对模型收敛具有潜在影响。 ZeRO 与这些努力正交，其优化不会改变模型优化方法或影响模型收敛，而是有效减少每个设备的优化器状态和梯度的显存占用。

2.3 Training Optimizers

   大模型如果要进行有效的训练以达到SOTA性能及准确率，那么自适应优化方法对其尤其重要 。 与 SGD 相比，自适应优化方法维护每个模型参数和梯度的细粒度一阶和二阶统计数据，代价是显著的显存占用。 ZeRO 可以将这些优化器的显存占用量减少几个数量级，使这些复杂的优化方法 适用于在具有适度设备显存的硬件上训练大型模型。 它还使得开发和使用更复杂、更消耗显存的优化器成为可能，这些优化器可能具有更好的收敛性。

3 显存去哪里了？

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

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

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

   混合精度训练 在当前一代 NVIDIA GPU 上训练大型模型的最先进方法是通过混合精度 (fp16/32) 训练，其中参数和激活信息存储为 fp16，从而能够使用 这些 GPU 上的高吞吐量的 tensor core。 在混合精度训练期间，前向和后向传播都是使用 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 亿个参数的模型，这导生产少需要 24GB 的显存需求，这明显高于单独保存 fp16 参数所需的区区 3GB 显存。

3.2 剩余显存消耗

   训练期间激活可能会占用大量显存。 作为一个具体例子，以序列长度为 1K、batch size为 32 训练的 1.5B 参数的GPT-2模型需要约 60GB 显存。 激活检查点（或激活重新计算）是一种常见的方法，可以将激活内存减少大约总激活的平方根，但代价是 33% 的重新计算开销。 这会将该模型的激活显存消耗减少至约 8 GB。

   尽管显著减少，但对于较大的模型，即使使用激活检查点，激活显存也会变得相当大。 例如，具有 1000 亿个参数的类 GPT 模型，batch size 为32时， 需要大约 60GB的显存。即使使用激活检查点也是如此。

   临时缓冲区 对于大型模型来说，用于存储中间结果的临时缓冲区会消耗大量显存。 梯度全归约或梯度范数计算等操作倾向于在应用操作之前将所有梯度融合到单个扁平缓冲区中，以提高吞吐量。 例如，跨设备的 all-reduce 带宽随着消息大小的增大而提高。 虽然梯度本身通常存储为 fp16 张量，但融合缓冲区可以是 fp32 张量，具体取决于操作。 当模型很大时，这些临时缓冲区大小就很重要。 例如，对于具有 1.5B 参数的模型，扁平化的 fp32 缓冲区将需要 6GB 内存。

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

4 ZeRO：见解和概述

   ZeRO 有两组优化：i) ZeRO-DP 旨在减少模型状态的显存占用，ii) ZeRO-R 旨在减少剩余显存消耗。 我们概述了优化和背后的见解，这使得 ZeRO 能够在保持高效的同时减少显存占用。 请注意，效率是这里的关键：如果没有这个约束，将所有参数状态移至 CPU 内存或任意增加 MP 并行度等简单解决方案都可以减少显存占用。

4.1 见解和概述：ZeRO-DP

ZeRO 支持的 DP 基于三个关键见解：

• DP 比 MP 具有更好的扩展效率，因为 MP 降低了计算粒度，同时也增加了通信开销。 超过某一点，较低的计算粒度会降低每个 GPU 的效率，而增加的通信开销会隐藏跨 GPU 的可扩展性，尤其是在跨越节点边界时。 相反，DP 具有更高的计算粒度和更低的通信量，从而具有更高的效率。
• DP 的显存效率较低，因为模型状态在所有数据并行进程中冗余存储。 相反，MP 对模型状态进行分区以获得显存效率。
• DP 和 MP 都保留整个训练过程中所需的所有模型状态，但并非始终需要所有状态。 例如，每层对应的参数只有在该层的前向传播和后向传播时才需要。

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

4.2 见解和概述： ZeRO-R

4.2.1 减少激活显存

两个关键点：

• MP 对模型状态进行分区，但通常需要复制激活信息显存。
  例如，如果我们垂直分割线性层的参数并在两个 GPU 上并行计算它们，则每个 GPU 需要整个激活信息来计算其分区
• 对于GPT-2或更大的模型，算术强度（每次迭代的计算量与每次迭代的激活检查点数量的比率）非常大（≥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}$: Optimizer状态分区

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

  内存节省： 如图1所示，优化状态分区后的内存消耗从$4\Psi +K\Psi$减少到$4\Psi +\frac{K\Psi }{N_d}$。如图1所示的具体示例，7.5B参数模型需要31.4GB的显存，使用64路DP（$N_d=64$）的$P_{os}$，而使用标准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$：梯度分区

  由于每个数据并行过程只更新其相应的参数分区，因此它只需要相应参数的reduced gradients。因此，当每一层的每个梯度在反向传播期间变得可用时，我们只在负责更新相应参数的数据并行过程中减少它们。在减少之后，我们不再需要梯度，并且可以释放它们的显存。这减少了将梯度从$2\Psi$字节保持到$\frac{2\Psi }{N_d}$所需的显存占用。

  实际上，这是一个Reduce-Scatter操作，其中对应于不同参数的梯度被分发到不同的处理过程。 为了在实践中提高效率，我们使用分桶策略，将与特定分区相对应的所有梯度分桶，并立即对整个桶执行分解（reduction）。 这在本质上类似于 NVIDIA 的 AMP 优化器如何将 all-reduce 梯度计算分桶 以整合通信和计算。 在我们的例子中，我们在分区边界执行reduce而不是all-reduce，以减少显存占用并整合计算和通信。

  内存节省： 通过消除梯度和优化器状态冗余，我们将内存占用进一步减少到 $2\Psi +\frac{14\Psi }{N_d}\approx 2\Psi$。 如图 1 所示，使用 64 路 DP（$N_d=64$）、$P_{os+g}$ 的 7.5 B 参数模型仅需要 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 中的示例，使用 64 路 DP ($N_d=64$) 的 $P_{os+p+g}$ 时，7.5 B 参数模型需要 1.9 GB 的模型状态显存，而使用标准 DP 则需要 120 GB。 这具有深远的意义：只要有足够数量的设备来共享模型状态，ZeRO 就可以使 DP 适应任意大小的模型。

5.4 对模型尺寸的影响

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

表 1：ZeRO-DP 中不同优化方式在每个设备上的显存消耗与 DP 并行度的关系。 粗体文本是模型可以装入 32GB V100 GPU 集群的组合。

6 深入了解 ZeRO-R

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

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

  节省内存 通过分区激活检查点，ZeRO将激活信息占用的空间 按照 MP并行度 成比例的进行了减小。考虑表4中所示的训练参数规模为100B模型，其中批量大小为32，序列长度为1024，MP并行数为16。如果我们为每个tansformer层检查单个激活，则每个GPU仅需要大约33GB的显存来存储激活检查点。但是使用ZeRO中的$P_a$，它可以减少到每个GPU大约2GB。此外，这2GB可以卸载到CPU，将激活的显存占用减少到几乎为零。

6.2 CB：恒定大小缓冲器

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

6.3 MD：显存碎片整理

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

表 2：通过显存分析得出的最大模型大小（左）以及使用 ZeRO-OS 运行时测得的模型大小（右）。 使用 Pos 测量的模型大小与理论最大值相匹配，这表明我们的显存分析分析出了模型大小的实际上限。

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

  ZeRO 通过为激活检查点和梯度预先分配连续的显存块，并在生成时将它们复制到预先分配的显存中，即时进行显存碎片整理。 MD不仅使ZeRO能够训练具有更大批量大小的大模型，而且还提高了在有限显存下的训练效率。

7 ZeRO-DP的通信分析

  由于 ZeRO 通过消除显存冗余来增加模型大小，因此很自然地会问我们是否用通信量来换取显存效率。 换句话说，与基线 DP 方法相比，ZeRO 支持的 DP 方法的通信量是多少？ 答案分为两部分：i) ZeRO-DP 使用 Pos 和 Pg 不会产生额外的通信，同时最多可减少 8 倍的内存，ii) 除了 Pos 之外，使用 Pp 时，ZeRO-DP 最多会产生 1.5 倍的通信 和Pg，同时进一步减少Nd倍的内存占用。 我们将在本节中进行分析。 首先，我们首先简要概述标准 DP 的通信量。

7.1 数据并行通

  在数据并行训练期间，在计算下一步的更新之前，在反向传播结束时对所有数据并行过程的梯度进行平均。 使用 all-reduce 通信集合来执行平均。 对于大模型大小，all-reduce 通信完全受通信带宽限制，因此，我们将分析限制为发送至每个数据并行进程和来自每个数据并行进程的总通信量。

  在最先进的 all-reduce 实现中，使用了两步方法，其中第一步是reduce-scatter 操作，它reduce了不同进程上不同的数据部分。 下一步是all-gather操作，其中每个进程收集所有进程的reduce数据。 这两个步骤的结果是all-reduce。 reduce-scatter和all-gather都是使用流水线方法实现的，这会导致数据中的所有元素都移动 $\Psi$次（对于具有 $\Psi$ 元素的数据）。 因此，标准 DP 在每个训练步骤期间都会产生 $2\Psi$ 次数据移动。

7.2 ZeRO DP通信量

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

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

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

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

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操作，两次all-reduce用于前向传播的重新计算，并在后向传播中执行另外两次all-reduce操作。 每个块的总通信量是 $12\times seq\_length\times hidden\_dim$，因为 all-reduce 的通信量是 $2\times message\_size$。

  当ZeRO-R对激活检查点进行分区时，在每个激活检查点上的反向传播的正向传播重新计算之前需要额外的all-gather操作。 一般来说，我们检查每个变压器块的输入激活，要求每个变压器块进行一次全收集。 因此，通信开销 $P_a$ 是 $seq\_length\ast hidden\_dim$，因为all-gather的通信量是$messag{e}_{s}ize$。 因此，$P_a$的总通信开销不到模型并行原始通信量的10%。

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

  最后，如果应用 $P_{a+cpu}$，则分区激活检查点将卸载到 CPU，将激活内存需求减少到几乎为零，但与 $P_a$ 相比，CPU 内存之间的数据移动量增加了 2 倍。在极端情况下，DP 通信量 即使使用 $P_a$，由于$batch\_size$较小也是主要瓶颈，只要 CPU 数据传输开销小于 DP 通信量开销，$P_{a+cpu}$ 就可以通过增大$batch\_size$来提高效率，对于较小的$batch\_size$通常如此 。

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

9 迈向 1 万亿个参数

  目前已发布的最大模型包含 100 亿个参数，这些参数的训练已经具有挑战性。 达到万亿个参数（大三个数量级）是不可避免的，但道路上将充满障碍、惊喜和创新。 虽然我们并不声称了解或解决所有这些问题，但 ZeRO 从系统角度解决了最基本的挑战之一：能够在当前硬件上拟合这种规模的模型，同时允许其以良好的系统可扩展性进行训练。

  与最先进技术的飞跃 最先进的Megatron框架，训练的最大模型，是在DGX-2 系统中的 16 - 20B 参数模型 。 由于节点间带宽有限，通过跨多个 DGX 节点进行模型并行进一步扩展会导致效率显着下降。

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

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

10 实施和评估

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

  实现 我们在PyTorch中实现了ZeRO-100B，包括Pos+g和ZeRO-R中的全套优化。它的接口与任何实现为torch.nn.module的模型都兼容。用户可以使用这个接口简单地包装他们的模型，并像使用经典DP一样利用ZeRO供电的DP。用户无需修改其模型。ZeRO驱动的DP可以与任何形式的MP结合，包括威震天LM。

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

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

  ZeRO 无MP的ZeRO实验，在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显示了使用带有MP的ZeRO-100B与单独使用Megatron MP时不同型号的每个GPU的吞吐量。对于具有8B到100B参数的模型，ZeRO-100B平均实现了15 PetaFlops的持续吞吐量（超过峰值的30%）。相比之下，基准MP性能随着型号大小的增加而迅速下降：MP导致GPU之间的通信量很高，超出单个节点以适应更大的型号会导致通信带宽从每链路300GB/秒（NVSwitch）下降到每链路12.5GB/秒，导致性能显著下降。ZeRO-100B实现了比基线高出10倍的加速，显著优于大型机型。

10.3 超线性可扩展性

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

10.4 民主化大模型训练

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

表3:ZeRO配置

10.5 显存和性能分析

  消耗和性能。这些优化在表中称为配置1至5（C1-C5）。

  最大的模型图6显示了通过对固定batch size 和 MP并行度为16的不同ZeRO优化实现的最大可训练模型。与C2相比，当使用C1训练时，由于使用$P_a$使激活显存减少了16x（MP度），模型大小从40B增加到60B，而使用C4跳到140B是由于启用$P_{os+g}$，与C2中的$P_{os}$相比，$P_{os+g}$使模型状态的显存需求减半。使用C5增加到150B仅仅是由于将分区的激活检查点卸载到CPU内存而进一步减少了激活显存。

  最大缓存内存 图7显示了PyTorch在40B和100B参数模型的每次训练迭代期间缓存的最大显存。从C1到C2，高速缓存存储器大小的减小是意料之中的。与激活显存相比，C2和C3之间的显存消耗的差异取决于模型状态的大小，并且当激活显存较大时增加，或者当模型状态较大时减少。值得注意的是，对于40B，高速缓存的显存不会从C4减少到C5，但对于100B，它会减少。这仅仅是因为100B的激活显存要大得多才能明显地减少。这使得当我们使用非常大的模型时，$P_{a+cpu}$成为一个适合更大批量的有价值的工具。在图8中，170B模型需要$P_{a+cpu}$才能在不耗尽显存的情况下执行。

表4：图2、3、4中不同型号尺寸、层数和隐藏层尺寸（HD）的配置。

  最大可实现性能 图8显示了不同优化集的最佳可实现性能。请注意，性能的提高对应于优化之间显存消耗的减少。如前所述，较低的显存消耗允许较大的batch size，从而提高性能。唯一需要注意的是60B参数模型在C4和C5之间的性能下降。尽管显存消耗较低，但C5会导致CPU之间的激活移动，这在大多数情况下会导致较差的性能，但少数情况除外，因为模型太大，没有C5，模型根本无法运行，或者没有C5，可以运行的batch size非常小（如图8中具有170B参数的模型）。在训练过程中，只有在有利的情况下，$P_{a+cpu}$才会使用。

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

  截至2020年5月12日，Turing-NLG是世界上最大的模型，参数超过17B。它以10.21 的 Webtext-103 困惑度实现了新的语言模型SOTA。使用ZeRO-100B对Turing-NLG进行端到端训练，图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成为未来大型模型训练研究的主要候选者。通过开源和社区反馈，我们计划让DL社区完全可以访问ZeRO，以促进大规模模型训练的发展和民主化。

致谢