自动化研究的小乐趣：

我们发现并帮助修复了一些可能影响vLLM和SGLang推理性能的边缘案例。

在上一篇博客文章中，我们介绍了为性能优化任务开发的一个奖励黑客评判器。

在将该评判器应用于一些新的自动化研究工作时，它发现FlashInfer（一个支撑vLLM和SGLang的库）的一些内核使用了硬编码的-50,000作为掩码注意力哨兵值——尽管有效的QK值可能更小。

这类数值上错误但静默的边缘案例可能带来巨大的调试难题，例如围绕闪存注意力(https://arxiv.org/abs/2405.02803)的历史性争论。

绝佳的AI应用案例。 https://github.com/flashinfer-ai/flashinfer/pull/4401…

闪存注意力是否稳定？

来源：https://arxiv.org/html/2405.02803 Alicia Golden Samuel Hsia Fei Sun Bilge Acun Basil Hosmer Yejin LeeZachary DeVito Jeff Johnson Gu-Yeon Wei David Brooks Carole-Jean Wu 隶属机构：Meta的FAIR / 哈佛大学

摘要

训练大规模机器学习模型带来独特的系统挑战，源于当今工作负载的规模与复杂性。近期，许多训练尖端生成式AI模型的机构报告了训练过程中的不稳定情况，常表现为损失尖峰。数值偏差已被视为这种训练不稳定性的潜在原因，但考虑到训练运行的高昂成本，量化这种偏差尤其困难。本研究开发了一种有原则的方法来理解数值偏差的影响，并构建了代理指标，以便在下游影响难以量化时为观察结果提供背景。作为案例研究，我们将此框架应用于分析广泛采用的闪存注意力优化。我们发现，在孤立的前向传播中以BF16测量时，闪存注意力相比基线注意力的数值偏差大约高出一个数量级。随后，我们使用基于Wasserstein距离的数据驱动分析，为这种数值偏差在训练中对模型权重的影响提供了上界，发现闪存注意力中的数值偏差比低精度训练的影响低2-5倍。

索引词：

生成式AI、数值偏差、训练不稳定性、注意力机制、Transformer模型

一、引言

随着机器学习趋向更大更复杂的模型，模型训练过程正变得越来越消耗计算资源和系统资源。生成式AI的出现进一步推动了模型开发的边界，大型语言模型（LLM）通常需要在数百或数千个GPU上持续训练数月。例如，LLaMA2的700亿参数模型就需要1,720,320个GPU小时进行训练[10 (https://arxiv.org/html/2405.02803#bib.bib10)]。在如此漫长的训练任务中，训练不稳定性变得日益棘手。如Google的PaLM模型报告所示，训练不稳定性常表现为训练过程中多达20次的损失尖峰[1 (https://arxiv.org/html/2405.02803#bib.bib1)]。这些损失尖峰代价高昂，因为它们通常导致训练过程中断，需要停止并重新启动训练。

虽然先前的工作从算法角度探讨了增加训练稳定性的可能缓解措施[6 (https://arxiv.org/html/2405.02803#bib.bib6),5 (https://arxiv.org/html/2405.02803#bib.bib5),9 (https://arxiv.org/html/2405.02803#bib.bib9),3 (https://arxiv.org/html/2405.02803#bib.bib3)]，但这种不稳定的根本原因仍不清楚。模型训练的随机性，加上许多此类影响仅在大规模下显现的事实，给全面理解这些不稳定性的本质带来了独特挑战。此外，考虑到成本和数据中心计算的高需求，反复训练这些大型模型以隔离影响因素通常是不可行的。

一个尚未充分探索的训练不稳定性潜在原因是数值偏差。优化及其对应基线之间的数值偏差可能导致错误的逐渐积累，在训练过程中有可能引发需要重置模型状态的损失尖峰[1 (https://arxiv.org/html/2405.02803#bib.bib1)]。这难以量化，因为训练的随机性意味着某种程度的数值偏差可能是可接受的，但确定训练变得不稳定的阈值却很困难。

在这项工作中，我们开发了一种有原则的定量方法来理解训练优化中的数值偏差。我们的方法包括两个阶段：(i) 开发微基准测试以扰动给定优化中的数值精度；(ii) 通过基于Wasserstein距离的数据驱动分析评估数值偏差如何转化为模型权重的变化。这最终使我们能够为给定优化的数值偏差量提供上界，并帮助将改进置于已知技术背景下进行评估。我们旨在使用这种有原则的分析来评估不同的尖端优化技术，并识别它们在用于训练大型模型时是否可能引入非预期的不稳定性。

作为案例研究，我们分析了尖端优化技术闪存注意力[2 (https://arxiv.org/html/2405.02803#bib.bib2)]，并量化了其引入的潜在数值偏差。闪存注意力是一种广泛采用的技术，用于加速注意力机制，这通常被视为Transformer模型中的系统瓶颈[11 (https://arxiv.org/html/2405.02803#bib.bib11)]。然而，尽管提供了加速和减少内存访问，闪存注意力依赖于算法优化，这些优化可能促成数值偏差的增加。具体而言，我们假设缩放因子的增加可能引入无意的近似，导致数值权衡，这可能后续影响训练稳定性。我们分析了闪存注意力在多模态文本到图像工作负载中的情况，以确定闪存注意力与其基线之间数值偏差的潜在重要性。

最终，我们引入了一个量化训练优化数值偏差及其下游影响的框架。我们的主要贡献如下：

• 我们设计了一个微基准测试，以隔离数值精度对数值偏差的影响。我们的微基准测试作为测量和量化传统黑盒优化（如闪存注意力）所导致数值偏差的技术。通过扰动通常通过内核不可用的方面，我们最初发现闪存注意力在低数值精度（BF16）下比基线注意力的数值偏差大约高出一个数量级。
• 我们进行了基于Wasserstein距离指标的数据驱动分析，以将观察到的数值偏差置于背景下，并形成对下游模型属性影响的上界。在我们的案例研究中，我们能够界定这种观察到的数值偏差的影响，并发现闪存注意力引入的模型权重偏差比低精度训练约少2-5倍。

我们的研究强调了开发一种有原则方法的重要性，不仅要量化，还要将训练优化对数值偏差的影响置于背景下。通过构建代理指标来将这种数值偏差置于背景中，我们旨在推理下游模型影响（即训练不稳定）的可能性，这些影响传统上难以测量。

二、背景

作为这项工作的案例研究，我们分析了尖端优化闪存注意力及其与基线注意力的潜在数值偏差。注意力操作一直是近期众多优化的焦点。注意力目前是Transformer架构的主要系统性能瓶颈，Transformer是现代机器学习算法中广泛采用的技术，以其在建模序列到序列任务方面的有效性而闻名[11 (https://arxiv.org/html/2405.02803#bib.bib11)]。

二-A、注意力作为系统性能瓶颈

注意力操作本质上是派生一个加权和，表示模型在生成新词元时应对所有先前词语的重视程度[11 (https://arxiv.org/html/2405.02803#bib.bib11)]。注意力的关键计算涉及矩阵乘法和softmax。三个表示（查询、键、值）首先从输入向量中派生出来，每个维度为N×d，其中N是序列长度，d是模型维度。然后，查询和键值的点积形成一个N×N矩阵，其大小随序列长度二次方缩放。这个所谓的相似度矩阵随后经过softmax操作，然后与值矩阵相乘以完成计算。完整操作描述如下：

Attention(Q,K,V)=softmax(QKᵀ/√dₖ)V

鉴于这种随序列长度的二次方缩放，已经提出了多种技术来加速注意力机制，包括系统感知方法如闪存注意力[2 (https://arxiv.org/html/2405.02803#bib.bib2)]。

二-B、理解闪存注意力

闪存注意力是最近提出的一种旨在加速Transformer特征性注意力瓶颈的技术[2 (https://arxiv.org/html/2405.02803#bib.bib2)]。作为一种IO感知技术，它旨在最小化注意力机制中常用的大N×N相似度矩阵的内存开销。它本质上使用传统的分块和重计算技术，以及在线softmax技巧，以便一次仅计算一个块的矩阵[2 (https://arxiv.org/html/2405.02803#bib.bib2)]。如图1 (https://arxiv.org/html/2405.02803#S2.F1)所示，分块的引入消除了物化大型相似度矩阵的需要。块/块的大小定义为Bc=⌈M/4d⌉和Br=min(⌈M/4d⌉,d)，其中M是SRAM的大小，d是模型维度。这确保了块能放入SRAM，从而消除了从HBM加载/存储大型矩阵的需要。

然而，由于在线softmax技术需要矩阵的全局信息，闪存注意力必须包含重新缩放因子，以确保计算的一致性，因为在单个块上进行计算时全局信息不再可用。虽然引入的开销相对较小，但这些额外的缩放因子确实引入了每个块计算的额外计算。

闪存注意力带来了时间性能加速以及更高效的资源利用率；我们发现它为我们示例的文本到图像模型带来了14%的前向+反向传播加速。然而，也有人假设，闪存注意力的缩放因子引入的额外计算可能在用于文本到图像模型训练时引入数值偏差，我们旨在下文对此进行调查。

参考图示图1：闪存注意力分块操作。闪存注意力使用分块和重计算来消除对大N×N相似度矩阵的需求。QKᵀ点积的输出改为按块计算，如绿色所示。

三、实验方法论

参考图示图2：实验方法论。(1) 我们实现了闪存注意力操作的数值微基准测试，允许进行不同数值精度的实验，以及测试算法中的各种优化。我们的框架允许直接比较基线注意力、闪存注意力和我们数值重实现的注意力矩阵输出。(2) 我们利用数据驱动的程序，通过检查训练过程中的模型权重变化，将这种数值差异置于背景下。

我们首先开发微基准测试以隔离和研究闪存注意力引起的数值偏差。我们微基准测试设计的摘要可以在图2 (https://arxiv.org/html/2405.02803#S3.F2)中找到。如图所示，我们数值重实现了闪存注意力，以分析不同的数值精度并在算法的每一步应用潜在的优化，这在原始CUDA实现中不易完成。这是必要的，因为闪存注意力内核目前仅支持FP16和BF16数字格式。闪存注意力内核也是一个包装好的CUDA代码API调用，使得扰动算法以检查对数值偏差的影响变得具有挑战性。相比之下，我们的微基准测试设计允许变化的精度输入和算法内部的修改。我们将微基准测试与原始闪存注意力内核进行了验证。

我们进一步设计了一种技术来比较模型执行期间每一步的注意力矩阵输出。我们修改模型代码，以便在每次调用注意力时同时计算基线注意力和闪存注意力，这允许相同的输入矩阵和精确的输出矩阵比较。为此背景，我们还通过相同且独立的训练运行量化了训练过程中的模型权重差异，使用最大差异和Wasserstein距离指标，这些在V节 (https://arxiv.org/html/2405.02803#S5)中有更详细的描述。

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

四、通过微基准测试量化数值偏差

我们首先分析闪存注意力在前向传播中的影响。我们利用微基准测试，在给定相同的随机初始化查询、键、值向量的情况下，检查不同的数值精度如何影响注意力计算的输出矩阵。

四-A、数值精度扫描

我们发现数值精度对闪存注意力的输出有影响，导致其偏离基线注意力的输出。我们通过逐元素计算注意力输出矩阵之间的最大差异来量化这一点，这作为可能偏差的上界。如图3 (https://arxiv.org/html/2405.02803#S4.F3)所示，当使用从BF16到FP64的不同数字格式时，闪存注意力与基线注意力之间的数值偏差随着尾数位数的增加而减少。这表明数值差异是较少尾数位固有近似的结果。

然后我们随后将其与基线注意力的行为进行比较。为了进行常见的比较，我们设置注意力的“黄金值”为FP64下的基线注意力。然后我们比较不同数字格式下注意力输出与这个黄金值的最大差异，如图4 (https://arxiv.org/html/2405.02803#S4.F4)所示。注意，我们绘制了闪存注意力输出与这个黄金值之间的最大差异（蓝色柱状图），同时将基线注意力输出与这个黄金值进行比较（红色柱状图）。我们发现闪存注意力的数值偏差大约是10倍。