一鍵部署 LLM 混合精度推理，端到端吞吐比 AWQ 最大提升 6 倍！
清華大學計算機系 PACMAN 實驗室發布開源混合精度推理系統——MixQ。
MixQ 支持 8 比特和 4 比特混合精度推理，可實現近無損的量化部署并提升推理的吞吐。
△圖 1 MixQ 吞吐與已有開源工作比較
MixQ 同時量化權重和激活，使用低精度張量核心（INT8/INT4 Tensor Core）實現推理加速；同時，MixQ 提取激活中少量的離群值，使用高精度張量核心（FP16 Tensor Core）保持推理準确性，通過系統優化掩蓋高精度訪存開銷。
不僅保持推理的準确性，而且通過使用低精度算力有效提升吞吐，充分發揮硬件計算潛力（圖 1）。
同時，研究團隊提供了基于 VLLM 和 Tensorrt-LLM 的混合精度推理，用戶可以方便地一鍵部署模型。
△圖 2 使用 VLLM 一鍵部署 4 比特和 8 比特混合精度量化并推理
MixQ 已支持多個主流大模型 LLaMA3，Qwen2，Baichuan2，ChatGLM 等。據了解，目前 MixQ 開源技術已被清程極智等 AI 行業公司應用在實際産品中。
該工作同時于高性能計算領域頂級國際會議 SC ’ 24 發表，第一作者清華大學博士後陳逸東、通訊作者爲翟季冬教授。
研究背景：已有量化技術總結
量化的主要技術路線有兩條，第一條是權重量化。
權重量化的理論加速比是 16/ 量化的比特數。例如，将模型壓縮成爲 4bit，那麽理論加速比爲 16/4=4 倍。
然而，當服務商面臨大量的用戶同時訪問時，權重量化的系統吞吐會低于 FP16 的吞吐，其主要原因是權重量化計算過程中将低精度權重恢複成 FP16 然後計算，這導緻權重量化并不使用低精度算力，當場景表現爲 compute bound 的時候，性能較低。
△圖 3 用戶請求多權重量化吞吐低于 FP16
第二條技術路線是量化權重和激活，使用低精度的張量核心來提升系統的吞吐。
直接将激活量化爲低比特可能會出現較大的精度損失。其原因在于激活矩陣中存在離群值（圖 4）。
一個有效的方法是 SmoothQuant，主要思想是通過平滑激活矩陣來降低量化激活的誤差。
△圖 4 激活矩陣中存在離群值
混合精度量化則是一類全新的量化方法，該方案先做了一個矩陣分解，對絕大部分權重和激活用低比特存儲，将離群值用 FP16 存儲，分别做矩陣乘法。
△圖 5 混合精度量化示意圖
混合精度量化的一個優勢就是可以實現近乎無損精度的量化。使用混合精度量化的 LlaMA 模型在 MMLU 20 個領域上的數據集進行推理準确率測試表明，采用 8bit 混合精度量化後的準确率下降不到 0.1%：
△圖 6 混合精度量化分類準确率
不過，此前已有的混合精度量化的系統的性能普遍不高，主要瓶頸在針對離群點進行查找、訪存和計算的開銷占比大。
以混合精度庫 Bitsandbytes 爲例，實測試表明，Bitsandbytes 在用戶請求數量爲 512 時僅有 1.08 倍的加速。
△圖 7 Bitsandbytes 的在 LLaMA70B 上的 Kernel 性能測試
△圖 8 Atomic operator 是混合精度推理系統的瓶頸之一
那麽，如何優化對離群點的查找、訪存和計算的開銷呢？
MixQ 的解決方案
MixQ 的核心思想是基于離群點的局部性對混合精度的計算圖做等價變換，使得變換後的混合精度的計算圖可以避免離群點查找的額外開銷；在此基礎上，通過圖層融合和設計高效的混合精度數據結構降低訪存開銷；最後通過 CUTLASS 生成高性能的混合精度算子，達到提升系統性能的效果。
MixQ 的設計基于以下的觀察：
離群點的局部性。對 LLM 的激活矩陣分析發現，在不同的 decode 階段的離群點的分布是有規律的。
如圖 9，紅色的點表示的是第一次出現的離群點，綠色的點表示的是重複出現的離群點，随着 decode 的進行，多數離群點出現在了固定的 channel。
△圖 9 decode 階段離群點的分布規律
因此，研究人員得到一個重要的結論：在大部分的 decode 階段是不需要重複檢測離群點的，也就是說我們可以避免檢查離群點的開銷。
剩下的問題是，如何知道哪些時候不需要重複檢查離群點呢？這個答案就隐藏在量化系數中。
在量化的過程中需要對矩陣進行 amax 的操作。因此，通過 amax 得到的結果可以判斷矩陣中是否存在離群點。如 amax 的值大于阈值，那矩陣中存在離群點。反之則不存在。
更重要的是，amax 操作可以和前一個操作融合。這樣不僅以極低的代價檢測離群點的存在，還通過對圖層進行融合來降低量化的開銷。
基于以上的分析，MixQ 的設計使用了三個關鍵技術：
一是對計算圖的等價變換。
針對混合精度的計算邏輯進行了等價變換以後，通過計算激活矩陣的 amax 的值，避免了檢測離群點的開銷。
△圖 10 優化混合精度的計算邏輯
二是設計混合精度數據結構。
MixQ 将離群點 " 拼接 " 成了一個新的矩陣。這一方法相較于 ATOM 采用的重排列（reorder）具有更低的開銷。
△圖 11 MixQ：order-reserved 數據結構
三是使用 CUTLASS 編寫高性能的混合精度的算子，這一關鍵技術的實現依賴于 NVIDIA 提供的高性能矩陣乘法模闆 CUTLASS 3.x。
MixQ 在寄存器中反量化低精度的計算結果并與高精度的結果進行相加。
△圖 12 融合 dequantize、scale 和 add 操作
下面來看 MixQ 的實驗結果，以 LLaMA 70B 爲例。
在準确率表現方面，MixQ 的準确率和 Bitsandbytes 一緻。
△圖 13 MixQ 的推理精度
在性能表現方面，MixQ 8bit kernel 是 Bitsandbytes 的 1.9 倍。
MixQ 4bit Kernel 的性能達 724TFLOPs，是 FP16 的 3.13 倍。
△圖 14 MixQ Kernel 性能
端到端測試下，MixQ 在 batch=512 相對 Bitsandbytes 和 AWQ 加速 1.78 和 6 倍。
△圖 15 多 batch 測試；上：MIXQ 的推理輸出（19.21it/s）；下：FP16 的推理輸出 （13.56it/s）
項目地址：
[ 1 ] https://github.com/Qcompiler/MixQ_Tensorrt_LLM
[ 2 ] https://github.com/Qcompiler/MIXQ
[ 3 ] https://github.com/Qcompiler/vllm-mixed-precision
— 完 —
投稿請發郵件到：
标題注明【投稿】，告訴我們：
你是誰，從哪來，投稿内容
附上論文 / 項目主頁鏈接，以及聯系方式哦
我們會（盡量）及時回複你
點這裏關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
>
