Ensure outer aliasing on DTensor matches inner aliasing #158954

ezyang · 2025-07-23T20:15:44Z

Stack from ghstack (oldest at bottom):

-> Ensure outer aliasing on DTensor matches inner aliasing #158954

Signed-off-by: Edward Z. Yang ezyang@meta.com

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta

[ghstack-poisoned]

pytorch-bot · 2025-07-23T20:15:48Z

🔗 Helpful Links

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

📄 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 2a5fd23 with merge base f33ce40 ():
💚 Looks good so far! There are no failures yet. 💚

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

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 26a291d Pull-Request: #158954

wconstab · 2025-07-23T20:30:28Z

torch/distributed/tensor/_api.py

@@ -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(


so.. what does this actually do? is it basically mapping the aliasing metadata from the inner tensor to the outer wrapper?

Yes. It will also make the "fake" storages match if needed such that our "standard" aliasing checks will pass.

albanD · 2025-07-23T21:15:16Z

torch/distributed/tensor/_api.py

@@ -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(


Yes. It will also make the "fake" storages match if needed such that our "standard" aliasing checks will pass.

ezyang · 2025-07-24T02:04:45Z

dang it, it makes some dtensor tests deadlock! Time for misery debugging

ezyang · 2025-08-10T18:38:21Z

Maybe it's unrelated? It doesn't repro locally

[ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 0829536 Pull-Request: #158954

ezyang · 2025-08-10T18:38:45Z

@pytorchbot merge

pytorchmergebot · 2025-08-10T18:40:30Z

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

pytorchmergebot · 2025-08-10T20:20:47Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-jammy-cuda12.8-py3.10-gcc11 / test (distributed, 1, 3, linux.g4dn.12xlarge.nvidia.gpu)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

ezyang · 2025-08-12T00:30:38Z

It is not unrelated. It's actually this:

RuntimeError: Process 1 exited with error code 10 and exception:
Traceback (most recent call last):
  File "/data/users/ezyang/b/pytorch/torch/testing/_internal/common_distributed.py", line 804, in run_test
    getattr(self, test_name)()
  File "/data/users/ezyang/b/pytorch/torch/testing/_internal/common_distributed.py", line 658, in wrapper
    fn()
  File "/data/users/ezyang/b/pytorch/torch/testing/_internal/common_utils.py", line 3224, in wrapper
    method(*args, **kwargs)
  File "/data/users/ezyang/b/pytorch/torch/testing/_internal/distributed/_tensor/common_dtensor.py", line 490, in wrapper
    raise e
  File "/data/users/ezyang/b/pytorch/torch/testing/_internal/distributed/_tensor/common_dtensor.py", line 487, in wrapper
    func(self, *args, **kwargs)  # type: ignore[misc]
  File "/data/users/ezyang/b/pytorch/torch/testing/_internal/distributed/checkpoint_utils.py", line 153, in wrapper
    func(self, *args, **kwargs)
  File "/data/users/ezyang/b/pytorch/torch/testing/_internal/common_distributed.py", line 215, in wrapper
    return func(*args, **kwargs)
  File "/data/users/ezyang/b/pytorch/test/distributed/checkpoint/test_dtensor_checkpoint.py", line 125, in test_distributed_tensor_planner
    (model, sharded_dt, replicated_dt) = self.create_dtensor_model(
  File "/data/users/ezyang/b/pytorch/test/distributed/checkpoint/test_dtensor_checkpoint.py", line 104, in create_dtensor_model
    model = MyTestModule(
  File "/data/users/ezyang/b/pytorch/test/distributed/checkpoint/test_dtensor_checkpoint.py", line 40, in __init__
    self.submesh_sdt = torch.nn.Parameter(submesh_sdt)
  File "/data/users/ezyang/b/pytorch/torch/nn/parameter.py", line 60, in __new__
    t = data.detach().requires_grad_(requires_grad)
  File "/data/users/ezyang/b/pytorch/torch/_compile.py", line 53, in inner
    return disable_fn(*args, **kwargs)
  File "/data/users/ezyang/b/pytorch/torch/_dynamo/eval_frame.py", line 1005, in _fn
    return fn(*args, **kwargs)
  File "/data/users/ezyang/b/pytorch/torch/distributed/tensor/_api.py", line 367, in __torch_dispatch__
    return return_and_correct_aliasing(
  File "/data/users/ezyang/b/pytorch/torch/utils/_python_dispatch.py", line 655, in return_and_correct_aliasing
    _correct_storage_aliasing(
  File "/data/users/ezyang/b/pytorch/torch/utils/_python_dispatch.py", line 513, in _correct_storage_aliasing
    alias_non_inplace_storage(args[arg_idx], outs[return_idx])
  File "/data/users/ezyang/b/pytorch/torch/utils/_python_dispatch.py", line 500, in alias_non_inplace_storage
    torch._functionalize_unsafe_set(ret, arg)
RuntimeError: Attempted to set the storage of a tensor on device "cpu" to a storage on different device "cuda:1".  This is no longer allowed; the devices must match.

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]

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: e0ba3bf Pull-Request: #158954

ezyang · 2025-08-12T02:24:29Z

@wanchaol this potentially adds overhead

ezyang · 2025-08-12T03:20:34Z

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]

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: b1b7f34 Pull-Request: #158954

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: b89a542 Pull-Request: #158954

ezyang · 2025-08-12T13:23:16Z

I fixed it in an easier way

wconstab · 2025-08-12T13:28:18Z

torch/distributed/tensor/_dispatch.py

@@ -164,7 +165,8 @@ def dispatch(
        assert output_sharding is not None, "output sharding should not be None"

        mesh = op_info.compute_mesh
-        if mesh.get_coordinate() is not None:
+        participating = mesh.get_coordinate() is not None


When is this false? Is this due to h even sharding leaving one rank 'empty'?

A DTensor is associated with a device mesh. The device mesh does NOT necessarily have to cover all of the GPUs that are actually known. We are still SPMD over everything, including things that are not in the mesh. Then those are not participating.

hmm. this is kinda news to me.

i would like to see what example this happened in to understand it better. can you link to it?

usually, we partition the world gpus into various submeshes, but each 'op' still happens on each gpu. and each submesh dim is an orthogonal partitioning of the world mesh. Maybe just in some unit tests there are like 8gpus but the test case only initializes 4 of them?

wconstab · 2025-08-12T13:42:44Z

torch/distributed/tensor/_op_schema.py

@@ -450,6 +450,12 @@ def is_out_variant_op(self) -> bool:
        # be entirely correct, but it's good enough for now.
        return "out" in self.op._schema.overload_name

+    def is_view_op(self) -> bool:
+        return any(
+            a.alias_info is not None and not a.alias_info.is_write


what does 'not is_write' have to do with it being a view? would a modified view not set this?

I'm being a little defensive here. Mutating ops have a write alias like mul_(Tensor(a!) self) -> Tensor(a!) self. The logic I'm applying here would work in this case, but it is not necessary because these have already been handled by DTensor manually.

wconstab

it looks like you sidestepped the issue of dealing with the cpu tensor by just fixing up the real tensors, which seems good to me. asked a couple of questions but assuming you have good answers i'm happy with the change.

ezyang · 2025-08-12T17:45:35Z

@pytorchbot merge

pytorchmergebot · 2025-08-12T17:47:23Z

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

Update

565943d

[ghstack-poisoned]

ezyang added a commit that referenced this pull request Jul 23, 2025

Ensure outer aliasing on DTensor matches inner aliasing

2dc4b5e

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 26a291d Pull-Request: #158954

pytorch-bot bot added ciflow/inductor oncall: distributed Add this issue/PR to distributed oncall triage queue labels Jul 23, 2025

github-actions bot requested review from albanD, antoniojkim, bdhirsh, miladm and SherlockNoMad July 23, 2025 20:15

wconstab reviewed Jul 23, 2025

View reviewed changes

albanD approved these changes Jul 23, 2025

View reviewed changes

ezyang added the topic: not user facing topic category label Jul 24, 2025

XilunWu mentioned this pull request Aug 5, 2025

[DTensor] add op support: aten.squeeze_.dim #159532

Draft

Update

abaee9d

[ghstack-poisoned]

ezyang added a commit that referenced this pull request Aug 10, 2025

Ensure outer aliasing on DTensor matches inner aliasing

f4b8611

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: 0829536 Pull-Request: #158954

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 10, 2025

pytorchmergebot added the merging label Aug 10, 2025

pytorchmergebot removed the merging label Aug 10, 2025

Update on "Ensure outer aliasing on DTensor matches inner aliasing"

34bb06f

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

ezyang added a commit that referenced this pull request Aug 12, 2025

Ensure outer aliasing on DTensor matches inner aliasing

226fac6

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: e0ba3bf Pull-Request: #158954

Update on "Ensure outer aliasing on DTensor matches inner aliasing"

63ca3e4

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

ezyang added a commit that referenced this pull request Aug 12, 2025

Ensure outer aliasing on DTensor matches inner aliasing

68872d6

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: b1b7f34 Pull-Request: #158954

Update on "Ensure outer aliasing on DTensor matches inner aliasing"

2a5fd23

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

ezyang added a commit that referenced this pull request Aug 12, 2025

Ensure outer aliasing on DTensor matches inner aliasing

c74d59e

Signed-off-by: Edward Z. Yang <ezyang@meta.com> ghstack-source-id: b89a542 Pull-Request: #158954

wconstab reviewed Aug 12, 2025

View reviewed changes

wconstab approved these changes Aug 12, 2025

View reviewed changes

pytorchmergebot added the merging label Aug 12, 2025

pytorchmergebot closed this in 3cec82a Aug 12, 2025

pytorchmergebot added Merged and removed merging labels Aug 12, 2025

Ensure outer aliasing on DTensor matches inner aliasing #158954

Ensure outer aliasing on DTensor matches inner aliasing #158954

Conversation

ezyang commented Jul 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 23, 2025 • 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/158954

✅ No Failures

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 Jul 24, 2025

Uh oh!

ezyang commented Aug 10, 2025

Uh oh!

ezyang commented Aug 10, 2025

Uh oh!

pytorchmergebot commented Aug 10, 2025

Merge started

Uh oh!

pytorchmergebot commented Aug 10, 2025

Merge failed

Uh oh!

ezyang commented Aug 12, 2025

Uh oh!

ezyang commented Aug 12, 2025

Uh oh!

ezyang commented Aug 12, 2025

Uh oh!

ezyang commented Aug 12, 2025

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!

wconstab left a comment

Choose a reason for hiding this comment

Uh oh!

ezyang commented Aug 12, 2025

Uh oh!

pytorchmergebot commented Aug 12, 2025

Merge started

Uh oh!

Uh oh!

ezyang commented Jul 23, 2025 •

edited

Loading

pytorch-bot bot commented Jul 23, 2025 •

edited

Loading