ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Introduction

并行性是深度学习中训练大模型的必要条件，目前的并行方法主要分为三个部分：Data parallelism（DP），Model parallelism（MP），Pipeline Parallelism（PP）。

• DP：模型部署在每个GPU上，需要训练的数据分割为多个子集分别在多个GPU上训练，得到的梯度经过收集归约，最终更新模型。
• MP：将模型的每一层的参数分割到多个GPU上，每一层的参数都需要大量的通信资源才能完成计算和训练。当在多卡单节点时，通信带宽足够时具有优势，当多节点时候无法充分利用计算资源。
• PP。 将模型按照Layer切分为多个部分，分布在多个GPU上，具有和MP相同的特点，同时能利用不同设备的不同的性能,但是因为其流水线模型，很容易出现气泡。

本文提出的ZeRO（Zero Redundancy Optimizer）能够有效利用设备性能，提高训练速度，并大量节省内存。它具有三个不同的等级。

• Stage 1：\(P_{os}\),Optimzer states Partition
• Stage2 : \(P_{os+g}\), + gradients
• Stage3 : \(P_{os+g+p}\) + parameters

除了上述的使用并行性来降低显存需求的技术，还有一些非并行性的技术

• Reducing Activation Memory。减少在forward和backward中activation的占用，activation是指为了计算出梯度，需要在forward后保存一些中间结果来供backward进行计算，这部分内存可以使用activation checkoint技术，以重新计算部分结果为代价显著降低。这些可以和ZeRO-R一起使用。
• CPU offload。将一些计算移动到CPU上。
• Memory Efficient Optimizer。deepspeed重写了部分optimizer，提供可显存友好，并且收敛更快的optimizer，实测非常好用。

Where Did All the Memory Go?

Model States:Optimizer States,Gradients and Parameters

在上述的三部分，往往optimizer占用的显存是最多的。
以Adam为具体例子，使用Adam对具有Ψ个参数的模型进行混合精度训练需要足够的内存来保存参数和梯度的f p16副本，分别需要2Ψ和2Ψ字节的内存。此外，还需要额外的内存来保存优化器状态：参数、动量和方差的fp32副本(为了更新参数)，分别需要4Ψ、4Ψ和4Ψ字节的内存。让我们用K表示优化器状态的内存倍增系数，即额外需要存储它们的内存为KΨ字节。混合精度Adam的K值为12。总体而言，这将导致16Ψ字节的内存需求。对于像GPT-2这样具有15亿个参数的模型，这将至少需要24GB的内存，远远高于仅存储fp16参数所需的3GB微不足道的内存。

Residual Memory Consumption

• activation :
  在训练过程中，数据通过激活函数得到的中间结果可能会占用大量的内存。以具体的例子来说，使用序列长度为1K和批处理大小为32进行训练的15亿参数的GPT-2模型需要约60GB的内存。激活检查点（或激活重新计算）是一种常见的方法，可以通过牺牲33%的重新计算开销来减少激活内存，大约降低总激活数的平方根[7]。这将将该模型的激活内存消耗减少到约8GB。

一个基于Transformer的模型的激活内存与Transformer层数、隐藏维度、序列长度和批处理大小成正比。对于类似GPT-2的架构，总激活数大约为12 × 隐藏维度 × 批处理大小 × 序列长度 × Transformer层数。

• Temporary buffers
  用于存储中间结果的临时缓冲区会消耗大型模型的相当数量的内存。诸如梯度全约简或梯度范数计算等操作通常会将所有梯度融合到一个扁平化的缓冲区中，然后再应用该操作，以提高吞吐量。例如，跨设备的全约简带宽会随着较大的消息大小而提高。虽然梯度本身通常以fp16张量的形式存储，但根据操作的不同，融合的缓冲区可以是fp32张量。当模型的大小较大时，这些临时缓冲区的大小就变得重要。例如，对于包含15亿参数的模型，一个扁平化的fp32缓冲区将需要6GB的内存。

• Memory Fragmentation:
  内存碎片化：到目前为止，我们讨论了训练过程中实际的内存消耗。此外，即使有足够的可用内存，也有可能出现可用内存不足的情况。这可能是由于内存碎片化引起的。如果没有足够的连续内存来满足内存请求，即使总可用内存大于所需内存，请求内存也会失败。我们观察到在训练非常大型模型时存在显著的内存碎片化问题，导致在某些极端情况下仍有超过30%的可用内存时出现内存不足的问题。

