![[LLM] Flash Attention](https://img1.daumcdn.net/thumb/R750x0/?scode=mtistory2&fname=https%3A%2F%2Fblog.kakaocdn.net%2Fdn%2FbChxRF%2FbtsDJkHmJXQ%2FdEWXiWTkXwNhyvxniriG81%2Fimg.png)
BackGround
- GPU Kernel
GPU의 커널은 GPU에서 실행되는 프로그램 또는 함수를 의미한다. 이들은 CUDA나 OpenCL과 같은 특수한 라이브러리등을 사용해서 병렬 처리와 GPU연산에 특화되게끔 프로그래밍 된다.
- GPU의 메모리 구조와 기존의 Attention 연산 방식
GPU의 메모리 구조는 다음과 같다.
GPU의 대부분 연산은 SRAM의 값에서 이뤄지고 CPU 연산의 DRAM의 역할을 GPU HBM이 한다고 생각할 수 있다. CPU RAM은 GPU입장에서는 다른 I/O작업이 필요하기 때문에 CPU의 SSD와 유사하다고 볼 수 있다.
기존 Attention 알고리즘에서 Query, Key, Value값들이 어떻게 메모리에서 load와 store가 되는지 알아보면 다음과 같은 순서로 진행된다.
1. HBM에서 Query, Key 값을 읽고 SRAM에서 이를 행렬곱을 하여 Similarity Score를 구하게 된다. 그 후 해당 값을 HBM에 저장하게 된다.
2. HBM에서 Similarity Score를 읽어와 Softmax를 통하여 이를 확률 분포(P)로 바꾸고 이를 HBM에 저장하게 된다.
3. HBM에서 확률 분포(P)와 Value값을 읽어와 행렬곱을 하여 최종 Output을 만들어낸다. 이를 HBM에 저장하고 Return하게 된다.
다음 과정에서 우리는 HBM에 값을 저장하는 횟수가 많은 것을 알 수 있다. 이는 GPU연산에서 bottleneck으로 작용하게 되고 전체적인 연산에서의 비효율을 발생시킨다.
- Softmax함수
Softmax함수는 행렬내의 모든 값들에 대해서 연산을 진행해야하기 때문에 행렬 크기만큼의 메모리가 필요하며 요소별 연산을 진행해야하기 때문에 메모리 부하와 연산의 bottleneck으로써 작용하게 된다.
Tilling
앞써 말한 Softmax 함수의 비효율성을 해결하기 위해 Flash Attention에서는 배열을 다음과 같이 분할하여 각각 softmax 함수를 적용하게 하여 필요한 메모리 공간을 효율적으로 줄일 수 있게 되었다.
이렇게 생긴 memory 사용량의 감소는 Attention에 해당하는 연산에 대한 Kernel fusion을 가능하게 했다.
Recomputation
기존 backward pass에서는 Query, Key, Value의 gradient를 구하기 위해 Similarity Score와 확률 분포를 다시 계산해야했다. 이 과정에서 HBM에 접근하는 횟수가 증가하면서 I/O로 인한 성능저하가 발생했다. Flash Attention은 Attention Output으로 부터 S, P를 recompute하여 dq, dk, dv를 구할 수 있게 backward pass를 수정하였다.
이 과정에서 Flop은 늘었지만 I/O의 횟수가 줄어들어 성능적인 측면에서의 향상을 야기할 수 있었다.
Performance
위의 기법을 통해 GFLOPs의 수는 늘어났지만 HBM에 대한 R/W가 획기적으로 줄면서 전체적인 학습과 추론 시간이 획기적으로 단축된 것을 볼 수 있다.
참고
https://arxiv.org/abs/2205.14135
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
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
arxiv.org
https://github.com/Dao-AILab/flash-attention
GitHub - Dao-AILab/flash-attention: Fast and memory-efficient exact attention
Fast and memory-efficient exact attention. Contribute to Dao-AILab/flash-attention development by creating an account on GitHub.
github.com
긴 글 읽어주셔서 감사합니다.
틀린 부분이 있으면 댓글을 달아주시면 감사하겠습니다.
📧 : realhwan1202@gmail.com
'인공지능 > LLM' 카테고리의 다른 글
| [RAG] Knowledge Graph Prompting for Multi-Document Question Answering (0) | 2024.04.07 |
|---|---|
| [LLM] SELF-RAG: Learning to Retrieve, Generate and Critique through Self-reflect (1) | 2024.02.13 |
| [LLM] Seven Failure Points When Engineering a Retrieval AugmentedGeneration System (0) | 2024.01.29 |
| [LLM] Paged Attention (1) | 2024.01.28 |