使用GPTQ量化8B生成式語言模型

服務器配置：4*3090

描述：使用四張3090，分別進行單卡量化，多卡量化。并使用SGLang部署量化后的模型，使用GPTQ量化
原來的模型精度為FP16，量化為4bit

首先下載gptqmodel量化工具

因為要離線量化，所以需要提前下載好模型權重文件，校準數據集使用allenai/c4，需要使用HF的鏡像來下載：參考我的這篇文章

huggingface-cli download --repo-type dataset --resume-download allenai/c4 --local-dir ./calibration

單卡量化

初始量化代碼：

from datasets import load_dataset
from gptqmodel import GPTQModel, QuantizeConfigmodel_id = "/home/zgq/Meta-Llama-3-8B-Instruct"
quant_path = "/home/zgq/sglang/Meta-Llama-3-8B-Instruct-gptqmodel-4bit"# 直接加載本地數據集，不再從 Hugging Face 下載
calibration_dataset = load_dataset("json",  # 指定數據格式data_files="calibration_data/en.noblocklist/c4-train.00001-of-01024.json.gz",  # 指定本地路徑split="train"
).select(range(1024))["text"]  # 選擇前 1024 條數據quant_config = QuantizeConfig(bits=4,# 量化為 4 位group_size=128  # 分組大小為 128) # quantization config
model = GPTQModel.load(model_id, quant_config) # load modelmodel.quantize(calibration_dataset, batch_size=2) # quantize
model.save(quant_path) # save model

量化過程出現了OOM
解決方案：
首先嘗試將batch_size降為1
優化顯存管理：export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
插入梯度檢查點：model.gradient_checkpointing_enable()

調整參數后第二次嘗試量化：

from datasets import load_dataset
from gptqmodel import GPTQModel, QuantizeConfig
import torch# torch.cuda.empty_cache()
# torch.cuda.set_per_process_memory_fraction(0.9)  # 保留 80% 顯存model_id = "/home/zgq/Meta-Llama-3-8B-Instruct"
quant_path = "/home/zgq/sglang/Meta-Llama-3-8B-Instruct-gptqmodel-4bit"# 直接加載本地數據集，不再從 Hugging Face 下載
calibration_dataset = load_dataset("json",  # 指定數據格式data_files="calibration_data/en.noblocklist/c4-train.00001-of-01024.json.gz",  # 指定本地路徑split="train"
).select(range(1024))["text"]  # 選擇前 1024 條數據quant_config = QuantizeConfig(bits=4,# 量化為 4 位group_size=128  # 分組大小為 128) # quantization config
model = GPTQModel.load(model_id, quant_config) # load modelmodel.gradient_checkpointing_enable()
model.quantize(calibration_dataset, batch_size=1) # quantize
model.save(quant_path) # save model

使用SGLang部署后發現：

這里KV 緩存的大小明顯增大了數倍，這能極大的提高并發數量
調整之后量化成功（可見降低了64.30%的顯存空間），但是總共耗時了將近一小時的時間，因此我想使用四卡并行量化

四卡量化

可以看到內存在急速分配，說明模型權重需要先搬運到內存，再傳遞到顯存中。
然后報了OOM的錯誤！于是更改了代碼：

import os
import torch
import torch.distributed as dist
from datasets import load_dataset
from gptqmodel import GPTQModel, QuantizeConfig# 初始化分布式環境
def init_distributed():dist.init_process_group(backend="nccl")print(f"Rank {dist.get_rank()}: Distributed environment initialized.")local_rank = int(os.environ["LOCAL_RANK"])torch.cuda.set_device(local_rank)return local_rankdef main():local_rank = init_distributed()# 模型路徑和量化保存路徑model_id = "/home/zgq/Meta-Llama-3-8B-Instruct"quant_path = "/home/zgq/Meta-Llama-3-8B-Instruct-gptqmodel-4bit"# 加載本地校準數據集（每個進程加載自己的分片）calibration_dataset = load_dataset("json",data_files="calibration_data/en.noblocklist/c4-train.00001-of-01024.json.gz",split=f"train[{local_rank * 256}:{(local_rank + 1) * 256}]")["text"]print(f"Rank {dist.get_rank()}: Loaded {len(calibration_dataset)} samples.")# 量化配置quant_config = QuantizeConfig(bits=4,group_size=128)# 分布式加載模型model = GPTQModel.load(model_id,quant_config,device_map={"": local_rank},torch_dtype=torch.float16,low_cpu_mem_usage=True)print(f"Rank {dist.get_rank()}: Model loaded on GPU {local_rank}.")# 啟用梯度檢查點model.gradient_checkpointing_enable()# 分布式量化（每個進程處理自己的數據分片）model.quantize(calibration_dataset, batch_size=1)print(f"Rank {dist.get_rank()}: Quantization completed.")# 主進程保存量化模型if local_rank == 0:model.save(quant_path)print(f"Rank {dist.get_rank()}: Model saved to {quant_path}.")dist.barrier()  # 等待所有進程完成if __name__ == "__main__":main()

程序啟動命令為：torchrun --nproc_per_node=4 torch-dist-gptq.py (指定每個節點使用 4 個 GPU,根據你的環境來修改)

可見四卡并行的方式要快很多！

以上圖片是運行時硬件狀態。

量化完成，和單卡量化效果一致，但是縮短了量化時間，效率約提升了4.5倍！