SageBwd: Uma Atenção Treinável de Baixos Bits
SageBwd: A Trainable Low-bit Attention
March 2, 2026
Autores: Jintao Zhang, Marco Chen, Haoxu Wang, Kai Jiang, Ion Stoica, Joseph E. Gonzalez, Jianfei Chen, Jun Zhu
cs.AI
Resumo
A atenção de baixa precisão, como a SageAttention, surgiu como uma abordagem eficaz para acelerar a inferência do modelo, mas sua aplicabilidade ao treinamento permanece pouco compreendida. Em trabalhos anteriores, introduzimos a SageBwd, uma atenção treinável em INT8 que quantiza seis das sete multiplicações matriciais de atenção, preservando o desempenho do fine-tuning. No entanto, a SageBwd exibiu uma lacuna de desempenho persistente em relação à atenção de precisão total (FPA) durante o pré-treinamento. Neste trabalho, investigamos por que essa lacuna ocorre e demonstramos que a SageBwd iguala a atenção de precisão total durante o pré-treinamento. Por meio de experimentos e análise teórica, chegamos a algumas conclusões e insights importantes: (i) a QK-norm é necessária para um treinamento estável com um grande número de tokens por etapa, (ii) os erros de quantização surgem principalmente do gradiente de pontuação dS na passagem backward, (iii) reduzir os tokens por etapa permite que a SageBwd iguale o desempenho da FPA no pré-treinamento, e (iv) o K-smoothing permanece essencial para a estabilidade do treinamento, enquanto o Q-smoothing oferece benefício limitado durante o pré-treinamento.
English
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.