【论文阅读】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 非线性函数也需要其大小为 $8sbh$ 的输入用于反向传播。最后，dropout 存储其大小为 $sbh$ 的掩码。总共，MLP 块需要 $19sbh$ 字节的存储空间。
层归一化：每个层归一化存储其大小为 $2sbh$ 的输入，因此总共需要 $4sbh$ 的存储空间。
将注意力、MLP 和层归一化所需的内存相加，存储单个 Transformer 层激活值所需的内存为：

上述方程是在没有应用任何形式的模型并行的情况下。

4.2 模型并行

在本节中，我们首先量化张量并行对每层所需激活内存的影响。然后，我们引入一种称为序列并行的新方法，该方法进一步减少了激活值每层所需的内存。在本节末尾，我们还讨论了流水线并行对激活内存的影响，并推导了激活值所需总内存的公式。

4.2.1 张量并行

我们使用 Shoeybi 等人开发的张量并行，并并行化注意力和 MLP 块，如图 4 所示。这种并行形式引入了两个额外的通信操作 $f$ 和$\overline{f}$。有关更多详细信息，请参阅论文。

张量并行不仅并行化了注意力和 MLP 块内的模型参数和优化器状态，还并行化了这些块内的激活值。请注意，这些块的输入激活值（例如，Q、K 和 V 矩阵乘法的输入或 线性层的输入）没有并行化，只有每个块内的激活值在张量并行组中进行分割。假设采用 $t$ 路张量并行，存储激活值所需的每层内存从公式（1）减少到：

4.2.2 序列并行

如图 4 所示，张量并行并行化了 Transformer 层中在训练期间耗时最多的部分，因此计算效率较高。然而，它保留了注意力和 MLP 块之后的层归一化以及 dropout，因此它们在张量并行组中被复制。这些元素不需要大量计算，但需要相当数量的激活内存。定量地说，公式（2）中的 $10sbh$ 部分是由于这些复制操作，因此它们不会被张量并行大小 $t$ 除。
我们注意到，在 Transformer 层的非张量并行区域中，操作在序列维度上是独立的。这个特性允许我们沿着序列维度 $s$ 对这些区域进行分区。沿着序列维度进行分区减少了激活值所需的内存。这种额外的并行级别在 $f$ 之前和 $\overline{f}$之后引入了新的通信集合操作，它们将充当序列并行区域和张量并行区域之间的转换器。例如，在前向传播中，我们需要在图 4 中的操作符 $f$ 之前进行额外的all-gather操作。这些额外的通信会引入开销并减慢训练速度。
为了避免这些额外的通信，我们将这些操作与 $f$ 和 $\overline{f}$操作符相结合，并引入新的操作 $g$ 和$\overline{g}$，如图 5 所示。可以看出，$g$ 和$\overline{g}$ 是序列并行区域和张量并行区域之间的转换器。我们在本节的剩余部分推导这些操作。

我们使用 MLP 块详细推导 $g$ 和$\overline{g}$ 。在非并行形式中，如图 2 所示，层归一化之后的 MLP 块可以表示为：
$Y=LayerNorm(X)$
$Z=GeLU(YA)$
$W=ZB$
$V=Dropout(W)$
其中 $X$ 是层归一化的输入，大小为$s\times b\times h$，$A$ 和 $B$ 是线性层的权重矩阵，大小分别为$h\times 4h$ 和$4h\times h$。上述操作的组合张量和序列并行形式如图 6 所示。下标表示在加速器之间的分割，上标表示分割所沿的维度。例如， 是 $X_1^s$ 在第一个加速器上沿 $s$ 维度（序列维度）分割的部分，而 是 $Z_2^h$ 在第二个加速器上沿 $h$ 维度（隐藏维度）分割的部分。

