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

📄 arXiv: 2310.16795v1 📥 PDF

作者: Elias Frantar, Dan Alistarh

分类: cs.LG

发布日期: 2023-10-25

💡 一句话要点

提出QMoE以解决大规模模型内存占用问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家 模型压缩 推理优化 大规模模型 GPU加速

📋 核心要点

  1. 现有的MoE模型如SwitchTransformer-c2048在推理时需要巨大的内存，限制了其实际应用。
  2. QMoE通过一种新的压缩算法，将MoE模型的参数压缩到每个参数不足1位，结合定制GPU解码内核以提高推理效率。
  3. QMoE能够将1.6万亿参数的模型压缩至160GB，且在运行时仅增加5%的开销，显著提升了模型的可用性。

📝 摘要（中文）

混合专家（MoE）架构通过稀疏路由为大型语言模型（LLMs）提供了一种通用解决方案，尽管其参数数量庞大，例如SwitchTransformer-c2048模型有1.6万亿个参数，运行时需要3.2TB的加速器内存，导致实际部署困难且成本高昂。本文提出了一种新的压缩和执行框架QMoE，能够将万亿参数的MoE模型压缩到每个参数不足1位，且与定制的GPU解码内核共同设计，以实现高效的端到端压缩推理，且相较于未压缩执行仅有轻微的运行时开销。具体而言，QMoE可以将1.6万亿参数的SwitchTransformer-c2048模型压缩到不足160GB（20倍压缩，0.8位每参数），且仅有轻微的准确性损失，压缩过程在单个GPU上不到一天完成。这使得在经济实惠的普通硬件上执行万亿参数模型成为可能。

🔬 方法详解

问题定义：本文旨在解决大规模MoE模型在推理时的内存占用问题，现有方法如SwitchTransformer-c2048需要3.2TB内存，导致实际部署困难。

核心思路：QMoE通过设计一种新的压缩算法，将每个参数的存储需求降低到不足1位，同时结合定制的GPU解码内核，以实现高效的推理过程。

技术框架：QMoE的整体架构包括压缩算法和解码内核两个主要模块。压缩算法负责将模型参数压缩，而解码内核则优化了压缩后的模型在推理时的执行效率。

关键创新：QMoE的主要创新在于其压缩算法能够在保持模型性能的同时，实现高达20倍的参数压缩，这在现有技术中尚属首次。

关键设计：在设计中，QMoE采用了特定的损失函数和网络结构，以确保在压缩过程中尽量减少准确性损失，同时优化了GPU解码内核以适应压缩后的模型结构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

QMoE在实验中成功将1.6万亿参数的SwitchTransformer-c2048模型压缩至160GB，达到了20倍的压缩率，且在推理时仅增加了5%的运行时开销。这一成果显著提升了大规模模型的可用性和经济性。

🎯 应用场景

QMoE的研究成果具有广泛的应用潜力，尤其是在需要处理大规模语言模型的场景中，如自然语言处理、机器翻译和对话系统等。通过降低内存需求，QMoE使得这些复杂模型能够在普通硬件上运行，从而降低了部署成本，推动了AI技术的普及与应用。

📄 摘要（原文）

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.