ZeRO: Insights and Overview

ZeRO powered DP基于以下三个关键见解：

• 相较于MP，DP具有更好的扩展效率，因为MP会降低计算的粒度，同时增加通信开销。在一定程度上，较低的计算粒度会降低每个GPU的效率，而增加的通信开销则阻碍了跨GPU的可扩展性，特别是在跨节点边界时。相反，DP具有更高的计算粒度和更低的通信量，从而实现更高的效率。
• DP在内存利用效率上存在问题，因为模型状态在所有数据并行进程中冗余存储。相反，MP通过对模型状态进行分区来提高内存利用效率。
• DP和MP都保留了整个训练过程中所需的所有模型状态，但并非所有状态始终都是必需的。例如，每个层的参数只在前向传播和反向传播时需要。
  基于这些见解，ZeRO-DP既保留了DP的训练效率，又实现了MP的内存利用效率。ZeRO-DP对模型状态进行分区而不是复制它们（第5节），并使用动态通信调度，利用模型状态的时间特性最小化通信量（第7节）。通过这样做，ZeRO-DP使模型的每个设备的内存占用与增加的DP度数成线性关系，同时保持通信量接近默认DP的水平，保持了效率。

ZeRO 主要分为两个部分：ZeRO-R：清楚模型训练中的冗余部分：包括activation、buffer和内存碎片等。ZeRO-DP：对模型状态进行分割，分别对应上述内存使用。

ZeRO-R

Reducing Activation Memory

两个关键见解如下：

• MP对模型状态进行分区，但通常需要复制激活内存。例如，如果我们将线性层的参数垂直拆分并在两个GPU上并行计算，则每个GPU都需要整个激活内存来计算其分区。
• 对于像GPT-2或更大的模型，每次迭代中的算术强度（每次迭代的计算量与激活检查点数量的比率）非常大（≥ 10K），并且随着隐藏维度的线性增加，这使得即使带宽较低，也能够隐藏激活检查点的数据移动成本。

ZeRO通过在GPU之间分区激活检查点，并使用allgather在需要时重新构造它们，从而消除了MP中的内存冗余。激活内存占用随MP度数成比例减少。对于非常大的模型，ZeRO甚至可以选择将激活分区移动到CPU内存中，而仍然由于这些模型中的高算术强度而实现良好的效率。

Managing Temporary buffers

ZeRO-R使用恒定大小的缓冲区，以避免随着模型大小的增加而导致临时缓冲区膨胀，同时使它们足够大以保持高效

Managing fragmented Memory

内存碎片化是短期内存对象和长期内存对象之间交错的结果。在前向传播过程中，激活函的是长期存在的，但重新计算的激活函数是短期存在的。类似地，在反向传播过程中，激活函数的梯度是短期存在的，而参数梯度是长期存在的。基于这个见解，ZeRO通过将激活检查点和梯度移动到预先分配的连续内存缓冲区中，实时进行内存整理。这不仅增加了内存可用性，还通过减少内存分配器寻找可用连续内存的时间来提高效率。

ZeRO-DP

下面将详细地介绍ZeRO-DP的三个不同stage的情况。

\(P_{os}\) : Optimizer State Partitioning

对于DP度数为\(N_d\)，我们将优化器状态分为\(N_d\)个相等的分区，使得第i个数据并行进程只更新与第\(i\)个分区对应的优化器状态。因此，每个数据并行进程只需要存储和更新总优化器状态的1/Nd，并且只更新总参数的\(\frac{1}{N_d}\)。在每个训练步骤结束时，我们执行一次全局收集（all-gather）操作，以获取所有数据并行进程中完全更新的参数。

内存节省：如图1所示，在优化状态分区后，内存消耗从\(4\Psi + K\Psi\)减少到\(4\Psi +\frac{K\Psi}{N_d}\)。以图1中的具体示例为例，一个包含75亿参数的模型，在使用64路DP（\(N_d\) = 64）时，使用Pos占用31.4GB的内存，而使用标准DP则需要120GB的内存。此外，当Nd较大时，模型状态的内存需求从\(4\Psi + 12\Psi = 16\Psi\)字节减少到\(4Ψ + 12Ψ/Nd ≈ 4Ψ\)字节，减少了4倍.

\(P_g\): Gradient Partitioning

