QMoE: Compressione pratica sotto il bit per modelli con trilioni di parametri

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

October 25, 2023

Autori: Elias Frantar, Dan Alistarh

cs.AI

Abstract

Le architetture Mixture-of-Experts (MoE) offrono una soluzione generale agli elevati costi di inferenza dei grandi modelli linguistici (LLM) attraverso il routing sparso, consentendo modelli più veloci e accurati, sebbene a costo di un numero enorme di parametri. Ad esempio, il modello SwitchTransformer-c2048 ha 1,6 trilioni di parametri, richiedendo 3,2 TB di memoria degli acceleratori per funzionare in modo efficiente, rendendo la distribuzione pratica complessa e costosa. In questo articolo, presentiamo una soluzione a questo problema di memoria, sotto forma di un nuovo framework di compressione ed esecuzione chiamato QMoE. Nello specifico, QMoE consiste in un algoritmo scalabile che comprime accuratamente MoE con trilioni di parametri a meno di 1 bit per parametro, in un formato personalizzato co-progettato con kernel di decodifica GPU su misura per facilitare un'inferenza compressa end-to-end efficiente, con sovraccarichi di runtime minimi rispetto all'esecuzione non compressa. In concreto, QMoE può comprimere il modello SwitchTransformer-c2048 da 1,6 trilioni di parametri a meno di 160 GB (compressione 20x, 0,8 bit per parametro) con una perdita di accuratezza minima, in meno di un giorno su una singola GPU. Ciò consente, per la prima volta, l'esecuzione di un modello con trilioni di parametri su hardware consumer accessibile, come un singolo server con 4x NVIDIA A6000 o 8x NVIDIA 3090 GPU, con un sovraccarico di runtime inferiore al 5% rispetto all'inferenza ideale non compressa. Il codice sorgente e i modelli compressi sono disponibili su 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.