Recipe and Input class definitions with e2e export #10034

tarun292 · 2025-04-10T00:14:10Z

Based on the discussion in #9027

This PR adds the executorch.export API and all the supporting components required for it. At a high level the executorch.export API takes in a model, example inputs and a recipe, then underneath the hood executes all the steps required to export, quantize and lower the model based on the recipe.

The pipeline consists of a staged setup where each major step in the process such as Export, Quantization etc. is listed as a separate stage and a chain of these is formed and then executed. The result of this will be that we'll get a PTE file which can then be executed on device.

The major new components added in this PR are:

class definition for ExportSession and ExportRecipe
Definitions for each stage in the process
executorch.export API which will return a session object that the user can then use to get access to the PTE file, run the model via pybindings, print delegation info etc.

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

pytorch-bot · 2025-04-10T00:14:14Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10034

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Large queue time for macos-m2-15 instances

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2025-04-10T00:14:59Z

This pull request was exported from Phabricator. Differential Revision: D71946730

digantdesai

Looks good at a high level. Left some comments.

digantdesai · 2025-04-11T23:24:59Z

export/_export.py

+    return manager
+
+
+class ExportSession:


This is a good idea, I like it. But just a pet peeve, let's refactor this before starting, and organize the logic so that it can scale better, else I fear it will end up becoming a giant bowl of spaghetti.

Suggestion, make a class Stage(ABC) and derive others from it like Quantize(Stage) and Export(Stage) and then Session just manages the pipeline with Stages.

See -

executorch/backends/xnnpack/test/tester/tester.py

Line 424 in 409447d

class Tester:

digantdesai · 2025-04-11T23:25:33Z

export/_recipe.py

+
+    name: Optional[str] = None
+    quantizer: Optional[Quantizer] = None
+    edge_compile_config: Optional[EdgeCompileConfig] = ( # pyre-ignore[11]: Type not defined


digantdesai · 2025-04-11T23:25:55Z

export/_recipe.py

+    )
+    mode: Mode = Mode.RELEASE
+
+    def get_quantizer(self) -> Optional[Quantizer]:


nit property?

digantdesai · 2025-04-11T23:28:56Z

export/_export.py

