Update

[ghstack-poisoned]

Author: fegin

torchft/fsdp_test.py

@@ -0,0 +1,70 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Tuple
+from unittest import skipUnless, TestCase
+from unittest.mock import Mock
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch._C._distributed_c10d import (
+    _resolve_process_group,
+    AllgatherOptions,
+    AllreduceOptions,
+    BroadcastOptions,
+    ReduceOp,
+)
+from torch.distributed import (
+    _functional_collectives,
+    get_world_size,
+    ReduceOp,
+    TCPStore,
+    Work,
+)
+from torch.distributed._composable.fsdp import fully_shard
+from torch.distributed.device_mesh import init_device_mesh
+from torch.testing._internal.common_distributed import MultiProcessTestCase
+
+from torchft.manager import Manager
+from torchft.process_group import ft_init_device_mesh, ManagedProcessGroup
+
+
+class FSDPTest(MultiProcessTestCase):
+    @property
+    def world_size(self):
+        return 4
+
+    def setUp(self):
+        super().setUp()
+        os.environ["TORCH_NCCL_DESYNC_DEBUG"] = "0"
+        self._spawn_processes()
+
+    def test_fsdp(self) -> None:
+        group_size = self.world_size // 2
+        group = self.rank // group_size
+        group_rank = self.rank % group_size
+
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = str(12346 + group)
+        os.environ["RANK"] = str(group_rank)
+        os.environ["WORLD_SIZE"] = str(group_size)
+
+        manager = Mock(spec=Manager)
+        device_mesh = ft_init_device_mesh(
+            device_type="cuda",
+            mesh_shape=(2, 2),
+            mesh_dim_names=("dp_replicate", "dp_shard"),
+            replicate_dim=0,
+            manager=manager,
+        )
+        manager.num_participants.return_value = 1
+        model = nn.Linear(128, 128).cuda()
+        batch = torch.randn(4, 128).cuda()
+        shard_model = fully_shard(model, mesh=device_mesh)
+        shard_model(batch).mean().backward()

torchft/process_group.py

@@ -250,9 +250,9 @@ class ProcessGroupGloo(ProcessGroupWrapper):
     This is a reconfigurable version of ProcessGroupGloo.
     """
 
-    PG_CLASS: Type[
-        BaseProcessGroup
-    ] = BaseProcessGroupGloo  # pyre-fixme[16]: no attribute ProcessGroupGloo
+    PG_CLASS: Type[BaseProcessGroup] = (
+        BaseProcessGroupGloo  # pyre-fixme[16]: no attribute ProcessGroupGloo
+    )
 
     def getBackendName(self) -> str:
         return "torchft-gloo"
@@ -269,9 +269,9 @@ class ProcessGroupNCCL(ProcessGroupWrapper):
     abort when reconfiguring, we need to ensure this is safe.
     """
 
-    PG_CLASS: Type[
-        BaseProcessGroup
-    ] = BaseProcessGroupNCCL  # pyre-fixme[16]: no attribute ProcessGroupNCCL
+    PG_CLASS: Type[BaseProcessGroup] = (
+        BaseProcessGroupNCCL  # pyre-fixme[16]: no attribute ProcessGroupNCCL
+    )
 
     def getBackendName(self) -> str:
         return "torchft-nccl"
@@ -745,9 +745,9 @@ class ProcessGroupBabyGloo(ProcessGroupBaby):
     ProcessGroupBabyNCCL.
     """
 
-    PG_CLASS: Type[
-        BaseProcessGroup
-    ] = BaseProcessGroupGloo  # pyre-fixme[16]: no attribute ProcessGroupGloo
+    PG_CLASS: Type[BaseProcessGroup] = (
+        BaseProcessGroupGloo  # pyre-fixme[16]: no attribute ProcessGroupGloo
+    )
 
     def getBackendName(self) -> str:
         return "torchft-baby-gloo"
@@ -769,9 +769,9 @@ class ProcessGroupBabyNCCL(ProcessGroupBaby):
     tensors may leak in the current PyTorch implementation. TODO fix
     """
 
-    PG_CLASS: Type[
-        BaseProcessGroup
-    ] = BaseProcessGroupNCCL  # pyre-fixme[16]: no attribute ProcessGroupNCCL
+    PG_CLASS: Type[BaseProcessGroup] = (
+        BaseProcessGroupNCCL  # pyre-fixme[16]: no attribute ProcessGroupNCCL
+    )
     WORK_CLASS = _BabyWorkNCCL
 
     def getBackendName(self) -> str:
@@ -807,27 +807,34 @@ def extend_device_mesh(
     )
 
 
-class ManagedDeviceMesh(DeviceMesh):
+class _ManagedDeviceMesh(DeviceMesh):
     def __init__(
         self,
         mesh: Optional[DeviceMesh],
         mesh_dim_names: Tuple[str],
         replicate_pg: ManagedProcessGroup,
         replicate_dim: int,
-        parent: Optional["ManagedDeviceMesh"],
+        parent: Optional["_ManagedDeviceMesh"],
     ):
