QuEST: Стабильное обучение LLM с весами и активациями 1 бит.

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

February 7, 2025

Авторы: Andrei Panferov, Jiale Chen, Soroush Tabesh, Roberto L. Castro, Mahdi Nikdan, Dan Alistarh

cs.AI

Аннотация

Один из подходов к снижению огромных затрат на большие языковые модели (LLM) заключается в использовании квантованных или разреженных представлений для обучения или развертывания. Хотя методы сжатия после обучения очень популярны, вопрос о получении еще более точных сжатых моделей путем прямого обучения на таких представлениях, то есть обучение с учетом квантования (QAT), остается открытым: например, недавнее исследование (arXiv:2411.04330v2) определило "оптимальное" количество битов, с которыми модели могут быть обучены с использованием QAT, оставаясь конкурентоспособными по точности с стандартной точностью FP16/BF16, на уровне весов и активаций в 8 бит. Мы продвигаем этот передовой метод с помощью нового метода под названием QuEST, который конкурентоспособен по Парето с FP16, то есть обеспечивает лучшую точность при меньшем размере модели, обучая модели с весами и активациями в 4 бита или менее. Более того, QuEST позволяет стабильное обучение с весами и активациями в 1 бит. QuEST достигает этого путем улучшения двух ключевых аспектов методов QAT: (1) точного и быстрого квантования (непрерывных) распределений весов и активаций с помощью нормализации Хадамарда и оптимальной подгонки по MSE; (2) нового оценщика доверия градиента на основе идеи явного минимизирования ошибки между шумным градиентом, вычисленным на квантованных состояниях, и "истинным" (но неизвестным) градиентом полной точности. Эксперименты на архитектурах типа Лама показывают, что QuEST вызывает стабильные законы масштабирования по всему диапазону поддерживаемых аппаратных точностей и может быть расширен на разреженные представления. Мы предоставляем поддержку ядра GPU, показывающую, что модели, созданные с помощью QuEST, могут быть эффективно выполнены. Наш код доступен по адресу 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.