pytorch / pytorch/pytorch

[DTensor] sharded view incorrectly passes when redistribution is needed

Open
#179,502 1 comment 1 reaction 0 assignees View on GitHub
bot-triaged high priority module: dtensor oncall: distributed triage review triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

This script performs a view over 4 GPUS: `x: dt[12, 8]: S(1); x.view(16, 6)` gives a `S(0)` result, when it requires redistribution. As a result, full-tensor numerics are incorrect.

```python
"""
Repro: DTensor view on a dim-1 sharded tensor gives wrong numerics without redistribution.

Run with: torchrun --nproc_per_node=4 agent_space/repro_view_shard.py
"""

import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor, Shard, distribute_tensor, init_device_mesh

def main():
dist.init_process_group(backend="gloo")
rank = dist.get_rank()
device_mesh = init_device_mesh("cpu", (4,))

torch.manual_seed(0)
tensor = torch.randn(12, 8)

dtensor = distribute_tensor(tensor, device_mesh, [Shard(1)])

# Reference: view on the full (global) tensor
ref = tensor.view(-1, 6)

# DTensor path: view on the sharded tensor, then collect
out = dtensor.view(-1, 6).full_tensor()

if rank == 0:
match = torch.allclose(ref, out)
print(f"Shapes: ref={ref.shape}, out={out.shape}")
print(f"Match: {match}")
if not match:
print(f"Max diff: {(ref - out).abs().max().item()}")
print(f"\nRef:\n{ref}")
print(f"\nOut:\n{out}")

dist.destroy_process_group()

if __name__ == "__main__":
main()

```

### Versions

nightly 4/6

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

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.