Квантование экспорта

Каждая конфигурация export может принимать параметр quantizer, который выполняет квантование модели в процессе экспорта. Квантизаторы находятся в модуле transformers.exporters.quantizers, и каждый из них запускается на определённом этапе экспорта: на графе FX, полученном от torch.export, до каких-либо преобразований, либо на модели, в которую бэкенд её преобразует.

Экспорт проверяет квантизатор на соответствие его бэкенду и установленным пакетам перед трассировкой, поэтому несоответствие обнаруживается на раннем этапе. Как и экспортеры, каждый квантизатор перечисляет версии пакетов, на которых он протестирован, и предупреждает, если ваши версии отличаются.

Пример ниже квантует классификатор тональности до int8 с помощью ONNX Runtime, а затем запускает его.

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 инкапсулирует квантизатор PT2E из torchao или целевого рантайма и запускает его на графе FX (prepare_pt2e → калибровка → convert_pt2e), поэтому он работает с любым бэкендом. Каждый квантизатор PT2E вставляет собственные операции квантизации/деквантизации, поэтому выбирайте тот, который поддерживает ваш рантайм:

Для ExecuTorch оберните собственный XNNPACKQuantizer из ExecuTorch:

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 содержит аргументы для генерации, такие как 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, с AWQ и оценкой масштаба для int4, а также коррекцией GPTQ и LoRA в OpenVINO). 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, которые насыщаются на CPU без поддержки VNNI (только Intel с AVX2 и AMD Zen 2 и более старые), поэтому ожидайте более низкую точность на таких процессорах. Сжатие только весов сохраняет умножения матриц в формате с плавающей запятой.

Калибровка

Атрибут calibration_dataset квантизатора — это итерируемый объект, содержащий входные данные модели, например список результатов токенизатора или DataLoader, чьи батчи формируют входные данные модели. Это обычный атрибут, поэтому его можно задать и после создания квантизатора (quantizer.calibration_dataset = samples). Если этот атрибут не указан, квантизатор, которому требуются данные для калибровки, выполнит её на собственных примерах входных данных экспорта, выдав предупреждение, так как один пример может исказить наблюдаемые диапазоны значений. Квантизаторы, которым данные не нужны (динамическая квантизация ONNX Runtime, сжатие весов до int8), никогда не обращаются к этому атрибуту.

Разный рецепт для каждого компонента

Передайте словарь {component: config} в ~HfExporter.export_for_generation, чтобы квантовать каждый компонент по своим условиям. Для модели выше: статический int8 для prefill и более лёгкий динамический int8 для decode:

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)

Словарь должен указывать каждый компонент, который создаёт decompose_for_generation() (мультимодальная модель добавляет свои энкодеры и может выделить language_model и lm_head). Компонент, в конфигурации которого не задан квантизатор, остаётся в полной точности. В этом словаре calibration_dataset квантизатора содержит forward-аргументы компонента, а не аргументы generate.

Добавьте квантизатор

Квантизаторы находятся в transformers.exporters.quantizers, по одному классу на каждый инструмент и представление. Наследуйте ExportQuantizer и объявите:

• quantizes_fx_graph: True для графа FX из torch.export до какого-либо преобразования (таким образом, каждый бэкенд получает исходный результат) или False (по умолчанию) для модели, в которую бэкенд преобразует граф (onnx_ir.Model, openvino.Model).
• 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 (traced_inputs для графа FX, onnx_feed для ONNX, openvino_feed для OpenVINO).

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))