SageBwd: Una Atención Entrenable de Bajo Bit
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
Resumen
La atención de baja precisión, como SageAttention, ha surgido como un enfoque efectivo para acelerar la inferencia del modelo, pero su aplicabilidad al entrenamiento sigue siendo poco comprendida. En trabajos previos, presentamos SageBwd, una atención entrenable en INT8 que cuantifica seis de las siete multiplicaciones matriciales de atención mientras preserva el rendimiento del ajuste fino. Sin embargo, SageBwd mostró una brecha de rendimiento persistente respecto a la atención de precisión completa (FPA) durante el pre-entrenamiento. En este trabajo, investigamos por qué ocurre esta brecha y demostramos que SageBwd iguala a la atención de precisión completa durante el pre-entrenamiento. Mediante experimentos y análisis teórico, llegamos a varias conclusiones e ideas importantes: (i) la QK-norma es necesaria para un entrenamiento estable con un gran número de tokens por paso, (ii) los errores de cuantización surgen principalmente del gradiente de puntuación dS en la pasada hacia atrás, (iii) reducir los tokens por paso permite que SageBwd iguale el rendimiento de FPA en el pre-entrenamiento, y (iv) el suavizado de K sigue siendo esencial para la estabilidad del entrenamiento, mientras que el suavizado de Q proporciona un beneficio limitado durante el pre-entrenamiento.
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.