一鍵部署LLM混合精度推理，端到端吞吐比AWQ最大提升6倍！

清華大學(xué)計(jì)算機(jī)系PACMAN實(shí)驗(yàn)室發(fā)布開源混合精度推理系統(tǒng)——MixQ。

MixQ支持8比特和4比特混合精度推理，可實(shí)現(xiàn)近無(wú)損的量化部署并提升推理的吞吐。

圖片

△圖1 MixQ吞吐與已有開源工作比較

MixQ同時(shí)量化權(quán)重和激活，使用低精度張量核心（INT8/INT4 Tensor Core）實(shí)現(xiàn)推理加速；同時(shí)，MixQ提取激活中少量的離群值，使用高精度張量核心（FP16 Tensor Core）保持推理準(zhǔn)確性，通過(guò)系統(tǒng)優(yōu)化掩蓋高精度訪存開銷。

不僅保持推理的準(zhǔn)確性，而且通過(guò)使用低精度算力有效提升吞吐，充分發(fā)揮硬件計(jì)算潛力（圖1）。

同時(shí)，研究團(tuán)隊(duì)提供了基于VLLM和Tensorrt-LLM的混合精度推理，用戶可以方便地一鍵部署模型。

圖2 使用VLLM一鍵部署4比特和8比特混合精度量化并推理

△

MixQ已支持多個(gè)主流大模型LLaMA3，Qwen2，Baichuan2，ChatGLM等。據(jù)了解，目前MixQ開源技術(shù)已被清程極智等AI行業(yè)公司應(yīng)用在實(shí)際產(chǎn)品中。

該工作同時(shí)于高性能計(jì)算領(lǐng)域頂級(jí)國(guó)際會(huì)議SC’24發(fā)表，第一作者清華大學(xué)博士后陳逸東、通訊作者為翟季冬教授。

圖片

研究背景：已有量化技術(shù)總結(jié)

量化的主要技術(shù)路線有兩條，第一條是權(quán)重量化。

權(quán)重量化的理論加速比是16/量化的比特?cái)?shù)。例如，將模型壓縮成為4bit，那么理論加速比為16/4=4倍。

然而，當(dāng)服務(wù)商面臨大量的用戶同時(shí)訪問(wèn)時(shí)，權(quán)重量化的系統(tǒng)吞吐會(huì)低于FP16的吞吐，其主要原因是權(quán)重量化計(jì)算過(guò)程中將低精度權(quán)重恢復(fù)成FP16然后計(jì)算，這導(dǎo)致權(quán)重量化并不使用低精度算力，當(dāng)場(chǎng)景表現(xiàn)為compute bound的時(shí)候，性能較低。

△圖3 用戶請(qǐng)求多權(quán)重量化吞吐低于FP16

第二條技術(shù)路線是量化權(quán)重和激活，使用低精度的張量核心來(lái)提升系統(tǒng)的吞吐。

直接將激活量化為低比特可能會(huì)出現(xiàn)較大的精度損失。其原因在于激活矩陣中存在離群值（圖4）。

一個(gè)有效的方法是SmoothQuant，主要思想是通過(guò)平滑激活矩陣來(lái)降低量化激活的誤差。

△圖4 激活矩陣中存在離群值

混合精度量化則是一類全新的量化方法，該方案先做了一個(gè)矩陣分解，對(duì)絕大部分權(quán)重和激活用低比特存儲(chǔ)，將離群值用FP16存儲(chǔ)，分別做矩陣乘法。

圖片

△圖5 混合精度量化示意圖

混合精度量化的一個(gè)優(yōu)勢(shì)就是可以實(shí)現(xiàn)近乎無(wú)損精度的量化。使用混合精度量化的LlaMA模型在MMLU 20個(gè)領(lǐng)域上的數(shù)據(jù)集進(jìn)行推理準(zhǔn)確率測(cè)試表明，采用8bit混合精度量化后的準(zhǔn)確率下降不到0.1%：

圖6 混合精度量化分類準(zhǔn)確率

不過(guò)，此前已有的混合精度量化的系統(tǒng)的性能普遍不高，主要瓶頸在針對(duì)離群點(diǎn)進(jìn)行查找、訪存和計(jì)算的開銷占比大。

以混合精度庫(kù)Bitsandbytes為例，實(shí)測(cè)試表明，Bitsandbytes在用戶請(qǐng)求數(shù)量為512時(shí)僅有1.08倍的加速。

圖7 Bitsandbytes的在LLaMA70B上的Kernel性能測(cè)試

圖8 Atomic operator是混合精度推理系統(tǒng)的瓶頸之一

