[ONNX] Restore readable names for parameters and buffers #104493

BowenBao · 2023-07-01T02:09:01Z

Stack from ghstack (oldest at bottom):

This PR introduces a new pass that restore parameter and buffer names from original module.

It is useful for readability of the exported ONNX graph. It restores the parameter and buffer names from the original module. For example, if the original module has a parameter named root.linear.0.weight, and the parameter is renamed to
_param_constant9 by FX, this pass will rename it back.

[ghstack-poisoned]

pytorch-bot · 2023-07-01T02:09:04Z

🔗 Helpful Links

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

📄 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

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

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

[ghstack-poisoned]

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

thiagocrepaldi

Are you adding unit test to check names are actually being preserved?

This PR introduces a new pass that restore parameter and buffer names from original module. It is useful for readability of the exported ONNX graph. It restores the parameter and buffer names from the original module. For example, if the original module has a parameter named `root.linear.0.weight`, and the parameter is renamed to `_param_constant9` by FX, this pass will rename it back. [ghstack-poisoned]

BowenBao · 2023-07-05T22:01:23Z

@wschin @thiagocrepaldi updated PR, PTAL.

test/onnx/test_fx_to_onnx.py

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

test/onnx/test_fx_to_onnx.py

wschin

LGTM.

This PR introduces a new pass that restore parameter and buffer names from original module. It is useful for readability of the exported ONNX graph. It restores the parameter and buffer names from the original module. For example, if the original module has a parameter named `root.linear.0.weight`, and the parameter is renamed to `_param_constant9` by FX, this pass will rename it back. [ghstack-poisoned]

thiagocrepaldi

nice

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

… buffers" This PR introduces a new pass that restore parameter and buffer names from original module. It is useful for readability of the exported ONNX graph. It restores the parameter and buffer names from the original module. For example, if the original module has a parameter named `root.linear.0.weight`, and the parameter is renamed to `_param_constant9` by FX, this pass will rename it back. [ghstack-poisoned]

This PR introduces a new pass that restore parameter and buffer names from original module. It is useful for readability of the exported ONNX graph. It restores the parameter and buffer names from the original module. For example, if the original module has a parameter named `root.linear.0.weight`, and the parameter is renamed to `_param_constant9` by FX, this pass will rename it back. [ghstack-poisoned]

thiagocrepaldi

Besides fx_module and nn_module, the implementation also created different references for them. original_module as another name for nn_module and module as another name for fx_name. The first two names are clearer

thiagocrepaldi · 2023-07-07T15:20:24Z

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

+
+        For each `get_attr` node, if the target is a str representing a parameter or buffer
+        under `self.module`, we rename the parameter or buffer to its original name.
+        The parameters and buffers between `self.module` and `self.original_module` refer


Suggested change

The parameters and buffers between `self.module` and `self.original_module` refer

The parameters and buffers between `self.fx_module` and `self.nn_module` refer

I'll fix these in a follow up to merge this first to unblock you.

nvm I think better to rebase to be safe.. Another PR with pass was merged.

thiagocrepaldi · 2023-07-07T15:20:49Z

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

