Avoid deep-copying the model during source transform (pytorch#21581)

Differential Revision: D114773308

Pull Request resolved: pytorch#21581

Author: JakeStevens

export/export.py

@@ -300,7 +300,10 @@ def _build_stages(self, stages: List[StageType]) -> Dict[StageType, Stage]:
         stage = None
         for stage_type in stages or self._get_default_pipeline():
             if stage_type == StageType.SOURCE_TRANSFORM:
-                stage = SourceTransformStage(self._quant_recipe)
+                stage = SourceTransformStage(
+                    self._quant_recipe,
+                    in_place=self._export_recipe.source_transform_in_place,
+                )
             elif stage_type == StageType.QUANTIZE:
                 stage = QuantizeStage(self._quant_recipe)
             elif stage_type == StageType.TORCH_EXPORT:

export/recipe.py

@@ -154,6 +154,9 @@ 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_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.
         lowering_recipe: Optional lowering recipe for model lowering and partitioning
         executorch_backend_config: Optional backend configuration for ExecuTorch
         pipeline_stages: Optional list of stages to execute, defaults to a standard pipeline.
@@ -166,6 +169,7 @@ class ExportRecipe:
     aten_transform_passes: Optional[
         List[Callable[[str, ExportedProgram], ExportedProgram]]
     ] = None
+    source_transform_in_place: bool = False
     lowering_recipe: Optional[LoweringRecipe] = None
     # pyre-ignore[11]: Type not defined
     executorch_backend_config: Optional[ExecutorchBackendConfig] = None

export/stages.py

@@ -315,8 +315,13 @@ class SourceTransformStage(Stage):
     Optional stage: Source transform stage: Apply source transformations to the model.
     """
 
-    def __init__(self, quantization_recipe: Optional[QuantizationRecipe]) -> None:
+    def __init__(
+        self,
+        quantization_recipe: Optional[QuantizationRecipe],
+        in_place: bool = False,
+    ) -> None:
         self._quantization_recipe = quantization_recipe
+        self._in_place = in_place
         self._transformed_models: Dict[str, nn.Module] = {}
 
     @property
@@ -347,8 +352,11 @@ def run(self, artifact: PipelineArtifact) -> None:
 
         assert isinstance(artifact.data, dict)
 
-        # Store the original models
-        self._transformed_models = copy.deepcopy(artifact.data)
+        # 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():

export/tests/test_export_stages.py

@@ -321,6 +321,25 @@ def test_run_with_ao_quantization_configs(
         # verify the result model is NOT the same object as the original
         self.assertIsNot(result_artifact.data["forward"], self.model)
 
+    @patch("executorch.export.stages.quantize_")
+    @patch("executorch.export.stages.unwrap_tensor_subclass")
+    def test_run_in_place_does_not_copy_the_model(
+        self, mock_unwrap: Mock, mock_quantize: Mock
+    ) -> None:
+        from torchao.core.config import AOBaseConfig
+
+        mock_recipe = Mock(spec=QuantizationRecipe)
+        mock_recipe.ao_quantization_configs = [
+            AOQuantizationConfig(ao_base_config=Mock(spec=AOBaseConfig))
+        ]
+
+        stage = SourceTransformStage(mock_recipe, in_place=True)
+        stage.run(PipelineArtifact(data=self.models_dict, context={}))
+
+        # The caller's own model is quantized rather than a copy of it.
+        self.assertIs(mock_quantize.call_args[0][0], self.model)
+        self.assertIs(stage.get_artifacts().data["forward"], self.model)
+
 
 class TestQuantizeStage(unittest.TestCase):
     def setUp(self) -> None: