Add support for models with mutated buffer on torch.onnx.dynamo_export #112272

thiagocrepaldi · 2023-10-27T18:15:29Z

Stack from ghstack (oldest at bottom):

This PR adds a unit test that leverages torch.export.ExportedProgram models that mutates registered buffers. Although the exporter already works out of the box in such scenario, the GraphModule and the exported ONNX model have extra outputs containing the mutated buffers. On future runs of the ONNX model, the mutated buffers are used as input to the model.

The aforementioned extra inputs and outputs are by design and the ONNXProgram.model_signature can be used to fetch detailed input/output schema for the exported model.

However, when we want to compare pytorch output to ONNX's, there is a mismatch between the schema because pytorch output does not include the mutated buffers present on the ONNX output.

This PR extends onnx_program.adapt_torch_outputs_to_onnx(torch_outputs) so that the mutated buffers are prepended to the Pytorch output, matching the ONNX schema.

`run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

pytorch-bot · 2023-10-27T18:15:34Z

🔗 Helpful Links

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

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit ce48b3d with merge base 88a8a0d ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / linux-focal-py3.8-clang10 / test (dynamo, 2, 2, linux.2xlarge) (gh)

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

`run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers ghstack-source-id: 4ef89ce32fecf1ecb3bd996b76ce65e1880c864b Pull Request resolved: #112272

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

`run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers ghstack-source-id: 979da739bfeeb9fcadc498e53d2afba1374de9c4 Pull Request resolved: #112272

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

test/onnx/onnx_test_common.py

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

…ynamo_export" This PR adds a unit test that leverages `torch.export.ExportedProgram` models that mutates registered buffers. Although the exporter already works out of the box in such scenario, the GraphModule and the exported ONNX model have extra outputs containing the mutated buffers. On future runs of the ONNX model, the mutated buffers are used as input to the model. The aforementioned extra inputs and outputs are by design and the `ONNXProgram.model_signature` can be used to fetch detailed input/output schema for the exported model. However, when we want to compare pytorch output to ONNX's, there is a mismatch between the schema because pytorch output does not include the mutated buffers present on the ONNX output. This PR extends `onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)` so that the mutated buffers are prepended to the Pytorch output, matching the ONNX schema. [ghstack-poisoned]

thiagocrepaldi · 2023-11-22T22:16:54Z

@pytorchbot merge

pytorchmergebot · 2023-11-22T22:18:46Z

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

BowenBao · 2023-11-22T22:19:07Z

test/onnx/onnx_test_common.py

-
+    # ONNXProgram holds a reference (not copy) to the original ref_model, including its state_dict.
+    # Thus, ONNXProgram() must run before ref_model() to prevent ref_model.forward() from changing the state_dict.
+    # Otherwise, the ref_model can change buffers on state_dict which would be used by ONNXProgram.__call__()


I recall there was a PR that will make ONNXProgram.call() update original pytorch model lol. This problem will arise again if that's added.

Btw there is a flag has_mutation up there that triggers model cloning. Might be helpful.

that pr was merged with this one.

huydhn · 2023-11-23T07:06:17Z

@pytorchbot revert -m 'Sorry for reverting you change but it is failing dynamo test in trunk https://hud.pytorch.org/pytorch/pytorch/commit/c4a22d6918b7ca218f2712d7e7e147aca7127fa3' -c weird

pytorchmergebot · 2023-11-23T07:07:54Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2023-11-23T07:08:01Z

@thiagocrepaldi your PR has been successfully reverted.

…mo_export (#112272)" This reverts commit c4a22d6. Reverted #112272 on behalf of https://github.com/huydhn due to Sorry for reverting you change but it is failing dynamo test in trunk https://hud.pytorch.org/pytorch/pytorch/commit/c4a22d6918b7ca218f2712d7e7e147aca7127fa3 ([comment](#112272 (comment)))

huydhn · 2023-11-23T08:16:05Z

It looks like the revert doesn't help, I will reland the change. Sorry for the churn.

huydhn · 2023-11-23T08:16:11Z

@pytorchbot merge

pytorchmergebot · 2023-11-23T08:18:19Z

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#112272) This PR adds a unit test that leverages `torch.export.ExportedProgram` models that mutates registered buffers. Although the exporter already works out of the box in such scenario, the GraphModule and the exported ONNX model have extra outputs containing the mutated buffers. On future runs of the ONNX model, the mutated buffers are used as input to the model. The aforementioned extra inputs and outputs are by design and the `ONNXProgram.model_signature` can be used to fetch detailed input/output schema for the exported model. However, when we want to compare pytorch output to ONNX's, there is a mismatch between the schema because pytorch output does not include the mutated buffers present on the ONNX output. This PR extends `onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)` so that the mutated buffers are prepended to the Pytorch output, matching the ONNX schema. Pull Request resolved: pytorch#112272 Approved by: https://github.com/titaiwangms, https://github.com/BowenBao

…mo_export (pytorch#112272)" This reverts commit c4a22d6. Reverted pytorch#112272 on behalf of https://github.com/huydhn due to Sorry for reverting you change but it is failing dynamo test in trunk https://hud.pytorch.org/pytorch/pytorch/commit/c4a22d6918b7ca218f2712d7e7e147aca7127fa3 ([comment](pytorch#112272 (comment)))

pytorch#112272) This PR adds a unit test that leverages `torch.export.ExportedProgram` models that mutates registered buffers. Although the exporter already works out of the box in such scenario, the GraphModule and the exported ONNX model have extra outputs containing the mutated buffers. On future runs of the ONNX model, the mutated buffers are used as input to the model. The aforementioned extra inputs and outputs are by design and the `ONNXProgram.model_signature` can be used to fetch detailed input/output schema for the exported model. However, when we want to compare pytorch output to ONNX's, there is a mismatch between the schema because pytorch output does not include the mutated buffers present on the ONNX output. This PR extends `onnx_program.adapt_torch_outputs_to_onnx(torch_outputs)` so that the mutated buffers are prepended to the Pytorch output, matching the ONNX schema. Pull Request resolved: pytorch#112272 Approved by: https://github.com/titaiwangms, https://github.com/BowenBao

Add torch.onnx.dynamo_export using ExportedProgram with mutated buffer

3e1a3c1

`run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

