Functionalize and compute joint simultaneously. #88063

ezyang · 2022-10-31T02:30:51Z

Stack from ghstack (oldest at bottom):

This also comes with some bug fixes that were uncovered from doing
this:

Forward device calls to inner tensor in FunctionalTensorWrapper
Make legacyExtractDispatchKey exclude Functionalize, so that
it can get at the real device type key. This is noncontroversial.
Stop stripping dense from key set. The reason for this is
FunctionalWrapperTensor may be used in contexts where people
query if it is dense or not. If it doesn't report this correctly
(from the dispatch key), it will cause errors. This caused some
torchbench models to fail when I did one-pass tracing.
Save and restore reapply views TLS correctly

Signed-off-by: Edward Z. Yang [email protected]

This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <[email protected]> [ghstack-poisoned]

pytorch-bot · 2022-10-31T02:30:53Z

🔗 Helpful Links

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

📄 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 our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

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

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

ezyang · 2022-10-31T13:16:36Z

Bah these tests are passing on the branch, need to find the fixes I need

ezyang · 2022-10-31T18:57:53Z

it's probably #87575

This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: 622ed0e Pull Request resolved: #88063

ezyang · 2022-11-01T14:28:35Z

This is blocked on #87647

This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

bdhirsh · 2022-11-02T15:59:27Z

aten/src/ATen/FunctionalTensorWrapper.cpp

-  // we need to know if we're dispatching to AutogradCPU or AutogradXLA).
-  // Instead, it's sufficient to remove the `Dense` dispatch key,
-  // which prevents us from accidentally trying to directly run a CPU/CUDA kernel.
-  key_set_ = key_set_.remove(c10::DispatchKey::Dense);


I am curious in what contexts we're querying whether or not a tensor is Dense 🤔

Probably user code that asks if the tensor is sparse or not?

Sorry, I don't remember, but yeah, in general it was stuff that was looking at the key set to figure out what type of tensor this was. It's too bad we had to get rid of this.

bdhirsh · 2022-11-02T15:59:33Z

aten/src/ATen/FunctionalTensorWrapper.cpp

  // We override a bunch of _custom(), so make sure they get called
  // TODO: metadata copying may not actually be necessary then
  set_custom_sizes_strides(SizesStridesPolicy::CustomSizes);
+  set_custom_device(true);


This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

bdhirsh · 2022-11-04T19:27:33Z

c10/core/TensorImpl.h

