[論文レビュー] Mixture of Quantized Experts (MoQE): Complementary Effect of Low-bit Quantization and Robustness

Young J. Kim, Raffy Fahim+1 more|arXiv (Cornell University)|Oct 3, 2023

Multimodal Machine Learning ApplicationsComputer Science被引用数 3

ひとこと要約

本稿では、Mixture-of-Quantized-Experts (MoQE) を提案する。MoQE は、Mixture-of-Experts (MoE) モデルのエキスパート重みにのみ、超低ビット (2ビット) の重みのみを用いた量子化手法である。エキスパート層の量子化に対する固有の耐性を活用することで、MoE モデルのモデルサイズを 79.6% 減少させ、A100 GPU 上で顕著な精度損失なしに 1.24× の推論速度向上を達成した。Quantization Aware Training (QAT) を用いた 2 ビット量子化において、密度型モデルを上回る性能を示した。

ABSTRACT

Large Mixture of Experts (MoE) models could achieve state-of-the-art quality on various language tasks, including machine translation task, thanks to the efficient model scaling capability with expert parallelism. However, it has brought a fundamental issue of larger memory consumption and increased memory bandwidth bottleneck at deployment time. In this paper, we propose Mixture of Quantized Experts (MoQE) which is a simple weight-only quantization method applying ultra low-bit down to 2-bit quantizations only to expert weights for mitigating the increased memory and latency issues of MoE models. We show that low-bit quantization together with the MoE architecture delivers a reliable model performance while reducing the memory size significantly even without any additional training in most cases. In particular, expert layers in MoE models are much more robust to the quantization than conventional feedforward networks (FFN) layers. In our comprehensive analysis, we show that MoE models with 2-bit expert weights can deliver better model performance than the dense model trained on the same dataset. As a result of low-bit quantization, we show the model size can be reduced by 79.6% of the original half precision floating point (fp16) MoE model. Combined with an optimized GPU runtime implementation, it also achieves 1.24X speed-up on A100 GPUs.

研究の動機と目的

• MoE モデル推論における高いメモリ帯域幅とレイテンシのオーバーヘッドを軽減すること。これは、FLOP の増加が非線形的であるにもかかわらず、展開効率を制限する要因である。
• MoE モデルのエキスパート層が、標準的なフィードフォワードネットワーク (FFN) やアテンション層と比較して、低ビット量子化に対してより耐性があるかどうかを調査すること。
• モデル品質を保持しながら、メモリフットプリントを著しく削減し、推論速度を向上させる、軽量な重みのみの量子化手法を開発すること。
• 2 ビットで量子化されたエキスパートを備えた MoE モデルが、機械翻訳タスクにおいて、フルプレシジョンの密度型モデルを上回ることを実証すること。
• 最新の GPU 硬件上で測定可能な高速化を達成するため、ランタイム最適化による推論性能の最適化を行うこと。

提案手法

• MoE モデルのエキスパート FFN 層にのみ、重みのみの量子化を適用し、アテンション層および密度型 FFN 層はフルプレシジョンのままに保つ。
• エキスパート重みに 2 ビット量子化を適用し、Quantization Aware Training (QAT) を用いてモデル精度を維持する。一方、QAT を用いない場合、3 ビット量子化を適用する。
• エキスパート層の重みが、外れ値が少なく、分布が狭く安定しているという観察を活用し、低ビット量子化に強いことを実証する。
• GPU に最適化された推論ランタイムを実装し、量子化された MoE モデルの推論を高速化することで、より高いスループットを達成する。
• 異なるモデル部品（エキスパート FFN、密度型 FFN、アテンション）における量子化の影響を比較し、エキスパート層の耐性を明確に分離する。
• WMT および社内バリデーションデータセットの BLEU スコアを用いて、複数の言語ペアおよびモデルサイズにおける量子化の影響を評価する。

(a) FFN linear 1 weight distribution across layers

実験結果

リサーチクエスチョン

• RQ1 低ビット量子化が MoE モデルの異なる部品（特にエキスパート FFN 層と密度型 FFN やアテンション層）に与える影響は何か？
• RQ2 追加のトレーニングやキャリブレーションデータなしに、エキスパート重みの 2 ビット量子化を適用可能か？その場合の精度への影響は？
• RQ3 2 ビットで量子化されたエキスパートを備えた MoE モデルは、モデル品質および効率面で、フルプレシジョンの密度型モデルをどの程度上回るか？
• RQ4 エキスパート重みの量子化と最適化された GPU 推論の組み合わせが、モデルスループットおよびメモリ消費量に与える総合的効果は何か？
• RQ5 標準的な FFN 層と比較して、エキスパート層の量子化に対する耐性はどの程度高いか？

主な発見

• MoE モデルのエキスパート FFN 層は、密度型 FFN やアテンション層と比較して、低ビット量子化に対して著しく耐性が高く、QAT を用いた 2 ビット量子化でも性能劣化はわずかである。
• 2 ビットで量子化されたエキスパート重みを備えた MoE モデルは、元の fp16 MoE モデルと比較して、モデルサイズが 79.6% 減少した。
• 最適化された GPU ランタイムを用いることで、MoQE は A100 GPU 上で元の fp16 MoE モデルと比較して 1.24× の高速化を達成した。
• 2 ビットで量子化されたエキスパートを備えた MoE モデルは、同じデータセットにおいて、フルプレシジョンの密度型モデルを BLEU スコアで上回った。
• 密度型 FFN 層を 2 ビットに量子化すると BLEU スコアがほぼゼロにまで低下し、エキスパート層が低ビット量子化に対して特異に耐性があることを確認した。
• 3 ビット量子化でも WMT データセットで平均 1.01 ポints の BLEU スコア低下にとどまるが、2 ビット量子化では平均 11.54 ポイントの低下にとどまる。これは、密度型モデルと比較して、依然として顕著に優れた性能を示している。

(b) FFN linear 2 weight distribution across layers