Exportações quantizadas

Cada configuração de export pode receber um quantizer, que quantiza o modelo durante o export. Os quantizadores ficam em transformers.exporters.quantizers, e cada um é executado em um ponto fixo do export: no grafo FX do torch.export, antes de qualquer conversão, ou no modelo para o qual um backend o converte.

A exportação verifica o quantizador em relação ao seu backend e aos pacotes instalados antes do tracing, de modo que uma incompatibilidade falha precocemente. Assim como os exportadores, cada quantizador lista as versões dos pacotes em que foi testado e emite um aviso quando as suas diferem.

O exemplo abaixo quantiza um classificador de sentimento para int8 com ONNX Runtime e, em seguida, o executa.

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 um quantizador PT2E, do torchao ou do runtime de destino, e o executa no grafo FX (prepare_pt2e → calibrar → convert_pt2e), de modo que funciona com todos os backends. Cada quantizador PT2E injeta suas próprias operações de quantize/dequantize, portanto, escolha aquele que seu runtime suporta:

Para o ExecuTorch, envolva o próprio XNNPACKQuantizer do 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))

Quantize modelos generativos por meio de ~HfExporter.export_for_generation. Aqui, calibration_dataset contém os kwargs de generate, como inputs, e cada componente realiza a calibração com base nos inputs capturados para ele. Um modelo apenas de texto possui um componente prefill e um componente decode; um modelo multimodal divide ainda mais seu prefill, em seus codificadores de modalidade (como image_encoder), seu projetor, seu language_model e seu lm_head (veja Como o export_for_generation funciona).

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": ...}

No OpenVINO, os pesos PT2E permanecem em precisão completa atrás de pares quantize/dequantize, que o OpenVINO converte para FakeQuantize e comprime para int8 por conta própria.

ONNX Runtime

OnnxRuntimeQuantizer executa quantize_static por padrão (ativações e pesos int8 QDQ, calibrados em calibration_dataset), conforme o exemplo acima, com os argumentos de um OnnxRuntimeStaticConfig. Com um OnnxRuntimeDynamicConfig, ele executa quantize_dynamic (pesos int8, ativações quantizadas em tempo de execução, sem calibração).

from transformers.exporters import OnnxRuntimeDynamicConfig, OnnxRuntimeStaticConfig

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

NNCF

Os três quantizadores NNCF na tabela acima executam nncf.quantize por padrão (ativações e pesos int8, calibrados com calibration_dataset, com model_type=nncf.ModelType.TRANSFORMER a menos que você passe outro), com os argumentos de um NNCFQuantizationConfig. Com um NNCFWeightCompressionConfig, eles executam nncf.compress_weights (somente pesos int8 ou int4, com AWQ e estimativa de escala para int4, além de correção GPTQ e LoRA no OpenVINO). O ExecuTorch reduz apenas o resultado somente de pesos do 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)

As ativações em Int8 executam multiplicações de matrizes em int8, que sofrem saturação em CPUs sem VNNI (Intel apenas com AVX2 e AMD Zen 2 ou mais antigo), portanto, espere uma precisão menor nesses casos. A compressão apenas de pesos mantém as multiplicações de matrizes em ponto flutuante.

Calibração

O calibration_dataset de um quantizador é um iterável de entradas do modelo, como uma lista de saídas do tokenizer ou um DataLoader cujos lotes são colados em entradas do modelo. É um atributo comum, portanto, você também pode defini-lo após a criação do quantizador (quantizer.calibration_dataset = samples). Se omitir, um quantizador que requer dados fará a calibração com as próprias entradas de amostra da exportação, com um aviso, pois uma única amostra pode distorcer os intervalos observados. Quantizadores que não requerem dados (quantização dinâmica do ONNX Runtime, compressão de pesos int8) nunca o leem.

Uma receita diferente para cada componente

Passe um dicionário {component: config} para ~HfExporter.export_for_generation para quantizar cada componente em seus próprios termos. Para o modelo acima, int8 estático em prefill e o int8 dinâmico mais leve em 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)

O dict deve nomear todos os componentes que decompose_for_generation() produz (um modelo multimodal adiciona seus codificadores e pode separar language_model e lm_head). Um componente cuja config não define um quantizador permanece em precisão total. Neste dict, o calibration_dataset de um quantizador contém os kwargs do forward do componente, não os kwargs de generate.

Adicionar um quantizador

Os quantizadores estão em transformers.exporters.quantizers, com uma classe para cada toolchain e representação. Subclasse ExportQuantizer e declare:

• quantizes_fx_graph: True para o grafo FX do torch.export, antes de qualquer conversão (de modo que cada backend recebe o resultado), ou False (o padrão) para o modelo em que um backend o converte (onnx_ir.Model, openvino.Model).
• supported_formats: os formatos de exportação que ele quantiza. A exportação rejeita qualquer outro formato antes do tracing.
• required_packages, e opcionalmente min_versions e tested_versions: o que ele importa. A exportação verifica se eles estão instalados antes do tracing, levanta um erro se a versão for inferior ao mínimo e emite um aviso quando a versão difere da testada.

Em seguida, implemente quantize(model, sample_inputs, export_format). model é o grafo FX ou o modelo de backend, e sample_inputs são as próprias entradas da exportação. Faça a calibração em um CalibrationSet: ele lê o calibration_dataset do quantizador (ou sample_inputs quando não há um) e mapeia cada amostra para as próprias entradas do modelo com um feed de transformers.exporters.quantizers.utils (traced_inputs para o 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))