[PT-D][TP] Fix TP API for FQN path based parallelization

fduwjj · fduwjj · commit ff492562a9eb · 2023-01-26T05:19:31.000Z
ghstack-source-id: 992954f Pull Request resolved: #93029
diff --git a/test/distributed/tensor/parallel/test_parallelize_api.py b/test/distributed/tensor/parallel/test_parallelize_api.py
@@ -3,7 +3,7 @@
 import torch
 from torch.distributed._tensor import DeviceMesh, DTensor, Replicate
 from torch.distributed.tensor.parallel._utils import _create_1d_device_mesh
-from torch.distributed.tensor.parallel.api import _parallelize_linear, _parallelize_mlp
+from torch.distributed.tensor.parallel.api import parallelize_module, _parallelize_linear, _parallelize_mlp
 from torch.distributed.tensor.parallel.style import (
     ColwiseParallel,
     make_input_replicate_1d,
@@ -77,6 +77,7 @@ def _compare_params(
         self,
         local_module,
         dist_module,
+        rank0_only,
         skip_rowwise_bias=False,
         compare_grad=False,
     ):
@@ -85,7 +86,7 @@ def _compare_params(
             dist_param = dist_module.get_parameter(name)
             param = param.grad if compare_grad else param
             dist_param = dist_param.grad if compare_grad else dist_param
-            if self.rank == 0 or (
+            if (not rank0_only) or (self.rank == 0) or (
                 name not in ["net2.bias"]
                 and not skip_rowwise_bias
                 or name not in ["bias", "net2.bias"]
@@ -95,15 +96,16 @@ def _compare_params(
                     dist_param.redistribute(
                         device_mesh=dist_param.device_mesh, placements=replicate
                     ).to_local(),
+                    f"{name} not equal between dist and non-dist"
                 )
 
-    def _compare_module(self, local_module, dist_module, inp_size, rowwise=False):
+    def _compare_module(self, local_module, dist_module, inp_size, rank0_only=True, rowwise=False):
         LR = 0.25  # the learning rate we use for testing
         local_optim = torch.optim.SGD(local_module.parameters(), lr=LR)
         dist_optim = torch.optim.SGD(dist_module.parameters(), lr=LR)
         torch.manual_seed(0)
         inp = torch.rand(*inp_size, device=self.device_type)
-        self._compare_params(local_module, dist_module)
+        self._compare_params(local_module, dist_module, rank0_only)
 
         # check forward correctness
         local_output = local_module(inp)
@@ -118,11 +120,11 @@ def _compare_module(self, local_module, dist_module, inp_size, rowwise=False):
         dist_output.sum().backward()
 
         # check backward and ensure gradients are same
-        self._compare_params(local_module, dist_module, rowwise, True)
+        self._compare_params(local_module, dist_module, rank0_only, rowwise, True)
 
         local_optim.step()
         dist_optim.step()
-        self._compare_params(local_module, dist_module, rowwise)
+        self._compare_params(local_module, dist_module, rank0_only, rowwise)
 
     @with_comms
     def test_parallelize_mlp(self):
@@ -141,6 +143,23 @@ def test_parallelize_mlp(self):
         model_tp = _parallelize_mlp(model_tp, device_mesh, PairwiseParallel())
         self._compare_module(model, model_tp, inp_size)
 
+    @with_comms
+    def test_parallelize_mlp_with_module_api(self):
+        inp_size = [12, 10]
+        model = MLPModule(self.device_type)
+        model_tp = MLPModule(self.device_type)
+
+        # Ensure model are initialized the same way.
+        self.assertEqual(model.net1.weight, model_tp.net1.weight)
+        self.assertEqual(model.net1.bias, model_tp.net1.bias)
+        self.assertEqual(model.net2.weight, model_tp.net2.weight)
+        self.assertEqual(model.net2.bias, model_tp.net2.bias)
+
+        # Parallelize module.
+        device_mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
+        model_tp = parallelize_module(model_tp, device_mesh, {"net1": ColwiseParallel(), "net2": ColwiseParallel()})
+        self._compare_module(model, model_tp, inp_size, rank0_only=False)
+
     @with_comms
     def test_parallelize_mlp_error(self):
         class DummyParallel(ParallelStyle):
@@ -177,7 +196,7 @@ def test_linear_row_wise_parallel(self):
 
         # let each rank generate unique local input
         torch.manual_seed(self.rank)
-        self._compare_module(model, model_tp, inp_size, True)
+        self._compare_module(model, model_tp, inp_size, rowwise=True)
 
     @with_comms
     def test_linear_col_wise_parallel(self):
diff --git a/torch/distributed/tensor/parallel/api.py b/torch/distributed/tensor/parallel/api.py
@@ -97,11 +97,12 @@ def parallelize_module(  # type: ignore[return]
         for module_path, parallelize_style in parallelize_plan.items():
             sub_module = module.get_submodule(module_path)
             module.register_module(  # type: ignore[call-arg] # pyre-ignore[20]
+                module_path,
                 parallelize_module(  # type: ignore[arg-type]
-                    module_path, sub_module, device_mesh, parallelize_style  # type: ignore[arg-type] # pyre-ignore[6]
-                )
+                    sub_module, device_mesh, parallelize_style  # type: ignore[arg-type] # pyre-ignore[6]
+                ),
             )
-            return module
+        return module
     else:
         raise RuntimeError(  # pyre-ignore[7]
             "Expect Union[ParallelStyle, Dict[str, ParallelStyle]] for"
