探索ChatGLM2在算能BM1684X上INT8量化部署，加速大模型商業(yè)落地

1. 背景介紹

在2023年7月時我們已通過靜態(tài)設計方案完成了ChatGLM2-6B在單顆BM1684X上的部署工作，量化模式F16，模型大小12GB，平均速度約為3 token/s,詳見《算豐技術揭秘｜探索ChatGLM2-6B模型與TPU部署》。為了進一步提升模型的推理效率與降低存儲空間，我們對模型進行了INT8量化部署，整體性能提升70%以上，模型大小降低到6.4GB，推理速度達到6.67 token/s。

2. 量化方案

首先TPU-MLIR原有的INT8量化方案并不適合直接應用于LLM。主要是因為無論PTQ的校準或者QAT的訓練對于LLM來說成本過高，對LLM的一輪PTQ的校準可能就需要1-2天時間；另外就是量化帶來的誤差在LLM上無法收斂，最終會導致模型精度大量損失。

在量化方案上我們沿用了ChatGLM2使用的W8A16策略，即只對GLMBlock中Linear Layer的權重進行per-channel量化存儲，在實際運算時仍將其反量化回F16進行運算。因為LLM中Linear Layer權重數值間差異非常小，對INT8量化較為友好，所以量化過后的結果與F16計算結果在余弦相似度上仍然能保持99%以上，精度上幾乎可以做到0損失。

W8A16 MatMul

3. TPU-MLIR實現

在Top到Tpu層的lowering階段，編譯器會自動搜尋模型中右矩陣輸入為權重，且該矩陣維數為2的MatMul，將其替換為W8A16MatMul算子。此處主要是為了與左右矩陣都為Acitvation的MatMul算子區(qū)分開（mm, bmm與linear layer在編譯器中會被統(tǒng)一轉換為MatMul算子）。以ChatGLM2中其中一個MatMul算子為例：L = (max_lengthx4096xf16), R = (4096x27392xf16)，量化后的權重由原來的214MB降為107MB，額外產生的Scale (4096xf16)只占了0.008MB的存儲空間，基本上可以達到減半的效果。算子替換源碼與權重量化源碼可在TPU-MLIR倉庫中查看。

Op Replacement in TPU-MLIR

4. 后端性能提升原理

前一節(jié)介紹的量化只實現了存儲空間減半的效果，而性能提升主要在于W8A16MatMul后端算子的實現。如果對TPU架構不熟悉可通過TPU原理介紹(1)和TPU原理介紹(2)兩期視頻了解（可關注b站“算能開發(fā)者”進行觀看）。按照算能當前的TPU架構，W8A16的計算過程主要分為5個步驟：

1. 從Global Memory中加載數據到Local Memory
2. 將INT8權重Cast為F16
3. 與Scale數據相乘完成反量化操作
4. 與Input Activation進行矩陣乘運算
5. 將計算結果存儲回Global Memory

W8A16Matmul Computation on TPU

因為Local Memory空間有限，對于大型數據通常需要進行切分，分批對數據進行加載、運算與存儲。為了提升效率，通常我們會利用GDMA與BDC指令并行，同時進行數據搬運與運算操作，所以Local Mmeory大致需要被需要被劃分為兩部分區(qū)域，同一個循環(huán)內一個區(qū)域用于數據運算，另一個區(qū)域存儲上一循環(huán)計算好的結果以及加載下一循環(huán)需要用到的數據，如下圖所示。

Local Memory Partition

矩陣乘等式如下：

當矩陣乘運算中左矩陣數據量較小時，性能瓶頸主要在于右矩陣的數據加載上，即數據加載時間遠比數據運算時間要長很多。W8A16通過量化能夠將右矩陣的數據搬運總量縮小為原來的一半，而且額外多出的Cast與Scale運算時間可以被數據搬運時間覆蓋住，因此并不會影響到整體runtime，如下圖所示。

GDMA and BDC parallel
總而言之，從后端角度來說，當越小，越大時，W8A16帶來的性能提升收益越大。

從LLM的角度來看，我們以ChatGLM2為例，一次推理的完整流程分為一輪prefill與多輪decode。在prefill階段，基于我們當前的靜態(tài)設計方案，輸入詞向量會被補位為當前模型所支持的最大文本長度max_length （e.g., 512, 1024, 2048）。而decode階段則固定只取前一輪生成的一個token作為輸入。

ChatGLM2 Inference
因此max_length越長，GLMBlock接收的輸入數據量越大，Linear Layer的也就越大，這就會導致W8A16的性能提升越有限。而decode階段始終保持為1，此時W8A16就能帶來明顯的性能提升。MatMuls in ChatGLM2 prefill and decode phase

5. 效果展示

將W8A16量化應用于ChatGLM2-6B后，整體性能如下所示：

1. 性能：整體性能得到70%以上的提升
2. 精度：與F16下的回答略有不同，但答案正確性仍然可以保證
3. 模型大?。河?2GB降為6.4GB

Result Comparison