Meta最新研究: Flash Attention 为何是系统性能瓶颈?

I. 引言

随着机器学习趋向于更大和更复杂的模型,模型训练过程变得越来越计算和资源密集。生成式AI的出现进一步推动了模型开发的边界,大型语言模型(LLMs)通常在数百或数千个GPU上训练数月。以LLaMA2的70-B参数模型为例,需要1,720,320 GPU小时来训练。对于如此长的训练作业,训练不稳定已成为日益严重的问题。正如Google的PaLM模型等工作中所报告的那样,训练不稳定通常表现为在整个训练过程中出现多达20次的损失峰值。这些损失峰值是代价高昂的,因为它们经常导致训练过程中断,需要停止并重新开始训练。

虽然之前的工作从算法角度研究了可能的缓解响应以提高训练稳定性,但这种不稳定的根本原因仍不清楚。模型训练的随机性,再加上许多这种效应只有在大规模下才会表现出来,给完全理解这些不稳定性的性质带来了独特的挑战。此外,考虑到数据中心计算的成本和高需求,重复训练如此大的模型以隔离影响因素通常是不可行的。

训练不稳定的一个尚未探索的潜在原因是数值偏差。优化与其相应基线之间的数值偏差可能导致误差的逐渐累积,在训练过程中有可能导致损失峰值,需要重置模型状态。这很难量化,因为训练的随机性表明某种程度的数值偏差可能是可以接受的,但确定训练变得不稳定的阈值却很困难。

在这项工作中,Meta研究开发了一种量化训练优化中数值偏差的原则性方法。该方法包括两个阶段:(i)开发微基准以扰动给定优化中的数值精度,以及(ii)通过基于Wasserstein距离的数据驱动分析评估数值偏差如何转化为模型权重的变化。这最终允许该研究为给定优化提供数值偏差的上限,并有助于将改进置于已知技术的背景下。该研究旨在使用这种原则性分析来评估不同的最先进优化技术,并确定当用于训练大型模型时,它们是否可能引入意外的不稳定性。

作为案例研究,该研究分析了最先进的优化技术Flash Attention,并量化了可能引入的数值偏差。Flash Attention是一种广泛采用的加速Attention机制的技术,通常被认为是Transformer模型的系统瓶颈。然而,虽然提供了更高的加速和更少的内存访问,但Flash Attention依赖于可能导致数值折衷的算法优化。具体而言,该研究假设重缩放因子的引入可能导致无意的近似,从而导致数值折衷,这可能在后期影响训练稳定性。该研究在多模态文本到图像工作负载的背景下分析Flash Attention,以确定Flash Attention和基线之间的数值偏差的潜在意义。

最终,该研究引入了一个框架来量化训练优化的数值偏差及其下游影响。主要贡献如下:

该研究设计了一个微基准来隔离数值精度对数值偏差的影响。该微基准可作为衡量和量化传统上不透明的优化(如Flash Attention)引起的数值偏差的技术。通过扰动通常无法通过提供的内核获得的方面,该研究最初发现,在低数值精度(BF16)下,与基线Attention相比,Flash Attention的数值偏差大约多一个数量级。
该研究进行了基于Wasserstein距离度量的数据驱动分析,以将观察到的数值偏差置于上下文中,并为其对下游模型属性的影响形成上限。在该案例研究中,该研究能够限制这种观察到的数值偏差的影响,发现与低精度训练相比,Flash Attention引入的模型权重偏差大约小2-5倍。

该研究的调查强调了开发原则性方法的重要性,不仅要量化训练优化对数值偏差的影响,还要将其置于上下文中。通过构建代理将这种数值偏差置于上下文中,该研究旨在推理下游模型效应(即训练不稳定)的可能性,而这在传统上很难衡量。

II. 背景

作为本工作的案例研究,该研究分析了最先进的优化Flash Attention及其与基线Attention的潜在数值偏差。Attention操作最近一直是无数优化的焦点。Attention目前是Transformer架构的主要系统性能瓶颈,Transformer是现代机器学习算法中广泛采用的一种技术,以其在序列到序列任务建模方面的有效性而著称。

A. 作为系统性能瓶颈的Attention

