Back out "[DeviceMesh] Add support for nD slicing (pytorch#119752)" (pytorch#121763)

Summary:
Original commit changeset: e52b8809c8d8

Original Phabricator Diff: D54778906

We have to backout this diff.
D54778906 seems to be causing test failures for APF blocking trunk health and hence release. Just starting to look at the issue. T182209248

Test Plan: Sandcastle

Reviewed By: satgera

Differential Revision: D54825114

Pull Request resolved: pytorch#121763
Approved by: https://github.com/osalpekar

Author: wz337

test/distributed/test_device_mesh.py

@@ -28,7 +28,6 @@
 )
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     DTensorTestBase,
-    skip_if_lt_x_gpu,
     skip_unless_torch_gpu,
     with_comms,
 )
@@ -79,10 +78,7 @@ def test_assert_invalid_mesh_tensor(self):
 
     @with_comms
     @run_with_both_funcol_impls
-    @skip_unless_torch_gpu
     def test_get_group(self):
-        # TODO: `test_get_group` still periodically timeout on cpu
-        # remove `@skip_unless_torch_gpu` after the problem is fixed.
         mesh_shape = (2, self.world_size // 2)
         mesh_2d = init_device_mesh(
             self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")
@@ -103,10 +99,7 @@ def test_get_group(self):
 
     @with_comms
     @run_with_both_funcol_impls
-    @skip_unless_torch_gpu
     def test_get_local_rank_raises_exception(self):
-        # TODO: `test_get_local_rank_raises_exception` still periodically timeout on cpu
-        # remove `@skip_unless_torch_gpu` after the problem is fixed.
         mesh_shape = (2, self.world_size // 2)
         mesh_2d = init_device_mesh(
             self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")
@@ -120,10 +113,7 @@ def test_get_local_rank_raises_exception(self):
 
     @with_comms
     @run_with_both_funcol_impls
-    @skip_unless_torch_gpu
     def test_get_local_rank(self):
-        # TODO: `test_get_local_rank_raises_exception` still periodically timeout on cpu
-        # remove `@skip_unless_torch_gpu` after the problem is fixed.
         mesh_shape = (2, self.world_size // 2)
         mesh_2d = init_device_mesh(
             self.device_type, mesh_shape, mesh_dim_names=("dp", "tp")
@@ -276,71 +266,47 @@ def world_size(self):
 
     @with_comms
     @run_with_both_funcol_impls
-    def test_raises_invalid_mesh_dim_names(self):
-        error_msg = "Invalid mesh_dim_name"
-        # Case 1: the DeviceMesh does not have a mesh_dim_names attribute
-        with self.assertRaisesRegex(
-            RuntimeError, "Cannot slice a DeviceMesh without mesh_dim_names."
-        ):
+    def test_raises_no_mesh_dim_found(self):
+        with self.assertRaisesRegex(KeyError, "No `mesh_dim_names` found."):
             mesh = init_device_mesh(self.device_type, (2, 4))
             child_mesh = mesh["DP"]
 
-        child_mesh_dim_names = "PP"
-        with self.assertRaisesRegex(ValueError, error_msg):
-            mesh_dim_names = ("DP", "TP")
-            mesh = init_device_mesh(
-                self.device_type, (2, 4), mesh_dim_names=mesh_dim_names
-            )
-            child_mesh = mesh[child_mesh_dim_names]
-
-        # Case 2
-        child_mesh_dim_names = ["PP", "CP"]
-        with self.assertRaisesRegex(ValueError, error_msg):
-            mesh_dim_names = ("DP", "TP")
-            mesh = init_device_mesh(
-                self.device_type, (2, 4), mesh_dim_names=mesh_dim_names
-            )
-            child_mesh = mesh[child_mesh_dim_names]
-
-        # Case 3: a given child_mesh_dim_name is not a contiguous subset of the parent mesh's mesh_dim_names.
-        child_mesh_dim_names = ("TP", "DP")
-        with self.assertRaisesRegex(ValueError, error_msg):
+    @with_comms
+    @run_with_both_funcol_impls
+    def test_raises_invalid_mesh_dim_name(self):
+        child_mesh_dim_name = "PP"
+        with self.assertRaisesRegex(
+            KeyError, f"Mesh dimension '{child_mesh_dim_name}' does not exist."
+        ):
             mesh_dim_names = ("DP", "TP")
             mesh = init_device_mesh(
                 self.device_type, (2, 4), mesh_dim_names=mesh_dim_names
             )
-            child_mesh = mesh[child_mesh_dim_names]
-
-        # Case 3
-        child_mesh_dim_names = ("PP", "TP")
-        with self.assertRaisesRegex(ValueError, error_msg):
-            mesh_dim_names = ("PP", "DP", "TP")
-            mesh = init_device_mesh(
-                self.device_type, (2, 2, 2), mesh_dim_names=mesh_dim_names
-            )
-            child_mesh = mesh[child_mesh_dim_names]
+            child_mesh = mesh[child_mesh_dim_name]
 
     @with_comms
     @run_with_both_funcol_impls
-    @skip_if_lt_x_gpu(8)
-    def test_get_item_2d(self):
-        # TODO: `test_get_item_2d` still periodically timeout on cpu
-        # remove `@skip_if_lt_x_gpu` after the problem is fixed.
+    def test_get_item(self):
         mesh_shape = (2, 4)
         mesh_dim_names = ("DP", "TP")
         mesh_2d = init_device_mesh(
             self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
         )
 
+        pg_ranks_by_dim_name = {}
+        for mesh_dim_name in mesh_dim_names:
+            mesh_dim = mesh_dim_names.index(mesh_dim_name)
+            pg_ranks_by_dim_name[mesh_dim_name] = mesh_2d.mesh.swapdims(
+                -1, mesh_dim
+            ).reshape(-1, mesh_2d.mesh.size(mesh_dim))
+
         tp_mesh = mesh_2d["TP"]
-        tp_group = [[0, 1, 2, 3], [4, 5, 6, 7]]
         tp_group_idx = self.rank // 4
-        self.assertEqual(tp_mesh.mesh.tolist(), tp_group[tp_group_idx])
+        self.assertEqual(tp_mesh.mesh, pg_ranks_by_dim_name["TP"][tp_group_idx])
 
         dp_mesh = mesh_2d["DP"]
-        dp_group = [[0, 4], [1, 5], [2, 6], [3, 7]]
         dp_group_idx = self.rank % 4
-        self.assertEqual(dp_mesh.mesh.tolist(), dp_group[dp_group_idx])
+        self.assertEqual(mesh_2d["DP"].mesh, pg_ranks_by_dim_name["DP"][dp_group_idx])
 
     @with_comms
     @run_with_both_funcol_impls
@@ -351,50 +317,14 @@ def test_get_item_1d(self):
         dp_mesh = mesh["dp"]
         self.assertEqual(dp_mesh, mesh)
 
-        with self.assertRaisesRegex(ValueError, "Invalid mesh_dim_name"):
+        with self.assertRaisesRegex(RuntimeError, "Invalid mesh_dim_name"):
             dp_mesh = mesh["dim0"]
 
-    @with_comms
-    @skip_if_lt_x_gpu(8)
-    def test_get_item_3d(self):
-        # TODO: `test_get_item_3d` still periodically timeout on cpu
-        # remove `@skip_if_lt_x_gpu` after the problem is fixed.
-        mesh_shape = (2, 2, 2)
-        mesh_dim_names = ("Replicate", "Shard", "TP")
-        mesh_3d = init_device_mesh(
-            self.device_type, mesh_shape, mesh_dim_names=mesh_dim_names
-        )
-
-        tp_group = [[0, 1], [2, 3], [4, 5], [6, 7]]
-        tp_group_idx = int(self.rank / 2)
-        self.assertEqual(mesh_3d["TP"].mesh.tolist(), tp_group[tp_group_idx])
-
-        shard_group = [[0, 2], [1, 3], [4, 6], [5, 7]]
-        shard_group_idx = self.rank % 2 + self.rank // 4 * 2
-        self.assertEqual(mesh_3d["Shard"].mesh.tolist(), shard_group[shard_group_idx])
-
-        replicate_group = [[0, 4], [1, 5], [2, 6], [3, 7]]
-        replicate_group_idx = self.rank % 4
-        self.assertEqual(
-            mesh_3d["Replicate"].mesh.tolist(), replicate_group[replicate_group_idx]
-        )
-
-        # We support both UX for nD slicing.
-        # mesh_3d[["Replicate", "Shard"]] or mesh_3d["Replicate", "Shard"]
-        hsdp_mesh_1 = mesh_3d[["Replicate", "Shard"]]
-        hsdp_mesh_2 = mesh_3d["Replicate", "Shard"]
-        hsdp_group = [[[0, 2], [4, 6]], [[1, 3], [5, 7]]]
-        hsdp_group_idx = self.rank % 2
-        self.assertEqual(hsdp_mesh_1.mesh.tolist(), hsdp_group[hsdp_group_idx])
-        self.assertEqual(hsdp_mesh_2.mesh.tolist(), hsdp_group[hsdp_group_idx])
-        self.assertEqual(hsdp_mesh_1, hsdp_mesh_2)
-
 
 @instantiate_parametrized_tests
 class TestMeshEnv(DTensorTestBase):
     @with_comms
     @run_with_both_funcol_impls
-    @skip_unless_torch_gpu
     def test_get_parent_mesh(self):
         mesh_shape = (2, self.world_size // 2)
         mesh_dim_names = ("DP", "TP")
@@ -415,7 +345,6 @@ def test_get_parent_mesh(self):
 
     @with_comms
     @run_with_both_funcol_impls
-    @skip_unless_torch_gpu
     def test_get_parent_mesh_dim_exist(self):
         mesh_shape = (2, self.world_size // 2)
         mesh_dim_names = ("DP", "TP")
@@ -428,7 +357,6 @@ def test_get_parent_mesh_dim_exist(self):
 
     @with_comms
     @run_with_both_funcol_impls
-    @skip_unless_torch_gpu
     def test_get_parent_mesh_dim_not_exist(self):
         mesh_shape = (self.world_size,)
         mesh = init_device_mesh(self.device_type, mesh_shape)
@@ -437,7 +365,6 @@ def test_get_parent_mesh_dim_not_exist(self):
 
     @with_comms
     @run_with_both_funcol_impls
-    @skip_unless_torch_gpu
     def test_get_mesh_dim_by_name(self):
         mesh_shape = (2, self.world_size // 2)
         mesh_dim_names = ("DP", "TP")

torch/distributed/device_mesh.py

@@ -67,43 +67,25 @@ def get_current_mesh(self) -> "DeviceMesh":
             return self.mesh_stack[-1]
 
         def create_child_mesh(
-            self,
-            device_mesh: "DeviceMesh",
-            mesh_dim_names: Tuple[str],
+            self, device_mesh: "DeviceMesh", mesh_dim: int, mesh_dim_name: str
         ) -> "DeviceMesh":
             # swap the current dim to the last dim then reshape to flatten out other
             # dims, so we can just extract the list of ranks which contains cur_rank.
-            mesh_dims = [
-                not_none(device_mesh.mesh_dim_names).index(mesh_dim_name)
-                for mesh_dim_name in mesh_dim_names
-            ]
             cur_rank = device_mesh.get_rank()
-            mesh = device_mesh.mesh
-            all_mesh_dims = list(range(mesh.ndim))
-            for mesh_dim in mesh_dims:
-                # remove not pop b/c we want the value of the ind removed not it's position in the list
-                # because this list dynamically changes.
-                all_mesh_dims.remove(mesh_dim)
-
-            mesh_sizes = [device_mesh.mesh.size(mesh_dim) for mesh_dim in mesh_dims]
-
-            pg_ranks_by_dim = device_mesh.mesh.permute(
-                *all_mesh_dims, *mesh_dims
-            ).reshape(-1, *mesh_sizes)
-
-            for mesh_nd in pg_ranks_by_dim:
-                if cur_rank in mesh_nd:
-                    sub_mesh = DeviceMesh(
-                        device_mesh.device_type,
-                        mesh_nd,
-                        mesh_dim_names=mesh_dim_names,
-                    )
-                    res_sub_mesh = sub_mesh
+            pg_ranks_by_dim = device_mesh.mesh.swapdims(-1, mesh_dim).reshape(
+                -1, device_mesh.mesh.size(mesh_dim)
+            )
 
-            res_sub_mesh._dim_group_infos = [  # type: ignore[possibly-undefined]
-                device_mesh._dim_group_infos[mesh_dim] for mesh_dim in mesh_dims
-            ]
+            for mesh_1d in pg_ranks_by_dim:
+                sub_mesh = DeviceMesh(
+                    device_mesh.device_type,
+                    mesh_1d,
+                    mesh_dim_names=(mesh_dim_name,),
+                )
+                if cur_rank in mesh_1d:
+                    res_sub_mesh = sub_mesh
 
+            res_sub_mesh._dim_group_infos = [device_mesh._dim_group_infos[mesh_dim]]  # type: ignore[possibly-undefined]
             # Assign the current DeviceMesh as the parent of the child DeviceMesh.
             self.child_to_parent_mapping[res_sub_mesh] = device_mesh
             return res_sub_mesh
@@ -228,7 +210,6 @@ def __init__(
             # private field to pre-generate DeviceMesh's hash
             self._flatten_mesh_list = tuple(self.mesh.flatten().tolist())
             self._hash = hash((self._flatten_mesh_list, self.mesh.shape, id(self)))
-            self._parent_mesh = _mesh_resources.get_parent_mesh(self)
 
             # Skip process group initialization if xla device.
             # TODO(yeounoh) implement DeviceMesh backend and register XLA backend.
@@ -237,8 +218,7 @@ def __init__(
                 # already. The world pg is used for device mesh identity (rank) on each
                 # process (we need to know if the current global rank is in the mesh or not).
                 self._get_or_create_default_group()
-                if not self._parent_mesh:
-                    self._init_process_groups()
+                self._init_process_groups()
 
         def _get_or_create_default_group(self):
             default_initialized = is_initialized()
@@ -360,7 +340,7 @@ def __eq__(self, other: object) -> bool:
                 and self._flatten_mesh_list == other._flatten_mesh_list
             )
 
-        def __getitem__(self, mesh_dim_names: Union[str, Tuple[str]]) -> "DeviceMesh":
+        def __getitem__(self, mesh_dim_name: str) -> "DeviceMesh":
             """
             Slice the current DeviceMesh based on the mesh_dim_name given to create a child
             DeviceMesh.
@@ -388,39 +368,17 @@ def __getitem__(self, mesh_dim_names: Union[str, Tuple[str]]) -> "DeviceMesh":
                 >>> # of cross-host(dim 0), and within-host (dim 1).
                 >>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]])
             """
-            if not self.mesh_dim_names:
-                raise RuntimeError("Cannot slice a DeviceMesh without mesh_dim_names.")
-
-            mesh_dim_names = (
-                (mesh_dim_names,) if isinstance(mesh_dim_names, str) else mesh_dim_names
-            )
-
-            error_msg = (
-                f"Invalid mesh_dim_name {mesh_dim_names} specified. "
-                f"Valid mesh_dim_names should be a contiguous subsequence of {self.mesh_dim_names}."
-            )
+            if self.mesh.ndim == 1:
+                if self.mesh_dim_names and mesh_dim_name == self.mesh_dim_names[0]:
+                    return self
+                else:
+                    raise RuntimeError(
+                        f"Invalid mesh_dim_name {mesh_dim_name} specified."
+                    )
 
-            # When the dimension slicing out is equal to the mesh dimensions of the current DeviceMesh,
-            # we simply return self if the given slicing is valid.
-            if mesh_dim_names == self.mesh_dim_names:
-                return self
-            # Check if the user-provided slicing is a valid contiguous subsequence of the mesh_dim_names
-            # of the current DeviceMesh.
-            elif len(mesh_dim_names) < len(self.mesh_dim_names):
-                outermost_dim_name = mesh_dim_names[0]
-                if outermost_dim_name not in self.mesh_dim_names:
-                    raise ValueError(error_msg)
-                outermost_dim_idx = self.mesh_dim_names.index(outermost_dim_name)
-                for i, j in zip(
-                    mesh_dim_names,
-                    self.mesh_dim_names[outermost_dim_idx : len(mesh_dim_names)],
-                ):
-                    if i != j:
-                        raise ValueError(error_msg)
-            else:
-                raise ValueError(error_msg)
+            mesh_dim = _mesh_resources.get_mesh_dim_by_name(self, mesh_dim_name)
+            submesh = _mesh_resources.create_child_mesh(self, mesh_dim, mesh_dim_name)
 
-            submesh = _mesh_resources.create_child_mesh(self, mesh_dim_names)
             return submesh
 
         def get_group(