SageBwd: 학습 가능한 저정밀 어텐션

SageBwd: A Trainable Low-bit Attention

SageAttention과 같은 저정밀 어텐션은 모델 추론 속도를 가속화하는 효과적인 방법으로 떠오르고 있지만, 학습에서의 적용 가능성에 대한 이해는 아직 부족합니다. 이전 연구에서, 우리는 6개의 7개 어텐션 행렬 곱셈을 양자화하면서도 미세 조정 성능을 유지하는 학습 가능한 INT8 어텐션인 SageBwd를 소개했습니다. 그러나 SageBwd는 사전 훈련 과정에서 여전히 전체 정밀도 어텐션(FPA)과의 성능 격차를 보였습니다. 본 연구에서는 이러한 격차가 발생하는 이유를 조사하고, SageBwd가 사전 훈련 과정에서 FPA와 동등한 성능을 달성함을 보여줍니다. 실험과 이론적 분석을 통해 다음과 같은 중요한 통찰력과 결론을 얻었습니다. (i) QK-norm은 큰 배치 크기로 훈련할 때 안정적인 학습에 필수적입니다. (ii) 양자화 오류는 주로 역전파 과정에서 발생하는 스코어 그래디언트 dS에서 비롯됩니다. (iii) 배치 크기를 줄이면 SageBwd가 사전 훈련에서 FPA와 동일한 성능을 달성할 수 있습니다. (iv) K-smoothing은 여전히 훈련 안정성에 필수적이며, Q-smoothing은 사전 훈련 과정에서 제한적인 이점을 제공합니다.

Low-bit attention, such as SageAttention, has emerged as an effective approach for accelerating model inference, but its applicability to training remains poorly understood. In prior work, we introduced SageBwd, a trainable INT8 attention that quantizes six of seven attention matrix multiplications while preserving fine-tuning performance. However, SageBwd exhibited a persistent performance gap to full-precision attention (FPA) during pre-training. In this work, we investigate why this gap occurs and demonstrate that SageBwd matches full-precision attention during pretraining. Through experiments and theoretical analysis, we reach a few important insights and conclusions: (i) QK-norm is necessary for stable training at large tokens per step, (ii) quantization errors primarily arise from the backward-pass score gradient dS, (iii) reducing tokens per step enables SageBwd to match FPA performance in pre-training, and (iv) K-smoothing remains essential for training stability, while Q-smoothing provides limited benefit during pre-training.