QMoE : Compression pratique en dessous de 1 bit pour des modèles à mille milliards de paramètres

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

October 25, 2023

Auteurs: Elias Frantar, Dan Alistarh

cs.AI

Résumé

Les architectures Mixture-of-Experts (MoE) offrent une solution générale aux coûts d'inférence élevés des grands modèles de langage (LLMs) grâce à un routage parcimonieux, permettant d'obtenir des modèles plus rapides et plus précis, au prix d'un nombre massif de paramètres. Par exemple, le modèle SwitchTransformer-c2048 possède 1,6 trillion de paramètres, nécessitant 3,2 To de mémoire d'accélérateur pour fonctionner efficacement, ce qui rend le déploiement pratique difficile et coûteux. Dans cet article, nous présentons une solution à ce problème de mémoire, sous la forme d'un nouveau cadre de compression et d'exécution appelé QMoE. Plus précisément, QMoE comprend un algorithme scalable qui compresse avec précision les MoE à trillion de paramètres à moins de 1 bit par paramètre, dans un format personnalisé co-conçu avec des noyaux de décodage GPU sur mesure pour faciliter une inférence compressée efficace de bout en bout, avec des surcharges d'exécution mineures par rapport à une exécution non compressée. Concrètement, QMoE peut compresser le modèle SwitchTransformer-c2048 à 1,6 trillion de paramètres à moins de 160 Go (compression de 20x, 0,8 bit par paramètre) avec une perte de précision minime, en moins d'une journée sur un seul GPU. Cela permet, pour la première fois, l'exécution d'un modèle à trillion de paramètres sur du matériel abordable, comme un seul serveur équipé de 4x NVIDIA A6000 ou 8x NVIDIA 3090 GPUs, avec une surcharge d'exécution inférieure à 5 % par rapport à une inférence non compressée idéale. Le code source et les modèles compressés sont disponibles sur 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.