[3/N][2D] Enable training with new 2D flow (#109553)

wz337 · pytorchmergebot · commit 217b37c023d5 · 2023-09-25T05:32:07.000Z
This PR enables training with new 2D flow and adds associated test. state_dict related changes would be in later PRs. cc. @fegin, @fduwjj, @wanchaol, @awgu Pull Request resolved: #109553 Approved by: https://github.com/fegin, https://github.com/awgu
diff --git a/test/distributed/tensor/parallel/test_fsdp_2d_parallel.py b/test/distributed/tensor/parallel/test_fsdp_2d_parallel.py
@@ -135,6 +135,15 @@ def _apply_optim_in_backward(param_group):
     )
 
 
+def _clean_up_fsdp_param_name(name):
+    return ".".join(
+        filter(
+            lambda name: name not in [FSDP_WRAPPED_MODULE, _CHECKPOINT_WRAPPED_MODULE],
+            name.split("."),
+        )
+    )
+
+
 class Test2dParallelIntegration(DTensorTestBase):
     @with_comms
     @skip_if_lt_x_gpu(4)
@@ -182,15 +191,6 @@ def _compare_params(self, m1, m2):
                         p2 = p2.redistribute(p2.device_mesh, [Replicate()]).to_local()
                     self.assertTrue(torch.allclose(p1, p2), f"{p1} vs {p2}")
 
