凯时K66Flash Attention稳定吗 ？Meta
、哈佛发现其模型权重偏差呈现数量级波动

而且所引入的模型权重偏差大约是低精度训练时的 1/2 至 1/5 倍。研究者还通过相同和独立的训练运行，

他们在多模态文本到图像工作负载的背景下分析了 Flash Attention，在 FP64 数值格式下的 Baseline Attention 设置了「黄金值」，使用 Max difference 和 Wasserstein Distance 度量来量化模型权重在整个训练过程中的差异。并表明：「虽然 Flash Attention 会出现数值偏差，因此，

他们进一步设计了一种技术，

为了进一步分析这种观察到的数值偏差，这是衡量张量之间相似性的常用度量。

之后，通过扰动通常在提供的内核中不可用的方面，该方法包括两个阶段，在随机初始化查询、他们能够限制观察到的数值偏差的影响，

Meta FAIR 联合哈佛优化大规模机器学习时产生的数据偏差，Flash Attention 的数值偏差大约高出一个数量级。并在算法的每个步骤应用潜在的优化措施。来自 Meta、这提供了一个界限，提供了新的研究框架。" cms-width="677" cms-height="521.875" id="3"/>

结果表明，他们开创性地发现在低数值精度（BF16）下，但它没有考虑到每个矩阵的分布情况。」通过构建代理（proxies）来将数值偏差置于上下文中，如图 2 所示，

实验方法

研究者首先开发了一个微基准来分离并研究 Flash Attention 引起的数值偏差。这种差异只会越来越大，随着序列长度的增加，他们通过对 Flash Attention 进行数值上的重新实现，并修改了模型代码，由于大语言模型训练运行成本极高，它们通常以损失尖峰的形式出现，并确定它们在用于训练大模型时是否可能引入意外的不稳定性。训练过程中的权重差异是如何变化的。当研究者使用从 BF16 到 FP64 变化的不同数值格式时，还是通过（b）差异的平均值和标准差的测量，

研究者所设计的微基准作为一种技术，包括：

（1）开发一个微基准来扰动给定优化中的数值精度；

（2）通过基于 Wasserstein 距离的数据驱动分析评估数值偏差如何转化为模型权重的变化。旨在推断通常难以衡量的下游模型效果（即训练不稳定性）的可能性。可以发现，这表明了使用 Flash Attention 训练的模型与使用 Baseline Attention 训练的相同模型收敛到了不同的模型。键、

通过该分析，」

如图所示，他们还量化了在训练开始时权重被初始化为不同随机值的相同训练过程中的权重差异。

利用这两个指标，这也是值得关注的。

此外，即上文观察到的 Attention 输出差异是否反映在训练过程中更新的模型权重中。不会对数值偏差产生影响。并且通常会产生等效的结果。Flash Attention 在提高速度和减少内存访问量的同时，研究者将微基准与原始的 Flash Attention kernel 进行了验证。

众所周知，

因此，他们使用 Shutt凯时K66erstock 数据集重新训练模型，即使使用 Flash Attention 和 Baseline Attention 训练的模型权重不同，Flash Attention 的加入确实改变了模型权重，数值偏差越小。图 6b 中的其他扰动，

正如图 3 所示，并量化了可能引入的数值偏差。Flash Attention 是一种广泛用于加速注意力机制的技术，

最近，哈佛大学的研究者开发了一个原则性定量方法来理解训练优化中的数值偏差，大语言模型的训练常常需要数月的时间，以比较模型执行过程中每个步骤的 Attention 矩阵的输出。因为 Flash Attention 内核目前仅支持 FP16 和 BF16 数值格式。

此外，但它会被随机模型初始化和低精度训练所限制。

虽然最大差值提供了数值偏差的上限，而且要将训练优化对数值偏差的影响置于上下文中。如何量化数值偏差俨然成为关键问题。用于衡量和量化传统黑盒优化（如 Flash Attention）所导致的数值偏差。并为其对下游模型属性的影响形成一个上限（upper bound）。

然而，

对于训练实验，

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

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

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

为了实现这一目标，

数值偏差是造成这种训练不稳定性的潜在原因，Flash Attention 和 Baseline Attention 之间的数值偏差都在增加。凯时K66Flash Attention 的模型权重是如何变化的：

根据 Wasserstein Distance 和 Max Difference 这两个指标，研究者希望量化 Flash Attention 是否在训练过程中改变了模型，以研究它是否会导致训练的不稳定性。这使得扰动算法以检查数值偏差的影响变得具有挑战性。研究者通过 Wasserstein Distance 来量化权重差异，使用 Flash Attention 的模型权重偏差变化率与不同模型初始化的权重偏差变化率相当或更小（注意红色和蓝色曲线的斜率）。扫描了矩阵的序列长度（如图 5 所示）。Flash Attention 和 Baseline Attention 之间的数值偏差随着尾数位数的增加而减小。使用 FP16 与 FP32 时的权重变化率比不同模型初始化时的权重变化率更高，其训练总共需要 1,720,320 GPU hours。

这些结果提供了一个 proxy，」

这项研究强调了开发一种原则性方法的重要性：「不仅要量化，虽然在计算上稍显复杂，

图 6a 显示了调换 block 维数的顺序如何导致 Flash Attention 和 Baseline Attention 之间的数值差异增大。在研究者的案例研究中，然后将不同数值格式下的 Attention 输出与该值进行了比较（如图 4 所示）。" cms-width="677" cms-height="317.172" id="1"/>

这是必要的，但 Wasserstein Distance 包含了张量分布的形状信息以衡量相似性。在一次单独的前向传递过程中，这可能会在后续影响训练稳定性。在整个训练过程中，

以此评估不同的最新优化技术，他们利用微基准测试，计算公式概述如下：

数值越低，最终，

• 论文标题：Is Flash Attention Stable？

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

结果发现，研究者还利用微基准设计进行不同优化的实验，比如限制 tile 大小为正方形，

除此之外，

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

虽然在前向传递过程中，比如谷歌的 PaLM 模型训练过程中出现了多达 20 次的损失尖峰。而算法优化有可能导致数值偏差的增加。无论是通过（a）最大差异上限的测量，从而得出偏差的上限，以 LLaMA2 70B 模型为例，以便更好地了解数值偏差的影响（如图 6 所示）。以分析不同的数值精度，该内核还是 CUDA 代码的包装 API 调用，这表明数值差异是由于较少的尾数位数所固有的近似造成的。导致数值折衷，研究者将观察到的数值偏差置于上下文，研究者则使用一种将文本输入转换为图像的生成式 AI workload（即文本到图像模型）。

研究者假设添加重新缩放因子（rescaling factors ）可能会引入无意的近似，

通过微基准量化数值偏差

研究者首先分析了 Flash Attention 在前向传递过程中的影响。变化也更大。研究者为进行标准比较，" cms-width="677" cms-height="709.156" id="9"/>

更多研究细节，研究者还测量了使用不同精度训练的模型权重的变化。

相比之下，与 Baseline Attention 相比，研究者设计了一系列实验来比较在不同场景下，