[Auto Parallel] fix hang caused by different process group initialization order (#68847)

deepllz · web-flow · commit 308c61d687cb · 2024-10-23T15:05:08.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -51,6 +51,7 @@
     convert_to_dims_mapping,
     fuse_param_func,
     get_dist_attr,
+    split_mesh,
     split_param_func,
     to_list,
 )
@@ -408,29 +409,6 @@ def backward(ctx, grad_tensor):
             return out
 
 
-def split_mesh(global_mesh: dist.ProcessMesh, sub_mesh_dim: int):
-    mesh_shape = global_mesh.shape
-    mesh_ndim = len(mesh_shape)
-    if sub_mesh_dim >= mesh_ndim or (
-        sub_mesh_dim < 0 and -sub_mesh_dim > mesh_ndim
-    ):
-        raise ValueError(
-            f"The sub_mesh_dim should between (-{mesh_ndim}, {mesh_ndim}]"
-        )
-    if sub_mesh_dim < 0:
-        sub_mesh_dim += mesh_ndim
-
-    process_ids = np.array(global_mesh.process_ids).reshape(mesh_shape)
-    splitted_process_ids = np.split(
-        process_ids, mesh_shape[sub_mesh_dim], axis=sub_mesh_dim
-    )
-    sub_mesh_list = []
-    for sub_process_ids in splitted_process_ids:
-        sub_mesh_list.append(dist.ProcessMesh(sub_process_ids))
-
-    return sub_mesh_list
-
-
 def _get_sub_meshes_and_local_placements(
     global_mesh, global_placements, sub_mesh_dim
 ):
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/nd_mesh_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/nd_mesh_reshard_func.py
@@ -15,7 +15,9 @@
 
 import paddle
 import paddle.distributed as dist
+from paddle.distributed.auto_parallel.static.utils import split_mesh
 
+from ..process_group import new_process_group
 from .base_reshard_func import (
     ReshardFunction,
     copy_dist_attr_with_new_member,
@@ -133,7 +135,9 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
             tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
                 src_value.type(), tmp_dst_dist_attr
             )
-
+            sub_mesh_list = split_mesh(process_mesh, in_mesh_axis)
+            for sub_mesh in sub_mesh_list:
+                new_process_group(sorted(sub_mesh.process_ids))
             # get the process_mesh on specific axis
             sub_mesh = get_1D_sub_process_mesh(process_mesh, in_mesh_axis)
 
@@ -197,7 +201,9 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
                 tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
                     src_value.type(), tmp_dst_dist_attr
                 )
-
+                sub_mesh_list = split_mesh(process_mesh, partial_dim)
+                for sub_mesh in sub_mesh_list:
+                    new_process_group(sorted(sub_mesh.process_ids))
                 # get the process_mesh on specific axis
                 sub_mesh = get_1D_sub_process_mesh(process_mesh, partial_dim)
 
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
@@ -58,10 +58,7 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
 
         cur_rank = paddle.distributed.get_rank()
 
-        group_ranks = sorted(dst_mesh.process_ids)
-        comm_group = new_process_group(group_ranks)
-
-        if cur_rank == root_rank:
+        if cur_rank in src_mesh.process_ids:
             # the root rank will broadcast the src_value to other ranks
             chunk_id = -1
             if src_value.get_defining_op().dist_attr:
diff --git a/python/paddle/distributed/auto_parallel/static/utils.py b/python/paddle/distributed/auto_parallel/static/utils.py
@@ -2678,3 +2678,26 @@ def split_param_func(
     else:
         # fuse_attention_ffn
         return split_fn(fused_param, split_nums, axis=-1)
+
+
+def split_mesh(global_mesh: ProcessMesh, sub_mesh_dim: int):
+    mesh_shape = global_mesh.shape
+    mesh_ndim = len(mesh_shape)
+    if sub_mesh_dim >= mesh_ndim or (
+        sub_mesh_dim < 0 and -sub_mesh_dim > mesh_ndim
+    ):
+        raise ValueError(
+            f"The sub_mesh_dim should between (-{mesh_ndim}, {mesh_ndim}]"
+        )
+    if sub_mesh_dim < 0:
+        sub_mesh_dim += mesh_ndim
+
+    process_ids = np.array(global_mesh.process_ids).reshape(mesh_shape)
+    splitted_process_ids = np.split(
+        process_ids, mesh_shape[sub_mesh_dim], axis=sub_mesh_dim
+    )
+    sub_mesh_list = []
+    for sub_process_ids in splitted_process_ids:
+        sub_mesh_list.append(ProcessMesh(sub_process_ids))
+
+    return sub_mesh_list
