diff --git a/conftest.py b/conftest.py index 773f33d59b72..02bbee10e0d5 100644 --- a/conftest.py +++ b/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") diff --git a/docs/source/en/exporters.md b/docs/source/en/exporters.md index a5897c971357..d53d3c500413 100644 --- a/docs/source/en/exporters.md +++ b/docs/source/en/exporters.md @@ -718,6 +718,78 @@ for (int64_t position = prompt_len; position < max_cache_len; ++position) { +## 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 diff --git a/src/transformers/exporters/base.py b/src/transformers/exporters/base.py index 5fcd1979f759..f78cfa2af00c 100644 --- a/src/transformers/exporters/base.py +++ b/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}' " diff --git a/src/transformers/exporters/configs.py b/src/transformers/exporters/configs.py index 1e58f4da16e9..ebbfa935f2f4 100644 --- a/src/transformers/exporters/configs.py +++ b/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 diff --git a/src/transformers/exporters/exporter_dynamo.py b/src/transformers/exporters/exporter_dynamo.py index 998f8f5c050a..59f4e5ad016a 100644 --- a/src/transformers/exporters/exporter_dynamo.py +++ b/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 diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index 690036c3ade5..4a28049931f3 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -17,8 +17,8 @@ Extends `DynamoExporter` to produce an `ExecutorchProgramManager` for mobile and edge deployment. The export pipeline runs: -1. **Backend preparation** (`_BACKEND_PREPARE`): `prepare_for_xnnpack` / `prepare_for_cuda` - move the model to the target device/dtype and build the partitioner list. +1. **Backend preparation** (`_BACKEND_PREPARE`): `prepare_for_xnnpack` / `prepare_for_cuda` / + `prepare_for_qnn` move the model to the target device/dtype and build the partitioner list. 2. **Torch patches** (`_PATCHES["executorch"]` via `apply_patches("executorch")`, plus the backend-specific `_PATCHES[f"executorch.{backend}"]`): reversibly swap `torch` ops the ExecuTorch backends can't accept (`split_copy`, `avg_pool2d`, …) with decomposed equivalents. @@ -139,21 +139,26 @@ def export( if prepare_for_backend is None: raise ValueError(f"Unsupported backend {config.backend} for ExecuTorch export") - model, sample_inputs, partitioner = prepare_for_backend(model, sample_inputs) + model, sample_inputs, partitioner = prepare_for_backend(model, sample_inputs, config) with apply_patches("executorch"), apply_patches(f"executorch.{config.backend}"): exported_program: ExportedProgram = super().export(model, sample_inputs, config=config) apply_fx_program_fixes("executorch", exported_program) apply_fx_node_fixes("executorch", exported_program.graph_module) - edge_program_manager: EdgeProgramManager = to_edge_transform_and_lower( - exported_program, partitioner=partitioner, compile_config=_get_edge_compile_config() - ) + edge_program_manager: EdgeProgramManager = _lower_to_edge(exported_program, partitioner, config.backend) executorch_programs_manager: ExecutorchProgramManager = edge_program_manager.to_executorch( config=_get_backend_config(config) ) return executorch_programs_manager + def _quantize(self, exported_program, config, sample_inputs, dynamic_shapes): + # QnnQuantizer's annotation passes assume every node carries a schema and crash on the + # `wrap_with_set_grad_enabled` HOP `torch.export` emits; lower to inference IR to inline it away. + if config.backend == "qnn": + exported_program = exported_program.run_decompositions({}) + return super()._quantize(exported_program, config, sample_inputs, dynamic_shapes) + def _get_edge_compile_config() -> EdgeCompileConfig: """Build the ``EdgeCompileConfig`` used for ``to_edge_transform_and_lower``. @@ -183,6 +188,29 @@ def _get_edge_compile_config() -> EdgeCompileConfig: ) +def _lower_to_edge(exported_program, partitioner, backend): + """Lower an (optionally quantized) program to the edge dialect and delegate it to ``partitioner``. + + QNN can't reuse the generic path: its HTP compiler needs its own edge-transform passes — notably + `FoldQDQ`, which folds the PT2E `quantize`/`dequantize` ops into HTP-native quant encodings. Without + them the raw dequantize ops reach the backend and graph finalization fails (`op validation ... 3110`). + Mirror `to_edge_transform_and_lower_to_qnn` for QNN; every other backend uses the plain path. + """ + if backend == "qnn": + from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager + from executorch.backends.qualcomm.utils.utils import qnn_edge_config + + aten_program = QnnPassManager().transform_for_export_pipeline(exported_program) + transform_passes = QnnPassManager().get_to_edge_transform_passes(exported_program) + return to_edge_transform_and_lower( + aten_program, transform_passes=transform_passes, partitioner=partitioner, compile_config=qnn_edge_config() + ) + + return to_edge_transform_and_lower( + exported_program, partitioner=partitioner, compile_config=_get_edge_compile_config() + ) + + def _get_backend_config(config): """Build the ``ExecutorchBackendConfig`` for ``to_executorch``, or ``None`` for defaults. @@ -222,7 +250,7 @@ def _make_contiguous(sample_inputs: dict[str, Any]) -> dict[str, Any]: return torch.utils._pytree.tree_map_only(torch.Tensor, lambda t: t.contiguous(), sample_inputs) -def prepare_for_xnnpack(model: PreTrainedModel, sample_inputs: dict[str, Any]): +def prepare_for_xnnpack(model: PreTrainedModel, sample_inputs: dict[str, Any], config: ExecutorchConfig): """CPU inference via XNNPACK. Moves the model to CPU: XNNPACK's partitioner/serializer and the edge-lowering passes all @@ -240,7 +268,47 @@ def prepare_for_xnnpack(model: PreTrainedModel, sample_inputs: dict[str, Any]): return model, _make_contiguous(sample_inputs), partitioner -def prepare_for_cuda(model: PreTrainedModel, sample_inputs: dict[str, Any]): +def prepare_for_qnn(model: PreTrainedModel, sample_inputs: dict[str, Any], config: ExecutorchConfig): + """On-device inference via the Qualcomm QNN (HTP/NPU) backend. + + Skipped in this repo's CI — the Qualcomm AI Engine Direct SDK isn't installed there (the + `executorch.backends.qualcomm` package won't even import without it). The SDK does install on a + plain x86 host, though, and the export path is verified against it (the `.pte` is built via the + SDK's HTP emulator, no Qualcomm device needed); see `test_qnn_export`. + + The model is traced on CPU (like XNNPACK), then delegated to the HTP via `QnnPartitioner`. Pass a + `QnnQuantizer` through `config.quantizer` for HTP int8/16 (the generic PT2E step quantizes the graph + before this lowering); with no quantizer the HTP compiler spec runs fp16 (`use_fp16=True`). The + `executorch.backends.qualcomm` imports are local because that package raises on import without the + SDK — a module-level import would break the ExecuTorch exporter for the XNNPACK/CUDA backends too. + """ + from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner + from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset + from executorch.backends.qualcomm.utils.utils import ( + generate_htp_compiler_spec, + generate_qnn_executorch_compiler_spec, + ) + + model.requires_grad_(False) + model = model.to(device="cpu") + + # The HTP has no rank-0 tensors, so QNN promotes scalars to rank-1 — but leaves mutable-buffer + # placeholders rank-0, so `StaticLayer.cumulative_length` (a scalar counter) fails its input-mutation + # writeback at `to_executorch`. Make it 1D up front (broadcast-equivalent) so the writeback matches. + for layer in getattr(sample_inputs.get("past_key_values"), "layers", ()): + cumulative_length = getattr(layer, "cumulative_length", None) + if isinstance(cumulative_length, torch.Tensor) and cumulative_length.ndim == 0: + layer.cumulative_length = cumulative_length.reshape(1) + + backend_options = generate_htp_compiler_spec(use_fp16=config.quantizer is None) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=getattr(QcomChipset, config.soc_model), backend_options=backend_options + ) + partitioner = [QnnPartitioner(compiler_specs)] + return model, _make_contiguous(sample_inputs), partitioner + + +def prepare_for_cuda(model: PreTrainedModel, sample_inputs: dict[str, Any], config: ExecutorchConfig): """GPU inference via the ExecuTorch CUDA backend, decoupled from the model's device. The backend requires bfloat16 (upcast here) and a visible GPU — it delegates ops to Triton @@ -262,6 +330,7 @@ def prepare_for_cuda(model: PreTrainedModel, sample_inputs: dict[str, Any]): _BACKEND_PREPARE = { "xnnpack": prepare_for_xnnpack, "cuda": prepare_for_cuda, + "qnn": prepare_for_qnn, } @@ -1111,6 +1180,33 @@ def _patch_squeeze_node_visitors(original): return new +@register_patch("executorch.qnn", "executorch.backends.qualcomm._passes.replace_inf_values.ReplaceInfValues.call") +def _patch_replace_inf_values(original): + """QNN's ``ReplaceInfValues`` pass ``setattr``s every float buffer back onto the graph module by name, + which ``register_buffer`` rejects for nested (dotted) names such as ``model.rotary_emb.inv_freq``. The + pass already clamps ``inf`` in place, so clamp the dotted float buffers ourselves, hide them from the + pass's ``setattr``, and restore them afterwards. + """ + + def call(self, graph_module): + hidden = [] + for name, buffer in list(graph_module.named_buffers()): + if "." in name and buffer.is_floating_point(): + buffer[buffer == float("inf")] = 255 + buffer[buffer == float("-inf")] = -255 + parent_path, _, attr = name.rpartition(".") + parent = graph_module.get_submodule(parent_path) + hidden.append((parent, attr, buffer)) + delattr(parent, attr) + try: + return original(self, graph_module) + finally: + for parent, attr, buffer in hidden: + parent.register_buffer(attr, buffer, persistent=False) + + return call + + # ── Stage 4: FX program fixes ───────────────────────────────────────────────── # `@register_fx_program_fix("executorch")` on `(exported_program) -> None` callables # applied in place between ``torch.export.export`` and ``to_edge_transform_and_lower``. diff --git a/src/transformers/exporters/exporter_onnx.py b/src/transformers/exporters/exporter_onnx.py index 504ac4085d5a..8a191fa6247e 100644 --- a/src/transformers/exporters/exporter_onnx.py +++ b/src/transformers/exporters/exporter_onnx.py @@ -942,6 +942,29 @@ def _aten_masked_fill(self, mask, value): return op.Where(mask, value_cast, self) +def _quantized_decomposed_dequantize_per_channel( + input, + scales, + zero_points, + axis: int, + quant_min: int, + quant_max: int, + dtype: int, + out_dtype: int = -1, +): + """ONNX translation for `quantized_decomposed.dequantize_per_channel`. + + onnxscript's torchlib registers only the per-*tensor* `quantized_decomposed` ops, so a PT2E graph + with per-channel weights (e.g. `X86InductorQuantizer`) has no ONNX function for the per-channel + dequant and fails to translate. ONNX `DequantizeLinear` handles it natively via `axis` (1-D + scale/zero-point along the channel axis); the zero-point must share the input's integer dtype. + """ + if zero_points is not None: + zero_points = op.CastLike(zero_points, input) + return op.DequantizeLinear(input, scales, zero_points, axis=axis) + return op.DequantizeLinear(input, scales, axis=axis) + + _ONNX_TRANSLATION_TABLE: dict[Any, Any] = {} if is_onnxscript_available(): _ONNX_TRANSLATION_TABLE.update( @@ -957,6 +980,19 @@ def _aten_masked_fill(self, mask, value): } ) + # The `quantized_decomposed` ops (torch's built-in decomposed-quant lib, not `torchao`) are only + # registered once that lib is imported; import it so the per-channel dequant overload resolves as a + # table key (onnxscript covers only the per-tensor variants). The entry is inert unless a + # quantized graph actually contains the op, so it costs nothing when quantization is unused. + try: + import torch.ao.quantization.fx._decomposed # noqa: F401 + + _ONNX_TRANSLATION_TABLE[torch.ops.quantized_decomposed.dequantize_per_channel.default] = ( + _quantized_decomposed_dequantize_per_channel + ) + except (ImportError, AttributeError): + pass + # ── Stage 5: ONNX IR fixes ──────────────────────────────────────────────────── # Post-export in-place fixes to the `ONNXProgram` IR for ORT compatibility. Each diff --git a/src/transformers/exporters/utils.py b/src/transformers/exporters/utils.py index 17bddd1e0adf..af1ef37b8f1e 100644 --- a/src/transformers/exporters/utils.py +++ b/src/transformers/exporters/utils.py @@ -68,13 +68,14 @@ # ── Patch and fix registries ──────────────────────────────────────────────── -# Single contract across exporters: `_PATCHES[backend]` lists `(obj, attribute, factory)` triples +# Single contract across exporters: `_PATCHES[backend]` lists `(obj_path, attribute, factory)` triples # to install reversibly, and `_FX_NODE_FIXES[backend]` lists `(gm, node) -> bool` fixers to # apply in place. Each exporter populates its slot at module load (via `@register_patch` / -# `@register_fx_node_fix` decorators, or direct list-append for cases that can't be expressed -# as dotted paths). The export pipeline drives them via the backend-keyed helpers below. +# `@register_fx_node_fix` decorators). `obj_path` is a dotted string resolved at apply time, not +# import time — resolving `executorch.backends.qualcomm`, say, would run the QNN SDK installer on +# `import transformers`. The export pipeline drives them via the backend-keyed helpers below. -_PATCHES: dict[str, list[tuple[Any, str, callable]]] = {} +_PATCHES: dict[str, list[tuple[str, str, callable]]] = {} _FX_NODE_FIXES: dict[str, list[callable]] = {} _FX_PROGRAM_FIXES: dict[str, list[callable]] = {} @@ -105,8 +106,14 @@ def patch_attributes(patches: list[tuple[Any, str, callable]]): @contextlib.contextmanager def apply_patches(backend: str): - """Install `_PATCHES[backend]` for the duration of the block.""" - with patch_attributes(_PATCHES.get(backend, [])): + """Install `_PATCHES[backend]` for the duration of the block, resolving each entry's dotted object + path now (importing submodules as needed). This runs only when `backend` is actually exporting, so + its modules are importable — an unresolvable path is a bug and raises rather than being skipped.""" + patches = [ + (_resolve_dotted_path(obj_path), attribute, factory) + for obj_path, attribute, factory in _PATCHES.get(backend, []) + ] + with patch_attributes(patches): yield @@ -145,10 +152,10 @@ def register_patch(backend: str, *paths: str): Each `path` is a dotted Python path like `"torch.where"`, `"torch.Tensor.unsqueeze"`, or `"transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier"`. - The rightmost segment is the attribute to swap; the rest is the object that owns it. - Paths are resolved at decoration time — submodules are imported as needed, falling - back to `getattr` for class attributes. A path that fails to resolve (e.g. the backend - isn't installed) is silently skipped so the module still imports. + The rightmost segment is the attribute to swap; the rest is the object that owns it. The + owner path is stored as-is and resolved by `apply_patches` at apply time — never at import, + so a path into an uninstalled or install-on-import backend (e.g. `executorch.backends.qualcomm`, + whose import runs the QNN SDK auto-installer) costs nothing until that backend actually runs. Passing multiple paths registers the SAME factory against each — useful for swapping the same method or torch op across several call sites (e.g. ``torch.unsqueeze`` + @@ -158,32 +165,26 @@ def register_patch(backend: str, *paths: str): def decorator(fn): for path in paths: obj_path, _, attribute = path.rpartition(".") - obj = _resolve_dotted_path(obj_path) - if obj is None: - continue - _PATCHES.setdefault(backend, []).append((obj, attribute, fn)) + _PATCHES.setdefault(backend, []).append((obj_path, attribute, fn)) return fn return decorator def _resolve_dotted_path(path: str): - """Resolve a dotted Python path to the actual object — importing submodules where - possible, falling back to `getattr` for class attributes (e.g. `torch.Tensor`). - Returns `None` if the path can't be resolved (e.g. the backend isn't installed).""" + """Resolve a dotted Python path to the actual object — importing submodules where possible, + falling back to `getattr` for class attributes (e.g. `torch.Tensor`). Raises `ImportError` / + `AttributeError` if the path can't be resolved — callers resolve only paths they expect to exist.""" import importlib parts = path.split(".") - try: - obj = importlib.import_module(parts[0]) - for part in parts[1:]: - try: - obj = importlib.import_module(f"{obj.__name__}.{part}") - except (ImportError, AttributeError): - obj = getattr(obj, part) - return obj - except (ImportError, AttributeError): - return None + obj = importlib.import_module(parts[0]) + for part in parts[1:]: + try: + obj = importlib.import_module(f"{obj.__name__}.{part}") + except (ImportError, AttributeError): + obj = getattr(obj, part) + return obj def apply_fx_node_fixes(backend: str, graph_module) -> None: @@ -915,3 +916,28 @@ def decompose_for_generation( components = decompose_multimodal(prefill_model, prefill_inputs) components["decode"] = stages["decode"] return components + + +def capture_calibration_inputs( + model: PreTrainedModel, + calibration_dataset: list[dict[str, Any]], + generation_config: Any = None, + multi_token_decode: bool = False, +) -> dict[str, list[dict]]: + """Capture per-component forward inputs for post-training quantization calibration. + + Reuses `decompose_for_generation`'s capture: each generate-style sample in `calibration_dataset` is + run through the decomposition, and every component's forward kwargs are collected. Returns + `{component_name: [forward_inputs, ...]}` — one list per component (same keys + `decompose_for_generation` produces), which the exporter feeds to that component's quantization + calibration. This is the input/output-capture "trick" applied to calibration: the user provides one + generate-level dataset, and each single-graph component gets its own inferred calibration set. + """ + calibration: dict[str, list[dict]] = {} + for sample in calibration_dataset: + components = decompose_for_generation( + model, sample, generation_config=generation_config, multi_token_decode=multi_token_decode + ) + for name, (_submodel, forward_inputs) in components.items(): + calibration.setdefault(name, []).append(forward_inputs) + return calibration diff --git a/tests/exporters/export_utils.py b/tests/exporters/export_utils.py new file mode 100644 index 000000000000..fcce9d7b6075 --- /dev/null +++ b/tests/exporters/export_utils.py @@ -0,0 +1,136 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared helpers to *run* exported programs (ONNX in ORT, ExecuTorch in its runtime). + +Used by both the per-model export tests (`test_export.py`, which check outputs against eager) and the +quantization tests (`test_quantization.py`, which only check the quantized graph executes). Kept out of +either test module so neither imports the other; the filename doesn't match pytest's `test_*` pattern, +so it isn't collected. +""" + +import re + +from transformers import set_seed +from transformers.exporters.utils import get_leaf_tensors +from transformers.utils.import_utils import is_torch_available + + +if is_torch_available(): + import torch + + +def run_onnx_program(onnx_program, inputs) -> dict: + """Run an ONNX program and return outputs as a `{name: tensor}` dict.""" + set_seed(1234) + onnx_inputs = get_leaf_tensors(inputs) + onnx_outputs = onnx_program(**onnx_inputs) + onnx_names = (re.sub(r"^output\.", "", node.name) for node in onnx_program.model_proto.graph.output) + return dict(zip(onnx_names, onnx_outputs)) + + +def run_executorch_program(program_manager, inputs): + """Load and run an ExecuTorch program, returning its outputs — or ``None`` to skip this component. + + ``None`` means "move on to the next component" and is returned when either: + - the export is valid but ExecuTorch's own runtime can't service it — a missing portable kernel + (``0x14``), an oversized arena (``0x21`` / ``bad_alloc``), or a portable-kernel / XNNPACK-delegate + failure at execute (``0x12`` / ``0x1``): a runtime limitation, not a transformers export defect; or + - the inputs couldn't be reconstructed for this program (a derived symint slot with no eager leaf). + + Otherwise the model's declared outputs are returned for the caller to check against eager. + ``torch.export`` also appends mutated inputs (in-place-modified ``pixel_values``, recurrent state, + …) to the program outputs; those are dropped here — keeping only ``USER_OUTPUT`` slots — so the + result matches eager's returned leaves. + + Inputs are bound *positionally* against the program's declared slots (``num_inputs`` / + ``input_tensor_meta``), filled in order from the eager pytree leaves — tensor leaves for tensor + slots, scalars for the rest. + """ + from executorch.runtime import Runtime, Verification + + set_seed(1234) + leaves = torch.utils._pytree.tree_leaves(inputs) + # The runtime rejects non-contiguous inputs, so materialise tensor leaves. `int` covers `bool`. + tensors = [t.contiguous() for t in leaves if isinstance(t, torch.Tensor)] + scalars = (t for t in leaves if isinstance(t, (int, float))) + + # Load — surfaces ExecuTorch resource limits (missing portable kernel / oversized arena). + try: + program = Runtime.get().load_program(program_manager.buffer, verification=Verification.Minimal) + method = program.load_method("forward") + except (RuntimeError, MemoryError) as e: + if is_executorch_runtime_limit(e): + return None + raise + + # Each slot declares its shape; match it to an eager tensor leaf of that shape so the right tensor + # lands in the right slot (count alone isn't enough — a wrong-shape tensor crashes conv/copy + # kernels at execute). Under dynamic shapes the declared shape is an upper bound and won't match a + # leaf, so fall back to the next unused leaf (leaf order tracks the program's input order). If a + # slot can't be filled — a derived symint, or no leaf of the right shape — reconstruction isn't + # possible; return None and rely on the load check rather than run with bogus inputs. + args = [] + for i in range(method.metadata.num_inputs()): + try: + shape = tuple(method.metadata.input_tensor_meta(i).sizes()) + except Exception: # non-tensor slot + args.append(next(scalars, None)) + else: + match = next((t for t in tensors if tuple(t.shape) == shape), tensors[0] if tensors else None) + if match is not None: + tensors.remove(match) + args.append(match) + if args[-1] is None: + return None + + try: + outputs = method.execute(args) + except (RuntimeError, MemoryError) as e: + if is_executorch_runtime_limit(e): + return None + raise + + # Drop `torch.export`'s appended mutated-input outputs, keeping only the model's `USER_OUTPUT`s + # (in program-output order). Then keep tensors only, mirroring eager's `get_leaf_tensors`, so the + # returned outputs line up with eager's returned leaves for the caller's count check. + exported_program = program_manager.exported_program + exported_program = exported_program() if callable(exported_program) else exported_program + output_kinds = [spec.kind.name for spec in exported_program.graph_signature.output_specs] + if len(output_kinds) == len(outputs): + outputs = [out for out, kind in zip(outputs, output_kinds) if kind == "USER_OUTPUT"] + return [out for out in outputs if isinstance(out, torch.Tensor)] + + +# ExecuTorch runtime error codes that mean "the export is valid (it produced a loadable program) but +# ExecuTorch's own portable runtime / XNNPACK backend can't service it" — a runtime limitation, not a +# transformers export defect (which surfaces earlier as a `torch.export` error or later as an output +# mismatch). Load: 0x14 missing portable kernel, 0x21 arena can't be allocated, 0x1 XNNPACK partition +# won't compile (`xnn_status_unsupported_parameter`). Execute: 0x12 portable-kernel InvalidArgument +# (constant_pad_nd/convolution/upsample_aa out-tensor sizing), 0x1 XNNPACK delegate failure, 0x10 +# XNNPACK delegate can't resize a static tensor to the runtime shape. The execute-phase codes surface +# from either `execute()` or `set_inputs()` (binding the runtime inputs is part of `Method.execute`). +_ET_LOAD_LIMIT_CODES = {"0x1", "0x14", "0x21"} +_ET_EXECUTE_LIMIT_CODES = {"0x1", "0x10", "0x12"} + + +def is_executorch_runtime_limit(exc): + """True if ``exc`` is a known ExecuTorch runtime limitation (missing kernel / arena / kernel bug).""" + msg = str(exc) + if isinstance(exc, MemoryError) or "bad_alloc" in msg: + return True + load = re.search(r"Failed to load method forward, error: 0x:?([0-9a-fA-F]+)", msg) + if load and f"0x{load.group(1)}" in _ET_LOAD_LIMIT_CODES: + return True + execute = re.search(r"(?:execute\(\)|set_inputs\(\) for method '\w+') failed with error 0x([0-9a-fA-F]+)", msg) + return bool(execute and f"0x{execute.group(1)}" in _ET_EXECUTE_LIMIT_CODES) diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 10c10d3de40b..3a6ac16b872e 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -22,6 +22,8 @@ import torch from parameterized import parameterized +from tests.exporters.export_utils import run_executorch_program as _run_executorch_program +from tests.exporters.export_utils import run_onnx_program as _run_onnx_program from transformers import GenerationConfig, set_seed from transformers.exporters.exporter_dynamo import DynamoConfig, DynamoExporter from transformers.exporters.exporter_executorch import ExecutorchConfig, ExecutorchExporter @@ -416,112 +418,6 @@ def _clean_inputs_for_export(inputs_dict, config): return inputs_dict -def _run_onnx_program(onnx_program, inputs) -> dict: - """Run an ONNX program and return outputs as a `{name: tensor}` dict.""" - set_seed(1234) - onnx_inputs = get_leaf_tensors(inputs) - onnx_outputs = onnx_program(**onnx_inputs) - onnx_names = (re.sub(r"^output\.", "", node.name) for node in onnx_program.model_proto.graph.output) - return dict(zip(onnx_names, onnx_outputs)) - - -def _run_executorch_program(program_manager, inputs): - """Load and run an ExecuTorch program, returning its outputs — or ``None`` to skip this component. - - ``None`` means "move on to the next component" and is returned when either: - - the export is valid but ExecuTorch's own runtime can't service it — a missing portable kernel - (``0x14``), an oversized arena (``0x21`` / ``bad_alloc``), or a portable-kernel / XNNPACK-delegate - failure at execute (``0x12`` / ``0x1``): a runtime limitation, not a transformers export defect; or - - the inputs couldn't be reconstructed for this program (a derived symint slot with no eager leaf). - - Otherwise the model's declared outputs are returned for the caller to check against eager. - ``torch.export`` also appends mutated inputs (in-place-modified ``pixel_values``, recurrent state, - …) to the program outputs; those are dropped here — keeping only ``USER_OUTPUT`` slots — so the - result matches eager's returned leaves. - - Inputs are bound *positionally* against the program's declared slots (``num_inputs`` / - ``input_tensor_meta``), filled in order from the eager pytree leaves — tensor leaves for tensor - slots, scalars for the rest. - """ - from executorch.runtime import Runtime, Verification - - set_seed(1234) - leaves = torch.utils._pytree.tree_leaves(inputs) - # The runtime rejects non-contiguous inputs, so materialise tensor leaves. `int` covers `bool`. - tensors = [t.contiguous() for t in leaves if isinstance(t, torch.Tensor)] - scalars = (t for t in leaves if isinstance(t, (int, float))) - - # Load — surfaces ExecuTorch resource limits (missing portable kernel / oversized arena). - try: - program = Runtime.get().load_program(program_manager.buffer, verification=Verification.Minimal) - method = program.load_method("forward") - except (RuntimeError, MemoryError) as e: - if _is_executorch_runtime_limit(e): - return None - raise - - # Each slot declares its shape; match it to an eager tensor leaf of that shape so the right tensor - # lands in the right slot (count alone isn't enough — a wrong-shape tensor crashes conv/copy - # kernels at execute). Under dynamic shapes the declared shape is an upper bound and won't match a - # leaf, so fall back to the next unused leaf (leaf order tracks the program's input order). If a - # slot can't be filled — a derived symint, or no leaf of the right shape — reconstruction isn't - # possible; return None and rely on the load check rather than run with bogus inputs. - args = [] - for i in range(method.metadata.num_inputs()): - try: - shape = tuple(method.metadata.input_tensor_meta(i).sizes()) - except Exception: # non-tensor slot - args.append(next(scalars, None)) - else: - match = next((t for t in tensors if tuple(t.shape) == shape), tensors[0] if tensors else None) - if match is not None: - tensors.remove(match) - args.append(match) - if args[-1] is None: - return None - - try: - outputs = method.execute(args) - except (RuntimeError, MemoryError) as e: - if _is_executorch_runtime_limit(e): - return None - raise - - # Drop `torch.export`'s appended mutated-input outputs, keeping only the model's `USER_OUTPUT`s - # (in program-output order). Then keep tensors only, mirroring eager's `get_leaf_tensors`, so the - # returned outputs line up with eager's returned leaves for the caller's count check. - exported_program = program_manager.exported_program - exported_program = exported_program() if callable(exported_program) else exported_program - output_kinds = [spec.kind.name for spec in exported_program.graph_signature.output_specs] - if len(output_kinds) == len(outputs): - outputs = [out for out, kind in zip(outputs, output_kinds) if kind == "USER_OUTPUT"] - return [out for out in outputs if isinstance(out, torch.Tensor)] - - -# ExecuTorch runtime error codes that mean "the export is valid (it produced a loadable program) but -# ExecuTorch's own portable runtime / XNNPACK backend can't service it" — a runtime limitation, not a -# transformers export defect (which surfaces earlier as a `torch.export` error or later as an output -# mismatch). Load: 0x14 missing portable kernel, 0x21 arena can't be allocated, 0x1 XNNPACK partition -# won't compile (`xnn_status_unsupported_parameter`). Execute: 0x12 portable-kernel InvalidArgument -# (constant_pad_nd/convolution/upsample_aa out-tensor sizing), 0x1 XNNPACK delegate failure, 0x10 -# XNNPACK delegate can't resize a static tensor to the runtime shape. The execute-phase codes surface -# from either `execute()` or `set_inputs()` (binding the runtime inputs is part of `Method.execute`). -_ET_LOAD_LIMIT_CODES = {"0x1", "0x14", "0x21"} -_ET_EXECUTE_LIMIT_CODES = {"0x1", "0x10", "0x12"} - - -def _is_executorch_runtime_limit(exc): - """True if ``exc`` is a known ExecuTorch runtime limitation (missing kernel / arena / kernel bug).""" - msg = str(exc) - if isinstance(exc, MemoryError) or "bad_alloc" in msg: - return True - load = re.search(r"Failed to load method forward, error: 0x:?([0-9a-fA-F]+)", msg) - if load and f"0x{load.group(1)}" in _ET_LOAD_LIMIT_CODES: - return True - execute = re.search(r"(?:execute\(\)|set_inputs\(\) for method '\w+') failed with error 0x([0-9a-fA-F]+)", msg) - return bool(execute and f"0x{execute.group(1)}" in _ET_EXECUTE_LIMIT_CODES) - - def _onnx_optimize_enabled(model_class, dynamic: bool) -> bool: """Return whether onnxscript optimisation should run for this model under this shape mode. diff --git a/tests/exporters/test_quantization.py b/tests/exporters/test_quantization.py new file mode 100644 index 000000000000..edf87620cb9b --- /dev/null +++ b/tests/exporters/test_quantization.py @@ -0,0 +1,399 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Post-training quantization export tests. + +PT2E quantization is a backend-agnostic recipe living in the shared Dynamo layer: pass a `quantizer` +(any PT2E `Quantizer` — `X86InductorQuantizer`, `XNNPACKQuantizer`, a vendor `QnnQuantizer`, …) on the +export config and the graph is quantized (`prepare_pt2e` → calibrate → `convert_pt2e`) before it's +returned/lowered — no hardcoded schemes. The tests cover: + +- **`test_quantized_{dynamo,onnx,executorch}`** — one per exporter backend (so each carries the right CI + marker), each over the architecture families (dense / MoE / SSM) with the quantizer(s) natural to that + backend: dynamo/onnx use the torchao-native `x86` quantizer (graph-level QDQ, no executorch dep); the + executorch backend uses the per-tensor `xnnpack` and vendor `qnn` quantizers it can delegate (per-channel + x86 has no delegated out variant). One structural check per cell — the artifact exports and carries the + quant ops (dynamo `quantize`/`dequantize`, ONNX QDQ that loads in ORT, int8 `.pte`) — driven entirely by + `config.quantizer`, no per-case code. QNN HTP's own gaps on MoE routing / SSM conv1d, and SDK/dep gaps, + skip with a reason. +- **`test_vlm_per_component_quantization`** — a VLM quantized component-by-component, each with its OWN + recipe (vision encoder static int8, decoder dynamic int8, `lm_head` fp32) via a per-component config dict. +- **calibration** — the generate-level `calibration_dataset` captured into a separate set per component, + and the single-sample fallback (with a warning) when it's omitted. + +Quantization runs on the decomposed generation components (whose attention mask is a precomputed input), +avoiding the in-graph mask construction that trips PT2E on a full-model forward. +""" + +import copy +import sys +import tempfile +import unittest +from unittest.mock import patch + +import pytest +from parameterized import parameterized + +from tests.exporters.export_utils import run_onnx_program +from transformers import GenerationConfig, LlamaConfig, LlamaForCausalLM +from transformers.exporters.utils import capture_calibration_inputs, decompose_for_generation +from transformers.testing_utils import ( + require_executorch, + require_torch, + require_torchao, + run_command, + slow, +) +from transformers.utils import is_torch_available + + +if is_torch_available(): + import torch + + +MAX_CACHE_LEN = 16 + + +def _qnn_available() -> bool: + """The QNN backend needs the Qualcomm AI Engine Direct SDK. Probe in a subprocess: importing + `executorch.backends.qualcomm` runs an auto-installer that mutates `LD_LIBRARY_PATH`, which would + corrupt the pytest process for the other tests. Use a script file rather than `python -c`: on an + old glibc the installer re-execs Python under a staged loader and only the file path survives.""" + with tempfile.NamedTemporaryFile("w", suffix=".py") as probe: + probe.write("import executorch.backends.qualcomm\n") + probe.flush() + try: + run_command([sys.executable, probe.name]) + return True + except Exception: + return False + + +def _has_quantize_ops(exported) -> bool: + """The exported FX graph carries PT2E quantize/dequantize nodes.""" + return any(n.op == "call_function" and "quantize" in str(n.target) for n in exported.graph.nodes) + + +def _has_dynamic_quant_ops(exported) -> bool: + """Activations are quantized dynamically (runtime `choose_qparams`), not with static calibrated scales.""" + return any(n.op == "call_function" and "choose_qparams" in str(n.target) for n in exported.graph.nodes) + + +def _has_onnx_quantize_ops(program) -> bool: + """The exported ONNX graph carries QDQ (`QuantizeLinear`) nodes.""" + return any(node.op_type == "QuantizeLinear" for node in program.model_proto.graph.node) + + +@slow +@require_torch +@require_torchao +class QuantizationExportTest(unittest.TestCase): + def _tiny_model(self): + torch.manual_seed(0) + config = LlamaConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=64, + max_position_embeddings=128, + ) + return LlamaForCausalLM(config).eval() + + def _moe_model(self): + """Tiny MoE — exercises expert-routing / expert-linear quantization.""" + from transformers import MixtralConfig, MixtralForCausalLM + + torch.manual_seed(0) + config = MixtralConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=64, + max_position_embeddings=128, + num_local_experts=2, + num_experts_per_tok=2, + ) + return MixtralForCausalLM(config).eval() + + def _ssm_model(self): + """Tiny SSM — exercises conv1d / SSM-projection quantization (no attention).""" + from transformers import Mamba2Config, Mamba2ForCausalLM + + torch.manual_seed(0) + config = Mamba2Config( + hidden_size=32, + num_hidden_layers=2, + vocab_size=64, + num_heads=8, + head_dim=8, + state_size=8, + n_groups=1, + chunk_size=8, + conv_kernel=4, + expand=2, + ) + return Mamba2ForCausalLM(config).eval() + + def _vlm_model(self): + """Tiny VLM (CLIP vision + Llama text) and its image+text inputs.""" + from transformers import CLIPVisionConfig, LlamaConfig, LlavaConfig, LlavaForConditionalGeneration + + torch.manual_seed(0) + vision_config = CLIPVisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + image_size=32, + patch_size=16, + num_channels=3, + ) + text_config = LlamaConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=64, + max_position_embeddings=128, + ) + model = LlavaForConditionalGeneration( + LlavaConfig(vision_config=vision_config, text_config=text_config, image_token_index=1) + ).eval() + # 4 image tokens = (image_size / patch_size) ** 2, then 2 text tokens + inputs = { + "input_ids": torch.tensor([[1, 1, 1, 1, 5, 6]]), + "attention_mask": torch.ones(1, 6, dtype=torch.long), + "pixel_values": torch.randn(1, 3, 32, 32), + } + return model, inputs + + def _generation_config(self): + return GenerationConfig(cache_implementation="static", max_cache_len=MAX_CACHE_LEN, do_sample=False) + + def _decode_component(self, model=None): + """Multi-token `decode` component against a fixed-size `StaticCache`. The mask is a precomputed + input, so PT2E has no in-graph mask construction to trip on.""" + model = model if model is not None else self._tiny_model() + inputs = {"input_ids": torch.randint(0, 64, (1, 4)), "attention_mask": torch.ones(1, 4, dtype=torch.long)} + return decompose_for_generation( + model, inputs, generation_config=self._generation_config(), multi_token_decode=True + )["decode"] + + def _quantizer(self, name, dynamic=False): + """Build the PT2E quantizer named `name`: + + - `x86`: torchao-native (no ExecuTorch dependency), static per-channel int8 — or, with + `dynamic=True`, dynamic int8 (runtime-quantized activations + int8 weights), the lighter recipe + typical for decoders (a true weight-only PT2E quantizer isn't available in torchao/executorch). + - `xnnpack`: per-tensor quantizer for the XNNPACK ExecuTorch backend. + - `qnn`: vendor Qualcomm HTP quantizer (only lowers via the QNN ExecuTorch backend). + + `dynamic` applies to `x86` only; `xnnpack`/`qnn` are static per-tensor. + """ + if name == "x86": + from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import ( + X86InductorQuantizer, + get_default_x86_inductor_quantization_config, + ) + + return X86InductorQuantizer().set_global(get_default_x86_inductor_quantization_config(is_dynamic=dynamic)) + if name == "xnnpack": + from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + XNNPACKQuantizer, + get_symmetric_quantization_config, + ) + + return XNNPACKQuantizer().set_global(get_symmetric_quantization_config()) + if name == "qnn": + from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer + + return QnnQuantizer() + raise ValueError(f"unknown quantizer {name}") + + def _quantization_target(self, family): + """The `(model, inputs)` to quantize for `family`, picked so the traced forward builds no attention + mask in-graph — PT2E's `make_fx` retrace trips on in-graph mask construction: + + - dense / MoE: the multi-token `decode` component, whose attention mask is a precomputed input; + - SSM: the plain forward — a state-space model has no attention mask to build. + """ + if family == "dense": + return self._decode_component(self._tiny_model()) + if family == "moe": + return self._decode_component(self._moe_model()) + if family == "ssm": + return self._ssm_model(), {"input_ids": torch.randint(0, 64, (1, 8))} + raise ValueError(f"unknown family {family}") + + # ──────────────────────────────── Dynamo ──────────────────────────────── + + @pytest.mark.torch_export_test + def test_calibration_defaults_to_sample_inputs_with_warning(self): + """With no `calibration_dataset`, calibration falls back to a single pass on the sample inputs and + warns (one sample can hurt accuracy) — quantization still applies.""" + from transformers.exporters import DynamoConfig, DynamoExporter, exporter_dynamo + + decode_model, decode_inputs = self._decode_component() + with patch.object(exporter_dynamo.logger, "warning_once") as warning_once: + exported = DynamoExporter().export( + decode_model, + copy.deepcopy(decode_inputs), + DynamoConfig(dynamic=True, quantizer=self._quantizer("x86")), + ) + warning_once.assert_called() + self.assertTrue(_has_quantize_ops(exported)) + + @pytest.mark.torch_export_test + def test_calibration_dataset_captured_per_component(self): + """One generate-level `calibration_dataset` becomes a separate calibration set for each generation + component: running it through `generate` captures every component's real inputs, giving one set per + component (each the same length as the input dataset) so each calibrates on its own activations.""" + model = self._tiny_model() + calibration = [ + {"input_ids": torch.randint(0, 64, (1, n)), "attention_mask": torch.ones(1, n, dtype=torch.long)} + for n in (3, 4, 5) + ] + captured = capture_calibration_inputs( + model, copy.deepcopy(calibration), generation_config=self._generation_config(), multi_token_decode=True + ) + self.assertEqual(set(captured), {"prefill", "decode"}) + self.assertTrue(all(len(inputs) == len(calibration) for inputs in captured.values())) + + @parameterized.expand([("dense",), ("moe",), ("ssm",)]) + @pytest.mark.torch_export_test + def test_quantized_dynamo(self, family): + """Every architecture family quantizes to a quantized FX graph on the Dynamo backend, via the + torchao-native x86 quantizer (the natural graph-level PT2E quantizer, no executorch dependency) — + the same `config.quantizer` mechanism, no per-case code. The family unit is the module carrying + that family's distinct ops (dense/MoE `decode` component, SSM full forward). Static export.""" + from transformers.exporters import DynamoConfig, DynamoExporter + + model, inputs = self._quantization_target(family) + exported = DynamoExporter().export( + model, + copy.deepcopy(inputs), + DynamoConfig(dynamic=False, quantizer=self._quantizer("x86"), calibration_dataset=[copy.deepcopy(inputs)]), + ) + self.assertTrue(_has_quantize_ops(exported)) + + @pytest.mark.torch_export_test + def test_vlm_per_component_quantization(self): + """A VLM is quantized component-by-component, each with its OWN recipe — the realistic pattern — + via a `{component: config}` dict on `export_for_generation` (multi-token decode). The vision + encoder and projector get static int8; the language model and the multi-token decode get the + lighter dynamic recipe; `lm_head` stays fp32. Recipes are chosen per component, not globally.""" + from transformers.exporters import DynamoConfig, DynamoExporter + + model, inputs = self._vlm_model() + # False = static int8 (vision side), True = dynamic int8 (language/decoder side) + recipes = { + "image_encoder": False, + "multi_modal_projector": False, + "language_model": True, + "decode": True, + } + config = { + name: DynamoConfig(dynamic=True, quantizer=self._quantizer("x86", dynamic=dyn)) + for name, dyn in recipes.items() + } + config["lm_head"] = DynamoConfig(dynamic=True) # per-component choice: keep the output head in fp32 + components = DynamoExporter().export_for_generation(model, inputs, config, multi_token_decode=True) + + for name in recipes: + self.assertTrue(_has_quantize_ops(components[name]), f"{name} should be quantized") + self.assertFalse(_has_quantize_ops(components["lm_head"]), "lm_head should stay fp32") + # the recipes really differ: the decoder side is dynamically quantized, the vision side is static + self.assertTrue(_has_dynamic_quant_ops(components["decode"]), "decode should be dynamically quantized") + self.assertFalse(_has_dynamic_quant_ops(components["image_encoder"]), "image_encoder should be static int8") + + # ──────────────────────────────── ONNX ────────────────────────────────── + + @parameterized.expand([("dense",), ("moe",), ("ssm",)]) + @pytest.mark.onnx_export_test + def test_quantized_onnx(self, family): + """The same x86-quantizer recipe, lowered to ONNX: every family produces a QDQ graph + (QuantizeLinear nodes) that runs in ONNX Runtime. Static export throughout.""" + from transformers.utils import is_onnxruntime_available, is_onnxscript_available + + if not (is_onnxruntime_available() and is_onnxscript_available()): + self.skipTest("requires onnxruntime + onnxscript") + + from transformers.exporters import OnnxConfig, OnnxExporter + + model, inputs = self._quantization_target(family) + program = OnnxExporter().export( + model, + copy.deepcopy(inputs), + OnnxConfig( + dynamic=False, + quantizer=self._quantizer("x86"), + calibration_dataset=[copy.deepcopy(inputs)], + external_data=False, + ), + ) + self.assertTrue(_has_onnx_quantize_ops(program)) + # the QDQ graph must run in ONNX Runtime, not just parse — quantization error rules out an + # eager-parity check, so assert it executes to finite outputs + outputs = run_onnx_program(program, copy.deepcopy(inputs)) + self.assertTrue(outputs) + self.assertTrue(all(o.isfinite().all() for o in outputs.values() if o.is_floating_point())) + + # ────────────────────────────── ExecuTorch ────────────────────────────── + + @parameterized.expand( + [(family, quantizer) for family in ("dense", "moe", "ssm") for quantizer in ("xnnpack", "qnn")] + ) + @require_executorch + @pytest.mark.executorch_export_test + def test_quantized_executorch(self, family, quantizer): + """The same `config.quantizer` recipe, lowered to an ExecuTorch `.pte`: every family × delegatable + quantizer produces a program. The x86 quantizer is absent — its per-channel q/dq ops have no out + variant, so they stay undelegated and fail `to_executorch`; the ExecuTorch backends want the + per-tensor `xnnpack`/`qnn` quantizers instead.""" + from transformers.exporters import ExecutorchConfig, ExecutorchExporter + + if quantizer == "qnn": + if not _qnn_available(): + self.skipTest("requires the Qualcomm QNN SDK") + if family == "moe": + # QNN's quantizer annotates the int64 routing `arange` for per-tensor quant, which its + # `quantize_per_tensor` meta kernel rejects (float-only). A QNN HTP limitation, not ours. + self.skipTest("QNN HTP quantizer can't annotate MoE integer routing tensors") + if family == "ssm": + # QNN's `CanonicalizeConv` pass unconditionally `unsqueeze`s a conv bias, which Mamba2's + # bias-less grouped `conv1d` doesn't have. A QNN HTP limitation, not ours. + self.skipTest("QNN HTP `CanonicalizeConv` pass can't lower Mamba's bias-less conv1d") + + et_backend = "qnn" if quantizer == "qnn" else "xnnpack" + model, inputs = self._quantization_target(family) + program = ExecutorchExporter().export( + model, + copy.deepcopy(inputs), + ExecutorchConfig( + backend=et_backend, + dynamic=False, + quantizer=self._quantizer(quantizer), + calibration_dataset=[copy.deepcopy(inputs)], + ), + ) + self.assertIsNotNone(program) + # Export + lowering is the check here, not runtime execution. The quantized `.pte` runs fine + # standalone, but executing it in-process aborts (native SIGABRT) when the pytest-rerunfailures + # plugin's background socket-server thread is live — which it is in this suite's config — so + # running it here would take down the whole worker. `test_export` sidesteps this by never + # executing a quantized program. diff --git a/tests/exporters/test_utils.py b/tests/exporters/test_utils.py index 29b6a7eab0f5..e251701601fb 100644 --- a/tests/exporters/test_utils.py +++ b/tests/exporters/test_utils.py @@ -28,7 +28,8 @@ - The **`decompose_prefill_decode` guard** against generators that bypass the top-level forward — real generators call ``forward`` many times, so the guard is dead code without a targeted test. -- **`register_patch`** unresolvable-path fallback — real registrations point at real paths. +- **`register_patch`** / **`apply_patches`** deferred resolution — real registrations point at real + paths, so the import-safe "store the string, resolve at apply, raise if unresolvable" path is untested. Everything below targets one of those gaps. """ @@ -57,6 +58,7 @@ from transformers import GenerationConfig from transformers.exporters.utils import ( + apply_patches, cast_leaf_tensors, decompose_prefill_decode, duplicate_leaf_tensors, @@ -203,10 +205,12 @@ def _bad_factory(original): self.assertEqual(a.method(), "original") self.assertEqual(b.method(), "original") - def test_register_patch_skips_unresolvable_path(self): - # Real backends only register paths that resolve; the silent-skip fallback is what lets - # `exporter_onnx.py` and `exporter_executorch.py` co-exist when only one backend is - # installed. If it ever started raising, one of the two backends would fail to import. + def test_register_patch_defers_resolution_to_apply(self): + # `register_patch` stores the dotted path as a string without importing it, so registering a path + # into an uninstalled backend is import-safe — this is what lets `exporter_onnx.py` and + # `exporter_executorch.py` co-exist when only one backend is present. Resolution happens in + # `apply_patches`, which runs only for the backend actually exporting; a genuinely unresolvable + # path there is a bug and raises rather than being silently skipped. backend = "_test_unresolvable" @register_patch(backend, "does.not.exist.at.all") @@ -214,7 +218,11 @@ def _patch(original): return original try: - self.assertEqual(exporter_utils._PATCHES.get(backend, []), []) + obj_path, attribute, fn = exporter_utils._PATCHES[backend][0] + self.assertEqual((obj_path, attribute), ("does.not.exist.at", "all")) # stored, not resolved + with self.assertRaises((ImportError, AttributeError)): + with apply_patches(backend): + pass finally: exporter_utils._PATCHES.pop(backend, None)