那么，如何優(yōu)化對(duì)離群點(diǎn)的查找、訪存和計(jì)算的開銷呢？

MixQ的解決方案

MixQ的核心思想是基于離群點(diǎn)的局部性對(duì)混合精度的計(jì)算圖做等價(jià)變換，使得變換后的混合精度的計(jì)算圖可以避免離群點(diǎn)查找的額外開銷；在此基礎(chǔ)上，通過(guò)圖層融合和設(shè)計(jì)高效的混合精度數(shù)據(jù)結(jié)構(gòu)降低訪存開銷；最后通過(guò)CUTLASS生成高性能的混合精度算子，達(dá)到提升系統(tǒng)性能的效果。

MixQ的設(shè)計(jì)基于以下的觀察：

離群點(diǎn)的局部性。對(duì)LLM的激活矩陣分析發(fā)現(xiàn)，在不同的decode階段的離群點(diǎn)的分布是有規(guī)律的。

如圖9，紅色的點(diǎn)表示的是第一次出現(xiàn)的離群點(diǎn)，綠色的點(diǎn)表示的是重復(fù)出現(xiàn)的離群點(diǎn)，隨著decode的進(jìn)行，多數(shù)離群點(diǎn)出現(xiàn)在了固定的channel。

因此，研究人員得到一個(gè)重要的結(jié)論：在大部分的decode階段是不需要重復(fù)檢測(cè)離群點(diǎn)的，也就是說(shuō)我們可以避免檢查離群點(diǎn)的開銷。

剩下的問(wèn)題是，如何知道哪些時(shí)候不需要重復(fù)檢查離群點(diǎn)呢？這個(gè)答案就隱藏在量化系數(shù)中。

在量化的過(guò)程中需要對(duì)矩陣進(jìn)行amax的操作。因此，通過(guò)amax得到的結(jié)果可以判斷矩陣中是否存在離群點(diǎn)。如amax的值大于閾值，那矩陣中存在離群點(diǎn)。反之則不存在。

更重要的是，amax操作可以和前一個(gè)操作融合。這樣不僅以極低的代價(jià)檢測(cè)離群點(diǎn)的存在，還通過(guò)對(duì)圖層進(jìn)行融合來(lái)降低量化的開銷。

基于以上的分析，MixQ的設(shè)計(jì)使用了三個(gè)關(guān)鍵技術(shù)：

一是對(duì)計(jì)算圖的等價(jià)變換。

針對(duì)混合精度的計(jì)算邏輯進(jìn)行了等價(jià)變換以后，通過(guò)計(jì)算激活矩陣的amax的值，避免了檢測(cè)離群點(diǎn)的開銷。

圖片

圖10 優(yōu)化混合精度的計(jì)算邏輯

二是設(shè)計(jì)混合精度數(shù)據(jù)結(jié)構(gòu)。

MixQ將離群點(diǎn)“拼接”成了一個(gè)新的矩陣。這一方法相較于ATOM采用的重排列（reorder）具有更低的開銷。

圖11 MixQ：order-reserved數(shù)據(jù)結(jié)構(gòu)

三是使用CUTLASS編寫高性能的混合精度的算子，這一關(guān)鍵技術(shù)的實(shí)現(xiàn)依賴于NVIDIA提供的高性能矩陣乘法模板CUTLASS 3.x。

MixQ在寄存器中反量化低精度的計(jì)算結(jié)果并與高精度的結(jié)果進(jìn)行相加。

圖12 融合dequantize、scale和add操作

下面來(lái)看MixQ的實(shí)驗(yàn)結(jié)果，以LLaMA 70B為例。

在準(zhǔn)確率表現(xiàn)方面，MixQ的準(zhǔn)確率和Bitsandbytes一致。

圖13 MixQ的推理精度

在性能表現(xiàn)方面，MixQ 8bit kernel是Bitsandbytes的1.9倍。

MixQ 4bit Kernel的性能達(dá)724TFLOPs，是FP16的3.13倍。

圖片

△圖14 MixQ Kernel性能

端到端測(cè)試下，MixQ在batch=512相對(duì)Bitsandbytes和AWQ加速1.78和6倍。

圖片

圖15 多batch測(cè)試；上：MIXQ的推理輸出（19.21it/s）；下：FP16的推理輸出 （1

項(xiàng)目地址：
[1]https://github.com/Qcompiler/MixQ_Tensorrt_LLM
[2]https://github.com/Qcompiler/MIXQ
[3]https://github.com/Qcompiler/vllm-mixed-precision