SageBwd: Un'Attenzione Addestrabile a Basso Numero di Bit
SageBwd: A Trainable Low-bit Attention
March 2, 2026
Autori: Jintao Zhang, Marco Chen, Haoxu Wang, Kai Jiang, Ion Stoica, Joseph E. Gonzalez, Jianfei Chen, Jun Zhu
cs.AI
Abstract
L'attenzione a basso bit, come SageAttention, è emersa come un approccio efficace per accelerare l'inferenza del modello, ma la sua applicabilità all'addestramento rimane poco compresa. In lavori precedenti, abbiamo introdotto SageBwd, un'attenzione addestrabile in INT8 che quantizza sei delle sette moltiplicazioni di matrici di attenzione preservando le prestazioni durante il fine-tuning. Tuttavia, SageBwd ha mostrato un divario di prestazioni persistente rispetto all'attenzione in precisione completa (FPA) durante la pre-training. In questo lavoro, indaghiamo le cause di questo divario e dimostriamo che SageBwd eguaglia l'attenzione in precisione completa durante la pre-training. Attraverso esperimenti e analisi teorica, raggiungiamo alcune importanti intuizioni e conclusioni: (i) la QK-norm è necessaria per un addestramento stabile con un numero elevato di token per step, (ii) gli errori di quantizzazione originano principalmente dal gradiente dello score dS nella backward-pass, (iii) ridurre i token per step permette a SageBwd di eguagliare le prestazioni della FPA nella pre-training, e (iv) il K-smoothing rimane essenziale per la stabilità dell'addestramento, mentre il Q-smoothing fornisce un beneficio limitato durante la pre-training.
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.