[DeviceMesh] Initialized mesh tensor with CPU context

Andrew Gu · Andrew Gu · commit 57ee17e4feec · 2024-04-23T12:45:26.000-07:00
[ghstack-poisoned]
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_init.py b/test/distributed/_composable/fsdp/test_fully_shard_init.py
@@ -487,6 +487,19 @@ def test_meta_device_1d_init(self):
                 self.assertEqual(param.device, torch.device("meta"))
             self._test_to_empty_and_reset_parameters(model, mesh, mlp_dim)
 
+        # Test that we can call `fully_shard` under meta-device context and
+        # that `init_device_mesh` call still works
+        mlp_dim = 8
+        with torch.device("meta"):
+            model = nn.Sequential(MLP(mlp_dim, with_buffer=True), MLP(mlp_dim))
+            for param in model.parameters():
+                self.assertEqual(param.device, torch.device("meta"))
+            for module in (model[0], model[1], model):
+                fully_shard(module)
+        for param in model.parameters():
+            self.assertEqual(param.device, torch.device("meta"))
+        self._test_to_empty_and_reset_parameters(model, mesh, mlp_dim)
+
     @unittest.skipIf(not TEST_CUDA, "no cuda")
     def test_meta_device_2d_init(self):
         assert self.world_size >= 4, f"{self.world_size}"
diff --git a/torch/distributed/device_mesh.py b/torch/distributed/device_mesh.py
@@ -553,7 +553,10 @@ def init_device_mesh(
                     f"Found len(mesh_dim_names): {len(mesh_dim_names)} and len(mesh_shape):{len(mesh_shape)}.",
                 )
 
-        mesh = torch.arange(math.prod(mesh_shape), dtype=torch.int).view(mesh_shape)
+        # Always initialize the mesh's tensor on CPU, regardless of what the
+        # external device type has been set to be (e.g. meta)
+        with torch.device("cpu"):
+            mesh = torch.arange(math.prod(mesh_shape), dtype=torch.int).view(mesh_shape)
         device_mesh = DeviceMesh(
             device_type=device_type,
             mesh=mesh,
