Update on "[device_mesh] add a private init backend option"

wanchaol · wanchaol · commit 08f0bae8710f · 2024-04-23T16:00:34.000-07:00
This PR adds a private init backend option, to tackle the issues sub
mesh creation:

in device mesh slicing we don't want to create process groups again,
so explicitly turn the group creation off it's useful

Also I think there might be more submesh creation functionality so
having this flag would ensure that there's no new group created

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k

[ghstack-poisoned]
diff --git a/test/distributed/test_device_mesh.py b/test/distributed/test_device_mesh.py
@@ -150,8 +150,9 @@ def test_device_mesh_init_backend(self):
         with self.assertRaisesRegex(RuntimeError, "process groups not initialized!"):
             mesh.get_group()
 
-        with self.assertRaisesRegex(AttributeError, "no attribute"):
-            mesh.get_coordinate()
+        # coordinates should always been populated when init_backend is False, as whenever
+        # we call init_backend we should make sure the default pg already created
+        mesh.get_coordinate()
 
     def test_fake_pg_device_mesh(self):
         fake_store = FakeStore()
