大型深度学习模型提供了显著的准确性提高，但训练数十亿到数万亿的参数具有挑战性.现有的解决方案（如数据并行和模型并行）在将这些模型适配到有限的设备存储器中并获得计算、通信和开发效率方面存在根本的局限性。本文作者开发了一种新的解决方案，零冗余优化器（Zero），以优化内存，大大提高训练速度，同时增加可以有效训练的模型大小。ZeRO消除了数据并行和模型并行训练中的内存冗余，同时保持了低通信量和高计算粒度，能够以持续的高效率按设备数量比例缩放模型大小。作者对内存需求和通信量的分析表明：使用当今的硬件，ZeRO有可能扩展到超过1万亿的参数。
作者实现并评估ZeRO：它在400个GPU上以超线性加速训练超过100B参数的大型模型，实现15个Petaflops的吞吐量。这意味着与现有技术相比，模型规模增加了8倍，可实现的性能增加了10倍。在可用性方面，ZeRO可以训练高达13B参数的大型模型，而不需要进行模型并行。
最后，研究人员利用ZeRO的系统突破创造了图灵NLG，这是当时世界上最大的语言模型（17B参数），精度创下了纪录。
数据并行不解决一块GPU放不下整个模型的问题。在当前32 GB内存的GPU上，如果参数超过1.4B，内存会耗尽。模型并行不能很有效地扩展。模型并行将模型垂直拆分，将每一层的计算和参数划分到多个设备上，需要在每一层之间进行大量通信。模型并行在单机多卡中工作得很好，但在多机多卡的情况下效率会迅速下降。作者在两个DGX-2节点上使用Megatron-LM训练了一个40B参数的模型，观察到每个V100 GPU的工作效率小于硬件峰值的5%。
如何解决这些问题并更有效地训练大型模型呢？为了回答这个问题，本文首先分析了现有系统在模型训练时的内存消耗：对于大型模型，大部分内存被模型状态占用，其中包括优化器状态（如Adam中的动量和方差）、梯度和参数。剩余内存被中间值、临时缓冲区和不可用的内存碎片消耗，统称为剩余状态。针对这两个部分所面临的不同挑战，作者提出并讨论了相应的解决方案。
并行化是大规模训练大型模型的关键策略。对于适合于设备存储器中进行训练的模型，数据并行（DP）用于将训练扩展到多个设备。在DP中，模型参数在每个设备上复制。在每个步骤中，在所有数据并行进程中均匀地划分一个小批量，这样每个进程都在不同的数据样本子集上执行前向和后向传播，并使用进程之间的平均梯度来本地更新模型。
当模型不适合设备存储器时，模型并行（MP）和流水线并行（PP）分别以垂直和水平方式在进程之间拆分模型。第二节讨论了ZeRO与DP和MP的关系。
流水线并行在不同的设备上水平分割模型，并使用微批处理来隐藏流水线气泡。由于水平拆分和微批处理，绑定权重和批处理标准化等模型功能难以实现。流行的PP实现（如G-pipe）对模型参数和总激活进行划分，但需要与流水线划分数量成比例的批处理大小来隐藏流水线气泡。大的批量大小会影响收敛速度，同时也需要大量内存来存储激活。
除了MP和PP之外，还有多个工作线旨在减少DL训练的内存开销。
（1） 减少激活内存：通过压缩、激活检查点或实时分析，多项工作都集中在减少激活的内存占用。这些努力是互补的，可以与ZeRO配合。
（2） CPU卸载：利用当今计算节点的异构性，分别通过算法设计和虚拟化内存将模型状态卸载到CPU内存。多达50%的培训时间将用于GPU-CPU传输。ZeRO的不同之处在于，它显著减少了内存消耗，而无需将模型状态存储到CPU内存中，而CPU内存的带宽因PCI-E而受到严重限制。在极少数情况下，ZeRO-R可能只卸载非常大型的激活检查点，以提高性能。
（3） 内存高效优化器：专注于通过保持模型参数和梯度的粗粒度统计来减少自适应优化方法的内存消耗，并对模型收敛保证产生潜在影响。ZeRO与这些工作是正交的，其优化不会改变模型优化方法或影响模型收敛，但有效地减少了每个设备的优化器状态和梯度的内存占用。
让我们退一步来检查当前训练系统的内存消耗。例如，1.5B参数的GPT-2模型需要3GB的内存来存储其16位精度的权重（或参数），然而，它不能使用Tensorflow或PyTorch在具有32GB内存的单个GPU上训练。人们可能会想，存储空间都去哪了。在模型训练期间，大部分内存被模型状态消耗，即由优化器状态、梯度和参数。除了这些模型状态之外，其余的内存都被激活、临时缓冲区和碎片内存消耗，作者称之为剩余状态。作者从两个方面详细了解内存消耗。
大多数设备内存在训练期间被模型状态消耗。例如，Adam，DL训练最流行的优化器之一。Adam需要存储两个优化器状态，（i）时间平均动量和（ii）计算更新的梯度方差。因此，要用Adam训练模型，必须有足够的内存来保存梯度动量和方差的副本。此外，需要有足够的内存来存储梯度和权重本身。在这三种类型的参数相关张量中，优化器状态通常消耗最多的内存，特别是在应用混合精度训练时。
混合精度训练：在当前的NVIDIA GPU上训练大型模型的最先进方法是通过混合精度（fp16/32）训练，其中参数和激活存储为fp16，从而能够在这些GPU上使用高通量张量核心单元。在混合精度训练期间，使用fp16权重和激活进行正向和反向传播。然而，为了在反向传播结束时有效地计算和更新权重，优化器保留参数的fp32副本以及所有其他优化器状态的fp32拷贝。
以Adam为例。使用Adam对参数量为Ψ的模型进行混合精度训练需要足够的内存来保存参数和梯度的fp16副本，其内存需求分别为2Ψ和2Ψ字节。此外，它需要保存优化器状态：参数、动量和方差的fp32副本，内存需求分别为4Ψ、4Ψ和4Ψ字节。用K表示优化器状态的内存乘数，即存储它们所需的额外内存为KΨ字节。混合精度Adam的K=12。总的来说，这导致2Ψ+2Ψ+KΨ=16Ψ字节的内存需求。对于具有15亿参数的GPT-2这样的模型，这导致至少24 GB的内存需求，这明显高于仅保存fp16参数所需的3 GB内存。
激活（中间变量）：在训练过程中，激活会占用大量内存。作为一个具体的例子，1.5B参数GPT-2模型以1K的序列长度和32的batch size进行训练需要60GB内存。激活检查点（或激活重新计算）是一种常见的方法，以33%的重新计算开销为代价，将激活内存减少大约总激活的平方根。这将使该模型的激活内存消耗减少到大约8GB。尽管显著减少，但对于更大的模型，即使使用激活检查点，激活内存也会增长得很大。例如，一个具有1000亿参数的类似GPT的模型需要大约60GB的内存用于批处理大小32，即使在使用激活检查点时也是如此。
临时缓冲区：用于存储中间结果的临时缓冲区消耗了大型模型的大量内存。这是因为诸如梯度all-reduce或梯度范数计算之类的操作倾向于在应用操作之前将所有梯度融合到单个缓冲区中，以提高吞吐量。
内存碎片：即使有足够的可用内存，也可能耗尽可用内存。这可能发生在内存碎片中。如果没有足够的连续内存来满足对内存的请求，即使总可用内存大于请求的内存，也会失败。作者观察到，在训练非常大的模型时，内存严重碎片，导致内存不足，在某些极端情况下，超过30%的内存仍然可用。
ZeRO有两组优化：（i）ZeRO-DP旨在减少模型状态的内存占用，（ii）ZeRO-R旨在减少剩余内存消耗。作者概述了优化及其背后的见解，这使ZeRO能够在保持效率的同时减少内存占用。请注意，效率是这里的关键：如果没有这个限制，将所有参数状态移动到CPU内存，或任意增加MP程度等微不足道的解决方案都可以减少内存占用。
三个关键见解如下：
（a）DP比MP具有更好的缩放效率，因为MP降低了计算的粒度，同时也增加了通信开销。超过某一点后，较低的计算粒度会降低每GPU的效率，而增加的通信开销会隐藏GPU的可扩展性，尤其是当跨越节点边界时。相反，DP具有更高的计算粒度和更低的通信量，从而允许更高的效率。
（b）DP的内存效率很低，因为模型状态在所有数据并行过程中都是冗余存储的。相反，MP对模型状态进行分区以获得内存效率。
（c）DP和MP都保留了整个训练过程中所需的所有模型状态，但并非总是需要所有状态。例如，仅在每层的正向传播和反向传播期间需要对应于每个层的参数。
基于这些见解，ZeRO-DP保留了DP的训练效率，同时实现了MP的存储效率。ZeRO-DP对模型状态进行划分，而不是完整复制它们，并使用动态通信调度，利用模型状态的内在时间性，同时最小化通信量。通过这样做，ZeRO-DP可以随着DP程度的增加而线性地减少模型的设备内存占用，同时保持通信量接近默认DP的通信量，从而保持效率。
两个个关键见解：
（a）MP对模型状态进行划分，但通常需要复制激活。例如，如果我们垂直分割线性层的参数，并在两个GPU之间并行计算它们，每个GPU都需要整个激活过程来计算其分区。
（b）对于GPT-2或更大的模型，算术强度（每次迭代的计算量与每次迭代的激活检查点数量之比）非常大(≥ 10K），并随着隐藏维度线性增加，即使在带宽较低的情况下，也可以隐藏激活检查点的数据移动成本。
ZeRO通过在GPU之间划分激活检查点来消除MP中的内存冗余，并使用all-collect按需重建它们。激活内存占用与MP程度成比例地减少。对于非常大的模型，ZeRO甚至可以选择将激活分区移动到CPU内存中，同时由于这些模型中的运算强度较大，因此仍然可以获得良好的效率。
（2）管理临时缓冲区：
ZeRO-R使用恒定大小的缓冲区来避免临时缓冲区随着模型大小的增加而变得非常大，同时使其足够大以保持效率。
（3）管理碎片内存：
内存碎片是短期和长期内存对象交错的结果。在前向传播期间，激活检查点是长期的，但重新计算的激活是短期的。类似地，反向计算，激活梯度是短期的，而参数梯度是长期的。基于这一见解，ZeRO通过将激活检查点和梯度移动到预先分配的连续内存缓冲区来执行内存碎片整理。这不仅提高了内存可用性，而且通过减少内存分配器查找空闲连续内存所需的时间来提高效率。
对于并行度为Nd的DP，作者将优化器状态分组为Nd个相等的分区，这样第i个数据并行过程只更新与第i个分区对应的优化器状态。因此，每个数据并行处理只需要存储和更新总优化器状态的1/Nd，然后只更新参数的1/Nd。在每个训练步骤结束时，在每一个数据并行过程上执行一次all-gather在所有数据并行过程中获得完全更新的参数。
内存节省：如图1所示，优化状态分区后的内存消耗从4Ψ+KΨ减少到4Ψ+KΨ/Nd。如图1中所示的具体示例，7.5B参数模型使用Pos和64路DP（Nd=64）需要31.4GB的内存，而使用标准DP需要120GB。此外，当Nd较大时，模型状态的内存需求从4Ψ+12Ψ=16Ψ字节减少到4Ψ+12ΨNd≈ 4Ψ字节，减少了4倍。
由于每个数据并行进程只更新其相应的参数分区，因此它只需要相应参数的减小的梯度。因此，当每个层的每个梯度在反向传播期间变得可用时，只在负责更新相应参数的数据并行过程中减少它们。在减少之后，不再需要梯度，可以释放它们的内存。
内存节省：通过消除梯度和优化器状态冗余，作者将内存占用进一步减少到2Ψ+14Ψ/Nd≈ 2Ψ. 如图1所示，使用Pos+g和64路DP（Nd=6 4），7.5 B参数模型只需要16.6 GB的内存，而使用标准DP则需要120 GB。当Nd较大时，模型状态的内存需求从2Ψ+14Ψ=16Ψ字节减少到2Ψ+1 4Ψ/Nd≈ 2Ψ字节，减少了8倍。
就像优化器状态和梯度一样，每个进程只存储与其分区对应的参数。当前向和后向传播需要分区外的参数时，通过广播从适当的数据并行处理接收这些参数。尽管乍一看，这可能会导致显著的通信开销，但作者表明，这种方法仅将基线DP系统的总通信量增加到1.5倍，同时实现与Nd成比例的内存减少。
内存节省：通过参数分区，作者将Ψ参数的模型的内存消耗从16Ψ减少到16Ψ/Nd。如图1所示，7.5 B参数模型需要1.9 GB的模型状态内存，使用Pos+p+g和64路DP（Nd=6 4），而标准DP需要120 GB。这具有深刻的含义：ZeRO支持DP以适合任意大小的模型，只要有足够数量的设备共享模型状态。作者在图2中展示了用Pos+p+g训练模型的伪代码。
Pos、Pos+g和Pos+g+p的三个阶段分别将模型状态下每个数据并行处理的内存消耗减少了4倍、8倍和Nd倍。表I分析了在ZeRO-DP优化的3个阶段下，不同DP程度的几个示例模型的模型状态内存消耗。如果没有ZeRO，无论DP程度如何，内存消耗都等于表中的第一行。注意，当Nd=64时，ZeRO可以分别使用Pos、Pos+g和Pos+g+p来训练具有高达7.5B、14B和128B参数的模型。当Nd=1024时，启用了所有优化（Pos+g+p）的ZeRO可以训练具有1万亿参数的模型。如果没有ZeRO，仅DP可以运行的最大模型的参数就不到15亿。
如IV-B中所讨论的，MP在设计上需要激活的复制，从而在模型并行GPU之间产生激活的冗余副本。ZeRO通过对激活进行划分来消除这种冗余，并且在计算中使用激活之前，一次只以一个激活层的复制形式将它们具体化。更具体地说，一旦计算了模型层的正向传播，就在所有模型并行过程中划分输入激活，直到在反向查询期间再次需要。作者将这种优化称为Pa。
ZeRO仔细选择临时数据缓冲区的大小，以平衡内存和计算效率。在训练期间，某些操作的计算效率可能高度依赖于输入大小，较大的输入实现更高的效率。例如，一个大的all-reduce操作比一个小的操作获得更高的带宽。因此，为了获得更好的效率，高性能库（如NVIDIA Apex或Megatron）在应用这些操作之前将所有参数融合到一个缓冲区中。然而，融合缓冲区的内存开销与模型大小成正比。例如，对于3B参数模型，32位融合缓冲区将需要12 GB的内存。为了解决这个问题，当模型变得太大时，作者使用性能高效的恒定大小融合缓冲区。通过这样做，缓冲区大小不取决于模型大小，通过保持足够大的缓冲区大小，仍然可以获得良好的效率。
模型训练中的内存碎片是激活检查点和梯度计算的结果。在带有激活检查点的前向传播期间，只有选定的激活被存储用于后向传播，而大多数激活被丢弃，因为它们可以在后向传播期间重新计算。这就产生了短期内存（丢弃的激活）和长期内存（检查点激活）的交错，导致内存碎片。对于使用有限内存运行的大型模型训练，内存碎片会导致两个问题：i）即使有足够的可用内存，由于缺乏连续内存，ii）由于存储器分配器花费大量时间来搜索连续的存储器以满足存储器请求，因此效率低下。
ZeRO通过为激活检查点和梯度预先分配连续内存块，并在生成时将其复制到预先分配的内存中，来实时进行内存碎片整理。MD不仅使ZeRO能够以更大的批量训练更大的模型，而且还可以在内存有限的情况下提高训练效率。
在数据并行训练期间，在计算下一步的更新之前，需要进行All-Reduce操作。最新的All-Reduce实现使用两步方法，其中第一步是Reduce-Scatter操作，用于进程上不同部分数据的Reduce。下一步是All-Gather操作，其中每个进程收集所有进程的梯度下降数据。Reduce-Scatter和All-Gather对于具有Ψ个元素的数据，将导致两次Ψ个元素的移动。因此，标准数据并行在每个训练步骤中都会引起2Ψ个数据移动。
使用梯度划分后，每个进程只存储梯度的一部分，这是更新其对应的参数分区所需的。因此，ZeRO不需要进行All-Reduce，只需要对梯度进行Scatter-Reduce操作量。在每个进程更新它负责的参数分区后，执行All-Gather从所有数据并行进程收集所有更新的参数。每个过程都需要Ψ的通信量，与基线数据并行完全相同。
使用参数划分后，每个数据并行进程只存储它更新的参数。在前向传播期间，它需要接收所有其他分区的参数。这可以通过流水线来避免内存开销。在计算与特定分区对应的模型部分的正向传播之前，负责该分区的数据并行进程可以将权重广播给所有数据并行进程。一旦完成了该分区的正向传播，就可以丢弃参数。因此，总通信量为Ψ。在前向和反向传播过程中执行All-Gather操作，参数使用完成之后丢弃并释放内存。
因此，使用Pos+g+p的通信量包括All-Gather操作的通信量和Reduce-Scatter操作的通信量即3Ψ，扩大了1.5倍。
作者将Zero-R中Pa的通信量与基线的模型并行相比较，并表明Pa引起的额外通信量一般小于基线模型并行的十分之一。此外，作者分析了与数据并行通信量相关的Pa通信开销，以便在不同场景下使用更大的批大小和更小的DP通信来提高效率。作者利用这种分析来决定何时应用Pa和Pa+cpu。
划分激活检查点的通信量取决于模型大小、检查点策略和MP策略。为了分享具体的见解，作者在分析了使用Megatron-LM的Transformer模型。在带有激活检查点的Megatron-LM中，每个Transformer块在正向传播中执行两次batch_size*seq_length*hidden_dim的All Reduce操作，两次用于正向重新计算的All Reduce操作，以及两次用于反向传播的All Reduce操作。每个块的总通信量是12*seq_length*hidden_dim，因为All Reduce的通信量是两倍的消息大小。
当ZeRO-R划分激活检查点时，在每个激活检查点上的反向传播的激活值重新计算之前，它需要一个额外的All-Gather操作。作者检查每个Transformer块的输入，并进行一次All Gather。通信开销Pa因此是seq_length*hidden_dim。因此，Pa的总通信开销小于原模型并行通信量的10%。
当MP与DP结合使用时，Pa可以以模型并行通信量增加10%为代价减少一个数量级的数据并行通信量，并在数据并行通信成为性能瓶颈时显著提高效率。注意，Pa通过模型并行度降低激活内存消耗，允许按比例增加批大小。对于大模型，MP可以大到16（DGX-2节点上的GPU数量）。数据并行训练的通信量与批处理大小成反比。因此，由于Pa导致的批大小的一个数量级的增加可能导致数据并行通信容量的一个数量级的下降。
最后，如果应用了Pa+cpu，划分激活检查点将被卸载到CPU，将激活内存需求降低到几乎为零，代价是与Pa相比，在CPU和GPU之间来回移动的数据增加了2倍。在极端情况下，DP通信容量是主要瓶颈，因为即使使用Pa的批大小较小，只要CPU数据传输开销小于DP通信容量开销，Pa+cpu可以通过增加批处理大小来提高效率。
目前已发表的最大模型在100亿个参数的范围内，这已经是训练的挑战。达到一万亿个参数，比这大3个数量级。
ZeRO极大地增加了可有效运行的模型大小。如表1所示，ZeRO在打开所有优化的情况下，仅使用DP就可以在1024个gpu上拟合超过1万亿的参数。或者，当与模型并行性结合时(如表2所示)，ZeRO可以在1024个GPU上拟合超过1万亿的参数，具有16路模型并行性(在每个DGX2节点内)和64个节点间的数据并行性。
作者的实现重点是支持具有100B参数的模型的高效训练，这比目前发布的最大模型(例如T5 11B)大一个数量级，同时在当前硬件(例如1K V100 GPU)的合理时间框架内可以训练。作者把这个实现称为ZeRO-100B。实验结果表明，ZeRO-100B可以有效地训练具有高达170B参数的模型，比SOTA大8倍，速度快10倍，并具有更好的可用性。
实现：作者在PyTorch中实现了ZeRO-100B，包括Pos+g和ZeRO-R的全套优化。它的接口与作为torch.nn.module实现的任何模型兼容。用户可以使用这个接口简单地包装他们的模型，并像使用经典DP一样利用ZeRO powered DP。用户不需要修改他们的模型。ZeRO-DP可以与任何形式的MP结合，包括Megatron-LM。
硬件：在400个V100 GPU(25个DGX-2节点)的集群上进行了实验，节点间通信带宽为800 Gbps。
基线：对于没有MP的实验，作者使用torch的分布式数据并行(DDP)作为基线。在MP实验中，作者使用Megatron-LM，使用的是来自NVIDIA 4的开源版本的Megatron-LM（2019年9月）。最新结果显示，使用32个DGX-2节点(共512个32GB V100 GPU)，可以扩展到16B参数模型。
ZeRO：没有MP的实验中，ZeRO-100B中使用ZeRO-DP实现。MP的实验中，将ZeRO-DP与Megatron的MP相结合。
模型配置：本节中介绍的模型是GPT-2类基于Transformer的模型。作者通过改变隐维数和层数来获得具有不同参数数的模型。
ZeRO-100B在400个gpu上有效地运行了高达170B参数的模型，比Megatron-LM大8倍以上。图3显示了使用带MP的ZeRO-100B与单独使用MP的不同模型大小的每个GPU的吞吐量。对于具有8B到100B参数的模型，ZeRO-100B平均实现了15个PetaFlops的持续吞吐量（超过峰值的30%）。相比之下，基线MP随着模型大小的增加性能下降迅速。ZeRO-100B实现了超过基线10倍的速度。
ZeRO-100B展示了超大模型的超线性可扩展性。图4显示了60B参数模型从64到400 GPU的可伸缩性结果，作者预计这种趋势将在更多GPU中继续。Pos+g降低了ZeRO-100B的每块GPU内存消耗，同时提高了DP程度，使ZeRO-100B能够适应每GPU更大的批处理大小，从而提高了运算强度，从而提高吞吐量
使用MP和PP对许多数据科学家来说是一个挑战，这是训练大型模型的一个众所周知的障碍。ZeRO不需要对模型本身进行任何更改，它可以像基线DP一样简单地使用，同时显著提高模型大小和速度。图5显示，ZeRO-100B可以在128个GPU上训练具有多达13B参数的模型，而无需MP，平均每个GPU的吞吐量超过40 TFlops。相比之下，在没有ZeRO的情况下，仅使用DP的最大可训练模型具有1.4B参数，每个GPU的吞吐量小于20 TFlops。此外，在没有来自MP的通信开销的情况下，这些模型可以用低端计算节点进行训练，而无需非常快速的内部节点互连（如NVLINK或NVSwitch）。
从HPC和系统的角度来看，作者认为ZeRO代表了大型模型训练领域的革命性变革。虽然作者在所做的实验中，用ZeRO-100B实现了模型尺寸的8倍增长，吞吐量提高了10倍以上，在现代GPU集群上实现了超线性加速，并训练了世界上最大的模型，但这仍然只是冰山一角。ZeRO整体上有可能将模型大小再增加一个数量级，从而实现未来万亿参数模型的训练。
此外，作者对ZeRO感到最乐观的是，与MP和PP等方法不同，它无需进行模型重构，而且它与标准DP一样易于使用，这使得ZeRO可能成为未来大型模型训练研究的首选。通过开源和社区反馈，作者计划让DL社区充分了解ZeRO，以促进大规模模型训练的发展。
联系客服
