[论文解读] QMoE: Practical Sub-1-Bit Compression of Trillion-Parameter Models

Elias Frantar, Dan Alistarh|arXiv (Cornell University)|Oct 25, 2023

Machine Learning and Algorithms被引用 6

一句话总结

QMoE 提出了一种可扩展的压缩框架，实现了万亿参数专家混合模型（MoE）的亚1比特量化（每参数0.8比特），将 SwitchTransformer-c2048 等模型的模型大小从 3.2TB（bfloat16）压缩至 160GB 以下，且精度损失极小。该框架通过自定义压缩格式和优化的 GPU 内核，实现了在单台通用 GPU 服务器上的高效端到端推理，运行时开销低于 5%。

ABSTRACT

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.

研究动机与目标

• 解决万亿参数专家混合（MoE）模型带来的高内存和计算成本问题，该问题限制了其实际部署和研究访问。
• 克服现有训练后量化技术在每参数低于1比特时难以实现高精度压缩的挑战。
• 设计一种可扩展的、无需微调的压缩流水线，在保持模型精度的同时，实现在通用硬件上的高效设备端推理。
• 开发自定义压缩格式和 GPU 优化的解码内核，以实现压缩 MoE 模型的快速、低开销推理。
• 证明在单台服务器上使用 4 张 A6000 或 8 张 3090 GPU 运行如 SwitchTransformer-c2048 这类大规模 MoE 模型的可行性，从而推动更广泛的研究与部署。

提出的方法

• 提出一种可扩展的、无需微调的压缩算法，利用专为低比特稀疏性和加速器效率设计的自定义压缩格式，将 MoE 权重量化至每参数低于 1 比特。
• 设计专用 GPU 内核，实现在推理过程中对权重进行实时解压缩，通过利用低比特量化权重中的天然稀疏性，最大限度降低运行时开销。
• 协同设计压缩格式与解码内核，确保与现有基于 HuggingFace 的推理流水线及 GPU 内存访问模式的兼容性。
• 将该压缩流水线应用于 1.6 万亿参数的 SwitchTransformer-c2048 模型，仅使用单张 GPU 在一天内完成压缩，压缩率高达 20 倍（从 3.2TB 压缩至 <160GB）。
• 优化量化过程以保持模型精度，仅导致预训练验证损失和零样本评估损失的轻微增加。
• 将该框架集成至 HuggingFace 生态系统，确保研究人员和实践者能够广泛使用并轻松部署。

Figure 1: Example of an MoE Transformer block. Each token is routed to a different fully-connected (FC) block.

实验结果

研究问题

• RQ1 能否在不造成显著精度下降的前提下，将万亿参数 MoE 模型压缩至每参数低于 1 比特？
• RQ2 是否可以使用无需微调的、训练后量化技术实现 MoE 模型的亚1比特压缩？
• RQ3 自定义压缩格式与 GPU 优化内核能否实现压缩 MoE 模型在通用硬件上的高效、低开销推理？
• RQ4 与高端硬件上的理想无压缩执行相比，压缩推理的运行时开销是多少？
• RQ5 压缩后的模型是否能够实现在单台服务器上使用消费级 GPU 的端到端部署与执行？

主要发现

• 使用 QMoE 将 1.6 万亿参数的 SwitchTransformer-c2048 模型从 3.2TB（bfloat16）压缩至 160GB 以下，实现了 20 倍的压缩率，每参数为 0.8 比特。
• 该压缩在单张 GPU 上仅用不到一天时间完成，采用无需微调的、训练后量化流水线，且在预训练验证和零样本基准测试中精度损失极小。
• 在配备 8 张 NVIDIA RTX 3090 或 4 张 NVIDIA A6000 GPU 的单台服务器上，成功实现了压缩模型的端到端推理，证明了其在通用硬件上的实际可部署性。
• 与理想化的无压缩执行相比，压缩推理的运行时开销低于 5%，而后者需要约 20 倍的 GPU 数量。
• 自定义压缩格式与 GPU 内核实现了高效的实时解码，避免了低效的离线加载，保持了高吞吐量。
• QMoE 框架为开源，并已集成至 HuggingFace，显著降低了大规模 MoE 模型研究与部署的门槛。

Figure 2: Illustration of the offloading execution for the sparse part of a Transformer block. An expert $E_{2}$ and its corresponding input tokens $X_{E}$ are fetched to GPU memory to produce $E^{\prime}_{2}$ , which together with the corresponding outputs $Y_{E}$ are written back to CPU again.