Introduce FXGraphExtractor into torch.onnx.dynamo_export #99940

thiagocrepaldi · 2023-04-24T22:50:55Z

The current API architecture can be seen as 3 independent exporters as shown below. The public API dynamo_export() defaults to one of the 3 variants and the other 2 must be used by instantiating private classes:

This PR refactors the API in a way that dynamo_export is the only way to use the ONNX exporter. It defaults to a FX tracer based on torch.export, but an internal-only idiom allows switching the FX tracer (aka FXGraphExtractor interface), as shown below:

Summary of changes:

Unifies all exporter variants under a single dynamo_export API
- ResolvedExportOptions was expanded to allow fx_tracer: FXGraphExtractor to be specified, selecting which FX graph extractor to use, according to the design proposal
- As a consequence, torch.onnx._internal.exporter.Exporter does not have to internally specialize for each type of FX API that the exporter might be used. This leads to a single Exporter with many FX graph extractors
- Before in red, after in green:
Input processing was moved from Exporter subclasses to FXGraphExtractor subclasses, where they are actually consumed
- Exporter is a [data]class that holds export options, model and input data in a single cohesive object. Specializing it means create different exporters instead of having one exporter capable of exporting models through different options.
- Exporter doesn't consume the model_args that caused it to specialize
Improved the circular dependency story.
- Delay torch.onnx import to after all dynamo [sub]components #99070 moves import torch.onnx to after all dynamo subcomponents, preventing torch.onnx to have circular depemndencies when torch.XXXX is imported during initialization
- There are other points we need to improve in subsequent PRs. APIs are organized in a way that it is easy to "import too much"
Refactored decomposition_table as an internal-only ResolvedExportOptions property.
- Similar to input processing, this helper is not actually consumed at tyhe Exporter layer. This PR moves it to the layer in which it is used
Demoted Exporter.model_signature to a simple standalone helper
- There is no need to have this as a exporter method; this is a standard inpect.signature usage without any state

Possible next steps are:

Decouple passes and dispatching from the cluttered export_fx_to_onnx
Further integration with http://github.com/pytorch/pytorch/pull/98421/ into FXGraphExtractor public API + helper for unit testing
- Some passes are changing input processing, which are not captured by the proposed input adapter

** COPILOT SUMMARY**

`🤖 Generated by Copilot at bdaba31`

Summary

📝🚀🔧

This pull request refactors the ONNX exporter code to use the FX graph representation and the new io_adapter module for input and output adaptation. It also adds support for custom export options and flattening HuggingFace model outputs in the ONNX test framework. It updates the ONNX dynamo exporter API tests and adds a new module torch/onnx/_internal/fx/dynamo_graph_extractor.py for exporting FX models to ONNX with dynamo support. It fixes some type annotations, imports, and formatting issues in the ONNX exporter code.

The ONNX exporter got a new look
With FX graph and dynamo hook
It uses io_adapter
And custom options matter
For HuggingFace models and model_signature book

Walkthrough

