QuEST: Stabiles Training von LLMs mit 1-Bit Gewichten und Aktivierungen

QuEST: Stable Training of LLMs with 1-Bit Weights and Activations

February 7, 2025

Autoren: Andrei Panferov, Jiale Chen, Soroush Tabesh, Roberto L. Castro, Mahdi Nikdan, Dan Alistarh

cs.AI

Zusammenfassung

Ein Ansatz zur Reduzierung der massiven Kosten großer Sprachmodelle (LLMs) besteht in der Verwendung von quantisierten oder spärlichen Darstellungen für das Training oder den Einsatz. Während post-Training-Kompressionsmethoden sehr beliebt sind, ist die Frage, ob noch genauere komprimierte Modelle durch direktes Training über solchen Darstellungen erzielt werden können, d.h. Quantisierungs-bewusstes Training (QAT), noch offen: Zum Beispiel hat eine kürzlich durchgeführte Studie (arXiv:2411.04330v2) die "optimale" Bit-Breite ermittelt, mit der Modelle unter Verwendung von QAT trainiert werden können, während sie konkurrenzfähig in Bezug auf Genauigkeit mit der Standard-FP16/BF16-Präzision bleiben, bei 8-Bit-Gewichten und Aktivierungen. Wir bringen diesen Stand der Technik mit einer neuen Methode namens QuEST voran, die Pareto-wettbewerbsfähig mit FP16 ist, d.h. sie bietet eine bessere Genauigkeit bei geringerer Modellgröße, während Modelle mit Gewichten und Aktivierungen in 4-Bit oder weniger trainiert werden. Darüber hinaus ermöglicht QuEST stabiles Training mit 1-Bit-Gewichten und Aktivierungen. QuEST erreicht dies durch Verbesserung zweier Schlüsselaspekte von QAT-Methoden: (1) genaue und schnelle Quantisierung der (kontinuierlichen) Verteilungen von Gewichten und Aktivierungen durch Hadamard-Normalisierung und MSE-optimale Anpassung; (2) ein neuer Vertrauensgradientenschätzer, der auf der Idee basiert, den Fehler zwischen dem über quantisierte Zustände berechneten rauschigen Gradienten und dem "wahren" (aber unbekannten) Vollpräzisionsgradienten explizit zu minimieren. Experimente an Llama-ähnlichen Architekturen zeigen, dass QuEST stabile Skalierungsgesetze über den gesamten Bereich der hardwareunterstützten Präzisionen induziert und auf spärliche Darstellungen erweitert werden kann. Wir bieten GPU-Kernelunterstützung an, die zeigt, dass von QuEST erzeugte Modelle effizient ausgeführt werden können. Unser Code ist unter https://github.com/IST-DASLab/QuEST verfügbar.

English

One approach to reducing the massive costs of large language models (LLMs) is the use of quantized or sparse representations for training or deployment. While post-training compression methods are very popular, the question of obtaining even more accurate compressed models by directly training over such representations, i.e., Quantization-Aware Training (QAT), is still open: for example, a recent study (arXiv:2411.04330v2) put the "optimal" bit-width at which models can be trained using QAT, while staying accuracy-competitive with standard FP16/BF16 precision, at 8-bits weights and activations. We advance this state-of-the-art via a new method called QuEST, which is Pareto-competitive with FP16, i.e., it provides better accuracy at lower model size, while training models with weights and activations in 4-bits or less. Moreover, QuEST allows stable training with 1-bit weights and activations. QuEST achieves this by improving two key aspects of QAT methods: (1) accurate and fast quantization of the (continuous) distributions of weights and activations via Hadamard normalization and MSE-optimal fitting; (2) a new trust gradient estimator based on the idea of explicitly minimizing the error between the noisy gradient computed over quantized states and the "true" (but unknown) full-precision gradient. Experiments on Llama-type architectures show that QuEST induces stable scaling laws across the entire range of hardware-supported precisions, and can be extended to sparse representations. We provide GPU kernel support showing that models produced by QuEST can be executed efficiently. Our code is available at https://github.com/IST-DASLab/QuEST.