QMoE: Compresión práctica de modelos de billones de parámetros en menos de 1 bit

QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models

October 25, 2023

Autores: Elias Frantar, Dan Alistarh

cs.AI

Resumen

Las arquitecturas Mixture-of-Experts (MoE) ofrecen una solución general a los altos costos de inferencia de los modelos de lenguaje grandes (LLMs) mediante el enrutamiento disperso, permitiendo modelos más rápidos y precisos, aunque a costa de un gran número de parámetros. Por ejemplo, el modelo SwitchTransformer-c2048 tiene 1.6 billones de parámetros, requiriendo 3.2TB de memoria en aceleradores para funcionar de manera eficiente, lo que hace que su implementación práctica sea desafiante y costosa. En este artículo, presentamos una solución a este problema de memoria, en forma de un nuevo marco de compresión y ejecución llamado QMoE. Específicamente, QMoE consiste en un algoritmo escalable que comprime con precisión MoEs de billones de parámetros a menos de 1 bit por parámetro, en un formato personalizado co-diseñado con núcleos de decodificación GPU específicos para facilitar una inferencia comprimida eficiente de extremo a extremo, con sobrecargas de tiempo de ejecución mínimas en comparación con la ejecución sin comprimir. Concretamente, QMoE puede comprimir el modelo SwitchTransformer-c2048 de 1.6 billones de parámetros a menos de 160GB (compresión de 20x, 0.8 bits por parámetro) con solo una mínima pérdida de precisión, en menos de un día en una sola GPU. Esto permite, por primera vez, la ejecución de un modelo de billones de parámetros en hardware asequible, como un solo servidor con 4x NVIDIA A6000 o 8x NVIDIA 3090 GPUs, con menos del 5% de sobrecarga de tiempo de ejecución en comparación con la inferencia ideal sin comprimir. El código fuente y los modelos comprimidos están disponibles en github.com/IST-DASLab/qmoe.

English

Mixture-of-Experts (MoE) architectures offer a general solution to the high inference costs of large language models (LLMs) via sparse routing, bringing faster and more accurate models, at the cost of massive parameter counts. For example, the SwitchTransformer-c2048 model has 1.6 trillion parameters, requiring 3.2TB of accelerator memory to run efficiently, which makes practical deployment challenging and expensive. In this paper, we present a solution to this memory problem, in form of a new compression and execution framework called QMoE. Specifically, QMoE consists of a scalable algorithm which accurately compresses trillion-parameter MoEs to less than 1 bit per parameter, in a custom format co-designed with bespoke GPU decoding kernels to facilitate efficient end-to-end compressed inference, with minor runtime overheads relative to uncompressed execution. Concretely, QMoE can compress the 1.6 trillion parameter SwitchTransformer-c2048 model to less than 160GB (20x compression, 0.8 bits per parameter) at only minor accuracy loss, in less than a day on a single GPU. This enables, for the first time, the execution of a trillion-parameter model on affordable commodity hardware, like a single server with 4x NVIDIA A6000 or 8x NVIDIA 3090 GPUs, at less than 5% runtime overhead relative to ideal uncompressed inference. The source code and compressed models are available at github.com/IST-DASLab/qmoe.