1. 写在前面
Llama 是 Meta AI 开源的一系列大型语言模型 (LLM)，在各种 NLP 任务上表现出色。然而，这些模型通常具有庞大的参数量，需要大量的计算资源和内存才能进行推理。为了降低 Llama 模型的部署成本，并提高其推理速度，我们可以采用模型量化 (Quantization) 技术。
本文将介绍 Llama 模型的量化方法，以及如何使用量化后的模型进行高效推理。
2. 模型量化概述
模型量化是一种将模型的权重和激活值从高精度 (例如 FP32, FP16) 转换为低精度 (例如 INT8, INT4) 的技术。通过降低模型的精度，可以：
- 减小模型大小：INT8 类型的数据大小是 FP32 的 1/4。
- 加快推理速度：低精度计算通常比高精度计算更快。
- 降低内存占用：更小的模型和中间激活值可以减少内存占用。
- 降低功耗：在某些硬件上，低精度计算可以降低功耗。
常见的量化方法：
- 训练后量化 (Post-Training Quantization, PTQ)：直接对训练好的模型进行量化，无需重新训练。
- 量化感知训练 (Quantization-Aware Training, QAT)：在训练过程中模拟量化操作，使模型更好地适应量化带来的精度损失。
量化的粒度：
- Per-tensor 量化：对整个张量使用一个缩放因子。
- Per-channel 量化：对张量的每个通道使用一个缩放因子，通常用于权重的量化。
3. Llama 模型的量化方法
Llama 模型通常使用 训练后量化 (PTQ) 方法进行量化，因为它更简单、更快速。常用的 PTQ 方法包括：
3.1. GPTQ
GPTQ 是一种针对 GPT 类模型的 PTQ 方法，它通过逐层量化并进行误差补偿，来减少精度损失。它将权重矩阵分解为多个块，并对每个块进行量化，同时使用一种称为 Optimal Brain Quantization (OBQ) 的技术来更新未量化的权重，以补偿量化带来的误差。
- 优点： 精度损失小，支持 INT4 和 INT3 量化。
- 缺点： 量化过程比较耗时。
使用 auto-gptq 库进行 GPTQ 量化：
# 安装依赖
# pip install auto-gptq
from transformers import AutoTokenizer, AutoModelForCausalLM
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
import torch
# 加载预训练模型和分词器
model_name = "meta-llama/Llama-2-7b-hf"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 定义量化配置
quantize_config = BaseQuantizeConfig(
bits=4, # 量化比特数
group_size=128, # 分组大小
desc_act=False, # 是否使用激活顺序
)
# 量化模型
quantized_model = AutoGPTQForCausalLM.from_pretrained(
model_name,
quantize_config=quantize_config,
device_map="auto",
use_safetensors=True,
trust_remote_code=False
)
# 使用量化后的模型进行推理
input_text = "The quick brown fox"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
# generate text
outputs = quantized_model.generate(input_ids)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
# 可选: 保存量化后的模型
# quantized_model.save_pretrained("llama-2-7b-4bit-gptq")
# tokenizer.save_pretrained("llama-2-7b-4bit-gptq")
更具体地，bits参数可以选择8 4 3 2来指定不同的量化精度，group_size越小精度损失越小，但占用的显存越大，不指定时默认group_size=128，关于激活顺序desc_act，通常设置为False
3.2. AWQ
AWQ (Activation-aware Weight Quantization) 是一种基于权重的训练后量化方法，AWQ 观察到，并非所有权重都同等重要：只有 1% 的权重对模型性能至关重要。于是 AWQ 会首先找到并保护这些重要的权重，根据激活值的统计信息进行量化。它通过观察激活值的分布，对权重进行缩放，使得缩放后的权重更适合量化，从而减少精度损失。
- 优点： 速度快，精度损失小，特别适用于文本生成任务。
- 缺点： 只支持 INT4 量化。
使用 auto-awq 库进行 AWQ 量化:
# 安装依赖
# pip install auto-awq
from transformers import AutoTokenizer, AutoModelForCausalLM
from awq import AutoAWQForCausalLM
# 加载预训练模型和分词器
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 量化模型
quantized_model = AutoAWQForCausalLM.from_pretrained(model_name)
quantized_model.quantize(tokenizer)
# 使用量化后的模型进行推理
input_text = "The quick brown fox"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
# generate text
outputs = quantized_model.generate(input_ids)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
# 可选: 保存量化后的模型
# quantized_model.save_quantized("llama-2-7b-4bit-awq")
# tokenizer.save_pretrained("llama-2-7b-4bit-awq")
3.3. 其他量化方法
除了 GPTQ 和 AWQ，还有一些其他的量化方法，例如：
- BitsAndBytes: 支持 INT8 量化。
- LLM.int8(): 提出了一种针对大模型的 INT8 量化方法。
4. 量化模型的推理
使用量化后的模型进行推理，通常需要使用支持量化推理的框架或库，例如：
- llama.cpp：一个用 C++ 编写的推理引擎，支持多种量化格式 (例如 GGUF、GGML)，并且可以在 CPU 和 GPU 上运行。
- ExLlama: 一个用 Python 和 CUDA 编写的推理引擎，支持 GPTQ 量化。
- vLLM：一个快速且易于使用的 LLM 推理和服务库，支持多种量化方法。
使用 llama.cpp 进行推理的示例：
- 安装
llama.cpp：
git clone https://github.com/ggerganov/llama.cpp.git
cd llama.cpp
make
- 转换模型为 GGUF 格式：
# 首先安装转换工具所需要的依赖
python3 -m pip install -r requirements.txt
# 假设你已经有了使用auto-gptq或者auto-awq量化的模型，假设模型目录为llama-2-7b-4bit-gptq
# 使用convert.py 脚本将模型转换为 GGUF 格式，例如：
python3 convert.py llama-2-7b-4bit-gptq \
--outfile llama-2-7b-4bit-gptq.gguf \
--outtype q4_0
这里 --outtype 可以指定输出的模型类型，例如q4_0表示4bit量化
- 使用
main程序进行推理：
./main -m llama-2-7b-4bit-gptq.gguf -p "The quick brown fox" -n 128
这里 -m 指定模型文件，-p 指定输入提示，-n 指定生成token的数量。