使用GPTQ量化8B生成式语言模型
服务器配置：4*3090
描述：使用四张3090，分别进行单卡量化，多卡量化。并使用SGLang部署量化后的模型，使用GPTQ量化
原来的模型精度为FP16，量化为4bit
首先下载gptqmodel量化工具
因为要离线量化，所以需要提前下载好模型权重文件，校准数据集使用allenai/c4，需要使用HF的镜像来下载：参考我的这篇文章
shell
huggingface-cli download --repo-type dataset --resume-download allenai/c4 --local-dir ./calibration单卡量化
初始量化代码：
python
from datasets import load_dataset
from gptqmodel import GPTQModel, QuantizeConfig
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 model
model.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()
调整参数后第二次尝试量化：
python
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 model
model.gradient_checkpointing_enable()
model.quantize(calibration_dataset, batch_size=1) # quantize
model.save(quant_path) # save model
使用SGLang部署后发现：
这里KV 缓存的大小明显增大了数倍，这能极大的提高并发数量
调整之后量化成功（可见降低了64.30%的显存空间），但是总共耗时了将近一小时的时间，因此我想使用四卡并行量化
四卡量化
可以看到内存在急速分配，说明模型权重需要先搬运到内存，再传递到显存中。
然后报了OOM的错误！于是更改了代码：
python
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_rank
def 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倍！