EfficientQAT: 大型语言模型的高效量化感知训练

安装

1. 克隆此仓库并导航到 EfficientQAT 文件夹

git clone https://github.com/OpenGVLab/EfficientQAT.git
cd EfficientQAT

1. 安装包

conda create -n efficientqat python==3.11

conda activate efficientqat

pip install -r requirements.txt

模型库

我们提供多个预量化的 EfficientQAT 模型，如下所示：

• WikiText2 PPL 是在 2048 上下文长度下测量的。
• 平均准确率表示在 5 个零样本推理任务（WinoGrande, PIQA, HellaSwag, Arc-Easy, Arc-Challenge）中的平均准确性，版本为 lm-eval v0.4.2。
• 1GB = 位
• Hub Link: EQAT 表示原始检查点。我们还将检查点转换为 GPTQ 和 BitBLAS 格式， 可以通过 GPTQModel 直接加载。(PS: GPTQModel 是 AutoGPTQ 的官方修复仓库，未来将并入 AutoGPTQ。)

Model	量化	WikiText2 PPL	平均准确率	模型大小 (GB)	Hub链接
Llama-2-7B	fp16	5.47	64.86	13.2	-
Llama-2-7B	w4g128	5.53	64.27	3.7	EQAT|GPTQ|BitBLAS
Llama-2-7B	w3g128	5.81	64.02	3.1	EQAT
Llama-2-7B	w2g64	6.86	60.14	2.3	EQAT|GPTQ|BitBLAS
Llama-2-7B	w2g128	7.17	59.50	2.2	EQAT|GPTQ|BitBLAS
Llama-2-13B	fp16	4.88	67.81	25.4	-
Llama-2-13B	w4g128	4.93	67.52	6.8	EQAT|GPTQ|BitBLAS
Llama-2-13B	w3g128	5.12	67.28	5.6	EQAT
Llama-2-13B	w2g64	5.96	64.88	4.0	EQAT|GPTQ|BitBLAS
Llama-2-13B	w2g128	6.08	63.88	3.8	EQAT|GPTQ|比特 BLAS
Llama-2-70B	fp16	3.32	72.41	131.6	-
Llama-2-70B	w4g128	3.39	72.62	35.8	EQAT|GPTQ|BitBLAS
Llama-2-70B	w3g128	3.61	71.76	29.1	EQAT
Llama-2-70B	w2g64	4.52	69.48	20.1	EQAT|GPTQ|BitBLAS
Llama-2-70B	w2g128	4.61	68.93	18.9	EQAT|GPTQ|BitBLAS
Llama-3-8B	fp16	6.14	68.58	13.0	-
Llama-3-8B	w4g128	6.47	68.43	5.4	EQAT|GPTQ|BitBLAS
Llama-3-8B	w3g128	7.09	67.35	4.7	EQAT
Llama-3-8B	w2g64	9.41	60.76	3.9	EQAT|GPTQ|BitBLAS
Llama-3-8B	w2g128	9.80	59.36	3.8	EQAT|GPTQ|BitBLAS
Llama-3-70B	fp16	2.85	75.33	137.8	-
Llama-3-70B	w4g128	3.17	74.57	38.9	EQAT|GPTQ|BitBLAS
Llama-3-70B	w3g128	4.19	72.42	32.2	EQAT
Llama-3-70B	w2g64	6.08	67.89	23.2	EQAT|GPTQ
Llama-3-70B	w2g128	6.38	67.57	22.0	EQAT|GPTQ|BitBLAS
Llama-3-8B-Instruct	fp16	8.29	68.43	13.0	-
Llama-3-8B-Instruct	w4g128	7.93	68.39	5.4	EQAT|GPTQ|BitBLAS
Llama-3-8B-Instruct	w3g128	8.55	67.24	4.7	EQAT
Llama-3-8B-Instruct	w2g64	11.19	60.66	3.9	EQAT|GPTQ|BitBLAS
Llama-3-8B-Instruct	w2g128	11.73	60.16	3.8	EQAT|GPTQ|BitBLAS
Llama-3-70B-Instruct	fp16	5.33	73.78	137.8	-
Llama-3-70B-Instruct	w4g128	5.35	73.47	38.9	EQAT|GPTQ|BitBLAS
Llama-3-70B-Instruct	w3g128	5.65	72.87	32.2	EQAT
Llama-3-70B-Instruct	w2g64	7.86	67.64	23.2	EQAT|GPTQ|BitBLAS
Llama-3-70B-Instruct	w2g128	8.14	67.54	22.0	EQAT|GPTQ|BitBLAS
Mistral-Large-Instruct-2407	fp16	2.74	77.76	228.5	-
Mistral-Large-Instruct-2407	w2g64	5.58	73.54	35.5	GPTQ

