高效深度学习：优化技术综合概述 👐 📚

社区文章发布于 2024 年 8 月 26 日

Daniil Suhoi

Isayoften

训练大型语言模型（LLM）需要大量的计算资源和时间。然而，通过优化训练过程，可以降低成本、加速开发并提高模型的整体性能。本指南全面探讨了各种优化策略，涵盖了从内存消耗基础知识到训练过程优化和分布式训练的所有内容。

我想指出，本文基本上是来自各种文章最相关摘录的组合，因此我能够以最高的质量和可靠性呈现材料。

0. 数据类型介绍

在深入模型训练的复杂细节之前，我们先简要探讨一下数字在计算机中的表示方式以及可用的不同数据表示类型。这些基础知识对于理解模型训练期间的内存消耗至关重要。

Int16/Int8/Int4

这些是标准整数类型。它们可以表示的值范围由 $[-2^{n-1}, 2^{n-1} - 1]$

Int16 的位布局示意图如下：1 位符号位和 15 位值位。

使用的位数越多，可以表示的值范围就越大。

Float32

在 Float32 中，位布局如下：1 位符号位、8 位指数位和 23 位尾数位。

值的公式为： $v = (-1)^{\text{sign}} \cdot 2^{E-127} \cdot \left(1 + \sum_{i=1}^{23} b_{23-i}2^{-i}\right)$

浮点类型背后的关键思想是，分配给指数的位数越多，可表示的值范围就越广，而分配给尾数的位数则决定了该范围内的精度。

Float16

Float16 格式使用 1 位符号位、5 位指数位和 10 位尾数位。

Float16 的主要缺点是其值范围有限，最大值为 65504，这使其容易在激活张量中发生溢出。

Bfloat16，或 brain float

Bfloat16 是 Google Brain 开发的一种专用数据格式。它可以被认为是 Float32 的近似值。其位布局为 1 位符号位、8 位指数位和 7 位尾数位。

请注意，指数位的数量与 Float32 相同，这意味着 bfloat16 可以表示相同的值范围，尽管精度较低。这降低了激活中的溢出风险。

bf16 的另一个优点是易于将值转换为 Float32。这得益于相似的位布局。然而，目前并非所有硬件都支持此类型，尤其是在移动设备中。

TensorFloat32

TensorFloat32 是 NVidia 推出的一种有趣的 19 位数据类型，在 NVidia Ampere (A-100) 架构上受支持。其位布局包括 1 位符号位、8 位指数位和 10 位尾数位。

主要特点

指数位的数量与 bfloat16 相同，因此也与 Float32 相同。
尾数位的数量与 Float16 相同。

这产生了一种不寻常但高效且精确的数据类型。它提供了出色的计算性能，适用于模型训练，尽管它仅在现代 NVidia GPU 上可用。

E4M3 和 E5M2

这些是 NVidia、ARM 和 Intel 在论文 FP8 Formats for Deep Learning 中引入的新型 8 位浮点类型。作者提出了两种可能的 8 位浮点格式：

E4M3：1 位符号位、4 位指数位和 3 位尾数位。
E5M2：1 位符号位、5 位指数位和 2 位尾数位。

实验表明，现代 LLM 和图像网络可以使用这些数据类型成功训练甚至进行推理。我们期待它们更广泛的采用和硬件支持。还有更激进的 4 位浮点格式的想法，例如 E2M1 和 E3M0。

1. 内存都去哪儿了？

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

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

训练期间，设备内存的大部分由模型状态消耗。例如，Adam 是深度学习训练中最流行的优化器之一，它需要存储两个优化器状态：1) 梯度的时间平均动量，和 2) 梯度的方差，用于计算更新。

因此，要使用 Adam 训练模型，必须有足够的内存来同时保存动量和梯度的方差副本。此外，还需要足够的内存来存储梯度和权重本身。在这三种参数相关张量中，优化器状态通常消耗的内存最多，尤其是在应用混合精度训练时。

混合精度训练 当前一代 NVIDIA GPU 上训练大型模型的最新方法是混合精度训练，其中参数和激活以 fp16 存储，从而能够利用这些 GPU 上的高吞吐量张量核心单元。在混合精度训练期间，前向和反向传播都使用 fp16 权重和激活执行。然而，为了有效地计算和应用反向传播结束时的更新，混合精度优化器保留了参数的 fp32 副本以及所有其他优化器状态的 fp32 副本。

让我们以 Adam 为具体示例。使用 Adam 对具有 Φ 参数的模型进行混合精度训练，需要足够的内存来保存参数和梯度的 fp16 副本，内存需求分别为 2Φ 和 2Φ 字节。此外，它还需要保存优化器状态：参数的 fp32 副本、动量和方差，内存需求分别为 4Φ、4Φ 和 4Φ 字节。

