【论文阅读】Reducing Activation Recomputation in Large Transformer Models

创新点：

针对Transformer结构，通过序列并行和选择性重计算激活值，在节省显存空间占用的情况下，不带来明显通信开销，同时减少重计算成本。
总的来说，就是在原有的张量并行的基础上，对LayerNorm和Dropout操作部分的Tensor进行序列并行，虽然形式上会增加通信量(正反向传播由4次all-reduce—》4次all-gather+4次Reduce-Scatter，但因为ring all-reduce=Reduce-Scatter+all-gather，因此两者的通信量一样)。
同时因为序列并行+张量并行，MLP中就只有进入第一个线性层的输入需要保存完整激活值用于梯度反传，这里选择了对激活值进行切分保存，反传的时候再all-gather一下，为了避免引入额外的通信开销，将all-gather通信和计算该激活值的梯度进行重叠。
对选择性重计算，只对空间占用大而重计算成本的低的部分进行checkpoint，从而做到空间节省和重计算成本之间的均衡。

摘要

训练大型 Transformer 模型是现代人工智能最重要的计算挑战之一。在本文中，我们展示了如何通过减少激活值的重新计算来显著加速大型 Transformer 模型的训练。激活值重新计算通常用于解决内存容量限制问题。传统上，为了节省内存，不存储用于反向传播的激活值，而是重新计算它们，但这增加了冗余计算。在这项工作中，我们表明大部分这种冗余计算是不必要的，因为我们可以在不进行冗余计算的情况下充分减少内存消耗。我们提出了两种新颖且非常简单的技术：序列并行和选择性激活值重新计算。结合张量并行，这些技术几乎消除了重新计算激活值的需要。我们在多达一万亿参数的语言模型上评估了我们的方法，并表明我们的方法将激活内存减少了 5 倍，同时将激活值重新计算带来的执行时间开销减少了 90% 以上。例如，在 2240 个 NVIDIA A100 GPU 上训练一个 530B 参数的 GPT - 3 风格模型时，我们实现了 54.2% 的模型浮点运算利用率，比使用重新计算时达到的 42.1% 快 29%。我们的实现将在 Megatron - LM 和 NeMo - Megatron 中提供。

1. 引言

随着 Transformer 模型朝着数万亿参数扩展，为了使模型参数、激活值和优化器状态能够放入设备内存并在实际时间内可训练，需要模型并行来在设备之间分配它们。尽管模型并行会线性减少每个设备上的参数数量，例如，当模型并行大小翻倍时，每个设备上的参数数量减半，但扩展模型并行存在限制。张量级模型并行增加了通信需求，并引入了更小且性能较低的矩阵乘法，这使得在大量设备上分割模型效率低下。因此，张量级模型并行通常仅限于一小群通过高速带宽连接的 GPU，例如在 DGX 服务器内通过 NVLink 连接的 GPU。流水线并行需要存储几个微批次的激活值以减少流水线气泡。因此，流水线并行只能帮助解决存储模型参数和优化器状态所需的内存，而不能在保持高设备利用率的同时减少激活值所需的内存。因此，激活值的存储很快成为扩展大型 Transformer 模型的关键问题。
为了量化这一点，图 1 显示了从 220 亿参数到 1 万亿参数的四种模型配置所需的内存（模型配置的详细信息在表 3 中提供）。可以看出，在所有这些情况下，基线情况下所需的内存都高于 NVIDIA A100 GPU 提供的 80GB 内存。缓解这种内存压力的标准方法是简单地不存储大多数激活值，并在反向传播期间根据需要重新计算它们以计算梯度。不幸的是，这种通常称为 “梯度检查点” 或 “激活值重新计算” 的方法会导致训练效率大幅降低。对于 Transformer 架构，大多数先前的工作在 Transformer 层边界处检查点或存储激活值，并在反向传播中重新计算其余必要的激活值。在本文中，我们将这种方法称为 “完全激活值重新计算”。在我们的训练运行中，我们观察到使用完全激活值重新计算时会有 30 - 40% 的执行时间开销。

在本文中，我们提出了新颖的技术，有助于缓解存储激活值的内存压力，从而减少重新计算激活值的需求。这些技术特定于 Transformer 架构，易于实现，并且对计算效率没有或只有极低的影响。正如我们在第 2 节中详细介绍的，还有其他几种技术可以减少训练大型模型的内存需求，例如在数据并行等级之间划分各种数据或将数据卸载到 CPU 内存。这些技术与本文提出的技术互补，可以额外使用以实现更大的内存节省；然而，一般来说，这些其他技术的实现成本更高，并且对计算效率的影响比本文提出的技术更大。将这些技术与我们的技术进行比较的分析超出了本文的范围，留待未来的工作。
我们首先简要回顾 Transformer 架构，然后建立一个用于存储单栈 Transformer 模型激活值所需内存的近似公式。使用这个公式，我们可以研究不同形式的模型并行如何影响激活内存需求。我们引入序列并行与张量并行一起，以防止在不利于标准张量并行的区域中冗余存储激活值。然后，我们表明通过选择性地保存哪些激活值和重新计算哪些激活值，我们可以在不使用重新计算时仅使用一小部分内存的情况下消除大部分重新计算成本。最后，我们进行了几个实验，测量这些技术对训练的各个组件以及整个训练吞吐量的改进。

2. 相关工作

