今天小編分享的科學經驗:清華開源混合精度推理系統MixQ,實現大模型近無損量化并提升推理吞吐,歡迎閱讀。
一鍵部署 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
— 完 —
投稿請發郵件到:
标題注明【投稿】,告訴我們:
你是誰,從哪來,投稿内容
附上論文 / 項目主頁鏈接,以及聯系方式哦
我們會(盡量)及時回復你
點這裡關注我,記得标星哦~
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
>