Skip to content

Add QAT support and pass hooks to QuantizationRecipe. #21935

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MartinPavella wants to merge 1 commit into
base: main
Choose a base branch
from
+964 −53
Open

Add QAT support and pass hooks to QuantizationRecipe. #21935

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 164 additions & 44 deletions export/recipe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from abc import ABCMeta, abstractmethod
from dataclasses import dataclass
from enum import Enum, EnumMeta
from typing import Callable, List, Optional
from typing import Callable, Iterable, List, Optional

import torch
from executorch.exir import EdgeProgramManager, ExportedProgram
Expand Down Expand Up @@ -88,16 +88,50 @@ class QuantizationRecipe:
"""
Configuration recipe for quantization.

This class holds the configuration parameters for quantizing a model.
This class holds the configuration parameters for quantizing a model, supporting
both post-training quantization (PTQ) and quantization-aware training (QAT).

Attributes:

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just wondering, will there be some documentation update on how to use the recipes with all the newly added options? Or a real world example?

quantizers: Optional list of quantizers for model quantization
quantizers: Optional list of quantizers for model quantization.
ao_quantization_configs: Optional list of AOQuantizationConfig objects that pair
AOBaseConfig with optional filter functions
AOBaseConfig with optional filter functions.
is_qat: If True, use the QAT flow (prepare_qat_pt2e -> train_fn -> convert_pt2e).
If False (default), use the PTQ flow (prepare_pt2e -> calibrate -> convert_pt2e).
calibration_inputs_fn: Optional callable returning an iterable of input tuples used for
PTQ calibration. When None (default), the example inputs are used.
Ignored when is_qat=True.
train_fn: Callable that receives the prepared GraphModule and trains it.
Required when is_qat=True; ignored otherwise.
pre_prepare_passes: Optional list of callables applied to the captured GraphModule
before prepare_pt2e / prepare_qat_pt2e.
Each callable receives a GraphModule and must return a GraphModule.
post_prepare_passes: Optional list of callables applied to the prepared GraphModule
after prepare_pt2e / prepare_qat_pt2e and before calibration / training.
Each callable receives a GraphModule and must return a GraphModule.
pre_convert_passes: Optional list of callables applied to the GraphModule after
calibration (PTQ) or training (QAT) and before convert_pt2e.
Each callable receives a GraphModule and must return a GraphModule.
post_convert_passes: Optional list of callables applied to the GraphModule after convert_pt2e.
Each callable receives a GraphModule and must return a GraphModule.
"""

quantizers: Optional[List[Quantizer]] = None
ao_quantization_configs: Optional[List[AOQuantizationConfig]] = None
is_qat: bool = False
calibration_inputs_fn: Optional[Callable[[], Iterable[tuple]]] = None
train_fn: Optional[Callable[["torch.fx.GraphModule"], None]] = None
pre_prepare_passes: Optional[

@novak-vaclav novak-vaclav Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: If possible, I'd forbid None for these passes and default to [], allowing you to remove the if ... is not None statements.

List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]]
] = None
post_prepare_passes: Optional[
List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]]
] = None
pre_convert_passes: Optional[
List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]]
] = None
post_convert_passes: Optional[
List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]]
] = None

def get_quantizers(self) -> Optional[List[Quantizer]]:
"""
Expand Down Expand Up @@ -249,60 +283,135 @@ def _combine_recipes( # noqa: C901
Returns:
Combined ExportRecipe for multi-backend deployment
"""
# Extract components from individual recipes
all_partitioners = []
all_quantizers = []
all_ao_quantization_configs = []
all_pre_edge_passes = []
all_transform_passes = []

# Scalar fields that must be identical across all recipes.
def _assert_agree(field_name: str, values: list) -> None:
unique = set(values)
if len(unique) > 1:
raise ValueError(
f"Cannot combine recipes with conflicting '{field_name}' values: {unique}"
)

# Collect all components.
all_partitioners: list = []
all_quantizers: list = []
all_ao_quantization_configs: list = []
all_pre_edge_passes: list = []
all_edge_transform_passes: list = []
all_edge_manager_transform_passes: list = []
all_pre_prepare_passes: list = []
all_post_prepare_passes: list = []
all_pre_convert_passes: list = []
all_post_convert_passes: list = []
combined_backend_config = None

is_qat_values: list = []
train_fn_values: list = []
calibration_inputs_fn_values: list = []
strict_values: list = []
mode_values: list = []
pipeline_stages_values: list = []
source_transform_in_place_values: list = []

for recipe in backend_recipes:
# Collect pre-edge transform passes
if recipe.aten_transform_passes:
all_pre_edge_passes.extend(recipe.aten_transform_passes)

# Collect partitioners from lowering recipes
if recipe.lowering_recipe and recipe.lowering_recipe.partitioners:
all_partitioners.extend(recipe.lowering_recipe.partitioners)

# Collect transform passes from lowering recipes
if recipe.lowering_recipe and recipe.lowering_recipe.edge_transform_passes:
all_transform_passes.extend(
recipe.lowering_recipe.edge_transform_passes
)

# Collect for quantize stage
if quantization_recipe := recipe.quantization_recipe:
# Collect PT2E quantizers
if quantization_recipe.quantizers:
all_quantizers.extend(quantization_recipe.quantizers)

# Collect source transform configs
if quantization_recipe.ao_quantization_configs:
all_ao_quantization_configs.extend(
quantization_recipe.ao_quantization_configs
if lr := recipe.lowering_recipe:
if lr.partitioners:
all_partitioners.extend(lr.partitioners)
if lr.edge_transform_passes:
all_edge_transform_passes.extend(lr.edge_transform_passes)
if lr.edge_manager_transform_passes:
all_edge_manager_transform_passes.extend(
lr.edge_manager_transform_passes
)

# Use the first backend config as base
if qr := recipe.quantization_recipe:
if qr.quantizers:
all_quantizers.extend(qr.quantizers)
if qr.ao_quantization_configs:
all_ao_quantization_configs.extend(qr.ao_quantization_configs)
is_qat_values.append(qr.is_qat)
train_fn_values.append(qr.train_fn)
calibration_inputs_fn_values.append(qr.calibration_inputs_fn)
if qr.pre_prepare_passes:
all_pre_prepare_passes.extend(qr.pre_prepare_passes)
if qr.post_prepare_passes:
all_post_prepare_passes.extend(qr.post_prepare_passes)
if qr.pre_convert_passes:
all_pre_convert_passes.extend(qr.pre_convert_passes)
if qr.post_convert_passes:
all_post_convert_passes.extend(qr.post_convert_passes)

strict_values.append(recipe.strict)
mode_values.append(recipe.mode)
pipeline_stages_values.append(
tuple(recipe.pipeline_stages) if recipe.pipeline_stages else None
)
source_transform_in_place_values.append(recipe.source_transform_in_place)

if combined_backend_config is None and recipe.executorch_backend_config:

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe this logic causes the first executorch_backend_config that is found to be used, however it's not enforced that they are the same across recipes. It may be desired behavior, if so, please add a comment.

combined_backend_config = copy.deepcopy(
recipe.executorch_backend_config
)

# Create combined quantization recipe
# Validate fields that must agree across all recipes.
_assert_agree("strict", strict_values)
_assert_agree("mode", mode_values)
_assert_agree("pipeline_stages", pipeline_stages_values)
_assert_agree("source_transform_in_place", source_transform_in_place_values)

# is_qat must agree across all recipes that carry a QuantizationRecipe.
_assert_agree("is_qat", is_qat_values)
# train_fn must have at most one non-None value across all recipes.
non_none_train_fns = [f for f in train_fn_values if f is not None]
if len(non_none_train_fns) > 1:

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also len(non_none_train_fns) can't be 0 I think

raise ValueError(
"Cannot combine recipes: more than one recipe provides a train_fn."
)
# Multiple calibration_inputs_fn values are chained into a single factory.
non_none_calib_fns = [f for f in calibration_inputs_fn_values if f is not None]
if len(non_none_calib_fns) > 1:
_fns = non_none_calib_fns

def _combined_calib_fn():
for _fn in _fns:
yield from _fn()

combined_calib_fn = _combined_calib_fn
else:
combined_calib_fn = non_none_calib_fns[0] if non_none_calib_fns else None

# Build combined QuantizationRecipe.
combined_quantization_recipe = None
if all_quantizers or all_ao_quantization_configs:
if (

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I'd add an INFO log message in the else branch of this if to let the user know the QuantizationRecipe is all default, which might be suspicious.

all_quantizers
or all_ao_quantization_configs
or all_pre_prepare_passes
or all_post_prepare_passes
or all_pre_convert_passes
or all_post_convert_passes
):
combined_quantization_recipe = QuantizationRecipe(
quantizers=all_quantizers if all_quantizers else None,
ao_quantization_configs=(
all_ao_quantization_configs if all_ao_quantization_configs else None
),
quantizers=all_quantizers or None,
ao_quantization_configs=all_ao_quantization_configs or None,
is_qat=is_qat_values[0] if is_qat_values else False,
train_fn=non_none_train_fns[0] if non_none_train_fns else None,
calibration_inputs_fn=combined_calib_fn,
pre_prepare_passes=all_pre_prepare_passes or None,
post_prepare_passes=all_post_prepare_passes or None,
pre_convert_passes=all_pre_convert_passes or None,
post_convert_passes=all_post_convert_passes or None,
)

# Create combined lowering recipe
combined_lowering_recipe = None
if all_partitioners or all_transform_passes:
if (

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Same thing here, I'd add a INFO log message in the else branch of this if to let the user know the LoweringRecipe is all default.

all_partitioners
or all_edge_transform_passes
or all_edge_manager_transform_passes
):
edge_compile_config = None

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know this is not part of your changes, but the handling of edge_compile_config is weird - similarily to executorch_backend_config, the first value found is used - it may be desired behavior, if so, please add a comment why, if you know.

for recipe in backend_recipes:
if (
Expand All @@ -313,10 +422,9 @@ def _combine_recipes( # noqa: C901
break

combined_lowering_recipe = LoweringRecipe(
partitioners=all_partitioners if all_partitioners else None,
edge_transform_passes=(
all_transform_passes if all_transform_passes else None
),
partitioners=all_partitioners or None,
edge_transform_passes=all_edge_transform_passes or None,
edge_manager_transform_passes=all_edge_manager_transform_passes or None,
edge_compile_config=edge_compile_config or EdgeCompileConfig(),
)

Expand All @@ -326,7 +434,19 @@ def _combine_recipes( # noqa: C901
return cls(
name=recipe_name,
quantization_recipe=combined_quantization_recipe,
aten_transform_passes=all_pre_edge_passes,
aten_transform_passes=all_pre_edge_passes or None,
lowering_recipe=combined_lowering_recipe,
executorch_backend_config=combined_backend_config,
strict=strict_values[0] if strict_values else True,
mode=mode_values[0] if mode_values else Mode.RELEASE,
pipeline_stages=(
list(pipeline_stages_values[0])
if pipeline_stages_values and pipeline_stages_values[0] is not None
else None
),
source_transform_in_place=(
source_transform_in_place_values[0]
if source_transform_in_place_values
else False
),
)
64 changes: 58 additions & 6 deletions export/stages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,11 @@
from torch._export.pass_base import PassType
from torch.fx.passes.infra.pass_manager import PassManager as GraphModulePassManager
from torchao.quantization import quantize_
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
from torchao.quantization.pt2e.quantize_pt2e import (
convert_pt2e,
prepare_pt2e,
prepare_qat_pt2e,
)
from torchao.quantization.pt2e.quantizer import (
ComposableQuantizer,
Quantizer as TorchAOPT2EQuantizer,
Expand Down Expand Up @@ -423,6 +427,15 @@ def _get_quantizer_for_prepare_pt2e(self, quantizers: List[Any]):
else:
raise ValueError("No quantizers detected")

@staticmethod
def _apply_passes(
model: "torch.fx.GraphModule",
passes: Optional[List[Callable]],
) -> "torch.fx.GraphModule":
for pass_fn in passes or []:
model = pass_fn(model)

@novak-vaclav novak-vaclav Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I'd implement some kind of error handling, or at least logging.

return model

def run(self, artifact: PipelineArtifact) -> None:
if not self._quantization_recipe or not self._quantization_recipe.quantizers:
logging.info(
Expand All @@ -433,6 +446,7 @@ def run(self, artifact: PipelineArtifact) -> None:

assert isinstance(artifact.data, dict)

recipe = self._quantization_recipe
models = artifact.data
example_inputs = artifact.get_context("example_inputs")

Expand All @@ -447,15 +461,53 @@ def run(self, artifact: PipelineArtifact) -> None:
inputs = example_inputs[method_name][0]
captured_graph = torch.export.export(model, inputs, strict=True).module()

quantizer = self._get_quantizer_for_prepare_pt2e(
self._quantization_recipe.quantizers # pyre-ignore
# Pass 1: pre-prepare passes.
captured_graph = self._apply_passes(
captured_graph, recipe.pre_prepare_passes
)
prepared_model = prepare_pt2e(captured_graph, quantizer)

for calibration_input in example_inputs[method_name]:
prepared_model(*calibration_input)
quantizer = self._get_quantizer_for_prepare_pt2e(recipe.quantizers)

if recipe.is_qat:
if recipe.train_fn is None:
raise ValueError("train_fn must be provided when is_qat=True")
prepared_model = prepare_qat_pt2e(captured_graph, quantizer)

# Pass 2: post-prepare passes.
prepared_model = self._apply_passes(
prepared_model, recipe.post_prepare_passes
)

recipe.train_fn(prepared_model)

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Model should be put into training mode and then moved back to eval mode.

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI pointed out that the train_fn is the same for all model methods, which seems too restrictive to me. I'd make the train_fn a Dict[str, Callable] mapping of method to a train function (similarily to example_inputs). If you decide to go with this suggestion, additional validation would also have to be implemented (checking there is a training function for each method...)

else:
prepared_model = prepare_pt2e(captured_graph, quantizer)

@novak-vaclav novak-vaclav Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Model should be put into eval mode just in case.


# Pass 2: post-prepare passes.
prepared_model = self._apply_passes(
prepared_model, recipe.post_prepare_passes
)

# Use custom calibration inputs when provided; fall back to example inputs.
if recipe.calibration_inputs_fn is not None:
calibration_inputs = recipe.calibration_inputs_fn()
else:
calibration_inputs = example_inputs[method_name]

for calibration_input in calibration_inputs:
prepared_model(*calibration_input)

# Pass 3: pre-convert passes.
prepared_model = self._apply_passes(
prepared_model, recipe.pre_convert_passes
)

quantized_model = convert_pt2e(prepared_model)

# Pass 4: post-convert passes.
quantized_model = self._apply_passes(
quantized_model, recipe.post_convert_passes
)

quantized_models[method_name] = quantized_model

self._artifact = artifact.copy_with_new_data(quantized_models)
Expand Down
Loading
Loading