[AutoParallel] Handle dynamic shape in InferGlobalShape and InferLocalShape (#71320)

pkuzyc · zhangbo9674 · web-flow · commit 1fc0ea7eebe5 · 2025-03-05T15:04:21.000+08:00
* fix

* fix

* handle dynamic shape in PIR infer_local_shape and infer_global_shape and add PIR nd_mesh_alltoall reshard function

* fix bugs in static auto parallel

---------

Co-authored-by: zhangbo9674 &lt;zhangbo54@baidu.com&gt;
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_type.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_type.cc
@@ -64,13 +64,15 @@ common::DDim InferLocalDDim(const common::DDim& global_ddim,
         int tensor_dim =
             dynamic_cast<const phi::distributed::Shard&>(*placements[i])
                 .get_dim();
+        if (local_ddim[tensor_dim] == -1) continue;
         auto dim_size = mesh_dim.at(i);
         local_ddim[tensor_dim] =
             (local_ddim[tensor_dim] + dim_size - 1) / dim_size;
       }
     }
   } else {
     for (size_t i = 0; i < dim_mapping.size(); ++i) {
+      if (local_ddim[i] == -1) continue;
       if (dim_mapping[i] != -1) {
         auto dim_size = mesh_dim.at(dim_mapping[i]);
         local_ddim[i] = (global_ddim[i] + dim_size - 1) / dim_size;
@@ -96,11 +98,29 @@ common::DDim InferGlobalDDim(const common::DDim& local_ddim,
                         local_ddim.size(),
                         dim_mapping.size()));
   common::DDim global_ddim(local_ddim);
-  for (size_t i = 0; i < dim_mapping.size(); ++i) {
-    if (dim_mapping[i] != -1) {
-      global_ddim[i] = local_ddim[i] * mesh_dim.at(dim_mapping[i]);
+  if (dist_attr.placements_attr().has_value()) {
+    PlacementsAttribute placements_attr = dist_attr.placements_attr().value();
+    const phi::distributed::Placements& placements =
+        placements_attr.placements();
+    for (size_t i = 0; i < placements.size(); i++) {
+      if (placements[i]->is_shard()) {
+        int tensor_dim =
+            dynamic_cast<const phi::distributed::Shard&>(*placements[i])
+                .get_dim();
+        if (global_ddim[tensor_dim] == -1) continue;
+        auto dim_size = mesh_dim.at(i);
+        global_ddim[tensor_dim] = local_ddim[tensor_dim] * dim_size;
+      }
+    }
+  } else {
+    for (size_t i = 0; i < dim_mapping.size(); ++i) {
+      if (global_ddim[i] == -1) continue;
+      if (dim_mapping[i] != -1) {
+        global_ddim[i] = local_ddim[i] * mesh_dim.at(dim_mapping[i]);
+      }
     }
   }
+
   return global_ddim;
 }
 
diff --git a/paddle/fluid/pybind/auto_parallel_py.cc b/paddle/fluid/pybind/auto_parallel_py.cc
@@ -832,6 +832,8 @@ void BindAutoParallel(py::module *m) {
     DistributedMapper mapper;
     return mapper.to_string();
   });
+
+  m->def("sub_mesh_dim", &phi::distributed::SubMeshDim);
 }
 
 static void parse_tensors(PyObject *obj,
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -2888,16 +2888,23 @@ CvtPlacements(Placements placements, int ndim) {
     if (placement->is_shard()) {
       auto shard_dim =
           dynamic_cast<const phi::distributed::Shard&>(*placement).get_dim();
-      PADDLE_ENFORCE_EQ(
-          dim_map[shard_dim],
-          -1,
-          common::errors::InvalidArgument(
-              "Tensor dim %lld is already sharded on mesh dim %lld,"
-              " DistTensor operator implementation does not support things "
-              "like hybrid"
-              " sharding strategies yet (i.e. [Shard(0), Shard(0)])",
-              shard_dim,
-              dim_map[shard_dim]));
+      if (dim_map[shard_dim] != -1) {
+        LOG(WARNING) << "WARNING: Tensor dim " << shard_dim
+                     << " is already sharded on "
+                     << "mesh dim" << dim_map[shard_dim]
+                     << ". Sharding a tensor dim with "
+                     << "multiple mesh dim is not supported yet.";
+      }
+      // PADDLE_ENFORCE_EQ(
+      //     dim_map[shard_dim],
+      //     -1,
+      //     common::errors::InvalidArgument(
+      //         "Tensor dim %lld is already sharded on mesh dim %lld,"
+      //         " DistTensor operator implementation does not support things "
+      //         "like hybrid"
+      //         " sharding strategies yet (i.e. [Shard(0), Shard(0)])",
+      //         shard_dim,
+      //         dim_map[shard_dim]));
       dim_map[shard_dim] = i;
     }
   }
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -487,6 +487,8 @@ def _cal_global_shape(local_shape, mesh, placements):
     for idx, placement in enumerate(placements):
         if placement.is_shard():
             shard_dim = placement.get_dim()
+            if global_shape[shard_dim] == -1:
+                continue
             local_dim_size = global_shape[shard_dim]
             global_shape[shard_dim] = local_dim_size * mesh.shape[idx]
     return global_shape
diff --git a/python/paddle/distributed/auto_parallel/moe_utils.py b/python/paddle/distributed/auto_parallel/moe_utils.py
@@ -26,7 +26,8 @@
 from paddle import Tensor
 from paddle.autograd import PyLayer
 
-from .placement_type import check_placements_equal
+from .placement_type import check_placements_equal, to_dim_map
+from .static.reshard_funcs.base_reshard_func import choose_reshard_func
 from .static.reshard_funcs.nd_mesh_reshard_func import get_1D_sub_process_mesh
 
 if TYPE_CHECKING:
@@ -44,8 +45,12 @@ def _specific_alltoall_dim(
         return None
 
     mesh_dim = None
-    src_mesh = dist_tensor.process_mesh
-    src_placements = dist_tensor.placements
+    if paddle.in_dynamic_mode():
+        src_mesh = dist_tensor.process_mesh
+        src_placements = dist_tensor.placements
+    elif paddle.framework.in_pir_mode():
+        src_mesh = dist_tensor.process_mesh
+        src_placements = dist_tensor.dist_attr().placements_attr
 
     if src_mesh != mesh or src_mesh.ndim == 1:
         return None
@@ -127,6 +132,59 @@ def _dtensor_from_local(
         )
 
 
+def _pir_nd_mesh_all2all(src_value, dst_type, mesh, placements, dim):
+    """
+    Use all to all communication in nd_mesh reshard.
+    """
+    # create value on sub 1D mesh
+    sub_value = paddle._C_ops.share_data(src_value)
+    sub_mesh = get_1D_sub_process_mesh(mesh, dim)
+    sub_placements = [src_value.dist_attr().placements_attr[dim]]
+    sub_value_shape = dist.auto_parallel.api._cal_global_shape(
+        src_value._local_shape, sub_mesh, sub_placements
+    )
+    sub_value_type = paddle.pir.create_shaped_type(
+        sub_value.type(), sub_value_shape
+    )
+    sub_dims_mapping, partial_status = to_dim_map(
+        sub_placements, len(sub_value_shape)
+    )
+    sub_value_dist_attr = (
+        paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+            sub_mesh, sub_dims_mapping, partial_status
+        )
+    )
+    sub_value_dist_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+        sub_value_type, sub_value_dist_attr
+    )
+    sub_value.set_type(sub_value_dist_type)
+
+    # 1D mesh reshard
+    dst_placements = [placements[dim]]
+    sub_dst_dims_mapping, partial_status = to_dim_map(
+        dst_placements, len(sub_value_shape)
+    )
+    sub_dst_type = paddle.pir.create_shaped_type(
+        sub_value.type(), sub_value_shape
+    )
+    sub_dst_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+        sub_mesh, sub_dst_dims_mapping, partial_status
+    )
+    sub_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+        sub_dst_type, sub_dst_dist_attr
+    )
+    reshard_func = choose_reshard_func(sub_value_dist_attr, sub_dst_dist_attr)
+    out = reshard_func.reshard(
+        sub_value_dist_attr, sub_dst_dist_attr, sub_value, sub_dst_type
+    )
+
+    # set the type of the output value with global mesh
+    if out is not None:
+        out.set_type(dst_type)
+
+    return out
+
+
 class _NdMeshAlltoAll(PyLayer):
     @staticmethod
     def forward(
@@ -144,7 +202,7 @@ def forward(
         ctx.out_placements = copy.deepcopy(placements)
 
         local_shape = _cal_local_shape(
-            dist_tensor.shape, mesh, dist_tensor.placements
+            dist_tensor.shape, sub_mesh, dist_tensor.placements
         )
         out = _dtensor_from_local(
             dist_tensor._local_value(),
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -329,21 +329,43 @@ def reshard_op_pass(dist_program, global_params_grads=None, block=None):
                     op.erase()
                     continue
 
-                reshard_func = choose_reshard_func(src_dist_attr, dst_dist_attr)
-                assert (
-                    reshard_func is not None
-                ), f'There is no reshard function that matches src_dist_attr: {src_dist_attr} and dst_dist_attr: {dst_dist_attr}, {var.get_defining_op()}'
-
                 paddle.pir.set_insertion_point(op)
                 ref_op_role = op.op_role
 
-                with pir_op_role_guard(ref_op_role):
-                    out_value = reshard_func.reshard(
-                        src_dist_attr,
-                        dst_dist_attr,
-                        op.operand_source(0),
-                        op.result(0).type(),
+                all_to_all_dim = (
+                    dist.auto_parallel.moe_utils._specific_alltoall_dim(
+                        var,
+                        dst_dist_attr.process_mesh,
+                        dst_dist_attr.placements_attr,
                     )
+                )
+
+                if all_to_all_dim is not None:
+                    with pir_op_role_guard(ref_op_role):
+                        out_value = (
+                            dist.auto_parallel.moe_utils._pir_nd_mesh_all2all(
+                                op.operand_source(0),
+                                op.result(0).type(),
+                                dst_dist_attr.process_mesh,
+                                dst_dist_attr.placements_attr,
+                                all_to_all_dim,
+                            )
+                        )
+                else:
+                    reshard_func = choose_reshard_func(
+                        src_dist_attr, dst_dist_attr
+                    )
+                    assert (
+                        reshard_func is not None
+                    ), f'There is no reshard function that matches src_dist_attr: {src_dist_attr} and dst_dist_attr: {dst_dist_attr}, {var.get_defining_op()}'
+
+                    with pir_op_role_guard(ref_op_role):
+                        out_value = reshard_func.reshard(
+                            src_dist_attr,
+                            dst_dist_attr,
+                            op.operand_source(0),
+                            op.result(0).type(),
+                        )
 
                 if out_value is not None:
                     op.result(0).replace_all_uses_with(out_value)
@@ -467,11 +489,13 @@ def prune_op(block):
                     continue
                 elif op.name() == "dist_op.dtensor_from_local":
                     dtensor_to_local_idx = idx
-                    while (
-                        reverse_block_ops[dtensor_to_local_idx].name()
-                        != "dist_op.dtensor_to_local"
-                    ):
-                        dtensor_to_local_idx += 1
+                    for i in range(idx, len(reverse_block_ops)):
+                        if (
+                            reverse_block_ops[i].name()
+                            == "dist_op.dtensor_to_local"
+                        ):
+                            dtensor_to_local_idx = i
+                            break
                     if (
                         op.dist_attr
                         and cur_rank
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/r_to_p_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/r_to_p_reshard_func.py
@@ -60,11 +60,13 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
         else:
             dst_value = paddle.assign(src_value)
 
-        src_chunk_id = src_value.get_defining_op().dist_attr.chunk_id
+        chunk_id = -1
+        if src_value.get_defining_op().dist_attr:
+            chunk_id = src_value.get_defining_op().dist_attr.chunk_id
         dst_value.set_type(dst_type)
         dst_value.get_defining_op().dist_attr = (
             paddle.base.libpaddle.pir.create_op_dist_attribute(
-                dst_mesh, [src_dist_attr], [dst_dist_attr], src_chunk_id
+                dst_mesh, [src_dist_attr], [dst_dist_attr], chunk_id
             )
         )
 
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/reshard_func_register.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/reshard_func_register.py
@@ -34,6 +34,7 @@
     SToRReshardFunction,
     SToRReshardFunctionCrossMesh,
 )
+from .s_to_s_reshard_func import SToSReshardFunction
 from .same_status_reshard_func import SameStatusReshardFunction
 from .sub_to_global_mesh_func import SubToGlobalMeshFunction
 
@@ -52,6 +53,7 @@ def register_reshard_funcs():
     register_reshard_func(NdMeshReshardFunctionCrossMesh())
     register_reshard_func(GlobaleToSubMeshFunction())
     register_reshard_func(SubToGlobalMeshFunction())
+    register_reshard_func(SToSReshardFunction())
 
 
 register_reshard_funcs()
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/s_to_s_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/s_to_s_reshard_func.py
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/sub_to_global_mesh_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/sub_to_global_mesh_func.py
diff --git a/test/auto_parallel/pir/semi_auto_parallel_simple_net_ep.py b/test/auto_parallel/pir/semi_auto_parallel_simple_net_ep.py
diff --git a/test/auto_parallel/pir/test_semi_auto_parallel_simple_net_ep.py b/test/auto_parallel/pir/test_semi_auto_parallel_simple_net_ep.py

Original file line number	Diff line number	Diff line change
`@@ -832,6 +832,8 @@ void BindAutoParallel(py::module *m) {`
`832`	`832`	`DistributedMapper mapper;`
`833`	`833`	`return mapper.to_string();`
`834`	`834`	`});`
	`835`	`+`
	`836`	`+ m->def("sub_mesh_dim", &phi::distributed::SubMeshDim);`
`835`	`837`	`}`
`836`	`838`
`837`	`839`	`static void parse_tensors(PyObject *obj,`
Original file line number	Diff line number	Diff line change
`@@ -60,11 +60,13 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):`
`60`	`60`	`else:`
`61`	`61`	`dst_value = paddle.assign(src_value)`
`62`	`62`
`63`		`- src_chunk_id = src_value.get_defining_op().dist_attr.chunk_id`
	`63`	`+ chunk_id = -1`
	`64`	`+ if src_value.get_defining_op().dist_attr:`
	`65`	`+ chunk_id = src_value.get_defining_op().dist_attr.chunk_id`
`64`	`66`	`dst_value.set_type(dst_type)`
`65`	`67`	`dst_value.get_defining_op().dist_attr = (`
`66`	`68`	`paddle.base.libpaddle.pir.create_op_dist_attribute(`
`67`		`- dst_mesh, [src_dist_attr], [dst_dist_attr], src_chunk_id`
	`69`	`+ dst_mesh, [src_dist_attr], [dst_dist_attr], chunk_id`
`68`	`70`	`)`
`69`	`71`	`)`
`70`	`72`