Let a QuantizationRecipe decompose the capture before quantization

backends/qualcomm/genai_pipeline/quant_utilities.py

@@ -214,6 +214,8 @@ def encoding_override(  # noqa: C901
     pbq_target = {
         torch.ops.torchao.dequantize_affine,
         torch.ops.torchao.quantize_affine,
+        torch.ops.torchao.dequantize_affine.default,
+        torch.ops.torchao.quantize_affine.default,
     }
     pcq_target = {
         torch.ops.quantized_decomposed.dequantize_per_channel.default,

examples/qualcomm/oss_scripts/llama/mix_precision_analyzer.py

@@ -67,11 +67,13 @@ def __init__(
         }
         self.q_ops = {
             torch.ops.torchao.quantize_affine,
+            torch.ops.torchao.quantize_affine.default,
             torch.ops.quantized_decomposed.quantize_per_channel.default,
             torch.ops.quantized_decomposed.quantize_per_tensor.default,
         }
         self.dq_ops = {
             torch.ops.torchao.dequantize_affine,
+            torch.ops.torchao.dequantize_affine.default,
             torch.ops.quantized_decomposed.dequantize_per_channel.default,
             torch.ops.quantized_decomposed.dequantize_per_tensor.default,
         }

examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py

@@ -915,6 +915,8 @@ def _encoding_override(  # noqa: C901
         pbq_target = {
             torch.ops.torchao.dequantize_affine,
             torch.ops.torchao.quantize_affine,
+            torch.ops.torchao.dequantize_affine.default,
+            torch.ops.torchao.quantize_affine.default,
         }
         pcq_target = {
             torch.ops.quantized_decomposed.dequantize_per_channel.default,

export/recipe.py

@@ -162,6 +162,11 @@ class QuantizationRecipe:
                             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.
+        preserve_ops: When set, the program captured for quantization is decomposed to
+                      core ATen, keeping these ops whole, before pre_prepare_passes and
+                      prepare_pt2e / prepare_qat_pt2e see it. Use it when the quantizer
+                      matches core ATen ops. When None (default), the capture is not
+                      decomposed.
     """
 
     quantizers: Optional[List[Quantizer]] = None
@@ -182,6 +187,7 @@ class QuantizationRecipe:
     post_convert_passes: Optional[
         List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]]
     ] = None
+    preserve_ops: Optional[List[torch._ops.OpOverload]] = None
 
     def get_quantizers(self) -> Optional[List[Quantizer]]:
         """
@@ -253,6 +259,7 @@ class _CombineAccumulator:
     dynamic_batch_size_values: list = field(default_factory=list)
     train_fn_values: list = field(default_factory=list)
     calibration_inputs_fn_values: list = field(default_factory=list)
+    preserve_ops_values: list = field(default_factory=list)
     strict_values: list = field(default_factory=list)
     mode_values: list = field(default_factory=list)
     pipeline_stages_values: list = field(default_factory=list)
@@ -405,6 +412,7 @@ def _combine_quantization_recipe(
         dynamic_batch_size_values: list,
         train_fn_values: list,
         calibration_inputs_fn_values: list,
+        preserve_ops_values: list,
         all_quantizers: list,
         all_ao_quantization_configs: list,
         all_pre_prepare_passes: list,
@@ -428,6 +436,15 @@ def _combine_quantization_recipe(
                 "Cannot combine recipes: more than one recipe provides a train_fn."
             )
 
+        # The capture is decomposed if any recipe asks for it, keeping every op any
+        # of them keeps.
+        preserve_ops_lists = [ops for ops in preserve_ops_values if ops is not None]
+        combined_preserve_ops = (
+            list(dict.fromkeys(op for ops in preserve_ops_lists for op in ops))
+            if preserve_ops_lists
+            else None
+        )
+
         # 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:
@@ -468,6 +485,7 @@ def _combined_calib_fn():
             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,
+            preserve_ops=combined_preserve_ops,
         )
 
     @classmethod
@@ -585,6 +603,7 @@ def _collect_quantization_fields(
         acc.dynamic_batch_size_values.append(qr.dynamic_batch_size)
         acc.train_fn_values.append(qr.train_fn)
         acc.calibration_inputs_fn_values.append(qr.calibration_inputs_fn)
+        acc.preserve_ops_values.append(qr.preserve_ops)
         if qr.pre_prepare_passes:
             acc.pre_prepare_passes.extend(qr.pre_prepare_passes)
         if qr.post_prepare_passes:
@@ -671,6 +690,7 @@ def _combine_recipes(
             dynamic_batch_size_values=acc.dynamic_batch_size_values,
             train_fn_values=acc.train_fn_values,
             calibration_inputs_fn_values=acc.calibration_inputs_fn_values,
+            preserve_ops_values=acc.preserve_ops_values,
             all_quantizers=acc.quantizers,
             all_ao_quantization_configs=acc.ao_quantization_configs,
             all_pre_prepare_passes=acc.pre_prepare_passes,

export/stages.py

@@ -23,6 +23,7 @@
 from executorch.export.recipe import LoweringRecipe, QuantizationRecipe
 from executorch.export.types import StageType
 from torch import nn
+from torch._decomp import remove_decompositions
 from torch._export.pass_base import PassType
 from torch.fx.passes.infra.pass_manager import PassManager as GraphModulePassManager
 from torchao.quantization import quantize_
@@ -598,18 +599,22 @@ def run(self, artifact: PipelineArtifact) -> None:
                     {0: batch} if isinstance(t, torch.Tensor) else None for t in inputs
                 )
 
-            # QAT requires the model to be in training mode at capture time so
-            # that batch_norm and dropout decompose with training-mode semantics.
-            if recipe.is_qat:
-                model.train()
+            # batch_norm and dropout are captured with the mode's semantics: QAT
+            # trains the captured graph, PTQ calibrates the graph as it will run.
+            model.train(recipe.is_qat)
 
             with torch.no_grad():
                 for hook in self._pre_trace_hooks or []:
                     hook(method_name, model)
 
-            captured_graph = torch.export.export(
+            program = torch.export.export(
                 model, inputs, dynamic_shapes=export_dynamic_shapes, strict=True
-            ).module()
+            )
+            if recipe.preserve_ops is not None:
+                decomp_table = torch.export.default_decompositions()
+                remove_decompositions(decomp_table, recipe.preserve_ops)
+                program = program.run_decompositions(decomp_table)
+            captured_graph = program.module()
 
             # Pass 1: pre-prepare passes.
             captured_graph = self._apply_passes(

export/tests/test_export_recipe.py

@@ -676,6 +676,47 @@ def test_single_train_fn_preserved(self) -> None:
         result = ExportRecipe.combine([r1, r2])
         self.assertIs(result.quantization_recipe.train_fn, fn)
 
+    def test_preserve_ops_are_combined(self) -> None:
+        linear = torch.ops.aten.linear.default
+        convolution = torch.ops.aten.convolution.default
+        r1 = ExportRecipe(
+            name="a",
+            quantization_recipe=QuantizationRecipe(
+                quantizers=[Mock()], preserve_ops=[linear]
+            ),
+        )
+        r2 = ExportRecipe(
+            name="b",
+            quantization_recipe=QuantizationRecipe(
+                quantizers=[Mock()], preserve_ops=[convolution, linear]
+            ),
+        )
+        result = ExportRecipe.combine([r1, r2])
+        self.assertEqual(result.quantization_recipe.preserve_ops, [linear, convolution])
+
+    def test_preserve_ops_from_one_recipe_decomposes(self) -> None:
+        r1 = ExportRecipe(
+            name="a",
+            quantization_recipe=QuantizationRecipe(
+                quantizers=[Mock()], preserve_ops=[]
+            ),
+        )
+        r2 = ExportRecipe(
+            name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()])
+        )
+        result = ExportRecipe.combine([r1, r2])
+        self.assertEqual(result.quantization_recipe.preserve_ops, [])
+
+    def test_preserve_ops_unset_does_not_decompose(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.preserve_ops)
+
     def test_single_calibration_inputs_fn_preserved(self) -> None:
         fn = Mock(return_value=[(1,), (2,)])
         r1 = ExportRecipe(

export/tests/test_export_stages.py

@@ -694,6 +694,7 @@ def test_run_with_quantizers(
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = False
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.calibration_inputs_fn = None
         mock_recipe.pre_prepare_passes = None
         mock_recipe.post_prepare_passes = None
@@ -746,6 +747,77 @@ 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_pt2e")
+    @patch("executorch.export.stages.ComposableQuantizer")
+    @patch("torch.export.export")
+    def test_run_with_preserve_ops(
+        self,
+        mock_torch_export: Mock,
+        mock_composable_quantizer: Mock,
+        mock_prepare_pt2e: Mock,
+        mock_convert_pt2e: Mock,
+    ) -> None:
+        """preserve_ops decomposes the capture, keeping those ops, before prepare."""
+        mock_recipe = Mock(spec=QuantizationRecipe)
+        mock_recipe.quantizers = [self.create_dummy_quantizer()]
+        mock_recipe.is_qat = False
+        mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = [torch.ops.aten.linear.default]
+        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
+        run_decompositions = mock_torch_export.return_value.run_decompositions
+
+        QuantizeStage(mock_recipe).run(
+            PipelineArtifact(data=self.models_dict, context=self.context)
+        )
+
+        (decomp_table,), _ = run_decompositions.call_args
+        self.assertNotIn(torch.ops.aten.linear.default, decomp_table)
+        self.assertIn(torch.ops.aten.conv2d.default, decomp_table)
+        mock_prepare_pt2e.assert_called_once_with(
+            run_decompositions.return_value.module.return_value,
+            mock_composable_quantizer.return_value,
+        )
+
+    def test_quantizer_sees_core_aten_except_preserve_ops(self) -> None:
+        """The quantizer sees core ATen plus preserve_ops, or plain export without."""
+
+        def annotated_targets(preserve_ops: list | None) -> list:
+            seen = []
+
+            class RecordingQuantizer(TorchAOPT2EQuantizer):
+                def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+                    seen.extend(
+                        node.target
+                        for node in model.graph.nodes
+                        if node.op == "call_function"
+                    )
+                    return model
+
+                def validate(self, model: torch.fx.GraphModule) -> None:
+                    pass
+
+            recipe = QuantizationRecipe(
+                quantizers=[RecordingQuantizer()], preserve_ops=preserve_ops
+            )
+            QuantizeStage(recipe).run(
+                PipelineArtifact(data=self.models_dict, context=self.context)
+            )
+            return seen
+
+        linear = torch.ops.aten.linear.default
+        self.assertIn(linear, annotated_targets(None))
+        decomposed = annotated_targets([])
+        self.assertIn(torch.ops.aten.addmm.default, decomposed)
+        self.assertNotIn(linear, decomposed)
+        kept = annotated_targets([linear])
+        self.assertIn(linear, kept)
+        self.assertNotIn(torch.ops.aten.addmm.default, kept)
+
     @patch("executorch.export.stages.allow_exported_model_train_eval")
     @patch("executorch.export.stages.move_exported_model_to_eval")
     @patch("executorch.export.stages.move_exported_model_to_train")
@@ -780,6 +852,7 @@ def test_run_qat_calls_prepare_qat_pt2e(
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = True
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.train_fn = train_fn
         mock_recipe.pre_prepare_passes = None
         mock_recipe.post_prepare_passes = None
@@ -829,6 +902,7 @@ def test_run_qat_missing_train_fn_raises(self, mock_torch_export: Mock) -> None:
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = True
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.train_fn = None
         mock_recipe.pre_prepare_passes = None
         mock_recipe.post_prepare_passes = None
@@ -870,6 +944,7 @@ def test_run_qat_model_put_in_train_mode_before_export(
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = True
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.train_fn = Mock()
         mock_recipe.pre_prepare_passes = None
         mock_recipe.post_prepare_passes = None
@@ -920,6 +995,7 @@ def test_run_ptq_does_not_call_prepare_qat_pt2e(
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = False
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.calibration_inputs_fn = None
         mock_recipe.pre_prepare_passes = None
         mock_recipe.post_prepare_passes = None
@@ -940,6 +1016,45 @@ def test_run_ptq_does_not_call_prepare_qat_pt2e(
         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_model_put_in_eval_mode_before_export(
+        self,
+        mock_torch_export: Mock,
+        mock_composable_quantizer: Mock,
+        mock_prepare_pt2e: Mock,
+        mock_convert_pt2e: Mock,
+    ) -> None:
+        """PTQ: model.eval() must be called before torch.export.export so that
+        calibration runs batch_norm and dropout as the deployed model does."""
+        mock_recipe = Mock(spec=QuantizationRecipe)
+        mock_recipe.quantizers = [self.create_dummy_quantizer()]
+        mock_recipe.is_qat = False
+        mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
+        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
+
+        self.model.train()
+        training_at_export_time = []
+
+        def capture_training_flag(model, *args, **kwargs):
+            training_at_export_time.append(model.training)
+            return Mock(spec=ExportedProgram)
+
+        mock_torch_export.side_effect = capture_training_flag
+
+        QuantizeStage(mock_recipe).run(
+            PipelineArtifact(data=self.models_dict, context=self.context)
+        )
+
+        self.assertEqual(training_at_export_time, [False])
+
     @patch("executorch.export.stages.move_exported_model_to_eval")
     @patch("executorch.export.stages.convert_pt2e")
     @patch("executorch.export.stages.prepare_pt2e")
@@ -968,6 +1083,7 @@ def pass_fn(m):
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = False
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.calibration_inputs_fn = None
         mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")]
         mock_recipe.post_prepare_passes = [make_pass("post_prepare")]
@@ -1023,6 +1139,7 @@ def pass_fn(m):
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = True
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.train_fn = Mock()
         mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")]
         mock_recipe.post_prepare_passes = [make_pass("post_prepare")]
@@ -1067,6 +1184,7 @@ def test_run_ptq_uses_calibration_inputs_fn_when_provided(
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = False
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.calibration_inputs_fn = calibration_inputs_fn
         mock_recipe.pre_prepare_passes = None
         mock_recipe.post_prepare_passes = None
@@ -1109,6 +1227,7 @@ def test_run_ptq_falls_back_to_example_inputs_when_no_calibration_fn(
         mock_recipe.quantizers = [mock_quantizer]
         mock_recipe.is_qat = False
         mock_recipe.dynamic_batch_size = False
+        mock_recipe.preserve_ops = None
         mock_recipe.calibration_inputs_fn = None
         mock_recipe.pre_prepare_passes = None
         mock_recipe.post_prepare_passes = None
@@ -1331,6 +1450,7 @@ def _make_recipe(is_qat: bool, dynamic_batch_size: bool) -> Mock:
         mock_recipe.quantizers = [Mock(spec=TorchAOPT2EQuantizer)]
         mock_recipe.is_qat = is_qat
         mock_recipe.dynamic_batch_size = dynamic_batch_size
+        mock_recipe.preserve_ops = None
         mock_recipe.calibration_inputs_fn = None
         mock_recipe.train_fn = Mock() if is_qat else None
         mock_recipe.pre_prepare_passes = None