Support quantization in Dynamo, ONNX and ET exporters

conftest.py

@@ -248,6 +248,8 @@ def pytest_configure(config):
     config.addinivalue_line("markers", "not_device_test: mark the tests always running on cpu")
     config.addinivalue_line("markers", "torch_compile_test: mark test which tests torch compile functionality")
     config.addinivalue_line("markers", "torch_export_test: mark test which tests torch export functionality")
+    config.addinivalue_line("markers", "onnx_export_test: mark test which tests ONNX export functionality")
+    config.addinivalue_line("markers", "executorch_export_test: mark test which tests ExecuTorch export functionality")
     config.addinivalue_line("markers", "flash_attn_test: mark test which tests flash attention functionality")
     config.addinivalue_line("markers", "flash_attn_3_test: mark test which tests flash attention 3 functionality")
     config.addinivalue_line("markers", "flash_attn_4_test: mark test which tests flash attention 4 functionality")

docs/source/en/exporters.md

@@ -718,6 +718,78 @@ for (int64_t position = prompt_len; position < max_cache_len; ++position) {
 
 </details>
 
+## Quantization
+
+Every export config accepts a `quantizer`. Set it to any PT2E
+[`Quantizer`](https://docs.pytorch.org/ao/main/pt2e_quantization/index.html) and the exporter runs post-training quantization on
+the traced graph (`prepare_pt2e` → calibrate → `convert_pt2e`) before the program is returned or
+lowered. Quantization happens on the graph rather than the modeling code, so a single `quantizer` works
+across every backend and architecture without per-model handling.
+
+Quantize through [`~HfExporter.export_for_generation`], which exports the decomposed generation components. Their attention mask is a precomputed graph input, which keeps PT2E away from the in-graph mask construction that trips its `make_fx` retrace on a full model forward.
+
+```python
+from transformers import LlamaForCausalLM
+from transformers.exporters import DynamoExporter, DynamoConfig
+from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import (
+    X86InductorQuantizer,
+    get_default_x86_inductor_quantization_config,
+)
+
+model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B").eval()
+inputs = ...  # forward kwargs
+
+quantizer = X86InductorQuantizer().set_global(get_default_x86_inductor_quantization_config())
+config = DynamoConfig(dynamic=True, quantizer=quantizer, calibration_dataset=[inputs])
+exported = DynamoExporter().export(model, inputs, config)  # quantize/dequantize ops folded into the graph
+```
+
+### Choosing a quantizer
+
+Each target runtime expects its own quantizer. Whichever you pass, the quantized graph is portable from there. It runs on inductor as int8, translates to ONNX `QuantizeLinear`/`DequantizeLinear` (per-channel included), or lowers to an ExecuTorch `.pte`.
+
+| Where you'll run | Quantizer to pass | Import from |
+| --- | --- | --- |
+| PyTorch inductor, or ONNX Runtime (QDQ) | `X86InductorQuantizer` | `torchao.quantization.pt2e.quantizer.x86_inductor_quantizer` |
+| ExecuTorch XNNPACK backend | `XNNPACKQuantizer` | `executorch.backends.xnnpack.quantizer.xnnpack_quantizer` |
+| ExecuTorch QNN backend (Qualcomm HTP) | `QnnQuantizer` | `executorch.backends.qualcomm.quantizer.quantizer` |
+
+### Calibration
+
+`calibration_dataset` is a list of forward-kwarg dicts run through the prepared graph to gather
+observer statistics. Omit it and the exporter falls back to a single pass over the export's own sample
+inputs, with a warning (one sample can skew the observed ranges).
+
+For generative models, set `calibration_dataset` on the config you pass to [`~HfExporter.export_for_generation`] and give it generate-style kwargs. Each sample runs through a short `generate`, and every component (`prefill`, `decode`, the encoders) is calibrated on the activations captured for it.
+
+### A different recipe per component
+
+Pass a `{component: config}` dict (instead of a single config) to
+[`~HfExporter.export_for_generation`], and each component is quantized on its own terms. The common
+multimodal case is static int8 on the vision tower, dynamic int8 on the language decoder, and a
+full-precision `lm_head`:
+
+```python
+def x86(dynamic):  # same quantizer family, static (per-channel) vs dynamic int8
+    return X86InductorQuantizer().set_global(get_default_x86_inductor_quantization_config(is_dynamic=dynamic))
+
+config = {
+    "image_encoder": DynamoConfig(dynamic=True, quantizer=x86(dynamic=False)),         # static int8
+    "multi_modal_projector": DynamoConfig(dynamic=True, quantizer=x86(dynamic=False)),
+    "language_model": DynamoConfig(dynamic=True, quantizer=x86(dynamic=True)),          # dynamic int8
+    "decode": DynamoConfig(dynamic=True, quantizer=x86(dynamic=True)),
+    "lm_head": DynamoConfig(dynamic=True),                                              # no quantizer → fp32
+}
+components = DynamoExporter().export_for_generation(model, inputs, config, multi_token_decode=True)
+```
+
+The dict must name every component [`~exporters.utils.decompose_for_generation`] produces; a component
+whose config sets no `quantizer` is left in full precision.
+
+> [!NOTE]
+> Quantization always runs on these decomposed components, whose attention mask is a precomputed input.
+> That keeps PT2E away from in-graph mask construction, which otherwise trips its `make_fx` retrace.
+
 ## Limitations and workarounds
 
 `torch.export`, `torch.onnx.export`, and ExecuTorch each have rough edges around specific

src/transformers/exporters/base.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 from abc import ABC, abstractmethod
 from collections.abc import MutableMapping
 from typing import TYPE_CHECKING
@@ -25,7 +26,7 @@
 from ..utils import logging
 from ..utils.import_utils import _is_package_available, is_torch_available
 from .configs import ExportConfigMixin
-from .utils import decompose_for_generation
+from .utils import capture_calibration_inputs, decompose_for_generation
 
 
 logger = logging.get_logger(__name__)
@@ -168,6 +169,14 @@ def export_for_generation(
                 classic single-token step (see [`~exporters.utils.decompose_for_generation`]). Only
                 stays dynamic under a dynamic-shape export (`config.dynamic=True`).
 
+        Quantization calibration: when a single `config` is passed (not a per-component dict) and it
+        carries a `quantizer`, its `calibration_dataset` is read as **generate** kwarg dicts (same level
+        as `sample_inputs` here) and fanned out — each sample is run through the decomposition to produce
+        a per-component calibration set that replaces each component's `config.calibration_dataset`
+        (per-graph forward kwargs). Leave it `None` to fall back to a single pass on each component's own
+        sample inputs (see [`DynamoConfig.calibration_dataset`]). A per-component `config` dict is left
+        untouched — set each component's `calibration_dataset` to its own forward kwargs directly.
+
         Returns:
             `dict[str, Any]`: `{component_name: backend_specific_artifact}` — same keys as
             [`~exporters.utils.decompose_for_generation`]. Values are whatever
@@ -189,13 +198,28 @@ def export_for_generation(
                     f"Expected one entry per component: {sorted(components)}."
                 )
             configs = config
+            calibration = {}
         else:
             configs = dict.fromkeys(components, config)
+            # a single config's `calibration_dataset` is generate-level here: fan it out into a
+            # per-component (forward-level) calibration set via the decomposition capture
+            calibration = {}
+            if getattr(config, "calibration_dataset", None):
+                calibration = capture_calibration_inputs(
+                    model,
+                    config.calibration_dataset,
+                    generation_config=generation_config,
+                    multi_token_decode=multi_token_decode,
+                )
 
         exported: dict[str, object] = {}
         for name, (submodel, subinputs) in components.items():
+            component_config = configs[name]
+            component_calibration = calibration.get(name)
+            if component_calibration is not None:
+                component_config = dataclasses.replace(component_config, calibration_dataset=component_calibration)
             try:
-                exported[name] = self.export(submodel, subinputs, config=configs[name])
+                exported[name] = self.export(submodel, subinputs, config=component_config)
             except Exception as e:
                 raise RuntimeError(
                     f"{type(self).__name__}.export failed on component '{name}' "

src/transformers/exporters/configs.py

@@ -96,6 +96,18 @@ class DynamoConfig(ExportConfigMixin):
             fine-grained ``Dim(min=, max=)`` bounds. Not needed with ``dynamic=True`` / ``Dim.AUTO``,
             where ``torch.export`` infers shape relations instead of verifying them against the
             user-stated bounds.
+        quantizer (`Quantizer`, *optional*):
+            Post-training quantization recipe — a PT2E `Quantizer` (e.g. `XNNPACKQuantizer(...)`,
+            `X86InductorQuantizer(...)`, a vendor `QnnQuantizer`, …). When set, the exported graph is
+            quantized (`prepare_pt2e` → calibrate → `convert_pt2e`) before it is returned/lowered.
+            Backend-agnostic: the resulting quantized `ExportedProgram` runs on inductor (int8), lowers
+            to ExecuTorch, or translates to ONNX QDQ. `None` (default) exports in full precision.
+        calibration_dataset (`list[dict]`, *optional*):
+            Forward-kwarg dicts run through the prepared graph to gather observer statistics for static
+            quantization. Ignored when `quantizer` is `None`. When `None`, calibration falls back to a
+            single pass on the export's own sample inputs (a warning is emitted — one sample can hurt
+            accuracy). For generative models, pass a generate-style dataset to `export_for_generation`'s
+            `calibration_dataset` instead — it fans out a per-component calibration set automatically.
     """
 
     export_format: ExportFormat = ExportFormat.DYNAMO
@@ -105,6 +117,9 @@ class DynamoConfig(ExportConfigMixin):
     dynamic_shapes: dict[str, Any] | None = None
     prefer_deferred_runtime_asserts_over_guards: bool = False
 
+    quantizer: Any = None
+    calibration_dataset: list[Any] | None = None
+
 
 @dataclass
 class OnnxConfig(DynamoConfig):
@@ -163,6 +178,10 @@ class ExecutorchConfig(DynamoConfig):
 
             - `"xnnpack"` — CPU inference via the XNNPACK library (default; runs anywhere).
             - `"cuda"` — GPU inference via the ExecuTorch CUDA backend.
+            - `"qnn"` — Qualcomm HTP (NPU) inference via the QNN backend. Requires the Qualcomm AI
+              Engine Direct SDK; pair with a `QnnQuantizer` via `quantizer` for int8/16 HTP (else fp16).
+        soc_model (`str`, *optional*, defaults to `"SM8650"`):
+            QNN only. Target Qualcomm chipset — a `QcomChipset` name (e.g. `"SM8650"`, `"SM8550"`).
         alloc_graph_input (`bool`, *optional*, defaults to `True`):
             Whether the memory-planning pass reserves arena memory for graph inputs. When `False`,
             the runtime uses the caller-provided input buffers directly instead of copying into the
@@ -180,6 +199,7 @@ class ExecutorchConfig(DynamoConfig):
     export_format: ExportFormat = ExportFormat.EXECUTORCH
 
     backend: str = "xnnpack"
+    soc_model: str = "SM8650"
     alloc_graph_input: bool = True
     alloc_graph_output: bool = True
     alloc_mutable_buffers: bool = True

src/transformers/exporters/exporter_dynamo.py

@@ -125,8 +125,60 @@ def export(
                 prefer_deferred_runtime_asserts_over_guards=config.prefer_deferred_runtime_asserts_over_guards,
             )
 
+        if config.quantizer is not None:
+            exported_program = self._quantize(exported_program, config, sample_inputs, dynamic_shapes)
+
         return exported_program
 
+    def _quantize(
+        self,
+        exported_program: ExportedProgram,
+        config: DynamoConfig,
+        sample_inputs: MutableMapping[str, Any],
+        dynamic_shapes: Any,
+    ) -> ExportedProgram:
+        """Post-training quantize the exported graph with PT2E, then re-export it.
+
+        Backend-agnostic: the recipe is the standard PT2E flow, identical for every model, and the only
+        backend-specific input is `config.quantizer` (`XNNPACKQuantizer`, `X86InductorQuantizer`, a vendor
+        `QnnQuantizer`, …). `prepare_pt2e` inserts observers, `config.calibration_dataset` (forward-kwarg
+        dicts) drives their statistics, and `convert_pt2e` folds them into `quantize`/`dequantize` ops.
+        Those two work on a `GraphModule`, so the converted graph is re-exported — with the same inputs and
+        dynamic-shape spec — back into an `ExportedProgram` that any downstream backend consumes (inductor
+        int8, ExecuTorch lowering, ONNX QDQ).
+        """
+        from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+
+        prepared = prepare_pt2e(exported_program.module(), config.quantizer)
+
+        # default to a single calibration pass on the export's own sample inputs
+        calibration_dataset = config.calibration_dataset
+        if not calibration_dataset:
+            logger.warning_once(
+                "Quantizing with no `calibration_dataset`; calibrating on the single sample input. Observer "
+                "statistics from one sample can hurt accuracy — set a representative `config.calibration_dataset` "
+                "(for generative models, `export_for_generation` fans a generate-level one out per component)."
+            )
+            calibration_dataset = [sample_inputs]
+
+        forward_keys = set(sample_inputs)  # traced forward signature (output flags already stripped)
+        for sample in calibration_dataset:
+            # keep only the traced forward kwargs (drop any captured generation flags), and deep-copy:
+            # a calibration forward writes the cache in place, so this stops it from mutating the
+            # caller's tensors — or `sample_inputs`, which the re-export below reuses
+            inputs = {name: copy.deepcopy(value) for name, value in sample.items() if name in forward_keys}
+            prepared(**inputs)
+
+        converted = convert_pt2e(prepared)
+        return torch.export.export(
+            converted,
+            args=(),
+            kwargs=copy.deepcopy(dict(sample_inputs)),
+            strict=config.strict,
+            dynamic_shapes=dynamic_shapes,
+            prefer_deferred_runtime_asserts_over_guards=config.prefer_deferred_runtime_asserts_over_guards,
+        )
+
 
 # ── Stage 1: Model signature patch ──────────────────────────────────────────
 # Replaces `model.forward` with a flat explicit signature derived from the