Issue Station
Sign In
Tools hereToolbox

pytorch - ✅(Solved) Fix DTensor indexing aliasing bug - causes overwrite of original dtensor spec [1 pull requests, 2 comments, 2 participants]

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…
GitHub stats
pytorch/pytorch#179290•Fetched 2026-04-08 02:43:39
View on GitHub
Comments
2
Participants
2
Timeline
49
Reactions
0
Author
jmaccarljmaccarl
Participants
aditvenkaditvenkjmaccarljmaccarl
Timeline (top)
mentioned ×20subscribed ×20labeled ×4commented ×2

Fix Action

Fix / Workaround

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 46 bits physical, 57 bits virtual Byte Order: Little Endian CPU(s): 128 On-line CPU(s) list: 0-127 Vendor ID: GenuineIntel Model name: Intel(R) Xeon(R) Platinum 8462Y+ CPU family: 6 Model: 143 Thread(s) per core: 2 Core(s) per socket: 32 Socket(s): 2 Stepping: 8 CPU max MHz: 4100.0000 CPU min MHz: 800.0000 BogoMIPS: 5600.00 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 invpcid_single cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts hfi vnmi avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr ibt amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities Virtualization: VT-x L1d cache: 3 MiB (64 instances) L1i cache: 2 MiB (64 instances) L2 cache: 128 MiB (64 instances) L3 cache: 120 MiB (2 instances) NUMA node(s): 2 NUMA node0 CPU(s): 0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52,54,56,58,60,62,64,66,68,70,72,74,76,78,80,82,84,86,88,90,92,94,96,98,100,102,104,106,108,110,112,114,116,118,120,122,124,126 NUMA node1 CPU(s): 1,3,5,7,9,11,13,15,17,19,21,23,25,27,29,31,33,35,37,39,41,43,45,47,49,51,53,55,57,59,61,63,65,67,69,71,73,75,77,79,81,83,85,87,89,91,93,95,97,99,101,103,105,107,109,111,113,115,117,119,121,123,125,127 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Retbleed: Not affected Vulnerability Spec rstack overflow: Not affected Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

PR fix notes

PR #179318: [DTensor] Fix select_int_strategy mutating input DTensorSpec

Description (problem / solution / changelog)

The output spec was aliased to the input spec (output_specs = input_specs) for the non-sharded case. When sharding propagation later updated tensor_meta on the output spec, it mutated the original DTensor's spec, causing repeated indexing (e.g. dt[0]; dt[1]) to progressively reduce the DTensor's dimensions.

Fix by using dataclasses.replace to create a fresh output spec.

Fixes https://github.com/pytorch/pytorch/issues/179290

<!-- ps-id: aeb77822-4dc5-407b-b5aa-7b240e6e6275 -->

Changed files

  • test/distributed/tensor/test_tensor_ops.py (modified, +16/-0)
  • torch/distributed/tensor/_ops/_tensor_ops.py (modified, +2/-1)

Code Example

import os
import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.device_mesh import DeviceMesh
from torch.distributed.tensor.placement_types import Replicate


def main(rank, world_size, master_port):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = str(master_port)
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    mesh = DeviceMesh("cuda", torch.arange(world_size), mesh_dim_names=("dp",))

    if rank == 0:
        # Test A: call dt[1]
        local_w = torch.randn(4, 32, 64, device="cuda")
        dt = DTensor.from_local(local_w.clone(), device_mesh=mesh,
                                placements=[Replicate()], run_check=False)
        print("=== Test A: dt[1] ===")
        print(f"  dt spec before: {dt._spec}")
        s = dt[1]
        print(f"  dt[1] shape={s.shape}, dim={s.dim()}")
        print(f"  dt spec after: {dt._spec}")

        # Test B: call dt[0] then dt[1] then dt[0] again
        dt2 = DTensor.from_local(local_w.clone(), device_mesh=mesh,
                                 placements=[Replicate()], run_check=False)
        print("=== Test B: dt[0] then dt[1] then dt[0] again ===")
        print(f"  dt spec before: {dt2._spec}")
        s0 = dt2[0]
        print(f"  dt[0] shape={s0.shape}, dim={s0.dim()}")
        s1 = dt2[1]
        print(f"  dt[1] shape={s1.shape}, dim={s1.dim()}")
        s0c = dt2[0]
        print(f"  dt[0] shape={s0c.shape}, dim={s0c.dim()}")
        print(f"  dt spec after: {dt2._spec}")

    dist.destroy_process_group()


if __name__ == "__main__":
    from torch.multiprocessing import spawn
    import socket
    s = socket.socket(); s.bind(('', 0)); port = s.getsockname()[1]; s.close()
    spawn(main, args=(2, port), nprocs=2, join=True)

