使用 Intel® Neural Compressor 轻松量化 PyTorch#
创建日期:2022 年 1 月 11 日 | 最后更新:2025 年 7 月 28 日 | 最后验证:未验证
概述#
大多数深度学习应用在推理时使用 32 位浮点精度。但低精度数据类型(如 fp8)因显著的性能提升而受到更多关注。采用低精度的一个关键问题是在满足预定要求的同时减轻精度损失。
Intel® Neural Compressor 旨在通过扩展 PyTorch 的面向精度的自动调优策略来解决上述问题,帮助用户在 Intel 硬件上快速找到最佳量化模型。
Intel® Neural Compressor 是一个开源项目,位于 Github。
特点#
易用 API: Intel® Neural Compressor 重新使用了 PyTorch 的
prepare、convertAPI 供用户使用。面向精度的调优: Intel® Neural Compressor 支持面向精度的自动调优过程,提供
autotuneAPI 供用户使用。量化类型: Intel® Neural Compressor 支持多种量化方法,包括经典的 INT8 量化、权重独占量化以及流行的 FP8 量化。Neural compressor 还提供了最新的模拟工作研究,例如 MX 数据类型仿真量化。有关更多详细信息,请参阅 支持矩阵。
入门#
安装#
# install stable version from pip
pip install neural-compressor-pt
注意:Neural Compressor 提供自动加速器检测,包括 HPU、Intel GPU、CUDA 和 CPU。要指定目标设备,建议使用 INC_TARGET_DEVICE,例如 export INC_TARGET_DEVICE=cpu。
示例#
本节展示了使用 Intel® Neural compressor 进行各种量化的示例
FP8 量化#
FP8 量化由 Intel® Gaudi®2&3 AI 加速器 (HPU) 支持。要准备环境,请参阅 Intel® Gaudi® 文档。
运行示例,
# FP8 Quantization Example
from neural_compressor.torch.quantization import (
FP8Config,
prepare,
convert,
)
import torch
import torchvision.models as models
# Load a pre-trained ResNet18 model
model = models.resnet18()
# Configure FP8 quantization
qconfig = FP8Config(fp8_config="E4M3")
model = prepare(model, qconfig)
# Perform calibration (replace with actual calibration data)
calibration_data = torch.randn(1, 3, 224, 224).to("hpu")
model(calibration_data)
# Convert the model to FP8
model = convert(model)
# Perform inference
input_data = torch.randn(1, 3, 224, 224).to("hpu")
output = model(input_data).to("cpu")
print(output)
权重独占量化#
权重独占量化也支持在 Intel® Gaudi®2&3 AI 加速器上运行。量化模型可以如下加载。
from neural_compressor.torch.quantization import load
# The model name comes from HuggingFace Model Hub.
model_name = "TheBloke/Llama-2-7B-GPTQ"
model = load(
model_name_or_path=model_name,
format="huggingface",
device="hpu",
torch_dtype=torch.bfloat16,
)
注意:Intel Neural Compressor 将在首次加载时将模型格式从 auto-gptq 转换为 hpu 格式,并将 hpu_model.safetensors 保存到本地缓存目录以便下次加载。因此,首次加载可能需要一些时间。
使用 PT2E 后端进行静态量化#
PT2E 路径使用 torch.dynamo 将 eager 模型捕获为 FX 图模型,然后在其上插入 observer 和 Q/QD 对。最后使用 torch.compile 执行模式匹配,并将 Q/DQ 对替换为优化的量化运算符。
使用 PT2E 后端执行 W8A8 静态量化有四个步骤:export、prepare、convert 和 compile。
import torch
from neural_compressor.torch.export import export
from neural_compressor.torch.quantization import StaticQuantConfig, prepare, convert
# Prepare the float model and example inputs for export model
model = UserFloatModel()
example_inputs = ...
# Export eager model into FX graph model
exported_model = export(model=model, example_inputs=example_inputs)
# Quantize the model
quant_config = StaticQuantConfig()
prepared_model = prepare(exported_model, quant_config=quant_config)
# Calibrate
run_fn(prepared_model)
q_model = convert(prepared_model)
# Compile the quantized model and replace the Q/DQ pattern with Q-operator
from torch._inductor import config
config.freezing = True
opt_model = torch.compile(q_model)
面向精度的调优#
要利用面向精度的自动调优,需要指定调优空间。 autotune 迭代调优空间,将配置应用于给定的高精度模型,然后记录并比较其评估结果与基线。调优过程在满足退出策略时停止。
from neural_compressor.torch.quantization import RTNConfig, TuningConfig, autotune
def eval_fn(model) -> float:
return ...
tune_config = TuningConfig(
config_set=RTNConfig(use_sym=[False, True], group_size=[32, 128]),
tolerable_loss=0.2,
max_trials=10,
)
q_model = autotune(model, tune_config=tune_config, eval_fn=eval_fn)
教程#
在 Intel® Neural Compressor 的官方 文档 中可以找到更详细的教程。