Modernize FakeTensorMode, throw on non-fake inputs by eellison · Pull Request #78516 · pytorch/pytorch

eellison · 2022-05-31T11:13:17Z

Stack from ghstack (oldest at bottom):

Modernizes FakeTensorMode by inheriting from TorchDispatchMode. FakeTensor and FakeTensorMode now call into a common helper function. I didn't see any existing idiomatic pattern for this so please feel free to comment if you think something else would be more ergonomic.

This also throws if a non-Fake tensor is an input to FakeTensorMode.. So it effectively only extends handling to constructors (for now, more general handling later in stack). This is because we need more careful logic to detect something like:

x = torch.rand([1, 1])
with FakeTensorMode()
    y = x.add_(3)
    x.resize_([4])  # y should be resized here as well, no way to support this, error

[ghstack-poisoned]

facebook-github-bot · 2022-05-31T11:13:23Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/78516
📄 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

✅ No Failures (0 Pending)

As of commit c7b545e (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.

[ghstack-poisoned]

Modernizes FakeTensorMode by inheriting from `TorchDispatchMode`. `FakeTensor` and `FakeTensorMode` now call into a common helper function. I didn't see any existing idiomatic pattern for this so please feel free to comment if you think something else would be more ergonomic. This also throws if a non-Fake tensor is an input to `FakeTensorMode`.. So it effectively only extends handling to constructors (for now, more general handling later in stack). [ghstack-poisoned]

ezyang · 2022-05-31T17:22:17Z

test/test_fake_tensor.py


    def test_constructor(self):
-        with enable_torch_dispatch_mode(FakeTensorMode):
+        with enable_torch_dispatch_mode(FakeTensorMode(inner=None)):


can we just write all of these tests using FakeTensorMode.push()? Or maybe @samdow's patch has landed so with FakeTensorMode() works now

Patch not landed yet (Richard has been hunting down a functorch/dispatch key bug and then I was going to ask him if he had anything to add)--and this actually doesn't work yet because push uses push_torch_dispatch_mode which causes the error Elias saw late last week (Creating a new Tensor subclass FakeTensor but the raw Tensor object is already associated to a python object of type Tensor)

On that note, I forgot to mention Richard debugged that and what's happening is that because push makes every mode have an inner mode (setting BaseTorchDispatchMode if there isn't one set), detach has a pyobj associated with it and we end up with this error). The patch fixes this because it removes BaseTorchDispatchMode but it will come up if we nest fake tensor mode with another mode. Long technical way of saying if we want FakeTensorMode to be composable, we should use push and add no_dispatch around all constructors in torch_dispatch. I'll try to flag everywhere this will come up

Can we have the constructors automatically apply no dispatch? Seems safer.

We can do it in _make_subclass, sure. I don't think there's any case where we wouldn't want this? (cc @zou3519)

Filed here: #78565, also FakeTensorMode.push() still errors

ezyang · 2022-05-31T17:23:52Z

torch/_subclasses/fake_tensor.py

-        return tree_map(partial(wrap, device=common_device), r)
+        def run_fn(func, types, args, kwargs):
+            return torch.Tensor.__torch_dispatch__(func, types, args, kwargs)
+        return torch_dispatch_impl(cls, func, types, args, kwargs, run_fn)


@samdow and I discussed this on Thursday and we think the right way to do code reuse is you put the real implementation in the mode, and then in the subclass implementation you (1) store the mode that allocated the subclass and (2) enable that mode before redispatching on the tensors as is.

I'm going to leave to follow up because changing this causes a few different errors (which i'll file issues / smaller repros for and/or debug) and i would like to unblock dynamo

samdow

LGTM! All of my points are small/cleanup stuff

torch/_subclasses/fake_tensor.py

samdow · 2022-05-31T18:08:11Z

torch/_subclasses/fake_tensor.py

+        conversion_made = False
+
+        def check_non_fake_tensor(x):
+            nonlocal conversion_made
+            conversion_made = conversion_made or (isinstance(x, torch.Tensor) and not isinstance(x, FakeTensor))
+
+        tree_map(check_non_fake_tensor, args)
+        tree_map(check_non_fake_tensor, kwargs)


nit: not sure which is more idiomatic but we could do any(tree_flatten(tree_map...)[0]) or even something like

conversion_made = False for x in tree_flatten(args)[0] + tree_flatten(kwargs)[0]: if isinstance(x, torch.Tensor) and not isinstance(x, FakeTensor): raise ...

this just saves us from using a nonlocal and lets us short circuit in the case that we run into the error

IMO it's not worth optimizing for failure modes for execution speed or when sacrificing readability..... but maybe this is more readable, will consider it. thanks!

I'm mentally filing an issue to do tree_iterate separately or something along those lines because there are a lot of places in the codebase where tree_map is used where the return value isnt used

samdow · 2022-05-31T18:11:34Z

torch/_subclasses/fake_tensor.py

+        def run_fn(func, types, args, kwargs):
+            return torch.Tensor.__torch_dispatch__(func, types, args, kwargs)


