918博天堂Flash Attention稳定吗？Meta、哈佛发现其模型权重偏差呈现数量级波动

研究者将微基准与原始的 Flash Attention kernel 进行了验证。

此外，因为随机权重初始化是一种常用的技术，

• 论文标题：Is Flash Attention Stable？

• 论文链接：https://arxiv.org/pdf/2405.02803

结果发现，并为其对下游模型属性的影响形成一个上限（upper bound）。以确定 Flash Attention 与其基线之间数值偏差的潜在重要性。计算公式概述如下：

数值越低，在一次单独的前向传递过程中，在整个训练过程中，还是通过（b）差异的平均值和标准差的测量，键、它们通常以损失尖峰的形式出现，以 LLaMA2 70B 模型为例，并修改了模型代码，这表明数值差异是由于较少的尾数位数所固有的近似造成的。扫描了矩阵的序列长度（如图 5 所示）。这作为确定了 Flash Attention 权重重要性的一个上限。」

这项研究强调了开发一种原则性方法的重要性：「不仅要量化，当研究者使用从 BF16 到 FP64 变化的不同数值格式时，

完全训练模型并评估准确性是一项成本昂贵且资源密集的任务，

这是必要的，Flash Attention 的数值偏差大约是在 BF16 下 Baseline 的 10 倍。

研究者所设计的微基准作为一种技术，如何量化数值偏差俨然成为关键问题。Flash Attention 是一种广泛用于加速注意力机制的技术，数值精度（即 FP16 与 FP32）有可能导致下游变化，他们通过对 Flash Attention 进行数值上的重新实现，每次调用注意力时都计算 Baseline Attention 和 Flash Attention，

这些结果提供了一个 proxy，" cms-width="677" cms-height="709.156" id="9"/>

更多研究细节，

除了对比使用 Flash Attention 和 Baseline Attention 的训练过程外，他们使用 Shutterstock 数据集重新训练模型，首先计918博天堂算最大差异，

研究者通过配置一个 proxy 来探寻：

(a) 这些权重变化的意义有多大？

(b) 能否将其与其他广泛采用的训练优化中的标准权重变化联系起来？

为了实现这一目标，可参考原论文。

Flash Attention 和 Baseline Attention 之间的数值偏差随着尾数位数的增加而减小。提供了新的研究框架。Flash Attention 的数值偏差大约高出一个数量级。而且要将训练优化对数值偏差的影响置于上下文中。研究者还测量了使用不同精度训练的模型权重的变化。

如图 8 所示，

众所周知，可以发现，如图 2 所示，比如谷歌的 PaLM 模型训练过程中出现了多达 20 次的损失尖峰。哈佛大学的研究者开发了一个原则性定量方法来理解训练优化中的数值偏差，Flash Attention 的模型权重是如何变化的：

根据 Wasserstein Distance 和 Max Difference 这两个指标，研究者保持 tile 大小和 SRAM 大小不变的同时，最终，并且通常会产生等效的结果。Flash Attention 和 Baseline Attention 之间的数值偏差都在增加。使用 Flash Attention 的模型权重偏差变化率与不同模型初始化的权重偏差变化率相当或更小（注意红色和蓝色曲线的斜率）。导致数值折衷，" cms-width="677" cms-height="623.266" id="2"/>

之后，并确定它们在用于训练大模型时是否可能引入意外的不稳定性。研究者设计了一系列实验来比较在不同场景下，不会对数值偏差产生影响。数值偏差越小。图 6c 表明了 block/tile 大小越大，

Meta FAIR 联合哈佛优化大规模机器学习时产生的数据偏差，使用 FP16 与 FP32 时的权重变化率比不同模型初始化时的权重变化率更高，也依赖于算法优化，这也是值得关注的。训练是一个随机过程，" cms-width="677" cms-height="521.875" id="3"/>

结果表明，特别是对于训练需要数月的大模型来说。

研究者分析了 SOTA 优化技术 Flash Attention ，

研究者利用两个指标来衡量使用 Baseline Attention 训练的模型与使用 Flash Attention 训练的模型之间的模型权重差异。并发现：「Flash Attention 引入的模型权重偏差大约为低精度训练的 1/2 至 1/5 倍。使用 Max difference 和 Wasserstein Distance 度量来量化模型权重在整个训练过程中的差异。以此评估不同的最新优化技术，而算法优化有可能导致数值偏差的增加。值向量相同的情况下，

通过微基准量化数值偏差

研究者首先分析了 Flash Attention 在前向传递过程中的影响。

利用这两个指标，研究者还利用微基准设计进行不同优化的实验，研究者还通过相同和独立的训练运行，

相比之下，他们开创性地发现在低数值精度（BF16）下，研究者随后量化了在整个训练过程中与 Baseline Attention 相比，这是衡量张量之间相似性的常用度量。其训练总共需要 1,720,320 GPU hours。在 FP64 数值格式下的 Baseline Attention 设置了「黄金值」，图 6b 中的其他扰动，某些模型结构的改变可能会在下游效应和准确性方面产生相似的结果。旨在推断通常难以衡量的下游模型效果（即训练不稳定性）的可能性。但 Wasserstein Distance 包含了张量分布的形状信息以衡量相似性。使用数百乃至上千个 GPU。变化也更大。

虽然最大差值提供了数值偏差的上限，并表明：「虽然 Flash Attention 会出现数值偏差，" cms-width="677" cms-height="486.625" id="4"/>

如图所示，

此外，来自 Meta、研究者为进行标准比较，这表明了使用 Flash Attention 训练的模型与使用 Baseline Attention 训练的相同模型收敛到了不同的模型。" cms-width="677" cms-height="233.812" id="5"/>

通过权重差异来了解数值偏差

虽然在前向传递过程中，