簡介

GPTQ 是透過 Post-Training 的方式對模型進行量化，其準確率與速度通常比 bitsandbytes (BNB) 4-Bit 好一些，是個相當受歡迎的量化方式。缺點是不像 BNB 可以快速完成模型量化，需要花上數十分鐘甚至幾個小時來進行量化。本篇文章介紹如何透過 AutoGPTQ 套件對模型權重進行量化。

GPTQ

GPTQ 最原始的 GitHub 為 IST-DASLab/gptq，後來出現了專門對 Llama 模型進行 GPTQ 量化的 qwopqwop200/GPTQ-for-LLaMa 專案，目前則是以 AutoGPTQ 專案為主，在 Hugging Face Transformers 裡面的 GPTQ 量化方法，也是整合 AutoGPTQ 而來。

但是 HF Transformers 目前只能使用固定資料集進行量化，像是 C4 或 Wikitext 之類的，對一般用途的 LLM 而言雖然已經夠用。但如果是特定語言或特殊領域的模型，使用該領域的資料可能較為合適，這時還是要回到 AutoGPTQ 才能使用自己的訓練資料。

關於 HF Transformers 的 GPTQ 可以參考這份文件。

準備校準資料集

筆者建議自己手動切好資料集然後存在硬碟上，就不用每次量化都要重新跑一遍 Tokenization，且固定的資料集也比較能夠重現結果，也能讓不同模型的量化比較能在同一條基準線上。

假設我們將所有文本資料放在 data.txt 裡面，可以參考以下程式碼建立資料集：

import json
import random

from transformers import LlamaTokenizerFast as TkCls

num_sample = 16  # 訓練資料的數量
seq_len = 128  # 每筆資料的長度

model_id = "TheBloke/Llama-2-7b-chat-fp16"
tk: TkCls = TkCls.from_pretrained(model_id)

with open("data.txt", "rt", encoding="UTF-8") as fp:
    content = fp.read()

tokens = tk.encode(content)
max_idx = len(tokens) - seq_len - 1

dataset = list()
for _ in range(num_sample):
    i = random.randint(0, max_idx)
    data = tokens[i : i + seq_len]
    dataset.append({"input_ids": data})

with open("dataset.json", "wt", encoding="UTF-8") as fp:
    json.dump(dataset, fp)

建議使用與目標任務領域相近的文本，例如模型是個中文的泛用模型，那可能可以選擇中文維基當作文本。如果目標是氣象知識問答之類的，那可以用氣象相關的文本等等。在 GPTQ 原論文裡面，是從預訓練資料裡面，隨機挑選 128 筆固定長度的文本做校準。這裡為了加快訓練速度，只挑選了 16 筆資料，每筆資料長度為 128。但如果要執行的完整一點，建議遵循原論文挑選至少 128 筆資料，並盡量把長度拉高。

進行量化

進行量化之前，先設定訊息紀錄方便我們確認量化的進度：

# 顯示量化過程的進度訊息
import logging

logging.basicConfig(
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    level=logging.INFO,
    datefmt="%Y-%m-%d %H:%M:%S",
)

接著將資料集整理成 AutoGPTQ 要求的格式，每筆資料都是一個字典，包含 input_ids 與 attention_mask 兩組張量，需為 (1, 序列長度) 的形狀：

# 讀取校準資料集
import json
import torch

with open("dataset.json", "rt", encoding="UTF-8") as fp:
    dataset = json.load(fp)

for item in dataset:
    # Shape 必須是 (1, Sequence Length)
    item["input_ids"] = torch.LongTensor([item["input_ids"]])
    item["attention_mask"] = torch.ones_like(item["input_ids"])

print(f"Data Size: {len(dataset)}")
print(dataset[0]["input_ids"].shape)
print(dataset[0]["attention_mask"].shape)

接著設定相關參數：

# 設定相關參數
from auto_gptq import BaseQuantizeConfig

source_model = "TheBloke/Llama-2-7b-chat-fp16"
output_model = "Llama-2-7b-chat-GPTQ"

quantize_config = BaseQuantizeConfig(
    bits=4,  # 可以是 2, 3, 4, 8
    group_size=128,
    desc_act=False,
    # 可以避免 torch._C._LinAlgError
    damp_percent=0.1,
)

參數 bits 為量化的位元大小，可以是 2, 3, 4 或 8 位元。參數 damp_percent 預設為 0.01，如果會遇到 torch._C._LinAlgError 錯誤的話，可以試試看把這個數值開高。

接著讀取模型並開始量化：

# 讀取模型並進行量化
from auto_gptq import AutoGPTQForCausalLM

model = AutoGPTQForCausalLM.from_pretrained(
    source_model,
    quantize_config,
    device_map="auto",
)

model.quantize(
    dataset,
    batch_size=1,
    use_triton=False,
    # 可以減少記憶體消耗
    cache_examples_on_gpu=False,
)

在官方文件雖然有提及 max_memory 參數，但目前這個參數的好處並不大，建議使用 device_map="auto" 就好。如果用來校準的資料集很大，把 cache_examples_on_gpu 關掉可以減少記憶體用量。

使用少量的校準資料，大約五六分鐘就能完成量化，然後將量化完的模型存下來：

# 儲存量化結果
model.save_quantized(
    output_model,
    # TGI 必須使用 Safetensors 格式
    use_safetensors=True,
)

HF Transformers 與 TGI 都支援 Safetensors 格式，所以直接存為 Safetensors 格式就好。為了讓 HF Transformers 可以使用這份 GPTQ 模型權重，我們還需要在 config.json 裡面加上一些設定：

{
  ...
  "vocab_size": 32000,
  "quantization_config": {
    "bits": 4,
    "group_size": 128,
    "damp_percent": 0.1,
    "desc_act": false,
    "sym": true,
    "true_sequential": true,
    "model_name_or_path": null,
    "model_file_base_name": "model",
    "quant_method": "gptq"
  }
}

quantization_config 的內容請根據自己實際量化時設定的參數填入。除此之外，還要將 gptq_model-4bit-128g.safetensors 改名為 model.safetensors，這樣就可以直接使用 HF Transformers 讀取了：

from transformers import LlamaForCausalLM

LlamaForCausalLM.from_pretrained(
    "Llama-2-7b-chat-GPTQ",
    device_map="auto",
)

"""
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=0)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaFlashAttention2(
          (rotary_emb): LlamaRotaryEmbedding()
          (k_proj): QuantLinear()
          (o_proj): QuantLinear()
          (q_proj): QuantLinear()
          (v_proj): QuantLinear()
        )
        (mlp): LlamaMLP(
          (act_fn): SiLUActivation()
          (down_proj): QuantLinear()
          (gate_proj): QuantLinear()
          (up_proj): QuantLinear()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
"""

另外，也可以把 Tokenizer 複製一份放進資料夾裡面，這樣就完成製作一份獨立的 GPTQ 模型權重了。

參考

• HF Docs: Quantization
• arXiv Paper: GPTQ
• GitHub: AutoGPTQ
• GitHub Issue: Quant with larger context length
• GitHub Issue: torch._C._LinAlgError