pytorch / pytorch/pytorch

[Devicemesh] Bug: `child_to_root_mapping` overwritten due to hash collision in DeviceMesh submeshes

Open
#173,789 5 comments 0 reactions 1 assignee Claimed by @arkadip-maitra View on GitHub
bot-triaged oncall: distributed oncall: distributed infra ptd-bot-triaged triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

When slicing multiple `DeviceMesh` instances that have the same shape and dimension names, the resulting submeshes can have identical hash values, causing entries in `child_to_root_mapping` to be overwritten. This leads to incorrect root mesh lookups via `get_root_mesh()`.

The `DeviceMesh.__hash__()` method computes the hash based on:
- `_flatten_mesh_list`
- `mesh.shape`
- `device_type`
- `mesh_dim_names`
- `_thread_id`

When two different parent meshes are sliced to create submeshes with identical properties (same shape, device type, dimension names, etc.), they produce the same hash value, breaking the parent-child relationship tracking.

## Reproduce
```py
# torchrun --nproc_per_node=4 ./issues/pytorch/device_mesh.py

import loguru
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
from torch.distributed.tensor.device_mesh import _mesh_resources

# Slicing again will overwrite child_to_root_mapping. since the hash of the two submeshes are identical.
slice_again = True

mesh1 = init_device_mesh(device_type='cuda', mesh_shape=(2, 2,), mesh_dim_names=('cp', 'tp',))
mesh2 = init_device_mesh(device_type='cuda', mesh_shape=(2, 2,), mesh_dim_names=('cp', 'tp',))

tp_mesh1 = mesh1['tp']
if slice_again:
tp_mesh2 = mesh2['tp']
loguru.logger.info(f'hash of tp_mesh1: {hash(tp_mesh1)}, hash of tp_mesh2: {hash(tp_mesh2)}')

assert _mesh_resources.get_root_mesh(tp_mesh1) is mesh1
```

### Versions

2.9.1

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @H-Huang

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.