模型并行使得能够在多个 GPU 上训练非常大的模型。这些模型的参数以及相关的优化器状态需要大量内存，无法在单个 GPU 上容纳。即使我们能够将模型放入单个 GPU 中（例如，通过在主机和设备内存之间交换参数），所需的大量计算操作也可能导致训练时间过长。这就需要并行性。通常使用两种形式的模型并行在 GPU 之间分配模型参数：1）张量并行，其中每层的参数分布在多个设备上；2）流水线并行，其中模型沿着网络的层维度进行分割。最近的一些方法结合了这两种类型的模型并行，以实现多达 1T 参数的大型模型的训练。
模型并行的替代方案是结合多种训练技术和数据并行来实现大规模模型训练。这种方法基于在数据并行等级之间分割优化器状态、梯度和参数。此外，最近的一项扩展使用 CPU 卸载技术，在少量 GPU 上实现数万亿参数模型的训练。与模型并行相比，这些基于数据并行的技术效率较低，并且在扩展到大量 GPU 时效果不佳，因此更适合在资源受限的环境中微调模型。本文仅关注模型并行优化。将这些技术与我们的技术进行比较的分析超出了本文的范围。
此外，Megatron - LM 中引入的张量并行在一定程度上有助于减少激活内存。在这种方法中，Transformer 的某些部分的激活值不会在张量并行等级之间分割，这增加了激活内存开销。如文献中所建议的序列并行，其中激活值在整个网络中沿着序列维度进行分区，可以缓解这个问题。然而，他们的方法与数据并行类似，需要在所有设备上复制参数和优化器状态，这使得它不适合大型模型训练。Sagemaker 和 GSPMD 提出了张量并行的内存高效版本，它在整个网络中沿着隐藏维度在设备之间分割激活值。这些方法的主要缺点是它们包含多设备层归一化，这在计算 / 通信方面效率非常低。与 Megatron - LM 方法中仅进行两次全规约通信相比，Sagemaker 在每个 Transformer 层进行四次规约通信。在本文中，我们提出了一种新技术，它利用了张量并行和序列并行的优点，而没有先前方法的缺点。换句话说，我们的技术将张量和序列并行混合，在没有任何额外计算、通信或内存开销的情况下显著减少激活内存。

3. Transformer 架构

在这项工作中，我们考虑一个具有 $L$ 层的简单栈堆叠 Transformer 编码器或解码器，如图 2 所示。在网络开始时，输入标记被馈送到大小为 $v\times h$ 的词嵌入表中，并且标记嵌入与大小为 $s\times h$ 的学习位置嵌入相结合，其中 $s$ 是序列长度，$h$ 是隐藏维度，$v$ 是词汇大小。嵌入层的输出，即 Transformer 块的输入，是大小为 $s\times b\times h$ 的三维张量，其中 $b$ 是微批次大小。每个 Transformer 层由一个具有 $a$ 个注意力头的自注意力块和一个具有两层的多层感知器（MLP）组成，该多层感知器将隐藏大小增加到 $4h$，然后再将其减少回 $h$。每个 Transformer 层的输入和输出具有相同的大小 $s\times b\times h$。最后一个 Transformer 层的输出被投影回词汇维度以计算交叉熵损失。我们假设词嵌入和输出层权重是共享的。变量名称列在表 1 中以供参考。

4. 激活内存

在本节中，我们推导一个用于存储单栈 Transformer 模型前向传播中激活值所需内存的近似公式。请注意，本文中的 “激活值” 是指在前向传播中创建且在反向传播中计算梯度所需的任何张量。因此，这不包括模型的主要参数和优化器状态，但包括例如 dropout 操作使用的掩码。
此外，我们仅考虑内存的主要贡献者并忽略小缓冲区。例如，对于层归一化块，计算梯度需要层的输入以及输入的均值和方差。输入包含 $sbh$ 个元素，而均值和方差每个只有 $sb$ 个元素。由于 h 很大（数千量级），我们有 $2sb\ll sbh$。因此，仅考虑存储输入所需的内存是一个很好的近似，即我们仅包括 $sbh$，而不是 $sbh+2sb$。
我们还假设网络和激活值以 16 位浮点格式存储，因此每个元素需要 2 字节的存储空间。唯一的例外是 dropout 掩码，每个元素仅需要 1 字节。请注意，本节中报告的所有大小均以字节为单位，除非明确提及，否则不是元素数量。

4.1 每个 Transformer 层的激活内存

如图 2 所示，每个 Transformer 层由一个注意力块和一个 MLP 块组成，它们通过两个层归一化连接。下面，我们推导存储这些元素的激活值所需的内存：
注意力块：包括自注意力、线性投影和注意力 dropout。线性投影存储其大小为 $2sbh$ 的输入激活值，注意力 dropout 需要一个大小为 $sbh$ 的掩码。如图 3 所示的自注意力由几个元素组成：
查询（Q）、键（K）和值（V）矩阵乘法：我们只需要存储它们的共享输入，大小为 $2sbh$。
矩阵乘法：它需要存储 Q 和 K，总大小为 $4sbh$。
Softmax：反向传播需要大小为 $2a{s}^{2}b$的 Softmax 输出。
Softmax dropout：只需要一个大小为 $a{s}^{2}b$的掩码。
对值（V）的注意力：我们需要存储 dropout 输出（$2a{s}^{2}b$）和值（$2sbh$），因此需要 $2a{s}^{2}b+2sbh$的存储空间。
将上述值相加，注意力块总共需要 字节的存储空间。
MLP：两个线性层存储其大小为 $2sbh$ 和 $8sbh$ 的输入。GeLU 非线性函数也需要其大小为 8