Exportaciones cuantizadas

Cada configuración de export puede aceptar un quantizer, que cuantiza el modelo durante la exportación. Los cuantizadores se encuentran en transformers.exporters.quantizers, y cada uno se ejecuta en un punto fijo de la exportación: sobre el grafo FX de torch.export, antes de cualquier conversión, o sobre el modelo al que un backend lo convierte.

La exportación verifica el cuantizador contra su backend y los paquetes instalados antes de realizar el trazado, por lo que una incompatibilidad se detecta de forma temprana. Al igual que los exportadores, cada cuantizador enumera las versiones de paquetes con las que se ha probado y emite una advertencia si las tuyas difieren.

El siguiente ejemplo cuantiza un clasificador de sentimiento a int8 con ONNX Runtime y luego lo ejecuta.

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 encapsula un cuantizador PT2E, de torchao o del runtime de destino, y lo ejecuta en el grafo FX (prepare_pt2e → calibrar → convert_pt2e), por lo que funciona con todos los backends. Cada cuantizador PT2E inyecta sus propias operaciones de cuantizar/descuantizar, así que elija el que su runtime soporte:

Para ExecuTorch, envuelva el propio XNNPACKQuantizer de 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))

Cuantice modelos generativos mediante ~HfExporter.export_for_generation. Aquí, calibration_dataset contiene los argumentos de generación, como inputs, y cada componente se calibra con las entradas capturadas para él. Un modelo solo de texto tiene un componente prefill y un componente decode; un modelo multimodal divide su prefill aún más, en sus codificadores de modalidad (como image_encoder), su proyector, su language_model y su lm_head (consulte Cómo funciona 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": ...}

En OpenVINO, los pesos de PT2E permanecen en precisión completa detrás de pares de cuantización/desquantización, que OpenVINO convierte en FakeQuantize y comprime a int8 por sí mismo.

ONNX Runtime

OnnxRuntimeQuantizer ejecuta quantize_static de forma predeterminada (activaciones y pesos int8 QDQ, calibrados con calibration_dataset), como en el ejemplo anterior, con los argumentos de una OnnxRuntimeStaticConfig. Con una OnnxRuntimeDynamicConfig ejecuta quantize_dynamic (pesos int8, activaciones cuantizadas en tiempo de ejecución, sin calibración).

from transformers.exporters import OnnxRuntimeDynamicConfig, OnnxRuntimeStaticConfig

config = OnnxConfig(
    quantizer=OnnxRuntimeQuantizer(OnnxRuntimeStaticConfig(per_channel=True), calibration_dataset=calibration_dataset)
)
config = OnnxConfig(quantizer=OnnxRuntimeQuantizer(OnnxRuntimeDynamicConfig()))

NNCF

Los tres cuantizadores de NNCF de la tabla anterior ejecutan nncf.quantize de forma predeterminada (activaciones y pesos int8, calibrados con calibration_dataset, con model_type=nncf.ModelType.TRANSFORMER a menos que se pase otro), con los argumentos de una NNCFQuantizationConfig. Con una NNCFWeightCompressionConfig ejecutan nncf.compress_weights (solo pesos int8 o int4, con AWQ y estimación de escala para int4, además de corrección GPTQ y LoRA en OpenVINO). ExecuTorch solo reduce el resultado de solo pesos de 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)

Las activaciones Int8 ejecutan matmuls int8, que se saturan en CPUs sin VNNI (Intel solo con AVX2 y AMD Zen 2 o anteriores), por lo que se espera una precisión menor en esos casos. La compresión solo de pesos mantiene los matmuls en punto flotante.

Calibración

El calibration_dataset de un cuantizador es un iterable de entradas del modelo, como una lista de salidas del tokenizador o un DataLoader cuyos lotes se combinan en entradas del modelo. Es un atributo simple, por lo que también puede configurarse después de crear el cuantizador (quantizer.calibration_dataset = samples). Si se omite, un cuantizador que necesite datos se calibrará con las propias entradas de muestra de la exportación, emitiendo una advertencia, ya que una sola muestra puede sesgar los rangos observados. Los cuantizadores que no requieren datos (cuantización dinámica de ONNX Runtime, compresión de pesos int8) nunca lo leen.

Una receta diferente por componente

Pase un diccionario {component: config} a ~HfExporter.export_for_generation para cuantizar cada componente por separado. Para el modelo anterior, int8 estático en prefill y el más ligero int8 dinámico en 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)

El diccionario debe nombrar cada componente que produce decompose_for_generation() (un modelo multimodal añade sus codificadores y puede separar language_model y lm_head). Un componente cuya configuración no establece ningún cuantizador se mantiene en precisión completa. En este diccionario, el calibration_dataset de un cuantizador contiene los kwargs de forward del componente, no los kwargs de generate.

Agregar un cuantizador

Los cuantizadores se encuentran en transformers.exporters.quantizers, con una clase por cada cadena de herramientas y representación. Subclase ExportQuantizer y declara:

• quantizes_fx_graph: True para el grafo FX de torch.export, antes de cualquier conversión (de modo que cada backend tome el resultado), o False (el valor predeterminado) para el modelo en el que un backend lo convierte (onnx_ir.Model, openvino.Model).
• supported_formats: los formatos de exportación que cuantiza. La exportación rechaza cualquier otro formato antes del trazado.
• required_packages, y opcionalmente min_versions y tested_versions: lo que importa. La exportación comprueba que estén instalados antes de trazar, lanza una excepción si la versión es inferior al mínimo y emite una advertencia cuando la versión difiere de la probada.

A continuación, implemente quantize(model, sample_inputs, export_format). model es el grafo FX o el modelo de backend, y sample_inputs son las propias entradas de la exportación. Calibre con un CalibrationSet: este lee el calibration_dataset del cuantizador (o sample_inputs cuando no hay ninguno) y mapea cada muestra a las propias entradas del modelo mediante un feed de transformers.exporters.quantizers.utils (traced_inputs para el grafo FX, onnx_feed para ONNX, openvino_feed para 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))