Introduce FXGraphExtractor into torch.onnx.dynamo_export #98893

thiagocrepaldi · 2023-04-12T00:37:43Z

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 to a single public API that can use different implementations of a FXGraphExtractor interface, as shown below:

Summary of changes:

Unifies all exporter variants under a single dynamo_export API
- ExportOptions 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"
Promoted decomposition_table as a ExportOptions.
- Currently, we have a method on Exporter to initialize such list, not allowing customization from users
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
Integrate [ONNX] Introduce Input/Ouptut adapter; Switch to 'DynamoExporter' #98421 into FXGraphExtractor public API + helper for unit testing

** COPILOT SUMMARY**

`🤖 Generated by Copilot at 1be2507`

Summary

📝🛠️🚀

This pull request refactors the ONNX exporter code to use the FX graph extractor engine and the new io_adapter module, which improve the modularity, readability, and performance of the export process. It also updates the test files and the utils.py file to reflect the changes in the exporter API and the export options. Additionally, it fixes some import issues, typos, and style inconsistencies.

To export models to ONNX with FX
We refactored the code and fixed the syntax
We added new classes and options
And improved the tests and docstrings
And avoided some imports that made conflicts

Walkthrough

Rename and refactor the exporter module and its subclasses to use the FXGraphExtractor abstract class and the io_adapter module for input and output adaptation (link,link,link,link,link,link,link,link,link,link,link,link,link,link,link,link,link,link,link)
Add the decomposition_table attribute and parameter to the ExportOptions and _ResolvedExportOptions classes, which allow customizing the decomposition of ATen operators into ONNX-friendly subgraphs (link,link,link,link,link)
Add the skip_fx_tracer decorator to the pytorch_test_common module, which allows skipping exporting tests for selected FX tracers based on a mapping from FX tracer class to skip reason (link,link)
Add the fx_tracer attribute and parameter to the TestFxToOnnxWithOnnxRuntime class and its subclasses, which allow parameterizing the FX tracer class and running tests with different FX graph extractors (link,link,link,link,link,link,link,link)
Add the concrete_args attribute and parameter to the FXSymbolicTracer class, which allow partially specializing the model inputs and removing control flow or data structures during tracing (link,link)
Add the model_signature function to the onnx_utils module and the utils module, which returns the signature of a PyTorch model or function (link,link)
Rename the FXSymbolicTraceExporter class to FXSymbolicTracer, which is a more descriptive name for the subclass of FXGraphExtractor that uses the torch.fx.symbolic_trace API to generate FX graphs (link)
Rename the FXGraphModuleExporter class to IOAdapter, which is a more descriptive name for the class that adapts the PyTorch model inputs and outputs to the exported ONNX model inputs and outputs format (link)
Rename the _ONNX_FRIENDLY_DECOMPOSITION_TABLE variable to _DEFAULT_ONNX_EXPORTER_DECOMPOSITION_TABLE, which is a more descriptive name for the subset of PyTorch's built-in aten-to-aten decomposition that is compatible with ONNX export (link)
Rename the torch/onnx/_internal/fx/fx_exporter.py file to torch/onnx/_internal/fx/io_adapter.py, which is a more descriptive name for the file that contains the classes for input and output adaptation (link)
Move the InputAdaptStep, InputAdapter, OutputAdaptStep, and OutputAdapter classes from the exporter module to the io_adapter module, which is a more logical place for them (link,link)
Move the DynamoOptimize and DynamoExport classes from the dynamo_exporter module to the dynamo_graph_extractor module, which is a more logical place for them (link)
Move the FXGraphModuleExporter class from the fx_exporter module to the io_adapter module, which is a more logical place for it (link)
Move the export_fx_to_onnx method from the FXSymbolicTraceExporter class to the FXGraphExtractor abstract class, which is a common logic that can be shared by different FX graph extractors (link)
Move the model_signature function from the fx_serialization module to the onnx_utils module, which is a more logical place for it (link)
Delete the torch/onnx/_internal/fx/__init__.py file, which is not needed (link)
Delete the torch/onnx/_internal/fx/dynamo_exporter.py file, which is not needed (link)
Use the io_adapter module instead of the exporter module to access the InputAdapter and OutputAdapter classes in various places, which reflects the refactoring of the input and output adapter logic (link,link,link)
Use the fx_context module instead of the fx module to access the FxToOnnxContext class in the _test_large_scale_exporter function, which is a more specific import that avoids importing the whole fx module (link)
Use the fx_serialization module instead of the fx module to access the save_model_with_external_data function in the _test_large_scale_exporter function, which is a more specific import that avoids importing the whole fx module (link)
Use the torch_ops module instead of the torch module to access the OpOverload and OpOverloadPacket classes in various places, which are more specific imports that avoid importing the whole torch module (link,link,link,link,link)
Use the export_options attribute instead of creating a new ExportOptions object in various places, which avoids duplicating the export options and uses the user-provided values (link,link)
Use the _ResolvedExportOptions class instead of the ResolvedExportOptions class in various places, which reflects the renaming and subclassing of the export options class (link,link,link,link,link,link)
Use the adapt_input and adapt_output methods instead of the _apply_input_adapt_step and _apply_output_adapt_step methods in various places, which reflect the renaming and refactoring of the input and output adaptation logic (link,link)
Use the FXSymbolicTracer class instead of the DynamoExporter class in the _run_test_with_fx_to_onnx_exporter_and_onnx_runtime function, which reflects the parameterization of the FX tracer class (link)
Use the FXGraphExtractor class instead of the Exporter class in the dynamo_export function, which reflects the renaming and refactoring of the exporter class (link)
Use the FXGraphExtractor class instead of the Exporter class in the export_options parameter of the __init__ method of the _ResolvedExportOptions class, which reflects the renaming and refactoring of the exporter class (link)
Use the FXGraphExtractor class instead of the Exporter class in the fx_tracer attribute and parameter of the TestFxToOnnxWithOnnxRuntime class and its subclasses, which reflects the renaming and refactoring of the exporter class (link,link,link,link,link,link)
Use the generate_fx method instead of the export method in various places, which reflects the renaming and refactoring of the FX graph extractor interface (link,link)
Use the skip_fx_tracer decorator instead of the xfail decorator in various places, which allows skipping tests only for selected FX tracers and providing reasons for skipping (link,link)
Use the skip_fx_tracer decorator to skip tests for the DynamoExport and DynamoOptimize FX tracers in various places, and provide the reasons for skipping (link,link,link,link,link,link,link)
Use the # type: ignore[arg-type] comment to suppress the type-checking error for the opset argument of the onnxscript.script decorator in various places, which is a false positive due to the dynamic nature of the onnxscript module (link,link)
Use a string annotation for the onnx_model parameter of the save_model_with_external_data function, which avoids importing onnx at the top level and causing conflicts with the user-installed onnx package (link)
Add a blank line to the import statement in torch/onnx/__init__.py, to follow the PEP 8 style guide for imports (link)
Remove the unused inspect module from the import statement in torch/onnx/_internal/exporter.py (link)
Remove the unused List annotation and add the Dict annotation from the import statement in torch/onnx/_internal/exporter.py (link)
Remove the unused torch module and add the Type annotation from the import statement in torch/onnx/_internal/exporter.py (link)
Remove the unused torch module from the import statement in torch/onnx/_internal/exporter.py (link)
Remove the unused ops module from the import statement in test/onnx/test_fx_to_onnx_with_onnxruntime.py (link)
Remove the model_signature property from the Exporter class, which is not used by the new FXGraphExtractor class (link)
Remove the logger property from the Exporter class, which is not used by the new dynamo_export function (link)
Remove the UnsatisfiedDependencyError class, which is not raised by the new dynamo_export function (link)
Remove the export abstract method from the Exporter class, which is replaced by the generate_fx abstract method in the FXGraphExtractor class (link)
Remove the export_fx_to_onnx method from the FXSymbolicTraceExporter class, which is moved to the FXGraphExtractor abstract class (link)
Remove the import torch._dynamo statement from the top level of torch/onnx/_internal/fx/passes/shape_inference.py, which could cause circular imports and unnecessary initialization of the dynamo module (link)
Fix a typo in the docstring of the _module_expansion_symbolic_trace function, which uses the torch.fx.symbolic_trace API to trace a model with module expansion (link)
Add the class name of the fx_tracer attribute to the file name of the ONNX model in the tearDown function of the TestFxToOnnxWithOnnxRuntime class, to distinguish the models exported by different FX tracers (link)
Add the opset argument to the onnxscript.script decorator in various places, which specifies the ONNX opset version to use for exporting the decorated functions (link,link)
Add the import torch._dynamo statement to the _run function of the ShapeInferenceWithFakeTensor class, which ensures that torch._dynamo is only imported when needed (link)
Add the fx module to the import statement in test/onnx/test_fx_to_onnx_with_onnxruntime.py, which contains the submodules for FX graph extraction, context, and serialization (link)
Add the exporter module to the import statement in test/onnx/pytorch_test_common.py, which contains the FXGraphExtractor abstract class that is used as a base class for the fx_tracer attribute of the TestFxToOnnxWithOnnxRuntime class (link)
Add the functools module and the Union annotation to the import statement in torch/onnx/_internal/fx/fx_symbolic_graph_extractor.py, which are used for implementing the FXSymbolicTracer class (link)
Add the onnx_utils module to the import statement in torch/onnx/_internal/fx/fx_symbolic_graph_extractor.py, which contains the model_signature function that is used for binding the model inputs with default values (link)
Add the Protocol and runtime_checkable annotations to the import statement in torch/onnx/_internal/fx/io_adapter.py, which are used for defining structural subtyping protocols that can be checked at runtime (link)
Add the Type annotation to the import statement in test/onnx/pytorch_test_common.py, which is used for type-checking the fx_tracer attribute of the TestFxToOnnxWithOnnxRuntime class (link)

pytorch-bot · 2023-04-12T00:37:46Z

🔗 Helpful Links

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

📄 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.

❗ 1 Active SEVs

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

ROCm jobs are failing with No space left on device

❌ 9 Failures

As of commit 19edea8:

BROKEN TRUNK - The following jobs failed but were present on the merge base d56adb1:

👉 Rebase onto the `viable/strict` branch to avoid these failures

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

BowenBao · 2023-04-13T20:56:50Z

I prefer landing #98421 first, since it's in a much more ready state. Please let me know if you have more concerns.

We'd need a discussion session for many points in this one. Aaron and I were talking about how we should expose Exporter/FXGraphExtractor the other day. One possible follow-up planned for #98421 is to actually delete DynamoOptimizeExporter (it is kept in history in working state after #98421 is landed should we revisit), since there really isn't anything that is not coverred by DynamoExporter. That leaves us with only FXSymbolicTraceExporter, which is itself a temorary solution and should go away if #95900 is done. Hence the design may not need to expose that much in public api, while retaining internal experiment flexibility.

thiagocrepaldi · 2023-04-13T22:00:14Z

I prefer landing #98421 first, since it's in a much more ready state. Please let me know if you have more concerns.

We'd need a discussion session for many points in this one. Aaron and I were talking about how we should expose Exporter/FXGraphExtractor the other day. One possible follow-up planned for #98421 is to actually delete DynamoOptimizeExporter (it is kept in history in working state after #98421 is landed should we revisit), since there really isn't anything that is not coverred by DynamoExporter. That leaves us with only FXSymbolicTraceExporter, which is itself a temorary solution and should go away if #95900 is done. Hence the design may not need to expose that much in public api, while retaining internal experiment flexibility.

