Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速

发布于 2024-11-1 15:37
浏览
0收藏

一、背景

本文中我们继续介绍一个 Intel 最新的关于 FP8 训练相关的工作,其在一定程度上分析并解决了 FP8 训练中的不收敛问题,进一步推进了 FP8 训练落地(尤其是在 H100/H800 GPU 上)的可行性。

对应的论文:[2409.12517] Scaling FP8 training to trillion-token LLMs [1]

二、摘要

本文中,作者首次在 2T Token 的数据集上使用 FP8 精度训练了 LLM,比以前的限制增加了 20 倍。通过这些扩展训练实验,作者发现了 FP8 训练中的关键不确定性,这些不确定性在早期持续时间较短的训练中是无法观察到的。作者进一步追溯到 SwiGLU 激活函数的异常值放大问题。有趣的是,从分析和经验上表明,这些放大只发生在较长的训练中,并将其与 SwiGLU 的权重对齐过程联系起来。

为了解决这个新发现的问题,作者引入了 Smooth-SwiGLU,这是一种新颖的修改,可确保稳定的 FP8 训练而不改变函数的行为。作者还首次演示了两个 Adam 优化器参数(一阶矩和二阶矩)的 FP8 量化。

结合这些创新,作者在 256 个 Intel Gaudi2 加速器上使用 FP8 精度成功训练了一个 7B 参数量的模型,实现了与 BF16 基线相当的结果,同时提供了高达 34% 的吞吐提升。

三、引言

3.1 浮点数值表示

我们之前的文章中提到过,虽然都是 FP8 精度,但是不同硬件上存在不同的表示方式,主要包括 E5M2 和 E4M3,其中 E 表示指数位(决定了动态范围),M 表示尾数位(决定了表示精度)。此外,虽然都是 E5M2 或者 E4M3,不同的硬件可能采用不同的格式。比如 NVIDIA GPU 上的 E5M2 符合 IEEE 754 Style,而 E4M3 却不符合 IEEE 754 Style,可以称为 ARM-Intel-Nvidia Style。此外,AMD-Graphcore-Qualcomm 的表示也有所不同。如下图所示,IEEE 754 Style 的 E4M3 范围为 [-240, 240],而 ARM-Intel-Nvidia Style 的 E4M3 范围是 [-448, 448]: 

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

PS:由于 Intel 和 NVIDIA 采用一致的表示方式,这也就意味着 Intel 的结论能够很容易扩展到 NVIDIA 的 GPU。

3.2 SwiGLU

在 [2002.05202] GLU Variants Improve Transformer [2] 中作者提出了在 Transformer 模型中使用各种 GLU 变体激活,也提到其在许多下游理解任务上获得了更好的结果。然而,作者也提到并没有解释为什么这些修改会有效,将其归功于上帝的恩赐。在后续 Google 的 PaLM,Meta 的 LLaMA 系列等模型中广泛采用了 SwiGLU 激活。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

如下图所示为 FFN 中应用 SwiGLU 的公式:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

其中 Swish 为激活函数,可以表示为:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

其 Pytorch 的实现也很简单,如下所示,这里用 silu 代替了 SwiGLU,对应 β 为 1:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

因为这里有三个参数:w1,w2,w3,为了保证总参数量和正常 FFN 一致,所以 LLaMA 中这里的 Hidden Dim 不是 4d,而是 4d*2/3=8d/3。

3.3 GLU 类激活的离群点

如下图 Figure 1 所示,在 [2405.14428] Mitigating Quantization Errors Due to Activation Spikes in GLU-Based LLMs [3] 中作者发现,各种 GLU 变体的激活函数容易在特定层(比如基于 SwiGLU 激活的 FFN 的最后一个 Liner 层的输入)出现激活的 Spike。此外,作者发现这些激活的 Spike 与中间层隐藏状态(Hidden Stage,每个 Transfomer Block 的输出)之间也存在高度相关性。并且 FFN 可能会通过残差连接中的加法运算放大 Hidden Stage。一旦 Hidden Stage 被放大,它就会在各层中持续存在,直到之后的层中再次遇到激活 Spike。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

