[AutoParallel] Reconstruct sharding mesh dimension inference logic - Part1 rename func parameter (#69393)

* rename

* update

* update test

* rename

* Update semi_auto_parallel_sharding_stage_1.py

* fix

Author: AndSonder

python/paddle/distributed/auto_parallel/api.py

@@ -992,7 +992,7 @@ def replicate_layer_params_and_buffers(
         )
 
 
-def get_placement_with_sharding(param, sharding_mesh_axis):
+def get_placement_with_sharding(param, sharding_axis):
     shard_axis = -1
     for placement in param.placements:
         if isinstance(placement, dist.Shard):
@@ -1011,7 +1011,7 @@ def get_placement_with_sharding(param, sharding_mesh_axis):
 
     new_placements = param.placements
     if placement_with_sharding is not None:
-        new_placements[sharding_mesh_axis] = placement_with_sharding
+        new_placements[sharding_axis] = placement_with_sharding
 
     return new_placements
 
@@ -1040,15 +1040,15 @@ def __init__(self, optimizer, shard_fn=None, gradient_accumulation_steps=1):
         ):
             self._shard_clip = True
         self._shard_fn = shard_fn
-        self._sharding_mesh_axis = None
+        self._sharding_axis = None
         self._sharding_degree = None
         self.gradient_accumulation_steps = gradient_accumulation_steps
 
         if isinstance(
             self._shard_fn, (ShardingStage1, ShardingStage2, ShardingStage3)
         ):
             self._set_and_check_sharding_prop_from_param()
-            self._shard_fn._set_sharding_mesh_axis(self._sharding_mesh_axis)
+            self._shard_fn._set_sharding_axis(self._sharding_axis)
 
         # Invoke register hook for sharding stage 2 strategy
         if isinstance(self._shard_fn, ShardingStage2):
@@ -1066,12 +1066,12 @@ def _set_and_check_sharding_prop_from_param(self):
             len(self._shard_fn._mesh._shape) == 1
         ):
             self._sharding_degree = self._shard_fn._mesh.get_dim_size(0)
-            self._sharding_mesh_axis = 0
+            self._sharding_axis = 0
         elif (self._shard_fn._mesh is not None) and (
             'dp' in self._shard_fn._mesh.dim_names
         ):
             self._sharding_degree = self._shard_fn._mesh.get_dim_size('dp')
-            self._sharding_mesh_axis = 0
+            self._sharding_axis = 0
         else:
             param_list = self._inner_opt._parameter_list
             for param in param_list:
@@ -1090,7 +1090,7 @@ def _set_and_check_sharding_prop_from_param(self):
                         for idx, placement in enumerate(placements):
                             if isinstance(placement, dist.Replicate):
                                 self._sharding_degree = mesh.dim_size(idx)
-                                self._sharding_mesh_axis = idx
+                                self._sharding_axis = idx
                                 break
                     elif any(
                         isinstance(placement, dist.Partial)
@@ -1100,12 +1100,12 @@ def _set_and_check_sharding_prop_from_param(self):
                 else:
                     # check the placement on sharding axis is Replicate
                     assert isinstance(
-                        placements[self._sharding_mesh_axis], dist.Replicate
-                    ), "The placement on sharding_mesh_axis should be Replicate"
+                        placements[self._sharding_axis], dist.Replicate
+                    ), "The placement on sharding_axis should be Replicate"
 
                     # check the sharding degree since it has already been set
                     assert (
-                        mesh.dim_size(self._sharding_mesh_axis)
+                        mesh.dim_size(self._sharding_axis)
                         == self._sharding_degree
                     ), "The sharding degree of all parameters must be equal currently."
 
@@ -1116,9 +1116,9 @@ def _set_and_check_sharding_prop_from_param(self):
         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._shard_fn._sharding_mesh_dim
             )
