LLM 混合精度推論をワンクリックで展開し、AWQ と比較してエンドツーエンドのスループットが最大 6 倍向上します。

清華大学コンピュータサイエンス学部PACMANラボは、オープンソースの混合精度推論システムであるMixQをリリースしました。

MixQ は 8 ビットと 4 ビットの混合精度推論をサポートし、ほぼロスレスの量子化展開を可能にし、推論スループットを向上させます。

図 1. MixQ のスループットと既存のオープンソース作業の比較。

MixQは、低精度テンソルコア（INT8/INT4テンソルコア）を用いて重みと活性化を同時に量子化し、推論を高速化します。同時に、MixQは活性化から少数の外れ値を抽出し、高精度テンソルコア（FP16テンソルコア）を用いて推論精度を維持し、システム最適化によって高精度メモリアクセスのオーバーヘッドを隠蔽します。

推論の精度を維持するだけでなく、低精度の計算能力を使用することでスループットを効果的に向上させ、ハードウェアの計算能力を最大限に活用します (図 1)。

一方、研究チームはVLLM と Tensorrt-LLM に基づく混合精度推論を提供し、ユーザーがワンクリックで簡単にモデルを展開できるようにしています。

△図2. VLLMを用いた4ビットおよび8ビット混合精度量子化と推論のワンクリック展開

MixQは既にLLaMA3、Qwen2、Baichuan2、ChatGLMなど、いくつかの主流の大規模モデルをサポートしています。MixQのオープンソース技術は、Qingcheng JizhiなどのAI企業によって既に実際の製品に応用されていることが分かっています。

この研究は、高性能コンピューティング分野におけるトップクラスの国際会議であるSC'24でも発表されました。筆頭著者は清華大学のポスドク研究員であるYidong Chen博士、責任著者はJidong Zhai教授です。

研究背景：既存の定量的手法の概要

量子化には主に2つの技術的アプローチがあります。1つ目は重み付け量子化です。

重み付き量子化による理論的な高速化は、量子化ビット数の16倍です。例えば、モデルを4ビットに圧縮すると、理論的な高速化は16/4 = 4倍になります。

しかし、サービスプロバイダーが多数のユーザーが同時にシステムにアクセスする場合、重み量子化のスループットはFP16よりも低くなります。主な理由は、重み量子化プロセスでは、計算前に低精度の重みをFP16に復元するため、重み量子化は低精度の計算能力を消費しないからです。計算能力が制限されるシナリオでは、パフォーマンスが低下します。

△図3 マルチ重み量子化スループットに対するユーザー要求はFP16よりも低い

2 番目の技術的アプローチは、低精度テンソル コアを使用してシステムのスループットを向上させる量子化重みとアクティベーションです。

活性化を低ビット数に直接量子化すると、精度が大幅に低下する可能性があります。これは、活性化行列に外れ値が存在するためです（図4）。

効果的なアプローチの 1 つは SmoothQuant です。その主なアイデアは、アクティベーション マトリックスを平滑化することによって量子化アクティベーションのエラーを減らすことです。

△図4. 活性化マトリックスに外れ値が存在する。

混合精度量子化は、全く新しいタイプの量子化手法です。この方式では、まず行列分解を行い、重みと活性化関数の大部分を低ビット形式で保存し、外れ値をFP16形式で保存した後、行列乗算を行います。

△図5 混合精度量子化の模式図

混合精度量子化の利点の一つは、ほぼロスレスな精度を実現できることです。20ドメインにわたるMMLUデータセットにおいて、混合精度量子化を用いたLlaMAモデルの推論精度テストを実施したところ、8ビット混合精度量子化を使用した場合と比較して、精度の低下は0.1%未満であることが示されました。

△図6. 混合精度量子化分類の精度

しかし、既存の混合精度量子化システムのパフォーマンスは一般的に高くなく、外れ値の検索、メモリへのアクセス、計算の実行におけるオーバーヘッドの割合が大きいことが主なボトルネックとなっています。

混合精度ライブラリ Bitsandbytes を例にとると、実際のテストでは、ユーザー要求の数が 512 のとき、Bitsandbytes の速度向上は 1.08 倍にしかならないことが示されています。

図 7 は、LLaMA70B における Bitsandbytes のカーネル パフォーマンス テストを示しています。

△図8. アトミック演算子は、混合精度推論システムにおけるボトルネックの1つです。

では、外れ値の検出、メモリへのアクセス、計算のオーバーヘッドをどのように最適化できるでしょうか?

MixQのソリューション

MixQの核となるアイデアは、外れ値の局所性に基づいて混合精度計算グラフに等価変換を行うことです。これにより、変換された混合精度計算グラフは外れ値検索による余分なオーバーヘッドを回避できます。この基盤に基づき、レイヤー融合と効率的な混合精度データ構造の設計によってメモリアクセスのオーバーヘッドが削減されます。最後に、CUTLASSを通じて高性能な混合精度演算子が生成され、システム性能が向上します。

MixQ の設計は、次の観察に基づいています。

外れ値の局所性LLMの活性化行列分析により、外れ値の分布は異なるデコード段階にわたって一定のパターンに従うことが明らかになった。

図9に示すように、赤い点は初めて出現する外れ値を表し、緑の点は繰り返し出現する外れ値を表します。デコードが進むにつれて、ほとんどの外れ値は固定されたチャネルに出現します。

△図9 デコード段階における外れ値の分布パターン

そのため、研究者たちは重要な結論に達しました。つまり、ほとんどのデコード段階では外れ値の検出は必要なく、外れ値をチェックするオーバーヘッドを回避できるということです。

残る疑問は、外れ値を繰り返しチェックする必要がないことをどうやって知るかということです。その答えは量子化係数にあります。

量子化処理では、行列に対してAmax演算を実行する必要があります。したがって、Amax演算の結果は、行列内に外れ値が存在するかどうかを判断できます。Amax値が閾値より大きい場合、行列内に外れ値が存在すると判断され、そうでない場合は外れ値は存在しません。

さらに重要なのは、amax演算を前の演算と統合できることです。これにより、非常に低コストで外れ値を検出できるだけでなく、レイヤーを統合することで量子化のオーバーヘッドも削減されます。

上記の分析に基づいて、MixQ の設計では次の 3 つの主要なテクノロジが活用されています。

1つは計算グラフの等価変換です。

混合精度の計算ロジックに同等の変換を実行した後、活性化行列の amax 値を計算することにより、外れ値検出のオーバーヘッドを回避します。

△図10 混合精度を最適化する計算ロジック

次に、混合精度データ構造を設計します。

MixQは外れ値を新しい行列に「連結」します。この方法は、ATOMの並べ替えよりもオーバーヘッドが低くなります。

△図11 MixQ: 注文予約データ構造

3つ目に、高性能な混合精度演算はCUTLASSを用いて記述されています。この主要技術の実装は、NVIDIAが提供する高性能行列乗算テンプレートであるCUTLASS 3.xに依存しています。

MixQ は、レジスタ内の低精度の計算結果を逆量子化し、それを高精度の結果に加算します。

△図12 逆量子化、スケール、加算演算の組み合わせ

LLaMA 70Bを例に、MixQの実験結果を見てみましょう。

精度の点では、MixQ の精度は Bitsandbytes と一致しています。

△図13 MixQの推論精度

パフォーマンスの点では、MixQ 8 ビット カーネルは Bitsandbytes よりも 1.9 倍高速です。

MixQ 4 ビット カーネルは、FP16 の 3.13 倍となる 724 TFLOP のパフォーマンスを誇ります。

△図14 MixQカーネルのパフォーマンス

エンドツーエンドのテストでは、バッチ=512 で MixQ は Bitsandbytes および AWQ よりも 1.78 倍高速です。

△図15 マルチバッチテスト；上：MIXQ推論出力（19.21 it/s）；下：FP16推論出力（13.56 it/s）

プロジェクトアドレス: [1]https://github.com/Qcompiler/..._Tensorrt_LLM [2]https://github.com/Qcompiler/... [3]https://github.com/Qcompiler/...