+        if mesh is None and parent is not None:
+            raise ValueError(
+                "_ManagedDeviceMesh doesn't support both mesh and parent are None."
+            )
         self.mesh = mesh
         self.mesh_dim_names = mesh_dim_names
         self.replicate_pg = replicate_pg
         self.replicate_dim = replicate_dim
         self.replicate_dim_name = mesh_dim_names[replicate_dim]
         self.parent = parent
         self.flatten_meshes = {}
+        self.device_type = mesh.device_type if mesh is not None else parent.device_type
+        self._flatten_mesh_list = tuple()
+        self._thread_id = None
 
     def __getitem__(self, mesh_dim_names: Union[str, Tuple[str, ...]]) -> DeviceMesh:
         if isinstance(mesh_dim_names, str):
             if mesh_dim_names == self.replicate_dim_name:
-                return ManagedDeviceMesh(
+                return _ManagedDeviceMesh(
                     mesh=None,
                     mesh_dim_names=(mesh_dim_names,),
                     replicate_pg=self.replicate_pg,
@@ -843,22 +850,25 @@ def __getitem__(self, mesh_dim_names: Union[str, Tuple[str, ...]]) -> DeviceMesh
             if self.replicate_dim_name in mesh_dim_names:
                 return self.mesh[mesh_dim_names]
             else:
-                return ManagedDeviceMesh(
+                return _ManagedDeviceMesh(
                     self.mesh[mesh_dim_names],
                     mesh_dim_names,
                     self.replicate_pg,
                     mesh_dim_name.index(self.replicate_dim_name),
                     parent=self,
                 )
 
+    def _real_mesh_dim(self, mesh_dim: int) -> int:
+        return mesh_dim - 1 if mesh_dim > self.replicate_dim else mesh_dim
+
     def get_group(self, mesh_dim: Optional[str] = None) -> BaseProcessGroup:
         if mesh_dim is None:
             assert self.mesh is None
             return self.replicate_pg
         elif mesh_dim == self.replicate_dim_name:
             return self.replicate_pg
         else:
-            return self.mesh.get_group(mesh_dim)
+            return self.mesh.get_group(self._real_mesh_dim(mesh_dim))
 
     def _flatten(self, mesh_dim_name: str) -> "DeviceMesh":
         flatten_mesh = _FlattenDeviceMesh(self)
@@ -877,7 +887,7 @@ def size(self, mesh_dim: Optional[int] = None) -> int:
         elif mesh_dim == self.replicate_dim:
             return self.replicate_pg.size()
         else:
-            return self.mesh.size(mesh_dim)
+            return self.mesh.size(self._real_mesh_dim(mesh_dim))
 
     @property
     def ndim(self) -> int:
@@ -904,14 +914,21 @@ def get_local_rank(self, mesh_dim: Optional[Union[int, str]] = None) -> int:
         elif mesh_dim in (self.replicate_dim, self.replicate_dim_name):
             return get_rank(self.replicate_pg)
         else:
-            return self.mesh.get_local_rank(mesh_dim)
+            return self.mesh.get_local_rank(self._real_mesh_dim(mesh_dim))
+
+    def get_coordinate(self) -> Optional[List[int]]:
+        """
+        Return the relative indices of this rank relative to all
+        dimensions of the mesh. If this rank is not part of the mesh, return None.
+        """
+        return self.mesh._coordinate_on_dim if self.mesh._coordinate_on_dim else None
 
     def get_all_groups(self) -> List[ProcessGroup]:
         raise NotImplementedError
 
 
 class _FlattenDeviceMesh(DeviceMesh):
-    def __init__(self, managed_mesh: ManagedDeviceMesh):
+    def __init__(self, managed_mesh: _ManagedDeviceMesh):
         self.managed_mesh = managed_mesh
 
     def __getitem__(self, mesh_dim_names: Union[str, Tuple[str, ...]]) -> DeviceMesh:
@@ -954,7 +971,7 @@ def ft_init_device_mesh(
     replicate_dim: int,
     manager: "Manager",
 ):
-    # We have to lie DeviceMesh that the replicate_dim has only
+    # We need to mislead DeviceMesh into thinking that replicate_dim has only
     # 1 rank.
     _mesh_shape = list(mesh_shape)
     _mesh_shape.pop(replicate_dim)
@@ -979,7 +996,7 @@ def ft_init_device_mesh(
     # the same backend has been registered.
     replicate_pg.register(mesh_dim_names[replicate_dim])
 
-    return ManagedDeviceMesh(
+    return _ManagedDeviceMesh(
         mesh=mesh,
         mesh_dim_names=mesh_dim_names,
         replicate_pg=replicate_pg,

torchft/process_group_test.py

@@ -33,7 +33,6 @@
 from torchft.manager import Manager
 from torchft.process_group import (
     ErrorSwallowingProcessGroupWrapper,
-    ManagedDeviceMesh,
     ManagedProcessGroup,
     ProcessGroup,
     ProcessGroupBabyGloo,