[論文レビュー] Fast Inference of Mixture-of-Experts Language Models with Offloading

Artyom Eliseev, D. Peter Mazur|arXiv (Cornell University)|Dec 28, 2023

Topic Modeling被引用数 7

ひとこと要約

本論文は、専用のエキスパートアクセスパターンと予測的プリロードを活用することで、コンsumerハードウェア上で大規模なMixture-of-Experts (MoE)言語モデルの推論を高速化するMoE特化型のオフロード戦略を提案する。混合量子化（エキスパートには2–3ビット、アテンション層には4–16ビット）とLRUキャッシュ、および予測的エキスパートロードを組み合わせることで、デスクトップグレードのGPUおよび無料のGoogle Colabインスタンス（T4）でも、Mixtral-8x7B-Instructを1秒あたり2–3トークンの速度でインタラクティブ推論が可能になる。

ABSTRACT

With the widespread adoption of Large Language Models (LLMs), many deep learning practitioners are looking for strategies of running these models more efficiently. One such strategy is to use sparse Mixture-of-Experts (MoE) - a type of model architectures where only a fraction of model layers are active for any given input. This property allows MoE-based language models to generate tokens faster than their dense counterparts, but it also increases model size due to having multiple experts. Unfortunately, this makes state-of-the-art MoE language models difficult to run without high-end GPUs. In this work, we study the problem of running large MoE language models on consumer hardware with limited accelerator memory. We build upon parameter offloading algorithms and propose a novel strategy that accelerates offloading by taking advantage of innate properties of MoE LLMs. Using this strategy, we build can run Mixtral-8x7B with mixed quantization on desktop hardware and free-tier Google Colab instances.

研究の動機と目的

• GPUメモリが限られたコンsumerグレードのハードウェアで、Mixtral-8x7B-Instructのような大規模なMoE言語モデルのインタラクティブ推論を可能にすること。
• 不規則なエキスパートアクセスパターンと高いメモリ帯域幅要件のため、Naiveなパラメータオフロードが非効率である点を是正すること。
• エキスパートの再利用やゲーティング関数の予測可能性といった内在的なアクセスの規則性を活用することで、MoEアーキテクチャに最適化されたオフロードを実現すること。
• 量子化と知的なオフロードの組み合わせにより、低価格のハードウェア、特にT4搭載の無料のGoogle Colabでも実用的なパフォーマンスを達成すること。
• 高価なデータセンターグレードGPUを必要とせず、ローカルで最新のMoEモデルを実行可能な、デプロイ可能でオープンソースのソリューションを提供すること。

提案手法

• 頻繁に使用されるエキスパートのGPU-RAM間の繰り返し転送を削減するため、LRUキャッシュを活用したMoEに特化したオフロードパイプラインを設計する。
• ゲーティング関数の出力を基に、計算と重ねてエキスパートをプリロードする予測的エキスパートプリロードを実装し、無駄な待機時間を削減する。
• 混合量子化を統合：MoEエキスパートには2–3ビットのHQQ量子化、アテンション層、埋め込み、正規化には4–16ビットの精度を適用する。
• GPUからCPUメモリにエキスパート重みを動的にオフロードするハイブリッド推論エンジンを採用し、ピン留めメモリと非同期転送による効率的なメモリ管理を実現する。
• 初期層の隠れ状態が、後続層で必要なエキスパートを予測できることを活用し、予測的ローディング戦略を可能にする。
• PCIeボトルネックが深刻なシステムを最適化するため、オフロード頻度を最小限に抑え、キャッシュとプリロードによるデータ再利用を最大化する。

Figure 1 : An example of expert loading pattern in Mixtral-8x7B-Instruct for select layers. Blue cells indicate that a certain expert was active when encoding a certain token; deeper blue indicates higher gating weight. Small gray squares show which experts are cached with an LRU cache for $k{=}2$ .

実験結果

リサーチクエスチョン

• RQ1 MoEモデルにおけるエキスパートアクセスパターンは、どのように活用することで、メモリ制限のあるハードウェア上でのオフロードオーバーヘッドを低減し、推論速度を向上させられるか？
• RQ2 コンsumerGPU上で動作する際、MoE層の量子化精度とモデル性能の最適なトレードオフは何か？
• RQ3 LRUキャッシュと予測的プリロードは、Naiveなオフロードと比較して、MoE推論におけるGPU-RAM通信をどの程度削減できるか？
• RQ4 オフロードと量子化のみを用いて、無料のColabやデスクトップグレードGPUでMixtral-8x7B-Instructをインタラクティブ（1秒あたり2–3トークン）に実行できるか？
• RQ5 低メモリ環境下において、標準的なオフロードフレームワーク（例：Hugging Face Accelerate）と比較して、MoE特化型オフロードのパフォーマンスはどの程度優れているか？

主な発見

• 提案されたMoE特化型オフロード戦略により、RTX 3060（12GB）、RTX 3080 Mobile（16GB）、T4（16GB）GPU上でも、Mixtral-8x7B-Instructの推論速度が1秒あたり2–3トークンに達した。
• 無料のGoogle Colab（T4搭載）上では、2ビット量子化されたエキスパートと完全なオフロードに加えプリロードとLRUキャッシュを適用した結果、1秒あたり2.092トークンの生成速度が達成された。
• LRUキャッシュと予測的プリロードの組み合わせにより、RTX 3060ではNaiveなオフロードと比較して推論時間が最大15%短縮された。
• エキスパートに3ビットのHQQ量子化、アテンション層に4ビットを適用した場合、周辺パープレクサリティの低下が最小限で、性能と精度のトレードオフに最も適した結果が得られた。
• 特にPCIe制限が深刻な環境では、Naiveなオフロード（例：Hugging Face Accelerate経由）と比較して、最大2.5倍の速度向上が達成された。
• VRAMが低いにもかかわらず、RTX 3060デスクトップはより高価なRTX 3080 Mobileとほぼ同等の性能を示しており、ホストからデバイスへの帯域幅が主なボトルネックであることが示された。

Figure 2 : (left) LRU cache hit ratio for different cache size $k$ ; (right) speculative loading recall when pre-loading a different number of experts. Regular lines represent loading 1 layer ahead; dashed line stands for 2 layers ahead; dotted line is 10 layers ahead.