Move the fx submodule from torch/onnx/_internal to torch/onnx/_internal/fx, and rename some of its modules ( link, link, link, link)
Add a new module torch/onnx/_internal/fx/dynamo_graph_extractor.py that defines a DynamoExport class for generating FX graphs using the torch._dynamo.export API (link)
Add a new module torch/onnx/_internal/fx/io_adapter.py that defines the input and output adapter classes and steps for the ONNX exporter, and a helper function to wrap models with output adapters (link, link, link, link)
Update the ResolvedExportOptions class in torch/onnx/_internal/exporter.py to inherit from the ExportOptions class, and to set the fx_tracer and decomposition_table attributes based on the dynamo_graph_extractor and function_dispatcher modules (link, link)
Update the Exporter class in torch/onnx/_internal/exporter.py to remove the export method and add a new abstract generate_fx method, and to use the fx_tracer attribute to generate and export the FX graph (link, link)
Update the FXSymbolicTraceExporter class in torch/onnx/_internal/fx/fx_symbolic_graph_extractor.py to be renamed to FXSymbolicTracer, and to inherit from exporter.FXGraphExtractor and implement the generate_fx method (link, link)
Update the export_fx_to_onnx method of the FXSymbolicTracer class to be renamed to _export_fx_to_onnx, and to be moved to the exporter.FXGraphExtractor class (link)
Update the dynamo_export function in torch/onnx/_internal/exporter.py to accept and return ResolvedExportOptions and Exporter objects, respectively (link)
Update the run_test_with_fx_to_onnx_exporter_and_onnx_runtime function in test/onnx/onnx_test_common.py to add a new parameter export_options for passing custom export options to the torch.onnx.dynamo_export function (link, link)
Update the test_log_sigmoid and _test_large_scale_exporter tests in test/onnx/test_fx_to_onnx_with_onnxruntime.py to use the updated run_test_with_fx_to_onnx_exporter_and_onnx_runtime function and the torch.onnx.dynamo_export function (link, link, link, link)
Update the test_raise_on_invalid_save_argument_type test in test/onnx/dynamo/test_exporter_api.py to use the io_adapter.InputAdapter and io_adapter.OutputAdapter classes instead of the exporter.InputAdapter and exporter.OutputAdapter classes (link)
Move the model_signature property from the Exporter class in torch/onnx/_internal/exporter.py to a standalone function in torch/onnx/utils.py, and update the references to it (link, link, link)
Move the UnsatisfiedDependencyError class from the Exporter class in torch/onnx/_internal/exporter.py to the top level of the module, and update the references to it (link)
Rename the _create_onnx_friendly_decomposition_table function and the _ONNX_FRIENDLY_DECOMPOSITION_TABLE dictionary in torch/onnx/_internal/fx/function_dispatcher.py to _create_default_onnx_decomposition_table and _DEFAULT_ONNX_EXPORTER_DECOMPOSITION_TABLE, respectively, and update the references to them (link, link)
Update the imports in torch/onnx/_internal/fx/function_dispatcher.py to use the torch._ops and torch._decomp modules instead of the torch.ops and torch.decomp modules, and to use aliases for accessing the onnxscript.function_libs.torch_aten.ops and torch._ops modules (link, link, link, link, link, link, link)
Update the ExportOutput class in torch/onnx/_internal/exporter.py to use the InputAdapter and OutputAdapter classes from io_adapter instead of the ones defined in the same module (link)
Update the type annotations in torch/onnx/_internal/fx/serialization.py and torch/onnx/_internal/exporter.py to fix some inconsistencies (link, link, link, link)
Remove an unused import of inspect from torch/onnx/_internal/exporter.py (link)
Remove an unused import of torch._dynamo from torch/onnx/_internal/fx/passes/shape_inference.py (link)
Add a comment to torch/onnx/_internal/fx/passes/shape_inference.py to explain why the import of torch._dynamo is done inside the _run method of the ShapeInferenceWithFakeTensor class (link)
Fix a typo in the docstring of the _module_expansion_symbolic_trace function in torch/onnx/_internal/fx/fx_symbolic_graph_extractor.py (link)
Add an empty line to torch/onnx/__init__.py for formatting purposes (link)
Delete the torch/onnx/_internal/fx/__init__.py file (link)

Fixes #ISSUE_NUMBER

cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

pytorch-bot · 2023-04-24T22:50:58Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

⏳ No Failures, 1 Pending

As of commit 640a5c7:
💚 Looks good so far! There are no failures yet. 💚

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

BowenBao · 2023-04-25T22:42:45Z

torch/onnx/_internal/exporter.py

+
+    # Private only attributes
+    decomposition_table: Dict[torch._ops.OpOverload, Callable]
+    """Dictionary with all ATen operators supported by the exporter."""


docstring not accurate, more like

Dictionary with ATen/Prims operators and their decompositions. Exporter runs decomposition based on this table after getting the FX graph from FXGraphExtractor. If not provided, the default decomposition table is based on the full PyTorch decomposition table, excluding Prims decomposition, and any ATen operator decomposition if the ATen operator has a matching OnnxFunction in torch_lib.