层归一化的输入沿着序列维度$X=[X_1^s,X_2^s]$并行化。因此，层归一化的输出也将沿着序列维度$Y=[Y_1^s,Y_2^s]$并行。具有 GeLU 非线性函数的线性层需要整个输入 $Y$(主要是矩阵运算需要完整的输入)，因此我们需要执行all-gather操作。这意味着 $g$ 在前向传播中是沿序列维度的all-gather操作。通过将 $A$ 沿其列分割（ $A_1^c$和 $A_2^c$）和 $B$ 沿其行分割（ $B_1^r$和$B_2^r$），我们避免了通信（更多详细信息请参阅论文[19]）并得到 $W_1$和$W_2$。这两个张量不再并行，需要在将它们馈送到 dropout 层之前进行求和$W=W_1+W_2$。然而，dropout 需要其输入在序列维度 s 上并行。与其先求和然后在序列维度上并行化，我们将这两个操作组合成一个reduce-scatter操作。因此， $\overline{g}$在前向传播中可以是单个reduce-scatter操作。综上所述，我们得到：

如果我们对反向传播进行类似的分解，我们发现 $g$ 和$\overline{g}$ 是彼此的共轭。$g$ 在前向传播中是all-gather操作，在反向传播中是reduce-scatter操作，而 $\overline{g}$在前向传播中是reduce-scatter操作，在反向传播中是all-gather操作。对 Transformer 层中层归一化之后的注意力部分进行类似的分解得到图 5。
张量并行在单个前向和反向传播中需要四次all-reduce操作，而张量与序列并行一起在单个前向和反向传播中需要四次all-gather和四次reduce-scatter操作。乍一看，似乎张量与序列并行相比张量并行需要更多的通信。然而，我们注意到ring all-reduce由两个步骤组成：reduce-scatter后接all-gather。因此，用于张量并行和张量与序列并行的通信带宽是相同的。因此，序列并行不会引入任何通信开销。
从公式（3）中，序列并行与张量并行一起将反向传播所需的所有激活值沿着并行维度进行分割，除了第一个线性操作所需的张量 $Y$。为了缓解这个问题，我们不在反向传播中存储完整的张量 $Y$。相反，我们仅在第 i 个张量并行等级上存储 $Y_i^s$部分，并在反向传播中执行额外的all-gather操作。为了消除此额外all-gather操作引入的延迟，我们将此通信与计算 $Y$ 的梯度所需的计算重叠，从而减少开销。
使用序列并行与张量并行，存储每个 Transformer 层激活值所需的内存从公式（2）减少到：

上述公式现在是公式（1）除以张量并行大小。这意味着，使用张量和序列并行，我们可以在张量并行组之间分配激活，并将所需的内存减少张量并行大小 $t$。

4.2.3 Pipeline Parallelism

流水线并行只是将 Transformer 的 $L$ 层划分为$L/p$组层，其中 $p$ 是流水线并行大小。然而，流水线并行并不会将激活所需的总内存均匀地除以 $p$。这是由于流水线并行调度为减少流水线气泡而引入的重叠所致。
为了量化这一点，我们考虑在 PipeDream 中开发的 1F1B 流水线调度。具有最小化流水线气泡的调度会给流水线的第一阶段带来最大的内存压力（流水线的第一阶段是指$L/p$第一组层，其中还包括输入嵌入）。附录 B 中展示了激活内存作为流水线阶段函数的可视化。为了保持流水线有压力并避免额外的空闲时间，第一阶段必须为 $p$ 个微批次存储激活值（更多详细信息见论文引用[13]的图 4 顶部）。每个阶段包含$L/p$层，所以无论流水线并行大小 p 如何，第一阶段都必须存储相当于 $p\times L/p=L$ 层的激活值。因此，在第一阶段存储激活值所需的总内存为：
对于其他流水线调度，所需的总内存会略有不同。例如，Megatron - LM 中开发的交错调度需要为$L(1+\frac{p-1}{pm})$层存储激活值，其中 $m$ 是交错阶段的数量。因此，如果使用交错调度，总的激活值内存大小将被缩放$(1+\frac{p-1}{pm})$。

4.3 Total Activations Memory

