요즘 고전이 된 논문들을 다시 공부하면서 커리어리에 정리해 보려고 한다. 오늘 소개할 논문은 2022년에 스탠퍼드 대학 박사과정 Tri Dao가 쓴 'FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness'(https://arxiv.org/abs/2205.14135) 이다.

이 논문의 주된 아이디어는 트랜스포머 모델의 훈련과 추론 속도를 높이기 위해, GPU 메모리 입출력(IO)을 줄이는 것이다. 이를 위해 연산량(FLOPS)이 증가하더라도 IO를 최소화하는 방식으로 어텐션 연산을 최적화한다.

먼저, GPU 메모리 계층 구조를 이해해야 한다. GPU 메모리는 빠르지만 용량이 작은 HBM과 느리지만 용량이 큰 SRAM으로 구성된다. 논문은 GPU 연산 중 HBM과 SRAM 사이의 데이터 이동이 주요 병목 현상이라고 지적한다.

참고로 딥러닝 연산은 크게 두 가지로 나뉜다:

1. 컴퓨팅 바운드 연산: 연산량이 메모리 IO보다 많은 경우 (예: 행렬곱)

2. 메모리 바운드 연산: 메모리 IO가 연산량보다 많은 경우 (예: 소프트맥스, 활성화 함수, 배치/레이어 정규화)

GPU는 주로 컴퓨팅 바운드 연산에 큰 이점을 가지고 있기 때문에, 느린 메모리 바운드 연산을 효율화하는 데 이 논문은 초점을 맞추었다.

트랜스포머의 일반적인 어텐션 연산은 다음과 같은 단계로 이루어진다:

1. 쿼리(Query)와 키(Key) 행렬의 곱셈

2. 행 단위 소프트맥스 적용

3. 밸류(Value) 행렬과의 곱셈

이 과정에서 HBM 읽기와 쓰기가 반복적으로 일어나게 된다.

반면, 플래시 어텐션은 다음과 같은 방식으로 어텐션 연산을 최적화한다:

1. 기존의 어텐션 연산을 SRAM 크기에 적합한 블록 단위로 분할

2. 각 블록에 대해 HBM에서 단 한 번만 데이터를 읽어와 SRAM에 저장

3. SRAM에서 행렬곱, 마스킹, 소프트맥스, 행렬곱을 모두 수행

4. 결과를 HBM에 다시 저장
5. 1 ~ 4번의 연산을 블록마다 반복

이것은 특정 벡터에 대한 소프트맥스 연산을 여러 서브 벡터에 대한 순차적인 누적 연산으로 분할하여 정확하게 계산할 수 있기 때문에 가능하다. 이러한 타일링(Tiling)은 입력이 동일한 다수의 연산을 합쳐서 처리량을 늘리는 커널 퓨전 기법의 일종이다.

또한, 플래시 어텐션은 역전파 시 필요한 중간값 행렬들을 저장하지 않고 그때그때 재연산하는 방식을 사용한다. 이 역시 연산량을 늘리지만 IO는 줄이는 효과가 있으며, 그래디언트 체크포인팅 기법이라고 할 수 있다.

이러한 플래시 어텐션 기법을 통해 훈련 및 추론 시간을 단축할 수 있지만, 별도의 CUDA 구현과 장비별 컴파일링 과정이 필요하다는 단점이 있다. 또한, 이론적인 성능에 비해 실제 속도 향상이 다소 제한적인 면이 있는데, 이에 대한 보완책으로 FlashAttention-2 논문이 발표되었다.

#Transformer #Optimization