簡介

Activation-Aware Weight Quantization (AWQ) 是類似於 GPTQ 的另外一種量化方法，同樣也是透過少量的校準資料集來進行更精準的量化，但是其量化模型的過程所需花費的時間，相較於 GPTQ 而言快了很多。量化後的模型結構，在推論階段時可以有相當快的速度。在社群上也有類似 AutoGPTQ 的 AutoAWQ 這種方便的工具可用，大大簡化了整體量化的流程。今天就來介紹 AWQ 的概念與如何使用 AutoAWQ 進行量化。

原理概念

AWQ 的論文指出，像是 GPTQ 這種透過反向傳播來進行量化的方法，可能會使 LLM 在校準資料集上發生 Overfitting 的情形。因此 AWQ 並不仰賴反向傳播，只透過前向運算來觀察輸出的分佈 (Activation Distribution) 並尋找重要的模型權重，這些權重在論文中被稱為 Salient Weights。作者指出，只要將其中 1% 的重要權重保留為 FP16 的資料型態，就能大幅降低量化帶來的誤差。

但如果為了保護這些分散在不同地方的重要權重，而使得整份權重量化完後是混精度的，則會使計算效率降低。因此作者引入 Activation-Aware Scaling 的技巧來對每組權重做量化，使得重要權重能夠獲得保護的同時，也能夠避免混精度權重以提昇硬體友善度，降低運算瓶頸，在實際推論上可以有相當快的速度。

GEMV vs GEMM

AWQ 分成 GEMV 與 GEMM 兩種矩陣乘法。因為文件沒有特別提到，所以筆者稍微雲了一下，GEMV 應該是指 General Matrix Vector Multiply，是矩陣對向量的乘法。而 GEMM 應該是指 General Matrix Multiply，是矩陣對矩陣的乘法。

GEMV 在單筆推論時，速度比 GEMM 快，但是當 Context 較大量時速度比較慢。也就是說 Prefill 的 Token 量越大，越適合使用 GEMM 的矩陣乘法。因此若是要進行 LLM Service 部署，一般而言選擇 GEMM 即可。

註：根據此 Issue 所述，目前 vLLM 並不支援 GEMV 因為他只能在 Batch Size 1 的情況下使用。

量化工具

AutoAWQ 是個專門對 LLM 進行 AWQ 量化的工具，因為量化的過程是一層一層的量化，而且 AutoAWQ 可以將模型權重本身放在 CPU 裡面，只有需要量化時才把其中一層的權重放在 GPU Memory 裡面，因此記憶體的消耗非常少，甚至能用單張 RTX 3090 24GB 來量化 70B 規模的模型。

最近正值 PyTorch 更換 CUDA 版本的日子，所以套件之間的相依關係開始有些複雜，因此建議要用 AutoAWQ 的話，最好開一個新的 Conda 環境：

conda create -n AutoAWQ python=3.10
conda activate AutoAWQ
pip install autoawq==0.1.7  # 此為筆者最後測試的版本

安裝好 AutoAWQ 後，可以使用 AutoAWQForCausalLM 類別讀取原始模型：

import torch
from awq import AutoAWQForCausalLM

model_dir = "Models/Llama2-7B"

model = AutoAWQForCausalLM.from_pretrained(
    model_dir,
    torch_dtype=torch.float16,
    device_map="cpu",  # 權重可以放在 CPU 就好
    safetensors=True,  # 如果有支援 Safetensors 的話就開啟
)

這裡我們將模型放在 CPU RAM 裡面，避免 FP16 權重佔用太多記憶體，以至於 AutoAWQ 沒有空間進行量化運算。接下來，我們需要一份校準資料集，預設會使用 mit-han-lab/pile-val-backup 這份資料集進行校準，但我們也可以自行準備，只要把資料以 String List 的形式存放即可：

# 此為示範用的校準資料集，請根據自身應用替換成適當資料集
# 只要是 list[str] 的型態且 Token 數量要夠多就能量化
calib_data = ["Hello, World! " * 8] * 16

如果給的文本太短太少的話，進行量化時會發生 RuntimeError: torch.cat(): expected a non-empty list of Tensors 的錯誤。這是因為在 awq/utils/calib_data.py 第 30 行，作者放了一個 Token 數量必須大於 512 才會拿來用的條件，可能是因為原本 AWQ 的實做就是如此的緣故。如果實際應用情況常常小於 512 Tokens

接下來放上相關設定，並開始進行量化：

from transformers import AutoTokenizer as TkCls

tokenizer = TkCls.from_pretrained(model_dir)

# 相關設定
quant_config = dict(
    zero_point=True,
    q_group_size=128,
    w_bit=4,
    version="GEMM",
)

# 進行量化
model.quantize(
    tokenizer,
    quant_config=quant_config,
    calib_data=calib_data,
)

# 儲存量化結果
output_dir = "Models/Llama2-7B-AWQ"
model.save_quantized(output_dir)
tokenizer.save_pretrained(output_dir)

筆者實測使用 RTX 3090 24GB 量化一個 7B 模型，大約只要七八分鐘就能完成。如果 CPU RAM 夠大，也能拿來量化 70B 的模型，大概要花費 100 分鐘左右。不過 70B 的模型即便量化到 4-Bit 也不能放在一張 24GB GPU 裡面使用就是了 QQ

這些量化完的模型，在 HF Transformers 裡面可以直接讀取進來使用：

from transformers import LlamaForCausalLM as ModelCls

model_dir = "Models/Llama2-7B-AWQ"
model: ModelCls = ModelCls.from_pretrained(
    model_dir,
    device_map="auto",
    use_safetensors=True,
)  # 約 4.5 GiB

不需要任何額外的設定就能讀取，雖然跟硬體效能有關，但 AWQ 模型在讀取速度上快滿多的。這種跟量化有關的技術，大神 TheBloke 同樣不會錯過。如果在 Hugging Face Hub 上搜尋 AWQ 的話，一字排開幾乎都是 TheBloke 上傳的模型。

從論文的描述來看，模型的精確度與校準資料集的內容相關性小的多，因此多數情況下直接拿別人上傳的 AWQ 模型來用理論上是沒問題的。但是實務上還是需要自行評測過才會知道是否沒問題，如果真的出問題了，或許就要考慮用自己的資料集來做校準量化了。

在推論的部份，與原本的用法並無二致：

from transformers import LlamaTokenizerFast as TkCls
from transformers import TextStreamer

tk: TkCls = TkCls.from_pretrained(model_dir)
ts = TextStreamer(tk)

prompt = "Hello, "
tokens = tk(prompt, return_tensors="pt").to("cuda")
model.generate(**tokens, max_new_tokens=16, streamer=ts)

實務上，更推薦使用 vLLM 或 TGI 來運行 AWQ 模型，請參考各框架的相關文件。

vLLM AutoAWQ Usage | TGI CLI Quantize

參考

• arXiv Paper - AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration
• GitHub - AWQ
• GitHub - AutoAWQ