Flash Attention 稳定吗？
机器学习
2024-05-07 v1 分布式、并行与集群计算
摘要
训练大规模机器学习模型带来了独特的系统挑战，考虑到当今工作负载的规模和复杂性。最近，许多训练最先进生成式AI模型的组织报告了训练过程中的不稳定情况，通常表现为损失尖峰。数值偏差已成为这种训练不稳定的潜在原因，尽管由于训练运行的高成本，量化这一点尤其具有挑战性。在这项工作中，我们开发了一种原则性方法来理解数值偏差的影响，并构建代理以在下游效应难以量化时将观察结果置于上下文中。作为案例研究，我们将此框架应用于分析广泛采用的Flash Attention优化。我们发现，在BF16下，Flash Attention在孤立前向传播中测量的数值偏差大约比基线注意力高一个数量级。然后，我们使用基于Wasserstein距离的数据驱动分析，提供了这种数值偏差在训练期间如何影响模型权重的上界，发现Flash Attention中存在的数值偏差比低精度训练的重要性低2-5倍。
引用
@article{arxiv.2405.02803,
title = {Is Flash Attention Stable?},
author = {Alicia Golden and Samuel Hsia and Fei Sun and Bilge Acun and Basil Hosmer and Yejin Lee and Zachary DeVito and Jeff Johnson and Gu-Yeon Wei and David Brooks and Carole-Jean Wu},
journal= {arXiv preprint arXiv:2405.02803},
year = {2024}
}