SageBwd: 학습 가능한 저비트 어텐션 메커니즘

대규모 언어 모델 학습 시 가장 큰 연산 병목인 Attention을 INT8로 양자화하면서도 Full-precision과 동일한 성능을 유지하는 방법을 제시한다. 특히 역전파 과정의 수치적 불안정성을 해결하여 소비자용 GPU에서도 대규모 모델 학습을 가속화할 수 있는 실질적인 가이드라인을 제공한다.

Transformer의 Attention은 문장 내 단어 간의 관계를 계산하기 위해 쿼리(Q)와 키(K)를 곱하는 과정을 거친다. 이 연산은 모델이 커질수록 기하급수적으로 늘어나며, 이를 INT8과 같은 낮은 정밀도로 계산하면 속도는 빠르지만 학습 과정에서 미세한 정보가 손실되어 모델이 제대로 수렴하지 못하는 한계가 있었다.

본 논문은 모델이 정답을 찾아가는 역전파(Backward pass) 과정에서 발생하는 dS(스코어 그래디언트)라는 값에 주목한다. dS는 시퀀스 길이가 길어질수록 값이 매우 작아지는데, 일반적인 양자화 방식은 이 작은 값을 0으로 처리하거나 노이즈로 오염시켜 전체 학습의 방향을 틀어버린다. 이는 마치 정밀한 지도가 필요한 상황에서 해상도가 낮은 사진을 사용하는 것과 같다.

이를 해결하기 위해 QK-norm을 도입하여 Q와 K의 값 범위를 일정하게 묶어줌으로써 양자화가 더 정밀하게 이루어지도록 돕는다. 또한 학습 시 한 번에 처리하는 토큰 양(TPS)을 적절히 조절하면, 학습 과정에서 발생하는 자연스러운 무작위 노이즈가 양자화로 인한 오류를 덮어주는 효과를 내어 최종적으로 성능 손실 없이 연산 속도만 1.67배 높일 수 있다.

SageBwd의 양자화 전략은 Attention 연산의 7개 MatMul 중 6개를 INT8로 수행하며, 역전파 시 가장 오차에 민감한 dP = dOV^T 연산만 FP16 정밀도를 유지하여 수치적 안정성을 확보한다. 이는 오차 증폭의 핵심 경로를 고정밀도로 보호하는 전략이다.

QK-norm을 통한 수치 안정화는 각 토큰 벡터에 RMSNorm을 적용하여 Q와 K의 스케일을 제어한다. [입력 벡터 → RMS 계산 및 정규화 → 출력 벡터] 과정을 거쳐 소프트맥스 입력값인 logits가 양자화에 적합한 수치 범위 내에 머물게 하여 정확도를 개선한다.

K-smoothing 전처리 기법은 양자화 직전 Key 행렬에서 토큰별 평균값을 차감한다. [Key 행렬 → 행별 평균 계산 및 차감 → 정규화된 Key] 순서로 연산하여 특정 채널에 치우친 이상치(outliers)의 영향을 제거하고 학습 안정성을 높인다.

OpenAI Triton을 활용한 커널 최적화는 블록 단위 양자화(per-block quantization)를 적용한다. 이를 통해 GPU의 Tensor Core 활용도를 극대화하고 메모리 대역폭 병목을 해결하여 실질적인 가속 성능을 달성한다.

325M Llama 모델을 78B 토큰으로 사전 학습한 결과, 260K TPS 설정에서 SageBwd는 Full-precision Attention(FPA)과 거의 일치하는 Loss 수렴 성능을 보였다. 이는 저비트 양자화가 사전 학습 단계에서도 충분히 실용적임을 입증한다.

RTX4090 GPU 환경에서 FlashAttention2 대비 최대 1.67배의 속도 향상을 달성했다. 이는 Triton 기반 구현임에도 불구하고 고도로 최적화된 기존 CUDA 커널보다 높은 효율성을 보여준 결과이다.

Ablation study 결과, QK-norm이 없는 대규모 배치(2.1M TPS) 학습에서는 양자화 오차 누적으로 인해 Loss가 발산하는 현상이 관찰되었다. 또한 K-smoothing이 없는 경우에도 학습 안정성이 크게 저하됨을 확인하여 각 구성 요소의 필수성을 검증했다.

dS 텐서의 크기 상한 분석을 통해 RMS(dS)가 시퀀스 길이 N의 제곱근에 반비례하여 작아짐을 수식적으로 증명했다. 이는 시퀀스가 길어질수록 dS의 신호 강도가 약해져 INT8 양자화 시 낮은 SNR을 유발하는 근본 원인이 된다.

역전파 그래디언트 전파 메커니즘 분석 결과, dQ = dSK 및 dK = dS^T Q 연산에서 dS의 양자화 오차가 Q와 K의 노름(norm)에 의해 증폭됨을 규명했다. QK-norm은 이 노름을 제한함으로써 상위 단계의 오차가 하위 단계로 전이되는 것을 억제한다.

TPS(Tokens Per Step)와 양자화 편향의 상관관계를 분석하여, 큰 배치 사이즈에서는 그래디언트가 결정론적으로 변해 양자화의 체계적 편향이 두드러지지만, 작은 배치에서는 확률적 노이즈가 이를 효과적으로 마스킹함을 발견했다.

Q-smoothing의 한계에 대해서도 고찰했다. Q-smoothing은 추가적인 바이어스 보정 항인 dK_bias = (dS^T 1) mu_Q^T를 필요로 하며, 이 과정에서 새로운 양자화 노이즈 경로가 추가되어 사전 학습 단계에서는 실질적인 이득이 적음을 확인했다.

매우 큰 배치 사이즈(2.1M TPS 이상)에서는 여전히 Full-precision Attention 대비 성능 격차가 존재한다. 배치 사이즈를 줄이지 않고도 dS 경로의 양자화 오차를 완화할 수 있는 추가적인 기법 연구가 향후 과제로 남아 있다.