diff --git a/export/recipe.py b/export/recipe.py index c205491f80e..2a6be3e6200 100644 --- a/export/recipe.py +++ b/export/recipe.py @@ -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 @@ -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: - 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[ + 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]]: """ @@ -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: 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: + 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 ( + 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 ( + all_partitioners + or all_edge_transform_passes + or all_edge_manager_transform_passes + ): edge_compile_config = None for recipe in backend_recipes: if ( @@ -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(), ) @@ -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 + ), ) diff --git a/export/stages.py b/export/stages.py index a68ad408493..5ed4c49f9c6 100644 --- a/export/stages.py +++ b/export/stages.py @@ -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, @@ -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) + return model + def run(self, artifact: PipelineArtifact) -> None: if not self._quantization_recipe or not self._quantization_recipe.quantizers: logging.info( @@ -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") @@ -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) + else: + prepared_model = prepare_pt2e(captured_graph, quantizer) + + # 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) diff --git a/export/tests/test_export_recipe.py b/export/tests/test_export_recipe.py index d22442371e2..0443f7c4e55 100644 --- a/export/tests/test_export_recipe.py +++ b/export/tests/test_export_recipe.py @@ -7,9 +7,16 @@ # pyre-strict import unittest -from typing import Any, Dict, Optional, Sequence - -from executorch.export.recipe import ExportRecipe, RecipeType +from typing import Any, Dict, List, Optional, Sequence +from unittest.mock import Mock + +from executorch.export.recipe import ( + ExportRecipe, + LoweringRecipe, + Mode, + QuantizationRecipe, + RecipeType, +) from executorch.export.recipe_provider import BackendRecipeProvider from executorch.export.recipe_registry import recipe_registry @@ -129,3 +136,436 @@ def test_get_recipe_with_kwargs_verification(self) -> None: # Verify that the kwargs were passed to the backend provider's create_recipe method self.assertIsNotNone(self.provider.last_kwargs) self.assertEqual(self.provider.last_kwargs, kwargs) + + +# --------------------------------------------------------------------------- +# Helpers shared by combine-recipe tests +# --------------------------------------------------------------------------- + + +def _make_pass(name: str, call_log: List[str]): + """Return a graph-module pass that appends *name* to *call_log*.""" + + def pass_fn(m): + call_log.append(name) + return m + + return pass_fn + + +class TestCombineRecipesEmpty(unittest.TestCase): + def test_empty_recipes_raises(self) -> None: + with self.assertRaises(ValueError): + ExportRecipe.combine([]) + + +class TestCombineRecipesSingleRecipe(unittest.TestCase): + def test_single_recipe_returned_unchanged(self) -> None: + recipe = ExportRecipe(name="solo") + result = ExportRecipe.combine([recipe]) + self.assertIs(result, recipe) + + +class TestCombineRecipesScalarFields(unittest.TestCase): + """Fields that must be identical across all combined recipes.""" + + def test_conflicting_strict_raises(self) -> None: + r1 = ExportRecipe(name="a", strict=True) + r2 = ExportRecipe(name="b", strict=False) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("strict", str(cm.exception)) + + def test_conflicting_mode_raises(self) -> None: + r1 = ExportRecipe(name="a", mode=Mode.DEBUG) + r2 = ExportRecipe(name="b", mode=Mode.RELEASE) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("mode", str(cm.exception)) + + def test_conflicting_source_transform_in_place_raises(self) -> None: + r1 = ExportRecipe(name="a", source_transform_in_place=True) + r2 = ExportRecipe(name="b", source_transform_in_place=False) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("source_transform_in_place", str(cm.exception)) + + def test_agreeing_scalar_fields_are_preserved(self) -> None: + r1 = ExportRecipe( + name="a", strict=False, mode=Mode.DEBUG, source_transform_in_place=True + ) + r2 = ExportRecipe( + name="b", strict=False, mode=Mode.DEBUG, source_transform_in_place=True + ) + result = ExportRecipe.combine([r1, r2]) + self.assertFalse(result.strict) + self.assertEqual(result.mode, Mode.DEBUG) + self.assertTrue(result.source_transform_in_place) + + def test_name_is_joined_from_input_recipe_names(self) -> None: + r1 = ExportRecipe(name="backend_a") + r2 = ExportRecipe(name="backend_b") + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.name, "backend_a_backend_b") + + def test_custom_recipe_name_is_used(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2], recipe_name="custom_name") + self.assertEqual(result.name, "custom_name") + + +class TestCombineRecipesAtenTransformPasses(unittest.TestCase): + def test_aten_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe(name="a", aten_transform_passes=[pass1]) + r2 = ExportRecipe(name="b", aten_transform_passes=[pass2]) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.aten_transform_passes, [pass1, pass2]) + + def test_aten_transform_passes_none_when_both_empty(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.aten_transform_passes) + + def test_aten_transform_passes_one_side_none(self) -> None: + pass1 = Mock() + r1 = ExportRecipe(name="a", aten_transform_passes=[pass1]) + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.aten_transform_passes, [pass1]) + + +class TestCombineRecipesLowering(unittest.TestCase): + def test_partitioners_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe(name="a", lowering_recipe=LoweringRecipe(partitioners=[p1])) + r2 = ExportRecipe(name="b", lowering_recipe=LoweringRecipe(partitioners=[p2])) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNotNone(result.lowering_recipe) + self.assertEqual(result.lowering_recipe.partitioners, [p1, p2]) + + def test_edge_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe( + name="a", lowering_recipe=LoweringRecipe(edge_transform_passes=[pass1]) + ) + r2 = ExportRecipe( + name="b", lowering_recipe=LoweringRecipe(edge_transform_passes=[pass2]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.lowering_recipe.edge_transform_passes, [pass1, pass2]) + + def test_edge_manager_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe( + name="a", + lowering_recipe=LoweringRecipe(edge_manager_transform_passes=[pass1]), + ) + r2 = ExportRecipe( + name="b", + lowering_recipe=LoweringRecipe(edge_manager_transform_passes=[pass2]), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual( + result.lowering_recipe.edge_manager_transform_passes, [pass1, pass2] + ) + + def test_lowering_recipe_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.lowering_recipe) + + def test_edge_compile_config_taken_from_first_recipe_with_one(self) -> None: + from executorch.exir.capture import EdgeCompileConfig + + config = EdgeCompileConfig() + r1 = ExportRecipe(name="a") + r2 = ExportRecipe( + name="b", + lowering_recipe=LoweringRecipe( + partitioners=[Mock()], edge_compile_config=config + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIs(result.lowering_recipe.edge_compile_config, config) + + +class TestCombineRecipesQuantization(unittest.TestCase): + def test_quantizers_merged(self) -> None: + q1 = Mock() + q2 = Mock() + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[q1]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[q2]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNotNone(result.quantization_recipe) + self.assertEqual(result.quantization_recipe.quantizers, [q1, q2]) + + def test_ao_quantization_configs_merged(self) -> None: + from executorch.export.recipe import AOQuantizationConfig + from torchao.core.config import AOBaseConfig + + cfg1 = AOQuantizationConfig(ao_base_config=Mock(spec=AOBaseConfig)) + cfg2 = AOQuantizationConfig(ao_base_config=Mock(spec=AOBaseConfig)) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(ao_quantization_configs=[cfg1]), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(ao_quantization_configs=[cfg2]), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual( + result.quantization_recipe.ao_quantization_configs, [cfg1, cfg2] + ) + + def test_quantization_recipe_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.quantization_recipe) + + def test_conflicting_is_qat_raises(self) -> None: + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=False), + ) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("is_qat", str(cm.exception)) + + def test_agreeing_is_qat_preserved(self) -> None: + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertTrue(result.quantization_recipe.is_qat) + + def test_two_train_fns_raises(self) -> None: + fn1 = Mock() + fn2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn2 + ), + ) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("train_fn", str(cm.exception)) + + def test_single_train_fn_preserved(self) -> None: + fn = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIs(result.quantization_recipe.train_fn, fn) + + def test_single_calibration_inputs_fn_preserved(self) -> None: + fn = Mock(return_value=[(1,), (2,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn + ), + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIs(result.quantization_recipe.calibration_inputs_fn, fn) + + def test_two_calibration_inputs_fns_chained(self) -> None: + fn1 = Mock(return_value=[(1,), (2,)]) + fn2 = Mock(return_value=[(3,), (4,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn2 + ), + ) + result = ExportRecipe.combine([r1, r2]) + combined_fn = result.quantization_recipe.calibration_inputs_fn + self.assertIsNotNone(combined_fn) + # Each factory must be called exactly once when the combined factory is consumed. + all_inputs = list(combined_fn()) + fn1.assert_called_once_with() + fn2.assert_called_once_with() + self.assertEqual(all_inputs, [(1,), (2,), (3,), (4,)]) + + def test_three_calibration_inputs_fns_chained_in_order(self) -> None: + fn1 = Mock(return_value=[(1,)]) + fn2 = Mock(return_value=[(2,)]) + fn3 = Mock(return_value=[(3,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn2 + ), + ) + r3 = ExportRecipe( + name="c", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn3 + ), + ) + result = ExportRecipe.combine([r1, r2, r3]) + combined_fn = result.quantization_recipe.calibration_inputs_fn + self.assertEqual(list(combined_fn()), [(1,), (2,), (3,)]) + + def test_no_calibration_inputs_fn_stays_none(self) -> None: + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.quantization_recipe.calibration_inputs_fn) + + def test_pre_prepare_passes_merged(self) -> None: + log: List[str] = [] + p1 = _make_pass("pre_a", log) + p2 = _make_pass("pre_b", log) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_prepare_passes, [p1, p2]) + + def test_post_prepare_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_prepare_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_prepare_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.post_prepare_passes, [p1, p2]) + + def test_pre_convert_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_convert_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_convert_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_convert_passes, [p1, p2]) + + def test_post_convert_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_convert_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_convert_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.post_convert_passes, [p1, p2]) + + def test_all_pass_lists_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + qr = result.quantization_recipe + self.assertIsNone(qr.pre_prepare_passes) + self.assertIsNone(qr.post_prepare_passes) + self.assertIsNone(qr.pre_convert_passes) + self.assertIsNone(qr.post_convert_passes) + + def test_pass_lists_preserved_when_only_one_recipe_contributes(self) -> None: + p = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p] + ), + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_prepare_passes, [p]) diff --git a/export/tests/test_export_stages.py b/export/tests/test_export_stages.py index e9acee0ea26..2269cda5b77 100644 --- a/export/tests/test_export_stages.py +++ b/export/tests/test_export_stages.py @@ -389,6 +389,12 @@ def test_run_with_quantizers( mock_quantizer = self.create_dummy_quantizer() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None stage = QuantizeStage(mock_recipe) # Mock the torch.export.export chain @@ -436,11 +442,304 @@ def test_run_with_quantizers( self.assertEqual(artifact.data["forward"], self.model) self.assertIsNot(result_artifact.data["forward"], self.model) + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_calls_prepare_qat_pt2e( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + ) -> None: + """QAT flow: prepare_qat_pt2e is called and train_fn is invoked with the prepared model.""" + mock_quantizer = self.create_dummy_quantizer() + train_fn = Mock() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.train_fn = train_fn + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_captured_graph = Mock() + mock_exported_program.module.return_value = mock_captured_graph + mock_torch_export.return_value = mock_exported_program + + mock_composed_quantizer = Mock() + mock_composable_quantizer.return_value = mock_composed_quantizer + mock_prepared_model = Mock() + mock_prepare_qat_pt2e.return_value = mock_prepared_model + mock_quantized_model = Mock() + mock_convert_pt2e.return_value = mock_quantized_model + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # prepare_qat_pt2e must be called, not prepare_pt2e + mock_prepare_qat_pt2e.assert_called_once_with( + mock_captured_graph, mock_composed_quantizer + ) + # train_fn must be called with the prepared model + train_fn.assert_called_once_with(mock_prepared_model) + # convert_pt2e must still be called after training + mock_convert_pt2e.assert_called_once_with(mock_prepared_model) + + result_artifact = stage.get_artifacts() + self.assertEqual(result_artifact.data["forward"], mock_quantized_model) + + @patch("torch.export.export") + def test_run_qat_missing_train_fn_raises(self, mock_torch_export: Mock) -> None: + """QAT flow with train_fn=None must raise ValueError.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.train_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + + with self.assertRaises(ValueError) as cm: + stage.run(artifact) + self.assertIn("train_fn must be provided when is_qat=True", str(cm.exception)) + + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_does_not_call_prepare_qat_pt2e( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + ) -> None: + """PTQ flow must not call prepare_qat_pt2e (regression guard).""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + mock_prepare_pt2e.assert_called_once() + mock_prepare_qat_pt2e.assert_not_called() + + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_four_passes_called_in_order( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + ) -> None: + """All four pass hooks are called at the correct points in the PTQ flow.""" + call_order = [] + + def make_pass(name): + def pass_fn(m): + call_order.append(name) + return m + + return pass_fn + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")] + mock_recipe.post_prepare_passes = [make_pass("post_prepare")] + mock_recipe.pre_convert_passes = [make_pass("pre_convert")] + mock_recipe.post_convert_passes = [make_pass("post_convert")] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_graph = Mock() + mock_exported_program.module.return_value = mock_graph + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + self.assertEqual( + call_order, + ["pre_prepare", "post_prepare", "pre_convert", "post_convert"], + ) + + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_four_passes_called_in_order( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + ) -> None: + """All four pass hooks are called at the correct points in the QAT flow.""" + call_order = [] + + def make_pass(name): + def pass_fn(m): + call_order.append(name) + return m + + return pass_fn + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.train_fn = Mock() + mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")] + mock_recipe.post_prepare_passes = [make_pass("post_prepare")] + mock_recipe.pre_convert_passes = [make_pass("pre_convert")] + mock_recipe.post_convert_passes = [make_pass("post_convert")] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + self.assertEqual( + call_order, + ["pre_prepare", "post_prepare", "pre_convert", "post_convert"], + ) + + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_uses_calibration_inputs_fn_when_provided( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + ) -> None: + """When calibration_inputs_fn is set, it is called and its output is used for calibration.""" + custom_input = (torch.randn(2, 10),) + calibration_inputs_fn = Mock(return_value=[custom_input]) + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = calibration_inputs_fn + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepared_model = Mock() + mock_prepare_pt2e.return_value = mock_prepared_model + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # calibration_inputs_fn must be called with no arguments + calibration_inputs_fn.assert_called_once_with() + # prepared model must be called with the custom calibration input + mock_prepared_model.assert_called_once_with(*custom_input) + + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_falls_back_to_example_inputs_when_no_calibration_fn( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + ) -> None: + """When calibration_inputs_fn is None, example inputs are used for calibration.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepared_model = Mock() + mock_prepare_pt2e.return_value = mock_prepared_model + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # The prepared model must be called with the example inputs (one tuple) + mock_prepared_model.assert_called_once_with(*self.example_inputs[0]) + def test_run_empty_example_inputs(self) -> None: """Test error when example inputs list is empty.""" mock_quantizer = Mock() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None stage = QuantizeStage(mock_recipe) context = {"example_inputs": {"forward": []}} artifact = PipelineArtifact(data=self.models_dict, context=context)