一
模型推理優化
随着模型在各種場景中的落地實踐，模型的推理加速早已成為AI工程化的重要内容。而近年基于Transformer架構的大模型繼而成為主流，在各項任務中取得SoTA成績，它們在訓練和推理中的昂貴成本使得其在合理的成本下的部署實踐顯得愈加重要。
大模型推理所面臨的挑戰主要有以下兩點：
- 巨大的記憶體（顯存）需求，主要來自于模型本身參數和推理的即時需求。
- 對于一個LLaMA2-30B的模型，載入顯存其模型本身需要約60GiB的顯存，推理過程中，單個token的KV cache 需要1.6MiB左右的顯存：6656(layer dim) * 52(layer num) *2 (K & V) * 2(fp16, 2bytes)；對于一個2048個token的請求則需要3.3GiB的顯存。
- 并行性較差，因為生成過程通常在時序上是一個串行的過程，導緻decoding的過程較難并行，成為計算的瓶頸。
常見的推理優化方式有知識蒸餾（Knowledge Distillation,KD），剪枝（Pruning）和量化（Quantization），以及針對LLM的記憶體優化而提出的各種方案（如Flash Attention、Paged Attention等）。
蒸餾指通過直接構造小模型，作為學生模型，通過軟标簽與原标簽結合的方式監督學習原模型的知識，進而使小模型具備與原模型相當的性能，最終用小模型代替大模型進而提高推理效率。
【圖檔出處：Knowledge Distillation: A survey,2021,p2】
剪枝則是通過靠剪除模型中不重要的權重進而給模型“瘦身”，提高模型的推理效率，為了保證模型的能力，通常剪枝過程也需要伴随着模型基于訓練資料的微調。根據剪除權重的次元不同，可以分為結構化剪枝（structured pruning）和非結構化剪枝（unstructured pruning）。
- 結構化剪枝：通常按權重張量的某一或多個次元成塊剪除不重要的通道，并保持正常的矩陣乘法；但因剪除的通道影響上下層的推理，需要檢查網絡的邏輯準确性。
- 非結構化剪枝：随機剪除權重張量中的不重要的元素，因而它通常會保持原本的權重結構，而造成稀疏的乘法計算，但并不能适配于通用的硬體，因而需要專用的硬體才能實作加速。
目前剪枝在LLM中的應用較少，如以下基于Activation-aware的剪枝工作[1]，主要是基于權重本身的的絕對值大小和輸入張量的絕對值大小做非結構化剪枝，使權重張量本身稀疏化，而模型的精度損失也并不能達到工程化的要求。
【圖檔出處：A simple and effective pruning approach for large language models,2021,p2】
再如下圖最近結構化剪枝的工作[2]，通過搜尋的方法尋找模型中的子結構，并通過重訓練以保持模型精度，剪枝後的模型的精度相比原模型有很大的降低，隻能跟同等參數量（剪枝後）的其他較小模型比較以顯示其方法的意義。
【圖檔出處: Sheared LLaMA: accelerating language model pre-training via structured pruning,2023,p3】
【圖檔出處: huggingface/Sheared-llama-1.3B】
而量化之是以會成為神經網絡以及LLM的首選，主要有以下的優點：
- 降低顯存的直覺展現。
- 一般LLM權重用FP16存儲，而權重量化為int4之後，則直覺上體積減小為原本的1/4（實際可能由于embeddings不量化，記憶體配置設定等一些原因會稍多一些），對顯存的資源需求大大降低。
- W4A16、W8A16等算子的加速，進而提升計算速度。
二
量化簡介
base
量化的本質通常是将模型的參數，或整個模型的推理過程從浮點轉化為整型。
量化參數通常由 scale 和 zero-point兩個值構成，前者為浮點，後者為整型。設x為一個張量（它可以為權重，也可以是推理的中間變量），其量化過程可以表示如下，
用b表示量化位寬，q{min}與q{max}分别表示整型值域的範圍，例如int-8量化可以取[-128,127]，即q{min}=-2^(b-1)=-128，q{max}=2^(b-1)-1=127，clamp(a;q{min},q{max})表示輸入值a基于[q{min}, q{max}]範圍的截斷操作，x{int}表示量化後的結果，s和z表示量化參數scale和zero-point。
【圖檔出處：A Survey of Quantization Methods for Efficient Neural Network Inference,2021,p5；An Introduction to Quantization of Large Language Models,p12】
而從整型到浮點的反量化過程如下，
關于量化參數，有很多算法基于搜尋，最優化，LKD(layer-by-layer 蒸餾)等各類算法計算其較優解，進而盡可能減少量化引起的精度損失；而最直接的計算scale 和方法即是基于張量元素min/max。
以下是一段簡單的代碼表示張量x從fp32量化到int8整型，再反量化回fp32的示例：
x->x{int}->x_hat的過程的一個示例如下：
量化前x：
量化後x_hat：
對稱/非對稱
相比于非對稱量化，對稱量化的定義是量化所映射的整型值域基于0值對稱，即上述公式的zero-point為0，qmax = -qmin，進而使量化的表達形式更為簡化。
非對稱量化有利于充分利用量化範圍。例如Conv+ReLU輸出的激勵張量，其值皆為正值，若使用對稱量化，則浮點将全部映射到[0~127]範圍，有一半的範圍未使用，其量化精度不如非對稱量化。
【圖檔出處：A Survey of Quantization Methods for Efficient Neural Network Inference,2021,p5】
實際中往往選擇對權重張量做對稱量化，而對輸入張量做非對稱量化。以下是來自qualcomm 的量化白皮書中的分析，如權重和輸入都選擇非對稱量化時，以Linear層的矩陣乘法為例，将表達式展開如下：
- 第一項是整型張量的乘法操作，是必須的即時操作；
- 第三、四項的操作包含了scale，zero和整型權重的乘法，這些都是提前預知的，因而可以事先計算作為偏置加上；
- 第二項的計算依賴x{int}，是每次推理需要即時計算的，而這會造成額外算力。
因而當我們将權重量化改為對稱量化時(zW=0)，則上式簡化為如下，即時計算時，隻需要計算第一項的矩陣乘法，第二項是預先算好的偏置項：
而當兩者都是對稱量化時的表達式，則簡化如下：
對比原模型中的浮點計算W{x}，W{int}x{int}是整型與整型之間的乘法，後者在Nvidia GPU上的運算速度遠快于前者，這是量化模型的推理速度大大加快的原因。
三
LLM的量化
Challenges in LLM Quantization
從模型表現的角度來講，量化自始至終要解決的一個前提是，如何保持量化後模型的精度，即讓模型的使用者覺得量化後的模型在推理效率提高的同時，還能保持原來的性能。
神經網絡中需要量化的操作主要是卷積層Conv(x;W)和全連接配接層Wx，即主要是按上一部分描述的操作分别對W和x做的權重量化（Weight Quantization,WQ）和激勵量化(Activation Quantization,AQ)。
而不同于CNN模型或者小型Transformer模型，基于Transformer的大模型的矩陣乘法産生的激勵張量通常有較多的離群值(outliers)，即離值分布的大多數點形成的點群較遠的值, 這些絕對值較大但占比較低的元素值增加了量化難度。而如何取舍outliers通常是量化工作中的一大難點，若過分考慮之，則會因量化範圍過大而降低量化的表達範圍，若過分截斷之，通常會因這些絕對值較大的值，在模型推理中對結果有較大影響，而導緻模型效果變差，而後者在LLM的量化則尤為明顯。
下圖分别是Resnet18與Opt-13B的某層輸入張量的元素值統計，sigma表示各自分布的标準差，Resnet18輸入的極大值約為28sigma，且絕對值6sigma以外的比例在0.05%；而Opt-13B網絡輸入的極大值越為325sigma，且絕對值6sigma以外的比例在0.2%。從量化效果而言，Resnet18的int-8精度基本無損失，而Opt-13B的int-8模型的精度已崩塌。
【圖檔出處：An Introduction to Quantization of Large Language Models,p20 】
在應對激勵量化的挑戰這方面，有一些方案嘗試降低量化精度，比如SmoothQuant提出的思路。
【圖檔出處：SmoothQuant,p4】
在矩陣乘法中，他們通過按比例縮小輸入張量X的值，而将縮小的比例補償給權重張量W，即把問題從量化X和W轉化為了量化 X·diag(s^(-1))和diag(s)·W。進而在保證乘法運算的積保持不變的前提下，降低張量X的量化難度。而在實際工程中，這種量化方案引起的量化誤差對大模型的推理效果仍然有比較明顯的影響，即使在int-8精度量化亦有明顯的誤差。如以下對Llama2-7B的SmoothQuant應用結果顯示其perplexity非常糟糕，難以在實際中應用。
是以在目前工程部署中的實用方案，大多以weight-only的量化方案為主，即放棄activation的量化。
GPTQ
GPTQ是最早被工程化部署所接受的量化方案，W8A16或W4A16的量化效果在多數場景中都有與原模型較為接近的表現，而且其量化過程非常快。
量化過程
以矩陣乘法的基本單元操作為例，基于 weight-only量化前後的乘積的均方差，可以寫出如下優化函數，
W 是在Transformer 中的Linear層權重，X表示其對應的輸入。離線量化的過程是逐子產品（Transformer）逐層（Q,K,V,O,Fc1,Fc2）做量化。
參數和資料定義如下：
- W∈R^{K×M}，X∈R^{M×N}，Y=W×X∈R^{K ×N}
- calibrate set：部分資料用作推理，用于檢視各層輸入張量的值範圍，并基于此量化。
具體量化過程如下：
- 計算Hessian（上述優化函數對于W_hat的Hessian，而非反向傳播中的Hessian），加入擾動項：
- act order sort（desc_act，值範圍相近的column一起做量化），基于diag(H)對W基于M次元作列重排，同理，對應地H在兩個次元上重排。
- 求逆H^(-1)（cholesky分解）。
- 對W沿次元M，從左到右逐塊量化，block size B=128，其右側還未量化部分基于H^(-1)更新，以補償量化損失。
- （inner loop）針對每個block内部，逐列量化，計算誤差，并對該block内部未量化的列，基于誤差更新。
- （outer loop）操作完該block，更新其後面的所有列：
group_size
- 若不指定group size，預設g=-1，以所有列為機關統計量化參數，并對每一行的權重做量化，對于W∈R^{K×M}，量化參數的數量為K×1。
- 若指定group size，例如g=128，則會以每128列為機關統計量化參數，并對每一行的權重做量化，對于W∈R^{K×M}，量化參數的數量為K×(M/g)。
重排desc_act
根據Hessian Matrix H，基于diag(H)對W基于M次元作列重排。其目的是優先量化絕對值較大的activaiton對應的weight的列，這些列在推理中被視為更為影響結果的重要的列，因而希望在量化這些列時盡可能産生較小的誤差，而将更多的量化誤差轉移到後面相對不重要的列中。
部分實驗表明desc_act對量化損失的效果在多數的任務中是有效的trick。
Perplexity of Pygmalion-7B with GPTQ [7]
【圖檔出處：https://huggingface.co/reeducator/vicuna-13b-free/discussions/22】
算子
嚴格來說基于weight-only的W4A16相比于原本的W16A16并沒有太多效率的提升，而且推理中還加入了quant/dequant過程；而随着weight-only成為LLM量化的主流且應用越來越多，有很多開源的工作基于W4A16高效算子的編寫為量化算法的推理提速賦能，比如GPTQ的python package AutoGPTQ已內建于開源工具exllama，後者基于triton和CUDA重寫了量化乘法的并行計算。在exllama/exllama_ext/matrix.cuh可以看到dot_product8_h對out=W_hat·x=(W{int}-z)s·x=(W{int}-z)x·s的實作。
【圖檔出處：https://github.com/turboderp/exllama/blob/3b013cd53c7d413cf99ca04c7c28dd5c95117c0d/exllama_ext/matrix.cuh#L86】
AWQ
相比于GPTQ從最優化問題出發設計方案，AWQ是基于搜尋提出的量化方案。
用Q(·)表示量化反量化過程，則修改前的量化過程如下：
修改後，量化過程如下，加入了對W的縮放：
搜尋
AWQ 的全稱為Activation-aware Weight Quantization, 即對Weight的量化過程考慮Activation的值的影響。其出發點也是基于在Weight的各個通道中，處理對應的Activtion的值較大的通道則相對重要，反之則相對不重要，進而通過乘以一個縮放系數Δ去展現其重要性，而Δ的值和範圍則通過輸入的activation的張量值設計。
搜尋的衡量标準依據Linear層量化前後輸出結果的比較，取MSE結果最小者為最優解。
效果
從模型表現效果方面，通過逐層 scale search 尋找最優的縮放系數，進而取量化誤差最小的解，以下來自AWQ paper的效果比較，從Perplexity的角度，顯示在兩代Llama的測試上其量化結果稍優于GPTQ及GPTQ的排序版。
【圖檔出處：AWQ, p6】
從實際任務的準确率來看，AWQ的準确率與GPTQ的act_order版本（GPTQ-R）相當，而速度優于後者。
【圖檔出處：AWQ, p5】
從模型的計算性能方面，GPTQ因為有reorder操作，矩陣乘法是MV（matrix×vector），為不連續的記憶體通路，而AWQ不存在reorder操作，矩陣乘法為（matrix×matrix），速度更快。
四
總結
關于LLM的量化工作目前的SOTA performance，基本上都是基于weight-only 的量化模式，模型在GPU運作所需的顯存降低是其主要的貢獻。
從模型的表現來看，因為存在不可避免的量化損失，且LLM模型通常比傳統的CNN模型對量化要敏感得多，雖然在很多任務上量化後的LLM表現與量化前差距不大，但是在一部分任務上可能依然無法勝任。
從模型的加速來看，weight-only的量化促使底層加速的工作基本上都在W4A16、W3A16、W8A16等乘法算子上的加速，從paper上提供的理論資料上來看通常相較于FP16模型隻有 1.x ~3.x 倍速度的提升，而實際部署效果可能低于此數值，其加速效果遠不如傳統量化方法的W4A4、W8A8等全整型的乘法算子。
總體來說，LLM領域的量化工作還很初步，若在實際任務中對模型的表現精度要求十分高，更推薦單純基于KV cache等方向提高機關顯存吞吐量的算法和工具，如Flash Attention-2、Paged Attention等。
五
Reference
1. A Simple and Effective Pruning Approach for Large Language Models, 2023.
2. Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning, 2023.
3. A White Paper on Neural Network Quantization, 2021.
4. SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models, 2023.
5. GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers, 2023.
6. AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration, 2023.
7. Some evaluation on GPTQ performance.
作者:xujiong
來源-微信公衆号:得物技術
出處:https://mp.weixin.qq.com/s/25gZPdn29YBpyEx_l-D09A