주의(attention)는 대형 언어 모델(LLMs)을 구동하는 변환기(transformer) 아키텍처의 기본 요소입니다. 그러나 LLM이 확대되고 점점 더 긴 입력 시퀀스를 관리함에 따라, 주의 계산의 계산 비용이 주요 병목 현상으로 부각되고 있습니다.
이 문제를 해결하기 위해 Colfax Research, Meta, Nvidia, Georgia Tech, Princeton University, Together AI의 협력팀이 FlashAttention-3를 도입했습니다. 이 첨단 기술은 Nvidia Hopper GPUs(H100 및 H800)에서 주의 계산을 대폭 가속화합니다.
LLMs에서의 주의 계산 문제 이해하기
변환기 모델의 주의 메커니즘은 입력 시퀀스 내 다양한 토큰 간의 관계를 평가할 수 있는 독특한 기능을 제공합니다. 그러나 이 메커니즘은 계산 집약적입니다. 입력 시퀀스의 길이가 증가함에 따라 주의 계산 비용은 제곱으로 증가해 LLM의 확장에 큰 병목 현상을 초래합니다.
또한, 현대 GPU는 주로 행렬 곱셈(matmul) 작업에 최적화되어 있으며, 이는 딥 러닝의 핵심입니다. 지수 승과 같은 다른 연산은 훨씬 느려 성능 문제를 더욱 악화시킵니다. 주의 계산은 행렬 곱셈과 더 복잡한 함수(예: 주의 가중치를 정규화하는 소프트맥스 함수)를 통합하므로 후자의 높은 계산 비용으로 제약을 받을 수 있습니다. 작업 부하를 효과적으로 스케줄링하여 운영 충돌을 피하고 메모리 자원을 효율적으로 활용하는 것이 매우 중요합니다.
FlashAttention으로 하드웨어 자원 활용 향상하기
2022년에 출시된 FlashAttention은 GPU의 고대역폭 메모리(HBM)와 정적 임의 접근 메모리(SRAM) 간의 메모리 전송을 최소화하여 주의 계산의 비효율성을 해결했습니다. 주의 가중치를 더 작은 단위인 "타일"로 처리함으로써 FlashAttention은 효율성을 향상시켜 LLM이 수천 개에서 잠재적으로 수백만 개의 토큰까지 컨텍스트 윈도를 확장할 수 있게 했습니다.
그러나 하드웨어 기능이 발전함에 따라 추가 최적화의 필요성도 증가했습니다. 2023년에 도입된 FlashAttention-2는 GPU 자원 활용을 최적화하여 Nvidia A100 GPU에서 최대 성능의 70%를 달성했지만 H100의 성능은 35%만 활용했습니다.
FlashAttention-3의 혁신
FlashAttention-3는 새로운 Nvidia Hopper GPU 기능을 활용하여 성능을 향상시킵니다. 이에는 행렬 곱셈의 처리량 개선과 메모리 세그먼트 간의 데이터 전송 속도 향상이 포함됩니다. 이로 인해 저정밀 연산에서도 더 나은 효율성을 제공합니다.
FlashAttention-3의 주요 혁신 사항은 다음과 같습니다:
1. 최적화된 스케줄링: 작업이 컴퓨테이션과 데이터 이동의 겹침을 극대화하도록 조직되어 GPU 유휴 시간을 감소시킵니다.
2. 원활한 동작 교차: 행렬 곱셈과 소프트맥스 연산을 결합해 FlashAttention-3는 병목 현상을 최소화합니다.
3. 향상된 양자화 모델 성능: 연산 내 특별한 배열이 저비트 표현 사용 시에도 빠르고 정확한 계산을 보장하여 양자화와 관련된 일반적인 정확도 저하 문제를 해결합니다.
연구 결과, FlashAttention-3는 H100 GPU의 최대 성능의 75%를 활용할 수 있으며, 이전 FlashAttention 버전보다 LLM 교육 및 배포에서 1.5-2배의 속도 향상을 제공합니다.
FlashAttention-3의 이점
FlashAttention-3가 제공하는 빠른 주의 계산은 LLM 개발 및 응용에 깊은 영향을 미칩니다:
- 가속화된 훈련: 향상된 효율성으로 훈련 시간을 크게 줄여 연구자들이 더 큰 모델과 데이터셋을 탐구할 수 있도록 합니다.
- 확장된 컨텍스트 윈도우: 긴 시퀀스를 효율적으로 처리할 수 있게 해 주며, LLM 응용 범위를 넓히고 있습니다.
- 비용 효율성: 높은 GPU 활용도로 인해 LLM 운영에 필요한 가속기가 줄어들어 생산 비용을 절감할 수 있습니다.
FlashAttention-3는 허가된 라이선스 하에 오픈소스로 제공되며, PyTorch 및 Hugging Face Transformers와 같은 인기 있는 딥 러닝 라이브러리와의 통합이 계획되어 있습니다. 이는 연구자와 개발자가 FlashAttention-3의 발전을 활용할 수 있도록 도와줄 것입니다. Together AI 블로그에서 연구진이 말한 것처럼, “하드웨어 기능을 활용한 알고리즘 설계는 효율성을 크게 향상시키고 새로운 모델 기능을 열 수 있습니다.” 이들은 LLM 추론과 다양한 하드웨어 아키텍처에서 자신들의 기술을 적용할 수 있는 추가 최적화를 기대하고 있습니다.
순위 목록
