QuEST: Entrenamiento Estable de LLMs con Pesos y Activaciones de 1 Bit

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

February 7, 2025

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

cs.AI

Resumen

Un enfoque para reducir los enormes costos de los grandes modelos de lenguaje (LLMs) es el uso de representaciones cuantizadas o dispersas para entrenamiento o implementación. Si bien los métodos de compresión posteriores al entrenamiento son muy populares, la cuestión de obtener modelos comprimidos aún más precisos mediante el entrenamiento directo sobre tales representaciones, es decir, el Entrenamiento Consciente de la Cuantización (QAT), sigue abierta: por ejemplo, un estudio reciente (arXiv:2411.04330v2) estableció el ancho de bits "óptimo" en el que los modelos pueden ser entrenados utilizando QAT, manteniendo la precisión competitiva con la precisión estándar FP16/BF16, en 8 bits para pesos y activaciones. Avanzamos en este estado del arte a través de un nuevo método llamado QuEST, que es competitivo en el sentido de Pareto con FP16, es decir, proporciona una mejor precisión con un tamaño de modelo más pequeño, al entrenar modelos con pesos y activaciones en 4 bits o menos. Además, QuEST permite un entrenamiento estable con pesos y activaciones de 1 bit. QuEST logra esto mejorando dos aspectos clave de los métodos QAT: (1) cuantización precisa y rápida de las distribuciones (continuas) de pesos y activaciones a través de la normalización de Hadamard y el ajuste óptimo de ECM; (2) un nuevo estimador de gradiente de confianza basado en la idea de minimizar explícitamente el error entre el gradiente ruidoso calculado sobre estados cuantizados y el gradiente de precisión completa "verdadero" (pero desconocido). Experimentos en arquitecturas tipo Llama muestran que QuEST induce leyes de escalado estables en todo el rango de precisiones admitidas por el hardware, y puede extenderse a representaciones dispersas. Proporcionamos soporte de kernel de GPU que muestra que los modelos producidos por QuEST pueden ejecutarse eficientemente. Nuestro código está disponible en 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.