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

调整参数后第二次尝试量化：

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的错误！于是更改了代码：

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倍！