FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

Статьи · arXiv 2205.14135

FlashAttention: быстрый и эффективный по памяти точный механизм внимания с учетом IO

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

Опубликовано 2022-05-27cs.LG arXiv PDF

Читать с переводом и ИИ

Кратко в трёх строках

🎯 Стандартное точное внимание материализует матрицу внимания сложности $O (N^{2})$ в медленной HBM GPU, из-за чего вычисления упираются в пропускную способность памяти и плохо масштабируются на длинные контексты.
🔧 FlashAttention вычисляет точное внимание внутри единого объединенного ядра CUDA с помощью тайлинга входных данных в быстрой SRAM на кристалле, инкрементального обновления softmax и перевычисления внимания на кристалле во время обратного прохода.
📈 FlashAttention обеспечивает до 3-кратного сквозного ускорения обучения GPT-2, обучается на 15% быстрее рекорда MLPerf 1.1 для BERT-large и позволяет масштабировать контекст до 64K токенов.

Аннотация

Трансформеры работают медленно и потребляют много памяти на длинных последовательностях, поскольку временная и пространственная сложность механизма self-attention квадратично зависят от длины последовательности. Приближенные методы внимания пытались решить эту проблему, жертвуя качеством модели ради снижения вычислительной сложности, однако зачастую они не приводят к реальному ускорению работы (wall-clock speedup). Мы утверждаем, что упущенным фундаментальным принципом является учет операций ввода-вывода (IO-awareness) в алгоритмах внимания — то есть учет операций чтения и записи между различными уровнями памяти GPU. Мы предлагаем FlashAttention — точный алгоритм внимания с учетом IO, использующий разбиение на блоки (tiling) для уменьшения количества операций чтения/записи между памятью высокой пропускной способности GPU (HBM) и локальной памятью SRAM на кристалле GPU. Мы анализируем сложность ввода-вывода FlashAttention и показываем, что ему требуется меньше обращений к HBM, чем стандартному вниманию, а также доказываем его оптимальность для широкого диапазона размеров SRAM. Мы также расширяем FlashAttention до block-sparse attention, получая алгоритм приближенного внимания, который работает быстрее любого существующего приближенного метода. FlashAttention обучает трансформеры быстрее существующих базовых решений: 15% сквозного ускорения на BERT-large (длина последовательности 512) по сравнению с рекордом скорости обучения MLPerf 1.1, ускорение в 3 $\times$ раза на GPT-2 (длина последовательности 1K) и ускорение в 2.4 $\times$ раза на long-range arena (длина последовательности 1K–4K). FlashAttention и block-sparse FlashAttention открывают путь к более длинному контексту в трансформерах, обеспечивая более высокое качество моделей (улучшение перплексии на 0.7 для GPT-2 и прирост на 6.4 пункта при классификации длинных документов), а также абсолютно новые возможности: первые трансформеры, показавшие результат выше случайного угадывания в бенчмарках Path-X (длина последовательности 16K, точность 61.4%) и Path-256 (длина последовательности 64K, точность 63.1%).

Оригинальная аннотация (англ.)

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

Разбор статьи

Проблема и мотивация

Трансформеры обладают квадратичной вычислительной сложностью и сложностью по памяти относительно длины последовательности $N$ . Хотя приближенные методы внимания снижают теоретическое количество операций FLOP, они часто не дают выигрыша по реальному времени выполнения, поскольку узким местом современных GPU являются задержки при передаче данных из памяти (IO), а не арифметические операции.

В стандартном механизме внимания при вычислении $S = Q K^{⊤}$ , $P = softmax (S)$ и $O = PV$ промежуточные матрицы размера $N \times N$ записываются в медленную память с высокой пропускной способностью (HBM) и считываются оттуда несколько раз. Это требует $Θ (N d + N^{2})$ обращений к HBM и приводит к расходу памяти $O (N^{2})$ , что на практике сильно ограничивает размер окна контекста.

Ключевые идеи

IO-Awareness (учет операций ввода-вывода): Явный учет иерархии памяти GPU: быстрая, но небольшая память SRAM на кристалле (например, 192 КБ на один SM с пропускной способностью около 19 ТБ/с на A100) в сравнении с большой, но медленной HBM (40–80 ГБ с пропускной способностью 1.5–2.0 ТБ/с).
Tiling (блочное вычисление softmax): Разбиение входных матриц $Q, K, V$ на блоки, полностью помещающиеся в SRAM. За счет отслеживания промежуточных статистик (максимум в строке и сумма нормализатора) редукция softmax вычисляется инкрементально по блокам без материализации полной матрицы $N \times N$ в HBM.
Перевычисление на обратном проходе (Recomputation in Backward Pass): Отказ от сохранения промежуточной матрицы внимания $N \times N$ для обратного распространения ошибки. Вместо этого сохраняются только выходной тензор $O$ и статистики softmax $(ℓ, m)$ , а $S$ и $P$ повторно вычисляются блок за блоком прямо в памяти SRAM во время обратного прохода.
Слияние ядер (Kernel Fusion): Объединение матричных умножений, маскирования, softmax и dropout в единое ядро CUDA для исключения избыточных циклов чтения/записи в HBM.

Метод (с ключевыми формулами)

Для $Q, K, V \in R^{N \times d}$ алгоритм FlashAttention задает размеры блоков $B_{c} = ⌈ M / (4 d)⌉$ и $B_{r} = min (⌈ M / (4 d)⌉, d)$ , где $M$ — объем SRAM.

Инкрементальное отслеживание Softmax

Для вектора, разбитого на подитерации $x^{(1)}$ и $x^{(2)}$ , текущие статистики обновляются следующим образом: $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

Для блока $j$ ключей/значений и блока $i$ запросов:

Вычисляются ненормализованные скоры в SRAM: $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$ .
Вычисляются локальные статистики блока: $\tilde{m}_{ij} = rowmax (S_{ij})$ , $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ и $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$ .
Обновляются текущие глобальные нормализаторы: $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
Обновляется текущий выходной блок $O_{i} \in R^{B_{r} \times d}$ : $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

Анализ сложности

Сложность по операциям ввода-вывода (IO Complexity): Стандартное внимание требует $Θ (N d + N^{2})$ обращений к HBM. FlashAttention требует $Θ (N^{2} d^{2} M^{- 1})$ обращений к HBM. Авторы доказывают, что ни один точный алгоритм внимания не может асимптотически достичь сложности $o (N^{2} d^{2} M^{- 1})$ для всех значений $M \in [d, N d]$ .
FLOP и память: FlashAttention выполняет $O (N^{2} d)$ FLOP (с небольшим увеличением количества операций на обратном проходе из-за перевычислений) и требует всего $O (N)$ дополнительной памяти помимо входных и выходных данных.

Расширение: Block-Sparse FlashAttention

Ограничение вычислений только ненулевыми блоками в соответствии с заданной маской разреженности $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ обеспечивает сложность ввода-вывода $Θ (N d + N^{2} d^{2} M^{- 1} s)$ , где $s$ — доля ненулевых блоков.

Результаты (с важнейшими показателями)

Ускорение обучения:
- BERT-large (длина последовательности 512): Превосходит рекорд скорости обучения MLPerf 1.1 на 15% (17.4 мин против 20.0 мин на $8 \times$ A100).
- GPT-2 (длина последовательности 1K): Обеспечивает ускорение до 3.5x по сравнению с реализацией HuggingFace и до 1.8x по сравнению с Megatron-LM.
- Long-Range Arena (LRA): Дает ускорение в 2.4x по сравнению со стандартным вниманием; Block-sparse FlashAttention ускоряет обучение в 2.8x раза.
Возможности работы с длинным контекстом:
- GPT-2 с контекстом 4K: Обучается быстрее, чем Megatron-LM с контекстом 1K, при этом улучшая перплексию на валидации на 0.7 пункта (17.5 против 18.2).
- Классификация длинных документов: Масштабирование контекста с 512 до 16K дает прирост в 4.3 micro- $F_{1}$ на MIMIC-III (с 52.8 до 57.1) и до 8.5 пунктов на ECtHR (с 72.2 до 80.7 при 8K).
- Path-X и Path-256: FlashAttention — первый трансформер, успешно решивший задачу Path-X (длина последовательности 16K) с точностью 61.4%. Block-sparse FlashAttention достигает 63.1% точности на Path-256 (длина последовательности 64K).

Ограничения и открытые вопросы

Сложность разработки на CUDA: Реализация алгоритмов с учетом IO требует низкоуровневого программирования на CUDA, а не использования высокоуровневых абстракций PyTorch или компиляторов, что усложняет переносимость кода между различными аппаратными архитектурами.
Многопроцессорный механизм внимания (Multi-GPU Attention): Оптимальный план операций ввода-вывода для одного узла не переносится напрямую на распределенное вычисление внимания на нескольких GPU через интерконнекты (NVLink, InfiniBand).
Нижняя граница сложности точного внимания: Хотя алгоритм оптимален для широкого диапазона $M$ , доказательство точных параметризованных нижних границ сложности относительно $M$ остается открытой теоретической задачей.

Ключевые термины

IO-awareness: Алгоритмический принцип проектирования операций с явным учетом чтения и записи между уровнями аппаратной памяти (например, быстрой SRAM на кристалле и более медленной HBM).
High Bandwidth Memory (HBM): Основная память GPU большого объема, обладающая высокой емкостью, но существенно меньшей пропускной способностью по сравнению с кэшем на кристалле.
SRAM: Быстрая статическая память на кристалле современных GPU, размещенная внутри каждого потокового мультипроцессора (SM).
Tiling: Реструктуризация вычислений путем разбиения входных матриц на меньшие блоки, помещающиеся в SRAM, с накоплением результатов по ходу итераций.
Selective gradient checkpointing: Метод, позволяющий не сохранять большую промежуточную матрицу внимания во время прямого прохода, а перевычислять ее на кристалле во время обратного прохода с использованием сохраненных статистик softmax.
Kernel fusion: Объединение нескольких отдельных операций (матричное умножение, применение маски, softmax, dropout) в единый запуск ядра GPU для устранения повторных операций чтения/записи в HBM.
Block-sparse FlashAttention: Расширение FlashAttention, вычисляющее только ненулевые блоки тайлов согласно блочной маске разреженности, что снижает сложность ввода-вывода пропорционально доле ненулевых элементов.
Path-X / Path-256: Бенчмарки визуального рассуждения на длинных последовательностях из Long-Range Arena (длина последовательностей 16K и 64K), на которых предыдущие трансформеры не могли показать точность выше случайного угадывания.

Вопросы

Является ли FlashAttention приближенным алгоритмом вычисления матрицы внимания?

Нет. FlashAttention вычисляет точный механизм самовнимания (exact self-attention), математически полностью эквивалентный стандартному вниманию с точностью до погрешности вычислений с плавающей запятой.

Почему FlashAttention работает быстрее, если на обратном проходе он выполняет больше FLOP?

Потому что стандартное внимание ограничено пропускной способностью памяти (memory-bound); перемещение матрицы $N \times N$ между HBM и SRAM занимает намного больше реального времени, чем выполнение дополнительных арифметических операций в быстрой памяти SRAM.

Какую вспомогательную память сохраняет FlashAttention для обратного прохода?

FlashAttention сохраняет только выходные данные прямого прохода $O$ , коэффициенты нормализации softmax $ℓ$ и $m$ (требующие $O (N)$ дополнительной памяти), а также состояние генератора псевдослучайных чисел для воспроизведения масок dropout на кристалле.

Как FlashAttention соотносится со стандартным вниманием по сложности ввода-вывода (IO complexity)?

Стандартное внимание требует $Θ (N d + N^{2})$ обращений к HBM, тогда как FlashAttention требует $Θ (N^{2} d^{2} M^{- 1})$ обращений, где $M$ — объем доступной памяти SRAM, а $d$ — размерность проекции (head dimension).

Обзор на этой странице написан ИИ по тексту статьи и может содержать ошибки. Подробности смотрите в самой статье.

Читать с переводом и ИИ