[Dist Dialect] Simple MoE training in PIR (#66750)

* refine and rename the MoE apis

* add error test in ut

* add pass and unit test for pir moe model

add sub_to_global_reshard function, the loss of PIR MoE demo is equal to dygraph auto parallel

refine and rename the MoE apis

* adapt arg type in shard_dataloader

* adapt grad_clip and CBroadcastOp in pir

* refine moe api to support multi-dimension process mesh

* refine moe api

* add spmd rule for ont_hot and cumsum in yaml

* revert shard_dataloader

* add unit test for reshape_grad with resharding x_grad

* fix bug in unit test

* revert shard_dataloader api args typing

* add

* add

---------

Co-authored-by: zhangbo9674 <[email protected]>

Author: pkuzyc

paddle/fluid/pir/dialect/operator/utils/utils.cc

@@ -39,6 +39,7 @@ const std::unordered_set<std::string> LegacyOpList = {
     LoadCombineOp::name(),
     CConcatOp::name(),
     CBroadcast_Op::name(),
+    CBroadcastOp::name(),
     CSyncCommStream_Op::name(),
     DistributedPushSparseOp::name(),
     SendV2Op::name(),

paddle/phi/ops/yaml/ops.yaml

@@ -1201,6 +1201,7 @@
   output : Tensor(out)
   infer_meta :
     func : CumScalarAxisInferMeta
+    spmd_rule : CumSumInferSpmdDynamic
   kernel :
     func : cumsum
     data_type : x
@@ -3511,6 +3512,7 @@
   output : Tensor(out)
   infer_meta :
     func : OneHotInferMeta
+    spmd_rule : OneHotInferSpmdDynamic
   kernel :
     func : one_hot
   traits : paddle::dialect::ForwardOnlyTrait

python/paddle/distributed/auto_parallel/api.py

@@ -329,6 +329,7 @@ def forward(
         ctx,
         local_tensor_list,
         local_mesh_list,
+        local_placements,
         idx,
         global_dims,
         mesh,
@@ -338,17 +339,15 @@ def forward(
         if local_tensor.is_dist():
             local_mesh = local_tensor.process_mesh
             local_val = local_tensor._local_value()
-            local_placement = local_tensor.placements[0]
         else:
             local_val = local_tensor
             local_mesh = None
-            local_placement = dist.Replicate()
 
         ctx.global_mesh = copy.deepcopy(mesh)
         ctx.placements = placements
         ctx.local_dims = local_tensor.shape
         ctx.local_mesh_list = copy.deepcopy(local_mesh_list)
-        ctx.local_placement = local_placement
+        ctx.local_placements = local_placements
 
         place = paddle.framework._current_expected_place()
         place = paddle.framework._get_paddle_place(place)
@@ -360,7 +359,7 @@ def forward(
             placements=placements,
             place=place,
         )
-        global_tensor.stop_gradient = False
+        global_tensor.stop_gradient = local_tensor.stop_gradient
         return global_tensor
 
     @staticmethod
@@ -377,91 +376,111 @@ def backward(ctx, grad_tensor):
                         grad_tensor._local_value(),
                         dims=ctx.local_dims,
                         process_mesh=local_mesh,
-                        placements=[ctx.local_placement],
+                        placements=ctx.local_placements,
                         place=place,
                     )
                 )
                 out[-1].get_tensor()._unsafe_set_skip_check_mesh(True)
             return out
 
 
-def get_sub_meshes_from_global_mesh(
-    global_mesh, global_placements, local_mesh_dim
-):
-    if (
-        global_mesh is not None
-        and local_mesh_dim is not None
-        and global_placements is not None
+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
     ):
-        mesh_shape = global_mesh.shape
-        mesh_ndim = len(mesh_shape)
-        if local_mesh_dim >= mesh_ndim or (
-            local_mesh_dim < 0 and -local_mesh_dim > mesh_ndim
-        ):
-            raise ValueError(
-                f"The local_mesh_dim should between (-{mesh_ndim}, {mesh_ndim}]"
-            )
-        if local_mesh_dim < 0:
-            local_mesh_dim += mesh_ndim
-    else:
         raise ValueError(
-            "the args global_mesh, global_placements and local_mesh_dim should all be set."
+            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[local_mesh_dim], axis=local_mesh_dim
+        process_ids, mesh_shape[sub_mesh_dim], axis=sub_mesh_dim
     )
-    local_mesh_list = []
-    for process_ids in splitted_process_ids:
-        local_mesh_list.append(dist.ProcessMesh(process_ids))
+    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
+):
+    if global_mesh is None or sub_mesh_dim is None or global_placements is None:
+        raise ValueError(
+            "the args global_mesh, global_placements and local_mesh_dim should all be set."
+        )
+
+    sub_mesh_list = split_mesh(global_mesh, sub_mesh_dim)
+
     local_placements = list(global_placements)
-    local_placements.pop(local_mesh_dim)
-    if local_placements == []:
-        local_placements.append(dist.Replicate())
-    return local_mesh_list, local_placements
+    if sub_mesh_dim < len(local_placements):
+        local_placements[sub_mesh_dim] = dist.Replicate()
+
+    return sub_mesh_list, local_placements
+
+
+def cal_global_shape(local_shape, mesh, placements):
+    # assume the each rank has the same tensor shape for now,
+    # just use the local shape to calculate the global shape
+    global_shape = list(local_shape)
+    for idx, placement in enumerate(placements):
+        if placement.is_shard():
+            shard_dim = placement.get_dim()
+            local_dim_size = global_shape[shard_dim]
+            global_shape[shard_dim] = local_dim_size * mesh.shape[idx]
+    return global_shape
 
 
 def moe_global_mesh_tensor(
     local_tensor_list, mesh, placements, local_mesh_dim=-1
 ):
-    # assume the each rank has the same tensor shape for now, just use the local shape to calculate the global shape
-    local_mesh_list, local_placements = get_sub_meshes_from_global_mesh(
+    local_mesh_list, local_placements = _get_sub_meshes_and_local_placements(
         mesh, placements, local_mesh_dim
     )
-
-    local_tensor_idx = mesh.process_ids.index(dist.get_rank())
+    process_ids = np.array(mesh.process_ids).reshape(mesh.shape)
+    local_coord = np.where(process_ids == dist.get_rank())
+    local_tensor_idx = local_coord[local_mesh_dim][0]
+    # local_tensor_idx = mesh.process_ids.index(dist.get_rank())
     local_tensor = local_tensor_list[local_tensor_idx]
-    global_dims = list(local_tensor.shape)
-    for idx, placement in enumerate(placements):
-        if placement.is_shard():
-            shard_dim = placement.get_dim()
-            local_dim_size = global_dims[shard_dim]
-            global_dims[shard_dim] = local_dim_size * mesh.shape[idx]
 
     if paddle.in_dynamic_mode():
+        global_dims = cal_global_shape(
+            local_tensor._local_value().shape, mesh, placements
+        )
         resharded_local_tensor_list = []
         for i, tensor in enumerate(local_tensor_list):
             tensor.get_tensor()._unsafe_set_skip_check_mesh(True)
             if (
-                tensor.placements != local_placements
+                not check_placements_equal(tensor.placements, local_placements)
                 or tensor.process_mesh != local_mesh_list[i]
             ):
                 resharded_local_tensor_list.append(
                     reshard(tensor, local_mesh_list[i], local_placements)
                 )
+                resharded_local_tensor_list[
+                    -1
+                ].get_tensor()._unsafe_set_skip_check_mesh(True)
             else:
                 resharded_local_tensor_list.append(tensor)
 
         return _moe_global_mesh_tensor.apply(
             resharded_local_tensor_list,
             local_mesh_list,
+            local_placements,
             local_tensor_idx,
             global_dims,
             mesh,
             placements,
         )
     elif paddle.framework.in_pir_mode():
+        global_dims = cal_global_shape(
+            local_tensor._local_shape, mesh, placements
+        )
         dist_tensor = paddle._C_ops.moe_global_mesh_tensor(
             local_tensor_list,
             local_mesh_list,
@@ -487,11 +506,13 @@ def forward(
         dist_tensor,
         local_mesh_list=None,
         local_placements=None,
+        local_mesh_dim=None,
         global_mesh=None,
         global_placements=None,
     ):
         ctx.local_mesh_list = copy.deepcopy(local_mesh_list)
         ctx.local_placements = local_placements
+        ctx.local_mesh_dim = local_mesh_dim
         ctx.global_mesh = copy.deepcopy(global_mesh)
         ctx.global_placements = global_placements
         ctx.global_shape = dist_tensor.shape
@@ -532,20 +553,24 @@ def forward(
                     place=place,
                 )
                 local_tensor.get_tensor()._unsafe_set_skip_check_mesh(True)
-                local_tensor.stop_gradient = False
+                local_tensor.stop_gradient = dist_tensor.stop_gradient
                 local_tensor_list.append(local_tensor)
             return local_tensor_list
 
     @staticmethod
     def backward(ctx, *grad_tensor):
         place = paddle.framework._current_expected_place()
         place = paddle.framework._get_paddle_place(place)
-        idx = ctx.global_mesh.process_ids.index(dist.get_rank())
-        local_grad = grad_tensor[idx]
+        # idx = ctx.global_mesh.process_ids.index(dist.get_rank())
+        mesh = ctx.global_mesh
+        process_ids = np.array(mesh.process_ids).reshape(mesh.shape)
+        local_coord = np.where(process_ids == dist.get_rank())
+        local_tensor_idx = local_coord[ctx.local_mesh_dim][0]
+        local_grad = grad_tensor[local_tensor_idx]
         global_tensor = paddle.Tensor(
             local_grad._local_value(),
             dims=ctx.global_shape,
-            process_mesh=ctx.global_mesh,
+            process_mesh=mesh,
             placements=ctx.global_placements,
             place=place,
         )
@@ -558,7 +583,7 @@ def moe_sub_mesh_tensors(
     """
     Get the local part of the ``dist_tensor`` on the specific ``local_mesh_dim``.
     """
-    local_mesh_list, local_placements = get_sub_meshes_from_global_mesh(
+    local_mesh_list, local_placements = _get_sub_meshes_and_local_placements(
         global_mesh, global_placements, local_mesh_dim
     )
 
@@ -567,6 +592,7 @@ def moe_sub_mesh_tensors(
             dist_tensor,
             local_mesh_list,
             local_placements,
+            local_mesh_dim,
             global_mesh,
             global_placements,
         )

python/paddle/distributed/auto_parallel/static/engine.py

@@ -769,7 +769,6 @@ def _parallel_pir(self, mode):
 
         # re-run apply_mix2dist_pass to dist accumulator.
         apply_mix2dist_pass(dist_program)
-        # print('program', startup_program, dist_program, flush=1)
 
         # Part 2: Parallelism search (for full auto-parallel)
         # NOTE make all parallelis search logic work as Pass,
@@ -791,7 +790,7 @@ def _parallel_pir(self, mode):
 
         # Part 3: Graph partition
         # TODO(JZ-LIANG) Step 3.1: Partition Pass
-        #   insert reshard op if operand tensor's placements if different from what the cumsumer op need.
+        #   insert reshard op if operand tensor's placements is different from what the cumsumer op need.
         #   Partition the computation graph into different pipeline stage if need.
         apply_partition_pass(dist_program)
 

python/paddle/distributed/auto_parallel/static/pir_pass.py

@@ -284,6 +284,88 @@ def _remove_other_rank_params_grads(dist_params_grads):
         dist_params_grads.pop(idx)
 
 
+# Replace the specific MoE-related dist op with the
+# executable op in the dense program. In expert parallelism
+# of the MoE model, the process mesh of each expert is
+# different. Two specific apis are used to transform the
+# input tensor's global process mesh to the experts' local
+# process meshes, which will add two dist ops in the program.
+# The following two functions are used to replace the two
+# dist ops with the executable share_data_ ops.
+def replace_moe_sub_mesh_tensors(op):
+    cur_rank = paddle.distributed.get_rank()
+    in_value = op.operand_source(0)
+    out_value = None
+    out_idx = -1
+    for idx, val in enumerate(op.results()):
+        val_mesh = val.dist_attr().process_mesh
+        if cur_rank in val_mesh.process_ids:
+            assert (
+                out_value is None
+            ), f'{op} has more than one results on rank {cur_rank}'
+            out_value = val
+            out_idx = idx
+
+    paddle.pir.set_insertion_point(op)
+    local_value = paddle._C_ops.share_data_(in_value)
+    local_value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+        out_value.type(), out_value.dist_attr()
+    )
+    local_value.set_type(local_value_type)
+    out_value.replace_all_uses_with(local_value)
+
+    op_dist_attr = op.dist_attr
+    share_data_op = local_value.get_defining_op()
+    share_data_op.dist_attr = (
+        paddle.base.libpaddle.pir.create_op_dist_attribute(
+            op_dist_attr.process_mesh,
+            [op_dist_attr.operand(0).as_tensor_dist_attr()],
+            [op_dist_attr.result(out_idx).as_tensor_dist_attr()],
+        )
+    )
+
+    assert all(val.use_empty() for val in op.results())
+    op.erase()
+
+
+def replace_moe_global_mesh_tensor(op):
+    cur_rank = paddle.distributed.get_rank()
+    out_value = op.result(0)
+    in_value = None
+    in_idx = -1
+    for idx, val in enumerate(op.operands_source()):
+        val_mesh = val.dist_attr().process_mesh
+        if cur_rank not in val_mesh.process_ids:
+            continue
+        assert (
+            in_value is None
+        ), f'{op} has more than one inputs on rank {cur_rank}'
+        in_value = val
+        in_idx = idx
+
+    paddle.pir.set_insertion_point(op)
+    local_value = paddle._C_ops.share_data_(in_value)
+    # local_value = paddle.assign(in_value)
+    local_value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+        out_value.type(), out_value.dist_attr()
+    )
+    local_value.set_type(local_value_type)
+    out_value.replace_all_uses_with(local_value)
+
+    op_dist_attr = op.dist_attr
+    share_data_op = local_value.get_defining_op()
+    share_data_op.dist_attr = (
+        paddle.base.libpaddle.pir.create_op_dist_attribute(
+            op_dist_attr.process_mesh,
+            [op_dist_attr.operand(in_idx).as_tensor_dist_attr()],
+            [op_dist_attr.result(0).as_tensor_dist_attr()],
+        )
+    )
+
+    assert all(val.use_empty() for val in op.results())
+    op.erase()
+
+
 # pruning op and value not belong to cur rank
 def remove_other_rank_op_pass(dist_program, dist_params_grads):
     cur_rank = paddle.distributed.get_rank()
@@ -298,6 +380,12 @@ def remove_other_rank_op_pass(dist_program, dist_params_grads):
             if can_delete:
                 op.erase()
             continue
+        if op.name() == "dist_op.moe_sub_mesh_tensors":
+            replace_moe_sub_mesh_tensors(op)
+            continue
+        if op.name() == "dist_op.moe_global_mesh_tensor":
+            replace_moe_global_mesh_tensor(op)
+            continue
         if cur_rank not in op.dist_attr.process_mesh.process_ids:
             op.erase()
         elif op.name() == "dist_op.reshard":