Keep pulling the thread on Tri Dao.
Increasing a model's context length by 4x in NVIDIA's Megatron-LM library reduces training efficiency by approximately 50%, even when normalizing for the total number of tokens.
Flash Attention enables Transformer training that is approximately 1.5x faster than NVIDIA's Megatron-LM at a sequence length of 2,000 tokens.
Flash Attention is 10 to 20 times more memory-efficient than standard attention implementations while being an exact algorithm with no approximation.
Models using Flash Attention achieved the fastest training times for BERT on the MLPerf 2.0 and 2.1 benchmarks.
Flash Attention is planned for inclusion in the official PyTorch 2.0 release, scheduled for March.
Flash Attention has been re-implemented or integrated by organizations including OpenAI, Meta, Hugging Face, MosaicML (in JAX), and is used by Stability AI for training diffusion models.
The core principle of Flash Attention is to exploit the asymmetry between fast, small SRAM and slower, larger HBM on GPUs by designing an I/O-aware algorithm that minimizes data movement.
To save memory, Flash Attention does not store the large N x N attention matrix from the forward pass and instead recomputes it on-the-fly in small blocks during the backward pass.
Flash Attention performs more floating-point operations (FLOPs) than standard attention due to recomputation, but it runs approximately 6x faster because it reduces memory reads and writes by about 9x.
The Flash Attention kernel provides a 2x to 4x speedup for the attention layer compared to a standard PyTorch implementation on an NVIDIA A100 GPU.
Flash Attention reduces the memory requirement of the attention mechanism from quadratic O(N^2) to linear O(N) with respect to sequence length N.
A BERT implementation using Flash Attention outperformed the previous MLPerf training record by 15%.