Back out "[DeviceMesh] Only include the real thread_id in DeviceMesh hash under threaded backend (#130495)" (#130676)
Summary:
Original commit changeset: 80c2ca639146
Original Phabricator Diff: D59612200
Test Plan: buck2 test 'fbcode//mode/opt' fbcode//apf/distributed/tests:pipeline_parallel_test_cpu -- --exact 'apf/distributed/tests:pipeline_parallel_test_cpu - apf.distributed.tests.pipeline_parallel_test_cpu.PipelineParallelContextTestCPU: test_stage_pg_creation_with_different_backends'
Differential Revision: D59719562
Pull Request resolved: https://github.com/pytorch/pytorch/pull/130676
Approved by: https://github.com/xunnanxu
diff --git a/torch/distributed/device_mesh.py b/torch/distributed/device_mesh.py
index 4d20186..0d2f26b 100644
--- a/torch/distributed/device_mesh.py
+++ b/torch/distributed/device_mesh.py
@@ -38,7 +38,6 @@
_find_pg_by_ranks_and_tag,
_get_default_group,
_get_group_tag,
- get_backend,
get_process_group_ranks,
get_rank,
get_world_size,
@@ -243,7 +242,7 @@
# private field to pre-generate DeviceMesh's hash
self._flatten_mesh_list = tuple(self.mesh.flatten().tolist())
self._parent_mesh: Optional[DeviceMesh] = None
- self._thread_id = None
+ self._thread_id = threading.get_ident()
# Skip process group initialization if xla device or init backend is False
# TODO(yeounoh) implement DeviceMesh backend and register XLA backend.
@@ -255,9 +254,6 @@
self._get_or_create_default_group()
self._init_process_groups()
- if get_backend() == "threaded":
- self._thread_id = threading.get_ident()
-
# calculate the coordinates of the current global rank on the mesh
rank_coords = (self.mesh == get_rank()).nonzero()
assert rank_coords.size(0) in (0, 1)
