QMoE: Compressão Prática de Modelos com Trilhões de Parâmetros em Sub-1-Bit

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

October 25, 2023

Autores: Elias Frantar, Dan Alistarh

cs.AI

Resumo

As arquiteturas Mixture-of-Experts (MoE) oferecem uma solução geral para os altos custos de inferência de grandes modelos de linguagem (LLMs) por meio de roteamento esparso, proporcionando modelos mais rápidos e precisos, ao custo de contagens massivas de parâmetros. Por exemplo, o modelo SwitchTransformer-c2048 possui 1,6 trilhões de parâmetros, exigindo 3,2TB de memória de acelerador para funcionar de forma eficiente, o que torna a implantação prática desafiadora e cara. Neste artigo, apresentamos uma solução para esse problema de memória, na forma de um novo framework de compressão e execução chamado QMoE. Especificamente, o QMoE consiste em um algoritmo escalável que comprime com precisão MoEs com trilhões de parâmetros para menos de 1 bit por parâmetro, em um formato personalizado co-desenhado com kernels de decodificação GPU sob medida para facilitar a inferência comprimida eficiente de ponta a ponta, com pequenos sobrecustos de tempo de execução em relação à execução não comprimida. Concretamente, o QMoE pode comprimir o modelo SwitchTransformer-c2048 de 1,6 trilhão de parâmetros para menos de 160GB (compressão de 20x, 0,8 bits por parâmetro) com apenas uma pequena perda de precisão, em menos de um dia em uma única GPU. Isso permite, pela primeira vez, a execução de um modelo com trilhões de parâmetros em hardware acessível, como um único servidor com 4x NVIDIA A6000 ou 8x NVIDIA 3090 GPUs, com menos de 5% de sobrecarga de tempo de execução em relação à inferência ideal não comprimida. O código-fonte e os modelos comprimidos estão disponíveis em 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.