[Devicemesh] Bug: `child_to_root_mapping` overwritten due to hash collision in DeviceMesh submeshes
- 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
Assessment
This issue has not been assessed yet.