Daily Arxiv

전 세계에서 발간되는 인공지능 관련 논문을 정리하는 페이지 입니다.
본 페이지는 Google Gemini를 활용해 요약 정리하며, 비영리로 운영 됩니다.
논문에 대한 저작권은 저자 및 해당 기관에 있으며, 공유 시 출처만 명기하면 됩니다.

SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training

Created by
  • Haebom

저자

Jintao Zhang, Jia Wei, Pengle Zhang, Xiaoming Xu, Haofeng Huang, Haoxu Wang, Kai Jiang, Jun Zhu, Jianfei Chen

개요

본 논문은 어텐션 메커니즘의 계산 복잡도를 해결하기 위해 두 가지 주요 기여를 제시합니다. 첫째, Blackwell GPU의 FP4 Tensor Cores를 활용하여 어텐션 연산을 가속화하는 방법을 제안하며, RTX 5090에서 기존 최고 성능 대비 5배 빠른 1038 TOPS의 성능을 달성했습니다. 다양한 모델의 추론 속도를 향상시키는 플러그 앤 플레이 방식을 제시합니다. 둘째, 저비트 어텐션을 훈련 과정에 적용하는 선구적인 연구를 진행했습니다. 기존의 FlashAttention3이나 SageAttention과 달리 추론뿐 아니라 훈련 효율 향상에도 초점을 맞춰, 정확하고 효율적인 8-bit 어텐션을 전방 및 역방향 전파에 적용했습니다. 실험 결과, 8-bit 어텐션은 미세 조정 작업에서는 손실 없는 성능을 달성하지만, 사전 훈련 작업에서는 수렴 속도가 느린 것으로 나타났습니다. 코드는 https://github.com/thu-ml/SageAttention 에서 공개될 예정입니다.

시사점, 한계점

시사점:
FP4 Tensor Cores 활용을 통한 어텐션 연산의 획기적인 속도 향상 (5배).
다양한 모델에 플러그 앤 플레이 방식으로 적용 가능한 FP4 어텐션 구현.
저비트 어텐션을 훈련 과정에 적용한 최초의 시도.
미세 조정 작업에서 손실 없는 성능을 보이는 8-bit 어텐션 제시.
한계점:
8-bit 어텐션은 사전 훈련 작업에서 수렴 속도 저하 문제 발생.
👍