SageBwd: Eine trainierbare Aufmerksamkeit mit niedriger Bittiefe
SageBwd: A Trainable Low-bit Attention
March 2, 2026
Autoren: Jintao Zhang, Marco Chen, Haoxu Wang, Kai Jiang, Ion Stoica, Joseph E. Gonzalez, Jianfei Chen, Jun Zhu
cs.AI
Zusammenfassung
Niedrigbit-Aufmerksamkeit, wie beispielsweise SageAttention, hat sich als effektiver Ansatz zur Beschleunigung der Modellinferenz erwiesen, doch ihre Anwendbarkeit auf das Training ist nach wie vor wenig verstanden. In vorhergehenden Arbeiten haben wir SageBwd vorgestellt, eine trainierbare INT8-Aufmerksamkeit, die sechs von sieben Aufmerksamkeitsmatrix-Multiplikationen quantisiert, ohne die Feinabstimmungsleistung zu beeinträchtigen. Allerdings wies SageBwd während des Vorabtrainings eine anhaltende Leistungslücke zur Vollpräzisions-Aufmerksamkeit (FPA) auf. In dieser Arbeit untersuchen wir, warum diese Lücke auftritt, und zeigen, dass SageBwd während des Vorabtrainings mit der Vollpräzisions-Aufmerksamkeit gleichzieht. Durch Experimente und theoretische Analysen gelangen wir zu einigen wichtigen Erkenntnissen und Schlussfolgerungen: (i) QK-Norm ist für stabiles Training bei vielen Token pro Schritt notwendig, (ii) Quantisierungsfehler entstehen hauptsächlich aus dem Score-Gradienten dS im Rückwärtsdurchlauf, (iii) eine Reduzierung der Token pro Schritt ermöglicht es SageBwd, die FPA-Leistung im Vorabtraining zu erreichen, und (iv) K-Glättung bleibt für die Trainingsstabilität entscheidend, während Q-Glättung während des Vorabtrainings nur begrenzten Nutzen bietet.
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.