总计，这导致内存需求为 16Φ 字节。对于像 GPT-2 这样拥有 15 亿参数的模型，这导致内存需求至少为 24 GB，这远高于仅保存 fp16 参数所需的 3 GB 内存。

1.2 残留内存消耗

在训练过程中，激活会占用大量内存。举个具体例子，一个拥有 1.5B 参数的 GPT-2 模型，在序列长度为 1K、批处理大小为 32 的情况下进行训练，需要大约 60 GB 的内存。

基于 Transformer 的模型的激活内存与*Transformer 层数* × *隐藏维度* × *序列长度* × *批处理大小*成正比。

激活检查点（或梯度检查点）是一种常见的减少激活内存的方法，通过增加 33% 的重新计算开销，将总激活内存减少约其平方根。这将使该模型的激活内存消耗从 60 GB 减少到约 8 GB。

尽管大幅减少，但对于更大的模型，即使使用激活检查点，激活内存也可能变得相当大。例如，一个拥有 1000 亿参数的 GPT 类模型，即使使用激活检查点，在批处理大小为 32 时也需要大约 60 GB 的内存。

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

内存碎片化：到目前为止，我们已经讨论了训练过程中实际的内存消耗。此外，即使有大量的可用内存，也可能出现可用内存不足的情况。这可能发生在内存碎片化时。如果无法满足内存请求，即使总可用内存大于请求内存，请求也会失败。我们观察到在训练非常大的模型时，内存碎片化现象显著，在某些极端情况下，即使仍有超过 30% 的内存可用，也会导致内存不足问题。

2. 量化

深度学习中的量化是指降低用于表示模型参数（权重）和计算的数字精度，通常从 32 位浮点（FP32）转换为 16 位浮点（FP16）、8 位整数（INT8）甚至更低的位宽格式。量化的主要目标是减小模型大小、减少内存使用，并通过使模型能够在计算资源有限的硬件上高效运行来加速推理。

一般来说，无法在量化模型上执行纯 4 位/8 位训练。但是，您可以利用参数高效微调方法 (PEFT) 训练这些模型，并在其上训练例如适配器。我们将在下一节深入探讨这种方法。

最简单的“量化”形式是将参数从 fp32 转换为 fp16。在训练期间，主权重始终存储在 FP32 中，但实际上，半精度权重在推理过程中通常提供与 fp32 对应物相似的质量——只有当模型接收到多个梯度更新时，才需要模型的精确参考。这意味着我们可以使用半精度权重并使用一半的 GPU 来达到相同的效果。

如果能进一步降低精度就太棒了，但推理质量在较低精度下会急剧下降。这就是为什么我们需要更巧妙的方法来做到这一点。

量化本质上是通过从一种数据类型“四舍五入”到另一种数据类型来完成的。例如，如果一种数据类型的范围是 0..9，另一种是 0..4，那么第一种数据类型中的值“4”将四舍五入到第二种数据类型中的“2”。但是，如果第一种数据类型中的值是“3”，它介于第二种数据类型的 1 和 2 之间，那么我们通常会四舍五入到“2”。这表明第一种数据类型中的值“4”和“3”在第二种数据类型中具有相同的值“2”。这突出表明量化是一个有噪声的过程，可能导致信息丢失，是一种有损压缩。

2.1 非对称和对称线性量化

我们先从插图开始

非对称：

对称：

本质上，我们正在将连续的实数范围映射到整数范围。该过程可以可视化如下：

此处，S 和 Z 是量化参数，在量化过程中计算得出。S（比例因子）决定了转换的比例，而 Z（零点）对应于量化域中的零值。

非对称
- $S = \frac {r_{max}-r_ {min}}{q_{max}-q_{min}}$
- $Z = \left[q_{min} - \frac{r_{min}}{S}\right]$
- $X_{quantized} = \left[\frac{X}{S} + Z\right]$
- $X_{dequantized} = S(X_{quantized} - Z)$
对称
- 量化范围由数据的最大绝对值决定。
- $S = \frac{|r|_{max}}{2^{N-1} - 1}$
- $Z = 0$
- $X_{quantized} = \left[\frac{X}{S}\right]$
- $X_{dequantized} = SX_{quantized}$
- 为了保持对称性，通常从量化数据类型中删除一个值。例如，带符号的 int8 范围 [-128, 127] 变为 [-127, 127]。

其中 $[*]$ 表示四舍五入。