I will look into #98421 one more time. The aspect that I was concerned is that we were using this adapter for internal and external purposes through the same mechanism/API. The external use causes some clutters on ExportOutput, while the internal use is actually needed.

Regarding this one, while we are in the experimental stage, having more than one FX extractor shouldn't be a problem; in fact, some things only the symbolic trace can do and we don't know how much time we will take to make dynamo export to do the same. When we can get rid of all and just stick to torch.export, we can remove the external API. But that doesn't mean the backend can't leverage a layered. The single export_fx_to_onnx, for instance, caused a hell of circular deps simply by exposing some classes to public API. If we layer things properly, dependency issues should be gone.

test/onnx/dynamo/test_exporter_api.py

BowenBao

Thanks for finishing what #95651 could not!

torch/onnx/__init__.py

BowenBao · 2023-04-21T17:34:40Z

test/onnx/test_fx_to_onnx_with_onnxruntime.py

@@ -654,6 +721,15 @@ def forward(self, x):
        )

    @pytorch_test_common.xfail("TypeError: missing a required argument: 'end'")


Many more xfails needs to be removed as being replaced by skip_fx_tracer.

One of the downside with parameterizing over tracer is skip over xfail which does not catch unexpected success.

I wonder if it is necessary to maintain the test coverage over DynamoOptimize. We are not doing it for FxSymbolicTracer. How to decide?

