Update ONNX's IO Adapter to support FakeTensor with ExportedProgram by thiagocrepaldi · Pull Request #114407 · pytorch/pytorch ·

@BowenBao

Stack from ghstack (oldest at bottom):

Currently, the ONNX exporter using torch.nn.Module as input can support
FakeTensor because the ONNX model stores all initializers

When using torch.export.ExportedProgram as input, the initializers are
lifted as inputs. In order to execute the ONNX model, we need to pass a
reference to the non-fake model to the
ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be
fetched from the model and fed to the ONNX model as input

ps: #115461 will track the API revision for the cases where additional model_with_state_dict are required to produce complete ONNX files exported with fake support. This is also tracked by the umbrella fake tensor issue #105464 FYI @BowenBao

Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input [ghstack-poisoned]

🔗 Helpful Links

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

📄 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, 4 Pending

As of commit 73d1f8b with merge base 441ecf0 ():
💚 Looks good so far! There are no failures yet. 💚

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

…edProgram" Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input [ghstack-poisoned]

Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input ghstack-source-id: 954c65e Pull Request resolved: #114407

torch/onnx/_internal/exporter.py

…edProgram" Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input [ghstack-poisoned]

Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input ghstack-source-id: fdfe139 Pull Request resolved: #114407

…edProgram" Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input [ghstack-poisoned]

Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input ghstack-source-id: 8011f60 Pull Request resolved: #114407

…edProgram" Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input [ghstack-poisoned]

torch/onnx/_internal/fx/torch_export_graph_extractor.py

…edProgram" Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input [ghstack-poisoned]

#115281) Currently (after #114407), the user has must pass the original user ``model`` to APIs such as ``ONNXProgram.__call__``, ``ONNXProgram.adapt_torch_inputs_to_onnx`` and ``ONNXProgram.adapt_torch_outputs_to_onnx`` APIs. This was needed because when the model is fakefied, a version of the non-fakefied model is needed so that the Initializers, buffers and constants can be extracted from a real model (and used as input to the ONNX model). That approach brings an unnecessary usability burden to the user when the model is not fakefied, because the model that was already passed to ``torch.onnx.dynamo_export`` could be used to extract ``state_dict``. This PR adds ``ONNXProgram._model_torch`` attribute to store the user model and demote ``model`` argument of the aforementioned APIs to optional, only (as opposed to required). As a result, for the fakefied model scenario, the user still need to pass the required model, but for non fakefied models, the persisted model is implicitly used to extract the model state_dict, making it easier to use. Pull Request resolved: #115281 Approved by: https://.com/BowenBao ghstack dependencies: #114407

…4762) Fixed by #113982 Pull Request resolved: #114762 Approved by: https://.com/BowenBao ghstack dependencies: #114407, #115281

@BowenBao

…ytorch#114407) Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input ps: pytorch#115461 will track the API revision for the cases where additional `model_with_state_dict` are required to produce complete ONNX files exported with fake support. This is also tracked by the umbrella fake tensor issue pytorch#105464 FYI @BowenBao Pull Request resolved: pytorch#114407 Approved by: https://.com/BowenBao

@BowenBao

…ytorch#114407) Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input ps: pytorch#115461 will track the API revision for the cases where additional `model_with_state_dict` are required to produce complete ONNX files exported with fake support. This is also tracked by the umbrella fake tensor issue pytorch#105464 FYI @BowenBao Pull Request resolved: pytorch#114407 Approved by: https://.com/BowenBao

pytorch#115281) Currently (after pytorch#114407), the user has must pass the original user ``model`` to APIs such as ``ONNXProgram.__call__``, ``ONNXProgram.adapt_torch_inputs_to_onnx`` and ``ONNXProgram.adapt_torch_outputs_to_onnx`` APIs. This was needed because when the model is fakefied, a version of the non-fakefied model is needed so that the Initializers, buffers and constants can be extracted from a real model (and used as input to the ONNX model). That approach brings an unnecessary usability burden to the user when the model is not fakefied, because the model that was already passed to ``torch.onnx.dynamo_export`` could be used to extract ``state_dict``. This PR adds ``ONNXProgram._model_torch`` attribute to store the user model and demote ``model`` argument of the aforementioned APIs to optional, only (as opposed to required). As a result, for the fakefied model scenario, the user still need to pass the required model, but for non fakefied models, the persisted model is implicitly used to extract the model state_dict, making it easier to use. Pull Request resolved: pytorch#115281 Approved by: https://.com/BowenBao ghstack dependencies: pytorch#114407

…orch#114762) Fixed by pytorch#113982 Pull Request resolved: pytorch#114762 Approved by: https://.com/BowenBao ghstack dependencies: pytorch#114407, pytorch#115281

@BowenBao