Attention操作本质上得出一个加权和,表示在生成新token时,模型应该对所有先前的单词放置多少重点。Attention中的关键计算涉及矩阵乘法和softmax。首先从输入向量中得出三个表示(Query、Key、Value),每个表示的维度为 $\times d$ ,其中 $N$ 是序列长度, $d$ 是模型维度。然后,Query和Key值的点积形成一个 $\times N$ 矩阵,其大小随序列长度呈二次缩放。这个所谓的相似性矩阵随后经历softmax操作,然后与Value矩阵相乘以完成计算。完整的操作如下所述:
$softmax(\frac{Q \dot K^T}{\sqrt{d_k}})V$
考虑到这种随序列长度的二次缩放,人们已经提出了无数加速Attention机制的技术,包括系统感知方法,如Flash Attention。

B. 理解Flash Attention

Flash Attention是最近提出的一种旨在加速Transformer特有的Attention瓶颈的技术。作为一种IO感知技术,它旨在最小化Attention机制中通常使用的大型 $\times N$ 相似性矩阵的内存开销。它本质上使用了传统的分块和重计算技术,以及在线softmax技巧,以便一次只计算矩阵的一个分块。如图1所示,分块的引入消除了实现大型相似性矩阵的需要。块/分块大小由 $B_c = \lceil M/4d \rceil$ 和 $B_r = \min(\lceil M/4d \rceil, d)$ 定义,其中 $M$ 是SRAM的大小, $d$ 是模型维度。这确保块适合SRAM,从而消除了需要从HBM加载/存储大矩阵。

在这里插入图片描述

然而,由于在线softmax技术需要有关矩阵的全局信息,Flash Attention必须合并重缩放因子,以便在计算单个块时允许一致的计算,因为全局信息不再可用。虽然引入了相对较小的开销,但这些额外的重缩放因子确实引入了每个分块计算的额外计算。

Flash Attention带来了时序性能加速以及更高效的资源利用;该研究发现,对于示例文本到图像模型,它在前向+后向传递方面产生了14%的加速。然而,人们还假设Flash Attention引入的额外计算在用于文本到图像模型训练时可能引入数值偏差,该研究旨在调查这一点。

III. 实验方法

该研究首先开发一个微基准来隔离和研究Flash Attention引起的数值偏差。微基准设计的摘要可以在图2中找到。如图所示,该研究在数值上重新实现了Flash Attention,以便分析不同的数值精度并在算法的每一步应用潜在的优化,这在原始CUDA实现中不容易完成。这是必要的,因为Flash Attention内核目前只支持FP16和BF16数字格式。

在这里插入图片描述

Flash Attention内核也是CUDA代码的包装API调用,这使得扰动算法以检查对数值偏差的影响变得具有挑战性。相比之下,该研究的微基准设计允许使用不同精度的输入和算法内的修改。该研究根据原始Flash Attention内核验证了微基准。

该研究进一步设计了一种技术来比较模型执行期间每一步注意力矩阵的输出。该研究修改了模型代码,以便在每次调用Attention时计算基线Attention和Flash Attention,这允许输入矩阵相同且精确比较输出矩阵。为了将其置于上下文中,该研究还通过相同和独立的训练运行,使用 $ma x d i ff ere n ce$ 和 $Wa ssers t e in D i s t an ce$ 指标量化了整个训练过程中的模型权重差异,这在第V节中有进一步详细说明。

对于训练实验,该研究使用了一种将文本输入转换为图像的生成式AI工作负载(即文本到图像模型)。该研究使用Shutterstock数据集重新训练模型,并在NVIDIA 80GB A100 GPU集群上运行实验。

IV. 通过微基准量化数值偏差

该研究首先分析Flash Attention在前向传递期间的影响。该研究利用微基准来检查不同的数值精度如何影响Attention计算的输出矩阵,给定相同的随机初始化查询、键、值向量。

A. 扫描数值精度

该研究发现,数值精度对Flash Attention的输出有影响,导致它偏离基线Attention的输出。该研究通过计算Attention输出矩阵之间的最大差异来量化这一点,该差异以元素方式计算,作为可能偏差的上限。如图3所示,当使用不同的数字格式(从BF16到FP64),Flash Attention和基线Attention之间的数值偏差随着尾数位数的增加而减小。这表明数值差异是由较少尾数位固有的近似引起的。