3.4 延迟缩放

在 FP8 的 Per Tensor Scaling 技术中,有两种常见的方式:Just-in-time Scaling 和 Delayed Scaling(可以参考 NVIDIA Transformer Engine 中的实现 Using FP8 with Transformer Engine [4])。

  • Just-in-time Scaling(实时缩放):直接计算 Tensor 绝对值的最大值(amax),然后得到 Scaling 值,再对 Tensor 进行 Scaling。此种方法更加精确,但是,额外引入的开销会大幅降低 FP8 带来的收益
  • Delayed Scaling(延迟缩放):核心思路是使用额外的 Tensor 来存储之前的 amax 历史,然后根据历史最大值估计当前的最大值。

如下图为 NVIDIA Transformer Engine 中的 Delayed Scaling 实现方案,amax history 最多可以存储 1024 个 history。在进行当前 Tensor 的 Scaling 操作时,使用当前 Tensor 之前的 amax history 来预测当前的 amax,然后再进行 Scaling 操作;Scaling 操作的同时会计算当前的 amax,并更新 amax history。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

四、方案

4.1 洞察

上面提到的 GLU 类激活引出的离群点(Outlier)问题会为 FP8 训练带来很多挑战。本文中,作者揭示,在大规模数据集上训练 LLM 的后期阶段,这些异常值变得尤为显著。

如下图 Figure 1 所示,(a)为训练的起始阶段,没有 Outlier;(b)为训练 200B Token 之后,出现偶发性的 Outlier。这些 Outlier 仅在训练中处理了很长一段时间才出现,此现象对维持训练中的数值稳定性带来极大挑战,进一步增加了 FP8 训练稳定性的难度,尤其是像 Megatron-LM 中广泛采用的延迟缩放方案中(如上述的介绍,这些方案假设了迭代期间的一致性)。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

如前所述,SwiGLU 激活函数可能导致 FFN 组件的最后一个 Linear 层的输入出现 Outlier。在使用 FP8 训练,并采用延迟缩放技术时,SwiGLU 引发的 Outlier 会打破延迟缩放的统计一致性假设,导致训练过程的不稳定性。如下图 Figure 3 所示,作者展示了在 FFN 的最后一个 Linear (也就是 SwiGLU 的输出)禁用量化后的训练收敛性,LLaMA2 FP8 的训练能够成功地在大规模数据集上收敛,从而解决先前观察到的发散问题,这也验证了 SwiGLU 对 FP8 训练稳定性的影响。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

4.2 SwiGLU 相关问题证明

如下图所示为 SwiGLU 的定义,其中,SwiGLU 由输入 x 与权重 w1、w2 进行乘积,分别得到 xTw1 和 xTw2。然后使用 Swish 激活函数对 xTw2 进行变换,将其与 xTw1 相乘。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

其他标准激活函数(如 ReLU、GeLU 和 Swish)在输入幅度较大时最多是线性的。这意味着当输入 u 逐渐趋于正无穷或负无穷时,这些激活函数的比值(即 ∣f(u)/u∣)会趋于小于等于 1 的某个值。然而,SwiGLU 是一个二次函数,可以达到更大的值,尤其在权重 w1、w2 相互“对齐”时(例如 w1=w2,并且 ||w1||=1)。

当 w1=w2 时,上式可以表示为:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

假设 xTw=c,则当 c 较大时,σ(c) 趋近于 1,此时下述结果约等于 c2,也就是上述所说的二次放大特性。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