+from ._recipe import ExportRecipe
+
+
+def export(


nit, WDYT about? I understand we want to do s/torch.export/executorch.export but it feels a bit confusing given we are returning something not analogous to the model which torch.export does.

Suggested change

def export(

def create_export_session(

# or

Session.__init__

# or

@classmethod

Session.create()

digantdesai · 2025-04-11T23:29:36Z

export/_export.py

+    model: Union[nn.Module, Dict[str, nn.Module]],
+    example_inputs: Union[List[tuple[torch.Tensor, ...]], Dict[str, List[tuple[torch.Tensor, ...]]]],
+    export_recipe: ExportRecipe,
+    name: Optional[str] = None,


what is the rationale for this?

digantdesai · 2025-04-11T23:30:55Z

export/_export.py

+        dynamic_shapes: Optional dynamic shape specifications
+        constant_methods: Optional dictionary of constant methods
+        artifact_dir: Optional directory to store artifacts
+        apply_quantization: Whether to apply quantization during export, defaults to False


this feels a bit confusing

digantdesai · 2025-04-11T23:31:46Z

export/_export.py

+    Returns:
+        A configured ExportSession instance with the export process completed if requested
+    """
+    manager = ExportSession(


nit

Suggested change

manager = ExportSession(

session = ExportSession(

digantdesai · 2025-04-11T23:32:57Z

export/_export.py

+    example_inputs: Union[List[tuple[torch.Tensor, ...]], Dict[str, List[tuple[torch.Tensor, ...]]]],
+    export_recipe: ExportRecipe,
+    name: Optional[str] = None,
+    dynamic_shapes: Optional[Union[Any, Dict[str, Any]]] = None,


why is this not on Session.export()?

mergennachin · 2025-05-01T16:28:21Z

export/_export.py

+
+from ._recipe import ExportRecipe
+
+


We should annotate with executorch.exir._warnings.experimental decorators

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Pull Request resolved: #10034 ghstack-source-id: 282298048 @exported-using-ghexport Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/)

facebook-github-bot · 2025-05-06T16:23:40Z

This pull request was exported from Phabricator. Differential Revision: D71946730

jackzhxng · 2025-05-06T18:04:27Z

export/export.py

+from .recipe import ExportRecipe
+
+
+class Stage(ABC):


I like this 👍🏻

jackzhxng · 2025-05-06T18:06:03Z

export/export.py

+        pass
+
+    @abstractmethod
+    def get_outputs(self) -> Any:


Suggested change

def get_outputs(self) -> Any:

def get_artifacts(self) -> Any:

Think this conveys the intent of this function better when the user hasn't read the docstring

jackzhxng · 2025-05-06T18:07:31Z

export/export.py

+
+    @property
+    def name(self) -> str:
+        return "quantize"


Would be worth to comment somewhere that quantization may also happen in SourceTransformStage

jackzhxng · 2025-05-06T18:12:34Z

export/export.py

+            )
+        return self._executorch_program_manager.buffer
+
+    def get_example_input(


I feel like this and run_method seem a bit confusing to have in the ExportSession?

The way i'm thinking is that export session lets the users do almost all the things they'd like to with an executorch model like load the model and run via pybindings, get reference inputs etc. and that's why I added these helper utils.

jackzhxng · 2025-05-06T18:16:16Z

export/export.py

+        elif stage_name == "executorch":
+            self._executorch_program_manager = result
+
+    def _run_pipeline(self, start_stage: str) -> None:


Why do we need this start_stage stuff, can we just not take any parameters and just run the all stages in self._stages?

Yeah I agree. We should not have this start_stage thing. We should just construct pipeline that can just be run with each stage in sequence. Construction of the pipeline should take care of apprpriate deps

Addressed in latest version.

jackzhxng · 2025-05-06T18:17:16Z

export/export.py

+        return self._executorch_program_manager
+
+
+class SourceTransformStage(Stage):


NIT: have this and QuantizeStage declared in the order of the export process, this one declared at the top

kimishpatel

Left quite a few comments

kimishpatel · 2025-05-07T00:01:49Z

export/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.


The directory structure is bit confusing. People coming to ET already know export and now we are teaching them of a different export.

Maybe executorch/lower or executorch/pipeline?

Thought quite a bit about this, the main reasoning to call this export too is because we're consuming the export call within the API. One of the main points of this API is that if someone wants to target ExecuTorch they don't even need to learn about torch.export (although many of them probably will).

kimishpatel · 2025-05-07T03:12:39Z

export/export.py

+    ) -> None:
+        self._exported_program: Dict[str, ExportedProgram] = {}
+        self._pre_edge_transform_passes = pre_edge_transform_passes
+        self._model_dict: Dict[str, nn.Module] = {}


is this to capture more than one Module to be exported? why do we need this?

Yep a user could possibly pass in multiple methods to be exported and wanted to have support for that from the beginning. Something we struggled with in ModAI by trying to add it later.

kimishpatel · 2025-05-07T03:13:32Z

export/export.py

+        self._model_dict: Dict[str, nn.Module] = {}
+        self._example_inputs_dict: Dict[str, List[tuple[torch.Tensor, ...]]] = {}
+        self._dynamic_shapes_dict: Dict[str, Any] = {}


If these are one to one correspondence than I would combine them into a single struct that can be passed in the init or separately initialized

I did think about that a bit, the thing is that even when doing this via a struct/dataclass the code tends to be equally cumbersome and this is more direct.

kimishpatel · 2025-05-07T03:16:27Z

export/export.py

+            self._dynamic_shapes_dict = export_config.get("dynamic_shapes", {})
+
+        # Check if we need to do export (if _exported_program is empty)
+        if not self._exported_program:


What if you want to rexport? For different input shapes or export args? If ExportStange is strictly for one export session and to export with different args you need a new ExportStage than I will make it immutable.

This wasn't needed, removed it. One export session is intended to be only for one set of input args.

kimishpatel · 2025-05-07T03:17:54Z

export/export.py

+        self._partitioners = partitioners
+        self._transform_passes = transform_passes


what is the relation between transformation passes and partitioners?

Transform passes are what we run after to_edge e.g. edge_manager.transform(), and partitioners are the ones we use in to_backend.

kimishpatel · 2025-05-07T03:44:34Z

export/export.py

+
+        # If quantization is available, add it after source transform
+        if "quantize" in self._stages:
+            self._pipeline.insert(0, "quantize")


As I said elsewhere I think quantize should not be responsible for running export. We should run export, run quantizer, prepare, convert, and run export again. This call can be managed within here by detecting the need to run export again if we have pt2e quantize requested

We need to store this because it's used in all the later stages such as to_edge etc.

kimishpatel · 2025-05-07T03:47:03Z

export/export.py

+        # Return the outputs
+        return self._stages[stage_name].get_outputs()
+
+    def _get_stage_inputs(self, stage_name: str) -> tuple[Any, Dict[str, Any]]:


Why do we need to do this here instead of just initializing each stage with the inputs they need. This way call to run is simply run?

I thought a lot about this too, in order to deal with quantize_ and later on other things such as custom op replacement in eager mode we do need to do it and it's quite a bit of complexity that'll be better hidden from the users.

deal with quantize_ and later on other things such as custom op replacement in eager mode

Again eager mode should not belong to this. It is coupling unrelated things together.

kimishpatel · 2025-05-07T03:48:47Z

export/export.py

+        for i in range(start_index, len(self._pipeline)):
+            stage_name = self._pipeline[i]
+            self._current_stage_index = i
+
+            # Get the primary input and configuration parameters for this stage
+            primary_input, config_params = self._get_stage_inputs(stage_name)
+
+            # Run the stage
+            result = self._run_stage(stage_name, primary_input, config_params)
+
+            # Store the result
+            self._store_stage_result(stage_name, result)


Yeah at thsi point i expect jsut iterate through stages and run without needing to figure out input and storage stage outputs. Outputs of each stage should be there in get_output/artifact if needed

This means that the quantization recipe doesn't necessarily have configs to use with quantize_ but might still have pt2e quantization. In that case it doesn't make sense to throw an error right?

Not sure I follow what you mean.

kimishpatel · 2025-05-07T03:49:36Z

export/export.py

+        elif stage_name == "executorch":
+            self._executorch_program_manager = result
+
+    def _run_pipeline(self, start_stage: str) -> None:


Yeah I agree. We should not have this start_stage thing. We should just construct pipeline that can just be run with each stage in sequence. Construction of the pipeline should take care of apprpriate deps

kimishpatel · 2025-05-07T03:51:05Z

export/export.py

+        # The original code expects this to be a tuple of tensors
+        return self._example_inputs[method_name][0]
+
+    def run_method(


why is this part of this class

I think that's a valid point, but i've felt like having all the quantization configuration related items in one class would generally be easier for recipe authors. We when creating the pipeline decide what source transforms need to be run and what is in pt2e. This isn't set in stone though, we'll definitely iterate on this as we add more recipes in XNNPack, CoreML etc. and we can arrive at a different structure over the next few weeks if the current one isn't flexible enough.

kimishpatel

I didnt "request changes" if you feel strongly about landing, but I think we should consider addressing these comments

jackzhxng · 2025-05-13T00:15:47Z

export/export.py

+    def run(
+        self,
+        models: Dict[str, Any],
+        export_config: Optional[Dict[str, Any]] = None,


Looking at this PR again, why isn't this a part of the ExportRecipe?

This contains example inputs and dynamic shapes which shouldn't be a part of the recipe.

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

tarun292 · 2025-05-14T03:30:06Z

@kimishpatel @jackzhxng thanks for the detailed review, really appreciate it! Addressed all of your comments:

In the latest version the major changes i made are:

Removed start_stage API and string input for stage name, rather just have the run_pipeline function now which runs through the entire pipeline. Changed quite a bit how the pipeline is constructed now.
The pipeline is a list of stages now and not a dictionary.
Export happens independently of quantization now.
Added get_artifact which gets the output of each stage.

digantdesai · 2025-05-14T05:31:34Z

Thanks @tarun292 - love the pipeline design.

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) ghstack-source-id: 276661724 Pull Request resolved: pytorch/executorch#10034

Pull Request resolved: pytorch/executorch#10034 ghstack-source-id: 284081137 @exported-using-ghexport Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/)

Recipe and Input class definitions with e2e export

e42c965

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

tarun292 requested review from JacobSzwejbka, lucylq and swolchok as code owners April 10, 2025 00:14

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 10, 2025

This was referenced Apr 10, 2025

Add some basic xnnpack recipes #10035

Open

Adding qualcomm recipes #10036

Open

Add coreml recipes #10037

Open

facebook-github-bot added the fb-exported label Apr 10, 2025

digantdesai approved these changes Apr 11, 2025

View reviewed changes

mergennachin reviewed May 1, 2025

View reviewed changes

tarun292 added 2 commits May 6, 2025 09:23

Update base for Update on "Recipe and Input class definitions with e2…

a9d80f1

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Update on "Recipe and Input class definitions with e2e export"

bba10a0

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

jackzhxng approved these changes May 6, 2025

View reviewed changes

kimishpatel reviewed May 7, 2025

View reviewed changes

jackzhxng reviewed May 13, 2025

View reviewed changes

tarun292 added 2 commits May 13, 2025 20:12

Update base for Update on "Recipe and Input class definitions with e2…

9964bf2

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Update on "Recipe and Input class definitions with e2e export"

7a9043e

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

tarun292 added the release notes: exir Changes to any dialects and passes on these dialects, such as memory planning label May 14, 2025

tarun292 added 4 commits May 14, 2025 14:51

Update base for Update on "Recipe and Input class definitions with e2…

e3076b7

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Update on "Recipe and Input class definitions with e2e export"

d565702

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Update base for Update on "Recipe and Input class definitions with e2…

edc5f11

…e export" Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

Update on "Recipe and Input class definitions with e2e export"

36c0298

Differential Revision: [D71946730](https://our.internmc.facebook.com/intern/diff/D71946730/) [ghstack-poisoned]

tarun292 changed the base branch from gh/tarun292/4/base to main May 15, 2025 00:00

Merge branch 'main' into gh/tarun292/4/head

2a65fab

tarun292 merged commit 1b593ad into main May 15, 2025
84 of 85 checks passed

tarun292 deleted the gh/tarun292/4/head branch May 15, 2025 00:48

-def export(
+def create_export_session(
+# or
+Session.__init__
+# or
+@classmethod
+Session.create()

	def get_outputs(self) -> Any:
	def get_artifacts(self) -> Any:

		return self._executorch_program_manager


		class SourceTransformStage(Stage):

		@@ -0,0 +1,24 @@
		# Copyright (c) Meta Platforms, Inc. and affiliates.

		self._partitioners = partitioners
		self._transform_passes = transform_passes

Recipe and Input class definitions with e2e export #10034

Recipe and Input class definitions with e2e export #10034

Uh oh!

Conversation

tarun292 commented Apr 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Apr 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10034

❗ 1 Active SEVs

Uh oh!

facebook-github-bot commented Apr 10, 2025

Uh oh!

digantdesai left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented May 6, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kimishpatel left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

tarun292 commented Apr 10, 2025 •

edited

Loading

pytorch-bot bot commented Apr 10, 2025 •

edited

Loading