When it becomes a heavy burden, we decide if these tests are bringing benefit to compensate the overhead. Right now it is too early to determine benefit or overhead

We are already losing functionality of xfail. Is it worth it?

Can you elaborate the issue and why it cannot be fixed in any other way?

torch/onnx/_internal/exporter.py

torch/onnx/utils.py

BowenBao · 2023-04-21T18:16:00Z

torch/onnx/_internal/fx/passes/shape_inference.py

+        import torch._dynamo  # NOTE: Do not import at top level.
+
+        # Even torch/__init__.py does it internally, only
+        # Causes circular when torch._dynamo.* surfaces public facing API during `import torch`


True but not necessary for files not intended to be imported during import torch right?

Why does this file need to be imported so early?

no sure.

before this PR, it was being globally imported in the beginning of the module. This change delays to only when it is going to be used

BowenBao · 2023-04-21T18:18:32Z

torch/onnx/_internal/fx/dynamo_graph_extractor.py

How can we mark which Extractor we are actively maintaining, and which Extractor is "almost immutable code"?

That is what ResolvedExportOptions do, right? it sets the default (most used and preferred) values.

Hmm but is it ok that such information is not revealed if just looking at these classes alone?

That was a design decision made when we decided that ExportOptions would default everything to None and only the private ResolvedExportOptions would actually assign it.

I am open to changes, but seems unrelated to this task

That was not what I meant. Simply put, should we add into docstring of DynamoOptimize and DynamoExport saying if it is maintained?

I will split this PR in 2. One for the layering fx_tracer in the exporter and another re-adding optimize.
The latter is delaying the former which adds more value :)

torch/onnx/_internal/exporter.py

BowenBao · 2023-04-22T00:31:31Z

Heads up PRs that changes function_dispather.py are in motion to merge. More rebase incoming ... :(

titaiwangms

I suggest using ghstack to sperate this lind of big PR to multiple subPRs that we could discuss different topic in different ones, and it's also easier for you that subPRs can be merged one by one.

titaiwangms · 2023-04-23T21:14:22Z

torch/onnx/_internal/fx/dynamo_graph_extractor.py

+        return graph_module
+
+
+class DynamoOptimize(exporter.FXGraphExtractor):


Could you put the reason in PR description why we need this back after #99202

titaiwangms · 2023-04-23T21:21:25Z

torch/onnx/_internal/fx/dynamo_graph_extractor.py

+        # NOTE: Do not import at top level.
+        # Even torch/__init__.py does it internally, only
+        # Causes circular when torch._dynamo.* surfaces public facing API during `import torch`
+        import torch._dynamo


from? give more details :)

thiagocrepaldi · 2023-04-25T01:47:40Z

@BowenBao @titaiwangms please review #99940 instead

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

pytorchbot added the open source label Apr 12, 2023

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 12, 2023

thiagocrepaldi mentioned this pull request Apr 12, 2023

[ONNX] Introduce Input/Ouptut adapter; Switch to 'DynamoExporter' #98421

Closed

thiagocrepaldi force-pushed the thiagofc/introduce-fxgraphextractor branch from 3d56ab1 to e433956 Compare April 13, 2023 17:56

