Add source transform hooks to ExportRecipe (pytorch#21747) (pytorch#21747)

Summary:

Add recipe-configurable source transforms and per-method pre-trace
hooks. Source transforms run once per distinct model, including shared
models, while pre-trace hooks run immediately before each method is
exported. Recipe combination preserves hooks and in-place behavior
without breaking existing positional arguments.

Differential Revision: D115602123

Author: JakeStevens

export/export.py

@@ -5,6 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from __future__ import annotations
+
 import logging
 import time
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -306,18 +308,21 @@ def _build_stages(self, stages: List[StageType]) -> Dict[StageType, Stage]:
             if stage_type == StageType.SOURCE_TRANSFORM:
                 stage = SourceTransformStage(
                     self._quant_recipe,
+                    source_transform_passes=list(
+                        self._export_recipe.source_transform_passes or ()
+                    ),
                     in_place=self._export_recipe.source_transform_in_place,
                 )
             elif stage_type == StageType.QUANTIZE:
-                stage = QuantizeStage(self._quant_recipe)
+                stage = QuantizeStage(
+                    self._quant_recipe,
+                    pre_trace_hooks=list(self._export_recipe.pre_trace_hooks or ()),
+                )
             elif stage_type == StageType.TORCH_EXPORT:
-                aten_transform_passes = None
-                if self._export_recipe.aten_transform_passes is not None:
-                    aten_transform_passes = list(
-                        self._export_recipe.aten_transform_passes
-                    )
                 stage = TorchExportStage(
-                    aten_transform_passes, strict=self._export_recipe.strict
+                    list(self._export_recipe.aten_transform_passes or ()),
+                    strict=self._export_recipe.strict,
+                    pre_trace_hooks=list(self._export_recipe.pre_trace_hooks or ()),
                 )
             elif stage_type == StageType.TO_EDGE_TRANSFORM_AND_LOWER:
                 stage = EdgeTransformAndLowerStage.from_recipe(self._lowering_recipe)

export/recipe.py

@@ -6,6 +6,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from __future__ import annotations
+
 import copy
 import dataclasses
 import logging
@@ -239,6 +241,8 @@ class _CombineAccumulator:
     quantizers: list = field(default_factory=list)
     ao_quantization_configs: list = field(default_factory=list)
     pre_edge_passes: list = field(default_factory=list)
+    source_transform_passes: list = field(default_factory=list)
+    pre_trace_hooks: list = field(default_factory=list)
     edge_transform_passes: list = field(default_factory=list)
     edge_manager_transform_passes: list = field(default_factory=list)
     pre_prepare_passes: list = field(default_factory=list)
@@ -273,6 +277,21 @@ class ExportRecipe:
         quantization_recipe: Optional quantization recipe for model quantization
         aten_transform_passes: Optional list of functions to apply transformation passes to the program before edge lowering.
                                These callables are invoked to modify and return the transformed program.
+        source_transform_passes: Optional list of nn.Module transforms applied once
+                               during the SOURCE_TRANSFORM stage, before quantization.
+                               Each is applied once per distinct model object, so a
+                               model shared by several methods is transformed once.
+                               Each pass must return the nn.Module to use, even when
+                               it mutates its input in place.
+        pre_trace_hooks: Optional list of (method_name, model) callables invoked
+                               before each method's first trace: in QUANTIZE for PT2E,
+                               otherwise in TORCH_EXPORT.
+                               Hooks may only change Python configuration, such as
+                               kernel selection, never tensor state or module structure.
+                               Tensor storage remains shared. Configuration persists
+                               on shared models, so hooks must overwrite the settings
+                               they own for each method. Return values are ignored.
+                               Hooks are not repeated on converted PT2E graphs.
         source_transform_in_place: Skip the defensive deepcopy in the SOURCE_TRANSFORM
                                stage and mutate the caller's model. Necessary for models
                                large enough that a second copy will not fit in memory.
@@ -295,6 +314,10 @@ class ExportRecipe:
     pipeline_stages: Optional[List[StageType]] = None
     mode: Mode = Mode.RELEASE
     strict: bool = True
+    source_transform_passes: Optional[
+        List[Callable[[torch.nn.Module], torch.nn.Module]]
+    ] = None
+    pre_trace_hooks: Optional[List[Callable[[str, torch.nn.Module], None]]] = None
 
     @classmethod
     def get_recipe(cls, recipe: "RecipeType", **kwargs) -> "ExportRecipe":
@@ -573,6 +596,10 @@ def _collect_recipe_fields(
         for recipe in backend_recipes:
             if recipe.aten_transform_passes:
                 acc.pre_edge_passes.extend(recipe.aten_transform_passes)
+            if recipe.source_transform_passes:
+                acc.source_transform_passes.extend(recipe.source_transform_passes)
+            if recipe.pre_trace_hooks:
+                acc.pre_trace_hooks.extend(recipe.pre_trace_hooks)
 
             if lr := recipe.lowering_recipe:
                 cls._collect_lowering_fields(acc, lr)
@@ -652,6 +679,8 @@ def _combine_recipes(
             name=recipe_name,
             quantization_recipe=combined_quantization_recipe,
             aten_transform_passes=acc.pre_edge_passes or None,
+            source_transform_passes=acc.source_transform_passes or None,
+            pre_trace_hooks=acc.pre_trace_hooks or None,
             lowering_recipe=combined_lowering_recipe,
             executorch_backend_config=acc.backend_config,
             pipeline_stages=shared_pipeline_stages,

export/stages.py

@@ -6,6 +6,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from __future__ import annotations
+
 import copy
 import logging
 from abc import ABC, abstractmethod
@@ -42,7 +44,7 @@
 
 
 def _drop_empty(
-    passes_by_method: Dict[str, List[PassType]]
+    passes_by_method: Dict[str, List[PassType]],
 ) -> Dict[str, List[PassType]]:
     return {method: p for method, p in passes_by_method.items() if p}
 
@@ -130,10 +132,12 @@ def __init__(
             List[Callable[[str, ExportedProgram], ExportedProgram]]
         ] = None,
         strict=True,
+        pre_trace_hooks: Optional[List[Callable[[str, nn.Module], None]]] = None,
     ) -> None:
         super().__init__()
         self._aten_transform_passes = aten_transform_passes
         self.strict = strict
+        self._pre_trace_hooks = pre_trace_hooks
 
     @property
     def stage_type(self) -> str:
@@ -151,6 +155,11 @@ def run(self, artifact: PipelineArtifact) -> None:
         models = artifact.data
         example_inputs = artifact.get_context("example_inputs")
         dynamic_shapes = artifact.get_context("dynamic_shapes", {})
+        pre_trace_hooks = (
+            None
+            if artifact.get_context("pre_trace_hooks_applied", False)
+            else self._pre_trace_hooks
+        )
 
         exported_programs = {}
 
@@ -163,6 +172,9 @@ def run(self, artifact: PipelineArtifact) -> None:
 
                 method_dynamic_shapes = dynamic_shapes.get(method_name)
 
+                for hook in pre_trace_hooks or []:
+                    hook(method_name, model)
+
                 # Export the model
                 exported_programs[method_name] = torch.export.export(
                     model,
@@ -373,15 +385,21 @@ def run(self, artifact: PipelineArtifact) -> None:
 
 class SourceTransformStage(Stage):
     """
-    Optional stage: Source transform stage: Apply source transformations to the model.
+    Apply source transformations and optional AO quantization to the model.
+
+    Each source transform must return the nn.Module to use.
     """
 
     def __init__(
         self,
         quantization_recipe: Optional[QuantizationRecipe],
         in_place: bool = False,
+        source_transform_passes: Optional[
+            List[Callable[[nn.Module], nn.Module]]
+        ] = None,
     ) -> None:
         self._quantization_recipe = quantization_recipe
+        self._source_transform_passes = source_transform_passes
         self._in_place = in_place
         self._transformed_models: Dict[str, nn.Module] = {}
 
@@ -401,46 +419,67 @@ def run(self, artifact: PipelineArtifact) -> None:
         """
         Apply source transformations to the model.
         """
-        if (
-            not self._quantization_recipe
-            or not self._quantization_recipe.ao_quantization_configs
-        ):
+        ao_configs = (
+            self._quantization_recipe.ao_quantization_configs
+            if self._quantization_recipe
+            else None
+        )
+        if not ao_configs and not self._source_transform_passes:
             logging.info(
-                "Quantization recipe is invalid to run SourceTransform, returning original artifact"
+                "No source transform passes and no AO quantization configs for SourceTransformStage, returning original artifact"
             )
             self._artifact = artifact
             return
 
         assert isinstance(artifact.data, dict)
 
+        if ao_configs and len(ao_configs) > 1:
+            raise ValueError(
+                "AO quantization configs cannot be reliably composed together, multiple quantization configs are disallowed for source transform at this point"
+            )
+
         # A second copy of the model is not affordable for every caller, so
         # large models can opt out and have their own model mutated instead.
-        self._transformed_models = (
-            artifact.data if self._in_place else copy.deepcopy(artifact.data)
-        )
-
-        # Apply torchao quantize_ to each model
-        for _, model in self._transformed_models.items():
-            # pyre-ignore
-            if len(self._quantization_recipe.ao_quantization_configs) > 1:
-                raise ValueError(
-                    "AO quantization configs cannot be reliably composed together, multiple quantization configs are disallowed for source transform at this point"
-                )
-
-            ao_config = self._quantization_recipe.ao_quantization_configs[0]
-            quantize_(model, ao_config.ao_base_config, ao_config.filter_fn)
-            unwrap_tensor_subclass(model)
+        models = artifact.data if self._in_place else copy.deepcopy(artifact.data)
+
+        # Several methods commonly share one model object. Transform each
+        # distinct object once, or a shared model is quantized once per method.
+        transformed_by_id: Dict[int, nn.Module] = {}
+        for method_name, model in list(models.items()):
+            original_id = id(model)
+            if original_id not in transformed_by_id:
+                for transform in self._source_transform_passes or []:
+                    if not callable(transform):
+                        raise TypeError("Source transform passes must be callable")
+                    model = transform(model)
+                    if not isinstance(model, nn.Module):
+                        raise TypeError(
+                            f"Source transform {transform!r} must return an nn.Module, "
+                            f"got {type(model).__name__}"
+                        )
+                if ao_configs:
+                    ao_config = ao_configs[0]
+                    quantize_(model, ao_config.ao_base_config, ao_config.filter_fn)
+                    unwrap_tensor_subclass(model)
+                transformed_by_id[original_id] = model
+            models[method_name] = transformed_by_id[original_id]
 
-        self._artifact = artifact.copy_with_new_data(self._transformed_models)
+        self._transformed_models = models
+        self._artifact = artifact.copy_with_new_data(models)
 
 
 class QuantizeStage(Stage):
     """
     Optional stage: Perform post-training quantization on the model.
     """
 
-    def __init__(self, quantization_recipe: Optional[QuantizationRecipe]) -> None:
+    def __init__(
+        self,
+        quantization_recipe: Optional[QuantizationRecipe],
+        pre_trace_hooks: Optional[List[Callable[[str, nn.Module], None]]] = None,
+    ) -> None:
         self._quantization_recipe = quantization_recipe
+        self._pre_trace_hooks = pre_trace_hooks
 
     @property
     def stage_type(self) -> str:
@@ -539,6 +578,10 @@ def run(self, artifact: PipelineArtifact) -> None:
             if recipe.is_qat:
                 model.train()
 
+            with torch.no_grad():
+                for hook in self._pre_trace_hooks or []:
+                    hook(method_name, model)
+
             captured_graph = torch.export.export(
                 model, inputs, dynamic_shapes=export_dynamic_shapes, strict=True
             ).module()
@@ -596,6 +639,9 @@ def run(self, artifact: PipelineArtifact) -> None:
             quantized_models[method_name] = quantized_model
 
         self._artifact = artifact.copy_with_new_data(quantized_models)
+        self._artifact.add_context(
+            "pre_trace_hooks_applied", bool(self._pre_trace_hooks)
+        )
 
 
 class ToEdgeStage(Stage):

export/tests/test_export_recipe.py

@@ -341,6 +341,44 @@ def test_combine_keeps_edge_transform_passes(self) -> None:
             combined.lowering_recipe.edge_transform_passes, [first, second]
         )
 
+    def test_new_hooks_do_not_change_existing_positional_arguments(self) -> None:
+        recipe = ExportRecipe(None, None, None, True)
+
+        self.assertTrue(recipe.source_transform_in_place)
+        self.assertIsNone(recipe.source_transform_passes)
+        self.assertIsNone(recipe.pre_trace_hooks)
+
+    def test_preserves_source_transform_passes_and_pre_trace_hooks(self) -> None:
+        first_source_transform = Mock()
+        second_source_transform = Mock()
+        first_pre_trace_hook = Mock()
+        second_pre_trace_hook = Mock()
+
+        combined = ExportRecipe.combine(
+            [
+                ExportRecipe(
+                    source_transform_passes=[first_source_transform],
+                    pre_trace_hooks=[first_pre_trace_hook],
+                    source_transform_in_place=True,
+                ),
+                ExportRecipe(
+                    source_transform_passes=[second_source_transform],
+                    pre_trace_hooks=[second_pre_trace_hook],
+                    source_transform_in_place=True,
+                ),
+            ]
+        )
+
+        self.assertEqual(
+            combined.source_transform_passes,
+            [first_source_transform, second_source_transform],
+        )
+        self.assertEqual(
+            combined.pre_trace_hooks,
+            [first_pre_trace_hook, second_pre_trace_hook],
+        )
+        self.assertTrue(combined.source_transform_in_place)
+
 
 # ---------------------------------------------------------------------------
 # Helpers shared by combine-recipe tests