[RFC][Short] More on ExportRecipes

[abhinaykukkadapu]

This discussion outlines few proposals regarding the ExportRecipes (currently experimental). These requirements have emerged from various discussions, and I want to consolidate them here for further discussion and prioritization:

1. API Naming

There is potential confusion between torch.export.export() and executorch.export(), especially since the latter encompasses additional functionalities such as quantization, lowering, and serialization. A previously proposed solution is to rename executorch.export(..) -> executorch.compile(..). This change could enhance clarity, but it may also be unnecessary. I would like to gather your thoughts on whether a renaming is warranted.

2. Support for Additional Backend Recipes

Support additional backend recipes

• XNNPack (supported)
• CoreML
• QNN

3. QAT Quantization support

Currently, our recipes support Post-Training Quantization (PTQ). Should we extend this support to include Quantization-Aware Training (QAT)? Additionally, which backends currently support QAT, or we should at least transparently error out if QAT based recipe is passed?

4. Differentiating Calibration Inputs vs Example Inputs

The proposal is to add a new parameter, calibration_inputs, to the API alongside example_inputs. This parameter would explicitly indicate to users that the data provided in calibration_inputs is used for quantization purposes, while example_inputs are solely for export.

5. Support for ExportedProgram

At present, the API requires an nn.Module and assumes that the export stage is always run within the scope of the recipe. The proposal is to decouple this to accept either an ExportedProgram or an nn.Module. This change would help use cases such as HuggingFace Optimum, where the ExportRecipe can be utilized for post export stages such as graph quantization, lowering. This would provide greater flexibility for models to use torch.export.export(...) in a manner that suits their needs but use recipes for just quantization and lowering.

6. Support to_edge -> to_backend -> to_executorch along with to_edge_transform_and_lower

Currently, the executorch/export supports only the to_edge_transform_and_lower API. Is there a need to support to_edge() -> to_backend() -> to_executorch()?

7. Modularization of ExportRecipe

Split current ExportRecipe into piece-meal recipes such asExportRecipe, QuantizationRecipe, LoweringRecipe and have a root recipe for example: CompileRecipe which takes all the above. This might require re-writing most of the current ExportRecipe component, so not sure if this is needed, we may just get away with doing just (5) - Supporting ExportedProgram which skips torch.export.export() and SourceTransform stages which might be just sufficient.

CC: @cbilgin, @mergennachin, @kimishpatel, @metascroy, @JacobSzwejbka, @guangy10, @digantdesai, @mcr229, @jackzhxng, @JakeStevens

[GregoryComer]

1. API Naming
   I'm definitely in favor of renaming. I think executorch.compile is clearer than executorch.export.

2. Support for Additional Backend Recipes
   If we can support XNNPACK, Core ML, Vulkan, and QNN, that will cover pretty much all mobile users.

4. Differentiating Calibration Inputs vs Example Inputs
   Seems reasonable to me, with one caveat. Dynamic quantization through pt2e seems to require at least one calibration sample, despite not actually needing the activation ranges and such. So for dynamic quant recipes, we'll probably want to fall back to running the post-pt2e quantized graph with at least the export example args.

5. Support for ExportedProgram
   This makes sense to me, with the caveat that it might limit certain quantization approach or pre-export manipulation. But I think, in general, being able to take an exported program would be nice.

[abhinaykukkadapu]

Seems reasonable to me, with one caveat. Dynamic quantization through pt2e seems to require at least one calibration sample

Thanks for the feedback, i thought we only need calibration for static quantization, Is it ok if we take example_input for that one calibration input that dynamic quantization expects?

This makes sense to me, with the caveat that it might limit certain quantization approach or pre-export manipulation. But I think, in general, being able to take an exported program would be nice.

If someone provides us a ExportedProgram it is already past the export stage and we wouldn't have to worry about pre-export source quantization?

[GregoryComer]

Yeah. On the dynamic quant, this is from memory so might be worth double checking. But I think it still uses the calibration run for some internal pt2e state, but not to actually calibrate observers. It can fail weirdly if there isn't a calibration run. But should be easy to handle in the recipe, so probably not as relevant to the RFC's questions.

On the ExportedProgram inputs, it sounds good to me. I was mainly thinking that we might not be able to accept it in certain recipes.

[JacobSzwejbka]

I think EP embedds the inputs it was exported on somewhere in the meta data so you could probably just use that to make some inputs for dynamic quant

[larryliu0820]

5. Support for ExportedProgram

Do we have a proper way to check which stage the ExportedProgram is? I'm a bit nervous when the API is taking arbitrary ExportedProgram but it may be in later stage of lowering. Like if we already run to_edge() with some config and pass the exported program to executorch.compile, and inside we are running to_edge() again with a different config, would that be a problem?

[abhinaykukkadapu]

Good point, can we limit to only run to_edge_transform_and_lower when an ExportedProgram is provided, if there is no way to tell what stage that ExportedProgram came from?

[jackzhxng]

Can't we simply just start the pipeline at the post export stage if the type of the input is ExportedProgram?

[abhinaykukkadapu]

I think Mengwei's point is that ExportedProgram may come after torch.export() or to_edge() or to_backend() and if we always run all the stages in recipe we might override the configs those were run with in the first place initially.

@larryliu0820 can we distinguish from the op's IR or i see dialect field but from my experimentation it just stays as "EDGE" all through the way, any other ideas?

[JacobSzwejbka]

@abhinaykukkadapu is it still EDGE immediately after export?

[abhinaykukkadapu]

@JacobSzwejbka dialect is "TRAINING" after torch.export.export(), but it stayed EDGE after doing to_edge, but my point is what if the exported program is passed after doing to_backend()?

[larryliu0820]

6. Support to_edge -> to_backend -> to_executorch along with to_edge_transform_and_lower

Yes please

[abhinaykukkadapu]

Can you elaborate on the reasons why would we want this support, would be great to learn status-quo on these API usages.

[JakeStevens]

Can you elaborate on the reasons why would we want this support, would be great to learn status-quo on these API usages.

to_edge -> to_backend -> to_executorch is used in Cadence backend. It may also be used on cortex, not sure that was decided yet. These are op library based, instead of having a true delegate.

• 1 to supporting both

[larryliu0820]

7. Modularization of ExportRecipe

Split current ExportRecipe into piece-meal recipes such asExportRecipe, QuantizationRecipe, LoweringRecipe

If you are designing this, please also look at export llm configs https://github.com/pytorch/executorch/blob/main/extension/llm/export/config/llm_config.py#L283

[abhinaykukkadapu]

Just took a brief look, it seems we are doing lots of customizations before and for each stage in the pipeline for etLLM, as the original motivation for recipes is to enable non-power users, is it worth including this config?

[jackzhxng]

I feel like this may be overkill to do this, but depends on the final implementation I suppose. I think I like 5 better

[Gasoonjia]

executorch.export() renaming

I do not have strong opinion on that. The meaning of export should rely on the envrionment it stays. If it living in torch.fx/torch.export, then torch.export.export should generate torch.export.ExportedProgram, and if it lives in executorch, generating executorch program makes sense.
Even if we rename to other name, there's still possible to be misunderstand by users. e.g if we called executorch.compile, there's still torch.compile in pytorch/pytorch space, which has diferent behavior with us.

Beyond than your proposals, I would love to have the native support for devtools, particularly the etrecord generation. Currently we hide all details under executorch.export which means it is impossible for executorch.export users to fetch essential information for profiling and debugging outside executorch.export. It will be great if we can make executorch.export support for etrecord generation. All updates we need for api interface is adding a "etrecord_path: Optional[str]" atrtibute.

[abhinaykukkadapu]

Agreed, it would be great to have devtools integrated, let me add this as one of the dimension. Thanks!

[jackzhxng]

Why do people think compile is better? I'm failing to see the link between the word compile and what's happening during this export process. I personally think export is fine as long as we do the documentation right and things are consistent

[abhinaykukkadapu]

Yeah the votes are divided right now and there is definitely not a strong opinion from anyone that we need to change it and as you and @Gasoonjia pointed out even the name compile has contender from torch.compile

[JakeStevens]

(2) the experimental cortex-m backend would be a great candidate as well. This will be favored more by embedded MCU developers, who may not have a ton of ML background and something more "plug and play" would be good

[JakeStevens]

I think (3) and (5) are intertwined. For QAT, export_for_training is used which returns an ExportedProgram. It would then be something like

prepare_for_qat() -> QAT performed outside of ET -> executorch.export/compile(qat_exported_program)

And so you need to take in an ExportedProgram

[guangy10]

5. Support for ExportedProgram
   At present, the API requires an nn.Module and assumes that the export stage is always run within the scope of the recipe. The proposal is to decouple this to accept either an ExportedProgram or an nn.Module. This change would help use cases such as HuggingFace Optimum, where the ExportRecipe can be utilized for post export stages such as graph quantization, lowering.

Instead of accepting an ExportedProgram, we can still accept a nn.Module for HuggingFace models (all we have enabled in Optimum-ExecuTorch). The way it works is that the model is already wrapped into a TorchExportableModuleForX (For decoder-only LLM, here is the code pointer https://github.com/huggingface/transformers/blob/v4.53.2/src/transformers/integrations/executorch.py#L30), the exportability for whatever models are using is covered in HuggingFace's CI.

In general I think the recipe should just take a nn.Module as input and the exportability of it is model's owners responsibility. The owner could be us, HuggingFace, or private models from end users. The recipe API can cover all these use-cases consistently.

[abhinaykukkadapu]

@guangy10 thanks for the feedback, are you saying that we should run model.export() if it exists in the module, within our recipe flow?