非对称量化的优势在于其能够更好地处理非对称数据分布，而对称量化则受益于其简单性和速度。使用对称量化，无需存储零点，去量化只是简单的乘以一个常数。

对称量化示例

结果是一个 8 位整数张量，量化常数为 23.5。这可以减少存储要求，并在必要时转换回原始 32 位浮点表示，尽管会有一些精度损失。

2.2 量化什么？

标准方法是量化模型的权重。这不需要额外的操作——只需应用公式。

您还可以量化层的输出，称为激活。为此，您需要估计激活张量中出现的值范围。这是通过将训练数据集中的数据通过训练好的神经网络并收集统计数据来完成的。利用这些信息，您确定量化参数。这种方法称为静态量化。

在动态量化中，激活在推理期间进行量化。这种方法可以产生更好的质量，但它也带来了挑战：在推理期间动态地找到量化参数使得该方法更加复杂和计算昂贵，尽管它确保参数始终是最新的。

2.3 何时量化？

网络量化可以在训练期间进行，称为量化感知训练。在这种方法中，特殊的块嵌入在神经网络中，并在训练期间模拟量化推理。

量化感知训练复杂且需要更多计算资源，但它生成的模型“适应”处理量化值，可能提供更高的准确性。

在训练后量化方法中，已经训练好的模型被量化。对于激活量化，您将校准数据集中的数据通过训练好的网络，收集张量统计信息，然后进行量化。如果您只量化权重，则无需额外数据，因为所有必要信息都已在张量中。这种方法比量化感知训练更简单、更快，但通常准确性较低。

2.4 粒度

量化可以以不同粒度级别应用。最基本的方法是同时量化整个网络，导致整个模型只有一个比例因子 S。这通常会导致不令人满意的结果。

更好的方法是单独量化张量，允许每个张量有自己的比例因子。您甚至可以更进一步，量化每个张量内的行或列，为每行（或列）设置自己的比例因子。尽管这会增加比例因子的存储要求，但它显著提高了计算的准确性。

您还可以将张量分成更小的块，这将产生更高的精度。这种方法有助于减轻矩阵中异常值的影响，我们将在后面进一步探讨这个主题。

总而言之，粒度越小，您需要存储的比例因子越少；反之，粒度越高，量化计算就越接近原始计算。

2.5 数据类型

量化神经网络模型通常涉及两种数据类型

量化类型 — 用于存储张量的类型。
计算类型 — 用于执行计算的类型。

不幸的是，这两种类型并不总是匹配。例如，您的硬件可能不支持特定量化类型中的操作。某些量化类型的高效矩阵乘法核可能不存在。在这种情况下，您需要在执行计算之前将矩阵转换为计算类型。计算类型还有助于避免激活中的溢出问题，因为将 8 位数字相乘很容易超出数据类型的限制。

2.6 离群点问题

考虑对称量化的例子

如果输入张量中存在离群点会发生什么？

权重被“压缩”到狭窄的范围，变得难以区分。模型的质量受到损害。在这种情况下，一个离群点毁掉了整个矩阵。

随着参数数量的增加，我们上面讨论的标准量化技术开始失效。当参数数量超过 67 亿时，量化模型会显著损失质量。这是由于矩阵中离群点数量的增加。

2.7 LLM.int8()

论文的作者提出了一种将大型模型（高达 1750 亿参数）从通常的 16 位或 32 位浮点权重量化为 8 位整数的方法，且质量损失最小。其关键思想是单独处理离群点，因为它们只占数据的一小部分（0.1-1% 的所有值），并且集中在激活张量的特定通道中。

我们考虑激活矩阵 X 乘以权重矩阵 W。X 的列分为两组：包含至少一个离群值的列和不包含任何离群值的列。这种划分导致从原始 W 派生出两个新的权重矩阵。

需要注意的是，激活 X 的第 i 列仅与权重 W 的第 i 行交互。因此，矩阵 W 也可以通过分离与 X 的离群列对应的行来分成两部分。

结果，我们得到了两组矩阵：一组带有离群值，一组不带。然后分别乘以每组，并将结果相加。这个和等同于通常的矩阵乘法。

大多数值将落入没有离群值的矩阵中，可以很容易地量化为 8 位，从而实现高效操作。包含离群值的矩阵保留其原始的 16 位类型，以确保计算保持准确。

然而，量化精度提高的代价是由于额外计算的开销导致性能下降。作者的基准测试显示，与 16 位默认值相比，BLOOM-176B 的推理速度降低了 15-23%。

2.8 GPTQ

量化技术正在迅速发展，不断涌现出新的高效方法。我们不再深入探讨这个主题，但会简要介绍另一种替代方法。

让我们重新审视这个问题：四舍五入到最接近的整数是最佳解决方案吗？也许不是。我们的实际目标是找到一个量化的权重矩阵 $\hat{W}$ ，当它与激活矩阵相乘时，产生的结果尽可能接近原始结果： $\min_{\hat{W}} \|XW - X\hat{W}\|_2^2$

这涉及到许多数学和工程解决方案，但其思想应该很清楚。更多细节可以参考原始论文

值得注意的是，到目前为止讨论的一切都只集中在将量化模型用于推理优化。那么训练呢？

3. PEFT（参数高效微调）、LoRA 和 QLoRa

PEFT 是一系列旨在通过只训练一小部分参数来高效调整大型模型的方法。这些方法显著降低了计算成本和内存需求，同时保持了与完全微调相当的质量。

3.1 LoRA：低秩适应

最流行和最有效的 PEFT 方法之一是 LoRa。

为了理解此图解，我们来深入探讨使此方法有效的基本观察结果。

神经网络包含许多执行矩阵乘法的密集层。这些层中的权重矩阵通常是满秩的。在适应特定任务时，Aghajanyan 等人 (2020) 的研究表明，预训练语言模型具有较低的“内在维度”，并且即使随机投影到较小的子空间也能有效学习。

这意味着，在训练广泛、复杂的任务时，神经网络中的权重矩阵是满秩的，这最大程度地减少了冗余。然而，当针对特定任务微调此通用模型时，原始模型中的所有知识并非都必要。因此，只需训练一小部分参数。简而言之，权重矩阵可以用参数更少的较小矩阵表示。因此，在完全微调期间，权重矩阵可以被认为是低秩的，这表明完全微调涉及一定程度的冗余。

受此启发，我们假设权重更新在适应过程中也具有较低的“内在秩”。

鉴于低秩权重矩阵足以对下游任务进行完全微调，因此可以合理地假设梯度更新本身可以通过低秩矩阵表示。

对于预训练的权重矩阵 $W_0 \in \mathbb{R}^{d\times d}$ ，我们通过将其表示为低秩分解 $W_0 + \Delta W = W_0 + BA$ 来限制其更新，其中 $B \in \mathbb{R}^{d\times r}$ ， $A \in \mathbb{R}^{r\times d}$ ，且秩 $r \ll d$ 。在训练期间， $W_{0}$ 被冻结，不接收梯度更新，而 $A$ 和 $B$ 包含可训练参数。请注意， $W_{0}$ 和 $\Delta W = BA$ 都与相同的输入相乘，它们各自的输出向量按坐标求和。对于 $h = W_{0} x$ ，我们修改后的前向传播产生： $h = W_0x + \Delta W x = W_0x + BAx$

本质上，我们冻结原始模型，在相关权重矩阵下插入低秩适配器，并训练这些适配器以模拟通常来自梯度的更新。有了这些概念和上述公式，您现在应该理解所提供的插图了。

最显著的好处是内存和存储使用量的减少。对于使用 Adam 训练的大型 Transformer，如果 $r \ll d$ ，我们可以将 VRAM 使用量减少多达 2/3，因为我们不需要存储冻结参数的梯度和优化器状态。与完全微调相比，我们还观察到 GPT-3 175B 的训练速度提高了 25%，因为我们不需要计算绝大多数参数的梯度。

3.2 QLoRA

QLoRA 使用 4 位量化来压缩预训练语言模型。然后冻结 LM 参数，并以低秩适配器的形式向模型添加相对少量可训练参数。在微调过程中，QLoRA 通过冻结的 4 位量化预训练语言模型将梯度反向传播到低秩适配器中。LoRA 层是训练期间唯一更新的参数。

QLoRA 具有一个用于基模型权重的存储数据类型（通常是 4 位 NormalFloat）和一个用于执行计算的计算数据类型（16 位 BrainFloat）。QLoRA 将权重从存储数据类型反量化到计算数据类型以执行前向和后向传播，但只计算使用 16 位 bfloat 的 LoRA 参数的权重梯度。权重只在需要时才解压缩，因此在训练和推理期间内存使用量保持较低。

让我们更详细地探讨这一点。首先，我们来看看论文作者使用的量化方法。正如我们从上一节中记住的，有许多不同的方法。

分块 k 位量化。 量化是将输入从包含更多信息的表示离散化到包含更少信息的表示的过程。它通常意味着将具有更多位的数据类型转换为更少位的数据类型，例如从 32 位浮点数转换为 8 位整数。为了确保充分利用低位数据类型的整个范围，输入数据类型通常通过对输入元素的绝对最大值进行归一化来重新缩放到目标数据类型范围，这些输入元素通常结构化为张量。例如，将 32 位浮点 (FP32) 张量量化为范围为 [-127, 127] 的 Int8 张量： $\mathbf{X}^{\text{Int8}} = \text{round}\left(\frac{127}{\text{absmax}(\mathbf{X}^{\text{FP32}})} \mathbf{X}^{\text{FP32}}\right) = \text{round}(c^{\text{FP32}} \cdot \mathbf{X}^{\text{FP32}}),$ 其中 c 是量化常数或量化尺度。反量化是逆过程： $\text{dequant}(c^{\text{FP32}}, \mathbf{X}^{\text{Int8}}) = \frac{\mathbf{X}^{\text{Int8}}}{c^{\text{FP32}}} = \mathbf{X}^{\text{FP32}}$

如我们所见，作者通过将矩阵分解成许多小块来解决我们之前讨论的离群值的重要问题，从而最大限度地减少单个量化块内的潜在方差。

此外，为了充分理解 QLoRA 的工作原理，我们还需要考虑两个更重要的概念。

双重量化。我们引入了双重量化（DQ），这是对量化常数进行量化以进一步节省内存的过程。虽然精确的 4 位量化需要较小的块大小（因为存在离群值），但它也会带来相当大的内存开销。例如，对于 W，使用 32 位常数和 64 的块大小，量化常数平均每个参数增加 32/64 = 0.5 位。双重量化有助于减少量化常数的内存占用。

这会将每个参数的内存占用从 0.5 位减少到 0.127 位。

Normal Float 4 (NF4) 利用预训练神经网络权重通常具有零中心正态分布的事实，该技术允许从 fp32 到 int4 进行更具信息性的映射，同时考虑到接近 0 的密度增加。

现在，我们准备好理解整个 QLoRA 过程（公式中的 L1 和 L2 对应于 B 和 A）。

QLoRA。使用上述组件，我们将带有一个 LoRA 适配器的量化基础模型中的单个线性层定义为 QLORA，如下所示： $\mathbf{Y}^{\text{BF16}} = \mathbf{X}^{\text{BF16}} \text{doubleDequant}(c_1^{\text{FP32}}, c_2^{k\text{-bit}}, \mathbf{W}^{\text{NF4}}) + \mathbf{X}^{\text{BF16}} \mathbf{L}_1^{\text{BF16}} \mathbf{L}_2^{\text{BF16}}$ ，其中 doubleDequant(·) 定义为： $\text{doubleDequant}(c_1^{\text{FP32}}, c_2^{k\text{-bit}}, \mathbf{W}^{k\text{-bit}}) = \text{dequant}(\text{dequant}(c_1^{\text{FP32}}, c_2^{k\text{-bit}}), \mathbf{W}^{4\text{bit}}) = \mathbf{W}^{\text{BF16}}$

对于参数更新，只需要适配器权重 $\frac{\partial E}{\partial \mathbf{L}_i}$ 的误差梯度，而不是 4 位权重 $\frac{\partial E}{\partial \mathbf{W}}$ 。然而， $\frac{\partial E}{\partial \mathbf{L}_i}$ 的计算需要通过将存储 $\mathbf{W}^{\text{NF4}}$ 数据类型反量化为计算数据类型 $\mathbf{W}^{\text{BF16}}$ 来计算 $\frac{\partial \mathbf{X}}{\partial \mathbf{W}}$ 的导数，该过程在 BFloat16 精度下进行。

总而言之，QLORA 具有一个存储数据类型（通常是 4 位 NormalFloat）和一个计算数据类型（16 位 BrainFloat）。我们将存储数据类型反量化为计算数据类型以执行前向和后向传递，但我们只计算使用 16 位 BrainFloat 的 LoRA 参数的权重梯度。

QLORA 将微调 65B 参数模型所需的平均内存从 >780GB GPU 内存减少到 <48GB，而与 16 位完全微调基线相比，运行时或预测性能没有下降。这标志着 LLM 微调可访问性的重大转变：现在最大的公开可用模型可以在单个 GPU 上进行微调。

4. 附加技术

4.1 Flash Attention

Transformer 架构的扩展受到自注意力机制的严重瓶颈，该机制具有二次时间复杂度和内存复杂性。加速器硬件的最新发展主要集中于提高计算能力，而不是内存和硬件之间的数据传输。这导致注意力操作存在内存瓶颈。

标准注意力机制使用高带宽内存（HBM）来存储、读取和写入键、查询和值。HBM 内存大，但处理速度慢；而 SRAM 内存小，但操作速度快。在标准注意力实现中，从 HBM 加载和写入键、查询和值的成本很高。它将键、查询和值从 HBM 加载到 GPU 片上 SRAM，执行注意机制的单个步骤，然后写回 HBM，并对每个注意力步骤重复此操作。

FlashAttention 是一种重新排序注意力计算并利用分块和重新计算以显著加速并减少内存使用（从序列长度的二次方降至线性）的算法。它使用分块将输入块从 HBM（GPU 内存）加载到 SRAM（快速缓存），对该块执行注意力并更新 HBM 中的输出。通过不将大型中间注意力矩阵写入 HBM，我们减少了内存读取/写入量，从而将挂钟时间加速 2-4 倍。

FlashAttention 前向传播图：通过分块和 softmax 重缩放，我们按块操作并避免从 HBM 读取/写入，同时获得正确的输出而无需近似。

对于 FP16

4.2. 梯度累积

梯度累积 是一种技术，您可以使用比机器通常能够容纳的内存更大的批次大小进行训练。这是通过在几个批次上累积梯度来实现的，并且仅在执行一定数量的批次后才更新优化器。

例如，如果梯度累积因子设置为 2，则过程如下：我们首先计算一个批次的梯度，这为我们提供了损失函数景观上的方向。我们不立即更新模型权重，而是计算下一个批次的另一个梯度，获得一个可能不同的方向。通过将这两个梯度相加，我们在损失景观中找到一条更精确的路径。为了确保最终的更新步骤正确缩放，我们将累积梯度除以批次数量，从而防止步长的人为膨胀。

当内存中只能容纳小批次大小（这可能导致过度嘈杂的更新和不稳定的训练）时，此技术尤其有用。

4.3 8 位优化器

有状态优化器会随时间维护梯度统计信息，例如过去梯度值的指数平滑和（带动量的 SGD）或平方和（Adam）。

与普通随机梯度下降相比，这种状态可以用来加速优化，但它会占用可能分配给模型参数的内存。因此，这限制了实践中可以训练的模型的最大大小。现在来看看可以训练的8 位优化器的最大模型。

如您所料，其想法是将优化器状态量化为 8 位。

为了克服由此产生的计算、量化和稳定性挑战，8 位优化器具有三个组件

分块量化：将输入张量分成更小的块，这些块独立量化，隔离异常值并更均匀地将误差分布到所有位。
动态量化：高精度量化小值和大值。
稳定嵌入层：提高了带有词嵌入的模型在优化过程中的稳定性。

有了这些组件，使用 8 位状态执行优化器更新变得很简单。在执行更新之前，8 位优化器状态被反量化为 32 位，然后这些状态被重新量化为 8 位进行存储。

8 位到 32 位的转换以元素为单位在寄存器中进行，这意味着不需要缓慢地复制到 GPU 内存或额外的临时内存来执行量化和反量化。对于 GPU 而言，这使得 8 位优化器比常规 32 位优化器快得多。

4.4 序列打包

在对大型语言模型进行全参数或参数高效微调时，由于数据管道效率低下，GPU 利用不足是一个常见问题。这是因为大多数微调数据集的序列长度分布不均匀，许多短序列和少数长序列遵循齐夫定律。

Transformer 模型只能接受固定长度的输入，因此输入必须用许多未使用的填充标记进行填充，这在两个方面是低效的：

对填充值执行的计算最终被模型输出忽略，导致 FLOPs 浪费。
微批次大小通常受包含较长序列的批次限制，因此大多数其他微批次的 GPU 内存利用不足。

序列打包是一种训练技术，将多个训练序列（示例）连接成一个长序列（包）。这消除了填充的需要，并允许在每个微批次中处理更多的标记，从而最大限度地利用 GPU 计算和 GPU 内存。

虽然预训练的序列可以天真地连接起来，但 SFT 和指令微调的情况并非如此，其中每个输入序列都应单独处理。传统的解决方案是构建一个扩展的注意力掩码来标记每个标记所属的序列 ID，并屏蔽序列之间的注意力值。

然而，这将注意力复杂度从 $\left(\sum_i s_i^2\right)$

相反，NeMo 提供了一个高度优化的序列打包版本，该版本利用 FlashAttention 和 TransformerEngine 中的可变长度注意力内核。这样，序列之间的注意力值永远不会计算，因此注意力的复杂性保持在 $\left(\sum_i s_i^2\right)$

总而言之，NeMo 的序列打包实现在 Llama 7B 上使用 Dolly 数据集提供了：

FLOPs 性能提升高达 10 倍
训练时间性能提升高达 6 倍
对模型收敛没有影响

4.5 torch.compile()

torch.compile 通过 JIT 编译 PyTorch 代码为优化内核，同时只需最少的代码更改，即可让 PyTorch 代码运行得更快。

每当您将模型包装在 torch.compile 下时，模型在执行前会经历以下步骤：

图获取：模型被分解并重写为子图。可以编译/优化的子图被展平，而不能编译的其他子图则回退到 Eager 模式。
图下沉：所有 PyTorch 操作都被分解为它们选择的后端特定内核。
图编译：所有后端内核都调用其相应的低级设备操作。

在来自不同库（例如 TIMM、TorchBench 和 Hugging Face）的 163 个开源模型上，torch.compile 在 NVIDIA A100 上提供了 30%-200% 的加速。

4.6 多查询注意力（MQA）和分组查询注意力（GQA）

多查询注意力（MQA）和分组查询注意力（GQA）是 Transformer 模型中传统多头注意力机制的改进。这些方法提高了注意力机制的效率和有效性。

MQA将所有注意力头视为一个单一组，从而降低了计算复杂性并缩短了训练时间。当模型可扩展性或计算资源有限时，MQA 尤其有用。
GQA将注意力头分组，每个组独立处理查询子集。这种方法在传统多头注意力的细节关注和 MQA 的广泛处理之间取得了平衡，增强了对细微输入数据的处理能力。

这些注意力变体提供了

降低计算负荷：这两种方法都减少了计算量，对大型模型有利。
提高处理速度：简化注意力机制可加快训练和推理速度。

5. 集体操作

在深入分布式训练之前，最好先了解多 GPU 和多节点通信中涉及的基本操作。

为此，我们将重点关注 NVIDIA NCCL

NVIDIA 集体通信库 (NCCL) 实现了针对 NVIDIA GPU 和网络优化的多 GPU 和多节点通信原语。NCCL 提供诸如 all-gather、all-reduce、broadcast、reduce、reduce-scatter 以及点对点发送和接收等例程，这些例程经过优化，可在节点内的 PCIe 和 NVLink 高速互连以及节点间的 NVIDIA Mellanox 网络上实现高带宽和低延迟。

Caffe2、Chainer、MxNet、PyTorch 和 TensorFlow 等领先的深度学习框架已集成 NCCL，以加速多 GPU 多节点系统上的深度学习训练。

集体操作必须针对每个等级（因此是 CUDA 设备）调用，才能形成一个完整的集体操作。否则将导致其他等级无限期等待。

5.1 AllReduce

AllReduce 操作对跨设备的数据执行规约（例如，求和、最小值、最大值），并将结果存储在每个等级的接收缓冲区中。

在 k 个等级之间的求和 allreduce 操作中，每个等级将提供一个包含 N 个值的输入数组，并在一个包含 N 个值的输出数组中接收相同的结果，其中 out[i] = in0[i]+in1[i]+…+in(k-1)[i]。

5.2 广播

广播操作将 N 个元素的缓冲区从根等级复制到所有等级。

重要提示：根参数是等级之一，而不是设备编号，因此受不同等级到设备映射的影响。

5.3 规约

规约操作执行与 AllReduce 相同的操作，但仅将结果存储在指定的根等级的接收缓冲区中。

重要提示：根参数是等级之一（而不是设备编号），因此受不同等级到设备映射的影响。

注意：规约操作后跟广播操作，等效于 AllReduce 操作。

5.4 规约散射

规约散射操作执行与规约相同的操作，只是结果在等级之间以等大小的块进行散射，每个等级根据其等级索引获得一块数据。

规约散射操作受不同等级到设备映射的影响，因为等级决定了数据布局。

5.5 全收集

全收集操作从 k 个等级收集 N 个值到一个大小为 k*N 的输出缓冲区，并将该结果分发给所有等级。

输出按等级索引排序。因此，全收集操作受不同等级到设备映射的影响。

注意：执行规约散射后跟全收集，等效于 AllReduce 操作。

6. 分布式训练

主要有两种并行方法——数据并行和模型并行。

6.1 DP - 数据并行

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

6.2 模型并行、张量并行、流水线并行

当模型不适合设备内存时，模型并行通过垂直或水平方式在进程之间拆分模型。

6.2.1 朴素模型并行

这种方法涉及通过将特定层分配给特定 GPU，将模型层组分发到多个 GPU。当数据流经这些层时，它会移动到与该层相同的 GPU，而其他层保持不变。

在此示例中，当数据流经一个 GPU 内的层时，与常规前向传播没有区别。但是，在不同 GPU 上的层之间移动数据会导致通信开销。如果参与的 GPU 位于同一计算节点（例如，同一物理机器）上，则此复制速度很快，但如果 GPU 分布在不同的计算节点（例如，多台机器）上，则通信开销可能会大大增加。

