Release intermediate export artifacts once they are consumed (pytorch#21748)

Summary:

Add an opt-in `release_intermediate_artifacts` recipe setting that drops each stage’s output after the next stage consumes it, reducing peak memory usage for large models. Released artifacts produce descriptive accessor errors, and stages remain reusable across subsequent exports. Recipe combination and positional compatibility are preserved.

Authored with Codex.

Reviewed By: psiddh

Differential Revision: D115607218

Pull Request resolved: pytorch#21748

Author: JakeStevens

export/export.py

@@ -9,7 +9,7 @@
 
 import logging
 import time
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import torch
 from executorch.exir import EdgeProgramManager
@@ -235,6 +235,7 @@ def __init__(
         }
 
         self._stage_to_artifacts: Dict[StageType, PipelineArtifact] = {}
+        self._released_stages: Set[StageType] = set()
 
     def _detect_model_type(
         self, model: Union[nn.Module, GraphModule, ExportedProgram, Dict]
@@ -452,8 +453,11 @@ def _run_pipeline(self) -> None:
         # After validation: a rejected pipeline must not destroy the previous
         # run. In place, since get_stage_artifacts() hands out this dict.
         self._stage_to_artifacts.clear()
+        self._released_stages.clear()
 
         current_artifact = PipelineArtifact(data=self._model, context=self._run_context)
+        release = self._export_recipe.release_intermediate_artifacts
+        previous_stage_type = None
 
         # Execute stages from registry in the order specified by pipeline_stages
         for stage_type in self._pipeline_stages:
@@ -473,6 +477,30 @@ def _run_pipeline(self) -> None:
 
             self._stage_to_artifacts[stage_type] = current_artifact
 
+            # The stage that just ran has consumed its input, so the previous
+            # stage's output is now unreachable except through the two
+            # references we hold ourselves.
+            if release and previous_stage_type is not None:
+                self._release_stage(previous_stage_type)
+            previous_stage_type = stage_type
+
+    def _release_stage(self, stage_type: StageType) -> None:
+        self._stage_to_artifacts.pop(stage_type, None)
+        stage = self._stage_registry.get(stage_type)
+        if stage is not None:
+            stage.release_artifact()
+        self._released_stages.add(stage_type)
+        logging.info(f"Released artifact for stage: {stage_type}")
+
+    def _raise_if_released(self, *stage_types: StageType) -> None:
+        released = [s for s in stage_types if s in self._released_stages]
+        if released:
+            names = ", ".join(str(s) for s in released)
+            raise RuntimeError(
+                f"Artifact for {names} was released to free memory. Disable "
+                "release_intermediate_artifacts on the recipe to retain it."
+            )
+
     def export(self) -> None:
         """
         Execute the full export process.
@@ -503,6 +531,7 @@ def get_exported_program(self, method_name: str = "forward") -> ExportedProgram:
             RuntimeError: If torch export stage has not been run
             KeyError: If the method name is not found in exported programs
         """
+        self._raise_if_released(StageType.TORCH_EXPORT)
         artifact = self._stage_to_artifacts.get(StageType.TORCH_EXPORT)
         if artifact is None or artifact.data is None:
             raise RuntimeError(
@@ -535,17 +564,19 @@ def get_edge_program_manager(self) -> "EdgeProgramManager":
             RuntimeError: If no edge stage has been run
         """
         # Check stages in order of preference
-        for stage_type in [
+        edge_stages = [
             StageType.TO_EDGE_TRANSFORM_AND_LOWER,
             StageType.TO_BACKEND,
             StageType.EDGE_PROGRAM_MANAGER_TRANSFORM,
             StageType.TO_EDGE,
-        ]:
+        ]
+        for stage_type in edge_stages:
             artifact = self._stage_to_artifacts.get(stage_type)
             if artifact is not None and artifact.data is not None:
                 logging.info(f"Returning edge program manager from stage {stage_type}")
                 return artifact.data
 
+        self._raise_if_released(*edge_stages)
         raise RuntimeError(
             "Edge program manager is not available. "
             "Run one of the edge stages first: TO_EDGE_TRANSFORM_AND_LOWER, TO_EDGE, EDGE_PROGRAM_MANAGER_TRANSFORM, or TO_BACKEND."
@@ -681,13 +712,16 @@ def print_delegation_info(self) -> None:
             & {StageType.TO_EDGE_TRANSFORM_AND_LOWER, StageType.TO_BACKEND}
         )
         if not lowering_stage:
-            RuntimeError(
+            raise RuntimeError(
                 "No delegation info available, atleast one of the lowering stages should be present"
             )
 
         stage_artifact = self._stage_to_artifacts.get(lowering_stage[0])
         if stage_artifact is None:
-            RuntimeError("No delegation info available, run the lowering stage first")
+            self._raise_if_released(lowering_stage[0])
+            raise RuntimeError(
+                "No delegation info available, run the lowering stage first"
+            )
 
         delegation_info_by_method = stage_artifact.get_context(
             "delegation_info_by_method", None

export/recipe.py

@@ -257,6 +257,7 @@ class _CombineAccumulator:
     mode_values: list = field(default_factory=list)
     pipeline_stages_values: list = field(default_factory=list)
     source_transform_in_place_values: list = field(default_factory=list)
+    release_intermediate_artifacts_values: list = field(default_factory=list)
     backend_config: object = None
     pre_partitioning_callbacks: list = field(default_factory=list)
 
@@ -295,6 +296,12 @@ class ExportRecipe:
         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.
+        release_intermediate_artifacts: Drop each stage's output once the next stage
+                               has consumed it, so only the final artifact stays
+                               resident. Necessary for models whose intermediates do
+                               not all fit in memory at once. Makes
+                               get_exported_program(), get_edge_program_manager(),
+                               and print_delegation_info() unavailable after export.
         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.
@@ -318,6 +325,7 @@ class ExportRecipe:
         List[Callable[[torch.nn.Module], torch.nn.Module]]
     ] = None
     pre_trace_hooks: Optional[List[Callable[[str, torch.nn.Module], None]]] = None
+    release_intermediate_artifacts: bool = False
 
     @classmethod
     def get_recipe(cls, recipe: "RecipeType", **kwargs) -> "ExportRecipe":
@@ -615,6 +623,9 @@ def _collect_recipe_fields(
             acc.source_transform_in_place_values.append(
                 recipe.source_transform_in_place
             )
+            acc.release_intermediate_artifacts_values.append(
+                recipe.release_intermediate_artifacts
+            )
 
             # Use the executorch_backend_config from the first recipe that supplies one.
             if acc.backend_config is None and recipe.executorch_backend_config:
@@ -645,6 +656,10 @@ def _combine_recipes(
         cls._assert_scalar_fields_agree(
             "source_transform_in_place", acc.source_transform_in_place_values
         )
+        cls._assert_scalar_fields_agree(
+            "release_intermediate_artifacts",
+            acc.release_intermediate_artifacts_values,
+        )
 
         combined_quantization_recipe = cls._combine_quantization_recipe(
             is_qat_values=acc.is_qat_values,
@@ -681,6 +696,11 @@ def _combine_recipes(
             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,
+            release_intermediate_artifacts=(
+                acc.release_intermediate_artifacts_values[0]
+                if acc.release_intermediate_artifacts_values
+                else False
+            ),
             lowering_recipe=combined_lowering_recipe,
             executorch_backend_config=acc.backend_config,
             pipeline_stages=shared_pipeline_stages,

export/stages.py

@@ -81,6 +81,7 @@ def __init__(self) -> None:
         Initialize the stage.
         """
         self._artifact = None
+        self._artifact_released = False
 
     @property
     @abstractmethod
@@ -116,9 +117,26 @@ def run(self, artifact: PipelineArtifact) -> None:
         pass
 
     def get_artifacts(self) -> "PipelineArtifact":
-        if self._artifact is None:
-            raise RuntimeError(f"Stage: {self.__class__.__name__} not executed")
-        return self._artifact
+        if self._artifact is not None:
+            self._artifact_released = False
+            return self._artifact
+        if self._artifact_released:
+            raise RuntimeError(
+                f"Stage: {self.__class__.__name__} artifact was released to free "
+                "memory. Disable release_intermediate_artifacts to retain it."
+            )
+        raise RuntimeError(f"Stage: {self.__class__.__name__} not executed")
+
+    def release_artifact(self) -> None:
+        """
+        Drop this stage's reference to its output.
+
+        The pipeline holds the only other reference, so releasing both lets a
+        stage's output be collected once the next stage has consumed it.
+        Subclasses that retain their own references must extend this.
+        """
+        self._artifact = None
+        self._artifact_released = True
 
 
 class TorchExportStage(Stage):
@@ -348,6 +366,7 @@ class ExecutorchStage(Stage):
     """
 
     def __init__(self, backend_config: Any) -> None:
+        super().__init__()
         self._backend_config = backend_config
 
     @property
@@ -398,11 +417,16 @@ def __init__(
             List[Callable[[nn.Module], nn.Module]]
         ] = None,
     ) -> None:
+        super().__init__()
         self._quantization_recipe = quantization_recipe
         self._source_transform_passes = source_transform_passes
         self._in_place = in_place
         self._transformed_models: Dict[str, nn.Module] = {}
 
+    def release_artifact(self) -> None:
+        super().release_artifact()
+        self._transformed_models = {}
+
     @property
     def stage_type(self) -> str:
         return StageType.SOURCE_TRANSFORM
@@ -478,6 +502,7 @@ def __init__(
         quantization_recipe: Optional[QuantizationRecipe],
         pre_trace_hooks: Optional[List[Callable[[str, nn.Module], None]]] = None,
     ) -> None:
+        super().__init__()
         self._quantization_recipe = quantization_recipe
         self._pre_trace_hooks = pre_trace_hooks
 

export/tests/test_export_recipe.py

@@ -348,6 +348,25 @@ def test_new_hooks_do_not_change_existing_positional_arguments(self) -> None:
         self.assertIsNone(recipe.source_transform_passes)
         self.assertIsNone(recipe.pre_trace_hooks)
 
+    def test_release_option_does_not_change_existing_positional_arguments(
+        self,
+    ) -> None:
+        lowering_recipe = Mock()
+        recipe = ExportRecipe(None, None, None, False, lowering_recipe)
+
+        self.assertIs(recipe.lowering_recipe, lowering_recipe)
+        self.assertFalse(recipe.release_intermediate_artifacts)
+
+    def test_combine_preserves_release_intermediate_artifacts(self) -> None:
+        combined = ExportRecipe.combine(
+            [
+                ExportRecipe(release_intermediate_artifacts=True),
+                ExportRecipe(release_intermediate_artifacts=True),
+            ]
+        )
+
+        self.assertTrue(combined.release_intermediate_artifacts)
+
     def test_preserves_source_transform_passes_and_pre_trace_hooks(self) -> None:
         first_source_transform = Mock()
         second_source_transform = Mock()
@@ -432,17 +451,33 @@ def test_conflicting_source_transform_in_place_raises(self) -> None:
             ExportRecipe.combine([r1, r2])
         self.assertIn("source_transform_in_place", str(cm.exception))
 
+    def test_conflicting_release_intermediate_artifacts_raises(self) -> None:
+        r1 = ExportRecipe(name="a", release_intermediate_artifacts=True)
+        r2 = ExportRecipe(name="b", release_intermediate_artifacts=False)
+        with self.assertRaises(ValueError) as cm:
+            ExportRecipe.combine([r1, r2])
+        self.assertIn("release_intermediate_artifacts", 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
+            name="a",
+            strict=False,
+            mode=Mode.DEBUG,
+            source_transform_in_place=True,
+            release_intermediate_artifacts=True,
         )
         r2 = ExportRecipe(
-            name="b", strict=False, mode=Mode.DEBUG, source_transform_in_place=True
+            name="b",
+            strict=False,
+            mode=Mode.DEBUG,
+            source_transform_in_place=True,
+            release_intermediate_artifacts=True,
         )
         result = ExportRecipe.combine([r1, r2])
         self.assertFalse(result.strict)
         self.assertEqual(result.mode, Mode.DEBUG)
         self.assertTrue(result.source_transform_in_place)
+        self.assertTrue(result.release_intermediate_artifacts)
 
     def test_name_is_joined_from_input_recipe_names(self) -> None:
         r1 = ExportRecipe(name="backend_a")

export/tests/test_export_session.py

@@ -111,6 +111,91 @@ def test_default_pipeline_execution_order(self) -> None:
         self.assertEqual(len(session._stage_to_artifacts), 5)
         self.assertEqual(set(session._stage_to_artifacts.keys()), set(stage_types))
 
+    def _run_five_stage_pipeline(self, recipe: ExportRecipe) -> ExportSession:
+        stage_types = [
+            StageType.SOURCE_TRANSFORM,
+            StageType.QUANTIZE,
+            StageType.TORCH_EXPORT,
+            StageType.TO_EDGE_TRANSFORM_AND_LOWER,
+            StageType.TO_EXECUTORCH,
+        ]
+        session = ExportSession(
+            model=self.model,
+            example_inputs=self.example_inputs,
+            export_recipe=recipe,
+        )
+        for stage_type in stage_types:
+            session.register_stage(stage_type, self._create_mock_stage(stage_type))
+        session.export()
+        return session
+
+    def test_intermediate_artifacts_retained_by_default(self) -> None:
+        session = self._run_five_stage_pipeline(ExportRecipe(name="test"))
+        self.assertEqual(len(session._stage_to_artifacts), 5)
+        self.assertEqual(session._released_stages, set())
+
+    def test_release_intermediate_artifacts_keeps_only_the_last(self) -> None:
+        session = self._run_five_stage_pipeline(
+            ExportRecipe(name="test", release_intermediate_artifacts=True)
+        )
+
+        self.assertEqual(
+            set(session._stage_to_artifacts.keys()), {StageType.TO_EXECUTORCH}
+        )
+        # Every stage but the last had both of its references dropped.
+        for stage_type in [
+            StageType.SOURCE_TRANSFORM,
+            StageType.QUANTIZE,
+            StageType.TORCH_EXPORT,
+            StageType.TO_EDGE_TRANSFORM_AND_LOWER,
+        ]:
+            self.assertIn(stage_type, session._released_stages)
+            session.get_registered_stage(
+                stage_type
+            ).release_artifact.assert_called_once()
+
+    def test_released_accessors_explain_why(self) -> None:
+        session = self._run_five_stage_pipeline(
+            ExportRecipe(name="test", release_intermediate_artifacts=True)
+        )
+
+        for accessor in (
+            session.get_exported_program,
+            session.get_edge_program_manager,
+            session.print_delegation_info,
+        ):
+            with self.assertRaises(RuntimeError) as cm:
+                accessor()
+            self.assertIn("released to free memory", str(cm.exception))
+
+    def test_reexport_clears_released_stage_state(self) -> None:
+        session = self._run_five_stage_pipeline(
+            ExportRecipe(name="test", release_intermediate_artifacts=True)
+        )
+        self.assertIn(StageType.TORCH_EXPORT, session._released_stages)
+
+        exported_program = Mock()
+        torch_export_stage = session.get_registered_stage(StageType.TORCH_EXPORT)
+        assert torch_export_stage is not None
+        torch_export_stage.get_artifacts.return_value.data = {
+            "forward": exported_program
+        }
+        session._pipeline_stages = [
+            StageType.SOURCE_TRANSFORM,
+            StageType.QUANTIZE,
+            StageType.TORCH_EXPORT,
+        ]
+
+        session.export()
+
+        self.assertIs(session.get_exported_program(), exported_program)
+
+    def test_release_does_not_touch_the_final_artifact(self) -> None:
+        session = self._run_five_stage_pipeline(
+            ExportRecipe(name="test", release_intermediate_artifacts=True)
+        )
+        self.assertIsNotNone(session.get_executorch_program_manager())
+
     def test_overriden_pipeline_execution_order(self) -> None:
         # Test when pipeline stages that are passed through recipe
         stage_types = [