+    if (is_nested()) {
+      return false;
+    }
+    if (key_set_.has(DispatchKey::Functionalize)) {


what ended up causing this to be necessary?

The old behavior seems a bit more correct - if the inner tensor advertises as strided, the outer tensor does too.

Oh, this is to get the fix where we don't generate as_strided calls in backwards. I made a policy decision for functionalization to always opt into view meta reconstruction.

This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

bdhirsh · 2022-11-04T19:47:17Z

functorch/_src/aot_autograd.py

+            # I do something unsound instead
+            # assert arg is new_arg, "input argument was mutated, this is not valid"
+            if arg is not new_arg:
+                assert arg.shape == new_arg.shape


So, there are existing tests in test_functionalization.py that mutate input metadata.

To preserve the old behavior, should this PR just do the same thing that functionalize() does, and
use .as_strided_() to mutate the input metadata when detected, and wait for
#82602 to add the asserts?

It's a moot point, because dynamo never calls AOTAutograd with ops that mutates inputs, and as_strided_ is in that set. (And direct user usage is unlikely to call this.)

bdhirsh · 2022-11-04T19:52:18Z

test/functorch/test_aotdispatch.py


    # Deleting this in a followup
    xfail('nn.functional.feature_alpha_dropout', 'with_train'),
-    xfail('nn.functional.pad', 'circular'),


Did this PR end up fixing these xfails? 🤔

I thought I fixed this one in particular with #88198. Maybe tracing with functinalization in one go sidesteps this issue.

It's fixed by https://github.com/pytorch/pytorch/pull/88063/files#r1014429009

bdhirsh · 2022-11-04T19:54:22Z

test/functorch/test_aotdispatch.py

+        def f(x):
+            return CustomFn.apply(x)
+
+        self.verify_aot_autograd(f, [torch.randn(3)])


What's actually going on here, when we pass a custom autograd.function object into aot_autograd to be traced?

(1) I thought that "support for custom autograd functions" was pretty unrelated to this change
(2) Does this actually work properly today? (do we trace "through" the custom function, and add its custom fwd and bwd ops into our traced graph)

We trace inside the custom autograd function, and end up with an AOTautograd custom function which is identity on backwards and plus one on backwards.

(1) My goal wasn't really to get custom autograd function working (and indeed, in the end to end Dynamo codepath I don't think there's enough wiring for this to work), but one side effect of doing the tracing all in one go is that custom autograd should just work... and it does!
(2) Yes.

bdhirsh · 2022-11-04T20:05:19Z

functorch/_src/aot_autograd.py

+    def inner(*args, **kwargs):
+        def to_fun(t):
+            if isinstance(t, Tensor):
+                r = torch._to_functional_tensor(t)


I was still thinking about the fact that we're calling .detach() now in this PR where previously we didn't have to, and I just want to lay out my understanding.

Basically, if we do inp_new = inp_old.detach().requires_grad(original_gradness), our new input can have different than the original input, BUT this metadata only matters in practice w.r.t. mutations. Since we expect to be guaranteed a mutation-free-graph, then these metadatas difference won't matter when we're tracing through autograd. In particular, the metadata fields that can differ are ._is_leaf and ._is_view().

For sanity though, are there any tests that exercise aot autograd on inputs that require grad, but are not leaves? Seem worth adding.

Yes, the point is that it doesn't matter if you don't mutate inputs.

The non-leaf input tensor case is exercised by some of the end-to-end dynamo tests, though I don't see any direct tests for it with aot autograd. I think my preference is to add the torture tests when we actually support input mutation; that's when those non-leaf/leaf tests are actually interesting.

bdhirsh

left a few comments but pre-stamping, generally lgtm

ezyang · 2022-11-05T03:50:58Z

@pytorchbot merge

pytorchmergebot · 2022-11-05T03:52:36Z

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

This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: pytorch#88063 Approved by: https://github.com/bdhirsh

Signed-off-by: Edward Z. Yang <[email protected]>

This also comes with some bug fixes that were uncovered from doing this: - Forward device calls to inner tensor in FunctionalTensorWrapper - Make legacyExtractDispatchKey exclude Functionalize, so that it can get at the real device type key. This is noncontroversial. - Stop stripping dense from key set. The reason for this is FunctionalWrapperTensor may be used in contexts where people query if it is dense or not. If it doesn't report this correctly (from the dispatch key), it will cause errors. This caused some torchbench models to fail when I did one-pass tracing. - Save and restore reapply views TLS correctly Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: pytorch#88063 Approved by: https://github.com/bdhirsh

ezyang requested a review from Chillee as a code owner October 31, 2022 02:30

github-actions bot requested review from Krovatkin, albanD, anjali411, antoniojkim, bdhirsh, miladm and wconstab October 31, 2022 02:31

ezyang mentioned this pull request Oct 31, 2022

aot_dispatch test fix: always use functionalization in symbolic tests #87647

Closed

ezyang added release notes: composability release notes category topic: not user facing topic category labels Oct 31, 2022

ezyang mentioned this pull request Nov 2, 2022

Add a reshape_copy operator. #88314

Closed

bdhirsh reviewed Nov 2, 2022

View reviewed changes

ezyang mentioned this pull request Nov 2, 2022

Put Python Dispatcher cache in dict, clear it on new registrations. #88329

Closed

ezyang mentioned this pull request Nov 2, 2022

Preserve reshapes in AOTAutograd #88361

Closed

ezyang added ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request labels Nov 2, 2022

ezyang added 2 commits November 2, 2022 13:46

ezyang added 2 commits November 2, 2022 20:44

bdhirsh reviewed Nov 4, 2022

View reviewed changes

bdhirsh approved these changes Nov 4, 2022

View reviewed changes

pytorchmergebot added the Merged label Nov 5, 2022

pytorchmergebot closed this in 0e3031f Nov 5, 2022

ezyang added a commit that referenced this pull request Nov 7, 2022

This was obsoleted by #88063

270cdb3

Signed-off-by: Edward Z. Yang <[email protected]>

ngimel mentioned this pull request Nov 15, 2022

With fixed index_add decomposition, compile_fx produces wrong graph pytorch/torchdynamo#1863

Closed

facebook-github-bot deleted the gh/ezyang/1482/head branch June 8, 2023 16:27

Functionalize and compute joint simultaneously. #88063

Functionalize and compute joint simultaneously. #88063

Uh oh!

Conversation

ezyang commented Oct 31, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

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

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

ezyang commented Oct 31, 2022

Uh oh!

ezyang commented Oct 31, 2022

Uh oh!

ezyang commented Nov 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!

Choose a reason for hiding this comment

Uh oh!

bdhirsh Nov 4, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

bdhirsh Nov 4, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

bdhirsh left a comment

Choose a reason for hiding this comment

Uh oh!

ezyang commented Nov 5, 2022

Uh oh!

pytorchmergebot commented Nov 5, 2022

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ezyang commented Oct 31, 2022 •

edited

Loading

pytorch-bot bot commented Oct 31, 2022 •

edited

Loading

bdhirsh Nov 4, 2022 •

edited

Loading

bdhirsh Nov 4, 2022 •

edited

Loading