有 CPU 就能跑大模型，性能甚至超過 NPU / GPU！
沒錯(cuò)，為了優(yōu)化模型端側(cè)部署，微軟亞洲研究院提出了一種新技術(shù) —— T-MAC。
這項(xiàng)技術(shù)主打性價(jià)比，不僅能讓端側(cè)模型跑得更快，而且資源消耗量更少。
咋做到的？？
在 CPU 上高效部署低比特大語言模型
一般來說，要想在手機(jī)、PC、樹莓派等端側(cè)設(shè)備上使用大語言模型，我們需要解決存儲(chǔ)和計(jì)算問題。
常見的方法是模型量化，即將模型的參數(shù)量化到較低的比特?cái)?shù)，比如 4 比特、3 比特甚至更低，這樣模型所需的存儲(chǔ)空間和計(jì)算資源就會(huì)減少。
不過這也意味著，在執(zhí)行推理時(shí)，需要進(jìn)行混合精度的矩陣乘法運(yùn)算（mpGEMM），即用低精度的權(quán)重和高精度的激活向量進(jìn)行計(jì)算。
然而，現(xiàn)有的系統(tǒng)和硬件并不原生支持這種混合精度的矩陣乘法，因此它們通常需要將低精度的權(quán)重轉(zhuǎn)換回高精度，這個(gè)過程叫做反量化 (dequantization)。
但這種方法不僅效率低，而且當(dāng)比特?cái)?shù)進(jìn)一步降低時(shí)，并不能帶來性能上的提升。
對(duì)此，新技術(shù) T-MAC 采用基于查找表（LUT）的計(jì)算范式，無需反量化，直接支持混合精度矩陣乘。
這樣，T-MAC 不僅提高了推理性能，還使得模型更加統(tǒng)一和可擴(kuò)展，尤其適合在資源受限的端側(cè)設(shè)備部署。
此外，T-MAC 不依賴于專用的硬件加速器 NPU 或 GPU，能夠僅利用 CPU 部署模型。甚至在某些情況下，它的推理速度可以超過專用加速器。
T-MAC 的關(guān)鍵創(chuàng)新在于采用基于查找表（LUT）的計(jì)算范式，而非傳統(tǒng)的乘累加（MAC）計(jì)算范式。
T-MAC 利用查找表直接支持低比特計(jì)算，從而消除了其他系統(tǒng)中必須的反量化操作，并且顯著減少了乘法和加法操作的數(shù)量。
經(jīng)過實(shí)驗(yàn)，T-MAC 展現(xiàn)出了卓越的性能：
在配備了最新高通 Snapdragon X Elite 芯片組的 Surface AI PC 上，3B BitNet-b1.58 模型的生成速率可達(dá)每秒 48 個(gè) token，2bit 7B llama 模型的生成速率可達(dá)每秒 30 個(gè) token，4bit 7B llama 模型的生成速率可達(dá)每秒 20 個(gè) token。
這甚至超越了 NPU 的性能！
當(dāng)部署 llama-2-7B-4bit 模型時(shí)，盡管使用 NPU 可以生成每秒 10.4 個(gè) token，但 CPU 在 T-MAC 的助力下，僅使用兩核便能達(dá)到每秒 12.6 個(gè) token，最高甚至可以飆升至每秒 22 個(gè) token。
這些都遠(yuǎn)超人類的平均閱讀速度，相比于原始的 llama.cpp 框架提升了 4~5 倍。
即使在較低端的設(shè)備如 Raspberry Pi 5 上，T-MAC 針對(duì) 3B BitNet-b1.58 也能達(dá)到每秒 11 個(gè) token 的生成速率。
同時(shí)，T-MAC 也具有顯著的功耗優(yōu)勢(shì)：
達(dá)到相同的生成速率，T-MAC 所需的核心數(shù)僅為原始 llama.cpp 的 1/4 至 1/6，降低能耗的同時(shí)也為其它應(yīng)用留下計(jì)算資源。
值得注意的是，T-MAC 的計(jì)算性能會(huì)隨著比特?cái)?shù)的降低而線性提高，這一現(xiàn)象在基于反量化去實(shí)現(xiàn)的 GPU 和 NPU 中是難以觀察到的。
這進(jìn)一步使得 T-MAC 能夠在 2 比特下實(shí)現(xiàn)單核每秒 10 個(gè) token，四核每秒 28 個(gè) token，大大超越了 NPU 的性能。
采用新的計(jì)算范式
好了，說完了效果，咱們接著展開 T-MAC 的技術(shù)細(xì)節(jié)。
1、矩陣乘不需乘，只需查表（LUT）
對(duì)于低比特參數(shù) (weights)，T-MAC 將每一個(gè)比特單獨(dú)進(jìn)行分組（例如，一組 4 個(gè)比特），這些比特與激活向量相乘，預(yù)先計(jì)算所有可能的部分和，然后使用 LUT 進(jìn)行存儲(chǔ)。
之后，T-MAC 采用移位和累加操作來支持從 1 到 4 的可擴(kuò)展位數(shù)。
通過這種方法，T-MAC 拋棄了 CPU 上效率不高的 FMA（乘加）指令，轉(zhuǎn)而使用功耗更低、效率也更高的 TBL / PSHUF（查表）指令。
2、以比特為核心的計(jì)算，取代以數(shù)據(jù)類型為核心的計(jì)算
傳統(tǒng)的基于反量化的計(jì)算，實(shí)際上是以數(shù)據(jù)類型為核心的計(jì)算，這種方式需要對(duì)每一種不同的數(shù)據(jù)類型單獨(dú)定制。
每種激活和權(quán)重的位寬組合，如 W4A16（權(quán)重 int4 激活 float16）和 W2A8，都需要特定的權(quán)重布局和計(jì)算內(nèi)核。
例如，W3 的布局需要將 2 位和另外 1 位分開打包，并利用不同的交錯(cuò)或混洗方法進(jìn)行內(nèi)存對(duì)齊或快速解碼。
然后，相應(yīng)的計(jì)算內(nèi)核需要將這種特定布局解包到硬件支持的數(shù)據(jù)類型進(jìn)行執(zhí)行。
而 T-MAC 通過從比特的視角觀察低比特矩陣乘計(jì)算，只需為單獨(dú)的一個(gè)比特設(shè)計(jì)最優(yōu)的數(shù)據(jù)結(jié)構(gòu)，然后通過堆疊的方式擴(kuò)展到更高的 2/3/4 比特。
同時(shí)，對(duì)于不同精度的激活向量（float16 / float32 / int8），僅有構(gòu)建表的過程需要發(fā)生變化，在查表的時(shí)候不再需要考慮不同的數(shù)據(jù)結(jié)構(gòu)。
同時(shí)，傳統(tǒng)基于反量化的方法，從 4-比特降低到 3/2/1-比特時(shí)，盡管內(nèi)存占用更少，但是計(jì)算量并未減小，而且由于反量化的開銷不減反增，性能反而可能會(huì)更差。
但 T-MAC 的計(jì)算量隨著比特?cái)?shù)降低能夠線性減少，從而在更低比特帶來更好加速，為最新的工作 BitNet, EfficientQAT 等發(fā)布的 2-比特模型提供了高效率的部署方案。
比如下圖展示了：
（1）使用不同端側(cè)設(shè)備 CPU 的單核，T-MAC 在 4 到 1 比特的混合精度 GEMV 算子相較 llama.cpp 加速 3-11 倍。
（2）T-MAC 的 GEMM 耗時(shí)能隨著比特?cái)?shù)減少線性減少，而基于反量化的 llama.cpp 無法做到（1 比特 llama.cpp 的算子性能由其 2 比特實(shí)現(xiàn)推算得到）。
3、高度優(yōu)化的算子實(shí)現(xiàn)
概括而言，基于比特為核心的計(jì)算具有許多優(yōu)勢(shì)，但將其實(shí)現(xiàn)在 CPU 上仍具有不小的挑戰(zhàn)：
與激活和權(quán)重的連續(xù)數(shù)據(jù)訪問相比，表的訪問是隨機(jī)的。
表在快速片上內(nèi)存中的駐留對(duì)于最終的推理性能尤為重要，然而，片上內(nèi)存是有限的，查找表（LUT）方法相比傳統(tǒng)的 mpGEMV 增大了片上內(nèi)存的使用。
這是因?yàn)椴檎冶硇枰４婕せ钕蛄颗c所有可能的位模式相乘的結(jié)果，這比激活本身要多得多。
為此，微軟亞洲研究院的研究員們深入探究了基于查表的計(jì)算數(shù)據(jù)流，為這種計(jì)算范式設(shè)計(jì)了高效的數(shù)據(jù)結(jié)構(gòu)和計(jì)算流程，其中包括：
1、將 LUT 存入片上內(nèi)存，以利用 CPU 上的查表向量指令 (TBL / PSHUF) 提升隨機(jī)訪存性能。
2、改變矩陣 axis 計(jì)算順序，以盡可能提升放入片上內(nèi)存的有限 LUT 的數(shù)據(jù)重用率。
3、為查表單獨(dú)設(shè)計(jì)最優(yōu)矩陣分塊 (Tiling) 方式，結(jié)合 autotvm 搜索最優(yōu)分塊參數(shù)
4、參數(shù) weights 的布局優(yōu)化：
a、weights 重排，以盡可能連續(xù)訪問并提升緩存命中率
b、weights 交錯(cuò)，以提升解碼效率
5、對(duì) Intel / ARM CPU 做針對(duì)性優(yōu)化，包括
a、寄存器重排以快速建立查找表
b、通過取平均數(shù)指令做快速 8-比特累加
研究員們?cè)谝粋€(gè)基礎(chǔ)實(shí)現(xiàn)上，一步步應(yīng)用各種優(yōu)化，最終相對(duì)于 SOTA 低比特算子獲得顯著加速。
例如，在實(shí)現(xiàn)各種優(yōu)化后，T-MAC 4-比特算子最終相對(duì)于 llama.cpp 獲得顯著加速：
最后，T-MAC 現(xiàn)已開源，相關(guān)論文已在 arXiv 公開，感興趣可以進(jìn)一步了解。
開源地址（含代碼）：
https://github.com/microsoft/T-MAC
論文：
https://www.arxiv.org/pdf/2407.00088
本文來自微信公眾號(hào)：量子位（ID：QbitAI），作者：微軟亞洲研究院，原標(biāo)題《手機(jī)跑大模型提速 4-5 倍！微軟亞研院開源新技術(shù)，有 CPU 就行》
廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。