[논문 리뷰] FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Tri Dao|arXiv (Cornell University)|2023. 07. 17.

Advanced Neural Network Applications인용 수 140

한 줄 요약

FlashAttention-2는 Transformer의 주의(attention)을 병렬 처리 및 작업 분할 재구성으로 가속화하여 FlashAttention 대비 약 2배의 속도 향상을 달성하고, 순방향에서 이론상 최대 FLOPs의 73%, 역방향에서 63%까지 달성하며, 엔드-투-엔드 GPT 스타일 훈련은 A100당 최대 225 TFLOPs/s까지 가능하다.

ABSTRACT

Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FlashAttention exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4$ imes$ compared to optimized baselines), with no approximation. However, FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40\% of the theoretical maximum FLOPs/s. We observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2$ imes$ speedup compared to FlashAttention, reaching 50-73\% of the theoretical maximum FLOPs/s on A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72\% model FLOPs utilization).

연구 동기 및 목표

• 주의 병목 현상을 줄여 Transformer를 더 긴 컨텍스트 길이로 확장하도록 동기를 부여한다.
• 스레드 블록과 워프 간 작업 분할을 재고하여 GPU 활용도를 향상시킨다.
• 비-매트멀 FLOPs를 줄여 대부분의 시간을 빠른 matmul 연산에 남겨두는 것을 목표로 한다.
• 시퀀스 길이, 배치, 헤드 전반에 걸친 병렬성을 늘려 점유율을 높인다.
• GPT 스타일 모델에서 엔드-투-엔드 훈련 속도 향상을 검증한다.

제안 방법

• 출력을 바꾸지 않으면서 비-매트멀 FLOPs를 줄이기 위해 FlashAttention의 순전파/역전파 알고리즘을 미세 조정한다.
• 점유율을 높이기 위해 배치와 헤드뿐 아니라 시퀀스 길이에도 주의(attention)을 병렬화한다.
• 스레드 블록 내 워프 간에 작업을 분배해 공유 메모리 트래픽을 최소화한다.
• 온라인 소프트맥스 타일링을 사용해 올바른 최종 출력을 가진 블록 단위 계산을 가능하게 한다.
• 적절한 블록을 건너뛰고 중복 마스킹 작업을 줄여 인과 마스킹을 효율적으로 적용한다.
• 블록/워프 분할이 포함된 순전파 및 역전파를 설명하는 알고리즘(Algorithm 1 및 Algorithm 2)을 제공한다.

실험 결과

연구 질문

• RQ1 FlashAttention-2가 FlashAttention에 비해 순전파와 역전파 주의 패스에서 더 높은 GPU 처리량을 달성할 수 있는가?
• RQ2 작업 분할 변경과 비-매트멀 FLOPs가 런타임과 메모리 트래픽을 얼마나 줄이는가?
• RQ3 다양한 컨텍스트 길이에 걸쳐 GPT 스타일 모델에서 FlashAttention-2를 사용할 때 어떤 엔드-투-엔드 훈련 속도 향상이 실현되는가?
• RQ4 현대 GPU(A100, H100)에서 이론상 최대 FLOPs/s에 얼마나 근접할 수 있는가?

주요 결과

• 벤치마크에서 FlashAttention 대비 약 2배의 속도향상을 제공한다.
• 순전파 처리량은 A100에서 이론상 최대 FLOPs/s의 최대 73%까지; 역전파는 최대 63%에 도달한다.
• GPT 스타일 모델의 엔드-투-엔드 훈련은 A100 GPU당 최대 225 TFLOPs/s에 도달하며 모델 FLOPs 활용도는 72%이다.
• GPT-3 1.3B/2.7B 규모에서 FlashAttention-2는 2k 및 8k 컨텍스트에서 베이스라인 및 FlashAttention 대비 상당한 개선을 달성한다.
• H100 GPU에서 순전파+역전파 속도향상과 원시 처리량은 새로운 하드웨어 기능으로 더 개선될 것으로 기대된다.