训练

EfficientQAT 包括两个连续的训练阶段：所有参数的块级训练 (Block-AP ) 和量化参数的端到端训练 (E2E-QP )。详细的训练脚本可以在 ./examples 中找到。我们在下面提供了关于 Llama-2-7B 的 w2g64 量化的训练脚本示例。

1. Block-AP

在运行以下命令之前，您应该将 --model 修改为脚本中全精度模型的文件夹。

bash examples/block_ap/Llama-2-7b/w2g64.sh

具体来说，--weight_lr 在我们的实验中对于 2 位是2e-5，而对于 3/4 位是1e-5。

一些其他重要的论点：

• --train_size: 训练数据样本数量，默认值为 4096
• --val_size：验证数据样本数量，默认值为 64
• --off_load_to_disk：将训练数据集保存到磁盘，节省 CPU 内存，但可能会降低训练速度

1. E2E-QP

然后，您可以加载 Block-AP 的量化模型以进行进一步的 E2E-QP。具体而言，E2E-QP 可以通过更改训练数据集适应不同的场景。您应该在运行以下命令之前，将 --quant_model_path 修改为脚本中量化模型的文件夹。

1. Train on RedPajama

bash examples/e2e_qp/Llama-2-7b/w2g64-redpajama.sh

2. 在阿尔帕卡上训练

bash examples/e2e_qp/Llama-2-7b/w2g128-redpajama.sh

具体而言，--learning_rate 在我们的实验中对 2 位数为 2e-5，而对 3/4 位数为 1e-5。您可以减少 --per_device_train_batch_size 来降低训练过程中的内存占用，并确保 --gradient_accumulation_steps 按相同比例增加，以保持相同的批量大小。

推断

1. 从 Huggingface 下载预量化的 EfficientQAT 模型

pip install huggingface_hub

huggingface-cli download ChenMnZ/Llama-2-7b-EfficientQAT-w2g64 --local-dir ./output/pre_quantized_models/Llama-2-7b-EfficientQAT-w2g64

1. 评估预量化的 EfficientQAT 模型

CUDA_VISIBLE_DEVICES=0 python main_block_ap.py \
--resume_quant ./output/pre_quantized_models/Llama-2-7b-EfficientQAT-w2g64 \
--net Llama-2 \
--wbits 2 \
--group_size 64 \
--output_dir ./output/inference_results/ \
--eval_ppl \
--eval_tasks  piqa,arc_easy,arc_challenge,hellaswag,winogrande

模型传输

首先，您应该安装gptqmodel包以支持 GPTQ 和 BitBLAS 量化格式：

git clone https://github.com/ModelCloud/GPTQModel.git && cd GPTQModel
bash install.sh

• 根据我们的经验，我们使用gptqmodel v0.9.8进行测试。

然后，我们提供以下三种转账方式：

1. 将 EfficientQAT 检查点转移到 GPTQ 格式

bash examples/model_transfer/efficientqat_to_gptq/llama-2-7b.sh

• 注意 : 目前 AutoGPTQ 存在用于非对称量化的溢出错误。因此，我们选择官方修复错误的版本 GPTQModel 来传输我们的非对称量化模型。因此，该库提供的 GPTQ 模型只能通过 GPTQModel 成功加载，否则是 AutoGPTQ。

1. 将 EfficientQAT 检查点转移到 BitBLAS 格式

bash examples/model_transfer/efficientqat_to_bitblas/llama-2-7b.sh

• 加速存在一些问题，详见此问题。

1. 将 EfficientQAT 检查点中的 fp32 数据转移到半精度对应项。一些参数在训练时以 fp32 保存，您可以将它们转移为半精度，以进一步减少训练后的模型大小。

bash examples/model_transfer/fp32_to_16/llama-2-7b.sh

其他格式的推断

以下是使用 GPTQ 或 BitBLAS 量化格式进行推理的示例。

from transformers import AutoTokenizer
from gptqmodel import GPTQModel

quant_dir = "ChenMnZ/Llama-2-7b-EfficientQAT-w2g128-GPTQ"
# quant_dir = "ChenMnZ/Llama-2-7b-EfficientQAT-w2g128-BitBLAS"
# or local path

tokenizer = AutoTokenizer.from_pretrained(quant_dir, use_fast=True)


# load quantized model to the first GPU
model = GPTQModel.from_quantized(quant_dir)

# inference with model.generate
print(tokenizer.decode(model.generate(**tokenizer("Model quantization is", return_tensors="pt").to(model.device))[0]))