+        """Restore parameter and buffer names from original module.
+
+        For each `get_attr` node, if the target is a str representing a parameter or buffer
+        under `self.module`, we rename the parameter or buffer to its original name.


Suggested change

under `self.module`, we rename the parameter or buffer to its original name.

under `self.fx_module`, we rename the parameter or buffer to its original name.

there is a bunch of old self.module going on here

can't do anything, this is in base class

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

thiagocrepaldi · 2023-07-07T15:31:13Z

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

+        nn_module: torch.nn.Module,
+    ):
+        super().__init__(diagnostic_context, fx_module)
+        self.original_module = nn_module


Suggested change

self.original_module = nn_module

self.nn_module = nn_module

thiagocrepaldi · 2023-07-07T15:31:39Z

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

+        # TODO(bowbao): fix #104670 and replace "." with "/" to avoid collision.
+        normalized_name = new_name.replace(".", "_")
+        attr_value = getattr(self.module, old_name)
+        setattr(self.module, normalized_name, attr_value)


Suggested change

setattr(self.module, normalized_name, attr_value)

setattr(self.fx_module, normalized_name, attr_value)

This PR introduces a new pass that restore parameter and buffer names from original module. It is useful for readability of the exported ONNX graph. It restores the parameter and buffer names from the original module. For example, if the original module has a parameter named `root.linear.0.weight`, and the parameter is renamed to `_param_constant9` by FX, this pass will rename it back. [ghstack-poisoned]

## Context prior to this PR #100017 was merged onto PyTorch `main` branch with the goal of enabling `torch._dynamo.export` to perform symbolic tracing. In that context, symbolic tracing is defined as tracing of a model using fake inputs and weights. An input is Fake when `torch.nn.Tensor` is replaced by `torch._subclasses.FakeTensor`, whereas a weight is fake when a `torch.nn.Parameter` is replaced by `torch._subclasses.FakeTensor`. For additional context, several strategies were discussed with Meta to enable this feature, including 1) calling `torch._dynamo.export` within a `torch._subclass.FakeTensorMode` context and 2) **fake**fying input and model as separate step and then call `torch._dynamo.export` without an active `torch._subclass.FakeTensorMode` context. At the end, 2) was preferred and implemented by #100017 to minimize the number of side-effects the fake tensor mode has on the code base. As a consequence, `torch._dynamo.export` API introduced a new argument called `fake_mode`. When symbolic tracing is used, the user must pass in the `fake_mode` used to fakefy both the input and the model. Internally, `torch._dynamo.export` will adopt this `fake_mode` instead of creating its own instance. This is needed because each instance of `FakeTensorMode` has metadata on the tensor/parameter it fakefied. Thus, using real tensor/model and specify a `fake_mode` to `torch._dynamo.export` is an error. Also, specify a `fake_mode` instance to `torch._dynamo.export` different than the one used to fakefy the model and input is also an error. ## Changes introduced from this PR This PR is intended to integrate `torch._dynamo.export(fake_mode=...)` through `torch.onnx.dynamo_export`. In essence, it * Introduces a new public API `ONNXFakeContext` which wraps a `FakeTensorMode` under the hood. This removes complexity from the user side while still allow the exporter to leverage the fake mode. * Adds a new public API `enable_fake_mode` *context manager* that instantiates and return a `ONNXFakeContext`. * Adds a new `ExportOptions.fake_context` that will be used to persist the `ONNXFakeContext` created by `enable_fake_mode` and plumb through until it reaches the call to `torch._dynamo.export`. * Adds a `model_state_dict` argument to `ExportOutput.save` API. * When model is exported with fake tensors, no actual data exist in the FX module and, therefore, in the ONNX graph. * In fact, `torch.fx.make_fx` lifts initializers as model input when fake tensors are used * #104493 is needed to enforce name matching between Parameters and inputs * A model checkpoint file or state_dict is needed to populate the ONNX graph with real initializers through `export_output.save(model_state_dict=...)` API Symbolic tracing, or onnx fake mode, is only enabled when the user instantiates the input and model within the `enable_fake_mode` context. Otherwise, real tracing is done, which preserves the current behavior. ## Usability Because symbolic tracing depends a lot on having changes made on Dynamo side before it can be consumed on ONNX exporter, this feature may have its API and assumptions changed as symbolic tracing matures upstream. Nonetheless, it is still important to have this feature merged ASAP on the ONNX exporter side to "lock" changes on Dynamo that would otherwise break ONNX exporter without warning. Example: ```python class Model(torch.nn.Module): def __init__(self) -> None: super().__init__() self.linear = torch.nn.Linear(2, 2) def forward(self, x): out = self.linear(x) return out with torch.onnx.enable_fake_mode() as fake_context: x = torch.rand(5, 2, 2) model = Model() # Export the model with fake inputs and parameters export_options = ExportOptions(fake_context=fake_context) export_output = torch.onnx.dynamo_export( model, x, export_options=export_options ) model_state_dict = Model().state_dict() # optional export_output.save("/path/to/model.onnx", model_state_dict=model_state_dict) ``` ## Next steps * Add unit tests running the exported model with ORT Today this is not possible yet because `make_fx` used by our Decomposition pass lifts initializers as model inputs. However, the initializer names are not preserved by FX tracing, causing a mismatch between the initializer and input name. #104493 and #104741 should fix the initializer mismatch, enabling model execution * Revisit `ONNXTorchPatcher` and how the ONNX initializers are saved in the graph as external data We can try to get rid of the PyTorch patcher. If we can't, we might prefer to create specific patchers, say `FXSymbolicTracePatcher` used specifically during an export using `torch.fx.symbolic_trace` and maybe a `ExportOutputSavePacther` used specifically for `ExportOutput.save` to prevent "patching too many pytorch API that we don't need ## References * [FakeTensor implementation](https://github.com/pytorch/pytorch/blob/main/torch/_subclasses/fake_tensor.py) * [PR that adds fake tensor support to torch._dynamo.export](#100017) * [Short fake tensor documentation](https://pytorch.org/torchdistx/latest/fake_tensor.html) Pull Request resolved: #103865 Approved by: https://github.com/BowenBao

[ONNX] Restore readable names for parameters and buffers

9b51e92

[ghstack-poisoned]

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

BowenBao mentioned this pull request Jul 1, 2023

[ONNX] Export module as function #101029

Closed

BowenBao added the topic: improvements topic category label Jul 1, 2023

pytorchbot added the open source label Jul 1, 2023

BowenBao marked this pull request as ready for review July 5, 2023 16:43

BowenBao requested review from abock and thiagocrepaldi as code owners July 5, 2023 16:43

BowenBao marked this pull request as draft July 5, 2023 16:45

Update on "[ONNX] Restore readable names for parameters and buffers"

27042ec

[ghstack-poisoned]

BowenBao marked this pull request as ready for review July 5, 2023 17:00

BowenBao requested a review from wschin July 5, 2023 17:00

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Outdated Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

thiagocrepaldi added the module: onnx Related to torch.onnx label Jul 5, 2023

thiagocrepaldi reviewed Jul 5, 2023

View reviewed changes

BowenBao mentioned this pull request Jul 5, 2023

[ONNX] Set node.target instead of node.name as ONNX initializer name #104670

Closed

thiagocrepaldi reviewed Jul 5, 2023

View reviewed changes

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

torch/onnx/_internal/fx/passes/readability.py Show resolved Hide resolved

wschin reviewed Jul 5, 2023

View reviewed changes

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

wschin approved these changes Jul 5, 2023

View reviewed changes

thiagocrepaldi approved these changes Jul 6, 2023

View reviewed changes

torch/onnx/_internal/fx/passes/readability.py Outdated Show resolved Hide resolved

torch/onnx/_internal/fx/passes/readability.py Outdated Show resolved Hide resolved

torch/onnx/_internal/fx/passes/readability.py Outdated Show resolved Hide resolved

thiagocrepaldi mentioned this pull request Jul 6, 2023

[ONNX] Add fake tensor support to torch.onnx.dynamo_export #103865

Closed

BowenBao mentioned this pull request Jul 6, 2023

[ONNX] Fix exported onnx initializer name #104741

Closed

BowenBao added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 7, 2023

thiagocrepaldi reviewed Jul 7, 2023

View reviewed changes

thiagocrepaldi mentioned this pull request Jul 7, 2023

Add unit tests for numerical validation for onnx export with fake mode enabled #104794

Closed

pytorchmergebot added the Merged label Jul 8, 2023

pytorchmergebot closed this in 012561f Jul 8, 2023

facebook-github-bot deleted the gh/BowenBao/257/head branch July 11, 2023 14:15

thiagocrepaldi mentioned this pull request Jul 11, 2023

Fix initializer naming at torch.onnx.ExportOutput.save_model_with_external_data #104976

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ONNX] Restore readable names for parameters and buffers #104493

[ONNX] Restore readable names for parameters and buffers #104493

BowenBao commented Jul 1, 2023 •

edited

pytorch-bot bot commented Jul 1, 2023 •

edited

thiagocrepaldi left a comment

BowenBao commented Jul 5, 2023

wschin left a comment

thiagocrepaldi left a comment

thiagocrepaldi left a comment

thiagocrepaldi Jul 7, 2023

BowenBao Jul 7, 2023

BowenBao Jul 7, 2023

thiagocrepaldi Jul 7, 2023

BowenBao Jul 7, 2023

thiagocrepaldi Jul 7, 2023

thiagocrepaldi Jul 7, 2023

	The parameters and buffers between `self.module` and `self.original_module` refer
	The parameters and buffers between `self.fx_module` and `self.nn_module` refer

	under `self.module`, we rename the parameter or buffer to its original name.
	under `self.fx_module`, we rename the parameter or buffer to its original name.

	setattr(self.module, normalized_name, attr_value)
	setattr(self.fx_module, normalized_name, attr_value)

[ONNX] Restore readable names for parameters and buffers #104493

[ONNX] Restore readable names for parameters and buffers #104493

Conversation

BowenBao commented Jul 1, 2023 • edited

pytorch-bot bot commented Jul 1, 2023 • edited

🔗 Helpful Links

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

✅ No Failures

thiagocrepaldi left a comment

Choose a reason for hiding this comment

BowenBao commented Jul 5, 2023

wschin left a comment

Choose a reason for hiding this comment

thiagocrepaldi left a comment

Choose a reason for hiding this comment

thiagocrepaldi left a comment

Choose a reason for hiding this comment

thiagocrepaldi Jul 7, 2023

Choose a reason for hiding this comment

BowenBao Jul 7, 2023

Choose a reason for hiding this comment

BowenBao Jul 7, 2023

Choose a reason for hiding this comment

thiagocrepaldi Jul 7, 2023

Choose a reason for hiding this comment

BowenBao Jul 7, 2023

Choose a reason for hiding this comment

thiagocrepaldi Jul 7, 2023

Choose a reason for hiding this comment

thiagocrepaldi Jul 7, 2023

Choose a reason for hiding this comment

BowenBao commented Jul 1, 2023 •

edited

pytorch-bot bot commented Jul 1, 2023 •

edited