1. 背景介紹

在2023年7月時我們已通過靜態設計方案完成了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算子區分開（mm, bmm與linear layer在編譯器中會被統一轉換為MatMul算子）。以ChatGLM2中其中一個MatMul算子為例：L = (max_lengthx4096xf16), R = (4096x27392xf16)，量化后的權重由原來的214MB降為107MB，額外產生的Scale (4096xf16)只占了0.008MB的存儲空間，基本上可以達到減半的效果。算子替換源碼與權重量化源碼可在TPU-MLIR倉庫中查看。

Op Replacement in TPU-MLIR

4. 后端性能提升原理

前一節介紹的量化只實現了存儲空間減半的效果，而性能提升主要在于W8A16MatMul后端算子的實現。如果對TPU架構不熟悉可通過TPU原理介紹(1)和TPU原理介紹(2)兩期視頻了解（可關注b站“算能開發者”進行觀看）。按照算能當前的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大致需要被需要被劃分為兩部分區域，同一個循環內一個區域用于數據運算，另一個區域存儲上一循環計算好的結果以及加載下一循環需要用到的數據，如下圖所示。

Local Memory Partition

矩陣乘等式如下：

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

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

從LLM的角度來看，我們以ChatGLM2為例，一次推理的完整流程分為一輪prefill與多輪decode。在prefill階段，基于我們當前的靜態設計方案，輸入詞向量會被補位為當前模型所支持的最大文本長度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. 模型大小：由12GB降為6.4GB

Result Comparison