作者也进一步通过理论分析证明了上述 w1、w2 相互“对齐”现象,这里不再赘述,具体可以查看论文。当然,作者也进一步通过实验验证了相关问题。如下图 Figure 2 所示:

  • (a):LLaMA2-7B 模型在 BF16 和 FP8 精度下的训练损失,其中 FP8 的训练在达到 200B Token 后开始出现明显的发散现象。
  • (b):展示了某一个特定 Channel 在训练过程中 w1 和 w2 范数的动态变化及其相关性。
  • (c):某一个 Outlier 通道在训练初期(8B Token)与后期(330B Token)w1 和 w2 元素散点图。可以看出,后期相关性显著提高,趋近于 w1=w2。
  • 某一个 Outlier 通道在训练初期(8B Token)与后期(330B Token)w1 的直方图分布。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

除了上图 (c) 那样的正相关性外,作者也观察到了明显的负相关性,也就是趋近于 w1=-w2,如下图所示:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

4.3 Smooth-SwiGLU

为了在解决 Outlier 问题的同时保持完整的 FP8 加速,作者提出了 Smooth-SwiGLU 方法。如下图 Figure 4 所示,其展示了 Smooth-SwiGLU 的核心理念:对 SwiGLU 函数的线性分支施加一个缩放因子,并在最后一个 Linear 后将其重新缩放。此方法防止了输入到最后一个 Linear 的量化过程中出现 Outlier,同时保留了 SwiGLU 激活函数的整体功能,使能够在整个网络中充分利用 FP8  精度。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

为降低计算开销,作者采用了一种高效的并行方法来计算缩放因子 si,也就是 Per Channel 量化:

1. 将 Tensor 分割成若干块,每块对应一个通道。

2. 对每个块(通道),并行计算其最大值。

3. 利用这些每个通道的最大值,确定各通道的独立缩放因子 si。

此方法实现了高效的逐通道缩放,因为每个通道的缩放因子是独立并行计算的。相较于 Linear 层中的矩阵乘法,这种方法的计算成本适中,尤其是在并行化处理下,即便在非优化实现中也是如此。在推理阶段,这些缩放因子可以合并到包含 SwiGLU 层及其后 Linear 的 FFN 的第一和第三Linear 的权重中。

五、实验&结果

5.1 FP8 优化器

Adam 优化器及其变体在深度学习中得到广泛应用,Adam 优化器的一个关键特征是其存储两个矩,传统上采用高精度(FP32),这显著增加了内存开销,尤其对于大规模模型而言。尽管先前研究中已证明将其一阶矩降至 FP8 精度的可行性,但仍保留二阶矩为 FP16。本文中作者则更进一步,成功将二阶矩也量化至 FP8,显著提升了大型语言模型优化器的效率。

如下图 Figure 5 所示,作者探索发现一阶矩 Mom1 采用 E4M3,二阶矩 Mom2 采用 E5M2 可以很好的维持训练精度:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

5.2 训练效果实验

如下图 Figure 6 所示,作者在 256 个 Gaudi2 上训练 LLaMA2 7B 模型,共训练了 330B 左右 Token。本文的 Smooth-SwiGLU + FP8 优化器可以和 BF16 训练维持相当的 Loss,而传统的 FP8 训练在 200B Token 时开始发散。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

如下图 Table 2 所示,作者进一步对比了下游任务的 Zero Shot 精度以及困惑度,可以看出,本文 FP8 训练出的模型可以很好的保持精度:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

5.3 训练速度

如下图 Table 3 所示,可以看出,本文方案在保持收敛性的同时获得了更高的加速比,相比 BF16 训练可以加速 33.52%(不及 FP8 主要是因为引入了一些额外的开销)。

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

虽然额外引入了一些计算开销,但显存并没有明显增加:

Intel Smooth-SwiGLU:FP8 LLM 训练,34% 加速-AI.x社区

六、参考链接

  1. ​https://arxiv.org/abs/2409.12517​
  2. ​https://arxiv.org/abs/2002.05202​
  3. ​https://arxiv.org/abs/2405.14428​
  4. ​https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html​

​本文转载自​​AI闲谈​​​,作者: AI闲谈 ​​​

收藏
回复
举报
回复
相关推荐