And based on the level of internal details, not sure why it needs to be exposed in this PR. Unless it is needed urgently for Prims experiments. cc @justinchuby. Ideally a much more natural api emerges after the dispatcher refactor.

This is not publicly exposed; ResolvedExportOptions is internal-only. This serves as an internal extension point and it is a natural Dynamo concept.

Using docstring from https://github.com/thiagocrepaldi/pytorch/blob/c8bf6f3c699bfd9833e9b9b65a1c39878f1ad440/torch/_dynamo/eval_frame.py#L705

"""A dictionary that maps operators to their decomposition functions."""

This is similar to the docstring for backend of DynamoOptimize that we discussed earlier. Both need code rework to match the same concept of its original fx/dynamo counterpart, which is what stated in the docstring. I won't block on this, will revisit in dispatcher theme.

I don't get why """A dictionary that maps operators to their decomposition functions.""" doesn't match the same concept.

I understand we are not giving a full description on how it is used, but it is still a truthful explanation; we do build a decomposition_table (with whatever our registry can support) and pass it to FX APIs o consume it, as they document it they should be.

Fair enough. I felt we are adding it before designing the purpose and the expected outcome under exporter context, which is what dispatcher design going to do anyway.

got your point.
not all changes here are final because the rest of the code is not final.
next redesigns may refactor some of the things done on previous refactoring until everything make sense :)

torch/onnx/_internal/exporter.py

BowenBao · 2023-04-25T22:51:54Z

torch/onnx/_internal/exporter.py

+        self.output_adapter: OutputAdapter = OutputAdapter()
+
+    @_beartype.beartype
+    def _export_fx_to_onnx(


This feels weird to be here. This class is FXGraphExtractor, but this method finishes everything off with a onnx ModelProto. Feels this method belongs to Exporter (or DynamoExporter(Exporter) for the sake of import dep problems).

To be honest, there is no right place on the repo for _export_fx_to_onnx. Next round of design can split it in several methods in their correct places, as we discussed offline.

it is the "fx passes" + "conversion to onnxscript" part. Again as discussed my hunch is they should be the general part of exporter. I'm just surprised that it is moved from DynamoExporter to Extractor. So now all the work e2e is completed in Extractor, that feels like just a rename of the previous FXExporter though. Anyway won't block on this too, as you mentioned there will be follow up designs.

class Exporter as proposed by Aaron is just a thread-safety container to keep user model+data+export options together so that we could run exporter in parallel (in theory). It was never meant to have hold code. The fact we did specialize and put a lot of helpers on it is the actual infraction we need to address.

_export_fx_to_onnx does too much (passes, graph building and protobuf serialization) and not even Exporter is the right place to it. In next steps, we will be able to break all three tasks in 3 different blocks, and neither Exporter nor FXGraphExtractor will need to hold any of it.

I agree these are 3 separate tasks. In my understanding the Exporter, instead of FXGraphExtractor, should drive these components.

Yeah, agreed. Exporter must drive the 3 classes that are to be born.
_export_fx_to_onnx will be gone and not inside fx_tracer very soon

torch/onnx/_internal/exporter.py

justinchuby · 2023-04-26T01:52:27Z

torch/onnx/_internal/exporter.py

@@ -96,6 +115,16 @@ def resolve(value: Optional[T], fallback: Union[T, Callable[[], T]]) -> T:

        self.opset_version = resolve(options.opset_version, _DEFAULT_OPSET_VERSION)
        self.dynamic_shapes = resolve(options.dynamic_shapes, False)
+        # TODO: Prevent circular dep + [onnx]script dep
+        import torch.onnx._internal.fx.dynamo_graph_extractor as dynamo_graph_extractor
+        from torch.onnx._internal.fx import (  # TODO: onnxscript is not available on Pytorch CI yet


We can maybe also delay the table evaluation?

ResolvedExportOptions is only evaluated when an export has started. This is as delayed as it can be.
decomposition_table is not evaluated when we do import torch or import torch.onnx, etc.

Sorry I meant the DEAFAULT_xxx table so as to avoid requiring onnxscript when importing these lines? Just a thought. May not be a good idea.

Yeah, that is something I also feel needs another round of refactoring. We shouldn't have code being called on global scope as decomp.py does.

This can be tackled along with the passes redesign, as decomposition is done as a pass today

torch/onnx/_internal/exporter.py

torch/onnx/_internal/fx/function_dispatcher.py

The current API architecture can be seen as 3 independent exporters as shown below. The public API `dynamo_export()` defaults to one of the 3 variants and the other 2 must be used by instantiating private classes: ![image](https://user-images.githubusercontent.com/5469809/231567368-ec899718-b7c1-4e59-b6a8-383142df245a.png) This PR refactors the API in a way that `dynamo_export` is the only way to use the ONNX exporter. A ``FXGraphExtractor`` layer was created to abstract all the FX extraction from a model out of the exporter. It defaults to a FX tracer based on ``torch.export``, but an internal-only idiom allows switching the FX tracer (aka `FXGraphExtractor` interface), as shown below: ![image](https://user-images.githubusercontent.com/5469809/231567495-3936362d-06de-4cfc-b752-6c2060701c08.png) Summary of changes: * Unifies all exporter variants under a single `dynamo_export` API * `ResolvedExportOptions` was expanded to allow `fx_tracer: FXGraphExtractor` to be specified, selecting which FX graph extractor to use, according to the design proposal * As a consequence, `torch.onnx._internal.exporter.Exporter` does not have to *internally* specialize for each type of FX API that the exporter might be used. This leads to a single `Exporter` with many `FX graph extractors` * Before in red, after in green: ![image](https://user-images.githubusercontent.com/5469809/232633531-4c67449b-4863-474d-9e18-78fc1d31b1bd.png) * `ResolvedExportOptions` was expanded to allow `decomposition_table` to be specified as an internal-only property. * Similar to input processing, this helper is not actually consumed at tyhe `Exporter` layer. This PR moves it to the layer in which it is used * Demoted `Exporter.model_signature` to a simple standalone helper * There is no need to have this as a exporter method; this is a standard `inpect.signature` usage without any state Possible next steps are: * Decouple `passes` and `dispatching` from the cluttered `export_fx_to_onnx` * Further integration with http://github.com/pytorch/pytorch/pull/98421/ into `FXGraphExtractor` public API + helper for unit testing * Some passes are changing input processing, which are not captured by the proposed input adapter ** COPILOT SUMMARY** copilot:all

BowenBao

LGTM, with separate follow-ups for dispatcher(decomposition_table) and _export_fx_to_onnx next.

thiagocrepaldi · 2023-04-26T17:34:44Z

@pytorchbot merge

pytorchmergebot · 2023-04-26T17:36:37Z

Merge failed

Reason: Approval needed from one of the following:
drdarshan, simpkins, nanoax, plahera, xianxl, ...

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

torch/_dynamo/eval_frame.py

Signed-off-by: Thiago Crepaldi <[email protected]>

thiagocrepaldi · 2023-04-26T22:02:00Z

@pytorchbot merge

pytorchmergebot · 2023-04-26T22:04:01Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Apr 24, 2023

pytorchbot added the open source label Apr 24, 2023

thiagocrepaldi force-pushed the thiagofc/thiagofc/introduce-fxgraphextractor-core branch from bdaba31 to 4d9c912 Compare April 25, 2023 01:05

github-actions bot added ciflow/inductor module: dynamo labels Apr 25, 2023

thiagocrepaldi force-pushed the thiagofc/thiagofc/introduce-fxgraphextractor-core branch from 3f92a81 to db4e8b9 Compare April 25, 2023 01:46

thiagocrepaldi marked this pull request as ready for review April 25, 2023 01:46

thiagocrepaldi requested review from BowenBao and abock as code owners April 25, 2023 01:46

thiagocrepaldi mentioned this pull request Apr 25, 2023

Introduce FXGraphExtractor into torch.onnx.dynamo_export #98893

Closed

thiagocrepaldi added module: onnx Related to torch.onnx enhancement Not as big of a feature, but technically not a bug. Should be easy to fix labels Apr 25, 2023

thiagocrepaldi requested a review from titaiwangms April 25, 2023 19:46

thiagocrepaldi assigned BowenBao Apr 25, 2023

BowenBao reviewed Apr 25, 2023

View reviewed changes

justinchuby reviewed Apr 26, 2023

View reviewed changes

torch/onnx/_internal/exporter.py Show resolved Hide resolved

justinchuby reviewed Apr 26, 2023

View reviewed changes

torch/onnx/_internal/exporter.py Show resolved Hide resolved

justinchuby reviewed Apr 26, 2023

View reviewed changes

torch/onnx/_internal/exporter.py Outdated Show resolved Hide resolved

justinchuby reviewed Apr 26, 2023

View reviewed changes

torch/onnx/_internal/fx/function_dispatcher.py Outdated Show resolved Hide resolved

thiagocrepaldi force-pushed the thiagofc/thiagofc/introduce-fxgraphextractor-core branch 2 times, most recently from c8bf6f3 to 10d5d75 Compare April 26, 2023 15:01

thiagocrepaldi requested review from BowenBao and justinchuby April 26, 2023 15:16

BowenBao approved these changes Apr 26, 2023

View reviewed changes

thiagocrepaldi force-pushed the thiagofc/thiagofc/introduce-fxgraphextractor-core branch from 10d5d75 to 7059a17 Compare April 26, 2023 16:34

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 26, 2023

pytorchmergebot added the merging label Apr 26, 2023

BowenBao reviewed Apr 26, 2023

View reviewed changes

torch/_dynamo/eval_frame.py Show resolved Hide resolved

thiagocrepaldi added the onnx-needs-import This PR is related to ONNX, but touches files outside of merge rule patterns, and hence needs import label Apr 26, 2023

jansel approved these changes Apr 26, 2023

View reviewed changes

thiagocrepaldi force-pushed the thiagofc/thiagofc/introduce-fxgraphextractor-core branch from 7059a17 to 9fed5a4 Compare April 26, 2023 19:57

Addressing comments

640a5c7

Signed-off-by: Thiago Crepaldi <[email protected]>

thiagocrepaldi force-pushed the thiagofc/thiagofc/introduce-fxgraphextractor-core branch from 9fed5a4 to 640a5c7 Compare April 26, 2023 22:01

pytorchmergebot added Merged and removed merging labels Apr 27, 2023

pytorchmergebot closed this in 32a67e4 Apr 27, 2023

thiagocrepaldi deleted the thiagofc/thiagofc/introduce-fxgraphextractor-core branch October 31, 2023 19:28

Introduce FXGraphExtractor into torch.onnx.dynamo_export #99940

Introduce FXGraphExtractor into torch.onnx.dynamo_export #99940

Uh oh!

Conversation

thiagocrepaldi commented Apr 24, 2023 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🤖 Generated by Copilot at bdaba31

Summary

Walkthrough

Uh oh!

pytorch-bot bot commented Apr 24, 2023 • 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/99940

⏳ No Failures, 1 Pending

Uh oh!

BowenBao Apr 25, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thiagocrepaldi Apr 26, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thiagocrepaldi Apr 26, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thiagocrepaldi Apr 26, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

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!

Uh oh!

Uh oh!

Uh oh!

BowenBao left a comment

Choose a reason for hiding this comment

Uh oh!

thiagocrepaldi commented Apr 26, 2023

Uh oh!

pytorchmergebot commented Apr 26, 2023

Merge failed

Uh oh!

Uh oh!

thiagocrepaldi commented Apr 26, 2023

Uh oh!

pytorchmergebot commented Apr 26, 2023

Merge started

Uh oh!

Uh oh!

thiagocrepaldi commented Apr 24, 2023 •

edited by pytorch-bot bot

Loading

`🤖 Generated by Copilot at bdaba31`

pytorch-bot bot commented Apr 24, 2023 •

edited

Loading

BowenBao Apr 25, 2023 •

edited

Loading

thiagocrepaldi Apr 26, 2023 •

edited

Loading

thiagocrepaldi Apr 26, 2023 •

edited

Loading

thiagocrepaldi Apr 26, 2023 •

edited

Loading