[AutoParallel] fix error when using pure sharding (#69214)

AndSonder · web-flow · commit bfdf6bc142d3 · 2024-11-13T16:31:19.000+08:00
* fix puring sharding error

* fix sd

* polish

* add unit test

* fix test

* polish

* Update semi_auto_parallel_sharding_stage_1.py

* Update semi_auto_parallel_sharding_stage_1.py

* Update semi_auto_parallel_sharding_stage_1.py

* fix code style
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -37,6 +37,7 @@
     in_pir_mode,
     use_pir_api,
 )
+from paddle.distributed import fleet
 from paddle.distributed.auto_parallel import Engine, strategy as auto_strategy
 from paddle.distributed.auto_parallel.interface import (
     shard_tensor as shard_tensor_static,
@@ -1044,6 +1045,7 @@ def __init__(self, optimizer, shard_fn=None, gradient_accumulation_steps=1):
                 self._shard_fn._shard_parameter(param)
 
     def _set_and_check_sharding_prop_from_param(self):
+        all_params_replicated_on_each_mesh = True
         if (self._shard_fn._mesh is not None) and (
             len(self._shard_fn._mesh._shape) == 1
         ):
@@ -1067,11 +1069,17 @@ def _set_and_check_sharding_prop_from_param(self):
                         isinstance(placement, dist.Shard)
                         for placement in placements
                     ):
+                        all_params_replicated_on_each_mesh = False
                         for idx, placement in enumerate(placements):
                             if isinstance(placement, dist.Replicate):
                                 self._sharding_degree = mesh.dim_size(idx)
                                 self._sharding_mesh_axis = idx
                                 break
+                    elif any(
+                        isinstance(placement, dist.Partial)
+                        for placement in placements
+                    ):
+                        all_params_replicated_on_each_mesh = False
                 else:
                     # check the placement on sharding axis is Replicate
                     assert isinstance(
@@ -1084,6 +1092,17 @@ def _set_and_check_sharding_prop_from_param(self):
                         == self._sharding_degree
                     ), "The sharding degree of all parameters must be equal currently."
 
+        # Note(luchang): When all parameters are replicated across meshes,
+        # we set sharding degree to dp degree and mesh axis to 0. This is sufficient
+        # because each device already holds a full copy of the parameters,
+        # so no additional sharding configuration is needed.
+        if self._sharding_degree is None and all_params_replicated_on_each_mesh:
+            global_mesh = fleet.auto.get_mesh()
+            self._sharding_degree = global_mesh.get_dim_size(
+                self._shard_fn._shard_dims
+            )
+            self._sharding_mesh_axis = 0
+
         assert (
             self._sharding_degree is not None
         ), "The sharding degree is None in ShardOptimizer"
@@ -1286,9 +1305,10 @@ def __setattr__(self, item, value):
 
 
 class _ShardingStageBase:
-    def __init__(self, mesh):
+    def __init__(self, mesh, shard_dims, shard_axis):
         self._mesh = mesh
-        self._sharding_mesh_axis = None
+        self._sharding_mesh_axis = shard_axis
+        self._shard_dims = shard_dims
 
     def _set_sharding_mesh_axis(self, sharding_mesh_axis):
         self._sharding_mesh_axis = sharding_mesh_axis
@@ -1331,6 +1351,8 @@ class ShardingStage1(_ShardingStageBase):
 
     Args:
         mesh(None|paddle.distributed.ProcessMesh): If mesh is not None, the `ProcessMesh` object describes the Cartesian topology of the used processes for dense type parameters. Note: Currently, only one mesh configuration is supported for all dense parameters. If there is a need for multiple mesh configurations, please configure them yourself in the upper layer networking code.
+        shard_dims(None|int|str): The sharding dimension in the mesh.
+        shard_axis(int): The sharding axis of the weight tensor.
 
     Examples:
         .. code-block:: python
@@ -1363,8 +1385,13 @@ class ShardingStage1(_ShardingStageBase):
             >>> # python -m paddle.distributed.launch --gpus=0,1 {test_case}.py
     """
 
-    def __init__(self, mesh: ProcessMesh | None = None) -> None:
-        super().__init__(mesh)
+    def __init__(
+        self,
+        mesh: ProcessMesh | None = None,
+        shard_dims: int | str | None = None,
+        shard_axis: int = 0,
+    ) -> None:
+        super().__init__(mesh, shard_dims, shard_axis)
 
     def __call__(self, key: str, param: Tensor, accumulator: Tensor) -> Tensor:
         if param.is_dist():
@@ -1418,6 +1445,8 @@ class ShardingStage2(_ShardingStageBase):
 
     Args:
         mesh(None|paddle.distributed.ProcessMesh): If mesh is not None, the `ProcessMesh` object describes the Cartesian topology of the used processes for dense type parameters. Note: Currently, only one mesh configuration is supported for all dense parameters. If there is a need for multiple mesh configurations, please configure them yourself in the upper layer networking code.
+        shard_dims(None|int|str): The sharding dimension name in the mesh.
+        shard_axis(int): The sharding axis of the weight tensor.
 
     Examples:
         .. code-block:: python
@@ -1450,8 +1479,13 @@ class ShardingStage2(_ShardingStageBase):
             >>> # python -m paddle.distributed.launch --gpus=0,1 {test_case}.py
     """
 
-    def __init__(self, mesh: ProcessMesh | None = None) -> None:
-        super().__init__(mesh)
+    def __init__(
+        self,
+        mesh: ProcessMesh | None = None,
+        shard_dims: int | str | None = None,
+        shard_axis: int = 0,
+    ) -> None:
+        super().__init__(mesh, shard_dims, shard_axis)
 
     def __call__(self, key: str, param: Tensor, accumulator: Tensor) -> Tensor:
         if param.is_dist():
@@ -1529,6 +1563,8 @@ class ShardingStage3(_ShardingStageBase):
 
     Args:
         mesh(None|paddle.distributed.ProcessMesh): If mesh is not None, the `ProcessMesh` object describes the Cartesian topology of the used processes for dense type parameters. Note: Currently, only one mesh configuration is supported for all dense parameters. If there is a need for multiple mesh configurations, please configure them yourself in the upper layer networking code.
+        shard_dims(None|int|str): The sharding dimension name in the mesh.
+        shard_axis(int): The sharding axis of the weight tensor.
 
     Examples:
         .. code-block:: python
@@ -1561,8 +1597,13 @@ class ShardingStage3(_ShardingStageBase):
             >>> # python -m paddle.distributed.launch --gpus=0,1 {test_case}.py
     """
 
-    def __init__(self, mesh: ProcessMesh | None = None) -> None:
-        super().__init__(mesh)
+    def __init__(
+        self,
+        mesh: ProcessMesh | None = None,
+        shard_dims: int | str | None = None,
+        shard_axis: int = 0,
+    ) -> None:
+        super().__init__(mesh, shard_dims, shard_axis)
 
     def _shard_parameter(self, param):
         if param.is_dense() and self._mesh is not None:
diff --git a/test/auto_parallel/semi_auto_parallel_sharding_stage_1.py b/test/auto_parallel/semi_auto_parallel_sharding_stage_1.py
@@ -27,6 +27,9 @@ def __init__(self):
         self._backend = os.getenv("backend")
         self._seed = eval(os.getenv("seed"))
         self._mesh = dist.ProcessMesh([0, 1], dim_names=["dp"])
+        self._multi_dim_mesh = dist.ProcessMesh(
+            [[0, 1]], dim_names=["pp", "dp"]
+        )
 
     def check_tensor_eq(self, a, b, rtol=1e-05, atol=0, verbose=True):
         np.testing.assert_allclose(a, b, rtol=rtol, atol=atol, verbose=verbose)
@@ -45,6 +48,7 @@ def get_single_card_rst(self):
         self.bias = linear.bias.numpy()
 
     def test_pure_sharding_stage_1(self):
+        paddle.distributed.auto_parallel.set_mesh(self._mesh)
         paddle.seed(self._seed)
         linear = paddle.nn.Linear(10, 10)
         batch = paddle.rand(shape=[10, 10])
@@ -61,7 +65,26 @@ def test_pure_sharding_stage_1(self):
         self.check_tensor_eq(self.weight, linear.weight.numpy())
         self.check_tensor_eq(self.bias, linear.bias.numpy())
 
+    def test_pure_sharding_multi_mesh_stage_1(self):
+        paddle.distributed.auto_parallel.set_mesh(self._multi_dim_mesh)
+        paddle.seed(self._seed)
+        linear = paddle.nn.Linear(10, 10)
+        batch = paddle.rand(shape=[10, 10])
+        # shard the input by sharding degree
+        batch = dist.shard_tensor(batch, self._mesh, [dist.Shard(0)])
+        # shard optimizer with stage 1 fn
+        opt = paddle.optimizer.AdamW(parameters=linear.parameters())
+        opt = dist.shard_optimizer(opt, dist.ShardingStage1(shard_dims="dp"))
+        for _ in range(5):
+            loss = linear(batch)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        self.check_tensor_eq(self.weight, linear.weight.numpy())
+        self.check_tensor_eq(self.bias, linear.bias.numpy())
+
     def test_sharding_stage_1_to_static(self):
+        paddle.distributed.auto_parallel.set_mesh(self._mesh)
         data_loader = create_data_loader()
         layer = DemoNet(self._mesh, "sharding_demonet")
         opt = paddle.optimizer.SGD(
@@ -94,6 +117,7 @@ def run_test_case(self):
         self.get_single_card_rst()
         self.test_pure_sharding_stage_1()
         self.test_sharding_stage_1_to_static()
+        self.test_pure_sharding_multi_mesh_stage_1()
 
 
 if __name__ == '__main__':
