W8A8量化
什么是W8A8量化
W8A8量化是一种将模型权重和激活值都量化为8位数据的技术。该技术把高位浮点数转为8位,通常是将权重和激活值从16位或32位浮点数转换为8位整数(int8)格式。量化后,模型权重体积会减少,同时使用int8格式数据进行矩阵乘法(MatMul)运算时,可减少计算量,进而提升推理性能。
W8A8量化方案能降低模型显存以及需要部署的卡数。也能同时降低首token时延和增量推理时延。
约束限制
- 支持W8A8量化的模型列表请参见支持的模型列表。
- 激活量化支持动态per-token,支持对称量化。
- 权重量化支持per-channel,支持对称量化。
获取量化后的模型权重
有两种方式获取量化模型权重:
方式一:可以在Huggingface开源社区获取量化后的模型权重,具体参考从开源社区下载发布的llm-compressor量化模型。
方式二:获取FP16/BF16的模型权重之后,通过llm-compressor工具进行量化,具体参考使用llm-compressor量化工具进行量化。
使用llm-compressor工具量化模型
本章节介绍如何在NPU的机器上使用开源量化工具llm-compressor量化模型权重,然后在NPU的机器上实现推理。了解更多开源量化工具信息,请参考llm-compressor。
- 使用该量化工具,需要切换conda环境,运行以下命令:
conda create --name llmcompressor --clone PyTorch-2.5.1 conda activate llmcompressor
- 安装llm-compressor
pip install llmcompressor==0.6.0 pip install transformers==4.51.3 pip install zstandard
- 配置校准数据集
W8A8量化需要一个校准数据集。默认使用在线数据集 HuggingFaceH4/ultrachat_200k。根据网络环境选择以下一种方式。
方式一:在线使用(环境可访问HuggingFace)
无需额外操作,脚本默认配置如下
DATASET_ID = "HuggingFaceH4/ultrachat_200k" DATASET_SPLIT = "train_sft" NUM_CALIBRATION_SAMPLES = 512 MAX_SEQUENCE_LENGTH = 2048
方式二:离线使用(环境无法访问HuggingFace)
手动下载数据集并上传至服务器,按以下操作。
- 下载数据集
通过浏览器访问 ultrachat_200k 数据集页面,下载data目录下的所有parquet文件,
- 在服务器上组织目录结构
上传后,必须保留原始的data目录,目录层级如下(<local-data-dir>是存放数据集的自定义目录):
<local-data-dir>/ └── ultrachat_200k/ # 数据集主文件夹 └── data/ # 原始数据目录(不要改名) ├── train_sft-00000.parquet ├── train_sft-00001.parquet └── ...注意:ultrachat_200k 是主文件夹名,其下必须直接包含 data 子目录。
- 修改脚本中的DATASET_ID
编辑 llm_compressor_W8A8.py,将 DATASET_ID 设置为服务器上 ultrachat_200k 文件夹的绝对路径(即 data 目录的上一级,不要包含 data)
# 正确示例(假设数据集存放在 /home/ma-user/datasets/ultrachat_200k) DATASET_ID = "/home/ma-user/datasets/ultrachat_200k" # 错误示例(路径中包含了 data,会导致无法识别) # DATASET_ID = "/home/ma-user/datasets/ultrachat_200k/data"
DATASET_SPLIT、NUM_CALIBRATION_SAMPLES、MAX_SEQUENCE_LENGTH 保持默认值即可,一般无需修改
完整的llm_compressor_W8A8.py脚本请参见llm_compressor_W8A8.py脚本。
- 下载数据集
- 运行"llm_compressor_W8A8.py"文件进行模型量化,量化时间和模型大小有关,预计30分钟~3小时。
# 根据机器卡的空闲情况指定量化使用的卡号,不设置则默认使用0号卡 export ASCEND_RT_VISIBLE_DEVICES=0 python llm_compressor_W8A8.py --model-path /home/ma-user/Qwen2.5-72B/ --quant-path /home/ma-user/Qwen2.5-72B-quant/
参数说明:
- --model-path:原始模型权重路径
- --quant-path:转换后保存权重路径
量化执行过程中如果打印如下WARNING日志,可以不用关注,不影响量化功能正常使用:
get_GPU_usage_nv | WARNING - Pynml library error: NVML Shared Library Not Found
- 启动在线推理服务,在启动服务时添加如下命令。
-q compressed-tensors 或者--quantization compressed-tensors
llm_compressor_W8A8.py脚本
llm_compressor_W8A8.py脚本完整代码如下
import argparse
import os
from functools import partial
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from datasets import load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier
from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
from llmcompressor.utils import dispatch_for_generation
from transformers import AutoModelForCausalLM, AutoTokenizer
os.environ['PYTORCH_NPU_ALLOC_CONF'] = 'expandable_segments:False'
# Select calibration dataset.
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"
# Select number of samples. 512 samples is a good place to start.
# Increasing the number of samples can improve accuracy.
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 2048
def preprocess(example, tokenizer):
return {
"text": tokenizer.apply_chat_template(
example["messages"],
tokenize=False,
)
}
# Tokenize inputs.
def tokenize(sample, tokenizer):
return tokenizer(
sample["text"],
padding=False,
max_length=MAX_SEQUENCE_LENGTH,
truncation=True,
add_special_tokens=False,
)
def quantize(model_id, quantized_path):
# Select model and load it
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Load dataset and preprocess.
ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
ds = ds.shuffle(seed=42)
preprocess_func = partial(preprocess, tokenizer=tokenizer)
ds = ds.map(preprocess_func)
tokenize_func = partial(tokenize, tokenizer=tokenizer)
ds = ds.map(tokenize_func, remove_columns=ds.column_names)
# Configure algorithms. In this case, we:
# * apply SmoothQuant to make the activations easier to quantize
# * quantize the weights to int8 with GPTQ (static per channel)
# * quantize the activations to int8 (dynamic per token)
recipe = [
SmoothQuantModifier(smoothing_strength=0.8),
GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"]),
]
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=MAX_SEQUENCE_LENGTH,
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)
# Save to disk compressed.
model.save_pretrained(quantized_path, save_compressed=True)
tokenizer.save_pretrained(quantized_path)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, required=True, help="Path to the input model")
parser.add_argument("--quant-path", type=str, required=True, help="Path to save the compressed model")
args = parser.parse_args()
quantize(args.model_path, args.quant_path)