…114407) (#115578) Currently, the ONNX exporter using torch.nn.Module as input can support FakeTensor because the ONNX model stores all initializers When using torch.export.ExportedProgram as input, the initializers are lifted as inputs. In order to execute the ONNX model, we need to pass a reference to the non-fake model to the ONNXProgram.adapt_torch_inputs_to_onnx API, so that initializers can be fetched from the model and fed to the ONNX model as input ps: #115461 will track the API revision for the cases where additional `model_with_state_dict` are required to produce complete ONNX files exported with fake support. This is also tracked by the umbrella fake tensor issue #105464 FYI @BowenBao Pull Request resolved: #114407 Approved by: https://.com/BowenBao

pytorch#115281) Currently (after pytorch#114407), the user has must pass the original user ``model`` to APIs such as ``ONNXProgram.__call__``, ``ONNXProgram.adapt_torch_inputs_to_onnx`` and ``ONNXProgram.adapt_torch_outputs_to_onnx`` APIs. This was needed because when the model is fakefied, a version of the non-fakefied model is needed so that the Initializers, buffers and constants can be extracted from a real model (and used as input to the ONNX model). That approach brings an unnecessary usability burden to the user when the model is not fakefied, because the model that was already passed to ``torch.onnx.dynamo_export`` could be used to extract ``state_dict``. This PR adds ``ONNXProgram._model_torch`` attribute to store the user model and demote ``model`` argument of the aforementioned APIs to optional, only (as opposed to required). As a result, for the fakefied model scenario, the user still need to pass the required model, but for non fakefied models, the persisted model is implicitly used to extract the model state_dict, making it easier to use. Pull Request resolved: pytorch#115281 Approved by: https://.com/BowenBao ghstack dependencies: pytorch#114407

#115281) (#115583) Currently (after #114407), the user has must pass the original user ``model`` to APIs such as ``ONNXProgram.__call__``, ``ONNXProgram.adapt_torch_inputs_to_onnx`` and ``ONNXProgram.adapt_torch_outputs_to_onnx`` APIs. This was needed because when the model is fakefied, a version of the non-fakefied model is needed so that the Initializers, buffers and constants can be extracted from a real model (and used as input to the ONNX model). That approach brings an unnecessary usability burden to the user when the model is not fakefied, because the model that was already passed to ``torch.onnx.dynamo_export`` could be used to extract ``state_dict``. This PR adds ``ONNXProgram._model_torch`` attribute to store the user model and demote ``model`` argument of the aforementioned APIs to optional, only (as opposed to required). As a result, for the fakefied model scenario, the user still need to pass the required model, but for non fakefied models, the persisted model is implicitly used to extract the model state_dict, making it easier to use. Pull Request resolved: #115281 Approved by: https://.com/BowenBao ghstack dependencies: #114407

thiagocrepaldi requested review from BowenBao, abock and wschin as code owners November 22, 2023 21:37

thiagocrepaldi mentioned this pull request Nov 22, 2023
Add support for models with mutated buffer on torch.onnx.dynamo_export #112272
Closed

pytorch-bot bot added the release notes: onnxtorch.onnx related changes that should show up in the release noteslabel Nov 22, 2023

pytorct added the open source label Nov 22, 2023

BowenBao reviewed Nov 22, 2023
View reviewed changes

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

titaiwangms self-requested a review November 28, 2023 00:31

thiagocrepaldi mentioned this pull request Nov 29, 2023
Enable builtin tests for ONNX Export with ExportedProgram models #114762
Closed

thiagocrepaldi commented Dec 5, 2023
View reviewed changes

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

thiagocrepaldi requested a review from BowenBao December 5, 2023 20:58

Thiago Crepaldi added 3 commits December 5, 2023 22:33

pytorchmergebot pushed a commit that referenced this pull request Dec 9, 2023
Enable builtin tests for ONNX Export with ExportedProgram models (#11…

13d2e3e
…4762) Fixed by #113982 Pull Request resolved: #114762 Approved by: https://.com/BowenBao ghstack dependencies: #114407, #115281

This was referencedDec 11, 2023
[Release 2.2][ONNX] Update ONNX's IO Adapter to support FakeTensor with ExportedProgram (… #115578
Merged
[v.2.2.0] Release Tracker #115300
Closed

thiagocrepaldi mentioned this pull request Dec 11, 2023
[Release 2.2][ONNX]Store user model to simplify ONNXProgram.{adapt_torch_*,__call__} AP #115583
Merged

facebook--bot deleted the gh/thiagocrepaldi/13/head branch December 12, 2023 15:30

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Conversation

Uh oh!

🔗 Helpful Links

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

⏳ No Failures, 4 Pending

Uh oh!

Uh oh!

Uh oh!

Uh oh!