量化导出
每个 export 配置都可以接受一个 quantizer,用于在导出过程中对模型进行量化。这些量化器位于 transformers.exporters.quantizers 中,并且每个量化器都在导出的固定阶段运行:要么在 torch.export FX 图上(在任何转换之前),要么在后端将其转换成的模型上。
| 量化器 | 量化 | Dynamo | ONNX | ExecuTorch | OpenVINO |
|---|---|---|---|---|---|
| PT2EQuantizer | FX 图,使用 PyTorch 2 Export (PT2E) 量化 | 🟢 | 🟢 | 🟢 | 🟢 |
| NNCFTorchFXQuantizer | FX 图,使用 NNCF | 🟢 | 🟢 | 🟢(仅权重) | 🟢 |
| OnnxRuntimeQuantizer | ONNX 模型,使用 onnxruntime.quantization | 🔴 | 🟢 | 🔴 | 🔴 |
| NNCFOnnxQuantizer | 使用 NNCF 的 ONNX 模型 | 🔴 | 🟢 | 🔴 | 🔴 |
| NNCFOpenVINOQuantizer | 带有 NNCF 的 OpenVINO 模型 | 🔴 | 🔴 | 🔴 | 🟢 |
导出操作会在进行追踪之前,检查量化器与其后端及已安装软件包是否匹配,因此任何不匹配都会尽早报错。 与导出器类似,每个量化器都会列出其经过测试的软件包版本,并在您的版本与之不一致时发出警告。
下面的示例使用 ONNX Runtime 将情感分类器量化为 int8,然后运行它。
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers.exporters import OnnxConfig, OnnxExporter, OnnxRuntimeQuantizer
model_id = "distilbert/distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)
texts = [
"A wonderful, heartfelt film.",
"The plot was dull and the acting worse.",
"I would watch it again tomorrow.",
"Two hours I will never get back.",
]
calibration_dataset = [tokenizer(text, return_tensors="pt") for text in texts]
inputs = tokenizer("An absolute delight from start to finish.", return_tensors="pt")
config = OnnxConfig(dynamic=True, quantizer=OnnxRuntimeQuantizer(calibration_dataset=calibration_dataset))
onnx_program = OnnxExporter().export(model, inputs, config)
logits = onnx_program(**inputs)[0]
print(model.config.id2label[int(logits.argmax())]) # POSITIVEPT2E
PT2EQuantizer 封装了来自 torchao 或目标运行时的 PT2E 量化器,并在 FX 图(prepare_pt2e → 校准 → convert_pt2e)上运行,因此适用于所有后端。每个 PT2E 量化器都会注入其自身的量化/反量化操作,因此请选择您的运行时所支持的量化器:
| 运行位置 | PT2E 量化器 | 从 |
|---|---|---|
| PyTorch inductor、ONNX Runtime (QDQ)、OpenVINO | X86InductorQuantizer | torchao.quantization.pt2e.quantizer.x86_inductor_quantizer |
| ExecuTorch XNNPACK 后端 | XNNPACKQuantizer | executorch.backends.xnnpack.quantizer.xnnpack_quantizer |
对于 ExecuTorch,请封装 ExecuTorch 自身的 XNNPACKQuantizer:
from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
XNNPACKQuantizer,
get_symmetric_quantization_config,
)
from transformers.exporters import ExecutorchConfig, PT2EQuantizer
xnnpack = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
config = ExecutorchConfig(backend="xnnpack", quantizer=PT2EQuantizer(xnnpack, calibration_dataset=calibration_dataset))通过 ~HfExporter.export_for_generation 对生成式模型进行量化。其中 calibration_dataset 包含 generate 的关键字参数,例如 inputs,每个组件都会使用为其捕获的输入进行校准。纯文本模型包含 prefill 和 decode 两个组件;多模态模型会进一步拆分其 prefill,包括其模态编码器(例如 image_encoder)、其投影器、其 language_model 以及其 lm_head(参见 export_for_generation 的工作原理)。
from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import (
X86InductorQuantizer,
get_default_x86_inductor_quantization_config,
)
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.exporters import DynamoConfig, DynamoExporter, PT2EQuantizer
model_id = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
prompts = ["The capital of France is", "Photosynthesis turns sunlight into", "The first moon landing was in"]
calibration_dataset = [tokenizer(prompt, return_tensors="pt") for prompt in prompts]
inputs = tokenizer("Once upon a time", return_tensors="pt")
x86 = X86InductorQuantizer().set_global(get_default_x86_inductor_quantization_config())
config = DynamoConfig(dynamic=True, quantizer=PT2EQuantizer(x86, calibration_dataset=calibration_dataset))
components = DynamoExporter().export_for_generation(model, inputs, config) # {"prefill": ..., "decode": ...}在 OpenVINO 上,PT2E 的权重以全精度形式保留在量化/反量化(quantize/dequantize)配对之后,OpenVINO 会将其转换为 FakeQuantize,并自行将其压缩为 int8。
ONNX Runtime
OnnxRuntimeQuantizer 默认运行 quantize_static(int8 QDQ 激活值和权重,基于 calibration_dataset 进行校准),如上文示例所示,并使用 OnnxRuntimeStaticConfig 的参数。若使用 OnnxRuntimeDynamicConfig,则运行 quantize_dynamic(int8 权重,激活值在运行时量化,无需校准)。
from transformers.exporters import OnnxRuntimeDynamicConfig, OnnxRuntimeStaticConfig
config = OnnxConfig(
quantizer=OnnxRuntimeQuantizer(OnnxRuntimeStaticConfig(per_channel=True), calibration_dataset=calibration_dataset)
)
config = OnnxConfig(quantizer=OnnxRuntimeQuantizer(OnnxRuntimeDynamicConfig()))NNCF
上表中的三个 NNCF 量化器默认运行 nncf.quantize(对激活值和权重进行 int8 量化,基于 calibration_dataset 进行校准,除非传入其他值,否则 model_type=nncf.ModelType.TRANSFORMER),其参数为 NNCFQuantizationConfig。使用 NNCFWeightCompressionConfig 时,它们运行 nncf.compress_weights(仅权重的 int8 或 int4 量化,int4 支持 AWQ 和 scale estimation,并在 OpenVINO 上支持 GPTQ 和 LoRA 校正)。ExecuTorch 仅支持 NNCFTorchFXQuantizer 的仅权重量化结果。
import nncf
from transformers.exporters import (
DynamoConfig,
NNCFOnnxQuantizer,
NNCFOpenVINOQuantizer,
NNCFTorchFXQuantizer,
NNCFWeightCompressionConfig,
OnnxConfig,
OpenVINOConfig,
OpenVINOExporter,
)
int4 = NNCFWeightCompressionConfig(mode=nncf.CompressWeightsMode.INT4_SYM, group_size=128)
# OpenVINO: int8 activations and weights, or weight-only int4
openvino_int8 = OpenVINOConfig(quantizer=NNCFOpenVINOQuantizer(calibration_dataset=calibration_dataset))
openvino_int4 = OpenVINOConfig(quantizer=NNCFOpenVINOQuantizer(int4))
# ONNX: weight-only int4, as ONNX Runtime's `MatMulNBits`
onnx_int4 = OnnxConfig(quantizer=NNCFOnnxQuantizer(int4))
# Dynamo: int8 on the FX graph
dynamo_int8 = DynamoConfig(quantizer=NNCFTorchFXQuantizer(calibration_dataset=calibration_dataset))
ov_model = OpenVINOExporter().export(model, inputs, openvino_int4)Int8 激活会执行 int8 矩阵乘法,而在没有 VNNI 的 CPU(仅支持 AVX2 的 Intel 以及 AMD Zen 2 及更早版本)上会发生饱和,因此请预期在这些平台上精度会降低。仅权重压缩会将矩阵乘法保持在浮点运算中。
校准
量化器的 calibration_dataset 是一个模型输入的迭代器,例如分词器输出的列表,或一个批次可整理为模型输入的 DataLoader。它是一个普通属性,因此也可以在创建量化器之后对其进行设置(quantizer.calibration_dataset = samples)。如果省略该属性,需要数据校准的量化器将使用导出自身的样本输入进行校准,并会发出警告,因为单个样本可能会扭曲观察到的范围。无需数据的量化器(动态 ONNX Runtime 量化、int8 权重压缩)永远不会读取该属性。
每个组件采用不同的配方
将 {component: config} 字典传递给 ~HfExporter.export_for_generation,即可按各自的条件对每个组件进行量化。对于上述模型,在 prefill 上使用静态 int8,在 decode 上使用更轻量的动态 int8:
dynamic_x86 = X86InductorQuantizer().set_global(get_default_x86_inductor_quantization_config(is_dynamic=True))
config = {
"prefill": DynamoConfig(dynamic=True, quantizer=PT2EQuantizer(x86)),
"decode": DynamoConfig(dynamic=True, quantizer=PT2EQuantizer(dynamic_x86)),
}
components = DynamoExporter().export_for_generation(model, inputs, config)该 dict 必须为 decompose_for_generation() 生成的每个组件命名(多模态模型会添加其编码器,并可能将 language_model 和 lm_head 拆分出来)。如果某个组件的 config 未设置量化器,则该组件保持全精度。在此 dict 中,量化器的 calibration_dataset 保存的是该组件的 forward kwargs,而非 generate kwargs。
添加量化器
量化器位于 transformers.exporters.quantizers 中,每个工具链和表示方式对应一个类。请继承 ExportQuantizer 并声明:
quantizes_fx_graph:对于torch.exportFX 图,在转换之前(因此每个后端都会接收该结果)设为True,或者对于后端将其转换成的模型(onnx_ir.Model、openvino.Model)设为False(默认值)。supported_formats:其进行量化的导出格式。在追踪之前,导出会拒绝任何其他格式。required_packages,以及可选的min_versions和tested_versions:即该模块导入的内容。导出时会在追踪之前检查它们是否已安装,若版本低于最低要求则抛出异常,若版本与经过测试的版本不一致则发出警告。
然后实现 quantize(model, sample_inputs, export_format)。model 是 FX 图或后端模型,sample_inputs 是导出自身的输入。在 CalibrationSet 上进行校准:它会读取量化器的 calibration_dataset(如果没有,则使用 sample_inputs),并通过 transformers.exporters.quantizers.utils 中的 feed 将每个样本映射到模型自身的输入(FX 图使用 traced_inputs,ONNX 使用 onnx_feed,OpenVINO 使用 openvino_feed)。
from transformers.exporters import ExportFormat
from transformers.exporters.quantizers import CalibrationSet, ExportQuantizer
from transformers.exporters.quantizers.utils import onnx_feed
class MyOnnxQuantizer(ExportQuantizer):
supported_formats = (ExportFormat.ONNX,)
required_packages = ("my_toolkit",)
tested_versions = {"my_toolkit": "1.0.0"}
def quantize(self, model, sample_inputs, export_format):
import my_toolkit
names = [graph_input.name for graph_input in model.graph.inputs]
calibration = CalibrationSet(self.calibration_dataset, sample_inputs, lambda sample: onnx_feed(names, sample))
return my_toolkit.quantize(model, calibration_feeds=list(calibration))