Transformers documentation

量化导出

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

量化导出

每个 export 配置都可以接受一个 quantizer,用于在导出过程中对模型进行量化。这些量化器位于 transformers.exporters.quantizers 中,并且每个量化器都在导出的固定阶段运行:要么在 torch.export FX 图上(在任何转换之前),要么在后端将其转换成的模型上。

量化器量化DynamoONNXExecuTorchOpenVINO
PT2EQuantizerFX 图,使用 PyTorch 2 Export (PT2E) 量化🟢🟢🟢🟢
NNCFTorchFXQuantizerFX 图,使用 NNCF🟢🟢🟢(仅权重)🟢
OnnxRuntimeQuantizerONNX 模型,使用 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())])  # POSITIVE

PT2E

PT2EQuantizer 封装了来自 torchao 或目标运行时的 PT2E 量化器,并在 FX 图(prepare_pt2e → 校准 → convert_pt2e)上运行,因此适用于所有后端。每个 PT2E 量化器都会注入其自身的量化/反量化操作,因此请选择您的运行时所支持的量化器:

运行位置PT2E 量化器从
PyTorch inductor、ONNX Runtime (QDQ)、OpenVINOX86InductorQuantizertorchao.quantization.pt2e.quantizer.x86_inductor_quantizer
ExecuTorch XNNPACK 后端XNNPACKQuantizerexecutorch.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.export FX 图,在转换之前(因此每个后端都会接收该结果)设为 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))
Update on GitHub