然后,该研究将其与基线Attention的行为进行比较。为了进行通用的比较,该研究将Attention的"黄金值"设置为FP64的基线Attention。然后,该研究将不同数字格式下Attention输出的最大差异与这个黄金值进行比较,如图4所示。请注意,该研究绘制了Flash Attention输出与该黄金值之间的最大差异(蓝色条),同时将基线Attention输出与该黄金值进行比较(红色条)。该研究发现,在BF16下,与基线Attention相比,Flash Attention的数值偏差大约多10倍。关于这种偏差水平是否显著的详细讨论可以在第V节中找到。

为了进一步分析这种观察到的数值偏差,该研究在保持分块大小和SRAM大小相同的情况下扫描矩阵的序列长度。如图5所示,随着序列长度的增加,通过(a)最大差异上限和(b)该差异的均值和标准差这两个指标测量,Flash Attention和基线Attention之间的数值偏差都在增加。由于更大的序列长度意味着必须分块的更大的 $\times N$ 中间矩阵,而分块大小保持不变,因此需要更多的重缩放。这提供了更多精度误差累积的机会,从而导致更多偏差。

B. 扫描算法变化

该研究进一步利用微基准设计来尝试不同的优化,以便更好地理解这种数值偏差的影响。图6显示了几种算法变化及其对观察到的数值偏差的相应影响。对于每个实验,该研究扫描块面积(根据第II节中引入的 $B_c$ 和 $B_r$ 维度定义),并绘制Attention矩阵输出之间相应的最大差异。该研究强调了在第IV-A节分析的四个精度下,这种差异如何变化。图6a显示了交换块维度的顺序如何导致Flash Attention和基线Attention之间更大的数值差异。

在这里插入图片描述

值得注意的是,该研究还发现,较大的块/分块大小导致较小的数值偏差(图6c)。这是因为较大的分块大小需要较少的重缩放计算,因为需要较少的分块来覆盖原始的 $\times N$ 矩阵。请注意,其他扰动,如将分块大小限制为方形,对数值偏差没有影响,因为这不会显著改变需要执行的重缩放计算的数量(图6b)。

V. 通过权重差异将数值偏差置于上下文中

虽然Flash Attention可能导致前向传递期间Attention输出的数值偏差,但该研究的最终目标是确定这是否对模型训练有任何影响,以调查它是否导致训练不稳定。因此,该研究旨在量化Flash Attention是否改变了训练期间的模型——即第IV节中观察到的Attention输出差异是否反映在训练期间更新的模型权重中。

该研究利用两个指标来衡量与使用基线Attention相比,使用Flash Attention训练的模型的模型权重差异。该研究首先通过找到权重矩阵之间差异的绝对值并取最大值来计算 $ma x d i ff ere n ce$ ,以给出偏差的上限,如下所示:

$torch.max(torch.abs(flash\_attn - baseline\_attn))$

虽然 $ma x d i ff ere n ce$ 提供了数值偏差的上限,但它没有考虑每个矩阵的分布。因此,该研究还通过Wasserstein距离量化权重差异,这是一种常用的度量张量相似性的指标。虽然计算复杂度稍高,但Wasserstein度量结合了关于张量分布形状的信息,以衡量相似性。Wasserstein距离的公式如下:

$\inf_{g \in \Pi(P,Q)} E_{(x,y) \sim g}[||x-y||]$

请注意,较低的值表示矩阵之间的相似性较高。

使用这两个指标,该研究随后量化了当实现Flash Attention时,与基线Attention相比,模型权重在整个训练过程中如何变化。如图7所示,无论是通过Wasserstein距离还是最大差异衡量,Flash Attention的合并确实会在整个训练过程中改变模型权重,并且随着训练的继续,这种差异只会增加。这表明使用Flash Attention训练的模型收敛到与使用基线Attention训练的相同模型不同的模型。

在这里插入图片描述

然而,训练是一个随机过程,一些模型架构的变化可以在下游效果和准确性方面产生可比的结果。因此,即使使用Flash Attention和基线Attention训练的模型权重不同,这是否重要?将模型训练到完成并评估准确性是一项代价高昂且资源密集的任务,尤其是对于训练需要数月的大型模型。因此,该研究制定了一个代理来理解(a)这些权重变化有多重要?以及(b)我们能否将其置于其他广泛采用的训练优化中的标准权重变化的背景下?