-    def _clean_up_fsdp_param_name(self, name):
-        return ".".join(
-            filter(
-                lambda name: name
-                not in [FSDP_WRAPPED_MODULE, _CHECKPOINT_WRAPPED_MODULE],
-                name.split("."),
-            )
-        )
-
     def _test_2d_e2e_flow(
         self,
         use_orig_params=False,
@@ -215,8 +215,7 @@ def _test_2d_e2e_flow(
             model_2d = input_reshard(model_2d, mesh_2d["TP"], 0)
         # Check named parameters are returning the same name at least.
         param_names_2d = [
-            self._clean_up_fsdp_param_name(name)
-            for name, _ in model_2d.named_parameters()
+            _clean_up_fsdp_param_name(name) for name, _ in model_2d.named_parameters()
         ]
         for name, _ in model.named_parameters():
             name = self._clean_up_fsdp_param_name(name)
@@ -314,8 +313,25 @@ def test_2d_fsdp_integration_fsdp_nested_param_groups(self) -> None:
 
 
 class TestNew2dParallelIntegration(DTensorTestBase):
+    # TODO: this is duplicate code from above, but once we remove the enable_2d_with_fsdp(),
+    # we will remove the above test class Test2dParallelIntegration.
+    def _compare_params(self, m1, m2):
+        with FSDP.summon_full_params(m1):
+            with FSDP.summon_full_params(m2):
+                for n_p1, n_p2 in zip(m1.named_parameters(), m2.named_parameters()):
+                    p1 = n_p1[1]
+                    p2 = n_p2[1]
+                    if n_p1[0] != n_p2[0]:
+                        self.assertTrue(n_p1[0] in n_p2[0])
+                    name = n_p1[0]
+                    if name == "net2.bias" and self.rank != 0:
+                        continue
+                    if type(p2) is DT:
+                        p2 = p2.redistribute(p2.device_mesh, [Replicate()]).to_local()
+                    self.assertTrue(torch.allclose(p1, p2), f"{p1} vs {p2}")
+
     @with_comms
-    @skip_if_lt_x_gpu(2)
+    @skip_if_lt_x_gpu(4)
     def test_2d_fsdp_state_enable_extension(self):
         mesh_2d = init_device_mesh(
             self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
@@ -327,6 +343,78 @@ def test_2d_fsdp_state_enable_extension(self):
         fsdp_state = _get_module_fsdp_state(model)
         self.assertEqual(fsdp_state._enable_extension, True)
 
+    def _test_2d_e2e_training(
+        self,
+        use_orig_params=False,
+        recompute_activation=False,
+    ) -> None:
+        torch.manual_seed(0)
+        model = SimpleModel().cuda(self.rank)
+        model = FSDP(model, use_orig_params=use_orig_params)
+        optim = torch.optim.Adam(model.parameters(), lr=0.01)
+
+        torch.manual_seed(0)
+        mesh_2d = init_device_mesh(
+            self.device_type, (2, self.world_size // 2), mesh_dim_names=("dp", "tp")
+        )
+        tp_mesh = mesh_2d["tp"]
+        dp_mesh = mesh_2d["dp"]
+        model_2d = parallelize_module(SimpleModel().cuda(), tp_mesh, PairwiseParallel())
+        model_2d = FSDP(
+            model_2d,
+            device_mesh=dp_mesh,
+            use_orig_params=use_orig_params,
+        )
+        optim_2d = torch.optim.Adam(model_2d.parameters(), lr=0.01)
+
+        if recompute_activation:
+            model_2d = input_reshard(model_2d, mesh_2d["tp"], 0)
+
+        # Check named parameters are returning the same name at least.
+        param_names_2d = [
+            _clean_up_fsdp_param_name(name) for name, _ in model_2d.named_parameters()
+        ]
+        for name, _ in model.named_parameters():
+            name = _clean_up_fsdp_param_name(name)
+            if name not in param_names_2d:
+                print(name, param_names_2d)
+            self.assertTrue(name in param_names_2d)
+        self._compare_params(model, model_2d)
+
+        # TODO: add additional tests for multi_param_group and optim_in_backward.
+
+        for i in range(5):
+            # Ensure all input across TP ranks are same.
+            # TODO: add a get_group_rank() to DeviceMesh.
+            torch.manual_seed(i + dist.get_rank(dp_mesh.get_dim_groups()[0]))
+            input = torch.rand(4, 5).cuda(self.rank)
+            output = model(input)
+            output_2d = model_2d(input)
+            self.assertEqual(output, output_2d)
+            output.sum().backward()
+            output_2d.sum().backward()
+            optim.step()
+            optim_2d.step()
+            self.assertEqual(model(input), model_2d(input))
+
+        # Ensure all params are still the same after optimizer update.
+        self._compare_params(model, model_2d)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_2d_e2e_training_default(self):
+        self._test_2d_e2e_training()
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_2d_e2e_training_use_orig_params(self):
+        self._test_2d_e2e_training(use_orig_params=True)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_2d_e2e_training_not_use_orig_params(self):
+        self._test_2d_e2e_training(recompute_activation=True)
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/torch/distributed/fsdp/_init_utils.py b/torch/distributed/fsdp/_init_utils.py
@@ -36,6 +36,7 @@
     clean_tensor_name,
     TrainingState,
 )
+from torch.distributed.fsdp._fsdp_extensions import _set_fsdp_extensions
 from torch.distributed.fsdp._limiter_utils import _FreeEventQueue
 from torch.distributed.fsdp.api import (
     BackwardPrefetch,
@@ -54,6 +55,7 @@
     HandleShardingStrategy,
 )
 from torch.distributed.fsdp.wrap import _Policy
+from torch.distributed.tensor.parallel.fsdp import DTensorExtensions
 from torch.distributed.utils import _sync_params_and_buffers
 from torch.utils.hooks import RemovableHandle
 
@@ -506,6 +508,27 @@ def _init_prefetching_state(
     return state
 
 
+@no_type_check
+def _init_extension(state: _FSDPState, device_mesh: DeviceMesh = None) -> _FSDPState:
+    # TODO: we need to add additional check once we support FSDP + PiPPy.
+    # This check is currently sufficient, since we only support FSDP + TP.
+    if device_mesh:
+        state._enable_extension = (
+            mesh_resources.get_parent_mesh(state._device_mesh) is not None
+        )
+
+        if state._enable_extension:
+            try:
+                _set_fsdp_extensions(DTensorExtensions())
+            except BaseException as e:
+                warnings.warn(
+                    "PyTorch doesn't have TensorFlattener extension point available"
+                    "2D parallelism won't work with FSDP"
+                    f"exception: {e}"
+                )
+    return state
+
+
 @no_type_check
 def _init_state_dict_state(state: _FSDPState) -> _FSDPState:
     state._state_dict_type = StateDictType.FULL_STATE_DICT
@@ -515,12 +538,6 @@ def _init_state_dict_state(state: _FSDPState) -> _FSDPState:
     unshard_params_ctx: Dict[nn.Module, Generator] = {}
     state._unshard_params_ctx = unshard_params_ctx
 
-    # TODO: we need to add additional check once we support FSDP + PiPPy.
-    # This check is currently sufficient, since we only support FSDP + TP.
-    if hasattr(state, "_device_mesh"):
-        state._enable_extension = (
-            mesh_resources.get_parent_mesh(state._device_mesh) is not None
-        )
     return state
 
 
diff --git a/torch/distributed/fsdp/fully_sharded_data_parallel.py b/torch/distributed/fsdp/fully_sharded_data_parallel.py
@@ -42,6 +42,7 @@
     _init_buffer_state,
     _init_core_state,
     _init_device_handle,
+    _init_extension,
     _init_ignored_module_states,
     _init_param_handle_from_module,
     _init_prefetching_state,
@@ -491,6 +492,8 @@ def __init__(
         _init_runtime_state(self)
         _init_prefetching_state(self, backward_prefetch, forward_prefetch)
         _init_buffer_state(self, module)
+        # extension needs to be set before `_init_param_handle_from_module()`
+        _init_extension(self, device_mesh)
         _init_param_handle_from_module(
             self,
             module,
diff --git a/torch/distributed/tensor/parallel/fsdp.py b/torch/distributed/tensor/parallel/fsdp.py
@@ -3,19 +3,55 @@
 
 import torch
 import torch.distributed as dist
-from torch.distributed._shard.sharded_tensor import (
-    Shard,
-)
+from torch.distributed._shard.sharded_tensor import Shard
+from torch.distributed.fsdp._fsdp_extensions import _set_fsdp_extensions, FSDPExtensions
 from torch.distributed.tensor.parallel._data_parallel_utils import (
     _chunk_tensor,
     _flatten_tensor,
     _pre_load_state_dict,
     _unflatten_tensor,
 )
 
-__all__ = ["enable_2d_with_fsdp"]
+__all__ = ["enable_2d_with_fsdp", "DTensorExtensions"]
+
+
+class DTensorExtensions(FSDPExtensions):
+    """
+    DTensorExtension is the TensorFlattener extension needed for 2D FSDP + TP.
+    This is the implementation for FSDPExtensions defined in
+    https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fsdp_extensions.py
+    """
+
+    def pre_flatten_transform(
+        self,
+        tensor: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Optional[Any]]:
+        return _flatten_tensor(tensor)
+
+    def post_unflatten_transform(
+        self, tensor: torch.Tensor, param_extension: Any
+    ) -> torch.Tensor:
+        return _unflatten_tensor(tensor, param_extension)
+
+    def chunk_tensor(
+        self,
+        tensor: torch.Tensor,
+        rank: int,
+        world_size: int,
+        num_devices_per_node: int,
+        pg: dist.ProcessGroup,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        return _chunk_tensor(tensor, rank, world_size, num_devices_per_node, pg)
 
+    def pre_load_state_dict_transform(
+        self,
+        tensor: torch.Tensor,
+    ) -> Tuple[torch.Tensor, List[Shard]]:
+        return _pre_load_state_dict(tensor)
 
+
+# TODO: remove enable_2d_with_fsdp() once we roll out the new 2D flow.
 def enable_2d_with_fsdp() -> bool:
     """
     The API registers the extension which is needed for Tensor Parallelism (TP)
@@ -27,43 +63,11 @@ def enable_2d_with_fsdp() -> bool:
         A `bool` indicated whether extension registration succeeds or not.
     """
 
-    torch._C._log_api_usage_once("torch.distributed.tensor.parallel.enable_2d_with_fsdp")
+    torch._C._log_api_usage_once(
+        "torch.distributed.tensor.parallel.enable_2d_with_fsdp"
+    )
 
     try:
-        from torch.distributed.fsdp._fsdp_extensions import (
-            _set_fsdp_extensions,
-            FSDPExtensions,
-        )
-
-        class DTensorExtensions(FSDPExtensions):
-            def pre_flatten_transform(
-                self,
-                tensor: torch.Tensor,
-            ) -> Tuple[torch.Tensor, Optional[Any]]:
-                return _flatten_tensor(tensor)
-
-            def post_unflatten_transform(
-                self, tensor: torch.Tensor, param_extension: Any
-            ) -> torch.Tensor:
-                return _unflatten_tensor(tensor, param_extension)
-
-            def chunk_tensor(
-                self,
-                tensor: torch.Tensor,
-                rank: int,
-                world_size: int,
-                num_devices_per_node: int,
-                pg: dist.ProcessGroup,
-                device: Optional[torch.device] = None,
-            ) -> torch.Tensor:
-                return _chunk_tensor(tensor, rank, world_size, num_devices_per_node, pg)
-
-            def pre_load_state_dict_transform(
-                self,
-                tensor: torch.Tensor,
-            ) -> Tuple[torch.Tensor, List[Shard]]:
-                return _pre_load_state_dict(tensor)
-
         _set_fsdp_extensions(DTensorExtensions())
         return True
 
