about
Is Flash Attention Stable? (No) (arxiv.org)
1 point by anewhnaccount2 on Nov 19, 2024 | hide | past | pdf | discuss on HN

In plain words: A framework measures how much rounding errors from a speedup trick can shift a model's weights, tested on Flash Attention. It found ten times more error than standard attention in one pass, though the effect on weights stayed several times smaller than low-precision training.

Abstract · Is Flash Attention Stable?

Training large-scale machine learning models poses distinct system challenges, given both the size and complexity of today's workloads. Recently, many organizations training state-of-the-art Generative AI models have reported cases of instability during training, often taking the form of loss spikes. Numeric deviation has emerged as a potential cause of this training instability, although quantifying this is especially challenging given the costly nature of training runs. In this work, we develop a principled approach to understanding the effects of numeric deviation, and construct proxies to put observations into context when downstream effects are difficult to quantify. As a case study, we apply this framework to analyze the widely-adopted Flash Attention optimization. We find that Flash Attention sees roughly an order of magnitude more numeric deviation as compared to Baseline Attention at BF16 when measured during an isolated forward pass. We then use a data-driven analysis based on the Wasserstein Distance to provide upper bounds on how this numeric deviation impacts model weights during training, finding that the numerical deviation present in Flash Attention is 2-5 times less significant than low-precision training.

Alicia Golden, Samuel Hsia, Fei Sun, Bilge Acun, Basil Hosmer, Yejin Lee, Zachary DeVito, Jeff Johnson, Gu-Yeon Wei, David Brooks, Carole-Jean Wu
arXiv:2405.02803 · cs.LG, cs.DC · submitted May 5, 2024
abstract · pdf · html

add comment on HN