Ensure outer aliasing on DTensor matches inner aliasing #158954
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/158954
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 3 Unrelated FailuresAs of commit 63ca3e4 with merge base d678674 ( NEW FAILURE - The following job has failed:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
torch/distributed/tensor/_api.py
Outdated
| @@ -355,10 +356,15 @@ def __coerce_same_metadata_as_tangent__(self, flatten_spec, expected_type=None): | |||
| # pyre-fixme[3]: Return type must be annotated. | |||
| # pyre-fixme[2]: Parameter must be annotated. | |||
| def __torch_dispatch__(cls, func, types, args=(), kwargs=None): # type: ignore[override] | |||
| return DTensor._op_dispatcher.dispatch( | |||
| return return_and_correct_aliasing( | |||
There was a problem hiding this comment.
so.. what does this actually do? is it basically mapping the aliasing metadata from the inner tensor to the outer wrapper?
There was a problem hiding this comment.
Yes. It will also make the "fake" storages match if needed such that our "standard" aliasing checks will pass.
torch/distributed/tensor/_api.py
Outdated
| @@ -355,10 +356,15 @@ def __coerce_same_metadata_as_tangent__(self, flatten_spec, expected_type=None): | |||
| # pyre-fixme[3]: Return type must be annotated. | |||
| # pyre-fixme[2]: Parameter must be annotated. | |||
| def __torch_dispatch__(cls, func, types, args=(), kwargs=None): # type: ignore[override] | |||
| return DTensor._op_dispatcher.dispatch( | |||
| return return_and_correct_aliasing( | |||
There was a problem hiding this comment.
Yes. It will also make the "fake" storages match if needed such that our "standard" aliasing checks will pass.
|
dang it, it makes some dtensor tests deadlock! Time for misery debugging |
|
Maybe it's unrelated? It doesn't repro locally |
|
@pytorchbot merge |
Merge startedYour 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 |
Merge failedReason: 1 mandatory check(s) failed. The first few are: Dig deeper by viewing the failures on hud |
|
It is not unrelated. It's actually this: it only happens on two ranks and not on the other two. |
Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]
|
@wanchaol this potentially adds overhead |
|
I know how to fix this but I don't have the patience to do it right now. So here's the deal. So the root cause of the test failure here is that occasionally DTensor will put garbage in local tensor, specifically when a DTensor is not participating in a collective. I poked a bit to see if I could enforce a more uniform invariant on local tensor, but I eventually concluded that it is not so easy. So OK, let's just only fix aliases when we participated in ranks. So while I was writing this I started thinking about the performance overhead of doing this. So then I cracked open return_and_correct_aliasing and started reading the implementation. It looks... kind of expensive, actually? So then I started snooping around to see if I could do it in a simple way. Along the way I noticed that return_and_correct_aliasing probably doesn't handle split correctly? And I also noticed that it had already manually reimplemented the mutation/out aliasing behavior. So it's just storage left. But there's no way to easily tell if you ran a view op. And to do it fast you want to have pre-computed the alias relationship instead of doing a doubly nested iteration... |
Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]
Stack from ghstack (oldest at bottom):
Signed-off-by: Edward Z. Yang ezyang@meta.com
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta