QuEST : Entraînement stable des LLM avec des poids et des activations à 1 bit

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

February 7, 2025

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

cs.AI

Résumé

Une approche pour réduire les coûts massifs des grands modèles de langage (LLM) est l'utilisation de représentations quantifiées ou éparses pour l'entraînement ou le déploiement. Alors que les méthodes de compression post-entraînement sont très populaires, la question d'obtenir des modèles compressés encore plus précis en entraînant directement sur de telles représentations, c'est-à-dire l'Entraînement Sensible à la Quantification (QAT), reste ouverte : par exemple, une étude récente (arXiv:2411.04330v2) a déterminé la largeur de bits "optimale" à laquelle les modèles peuvent être entraînés en utilisant QAT, tout en restant compétitifs en termes de précision avec la précision FP16/BF16 standard, à 8 bits pour les poids et les activations. Nous faisons progresser cet état de l'art grâce à une nouvelle méthode appelée QuEST, qui est compétitive selon le critère de Pareto avec FP16, c'est-à-dire qu'elle offre une meilleure précision avec une taille de modèle plus petite, tout en entraînant des modèles avec des poids et des activations sur 4 bits ou moins. De plus, QuEST permet un entraînement stable avec des poids et des activations sur 1 bit. QuEST parvient à cela en améliorant deux aspects clés des méthodes QAT : (1) la quantification précise et rapide des distributions (continues) des poids et des activations via la normalisation de Hadamard et un ajustement optimal MSE ; (2) un nouvel estimateur de gradient de confiance basé sur l'idée de minimiser explicitement l'erreur entre le gradient bruité calculé sur des états quantifiés et le gradient de "vraie" précision (mais inconnu). Des expériences sur des architectures de type Llama montrent que QuEST induit des lois d'échelle stables sur l'ensemble de la gamme de précisions supportées par le matériel, et peut être étendu à des représentations éparses. Nous fournissons un support de noyau GPU montrant que les modèles produits par QuEST peuvent être exécutés efficacement. Notre code est disponible sur https://github.com/IST-DASLab/QuEST.

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.