---

=== Test A: dt[1] ===
  dt spec before: Spec(R on (4, 32, 64))
  dt[1] shape=torch.Size([32, 64]), dim=2
  dt spec after: Spec(R on (32, 64))
=== Test B: dt[0] then dt[1] then dt[0] again ===
  dt spec before: Spec(R on (4, 32, 64))
  dt[0] shape=torch.Size([32, 64]), dim=2
  dt[1] shape=torch.Size([64]), dim=1
  dt[0] shape=torch.Size([]), dim=0
  dt spec after: Spec(R on ())
RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

Hi - first issue here 👋

I'm encountering unexpected modifications to DTensor dims when indexing into them. Here's a simple repro:

import os
import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.device_mesh import DeviceMesh
from torch.distributed.tensor.placement_types import Replicate


def main(rank, world_size, master_port):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = str(master_port)
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    mesh = DeviceMesh("cuda", torch.arange(world_size), mesh_dim_names=("dp",))

    if rank == 0:
        # Test A: call dt[1]
        local_w = torch.randn(4, 32, 64, device="cuda")
        dt = DTensor.from_local(local_w.clone(), device_mesh=mesh,
                                placements=[Replicate()], run_check=False)
        print("=== Test A: dt[1] ===")
        print(f"  dt spec before: {dt._spec}")
        s = dt[1]
        print(f"  dt[1] shape={s.shape}, dim={s.dim()}")
        print(f"  dt spec after: {dt._spec}")

        # Test B: call dt[0] then dt[1] then dt[0] again
        dt2 = DTensor.from_local(local_w.clone(), device_mesh=mesh,
                                 placements=[Replicate()], run_check=False)
        print("=== Test B: dt[0] then dt[1] then dt[0] again ===")
        print(f"  dt spec before: {dt2._spec}")
        s0 = dt2[0]
        print(f"  dt[0] shape={s0.shape}, dim={s0.dim()}")
        s1 = dt2[1]
        print(f"  dt[1] shape={s1.shape}, dim={s1.dim()}")
        s0c = dt2[0]
        print(f"  dt[0] shape={s0c.shape}, dim={s0c.dim()}")
        print(f"  dt spec after: {dt2._spec}")

    dist.destroy_process_group()


if __name__ == "__main__":
    from torch.multiprocessing import spawn
    import socket
    s = socket.socket(); s.bind(('', 0)); port = s.getsockname()[1]; s.close()
    spawn(main, args=(2, port), nprocs=2, join=True)

Output

=== Test A: dt[1] ===
  dt spec before: Spec(R on (4, 32, 64))
  dt[1] shape=torch.Size([32, 64]), dim=2
  dt spec after: Spec(R on (32, 64))
=== Test B: dt[0] then dt[1] then dt[0] again ===
  dt spec before: Spec(R on (4, 32, 64))
  dt[0] shape=torch.Size([32, 64]), dim=2
  dt[1] shape=torch.Size([64]), dim=1
  dt[0] shape=torch.Size([]), dim=0
  dt spec after: Spec(R on ())

It doesn't matter what indexes are used, any indexing seems to modify the DTensor dimension -1 each time you use it. I believe this is due to aliasing used in select_int_strategy between the input and output DTensor spec, since making a copy here solved the issue: https://github.com/pytorch/pytorch/blob/168b7f3c3223995cc088bd81e254e6241a94196e/torch/distributed/tensor/_ops/_tensor_ops.py#L368

Versions

PyTorch version: 2.8.0+cu128 Is debug build: False CUDA used to build PyTorch: 12.8 ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64) GCC version: (conda-forge gcc 14.3.0-18) 14.3.0 Clang version: Could not collect CMake version: Could not collect Libc version: glibc-2.35

Python version: 3.12.12 | packaged by conda-forge | (main, Jan 26 2026, 23:51:32) [GCC 14.3.0] (64-bit runtime) Python platform: Linux-6.5.13-65-650-4141-22041-coreweave-amd64-85c45edc-x86_64-with-glibc2.35 Is CUDA available: True CUDA runtime version: 12.8.93 CUDA_MODULE_LOADING set to: LAZY GPU models and configuration: GPU 0: NVIDIA H100 80GB HBM3 GPU 1: NVIDIA H100 80GB HBM3

