FlashAttention

AI가 문장이나 이미지를 처리할 때 자주 쓰는 어텐션 계산을 4비트 숫자로 더 빠르게 돌리려는 연구다. 무엇이 문제였나 — 문제: GPU는 4비트 곱셈을 아주 빨리 할 수 있지만, 어텐션 중간에서 확률값을 만드는 단계가 발목을 잡았다. 어떻게 풀었나 — 해결: 확률을 정밀하게 만든 뒤 줄이는 대신, 처음부터 GPU가 바로 쓸 수 있는 4비트 코드에 점수를 넣는다. 그래서 뭐가 좋아졌나 — 결과: 좋은 조건에서는 어텐션 forward가 BF16보다 최대 2.13배 빨랐고, 8B급 모델의 한 번 학습 업데이트도 최대 1.14배 빨라졌다. ...

학습 중에만 어텐션을 압축했다가 마지막에 원래 방식으로 돌려놓으면, 긴 문맥을 훨씬 빠르고 저렴하게 학습하면서도 정상적인 모델을 얻는다. 무엇이 문제였나 — 긴 문맥(10만100만 토큰) 학습은 어텐션의 제곱 비용 때문에 GPU 시간과 메모리가 폭발한다. 어떻게 풀었나 — 질문·키·값을 같이 여러 해상도로 압축하고 중요한 것만 골라 표준 어텐션을 돌린 뒤, 마지막에 잠깐 원래 어텐션으로 마무리한다. 그래서 뭐가 좋아졌나 — 동일 토큰 예산에서 dense baseline보다 loss가 더 낮거나 같으면서 1.41.7× wall-clock을 절약하고, 512K에서는 21× 빨라진다. ...