Attach ProxyTorchDispatchMode to ProxyTensor and use the mode in __torch_dispatch__ by IvanYashchuk · Pull Request #82549 · pytorch/pytorch

IvanYashchuk · 2022-07-31T13:23:41Z

Migrates ProxyTensors to always invoke ProxyTorchDispatchMode instead of calling into proxy_call.

facebook-github-bot · 2022-07-31T13:23:48Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/82549
❔ Python docs are pending
📄 Preview C++ docs built from this PR
❓Need help or want to give feedback on the CI? Visit our office hours

✅ No Failures (33 Pending)

As of commit 5f02e1a (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

ezyang · 2022-07-31T15:13:59Z

torch/fx/experimental/proxy_tensor.py

            raise AssertionError(f"Unexpected tracing type: {tracing_mode}")

-        proxy_mode = ProxyTorchDispatchMode(fx_tracer) if trace_factory_functions else nullcontext()
+        proxy_mode = ProxyTorchDispatchMode(fx_tracer, trace_factory_functions=trace_factory_functions)


nb: from a perf perspective this is worse than before, as you're forced to go into the torch dispatch hook even though you don't do anything with it. It doesn't matter too much though since trace factory functions is usually true.

Well, now we always need to create a ProxyTorchDispatchMode to be attached to ProxyTensor. When trace_factory_functions is false we still need the dispatch mode and not a nullcontext to pass it to the wrap_key function.

What breaks if we don't enable the mode (keep this line as it was previously), and don't have the trace_factory_functions argument? In my mind, if we don't have the mode on shouldn't it go:

Factory Function -> No Python Key (since no arguments have the python key and there's no mode) -> Not Traced

Non-factory function -> Python Key (from arguments) -> Mode -> redispatch on python key -> mode's torch dispatch

AttributeError: 'nullcontext' object has no attribute 'restore' when called with proxy_mode.restore(). We can try to intercept this case and call proxy_call with nullcontext passed as mode. Then at least one test fails TestProxyTensorCPU.test_mode_tracing_factory_function_no_factory_function_cpu with
TypeError: unsupported operand type(s) for +: 'ProxyTensor' and 'Tensor'.

ezyang · 2022-07-31T15:16:28Z

torch/fx/experimental/proxy_tensor.py

+                else:
+                    assert proxy_mode is arg.proxy_mode, "All arguments must be in the same proxy mode"
+
+        with enable_torch_dispatch_mode(proxy_mode):


Looks like this is latent bug from the fake tensor code; does proxy_mode.restore() work?

When using proxy_mode.restore() TestProxyTensorCPU.test_decomposition_interpreter_cpu fails with

RuntimeError: <torch.fx.experimental.proxy_tensor.ProxyTorchDispatchMode object at 0x7f084183ddb0> does not have any ancestors. Use the standard version instead of restore

this fixes this error

diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py index a7e38074eb6..d5bafefc0f2 100644 --- a/torch/fx/experimental/proxy_tensor.py +++ b/torch/fx/experimental/proxy_tensor.py @@ -476,6 +476,9 @@ class DecompositionInterpreter(torch.fx.Interpreter): return out def run(self, *args, **kwargs): + # snapshot the current mode stack + with self.mode: + pass with decompose(self.decomposition_table): return super().run(*args, **kwargs)

I need to think more about what went wrong on the API design side though.

cc @samdow: what's going on here is people are allocating mode objects, but then never actually entering them (and manually calling API functions that make use of the mode object.)

Ahh yeah that's a hard to understand error message. I'll change it so it's clearer that that's what happening

ezyang · 2022-08-01T15:44:50Z

I see you hacking it up. I'd like to properly diagnose why the underlying API isn't working in this case. As a stopgap, instead of relying on storing the mode on the proxy tensor for dispatch, you could solely use it to determine if tracing is enabled or not.

samdow · 2022-08-01T19:35:52Z

Okay full disclosure that I did have #81708 that I just haven't had the bandwidth to clean up and look at the failures. I'll leave some comments but happy to use either and close the other

samdow · 2022-08-01T19:43:46Z

torch/fx/experimental/proxy_tensor.py

+        # Verify that the proxy mode for all arguments is the same
+        proxy_mode = None
+        for arg in pytree.tree_flatten((args, kwargs))[0]:
+            if isinstance(arg, ProxyTensor):
+                if proxy_mode is None:
+                    proxy_mode = arg.proxy_mode
+                else:
+                    assert proxy_mode is arg.proxy_mode, "All arguments must be in the same proxy mode"


Small nit that Proxy also checks that all the tracers are the same (equivalent of this). I'm not attached to either one but if we want to have the same error message as before, we can just break after we find the first proxy mode and assume that they're all the same and let Proxy deal with it if that's not the case

Sure, I can change that.

Actually, @samdow, didn't you write some utilities for doing this? Or am I misremembering

Oh actually yeah--if we need to do something like this in the future, we just need to map to get all the proxy_modes and then call all_modes_same_scope(modes) (https://github.com/pytorch/pytorch/blob/master/torch/utils/_mode_utils.py#L124)

ezyang

🤞 no reference cycles lol (cc @eellison, the only reference cycle from mode on tensor was the converter cache right?)

…or-mode

IvanYashchuk · 2022-08-02T10:25:07Z

@pytorchbot merge

pytorchmergebot · 2022-08-02T10:28:28Z

@pytorchbot successfully started a merge job. Check the current status here

pytorchmergebot · 2022-08-02T10:33:32Z

Merge failed due to Refusing to merge as mandatory check(s) Lint failed for rule superuser
Raised by https://github.com/pytorch/pytorch/actions/runs/2781445565

IvanYashchuk · 2022-08-02T10:59:06Z

@pytorchbot merge

pytorchmergebot · 2022-08-02T11:02:03Z

@pytorchbot successfully started a merge job. Check the current status here

pytorchmergebot · 2022-08-02T11:12:08Z

Merge failed due to New commits were pushed while merging. Please rerun the merge command.
Raised by https://github.com/pytorch/pytorch/actions/runs/2781632681

IvanYashchuk · 2022-08-02T11:25:42Z

@pytorchbot merge

pytorchmergebot · 2022-08-02T11:27:00Z

@pytorchbot successfully started a merge job. Check the current status here

github-actions · 2022-08-02T14:03:23Z

Hey @IvanYashchuk.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

…rch_dispatch__ (#82549) (#82549) Summary: Migrates ProxyTensors to always invoke ProxyTorchDispatchMode instead of calling into `proxy_call`. Pull Request resolved: #82549 Approved by: https://github.com/ezyang Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/a3316cb3c7d7295de0ff686bf7b7888f2c4665ef Reviewed By: kit1980 Differential Revision: D38359546 fbshipit-source-id: c7747e20f578e1ea2ec15c5e65d2a9ff4dbbf9a7

Attach proxy_mode to ProxyTensor

170b73c

IvanYashchuk added the module: fx label Jul 31, 2022

IvanYashchuk requested a review from ezyang July 31, 2022 13:23

facebook-github-bot added the cla signed label Jul 31, 2022

facebook-github-bot added the fx label Jul 31, 2022

pytorchbot added the open source label Jul 31, 2022

IvanYashchuk mentioned this pull request Jul 31, 2022

Add context manager for conditional rewrites of torch.* to torch._refs.* calls #81764

Closed

ezyang reviewed Jul 31, 2022

View reviewed changes

Set proxy_mode.ancestors=set() and use proxy_mode.restore()

b547a91

Should enter the mode for being able to restore it later

bbeb147

samdow reviewed Aug 1, 2022

View reviewed changes

ezyang approved these changes Aug 1, 2022

View reviewed changes

IvanYashchuk added 2 commits August 2, 2022 13:22

Get only the first proxy_mode

47767a7

Merge remote-tracking branch 'upstream/viable/strict' into proxy-tens…

52f7a49

…or-mode

Add proxy_mode kwarg to wrap_output

1547d3a

Merge remote-tracking branch 'upstream/master' into proxy-tensor-mode

5f02e1a

pytorchmergebot added the Merged label Aug 2, 2022

pytorchmergebot closed this in a3316cb Aug 2, 2022

Conversation

IvanYashchuk commented Jul 31, 2022

Uh oh!

facebook-github-bot commented Jul 31, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

✅ No Failures (33 Pending)

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!

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!

ezyang commented Aug 1, 2022

Uh oh!

samdow commented Aug 1, 2022

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!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

IvanYashchuk commented Aug 2, 2022

Uh oh!

pytorchmergebot commented Aug 2, 2022

Uh oh!

pytorchmergebot commented Aug 2, 2022

Uh oh!

IvanYashchuk commented Aug 2, 2022

Uh oh!

pytorchmergebot commented Aug 2, 2022

Uh oh!

pytorchmergebot commented Aug 2, 2022

Uh oh!

IvanYashchuk commented Aug 2, 2022

Uh oh!

pytorchmergebot commented Aug 2, 2022

Uh oh!

github-actions bot commented Aug 2, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

facebook-github-bot commented Jul 31, 2022 •

edited

Loading