[論文レビュー] 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は、SwitchTransformer-c2048などのトロイオンパラメータのMixture-of-Experts（MoE）モデルを1パラメータあたり0.8ビット未塔のサブ1ビット量子化によりスケーラブルな圧縮フレームワークを実現し、モデルサイズをbfloat16の3.2TBから160GB未塔に削減しながら、精度の損失を最小限に抑えました。カスタム圧縮フォーマットと最適化されたGPUカーネルを活用し、1台のコンsumer 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.

研究の動機と目的

• トロイオンパラメータのMixture-of-Experts（MoE）モデルの高いメモリおよび計算コストに対処し、実用的導入と研究アクセスを制限する要因を解消すること。
• 既存のトレーニング後量子化技術では達成が難しい、1パラメータあたり1ビット未塔の高精度圧縮を達成すること。
• モデル精度を維持しながら、再訓練を不要とするスケーラブルな圧縮パイプラインを設計し、コンsumerハードウェア上で効率的なオンデバイス推論を可能にすること。
• 高速で低オーバーヘッドの推論を実現するため、カスタム圧縮フォーマットとGPU最適化されたデコードカーネルを開発すること。
• 4枚のA6000または8枚の3090 GPUを搭載した1台のサーバーで、SwitchTransformer-c2048のような大規模なMoEモデルを実行可能にし、研究および導入の広がりを促進すること。

提案手法

• 低ビット幅スパーシティとアクセラレータ効率に最適化されたカスタム圧縮フォーマットを設計し、その上でMoE重みを1パラメータあたり1ビット未塔に量子化するスケーラブルで再訓練不要の圧縮アルゴリズムを提案すること。
• 低ビット量子化された重みに内在する自然なスパーシティを活用し、リアルタイムでの復元と推論を実現する独自のGPUカーネルを設計することで、実行時間のオーバーヘッドを最小限に抑えること。
• 圧縮フォーマットとデコードカーネルを共同で設計し、既存のHuggingFaceベースの推論パイプラインおよびGPUメモリアクセスパターンと互換性を持たせること。
• 1.6兆パラメータのSwitchTransformer-c2048モデルに圧縮パイプラインを適用し、単一のGPUで1日未塔の時間で3.2TB（bfloat16）から160GB未塔への20倍圧縮を達成したこと。
• 量子化プロセスを最適化し、微増の事前学習検証損失とゼロショット評価損失の増加でモデル精度を維持したこと。
• 研究者および実務家が広く利用可能で容易にデプロイ可能なように、HuggingFaceエコシステムにフレームワークを統合したこと。

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

実験結果

リサーチクエスチョン

• RQ1 1パラメータあたり1ビット未塔に圧縮されたトロイオンパラメータのMoEモデルは、顕著な精度劣化なしに実現可能か？
• RQ2 再訓練を不要とするトレーニング後量子化技術で、MoEモデルのサブ1ビット圧縮が可能か？
• RQ3 カスタム圧縮フォーマットとGPU最適化カーネルにより、コンsumerハードウェア上で効率的かつ低オーバーヘッドの推論が可能か？
• RQ4 高価なハードウェアで理想化された非圧縮実行と比較して、圧縮推論の実行時間オーバーヘッドはどの程度か？
• RQ5 圧縮されたモデルは、コンsumerグレードのGPUを搭載した1台のサーバー上でエンドツーエンドにデプロイ・実行可能か？

主な発見

• 1.6兆パラメータのSwitchTransformer-c2048モデルは、QMoEによりbfloat16の3.2TBから160GB未塔に圧縮され、20倍の圧縮率と1パラメータあたり0.8ビットを達成した。
• 単一のGPUで1日未塔の時間で再訓練不要のトレーニング後量子化パイプラインを用いて圧縮が達成され、事前学習検証およびゼロショットベンチマークでの精度損失は最小限に抑えられた。
• 8枚のNVIDIA RTX 3090または4枚のNVIDIA A6000 GPUを搭載した1台のサーバー上で、圧縮モデルのエンドツーエンド推論が正常に実行され、コンsumerハードウェアへの実用的導入を実証した。
• 圧縮推論の実行時間オーバーヘッドは、約20倍のGPUを必要とする理想化された非圧縮実行と比較して5%未塔であった。
• カスタム圧縮フォーマットと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.