thiagocrepaldi requested review from BowenBao, abock and wschin as code owners October 27, 2023 18:15

This was referenced Oct 27, 2023

Rename torch.onnx.ExportOutput* to ONNXProgram* #112263

Closed

Add torch.onnx.dynamo_export test using ExportedProgram from file #112271

Closed

pytorch-bot bot added the topic: not user facing topic category label Oct 27, 2023

pytorchbot added the open source label Oct 27, 2023

thiagocrepaldi marked this pull request as draft October 27, 2023 18:16

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

09315ea

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

thiagocrepaldi mentioned this pull request Oct 27, 2023

Extend _TestONNXRuntime to reuses all tests for new model format #112289

Closed

Thiago Crepaldi added 2 commits October 30, 2023 21:24

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

c1bf845

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

b7ee97b

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

thiagocrepaldi mentioned this pull request Oct 30, 2023

Add decomposition for dynamo_export + ExportedProgram and remove None from input #112444

Closed

Thiago Crepaldi added 2 commits October 31, 2023 15:21

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

4a7e5de

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

109a19b

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

thiagocrepaldi marked this pull request as ready for review October 31, 2023 18:51

BowenBao reviewed Nov 1, 2023

View reviewed changes

test/onnx/onnx_test_common.py Outdated Show resolved Hide resolved

Thiago Crepaldi added 3 commits November 1, 2023 23:54

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

fe32225

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

48e3fd0

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

9582b0c

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

thiagocrepaldi requested a review from BowenBao November 8, 2023 23:41

Update on "Add torch.onnx.dynamo_export using ExportedProgram with mu…

c497957

…tated buffer" `run_test_with_fx_to_onnx_exporter_and_onnx_runtime` had to be extended to allow extra outputs on the ONNX graph due to extra outputs due to mutated inputs and buffers [ghstack-poisoned]

thiagocrepaldi mentioned this pull request Nov 22, 2023

Update ONNX's IO Adapter to support FakeTensor with ExportedProgram #114407

Closed

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

pytorchmergebot added the merging label Nov 22, 2023

BowenBao reviewed Nov 22, 2023

View reviewed changes

BowenBao approved these changes Nov 22, 2023

View reviewed changes

pytorchmergebot added Merged and removed merging labels Nov 23, 2023

pytorchmergebot closed this in c4a22d6 Nov 23, 2023

pytorchmergebot added the Reverted label Nov 23, 2023

pytorchmergebot reopened this Nov 23, 2023

pytorchmergebot added the merging label Nov 23, 2023

pytorchmergebot removed the merging label Nov 23, 2023

pytorchmergebot closed this in a76bb5d Nov 23, 2023

facebook-github-bot deleted the gh/thiagocrepaldi/7/head branch November 26, 2023 15:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for models with mutated buffer on torch.onnx.dynamo_export #112272

Add support for models with mutated buffer on torch.onnx.dynamo_export #112272

thiagocrepaldi commented Oct 27, 2023 •

edited

pytorch-bot bot commented Oct 27, 2023 •

edited

thiagocrepaldi commented Nov 22, 2023

pytorchmergebot commented Nov 22, 2023

BowenBao Nov 22, 2023

thiagocrepaldi Nov 27, 2023

huydhn commented Nov 23, 2023

pytorchmergebot commented Nov 23, 2023

pytorchmergebot commented Nov 23, 2023

huydhn commented Nov 23, 2023

huydhn commented Nov 23, 2023

pytorchmergebot commented Nov 23, 2023

Add support for models with mutated buffer on torch.onnx.dynamo_export #112272

Add support for models with mutated buffer on torch.onnx.dynamo_export #112272

Conversation

thiagocrepaldi commented Oct 27, 2023 • edited

pytorch-bot bot commented Oct 27, 2023 • edited

🔗 Helpful Links

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

✅ You can merge normally! (1 Unrelated Failure)

thiagocrepaldi commented Nov 22, 2023

pytorchmergebot commented Nov 22, 2023

Merge started

BowenBao Nov 22, 2023

Choose a reason for hiding this comment

thiagocrepaldi Nov 27, 2023

Choose a reason for hiding this comment

huydhn commented Nov 23, 2023

pytorchmergebot commented Nov 23, 2023

pytorchmergebot commented Nov 23, 2023

huydhn commented Nov 23, 2023

huydhn commented Nov 23, 2023

pytorchmergebot commented Nov 23, 2023

Merge started

thiagocrepaldi commented Oct 27, 2023 •

edited

pytorch-bot bot commented Oct 27, 2023 •

edited