thiagocrepaldi marked this pull request as ready for review April 13, 2023 20:05

thiagocrepaldi requested review from BowenBao and abock as code owners April 13, 2023 20:05

thiagocrepaldi force-pushed the thiagofc/introduce-fxgraphextractor branch 2 times, most recently from 36343b7 to 976fb4c Compare April 14, 2023 18:13

janeyx99 added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Apr 14, 2023

thiagocrepaldi force-pushed the thiagofc/introduce-fxgraphextractor branch 6 times, most recently from 5c0e01e to 1be2507 Compare April 20, 2023 16:52

justinchuby reviewed Apr 20, 2023

View reviewed changes

test/onnx/dynamo/test_exporter_api.py Show resolved Hide resolved

Thiago Crepaldi added 4 commits April 21, 2023 17:15

Introduce FXGraphExtractor into torch.onnx.dynamo_export

e0c1eac

Fix import issue

c80b75a

Remove fx tracer from public API

a420566

Remove tracing_mode from torch.export call

65a10c8

thiagocrepaldi force-pushed the thiagofc/introduce-fxgraphextractor branch from 1be2507 to 65a10c8 Compare April 21, 2023 18:05

BowenBao reviewed Apr 21, 2023

View reviewed changes

Thiago Crepaldi added 2 commits April 21, 2023 20:26

Address comments

b5611a3

Move input/output adapter into FXGraphExtractor

19edea8

thiagocrepaldi requested a review from BowenBao April 21, 2023 22:02

titaiwangms reviewed Apr 23, 2023

View reviewed changes

thiagocrepaldi closed this Apr 25, 2023

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduce FXGraphExtractor into torch.onnx.dynamo_export #98893

Introduce FXGraphExtractor into torch.onnx.dynamo_export #98893

thiagocrepaldi commented Apr 12, 2023 •

edited by ghost

pytorch-bot bot commented Apr 12, 2023 •

edited

BowenBao commented Apr 13, 2023

thiagocrepaldi commented Apr 13, 2023

BowenBao left a comment

BowenBao Apr 21, 2023 •

edited

thiagocrepaldi Apr 21, 2023

BowenBao Apr 21, 2023

thiagocrepaldi Apr 21, 2023

BowenBao Apr 21, 2023

thiagocrepaldi Apr 21, 2023

BowenBao Apr 21, 2023

thiagocrepaldi Apr 21, 2023

BowenBao Apr 21, 2023

thiagocrepaldi Apr 21, 2023

BowenBao Apr 21, 2023

thiagocrepaldi Apr 24, 2023

BowenBao commented Apr 22, 2023

titaiwangms left a comment

titaiwangms Apr 23, 2023

titaiwangms Apr 23, 2023

thiagocrepaldi Apr 24, 2023

thiagocrepaldi commented Apr 25, 2023

		@@ -654,6 +721,15 @@ def forward(self, x):
		)

		@pytorch_test_common.xfail("TypeError: missing a required argument: 'end'")

		return graph_module


		class DynamoOptimize(exporter.FXGraphExtractor):

Introduce FXGraphExtractor into torch.onnx.dynamo_export #98893

Introduce FXGraphExtractor into torch.onnx.dynamo_export #98893

Conversation

thiagocrepaldi commented Apr 12, 2023 • edited by ghost

🤖 Generated by Copilot at 1be2507

Summary

Walkthrough

pytorch-bot bot commented Apr 12, 2023 • edited

🔗 Helpful Links

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

❗ 1 Active SEVs

❌ 9 Failures

BowenBao commented Apr 13, 2023

thiagocrepaldi commented Apr 13, 2023

BowenBao left a comment

Choose a reason for hiding this comment

BowenBao Apr 21, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BowenBao commented Apr 22, 2023

titaiwangms left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

thiagocrepaldi commented Apr 25, 2023

thiagocrepaldi commented Apr 12, 2023 •

edited by ghost

`🤖 Generated by Copilot at 1be2507`

pytorch-bot bot commented Apr 12, 2023 •

edited

BowenBao Apr 21, 2023 •

edited