-            self._sharding_mesh_axis = 0
+            self._sharding_axis = 0
 
         assert (
             self._sharding_degree is not None
@@ -1176,7 +1176,7 @@ def _reset_placements(self, param):
             # in pir mode, reshard pass will automatically handle inplace case, so no extra work is required here.
             if not isinstance(param, pir.Value):
                 new_placement = param.placements
-                new_placement[self._sharding_mesh_axis] = dist.Replicate()
+                new_placement[self._sharding_axis] = dist.Replicate()
                 out_param = dist.reshard(
                     param, param.process_mesh, new_placement
                 )
@@ -1322,13 +1322,13 @@ def __setattr__(self, item, value):
 
 
 class _ShardingStageBase:
-    def __init__(self, mesh, shard_dims, shard_axis):
+    def __init__(self, mesh, sharding_mesh_dim):
         self._mesh = mesh
-        self._sharding_mesh_axis = shard_axis
-        self._shard_dims = shard_dims
+        self._sharding_axis = 0
+        self._sharding_mesh_dim = sharding_mesh_dim
 
-    def _set_sharding_mesh_axis(self, sharding_mesh_axis):
-        self._sharding_mesh_axis = sharding_mesh_axis
+    def _set_sharding_axis(self, sharding_axis):
+        self._sharding_axis = sharding_axis
 
     def shard_master_weight(
         self, param: Tensor, master_weight: Tensor
@@ -1340,7 +1340,7 @@ def shard_master_weight(
                     data_op.name() == "pd_op.data"
                 ), "The master weight must be a result of data op."
                 placements = get_placement_with_sharding(
-                    param, self._sharding_mesh_axis
+                    param, self._sharding_axis
                 )
                 dim_map, partial_status = to_dim_map(
                     placements, len(master_weight.shape)
@@ -1368,8 +1368,7 @@ 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.
+        sharding_mesh_dim(None|int|str): The sharding dimension in the mesh.
 
     Examples:
         .. code-block:: python
@@ -1405,17 +1404,16 @@ class ShardingStage1(_ShardingStageBase):
     def __init__(
         self,
         mesh: ProcessMesh | None = None,
-        shard_dims: int | str | None = None,
-        shard_axis: int = 0,
+        sharding_mesh_dim: int | str | None = None,
     ) -> None:
-        super().__init__(mesh, shard_dims, shard_axis)
+        super().__init__(mesh, sharding_mesh_dim)
 
     def __call__(self, key: str, param: Tensor, accumulator: Tensor) -> Tensor:
         if param.is_dist():
             # Only deal with momentum in optimizer, beta should be replicated cross param's mesh
             if 'beta' not in key:
                 placements = get_placement_with_sharding(
-                    param, self._sharding_mesh_axis
+                    param, self._sharding_axis
                 )
             else:
                 placements = [
@@ -1462,8 +1460,7 @@ 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.
+        sharding_mesh_dim(None|int|str): The sharding dimension name in the mesh.
 
     Examples:
         .. code-block:: python
@@ -1499,17 +1496,16 @@ class ShardingStage2(_ShardingStageBase):
     def __init__(
         self,
         mesh: ProcessMesh | None = None,
-        shard_dims: int | str | None = None,
-        shard_axis: int = 0,
+        sharding_mesh_dim: int | str | None = None,
     ) -> None:
-        super().__init__(mesh, shard_dims, shard_axis)
+        super().__init__(mesh, sharding_mesh_dim)
 
     def __call__(self, key: str, param: Tensor, accumulator: Tensor) -> Tensor:
         if param.is_dist():
             # Only deal with momentum in optimizer, beta should be replicated cross param's mesh
             if 'beta' not in key:
                 placements = get_placement_with_sharding(
-                    param, self._sharding_mesh_axis
+                    param, self._sharding_axis
                 )
             else:
                 placements = [
@@ -1580,8 +1576,7 @@ 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.
+        sharding_mesh_dim(None|int|str): The sharding dimension name in the mesh.
 
     Examples:
         .. code-block:: python
@@ -1617,10 +1612,9 @@ class ShardingStage3(_ShardingStageBase):
     def __init__(
         self,
         mesh: ProcessMesh | None = None,
-        shard_dims: int | str | None = None,
-        shard_axis: int = 0,
+        sharding_mesh_dim: int | str | None = None,
     ) -> None:
-        super().__init__(mesh, shard_dims, shard_axis)
+        super().__init__(mesh, sharding_mesh_dim)
 
     def _shard_parameter(self, param):
         if param.is_dense() and self._mesh is not None:
@@ -1630,7 +1624,7 @@ def _shard_parameter(self, param):
             param._to_dist_(placements, self._mesh)
         if param.is_dist():
             new_placements = get_placement_with_sharding(
-                param, self._sharding_mesh_axis
+                param, self._sharding_axis
             )
             shard_param = dist.reshard(
                 param, param.process_mesh, new_placements
@@ -1641,8 +1635,8 @@ def _shard_parameter(self, param):
     def _unshard_parameter(self, param):
         if param.is_dist():
             new_placements = param.placements
-            if isinstance(new_placements[self._sharding_mesh_axis], dist.Shard):
-                new_placements[self._sharding_mesh_axis] = dist.Replicate()
+            if isinstance(new_placements[self._sharding_axis], dist.Shard):
+                new_placements[self._sharding_axis] = dist.Replicate()
 
             new_param = dist.reshard(param, param.process_mesh, new_placements)
             param.get_tensor()._share_data_with(new_param.get_tensor())
@@ -1657,7 +1651,7 @@ def __call__(self, key: str, param: Tensor, accumulator: Tensor) -> Tensor:
                     for placement in placements
                 ):
                     placements = get_placement_with_sharding(
-                        param, self._sharding_mesh_axis
+                        param, self._sharding_axis
                     )
 
             else:

python/paddle/distributed/auto_parallel/sharding.py

@@ -106,8 +106,8 @@ def __init__(self, optimizer, shard_fn=None, strategy=None):
         else:
             self.pp_meshes.add(mesh)
 
-        self._sharding_mesh_axis = mesh._dim_names.index("dp")
-        self._sharding_degree = mesh._shape[self._sharding_mesh_axis]
+        self._sharding_axis = mesh._dim_names.index("dp")
+        self._sharding_degree = mesh._shape[self._sharding_axis]
         self._mp_mesh_axis = -1
         self._mp_degree = 1
         if "mp" in mesh._dim_names:
@@ -151,7 +151,7 @@ def apply_gradients(self, params_grads):
 
             if dist.get_rank() in param_dist_attr.process_mesh.process_ids:
                 sub_mesh = get_1D_sub_process_mesh(
-                    param_dist_attr.process_mesh, self._sharding_mesh_axis
+                    param_dist_attr.process_mesh, self._sharding_axis
                 )
                 assert (
                     sorted(sub_mesh.process_ids) == self._sharding_group.ranks
@@ -178,7 +178,7 @@ def apply_gradients(self, params_grads):
                 param._local_shape == grad._local_shape
             ), f"Parameter and grad should have same local shape. but received name:{param.name}, parameter:{param}, grad: {grad}."
 
-            if self._sharding_mesh_axis not in grad_dist_attr.partial_dims:
+            if self._sharding_axis not in grad_dist_attr.partial_dims:
                 new_params_grads.append((param, grad))
                 if param.optimize_attr is None:
                     param.optimize_attr = {'no_fusion': True}
@@ -361,7 +361,7 @@ def apply_gradients(self, params_grads):
                     partail_status = (
                         group_grad_list[index].dist_attr().partial_status
                     )
-                    partail_status.pop(self._sharding_mesh_axis)
+                    partail_status.pop(self._sharding_axis)
                     slice_grad_dist_attr = pir.create_tensor_dist_attribute(
                         slice_grad.process_mesh, [-1], partail_status
                     )

test/auto_parallel/semi_auto_parallel_sharding_stage_1.py

@@ -74,7 +74,9 @@ def test_pure_sharding_multi_mesh_stage_1(self):
         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"))
+        opt = dist.shard_optimizer(
+            opt, dist.ShardingStage1(sharding_mesh_dim="dp")
+        )
         for _ in range(5):
             loss = linear(batch)
             loss.backward()