DeviceMesh BC fix (#387)

fix BC issues There's another pipeline bc issue :(
pytorch · Jun 9, 2024 · ccf1ed8 · ccf1ed8
1 parent 104bd6c
commit ccf1ed8
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/torchtitan/utils.py b/torchtitan/utils.py
@@ -70,9 +70,7 @@ def set_pg_timeouts(timeout, world_mesh):
     torch.distributed.barrier()
     torch.cuda.synchronize()
 
-    groups = (
-        [world_mesh.get_group()] if world_mesh.ndim == 1 else world_mesh.get_group()
-    )
+    groups = [world_mesh.get_group(mesh_dim) for mesh_dim in range(world_mesh.ndim)]
 
     # None represents the 'default' PG, not part of the mesh
     groups.append(None)