CPUで低ビット大規模言語モデルを効率的に展開する
T-MACの主要な革新は、従来の乗算累積(MAC)計算パラダイムではなく、ルックアップテーブル(LUT)ベースの計算パラダイムを採用したことです。T-MACはルックアップテーブルを使用して低ビット計算を直接サポートし、他のシステムで必要な逆量子化操作を排除し、乗算と加算の数を大幅に削減します。
実験の結果、T-MACは優れたパフォーマンスを示しました:最新のQualcomm Snapdragon X E liteチップセットを搭載したSurface AI PCで、3B BitNet-b1.58モデルは毎秒48トークン、2ビット7B llamaモデルは毎秒30トークン、4ビット7B llamaモデルは毎秒20トークンの生成速度を達成しました。
これはNPUのパフォーマンスさえも上回ります!
llama-2-7b-4bitモデルを展開する際、NPUは毎秒10.4トークンを生成できますが、T-MACを活用したCPUは2コアだけで毎秒12.6トークンを達成し、最大で毎秒22トークンまで上昇します。
これらはすべて人間の平均読書速度をはるかに超え、オリジナルのllama.cppフレームワークと比較して4〜5倍の向上を示しています。
Raspberry Pi 5のような低スペックのデバイスでも、T-MACは3B BitNet-b1.58で毎秒11トークンの生成速度を達成します。T-MACは電力効率も大幅に向上させています:同じ生成速度を達成するのに必要なコア数は、オリジナルのllama.cppの1/4から1/6に減少し、エネルギー消費を削減しながら他のアプリケーションのための計算リソースも確保しています。
注目すべきは、T-MACの計算性能がビット数の減少に伴って線形に向上することです。これは逆量子化に基づいて実装されたGPUやNPUでは観察しにくい現象です。しかしT-MACは2ビットで単コアあたり毎秒10トークン、4コアで毎秒28トークンを実現し、NPUのパフォーマンスを大きく上回ります。
行列乗算に乗算は不要、テーブル参照(LUT)のみ
低ビットパラメータ(重み)に対して、T-MACは各ビットを個別にグループ化し(例:4ビットのグループ)、これらのビットを活性化ベクトルと乗算し、可能なすべての部分和を事前計算してLUTに格納します。
その後、T-MACはシフトと累積操作を採用して1から4までのスケーラブルなビット数をサポートします。この方法により、T-MACはCPU上で効率の低いFMA(乗加算)命令を排除し、代わりに消費電力が低く効率の高いTBL/PSHUF(テーブル参照)命令を使用します。
データ型中心の計算からビット中心の計算へ
従来の逆量子化ベースの計算は、実際にはデータ型中心の計算であり、この方法では各データ型に対して個別にカスタマイズする必要があります。
W4A16(重みint4、活性化float16)やW2A8など、活性化と重みのビット幅の組み合わせごとに、特定の重みレイアウトと計算カーネルが必要です。
例えば、W3のレイアウトでは2ビットと1ビットを別々にパッキングし、異なるインターリーブまたはシャッフル方法を使用してメモリアライメントまたは高速デコードを行う必要があります。そして、対応する計算カーネルはこの特定のレイアウトをハードウェアがサポートするデータ型に展開して実行する必要があります。
一方、T-MACは低ビット行列乗算をビットの観点から見ることで、単一のビットに対して最適なデータ構造を設計し、それを積み重ねて2/3/4ビットに拡張するだけで済みます。
同時に、異なる精度の活性化ベクトル(float16/float32/int8)に対しては、テーブル構築プロセスのみが変更され、テーブル参照時に異なるデータ構造を考慮する必要がなくなります。
また、従来の逆量子化ベースの方法では、4ビットから3/2/1ビットに減少させても、メモリ使用量は減少しますが、計算量は減少せず、逆量子化のオーバーヘッドが増加するため、パフォーマンスが悪化する可能性があります。
しかしT-MACの計算量はビット数の減少に伴って線形に減少するため、より低いビットでより良い加速を実現し、最新のBitNet、EfficientQATなどで発表された1ビット/2ビットモデルに高効率な展開ソリューションを提供します。
高度に最適化された演算子の実装
ビット中心の計算には多くの利点がありますが、CPUでの実装にはまだ大きな課題があります:
(1)活性化と重みの連続的なデータアクセスと比較して、テーブルへのアクセスはランダムです。高速オンチップメモリ内のテーブルの常駐は、最終的な推論パフォーマンスにとって特に重要です。
(2)しかし、オンチップメモリは限られており、ルックアップテーブル(LUT)方式は従来のmpGEMVと比較してオンチップメモリの使用量を増加させます。これは、ルックアップテーブルが活性化ベクトルとすべての可能なビットパターンの乗算結果を保存する必要があるためです。これは活性化自体よりもはるかに多くなります。
これに対処するため、マイクロソフトアジア研究所の研究者たちは、テーブル参照ベースの計算データフローを深く探求し、この計算パラダイムに効率的なデータ構造と計算フローを設計しました。これには以下が含まれます:
-
LUTをオンチップメモリに格納し、CPU上のテーブル参照ベクトル命令(TBL/PSHUF)を利用してランダムアクセスのパフォーマンスを向上させる。
-
行列軸の計算順序を変更し、オンチップメモリに格納された限られたLUTのデータ再利用率を最大限に高める。
-
テーブル参照に特化した最適な行列タイリング方法を設計し、autotvmを使用して最適なタイリングパラメータを検索する。
-
パラメータ(重み)のレイアウト最適化
a) 重みの再配置によるキャッシュヒット率の向上
b) 重みのインターリーブによるデコード効率の向上
- Intel/ARM CPUに対する特定の最適化
a) レジスタの再配置による高速なルックアップテーブルの構築
b) 平均化命令を使用した高速8ビット累積
研究者たちは基本的な実装に段階的に様々な最適化を適用し、最終的にSOTA低ビット演算子と比較して大幅な高速化を達成しました: