器集群產(chǎn)品圖..png)
器產(chǎn)品圖.png)
器產(chǎn)品圖.png)
器產(chǎn)品圖.png)
器產(chǎn)品圖1-8080.png)
品主圖.png)
算豐技術(shù)揭秘｜探索ChatGLM2在算能BM1684X上INT8量化部署，加速大模型商業(yè)落地
?
1. 背景介紹
在2023年7月時(shí)我們已通過(guò)靜態(tài)設計方案完成了ChatGLM2-6B在單顆BM1684X上的部署工作，量化模式F16，模型大小12GB，平均速度約為3 token/s,詳見(jiàn)《算豐技術(shù)揭秘｜探索ChatGLM2-6B模型與TPU部署》。為了進(jìn)一步提升模型的推理效率與降低存儲空間，我們對模型進(jìn)行了INT8量化部署，整體性能提升70%以上，模型大小降低到6.4GB，推理速度達到6.67 token/s。
?
2. 量化方案
首先TPU-MLIR原有的INT8量化方案并不適合直接應用于LLM。主要是因為無(wú)論PTQ的校準或者QAT的訓練對于LLM來(lái)說(shuō)成本過(guò)高，對LLM的一輪PTQ的校準可能就需要1-2天時(shí)間；另外就是量化帶來(lái)的誤差在LLM上無(wú)法收斂，最終會(huì )導致模型精度大量損失。
在量化方案上我們沿用了ChatGLM2使用的W8A16策略，即只對GLMBlock中Linear Layer的權重進(jìn)行per-channel量化存儲，在實(shí)際運算時(shí)仍將其反量化回F16進(jìn)行運算。因為L(cháng)LM中Linear Layer權重數值間差異非常小，對INT8量化較為友好，所以量化過(guò)后的結果與F16計算結果在余弦相似度上仍然能保持99%以上，精度上幾乎可以做到0損失。
W8A16 MatMul?
3. TPU-MLIR實(shí)現
在Top到Tpu層的lowering階段，編譯器會(huì )自動(dòng)搜尋模型中右矩陣輸入為權重，且該矩陣維數為2的MatMul，將其替換為W8A16MatMul算子。此處主要是為了與左右矩陣都為Acitvation的MatMul算子區分開(kāi)（mm, bmm與linear layer在編譯器中會(huì )被統一轉換為MatMul算子）。以ChatGLM2中其中一個(gè)MatMul算子為例：L = (max_lengthx4096xf16), R = (4096x27392xf16)，量化后的權重由原來(lái)的214MB降為107MB，額外產(chǎn)生的Scale (4096xf16)只占了0.008MB的存儲空間，基本上可以達到減半的效果。算子替換源碼與權重量化源碼可在TPU-MLIR倉庫中查看。
?
Op Replacement in TPU-MLIR?
4. 后端性能提升原理
前一節介紹的量化只實(shí)現了存儲空間減半的效果，而性能提升主要在于W8A16MatMul后端算子的實(shí)現。如果對TPU架構不熟悉可通過(guò)TPU原理介紹(1)和TPU原理介紹(2)兩期視頻了解（可關(guān)注b站“算能開(kāi)發(fā)者”進(jìn)行觀(guān)看）。按照算能當前的TPU架構，W8A16的計算過(guò)程主要分為5個(gè)步驟：
1. 從Global Memory中加載數據到Local Memory
2. 將INT8權重Cast為F16
3. 與Scale數據相乘完成反量化操作
4. 與Input Activation進(jìn)行矩陣乘運算
5. 將計算結果存儲回Global Memory
????W8A16Matmul Computation on TPU
?
因為L(cháng)ocal Memory空間有限，對于大型數據通常需要進(jìn)行切分，分批對數據進(jìn)行加載、運算與存儲。為了提升效率，通常我們會(huì )利用GDMA與BDC指令并行，同時(shí)進(jìn)行數據搬運與運算操作，所以L(fǎng)ocal Mmeory大致需要被需要被劃分為兩部分區域，同一個(gè)循環(huán)內一個(gè)區域用于數據運算，另一個(gè)區域存儲上一循環(huán)計算好的結果以及加載下一循環(huán)需要用到的數據，如下圖所示。
Local Memory Partition?
矩陣乘等式如下：
?
當矩陣乘運算中左矩陣數據量較小時(shí)，性能瓶頸主要在于右矩陣的數據加載上，即數據加載時(shí)間遠比數據運算時(shí)間要長(cháng)很多。W8A16通過(guò)量化能夠將右矩陣的數據搬運總量縮小為原來(lái)的一半，而且額外多出的Cast與Scale運算時(shí)間可以被數據搬運時(shí)間覆蓋住，因此并不會(huì )影響到整體runtime，如下圖所示。
GDMA and BDC parallel總而言之，從后端角度來(lái)說(shuō)，當越小，越大時(shí)，W8A16帶來(lái)的性能提升收益越大。
從LLM的角度來(lái)看，我們以ChatGLM2為例，一次推理的完整流程分為一輪prefill與多輪decode。在prefill階段，基于我們當前的靜態(tài)設計方案，輸入詞向量會(huì )被補位為當前模型所支持的最大文本長(cháng)度max_length （e.g., 512, 1024, 2048）。而decode階段則固定只取前一輪生成的一個(gè)token作為輸入。
ChatGLM2 Inference因此max_length越長(cháng)，GLMBlock接收的輸入數據量越大，Linear Layer的也就越大，這就會(huì )導致W8A16的性能提升越有限。而decode階段始終保持為1，此時(shí)W8A16就能帶來(lái)明顯的性能提升。
MatMuls in ChatGLM2 prefill and decode phase5. 效果展示
將W8A16量化應用于ChatGLM2-6B后，整體性能如下所示：
-
性能：整體性能得到70%以上的提升
-
精度：與F16下的回答略有不同，但答案正確性仍然可以保證
-
模型大?。河?2GB降為6.4GB
Result Comparison?