由于每个数据并行进程只更新其对应的参数分区，它只需要相应参数的减少梯度。因此，在反向传播期间，当每个层的梯度可用时，我们只在负责更新相应参数的数据并行进程上进行减少操作。减少完成后，我们不再需要梯度，它们所占用的内存可以释放。这将梯度所需的内存占用从\(2Ψ\)字节降低到\(\frac{2\Psi}{N_d}\)字节。

\(P_p\): Parameter Partitioning

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

内存节省：通过参数分区，我们将一个Ψ参数模型的内存消耗从16Ψ减少到16Ψ/Nd。如图1中的示例，一个包含75亿参数的模型，在使用64路DP（Nd = 64）时，使用Pos+p+g占用1.9GB的模型状态内存，而使用标准DP则需要120GB的内存。这具有深远的意义：只要有足够的设备来共享模型状态，ZeRO可以使DP适应任意大小的模型.

ZeRO-R详细介绍

详细介绍的ZeRO-R中采用的三种主要的策略：Partitioned Activation Checkpointing,Constant Size Buffers,Memory Defragmentation

\(P_a\):Partitioned Activation Checkpointing

正如第4.2节讨论的那样，MP设计上需要复制激活值，在模型并行的GPU之间产生冗余的激活值副本。ZeRO通过对激活值进行分区，并仅在计算中使用激活值之前逐个层将其以复制的形式实现。具体来说，一旦计算了模型的一层的前向传播，输入激活值就会在所有模型并行进程之间进行分区，直到在反向传播过程中再次需要它们。此时，ZeRO使用全局收集操作重新生成一个复制的激活值副本。我们称这个优化为Pa。它与激活检查点[7]结合使用，仅存储分区的激活检查点，而不是复制的副本。此外，在模型非常大且设备内存非常有限的情况下，这些分区的激活检查点也可以转移到CPU上，以几乎零的激活内存开销为代价，这将在第7节中讨论。我们称之为Pa+cpu。

内存节省：通过分区的激活检查点，ZeRO将激活值的占用空间减少了与MP度数成比例的因子。考虑使用表4中显示的100B模型进行训练，批量大小为32，序列长度为1024，MP度数为16。如果我们为每个Transformer层检查点一个单独的激活值，每个GPU需要约33GB的内存来存储激活检查点。但是使用ZeRO中的Pa，可以将其减少到约2GB每个GPU。此外，这2GB可以转移到CPU上，从而将激活值的内存占用减少到几乎零。

\(M_D\):Memory Defragmentation

在模型训练过程中，由于激活检查点和梯度计算，会出现内存碎片问题。在使用激活检查点的前向传播中，只有选择性地存储了一部分激活值用于反向传播，而大多数激活值在后向传播期间被丢弃，因为它们可以在后向传播期间重新计算。这样就产生了短期内存（被丢弃的激活值）和长期内存（检查点激活值）的交替存在，导致内存碎片化。同样，在反向传播过程中，参数梯度是长期的，而激活梯度和用于计算参数梯度的其他缓冲区是短期的。再次出现短期内存和长期内存的交替，导致内存碎片化

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

ZeRO通过动态进行内存整理来解决这个问题，它预先分配连续的内存块用于激活检查点和梯度，并在生成它们时将它们复制到预分配的内存中。MD不仅使ZeRO能够使用更大的批量大小训练更大的模型，而且在使用有限内存进行训练时提高了效率。

\(C_B\):Constant Size Buffers

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

ZeRO-DP 详细介绍

eRO通过消除内存冗余来提升模型大小，因此很自然地会问我们是否在交换通信量以获得内存效率。换句话说，与基线DP方法相比，ZeRO驱动的DP方法的通信量是多少？答案分为两部分：

• 使用\(P_{os}\)和\(P_g\)的ZeRO-DP不会产生额外的通信开销，同时可以实现高达8倍的内存减少；
• 使用\(P_p\)，在使用\(P_{os}\)和\(P_g\)之外，ZeRO-DP最多会增加1.5倍的通信开销，同时进一步减少\(N_d\)倍的内存占用。我们在本节中进行了详细分析。首先，我们对标准DP的通信量进行简要概述。

在数据并行训练中，在进行下一步更新之前，会对所有数据并行进程上的梯度进行平均化处理。平均化操作使用了一个all-reduce通信集合。对于大型模型，all-reduce通信完全受到通信带宽的限制，因此我们将分析限制在发送给每个数据并行进程和从每个数据并行进程接收的总通信量上。

