FlashAttention: IO-aware attention that speeds up transformer training
In one sentence Tri Dao and colleagues (Stanford, University at Buffalo) publish FlashAttention: exact IO-aware attention that avoids writing the attention matrix to HBM, with memory linear in sequence length and up to 3× faster GPT-2 training.
Training large language models requires computing attention between every pair of tokens in a text — an operation that grows quadratically with sequence length. Until 2022, this operation wrote and read massive matrices to chip memory, slow and expensive.
Tri Dao, a PhD student at Stanford, with Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Ré, starts from an observation: the bottleneck is not the number of math operations but the traffic between HBM memory (slow) and SRAM (fast). FlashAttention reorganizes the computation into blocks that stay entirely in SRAM, avoiding materializing the full attention matrix.
The result is faster attention that uses far less memory (10× less at 2K tokens, 20× at 4K), with no approximation: the output is the same. The code is open source and FlashAttention quickly becomes a standard component for training transformers.
Companies
Stanford University, University at Buffalo
Tools
FlashAttention, CUDA, PyTorch
Tags
Sources