在CPU上高效部署低比特大語言模型
T-MAC的關鍵創新在於採用基於查找表(LUT)的計算範式,而非傳統的乘累加(MAC)計算範式。T-MAC 利用查找表直接支持低比特計算,從而消除了其他系統中必須的反量化(dequantization)操作,並且顯著減少了乘法和加法操作的數量。
經過實驗,T-MAC展現出了卓越的性能:在配備了最新高通Snapdragon X E lite晶片組的Surface AI PC 上,3B BitNet-b1.58模型的生成速率可達每秒48個token,2bit 7B llama模型的生成速率可達每秒30個token,4bit 7B llama模型的生成速率可達每秒20個token。
這甚至超越了NPU的性能!
當部署llama-2-7b-4bit模型時,儘管使用NPU可以生成每秒10.4個token,但CPU在T-MAC的助力下,僅使用兩核便能達到每秒12.6個token,最高甚至可以飆升至每秒22個token。
這些都遠超人類的平均閱讀速度,相比於原始的llama.cpp框架提升了4至5倍。
即使在較低端的設備如Raspberry Pi 5上,T-MAC針對3B BitNet-b1.58也能達到每秒11個token的生成速率。T-MAC也具有顯著的功耗優勢:達到相同的生成速率,T-MAC所需的核心數僅為原始llama.cpp的1/4至1/6,降低能耗的同時也為其它應用留下計算資源。
值得注意的是,T-MAC的計算性能會隨著比特數的降低而線性提高,這一現象在基於反量化去實現的GPU和NPU中是難以觀察到的。但T-MAC能夠在2比特下實現單核每秒10個token,四核每秒28個token,大大超越了NPU的性能。
矩陣乘不需乘,只需查表 (LUT)
對於低比特參數 (weights),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)提升隨機訪存性能。
-
改變矩陣axis計算順序,以盡可能提升放入片上內存的有限LUT的數據重用率。
-
為查表單獨設計最優矩陣分塊(Tiling)方式,結合autotvm搜索最優分塊參數
-
參數weights的布局優化
a)weights重排,以盡可能連續訪問並提升緩存命中率
b)weights交錯,以提升解碼效率
- 對Intel/ARM CPU 做針對性優化,包括
a)寄存器重排以快速建立查找表
b)通過取平均數指令做快速8-比特累加
研究員們在一個基礎實現上,一步步應用各種優化,最終相對於SOTA低比特算子獲得顯著加速: