使用GPTQ量化Llama-3-8B大模型

使用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()