frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[DeviceMesh] Use user set backend and pg option even for the global mesh (#157501)

Browse Source 
Short term solution to https://github.com/pytorch/pytorch/issues/156593.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/157501
Approved by: https://github.com/fegin, https://github.com/lw
This commit is contained in:
fduwjj
2025-07-02 14:30:55 -07:00
committed by PyTorch MergeBot
parent bf1ebe0531
commit 2b8d3b1b2b
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/distributed/device_mesh.py
View File

@ -529,7 +529,11 @@ else:
dim_group_names: list[str] = []
default_group = _get_default_group()
if self.mesh.ndim == 1 and self.mesh.numel() == get_world_size():
if (
self.mesh.ndim == 1
and self.mesh.numel() == get_world_size()
and 0 not in _mesh_resources.mesh_dim_group_options
):
# Append the default pg to the first dim groups only if the default pg is compatible with `self.device_type`.
# Otherwise, create new pg.
ranks = list(range(get_world_size()))