朴素模型并行化的主要问题是，在任何给定时刻，除了一个 GPU 外，所有 GPU 都处于空闲状态，效率非常低。

6.2.2 流水线并行

PP 与朴素 MP 几乎相同，但它通过将传入批次分块为微批次并人工创建流水线来解决 GPU 空闲问题，这允许不同的 GPU 并发参与计算过程。

但这带来了大量的技术复杂性。

6.2.3 张量并行

在张量并行中，每个 GPU 处理张量的一个切片，并且只在需要时聚合完整的张量进行操作。因此，与模型并行（MP）不同，我们不必等待前一个 GPU 完成模型前一个层的处理。这允许更高效的处理和减少空闲时间。

任何 Transformer 的主要组成部分是完全连接的 nn.Linear，然后是非线性激活 GeLU。根据 Megatron 论文的表示法，它的点积部分可以写成 Y = GeLU(XA)，其中 X 是输入向量，Y 是输出向量，A 是权重矩阵。

如果我们将计算以矩阵形式查看，您可以看到矩阵乘法如何在多个 GPU 之间分割

如果我们将权重矩阵 A 按列在 N 个 GPU 之间分割，并并行执行矩阵乘法 XA_1 到 XA_n，那么我们将得到 N 个输出向量 Y_1, Y_2, ..., Y_n，这些向量可以独立地输入到 GeLU 中： $[Y_1, Y_2] = [\text{GeLU}(XA_1), \text{GeLU}(XA_2)]$

利用此原理，我们可以更新任意深度的多层感知器，而无需在最后进行任何 GPU 之间的同步，因为我们需要从分片重建输出向量。Megatron-LM 论文的作者为此提供了一个有用的插图

多头注意力层的并行化甚至更简单，因为它们由于具有多个独立的头而本身就具有并行性！

6.4 FSDP - 完全分片数据并行

FSDP 在分布式数据并行的基础上进行了扩展，不仅并行化了数据，还并行化了模型参数、优化器状态和与模型相关的梯度。具体来说——每个 GPU 只存储整个模型的一个子集以及相关的优化器状态和梯度子集。

6.4.1 FSDP 单元

FSDP 将模型实例分解为更小的单元，然后将每个单元内的所有参数展平并分片。分片参数在计算之前按需通信和恢复，然后立即丢弃。这种方法确保 FSDP 一次只需要具体化一个单元的参数，这显著降低了峰值内存消耗。

6.4.2 FSDP 工作流

让我们以包含 [layer1,layer2] 的 FSDP unit1 为例来解释这个过程。

前向传播:

在前向计算进入 layer1 之前，FSDP 通过从其他对等等级收集分片来收集 layer1 和 layer2 的未分片参数。
使用未分片参数，FSDP 运行这些层的局部计算
然后释放刚刚收集的对等分片以减少内存占用

因此，在整个前向传播过程中，FSDP 一次只需要完全实例化一个单元，而所有其他单元都可以保持分片状态。

反向传播:

类似地，在反向计算期间，FSDP unit1 在反向传播到达 layer2 之前恢复 layer1 和 layer2 的未分片参数
当自动求导引擎完成这两个层的反向计算时，FSDP 释放对等分片并启动 ReduceScatter 以规约和分片梯度。
因此，在反向计算之后，每个等级只保留参数和梯度的一个分片

整个工作流可以可视化如下

虽然 FSDP 通过参数分片显著优化了内存使用，但由于频繁需要跨 GPU 收集和分散参数和梯度，它引入了一些通信开销。这种开销是内存占用减少的权衡，其影响可能因 GPU 之间的网络带宽和延迟而异。高效实现收集和分散操作，以及优化（例如将通信与计算重叠），可以帮助减轻这种开销，以保持高训练吞吐量。

6.4.3 分片策略

分片策略是 FSDP 中的一个重要元素，对内存占用和通信开销起着重要作用。FSDP 提供多种分片策略，从完全复制到完全分片。

6.4.4 结果

FSDP 通过一系列先进技术实现了可用性和效率，包括延迟初始化、灵活的分片策略、通信重叠和预取，以及通信集合的速率限制。所有这些技术都与 PyTorch 的其他关键组件紧密协同设计，以确保解决方案的健全性和鲁棒性。评估表明，FSDP 可以促进大型语言模型的近线性可扩展性。

FSDP 是一个庞大而复杂的主题，需要完全理解。如果您对此感兴趣，最好研究原始文章，从中您可以了解更多关于工作流程的详细信息，以及模型初始化、参数分片和通信优化是如何进行的。

参考

社区

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录以评论