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

使用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_datasetfrom 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 configmodel = GPTQModel.load(model_id, quant_config) # load modelmodel.quantize(calibration_dataset, batch_size=2) # quantizemodel.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_datasetfrom gptqmodel import GPTQModel, QuantizeConfigimport 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 configmodel = GPTQModel.load(model_id, quant_config) # load modelmodel.gradient_checkpointing_enable()model.quantize(calibration_dataset, batch_size=1) # quantizemodel.save(quant_path) # save model

使用SGLang部署后发现：

这里KV 缓存的大小明显增大了数倍，这能极大的提高并发数量
调整之后量化成功（可见降低了64.30%的显存空间），但是总共耗时了将近一小时的时间，因此我想使用四卡并行量化

四卡量化

可以看到内存在急速分配，说明模型权重需要先搬运到内存，再传递到显存中。
然后报了OOM的错误！于是更改了代码：

import osimport torchimport torch.distributed as distfrom datasets import load_datasetfrom 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倍！