SageBwd: Een Trainbare Low-bit Attention
SageBwd: A Trainable Low-bit Attention
March 2, 2026
Auteurs: Jintao Zhang, Marco Chen, Haoxu Wang, Kai Jiang, Ion Stoica, Joseph E. Gonzalez, Jianfei Chen, Jun Zhu
cs.AI
Samenvatting
Low-bit aandacht, zoals SageAttention, is naar voren gekomen als een effectieve benadering voor het versnellen van modelinferentie, maar de toepasbaarheid ervan op training blijft slecht begrepen. In eerder werk introduceerden we SageBwd, een trainbare INT8-aandacht die zes van de zeven aandachtmatrixvermenigvuldigingen kwantiseert terwijl de fine-tuningprestaties behouden blijven. SageBwd vertoonde echter een hardnekkige prestatiekloof met full-precision aandacht (FPA) tijdens voorafgaande training. In dit werk onderzoeken we waarom deze kloof optreedt en tonen we aan dat SageBwd gelijkwaardig presteert aan full-precision aandacht tijdens de voorafgaande training. Door experimenten en theoretische analyse komen we tot enkele belangrijke inzichten en conclusies: (i) QK-norm is noodzakelijk voor stabiele training bij een groot aantal tokens per stap, (ii) kwantiseringsfouten ontstaan voornamelijk vanuit de scoregradiënt dS in de backward-pass, (iii) het verminderen van tokens per stap stelt SageBwd in staat om de FPA-prestaties in voorafgaande training evenaren, en (iv) K-smoothing blijft essentieel voor trainstabiliteit, terwijl Q-smoothing beperkt voordeel biedt tijdens voorafgaande 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.