Support send/recv when mesh is pp_mesh (#71433)

xuxinyi389 · web-flow · commit 9f50fc923b55 · 2025-03-10T11:35:11.000+08:00
* Support send/recv when mesh is pp_mesh

* fix_bugs

* remove_unreleated_code

* fix_error
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_mesh_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_mesh_reshard_function.cc
@@ -111,22 +111,35 @@ void SubMeshToGlobalReshardFunction::Eval(phi::DeviceContext* dev_ctx,
     const DenseTensor& in_dense_value = in.value();
     std::vector<int64_t>& recv_vec = send2recv_map[cur_global_rank];
     for (int64_t recv_id : recv_vec) {
+      auto relative_recv_rank = recv_id;
+      for (size_t i = 0; i < all_process_ids.size(); ++i) {
+        if (all_process_ids[i] == recv_id) {
+          relative_recv_rank = i;
+        }
+      }
       RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
                                 PSendKernel,
                                 dtype,
                                 all_process_ids,
                                 in_dense_value,
-                                recv_id,
+                                relative_recv_rank, /*peer*/
                                 true /*dynamic_shape*/);
     }
     SetValue(out, in_dense_value);
   } else {
     int64_t send_id = recv2send_map[cur_global_rank];
+    auto relative_send_rank = send_id;
+    for (size_t i = 0; i < all_process_ids.size(); ++i) {
+      if (all_process_ids[i] == send_id) {
+        relative_send_rank = i;
+      }
+    }
+
     RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
                               PRecv,
                               dtype,
                               all_process_ids,
-                              send_id,
+                              relative_send_rank, /*peer*/
                               {} /*out_shape*/,
                               true /*dynamic_shape*/,
                               GetMutableTensor(out));
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/r_to_x_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/r_to_x_reshard_function.cc
@@ -71,12 +71,20 @@ void RToXExpandReshardFunction::Eval(phi::DeviceContext* dev_ctx,
   if (root_rank == cur_global_rank) {
     for (const auto& out_process_id : out_process_ids) {
       if (out_process_id != root_rank) {
+        // Should acculate the relative rank in communication.
+        auto relative_rank = out_process_id;
+        for (size_t i = 0; i < all_process_ids.size(); ++i) {
+          if (all_process_ids[i] == out_process_id) {
+            relative_rank = i;
+          }
+        }
+
         RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
                                   PSendKernel,
                                   dtype,
                                   all_process_ids,
                                   in.value(),
-                                  out_process_id,
+                                  /*peer*/ relative_rank,
                                   /*dynamic_shape=*/true);
       }
     }
diff --git a/python/paddle/distributed/utils/moe_utils.py b/python/paddle/distributed/utils/moe_utils.py
@@ -14,6 +14,7 @@
 
 from paddle import _legacy_C_ops
 from paddle.common_ops_import import check_variable_and_dtype
+from paddle.distributed import fleet
 from paddle.framework import LayerHelper, in_dynamic_mode
 
 
@@ -277,3 +278,30 @@ def global_gather(
             },
         )
         return out
+
+
+def get_complete_pp_mesh(mesh):
+    """
+    Get complete pp mesh with given mesh.
+
+    Args:
+        mesh (Mesh): Mesh object.
+
+    Returns:
+        Mesh: Complete mesh.
+
+    """
+    process_id = mesh.process_ids[0]
+    global_mesh = fleet.auto.get_mesh()
+
+    if global_mesh and "pp" in global_mesh.dim_names:
+        pp_degree = global_mesh.get_dim_size("pp")
+        for i in range(pp_degree):
+            pp_mesh = global_mesh.get_mesh_with_dim("pp", i)
+            if process_id in pp_mesh.process_ids:
+                return pp_mesh
+        AssertionError(
+            f"Current mesh: {mesh} not found in global mesh {global_mesh}"
+        )
+    else:
+        return mesh
diff --git a/python/paddle/nn/clip.py b/python/paddle/nn/clip.py
@@ -26,6 +26,7 @@
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.libpaddle import DataType
 from paddle.common_ops_import import Variable, check_type, default_main_program
+from paddle.distributed.utils.moe_utils import get_complete_pp_mesh
 from paddle.framework import (
     LayerHelper,
     in_dynamic_mode,
@@ -729,6 +730,12 @@ def _dygraph_clip(self, params_grads):
             # if the gradient mesh is not equal to src mesh
             # do reshard to get the result of squared_l2 from other pp stage mesh
             if src_mesh is not None and g.process_mesh != src_mesh:
+                pp_mesh = get_complete_pp_mesh(g.process_mesh)
+                if set(g.process_mesh.process_ids) < set(pp_mesh.process_ids):
+                    sum_square = dist.reshard(
+                        sum_square, pp_mesh, sum_square.placements
+                    )
+
                 sum_square = dist.reshard(
                     sum_square, src_mesh, sum_square.placements
                 )
@@ -821,6 +828,15 @@ def async_add_n(var_list):
                                 "Reshard a sharded tensor from a local mesh to a global mesh is not supported"
                             )
                     else:
+                        pp_mesh = get_complete_pp_mesh(g.process_mesh)
+
+                        if set(g.process_mesh.process_ids) < set(
+                            pp_mesh.process_ids
+                        ):
+                            clip_input = dist.reshard(
+                                clip_input, pp_mesh, clip_input.placements
+                            )
+
                         clip_input = paddle.distributed.reshard(
                             clip_input, g.process_mesh, clip_input.placements
                         )