现代的all-reduce实现使用了一个两步的方法，第一步是reduce-scatter操作，将数据的不同部分在不同的进程上进行归约。接下来是全收集（all-gather）操作，每个进程都收集所有进程上的归约数据。这两个步骤的结果就是all-reduce。reduce-scatter和all-gather都使用流水线化的方法实现，每个操作对应一个数据移动量为Ψ个元素（对于包含Ψ个元素的数据）。因此，在每个训练步骤中，标准DP需要进行2Ψ个元素的数据移动。

让我们详细解释一下流水线操作在reduce-scatter阶段的应用：

Reduce-Scatter阶段：在这个阶段，数据被划分为多个部分，并在不同的进程之间进行归约。然而，为了减少通信延迟，流水线操作被引入。具体来说，reduce-scatter被划分为多个阶段，每个阶段都有自己的处理步骤。
阶段一：首先，第一个进程将它的数据发送给下一个进程，同时也接收到前一个进程发来的数据。
阶段二：然后，每个进程将接收到的数据与自己的数据进行归约操作，并将结果发送给下一个进程，同时也接收到前一个进程发来的数据。
这些阶段依次进行，直到最后一个进程完成归约操作。通过这种方式，每个进程只需要等待前一个进程发送数据，而不是等待所有进程都完成归约操作。这样就能够充分利用网络带宽和计算资源，提高效率。
接下来，我们解释一下流水线操作在all-gather阶段的应用：

All-Gather阶段：在这个阶段，每个进程需要将自己的归约结果发送给其他所有进程，并接收其他进程发送的数据。为了减少通信延迟，同样采用了流水线操作。
阶段一：首先，第一个进程将它的归约结果发送给下一个进程，同时也接收到前一个进程发来的归约结果。
阶段二：然后，每个进程将接收到的归约结果与自己的归约结果合并，并将结果发送给下一个进程，同时也接收到前一个进程发来的归约结果。
这些阶段依次进行，直到最后一个进程完成数据的合并和发送操作。通过这种方式，每个进程只需要等待前一个进程发送数据，而不是等待所有进程都完成数据的合并和发送操作。
通过使用流水线操作，reduce-scatter和all-gather过程可以在并行进行，每个阶段之间不存在明显的等待时间。这允许在数据传输和计算方面的重叠，从而提高整体的效率和性能。流水线操作的好处是可以充分利用网络带宽和计算资源，减少通信延迟，降低训练时间，从而加速深度学习模型的训练过程。

\(P_{os+g}\)内存分析

在数据并行训练中，在进行下一步更新之前，会对所有数据并行进程上的梯度进行平均化处理。平均化操作使用了一个all-reduce通信集合。对于大型模型，all-reduce通信完全受到通信带宽的限制，因此我们将分析限制在发送给每个数据并行进程和从每个数据并行进程接收的总通信量上。

现代的all-reduce实现使用了一个两步的方法，第一步是reduce-scatter操作，将数据的不同部分在不同的进程上进行归约。接下来是全收集（all-gather）操作，每个进程都收集所有进程上的归约数据。这两个步骤的结果就是all-reduce。reduce-scatter和all-gather都使用流水线化的方法实现，每个操作对应一个数据移动量为Ψ个元素（对于包含Ψ个元素的数据）。因此，在每个训练步骤中，标准DP需要进行2Ψ个元素的数据移动

\(P_p\)内存分析

在参数分区之后，每个数据并行进程只存储它更新的参数。因此，在前向传播期间，它需要接收所有其他分区的参数。然而，这可以通过流水线化来避免内存开销。在计算与某个特定分区对应的模型部分的前向传播之前，负责该分区的数据并行进程可以将权重广播到所有数据并行进程。完成该分区的前向传播后，可以丢弃这些参数。总通信量因此为Ψ×Nd/Nd = Ψ。换句话说，我们通过将参数全收集重新安排在整个前向传播中，并在使用后丢弃参数来优化内存。然而，请注意在反向传播过程中，这个全收集需要以相反的顺序再次进行。因此，总通信量是这些全收集产生的通信量与梯度reduce-scatter产生的通信量之和。总体通信量因此为3Ψ，相比基准模型增加了1.5倍。梯度和参数分区都利用了这样一个观点——并非始终需要所有梯度和参数的状态——通过明智地通信来优化内存。

后面就是一些在LLM上的分析和训练，大部分在展示ZeRO在SOTA LLM上的表现。