We should be okay running func(*args, **kwargs) in all cases. This way we can also remove the run_function argument

This actually causes an error:

x = FakeTensor.from_tensor(torch.tensor(0.0)) y = FakeTensor.from_tensor(torch.rand([4, 4], device="cuda")) out = x + y ... ERROR: test_zero_dim (__main__.FakeTensorTest) ---------------------------------------------------------------------- Traceback (most recent call last): File "test/test_fake_tensor.py", line 32, in test_zero_dim out = x + y TypeError: unsupported operand type(s) for +: 'FakeTensor' and 'FakeTensor'

Although it's a mute point because I'm going to do the refactoring of tensors storing mode and calling into it

samdow · 2022-05-31T18:48:29Z

test/test_fake_tensor.py


    def test_constructor(self):
-        with enable_torch_dispatch_mode(FakeTensorMode):
+        with enable_torch_dispatch_mode(FakeTensorMode(inner=None)):


Patch not landed yet (Richard has been hunting down a functorch/dispatch key bug and then I was going to ask him if he had anything to add)--and this actually doesn't work yet because push uses push_torch_dispatch_mode which causes the error Elias saw late last week (Creating a new Tensor subclass FakeTensor but the raw Tensor object is already associated to a python object of type Tensor)

On that note, I forgot to mention Richard debugged that and what's happening is that because push makes every mode have an inner mode (setting BaseTorchDispatchMode if there isn't one set), detach has a pyobj associated with it and we end up with this error). The patch fixes this because it removes BaseTorchDispatchMode but it will come up if we nest fake tensor mode with another mode. Long technical way of saying if we want FakeTensorMode to be composable, we should use push and add no_dispatch around all constructors in torch_dispatch. I'll try to flag everywhere this will come up

torch/_subclasses/fake_tensor.py

Modernizes FakeTensorMode by inheriting from `TorchDispatchMode`. `FakeTensor` and `FakeTensorMode` now call into a common helper function. I didn't see any existing idiomatic pattern for this so please feel free to comment if you think something else would be more ergonomic. This also throws if a non-Fake tensor is an input to `FakeTensorMode`.. So it effectively only extends handling to constructors (for now, more general handling later in stack). This is because we need more careful logic to detect something like: ``` x = torch.rand([1, 1]) with FakeTensorMode() y = x.add_(3) x.resize_([4]) # y should be resized here as well, no way to support this, error ``` [ghstack-poisoned]

eellison · 2022-06-01T21:42:44Z

@pytorchbot merge this please

github-actions · 2022-06-01T21:44:30Z

Hey @eellison.
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.

Summary: Pull Request resolved: #78516 Approved by: https://github.com/samdow Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/6671b504f7e3934bb26df93fd9a02d4081ba1713 Reviewed By: b0noI Differential Revision: D36854261 Pulled By: eellison fbshipit-source-id: 41c3f0d74b592561a2f9a3262ae2bc421a6c43af

Modernize FakeTensorMode, throw on non-fake inputs

860d421

[ghstack-poisoned]

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

eellison mentioned this pull request May 31, 2022

Fake Tensor Part 1 #77969

Closed

Update on "Modernize FakeTensorMode, throw on non-fake inputs"

f770a99

[ghstack-poisoned]

This was referenced May 31, 2022

Add CPU Fallback #78522

Closed

refactor op handling to use register pattern #78523

Closed

Update on "Modernize FakeTensorMode, throw on non-fake inputs"

849ffc4

[ghstack-poisoned]

eellison mentioned this pull request May 31, 2022

[WIP] non-fake inputs #78524

Closed

Update on "Modernize FakeTensorMode, throw on non-fake inputs"

bc388ce

[ghstack-poisoned]

eellison requested review from ezyang and samdow May 31, 2022 16:21

eellison mentioned this pull request May 31, 2022

add non-kwarg device and _like constructors #78536

Closed

ezyang reviewed May 31, 2022

View reviewed changes

samdow approved these changes May 31, 2022

View reviewed changes

eellison mentioned this pull request Jun 1, 2022

Migrate FakeTensors to always call into FakeTensorMode and have them hold a reference #78677

Closed

pytorchmergebot added the Merged label Jun 1, 2022

pytorchmergebot closed this in 6671b50 Jun 1, 2022

facebook-github-bot deleted the gh/eellison/298/head branch June 5, 2022 14:16

		def run_fn(func, types, args, kwargs):
		return torch.Tensor.__torch_dispatch__(func, types, args, kwargs)

Conversation

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

Uh oh!

Uh oh!

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

Uh oh!

🔗 Helpful links

✅ No Failures (0 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!

eellison May 31, 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!

samdow left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison May 31, 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!

Uh oh!

Uh oh!

eellison commented Jun 1, 2022

Uh oh!

github-actions bot commented Jun 1, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

eellison commented May 31, 2022 •

edited

Loading

facebook-github-bot commented May 31, 2022 •

edited

Loading

eellison May 31, 2022 •

edited

Loading

eellison May 31, 2022 •

edited

Loading