Nvidia driver version: 580.95.05 cuDNN version: Probably one of the following: /usr/lib/x86_64-linux-gnu/libcudnn.so.9.13.0 /usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.13.0 /usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.13.0 /usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.13.0 /usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.13.0 /usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.13.0 /usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.13.0 /usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.13.0 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 46 bits physical, 57 bits virtual Byte Order: Little Endian CPU(s): 128 On-line CPU(s) list: 0-127 Vendor ID: GenuineIntel Model name: Intel(R) Xeon(R) Platinum 8462Y+ CPU family: 6 Model: 143 Thread(s) per core: 2 Core(s) per socket: 32 Socket(s): 2 Stepping: 8 CPU max MHz: 4100.0000 CPU min MHz: 800.0000 BogoMIPS: 5600.00 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 invpcid_single cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts hfi vnmi avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr ibt amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities Virtualization: VT-x L1d cache: 3 MiB (64 instances) L1i cache: 2 MiB (64 instances) L2 cache: 128 MiB (64 instances) L3 cache: 120 MiB (2 instances) NUMA node(s): 2 NUMA node0 CPU(s): 0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52,54,56,58,60,62,64,66,68,70,72,74,76,78,80,82,84,86,88,90,92,94,96,98,100,102,104,106,108,110,112,114,116,118,120,122,124,126 NUMA node1 CPU(s): 1,3,5,7,9,11,13,15,17,19,21,23,25,27,29,31,33,35,37,39,41,43,45,47,49,51,53,55,57,59,61,63,65,67,69,71,73,75,77,79,81,83,85,87,89,91,93,95,97,99,101,103,105,107,109,111,113,115,117,119,121,123,125,127 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Retbleed: Not affected Vulnerability Spec rstack overflow: Not affected Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

Versions of relevant libraries: [pip3] cuequivariance-ops-torch-cu12==0.9.0 [pip3] cuequivariance-torch==0.9.0 [pip3] msgpack-numpy==0.4.8 [pip3] mypy==1.19.1 [pip3] mypy_extensions==1.1.0 [pip3] nccl-profiler-inspector-cu12==2.29.7 [pip3] numpy==1.26.4 [pip3] nvidia-cublas-cu12==12.8.4.1 [pip3] nvidia-cuda-cupti-cu12==12.8.90 [pip3] nvidia-cuda-nvrtc-cu12==12.8.93 [pip3] nvidia-cuda-runtime-cu12==12.8.90 [pip3] nvidia-cudnn-cu12==9.10.2.21 [pip3] nvidia-cufft-cu12==11.3.3.83 [pip3] nvidia-curand-cu12==10.3.9.90 [pip3] nvidia-cusolver-cu12==11.7.3.90 [pip3] nvidia-cusparse-cu12==12.5.8.93 [pip3] nvidia-cusparselt-cu12==0.7.1 [pip3] nvidia-nccl-cu12==2.29.7 [pip3] nvidia-nvjitlink-cu12==12.8.93 [pip3] nvidia-nvtx-cu12==12.8.90 [pip3] onnx==1.20.1 [pip3] onnx-ir==0.2.0 [pip3] onnxruntime==1.24.2 [pip3] onnxscript==0.6.2 [pip3] pytorch-lightning==2.6.1 [pip3] torch==2.8.0+cu128 [pip3] torch_c_dlpack_ext==0.1.5 [pip3] torch-geometric==2.7.0 [pip3] torchaudio==2.8.0+cu128 [pip3] torchdata==0.11.0 [pip3] torchmetrics==1.8.2 [pip3] torchvision==0.23.0+cu128 [pip3] triton==3.4.0 [conda] Could not collect

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @xmfan @tianyu-l @XilunWu @SherlockNoMad @ppwwyyxx

extent analysis

TL;DR

The issue can be resolved by making a copy of the DTensor spec in the select_int_strategy function to prevent aliasing between the input and output DTensor specs.

Guidance

  • The problem seems to be caused by aliasing in the select_int_strategy function, which can be fixed by making a copy of the DTensor spec.
  • To verify the fix, run the provided test code and check if the DTensor dimensions are modified as expected after indexing.
  • The issue is specific to the PyTorch version 2.8.0+cu128, and it may be resolved in future versions.
  • To mitigate the issue, users can try making a copy of the DTensor spec in their own code before indexing into it.

Example

No code example is provided as the fix is specific to the PyTorch internal implementation.

Notes

The issue is caused by a specific implementation detail in PyTorch, and the fix may not be applicable to all users. The provided test code can be used to verify the issue and the fix.

Recommendation

Apply the workaround by making a copy of the DTensor spec in the select_int_strategy function, as this is the most straightforward way to resolve the issue without waiting for a future PyTorch version that may include the fix.

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…
#api#ssr#installation#tensor shape#request timeout

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING

pytorch - ✅(Solved) Fix DTensor indexing aliasing bug - causes overwrite of original dtensor spec [1 pull requests, 2 comments, 2 participants]