所需激活内存的大部分由公式（5）涵盖。然而，该公式没有考虑图 2 中所示的输入嵌入、最后一层归一化和输出层所需的激活内存。
位置和词嵌入在反向传播中不需要存储大量激活值。然而，dropout 需要存储。嵌入层中的 dropout 也沿着序列维度并行化。因此，它将需要$sbhp/t$的存储空间。请注意，因子 $p$ 来自流水线并行，并且我们需要存储 $p$ 个微批次（见第 4.2.3 节）。
输出层之前的层归一化也使用序列并行，因此需要$2sbh/t$的存储空间。输出层投影到词汇维度将需要存储其大小为$2sbh/t$的输入。最后，交叉熵损失需要存储以 32 位浮点计算的对数，因此将需要$4sbv/t$的存储空间。请注意，由于我们只考虑流水线第一阶段的激活值，上述激活值，即总共$4sbh/t(1+v/h)$，仅在没有流水线并行（$p=1$）的情况下才包括在内。

加上上述内存，由于输入嵌入、最后一层归一化和输出层导致的额外内存为：

，其中${\delta }_{p=1}$在$p=1$时为 1，否则为 0。我们注意到，与公式（5）中的$34+5\frac{as}{h}$项相比，$p/L$和$4/L(1+v/h)$都可以忽略不计。例如，对于一个具有 220 亿参数的模型，这些额外项占总激活内存需求的不到 0.01%。因此，公式（5）是所需总激活内存的一个很好的近似，我们将在本文的其余部分使用它。

5. 选择性激活值重新计算

公式（5）中的总所需激活内存对于大型模型仍然可能相当大。激活值重新计算通过存储（或 “checkpoint”）一组层的输入激活值，并在反向传播期间使用额外的前向传播重新计算其他所需的激活值来克服此内存限制（本文中将此称为完全激活值重新计算）。假设组仅包含单个层，并且忽略 Transformer 层之外的激活值，此方法将激活值所需的总内存减少到 $2sbhL$。我们注意到，如果我们仅在每个张量并行等级中存储一部分激活值，则所需内存可以进一步减少到 $2sbhL/t$。然而，这种方法每层需要额外的all-gather操作，会增加通信开销，因此我们不考虑这种方法。
与存储所有激活值（公式（5））相比，检查点所有 Transformer 层显著减少了训练模型所需的内存量。这种减少是以重新计算（额外的前向传播）为代价的，这可能会引入高达 30 - 40% 的计算时间开销。为了平衡内存节省和计算开销，理想的情况是仅检查点足够的激活值，以允许给定的模型并行配置在设备内存的限制下进行训练。序列并行提供的内存节省允许比以前更多的配置在不进行重新计算的情况下进行训练，但是大型模型的最佳模型并行配置通常仍然需要一些激活值的保存和重新计算。一种选择存储与重新计算的激活值数量的简单方法是仅checkpoint一些 Transformer 层并存储其他层的所有激活值。这种方法在扩展到大型模型时效果不佳；例如，在训练 MT - NLG 时，每个设备只有三层，这限制了你可以平衡内存与计算的粒度。此外，我们注意到并非所有激活值都需要相同数量的操作来重新计算，因此更智能地选择存储哪些激活值和重新计算哪些激活值是有益的。
我们不checkpoint和重新计算整个 Transformer 层，而是建议仅checkpoint和重新计算每个 Transformer 层中占用大量内存但重新计算计算成本不高的部分，即选择性激活值重新计算。为此，我们注意到公式（5）中的$5\frac{as}{h}$项是由于在通过计算 Q、K 和 V 值的线性层增加网络宽度之后的注意力操作；即矩阵乘法、softmax、softmax dropout 和对 V 的注意力，如图 3 所示。这些操作通常具有较大的输入大小，因此激活值也较大，但是每个输入元素的浮点运算（FLOPs）数量非常低。Transformer 层的其余部分占公式（5）中的 34 项。因此，对于大型模型，其中$5as/h>34$，如果我们检查点并重新计算 Transformer 层的这部分，我们存储的激活值不到一半，并且重新计算未存储的激活值的成本也不高。
为了量化这一点，让我们考虑 GPT - 3 和 MT - NLG 模型，它们是迄今为止训练的一些最大的模型。对于 GPT - 3，$a=96,s=2048,andh=12288$，因此$5as/h=80$。对于 MT - NLG，$a=128,s=2048,andh=20480$，所以$5as/h=64$。将这些数字与 34（层其余部分的因子）进行比较，我们可以看到这些激活值占总激活值的很大一部分。因此，通过使用选择性激活值重新计算，我们可以分别为 GPT - 3 和 MT - NLG 模型节省 70% 和 65% 的激活值所需内存。这些激活值的重新计算仅为这两个模型引入了 2.7% 和 1.6% 的 FLOPs 开销。有关 FLOPs 计算的更多详细信息，请参阅附录 A。
使用这种形式的选择性激活值重新计算，存储激活值所需的内存从公式（5）减少到：

上述方程表明，使用选择性激活值重新计算允许所需的激活内存与序列长度线性缩放，并且与注意力头的数量无关。如第 4.2.3 节所述，在交错流水线调度的情况下，上述方程需要乘以$(1+\frac{p-1}{pm})$。
当使用流水线并行时，如第 4.2.3 节所述，即使给定设备只有$L/p$层，第一阶段仍必须存储相当于 L 层的激活值，因为它必须为 p 个微批次存储激活值以保持流水线有压力。在这种情况下，可以采用的另一种减少重新计算成本的技术是在给定可用设备内存的情况下，为尽可能多的微批次存储所有激活值，并对其余部分进行完全或选择性重新计算。在实践中，我们发现，在应用序列并行和选择性激活值重新计算之后，重新计算开销足够小，以至于这种额外技术提供的改进非常有限。这种技术在附录 C 中有更详细的描述和分析。

6. 评估

在本节中，我们评估所提出方法对训练的内存使用和执行速度的影响。表 3 列出了评估中使用的模型配置。我们考虑多达一万亿参数的模型，对于所有这些模型，张量并行大小设置为 8。我们对 175B 和 530B 模型使用具有三个交错阶段（$m=3$）的交错调度。在所有情况下，序列长度设置为$s=2048$，词汇大小设置为$v=51200$。我们还注意到，在这些评估中不考虑数据并行，因为我们的方法与数据并行无关。因此，我们分析中使用的批量大小远低于端到端训练中使用的批量大小。我们所有的结果都是在 Selene 超级计算机上以混合精度运行的。每个集群节点有 8 个 NVIDIA 80GB A100 GPU，通过 NVLink 和 NVSwitch 相互连接。每个节点有八个 NVIDIA Mellanox 200Gbps HDR Infiniband HCAs 用于应用程序通信。

6.1 内存使用

表 2 总结了本文讨论的不同技术所需的内存。为了量化这一点，图 7 显示了不同技术使用的激活内存占将所有激活值在张量并行等级之间分割所需内存（即公式（2））的百分比。单独来看，这两种技术都将内存需求几乎减半，并且结合起来提供了 5 倍的减少，使内存需求降至 20% 以下。这仅约为完全激活值重新计算的 2 倍，完全激活值重新计算为基线的 10%。如果没有序列并行和选择性重新计算共同提供的内存节省，这些模型都无法放入内存。请注意，所有这些结果都包括附录 B 中描述的内存优化。

6.2 每层执行时间

表 4 显示了 22B 模型的一个 Transformer 层在各种实验中的前向和反向传播执行时间。前两行表明序列并行对完成一个 Transformer 层的时间有适度的改进，将前向时间从 7.7ms 减少到 7.2ms，速度提高了 6%。这种改进来自于在的数据上执行层归一化和 dropout 层。我们还发现，尽管移动的数据量相同，但reduce-scatter和all-gather操作的组合执行比单独的all-reduce操作慢，从而降低了序列并行带来的改进。请注意，这种加速是使用序列并行的主要优势（即允许减少激活值重新计算的内存节省）之外的额外好处。