为了实现这一点,该研究设计了一系列实验来比较在不同情况下训练过程中权重差异如何变化。除了比较使用Flash和基线Attention的训练运行外,该研究还量化了在训练开始时权重初始化为不同随机值的相同训练运行的权重差异。这提供了一个界限,因为随机权重初始化是一种常用技术,通常会产生等效的结果。此外,该研究还测量了使用不同精度训练的模型的模型权重变化。数值精度(即FP16与FP32)有可能导致下游变化,这可作为确定Flash Attention权重重要性的上限。

在这里插入图片描述

图8显示了使用Wasserstein距离度量测量的训练过程中的相对权重差异。该研究发现,使用Flash Attention的模型的权重偏差变化率与不同模型初始化的偏差相当或更小(注意红色和蓝色曲线的斜率)。此外,该研究看到,使用FP16与FP32时的权重变化率高于不同模型初始化的变化率,并且更加可变。这些结果提供了一个代理,表明尽管Flash Attention存在数值偏差,但它受随机模型初始化和低精度训练的限制,并且与低精度训练相比,引入的模型权重偏差大约小2-5倍。

VI. 讨论与未来工作

该研究朝着回答"Flash Attention是否稳定"这个问题迈出了第一步,但在得出确定的答案之前,还有更多工作需要完成。由于训练不稳定性很难且成本高昂,该研究探讨了数值偏差,这被认为是潜在的原因。通过原则性的数值偏差分析,该研究通过开发框架来量化数值偏差并开发代理以根据模型权重界定这种偏差的影响,从而朝着这一目标取得了进展。然而,最终将这种数值偏差与训练不稳定性联系起来需要进一步调查。

该研究的数值量化方法开启了一系列更广泛的问题,包括了解各种其他优化如何影响数值偏差。尽管这项分析侧重于Flash Attention,但未来的工作应该旨在扩大范围,并调查额外的训练优化及其与适当基线的相应数值偏差。例如,调查Winograd算法相对于传统卷积引起的数值偏差,或各种其他Attention优化、内核融合技术等。

更广泛地说,这项工作提出了一组更大的研究问题,值得关注,特别是关于训练不稳定性。例如,调查训练中损失峰值发生的确切点,然后将其与此时测量的权重偏差联系起来,这将很有趣。总的来说,考虑到这些实验的成本高昂,进一步开发代理来理解训练不稳定性至关重要。探索以下方面也会很有趣:

训练不稳定性和可靠性。训练不稳定性会导致训练中断,通常表现为损失峰值。然而,这并不是中断的唯一原因。例如,硬件故障也会导致数据中心的模型训练启动/停止。未来的工作应该调查这种硬件可靠性、检查点和不稳定性之间的关系。
训练不稳定性和系统开销。该研究还有兴趣量化损失峰值发生后重新排队训练作业所带来的系统开销。随着这些不稳定性而来的额外开销只会在大规模训练时放大,因此量化它们很重要。
训练不稳定性和可持续性。训练程序中的频繁中断也会导致数据中心的重大功率波动。例如,中断会在训练恢复时导致功耗急剧上升。这对于设计数据中心的低碳、节能基础设施具有可持续性影响。

VII. 结论

在这项工作中,该研究开发了一种原则性方法来理解数值偏差的影响,并开发代理以在下游效应难以衡量时将观察结果置于上下文中。该研究以Flash Attention为案例研究,在量化数值偏差方面取得了进展。该研究希望通过分享方法,能够让其他人以类似的方式调查未来的研究问题,并鼓励其他人调查这一具有挑战性的训练不稳定性问题。

参考论文: https://arxiv.org/pdf/2405.02803

我的观感:

这篇论文研究的问题非常有意义和前沿。随着深度学习模型的不断发展,训练效率和稳定性已经成为亟待解决的关键问题。论文聚焦Flash Attention引入的数值偏差及其对模型训练稳定性的潜在影响,切入点新颖独特。
作者提出了一套行之有效的研究方法,包括:

设计微基准,通过扰动Flash
Attention算法中的数值精度、分块大小等,来分离研究数值偏差的影响因素,这种做法可以清晰地界定性能瓶颈。
基于Wasserstein距离等指标,将Flash
Attention引入的数值偏差置于不同训练技术的背景下考察,如低精度训练、随机初始化等。通过横向比较,可以更好地理解数值偏差的相对影响。
系统地实验和分析,包括改变序列长度、算法细节等,并在真实的文本到图像任务上进行评估,保证了结论的可靠性。

这些方法对于类似的系统优化研究具有很好的借鉴意义。
另一方面,论文目前还存在一些局限和改进空间: