QuEST: Treinamento Estável de LLMs com Pesos e Ativações 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

Resumo

Uma abordagem para reduzir os altos custos dos grandes modelos de linguagem (LLMs) é o uso de representações quantizadas ou esparsas para treinamento ou implantação. Embora os métodos de compressão pós-treinamento sejam muito populares, a questão de obter modelos comprimidos ainda mais precisos por meio do treinamento direto sobre tais representações, ou seja, Treinamento Consciente de Quantização (QAT), ainda está em aberto: por exemplo, um estudo recente (arXiv:2411.04330v2) determinou a largura de bits "ótima" na qual os modelos podem ser treinados usando QAT, mantendo a precisão competitiva com a precisão padrão FP16/BF16, em 8-bits para pesos e ativações. Avançamos este estado-da-arte por meio de um novo método chamado QuEST, que é competitivo em Pareto com FP16, ou seja, fornece melhor precisão com tamanho de modelo menor, treinando modelos com pesos e ativações em 4-bits ou menos. Além disso, o QuEST permite treinamento estável com pesos e ativações de 1-bit. O QuEST alcança isso melhorando dois aspectos-chave dos métodos QAT: (1) quantização precisa e rápida das distribuições (contínuas) de pesos e ativações por meio de normalização de Hadamard e ajuste MSE-ótimo; (2) um novo estimador de gradiente de confiança baseado na ideia de minimizar explicitamente o erro entre o gradiente ruidoso calculado sobre estados quantizados e o gradiente de precisão total "verdadeiro" (mas desconhecido). Experimentos em arquiteturas do tipo Llama mostram que o QuEST induz leis de escalonamento estáveis em toda a gama de precisões suportadas pelo hardware e pode ser estendido a representações esparsas. Fornecemos suporte de kernel GPU mostrando que os modelos produzidos pelo QuEST podem ser executados de forma eficiente. Nosso código está disponível em 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.