表 4 中的接下来两行表明，如果我们有选择地重新计算操作（由于序列并行，我们可以在更多配置中这样做），我们可以显著减少反向传播中的重新计算开销。选择性重新计算的开销为 1.3ms，占 11.9ms 基线的 11%，而重新计算整个层的开销为 7.6ms，占 64%。对于前向和反向传播的总时间，开销为 7%，而完全重新计算的开销为 39%。请注意，重新计算整个层的开销为 39%（而不是预期的 33%）是由于在反向传播中我们将a通信ll-reduce与线性权重的梯度计算重叠的优化。正如我们稍后看到的，这种好处随着模型大小的增加而增加。表 4 中的最后一行显示了选择性重新计算和序列并行的综合好处。当这两种技术一起使用时，开销降至仅 4%。
图 8 显示了我们所有测试案例的相同细分。我们看到，随着模型大小的增长，开销的减少也增加。对于 530B 和 1T 的情况，开销仅为 2%，而完全重新计算的开销为 36%。

6.3 端到端迭代时间

表 5 列出了表 3 中列出的四种配置的完整端到端迭代时间。我们发现，对于所有测试的配置，本文提出的技术在吞吐量方面比不使用序列并行的完全重新计算提高了 29.0% 到 32.1%。这些节省将直接转化为更短的训练时间。

我们类似于 Chowdhery 等人定义模型 FLOPs 利用率（MFU）和硬件 FLOPs 利用率（HFU）。模型 FLOPs 是执行单个前向和反向传播（单次迭代）所需的浮点运算数量，与实现和硬件限制无关。因此，模型 FLOPs 与硬件和实现无关，仅取决于基础模型。另一方面，硬件 FLOPs 表示每次迭代在硬件上实际执行的浮点运算数量。因此，如果实现需要激活值重新计算（例如我们的实现），那么硬件 FLOPs 将大于模型 FLOPs。我们在附录 A 中提供了模型和硬件 FLOPs 的紧密下限公式。对于我们的方法，硬件与模型 FLOPs 的比率约为$1+s/6h$。
随后，我们将模型和硬件 FLOPs 每秒定义为模型和硬件 FLOPs 分别除以迭代时间。使用这些定义，MFU 和 HFU 定义为模型和硬件 FLOPs 每秒除以加速器理论峰值 FLOPs 每秒。表 5 中提供了所有四种配置的 MFU 和 HFU。随着模型大小的增加，我们实现了更好的 GPU 利用率，对于一万亿参数模型，我们分别达到了 56.3% 和 57.0% 的 MFU 和 HFU。
我们应该注意到，表 5 中的结果没有使用任何数据并行。数据并行由于数据并行组之间所需的梯度全规约会引入一些开销。然而，对于大型 Transformer 模型，这种开销并不大。例如，如果我们将 530B 模型扩展到 8 路数据并行（2240 个 GPU），同时保持每个模型实例的批量大小不变（即批量大小也乘以数据并行大小），每次迭代的时间从 37.83 秒略微增加到 39.15 秒。这导致 MFU 从 56.0% 下降到 54.2%，这并不显著。我们注意到我们没有使用梯度全规约与反向传播的任何重叠，有效的重叠可以几乎完全消除迭代时间的这种增加。

7. 结论和未来工作

在这项工作中，我们提出了两种新颖且简单的技术，可减少存储激活值带来的内存压力，从而减少重新计算激活值的需求。我们表明，将序列并行与张量并行一起使用可以显著减少所需的激活内存。结合选择性激活值重新计算，我们表明我们可以实现 5 倍的内存减少，并恢复使用完全激活值重新计算引入的超过 90% 的计算开销。
在未来，我们计划通过解决大微批次内存碎片和流水线并行导致的非均匀内存分配问题，进一步减少激活内存。此外，我们计划研究可以减少流水线第一阶段内存压力的方法。