chore: import upstream snapshot with attribution

python/paddle/distributed/auto_parallel/static/operators/__init__.py

@@ -0,0 +1,61 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+import os
+
+from . import (  # noqa: F401
+    dist_assign,
+    dist_check_finite_and_unscale,
+    dist_concat,
+    dist_default,
+    dist_dropout,
+    dist_eltwise,
+    dist_embedding,
+    dist_expand_as,
+    dist_fill_constant_batch_size_like,
+    dist_flash_attn,
+    dist_fused_attention,
+    dist_fused_dropout_add,
+    dist_fused_feedforward,
+    dist_fused_rms_norm,
+    dist_fused_rope,
+    dist_gather_nd,
+    dist_layer_norm,
+    dist_matmul,
+    dist_pnorm,
+    dist_reduce_sum_p,
+    dist_reshape,
+    dist_scale,
+    dist_shape,
+    dist_slice,
+    dist_softmax,
+    dist_split,
+    dist_stack,
+    dist_strided_slice,
+    dist_tile,
+    dist_transpose,
+    dist_unsqueeze2,
+    dist_update_loss_scaling,
+)
+from .common import (  # noqa: F401
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    find_compatible_distributed_operator_impls,
+    find_distributed_operator_impl_container,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+
+parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")
+if parallel_ce == "true":
+    from . import dist_cross_entropy  # noqa: F401

python/paddle/distributed/auto_parallel/static/operators/common.py

@@ -0,0 +1,838 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import abc
+import logging
+import warnings
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.log_helper import get_logger
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+
+from ..dist_attribute import OperatorDistAttr
+from ..process_group import new_process_group
+from ..utils import (
+    _get_comm_group,
+    _get_corresponding_rank,
+    compute_compatible_dims_mapping,
+    is_optimize_op,
+    set_dist_op_desc_original_id,
+)
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+_g_distributed_operator_impl_containers = {}
+
+_g_elementwise_ops = [
+    "assign",
+    "elementwise",
+    "gelu",
+    # "dropout",
+    "scale",
+    "relu",
+    "cast",
+    # "gather",
+    # "concat",
+    "silu",
+    "fused_softmax_mask_upper_triangle",
+]
+BACKWARD_ONLY_DIST_OPS = {'check_finite_and_unscale', 'update_loss_scaling'}
+
+_gradient_sync_by_partial_ops = [
+    "matmul_v2_grad",
+    "elementwise_add_grad",
+    "layer_norm_grad",
+    "lookup_table_v2_grad",
+    # "conv",
+]
+
+
+class ParallelMode:
+    """
+    the parallel mode for communication or auxiliary operator
+    """
+
+    DataParallel = "auto_parallel/data_parallel"
+    TensorParallel = "auto_parallel/tensor_parallel"
+    PipelineParallel = "auto_parallel/pipeline_parallel"
+    MoEParallel = "auto_parallel/moe_parallel"
+
+
+class SyncMode:
+    """
+    the synchronization mode for communication or auxiliary operator
+    """
+
+    AmpFlagSync = "auto_parallel/amp_flag_synchronization"
+    GlobalNormSync = "auto_parallel/global_norm_synchronization"
+
+
+def is_elementwise_op(op_type):
+    if op_type in _g_elementwise_ops:
+        return True
+    if "elementwise" in op_type:
+        return True
+    return False
+
+
+class DistributedOperatorImplContainer(abc.ABC):
+    def __init__(self, op_type):
+        self._type = op_type
+        self._impls = []
+
+    @property
+    def type(self):
+        return self._type
+
+    @type.setter
+    def type(self, op_type):
+        self._type = op_type
+
+    @property
+    def impls(self):
+        return self._impls
+
+    def register_impl(self, dist_impl):
+        assert self.type == dist_impl.type, (
+            "Op type of container must be same as that of the implementation."
+        )
+        impl_idx = len(self.impls)
+        dist_impl.idx = impl_idx
+        self._impls.append(dist_impl)
+
+    def get_impl(self, impl_idx):
+        return self._impls[impl_idx]
+
+    def get_input_compatible_impls(self, dist_op):
+        compatible_impls = []
+        for impl in self.impls:
+            if impl.is_input_compatible(dist_op):
+                compatible_impls.append(impl)
+        return compatible_impls
+
+    def get_output_compatible_impls(self, dist_op):
+        compatible_impls = []
+        for impl in self.impls:
+            if impl.is_output_compatible(dist_op):
+                compatible_impls.append(impl)
+        return compatible_impls
+
+    def get_compatible_impls(self, dist_op):
+        compatible_impls = []
+        for impl in self.impls:
+            if impl.is_auto_compatible(dist_op):
+                compatible_impls.append(impl)
+        return compatible_impls
+
+    # (NOTE) Currently, both DistributedOperatorImplContainer and DistributedOperatorImpl have update_dims_mapping method.
+    # But this method is supposed to be maintained by DistributedOperatorImplContainer, and we are ongoing adding method
+    # to DistributedOperatorImplContainer and removing those in DistributedOperatorImpl.
+    # @abc.abstractmethod
+    def update_dims_mapping(self, dist_op):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+    # (NOTE) Currently we has limited DistributedOperatorImpls for an op to deal with different parallel patterns of this op.
+    # This function help to choose the correct DistributedOperatorImpl based on the result from InferSPMD.
+    # @abc.abstractmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+
+class DistributedOperatorImpl(abc.ABC):
+    def __init__(self, name):
+        self._name = name
+        self._type = None
+        self._idx = None
+        self._forward_implemented = False
+        self._backward_implemented = False
+
+    @property
+    def name(self):
+        return self._name
+
+    @name.setter
+    def name(self, name):
+        self._name = name
+
+    @property
+    def type(self):
+        return self._type
+
+    @type.setter
+    def type(self, op_type):
+        self._type = op_type
+
+    @property
+    def idx(self):
+        return self._idx
+
+    @idx.setter
+    def idx(self, impl_idx):
+        self._idx = impl_idx
+
+    # to be deprecated
+    @abc.abstractmethod
+    def is_input_compatible(self, dist_op):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+    # to be deprecated
+    @abc.abstractmethod
+    def is_output_compatible(self, dist_op):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+    # to be deprecated
+    @abc.abstractmethod
+    def is_auto_compatible(self, dist_op):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+    @staticmethod
+    @abc.abstractmethod
+    def forward(dist_ctx, *args, **kwargs):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+    @staticmethod
+    @abc.abstractmethod
+    def backward(dist_ctx, *grad_outputs, **kwargs):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+    # to be deprecated
+    def update_dims_mapping(self, dist_op):
+        raise NotImplementedError("Please Implement this method in Subclass.")
+
+
+def register_distributed_operator_impl_container(container):
+    global _g_distributed_operator_impl_containers
+    _g_distributed_operator_impl_containers[container.type] = container
+
+
+def get_distributed_operator_impl_container(op_type):
+    global _g_distributed_operator_impl_containers
+    return _g_distributed_operator_impl_containers.get(op_type, None)
+
+
+def register_distributed_operator_impl(op_type, dist_impl):
+    dist_op_impl_container = get_distributed_operator_impl_container(op_type)
+    if dist_op_impl_container is not None:
+        dist_impl.type = op_type
+        dist_op_impl_container.register_impl(dist_impl)
+    else:
+        raise AssertionError(
+            "Must register distributed operator registry first."
+        )
+
+
+def find_compatible_distributed_operator_impls(dist_op, fwd=True, partial=True):
+    """
+    Here just return the first compatible implementation.
+    This will be improved by cost model in the future.
+    """
+    op_type = dist_op.serial_op.type
+    dist_op_impl_container = get_distributed_operator_impl_container(op_type)
+    dist_op_eltwise_impl_container = get_distributed_operator_impl_container(
+        "elementwise"
+    )
+    dist_op_default_impl_container = get_distributed_operator_impl_container(
+        "default"
+    )
+    compatible_impls = []
+    if partial:
+        if fwd:
+            # First, find impls in the corresponding container
+            if dist_op_impl_container:
+                compatible_impls.extend(
+                    dist_op_impl_container.get_input_compatible_impls(dist_op)
+                )
+            # Second, find impls in the elementwise container
+            if dist_op_eltwise_impl_container and is_elementwise_op(op_type):
+                compatible_impls.extend(
+                    dist_op_eltwise_impl_container.get_input_compatible_impls(
+                        dist_op
+                    )
+                )
+            # Third, find impls in the default container
+            if dist_op_default_impl_container:
+                compatible_impls.extend(
+                    dist_op_default_impl_container.get_input_compatible_impls(
+                        dist_op
+                    )
+                )
+        else:
+            # First, find impls in the corresponding container
+            if dist_op_impl_container:
+                compatible_impls.extend(
+                    dist_op_impl_container.get_output_compatible_impls(dist_op)
+                )
+            # Second, find impls in the elementwise container
+            if dist_op_eltwise_impl_container and is_elementwise_op(op_type):
+                compatible_impls.extend(
+                    dist_op_eltwise_impl_container.get_output_compatible_impls(
+                        dist_op
+                    )
+                )
+            # Third, find impls in the default container
+            if dist_op_default_impl_container:
+                compatible_impls.extend(
+                    dist_op_default_impl_container.get_output_compatible_impls(
+                        dist_op
+                    )
+                )
+    else:
+        # First, find impls in the corresponding container
+        if dist_op_impl_container:
+            compatible_impls.extend(
+                dist_op_impl_container.get_compatible_impls(dist_op)
+            )
+        # Second, find impls in the elementwise container
+        if dist_op_eltwise_impl_container and is_elementwise_op(op_type):
+            compatible_impls.extend(
+                dist_op_eltwise_impl_container.get_compatible_impls(dist_op)
+            )
+        # Third, find impls in the default container
+        if dist_op_default_impl_container:
+            compatible_impls.extend(
+                dist_op_default_impl_container.get_compatible_impls(dist_op)
+            )
+
+    if compatible_impls:
+        # For now, just return the first compatible impl
+        # best_compatible_impl = compatible_impls[0]
+        best_compatible_impl = compatible_impls
+    else:
+        best_compatible_impl = None
+    return best_compatible_impl
+
+
+def find_distributed_operator_impl_container(dist_op):
+    """
+    Return a unique container for dist op.
+    If not specific container found, default container will be return.
+    """
+    op_type = dist_op.serial_op.type
+
+    # Op has a  match container
+    dist_op_impl_container = get_distributed_operator_impl_container(op_type)
+    if dist_op_impl_container is None:
+        # if op is register to elemwise spmd rule and has NO specific container implemented
+        if is_elementwise_op(op_type):
+            dist_op_impl_container = get_distributed_operator_impl_container(
+                "elementwise"
+            )
+        # default container for all bottom line cases
+        else:
+            dist_op_impl_container = get_distributed_operator_impl_container(
+                "default"
+            )
+
+    _logger.debug(
+        f"Op [{op_type}] Complete DistAttr using {type(dist_op_impl_container).__name__}"
+    )
+    return dist_op_impl_container
+
+
+def is_parameter_related(varname, block, dist_context=None):
+    # TODO(zhaoyingli): maintain a dict in dist_context to record all variables which are be renamed
+    if ".subprog_" in varname:
+        varname = varname[: varname.index(".subprog_")]
+    if ".cast_fp" in varname:
+        varname = varname[: varname.index(".cast_fp")]
+    if ".cast_bf" in varname:
+        varname = varname[: varname.index(".cast_bf")]
+    if ".quantized" in varname:
+        varname = varname[: varname.index(".quantized")]
+    assert block._find_var_recursive(varname), (
+        f"cannot find var {varname} in cur block"
+    )
+    var = block._var_recursive(varname)
+    # NOTE(hack method): to find the param which is resharded
+    if dist_context and "@RESHARD" in varname:
+        varname = varname[: varname.index("@RESHARD")]
+        serial_program = dist_context.serial_main_program
+        var = serial_program.global_block()._find_var_recursive(varname)
+        if var is None:
+            return False
+    # NOTE(liym27): when Y_var is not a parameter, but Y_var is resharded by a parameter.
+    elif "reshard_api" in varname:
+        for op in block.ops:
+            if op.type == "assign" and varname in op.output("Out"):
+                in_varname = op.input("X")[0]
+                var = block._find_var_recursive(in_varname)
+                if var is not None and var.is_parameter:
+                    return True
+    return var.is_parameter
+
+
+def infer_shape(block, src_var, src_var_dist_attr, op_input_dist_attr):
+    var_shape = block._var_recursive(src_var.name).shape
+    var_topology = src_var_dist_attr.process_mesh.shape
+    var_dims_mapping = src_var_dist_attr.dims_mapping
+
+    complete_shape = []
+    for idx, shape in enumerate(var_shape):
+        if var_dims_mapping[idx] == -1:
+            complete_shape.append(shape)
+        else:
+            new_shape = shape * var_topology[var_dims_mapping[idx]]
+            complete_shape.append(new_shape)
+
+    exact_shape = []
+    input_topology = op_input_dist_attr.process_mesh.shape
+    input_dims_mapping = op_input_dist_attr.dims_mapping
+    for idx, shape in enumerate(complete_shape):
+        if input_dims_mapping[idx] == -1:
+            exact_shape.append(shape)
+        else:
+            new_shape = shape // input_topology[input_dims_mapping[idx]]
+            exact_shape.append(new_shape)
+
+    return exact_shape
+
+
+def set_comm_op_dist_attr_for_program(
+    new_op, process_mesh, tensor_dist_attr, ctx, **kwargs
+):
+    assert process_mesh is not None
+    assert tensor_dist_attr is not None
+
+    new_op_dist_attr = OperatorDistAttr()
+    new_op_dist_attr.process_mesh = process_mesh
+    if "chunk_id" in kwargs:
+        new_op_dist_attr.chunk_id = kwargs["chunk_id"]
+    for input_varname in new_op.desc.input_arg_names():
+        new_op_dist_attr.set_input_dist_attr(input_varname, tensor_dist_attr)
+    for output_varname in new_op.desc.output_arg_names():
+        new_op_dist_attr.set_output_dist_attr(output_varname, tensor_dist_attr)
+    ctx.set_op_dist_attr_for_program(new_op, new_op_dist_attr)
+
+
+def naive_copy_op_dist_attr_for_program(new_op, ref_op, ctx):
+    ref_dist_attr = ctx.get_op_dist_attr_for_program(ref_op)
+    new_op_dist_attr = OperatorDistAttr()
+    new_op_dist_attr.process_mesh = ref_dist_attr.process_mesh
+    new_op_dist_attr.impl_type = ref_dist_attr.impl_type
+    new_op_dist_attr.impl_idx = ref_dist_attr.impl_idx
+    new_op_dist_attr.chunk_id = ref_dist_attr.chunk_id
+
+    for input_name in ref_op.input_names:
+        assert input_name in new_op.input_names
+        assert len(ref_op.input(input_name)) == 1
+        assert len(new_op.input(input_name)) == 1
+
+        ref_tensor_dist_attr = ref_dist_attr.get_input_dist_attr(
+            ref_op.input(input_name)[0]
+        )
+        new_op_dist_attr.set_input_dist_attr(
+            new_op.input(input_name)[0], ref_tensor_dist_attr
+        )
+
+    for output_name in ref_op.output_names:
+        assert output_name in new_op.output_names
+        assert len(ref_op.output(output_name)) == 1
+        assert len(new_op.output(output_name)) == 1
+
+        ref_tensor_dist_attr = ref_dist_attr.get_output_dist_attr(
+            ref_op.output(output_name)[0]
+        )
+        new_op_dist_attr.set_output_dist_attr(
+            new_op.output(output_name)[0], ref_tensor_dist_attr
+        )
+
+    ctx.set_op_dist_attr_for_program(new_op, new_op_dist_attr)
+
+
+def get_data_parallel_group(dist_ctx, op, act_grad_names, rank):
+    """
+    deduce the data parallel communication group for current operator.
+
+    Args:
+        dist_ctx (DistributedContext): dist context.
+        op (Operator): the current (backward) operator which might need.
+        act_grad_names (list): list of input activation grads variable name to the current operator.
+        rank (int): global ranks index for current process.
+    """
+    dp_group = None
+
+    op_dist_attr = dist_ctx.get_op_dist_attr_for_program(op)
+    process_mesh = op_dist_attr.process_mesh
+    mesh_shape = process_mesh.shape
+    # FIXME Hack for Pipeline Parallelism where the current operator
+    # not belong to the mesh the current rank belong to.
+    if rank not in process_mesh.process_ids:
+        rank = _get_corresponding_rank(dist_ctx, process_mesh, rank)
+
+    for var_name in act_grad_names:
+        var_dim_mapping = op_dist_attr.get_input_dims_mapping(var_name)
+        # consider that the variable's shape is [], which is 0-D
+        # TODO utilize the batch_dim attr instead of "0" in future
+        batch_size_axis = var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+
+        if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+            group_ranks = _get_comm_group(
+                process_mesh.process_ids,
+                process_mesh.shape,
+                batch_size_axis,
+                rank,
+            )
+            dp_group = new_process_group(group_ranks)
+            break
+    if dp_group is not None:
+        return [dp_group]
+    else:
+        return []
+
+
+def sync_and_scale_gradients(dist_ctx, op, groups, allreduce_var_names):
+    """
+    insert the allreduce and scale ops for gradients of model
+    parameters for operator in data parallelism.
+
+    Args:
+        dist_ctx (DistributedContext): dist context.
+        op (Operator): the current (backward) operator which might need.
+        allreduce_var_names (list): list of the parameter's grads variable name in the current operator output.
+    """
+
+    op_dist_attr = dist_ctx.get_op_dist_attr_for_program(op)
+    process_mesh = op_dist_attr.process_mesh
+    chunk_id = op_dist_attr.chunk_id
+    dist_op_context = dist_ctx.dist_op_context
+    main_block = dist_op_context.work_block
+
+    reduce_type = dist.ReduceOp.SUM
+    need_scale = dist_ctx.gradient_scale
+
+    for group in groups:
+        group_size = len(group.ranks)
+
+        for var_name in allreduce_var_names:
+            added_ops = []
+            grad_var = main_block.var(var_name)
+            allreduce_op = main_block.append_op(
+                type='all_reduce',
+                inputs={'x': [grad_var]},
+                outputs={'out': [grad_var]},
+                attrs={
+                    'ring_id': group.id,
+                    'reduce_type': reduce_type,
+                    OP_ROLE_KEY: OpRole.Backward,
+                },
+            )
+            allreduce_op._set_attr(
+                'op_namescope', '/' + ParallelMode.DataParallel
+            )
+            added_ops.append(allreduce_op)
+
+            if need_scale:
+                scale_op = main_block.append_op(
+                    type='scale',
+                    inputs={'X': grad_var},
+                    outputs={'Out': grad_var},
+                    attrs={
+                        'scale': 1.0 / group_size,
+                        OP_ROLE_KEY: OpRole.Backward,
+                    },
+                )
+                scale_op._set_attr(
+                    'op_namescope', '/' + ParallelMode.DataParallel
+                )
+                added_ops.append(scale_op)
+
+            dims_mapping = op_dist_attr.get_output_dims_mapping(grad_var.name)
+            assert dims_mapping is not None, (
+                f"Unexpected: dims_mapping of output [{grad_var.name}] of op [{op_dist_attr.op_type}] is None"
+            )
+            # NOTE auxiliary op's dist attr should follow dist_op not dist_tensor
+            for new_op in added_ops:
+                new_op_attr = OperatorDistAttr()
+                new_op_attr.process_mesh = process_mesh
+                new_op_attr.chunk_id = chunk_id
+                new_op_attr.set_output_dims_mapping(grad_var.name, dims_mapping)
+                new_op_attr.set_input_dims_mapping(grad_var.name, dims_mapping)
+                dist_ctx.set_op_dist_attr_for_program(new_op, new_op_attr)
+
+
+def get_partial_groups(dist_ctx, op, out_grad_names, rank):
+    """
+    deduce the partial communication group for current operator output vars.
+
+    Args:
+        dist_ctx (DistributedContext): dist context.
+        op (Operator): the current (backward) operator which might need.
+        out_grad_names (list): list of the output parameter's grads variable name of the current operator.
+        rank (int): global ranks index for current process.
+    """
+    op_dist_attr = dist_ctx.get_op_dist_attr_for_program(op)
+    process_mesh = op_dist_attr.process_mesh
+    mesh_shape = process_mesh.shape
+
+    groups = []
+
+    partial_dims = None
+    for var_name in out_grad_names:
+        var_dist_attr = op_dist_attr.get_output_dist_attr(var_name)
+        if partial_dims is None:
+            partial_dims = var_dist_attr._partial_dims()
+        else:
+            assert partial_dims == var_dist_attr._partial_dims(), (
+                f"Partial dims of outputs {out_grad_names} of op [{op.type}] is not consistent"
+            )
+
+    partial_dims = list(partial_dims)
+    partial_dims.sort()
+
+    # FIXME Hack for Pipeline Parallelism where the current operator
+    # not belong to the mesh the current rank belong to.
+    if rank not in process_mesh.process_ids:
+        rank = _get_corresponding_rank(dist_ctx, process_mesh, rank)
+
+    for dim in partial_dims:
+        if mesh_shape[dim] > 1:
+            group_ranks = _get_comm_group(
+                process_mesh.process_ids,
+                process_mesh.shape,
+                dim,
+                rank,
+            )
+            groups.append(new_process_group(group_ranks))
+
+    return groups
+
+
+def gradient_synchronization(
+    dist_ctx, op, act_grad_names, out_grad_names, rank
+):
+    """
+    conduct the allreduce and scaling for gradients of model
+    parameters for operator in parallelism train.
+
+    Args:
+        dist_ctx (DistributedContext): dist context.
+        op (Operator): the current (backward) operator which might need.
+        act_grad_names (list): list of input activation grads variable name to the current operator.
+        out_grad_names (list): list of the output parameter's grads variable name of the current operator.
+        rank (int): global ranks index for current process.
+    """
+
+    if not is_in_backward_phase(dist_ctx):
+        return
+
+    if (
+        is_optimize_op(op)
+        or len(act_grad_names) == 0
+        or len(out_grad_names) == 0
+    ):
+        return
+
+    if op.type in _gradient_sync_by_partial_ops:
+        sync_groups = get_partial_groups(dist_ctx, op, out_grad_names, rank)
+    # NOTE we reverse the following old branch to support operators (e.g. fuse operators) that haven't been adopted for partial inferspmd,
+    # and remove this branch after all operators are adopted for partial inferspmd.
+    else:
+        sync_groups = get_data_parallel_group(
+            dist_ctx, op, act_grad_names, rank
+        )
+
+    if len(sync_groups) < 1:
+        return
+
+    sync_and_scale_gradients(dist_ctx, op, sync_groups, out_grad_names)
+
+
+def is_data_parallel_scale_op(op):
+    return (
+        op.type == "scale"
+        and op.desc.has_attr("op_namescope")
+        and ParallelMode.DataParallel in op.desc.attr("op_namescope")
+    )
+
+
+def is_data_parallel_reduce_op(op):
+    is_allreduce_op = op.type in [
+        "c_allreduce_sum",
+        "c_allreduce_avg",
+    ]
+    is_all_reduce_op = op.type == "all_reduce" and op.desc.attr(
+        "reduce_type"
+    ) in [
+        dist.ReduceOp.SUM,
+        dist.ReduceOp.AVG,
+    ]
+    is_reduce_op = op.type == "reduce" and op.desc.attr("reduce_type") in [
+        dist.ReduceOp.SUM,
+        dist.ReduceOp.AVG,
+    ]
+    return (
+        (is_allreduce_op or is_all_reduce_op or is_reduce_op)
+        and op.desc.has_attr("op_namescope")
+        and ParallelMode.DataParallel in op.desc.attr("op_namescope")
+    )
+
+
+def is_amp_flag_sync_op(op):
+    return (
+        op.type == "all_reduce"
+        and op.desc.attr("op_type") == paddle.distributed.ReduceOp.MAX
+        and op.desc.has_attr("op_namescope")
+        and SyncMode.AmpFlagSync in op.desc.attr("op_namescope")
+    )
+
+
+def is_global_norm_sync_op(op):
+    return (
+        op.type == "all_reduce"
+        and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
+        and op.desc.has_attr("op_namescope")
+        and SyncMode.GlobalNormSync in op.desc.attr("op_namescope")
+    )
+
+
+def is_in_backward_phase(dist_ctx):
+    # NOTE currently high-order differential in Paddle dose NOT distinguish gradient computation operators
+    # in Forward phase and operators in Backward phase (both with op_role=1), which will mislead
+    # auto parallel to add gradient synchronization for gradient computation operators in Forward phase.
+    # we use this FLAG to distinguish these two phases temporarily.
+
+    return dist_ctx.dist_op_context.in_backward_phase()
+
+
+def merge_forward_backward_dims_mapping(fw_results, bw_results):
+    flatten_fw_inputs = paddle.utils.flatten(fw_results[0])
+    flatten_fw_outputs = paddle.utils.flatten(fw_results[1])
+    flatten_bw_inputs = paddle.utils.flatten(bw_results[0])
+    flatten_bw_outputs = paddle.utils.flatten(bw_results[1])
+    ninputs = len(flatten_fw_inputs)
+    noutputs = len(flatten_fw_outputs)
+    inferred_input_dims_mappings = []
+    inferred_output_dims_mappings = []
+
+    for i in range(ninputs):
+        compatible_dims_mapping = compute_compatible_dims_mapping(
+            [
+                flatten_fw_inputs[i].dims_mapping,
+                flatten_bw_inputs[i].dims_mapping,
+            ]
+        )
+        inferred_input_dims_mappings.append(compatible_dims_mapping)
+
+    for i in range(noutputs):
+        compatible_dims_mapping = compute_compatible_dims_mapping(
+            [
+                flatten_fw_outputs[i].dims_mapping,
+                flatten_bw_outputs[i].dims_mapping,
+            ]
+        )
+        inferred_output_dims_mappings.append(compatible_dims_mapping)
+    return inferred_input_dims_mappings, inferred_output_dims_mappings
+
+
+def update_op_dims_mapping(
+    dist_op, input_arg_names, output_arg_names, fw_results, bw_results
+):
+    (
+        inferred_input_dims_mappings,
+        inferred_output_dims_mappings,
+    ) = merge_forward_backward_dims_mapping(fw_results, bw_results)
+
+    op_dist_attr = dist_op.dist_attr
+    changed = False
+    if len(input_arg_names) != len(inferred_input_dims_mappings):
+        warnings.warn(
+            f"dims mapping is NOT Match, inferred [{len(inferred_input_dims_mappings)}], original: [{len(input_arg_names)}]; dist op: [{dist_op}]"
+        )
+    if len(output_arg_names) != len(inferred_output_dims_mappings):
+        warnings.warn(
+            f"dims mapping is NOT Match, inferred [{len(inferred_output_dims_mappings)}], original: [{len(output_arg_names)}]; dist op: [{dist_op}]"
+        )
+
+    for i in range(len(input_arg_names)):
+        original_dims_mapping = op_dist_attr.get_input_dims_mapping(
+            input_arg_names[i]
+        )
+        inferred_dims_mapping = inferred_input_dims_mappings[i]
+        if (inferred_dims_mapping is not None) and (
+            original_dims_mapping != inferred_dims_mapping
+        ):
+            _logger.debug(
+                f"Changed: Op [{dist_op.serial_op.type}], name [{input_arg_names[i]}], Original [{original_dims_mapping}], Inferred [{inferred_dims_mapping}]"
+            )
+            changed = True
+            op_dist_attr.set_input_dims_mapping(
+                input_arg_names[i], inferred_dims_mapping
+            )
+        # TODO support partial for inputs
+
+    for i in range(len(output_arg_names)):
+        original_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            output_arg_names[i]
+        )
+        inferred_dims_mapping = inferred_output_dims_mappings[i]
+        if (inferred_dims_mapping is not None) and (
+            original_dims_mapping != inferred_dims_mapping
+        ):
+            _logger.debug(
+                f"Changed: Op [{dist_op.serial_op.type}], name [{output_arg_names[i]}], Original [{original_dims_mapping}], Inferred [{inferred_dims_mapping}]"
+            )
+            changed = True
+            op_dist_attr.set_output_dims_mapping(
+                output_arg_names[i], inferred_dims_mapping
+            )
+
+        # NOTE in partial stage-I, we infer partial for output in infer_forward only
+        output_dist_attr = op_dist_attr.get_output_dist_attr(
+            output_arg_names[i]
+        )
+        output_idx = output_arg_names.index(output_arg_names[i])
+        if (
+            fw_results[1][output_idx]._partial_dims()
+            != output_dist_attr._partial_dims()
+        ):
+            # _logger.info(
+            #     "Changed: Op [{}], tensor name [{}], Original partial on [{}], Inferred partial on [{}]".format(
+            #         dist_op.serial_op.type,
+            #         output_arg_names[i],
+            #         output_dist_attr._partial_dims(),
+            #         fw_results[1][output_idx]._partial_dims(),
+            #     )
+            # )
+            output_dist_attr._clean_partial_status()
+            output_dist_attr._set_partial_dims(
+                list(fw_results[1][0]._partial_dims())
+            )
+            changed = True
+
+    return changed
+
+
+def get_default_distributed_operator_impl():
+    dist_op_default_impl_container = get_distributed_operator_impl_container(
+        "default"
+    )
+    num_impls = len(dist_op_default_impl_container.impls)
+    assert num_impls == 1, f"Default dist op has [{num_impls}] impls"
+    return dist_op_default_impl_container.get_impl(0)
+
+
+def copy_op_without_infer_shape(src_op, block, ctx, varname_kwargs):
+    new_op = block.append_op(type='nop')
+    new_op_desc = new_op.desc
+    new_op_desc.copy_from(src_op.desc)
+    set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
+    for input_name in src_op.desc.input_names():
+        new_op_desc.set_input(input_name, varname_kwargs[input_name])
+    for output_name in src_op.desc.output_names():
+        new_op_desc.set_output(output_name, varname_kwargs[output_name])
+    # TODO: should we add a new dist attr for the new op here?
+    return new_op

python/paddle/distributed/auto_parallel/static/operators/dist_assign.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..utils import compute_compatible_and_update_dim_mapping
+from .common import DistributedOperatorImpl, DistributedOperatorImplContainer
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedAssign(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+# TODO remove assign dist op
+# register_distributed_operator_impl_container(DistributedAssign("assign"))
+
+
+class DistributedAssignImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        if x_dims_mapping != out_dims_mapping:
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        for i in range(len(x_dims_mapping)):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [x_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                changed = True
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+# register_distributed_operator_impl("assign", DistributedAssignImpl("assign"))

python/paddle/distributed/auto_parallel/static/operators/dist_check_finite_and_unscale.py

@@ -0,0 +1,206 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import paddle
+from paddle.distributed.auto_parallel.static.process_group import (
+    get_world_process_group,
+)
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+from paddle.framework import core
+
+from ..dist_attribute import OperatorDistAttr
+from ..process_group import new_process_group
+from ..utils import set_dist_op_desc_original_id, set_var_dist_attr
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    SyncMode,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+
+world_process_group = get_world_process_group()
+
+
+class DistributedCheckFiniteAndUnscale(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(
+    DistributedCheckFiniteAndUnscale("check_finite_and_unscale")
+)
+
+
+class DistributedCheckFiniteAndUnscaleImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._name = name
+        self._forward_implemented = False
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        raise RuntimeError(
+            "DistributedCheckFiniteAndUnscaleImpl's is_input_compatible should not be called !"
+        )
+
+    def is_output_compatible(self, dist_op):
+        raise RuntimeError(
+            "DistributedCheckFiniteAndUnscaleImpl's is_output_compatible should not be called !"
+        )
+
+    def is_auto_compatible(self, dist_op):
+        raise RuntimeError(
+            "DistributedCheckFiniteAndUnscaleImpl's is_auto_compatible should not be called !"
+        )
+
+    def update_dims_mapping(self, dist_op):
+        raise RuntimeError(
+            "DistributedCheckFiniteAndUnscaleImpl's update_dims_mapping should not be called !"
+        )
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        raise RuntimeError(
+            "DistributedCheckFiniteAndUnscaleImpl's forward should not be called !"
+        )
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        # by now the backward function only insert the gradient allreduce for dist op itself
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.main_block
+        backward_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
+        assert dist_attr is not None, (
+            f"backward op [{backward_op}] don't have dist attribute !"
+        )
+
+        assert rank_id in dist_attr.process_mesh.process_ids
+
+        assert 'X' in kwargs, "input [{}] is not given".format('X')
+        assert 'Scale' in kwargs, "input [{}] is not given".format('Scale')
+        assert 'Out' in kwargs, "input [{}] is not given".format('Out')
+        assert 'FoundInfinite' in kwargs, "output [{}] is not given".format(
+            'FoundInfinite'
+        )
+
+        assert len(kwargs['Scale']) == 1, (
+            "check_finite_and_unscale input Scale take 1 variable but got {}".format(
+                kwargs['Scale']
+            )
+        )
+        assert len(kwargs['FoundInfinite']) == 1, (
+            "check_finite_and_unscale input FoundInfinite take 1 variable but got {}".format(
+                kwargs['FoundInfinite']
+            )
+        )
+        assert len(kwargs['X']) == len(kwargs['Out']), (
+            "check_finite_and_unscale got [{}] X and [{}] Out, which are supposed to be equal".format(
+                len(kwargs['X']), len(kwargs['Out'])
+            )
+        )
+
+        filter_vars = []
+        for varname in kwargs['X']:
+            if (
+                rank_id
+                in ctx.get_tensor_dist_attr_for_program(
+                    main_block._var_recursive(varname)
+                ).process_mesh.process_ids
+            ):
+                filter_vars.append(varname)
+
+        # replicate op in dist program
+        dist_op = main_block.append_op(type='nop')
+        dist_op_desc = dist_op.desc
+        dist_op_desc.copy_from(backward_op.desc)
+        set_dist_op_desc_original_id(dist_op_desc, backward_op.desc, ctx)
+        dist_op_desc.set_input('X', filter_vars)
+        dist_op_desc.set_output('Out', filter_vars)
+        # TODO: should we add a new dist attr for the new op here?
+
+        # sync result
+        group = new_process_group(world_process_group.ranks)
+
+        inf_var = main_block._var_recursive(kwargs['FoundInfinite'][0])
+        inf_var_int32 = main_block.create_var(
+            name=inf_var.name + "@cast_int32",
+            shape=inf_var.shape,
+            dtype=core.VarDesc.VarType.INT32,
+        )
+        set_var_dist_attr(
+            ctx,
+            inf_var_int32,
+            ctx.get_tensor_dist_attr_for_program(inf_var).dims_mapping,
+            ctx.get_tensor_dist_attr_for_program(inf_var).process_mesh,
+        )
+        cast_op1 = main_block.append_op(
+            type='cast',
+            inputs={'X': inf_var},
+            outputs={'Out': inf_var_int32},
+            attrs={
+                "in_dtype": inf_var.dtype,
+                "out_dtype": inf_var_int32.dtype,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+        allreduce_op = main_block.append_op(
+            type='all_reduce',
+            inputs={'x': inf_var_int32},
+            outputs={'out': inf_var_int32},
+            attrs={
+                'ring_id': group.id,
+                'op_type': paddle.distributed.ReduceOp.MAX,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+        allreduce_op._set_attr('op_namescope', '/' + SyncMode.AmpFlagSync)
+        cast_op2 = main_block.append_op(
+            type='cast',
+            inputs={'X': inf_var_int32},
+            outputs={'Out': inf_var},
+            attrs={
+                "in_dtype": inf_var_int32.dtype,
+                "out_dtype": inf_var.dtype,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+
+        for op in [cast_op1, allreduce_op, cast_op2]:
+            new_op_dist_attr = OperatorDistAttr()
+            for varname in op.input_arg_names:
+                var_dist_attr = ctx.get_tensor_dist_attr_for_program(
+                    main_block._var_recursive(varname)
+                )
+                assert var_dist_attr is not None
+                new_op_dist_attr.set_input_dims_mapping(
+                    varname, var_dist_attr.dims_mapping
+                )
+            for varname in op.output_arg_names:
+                var_dist_attr = ctx.get_tensor_dist_attr_for_program(
+                    main_block._var_recursive(varname)
+                )
+                new_op_dist_attr.set_output_dims_mapping(
+                    varname, var_dist_attr.dims_mapping
+                )
+            new_op_dist_attr.process_mesh = var_dist_attr.process_mesh
+            ctx.set_op_dist_attr_for_program(op, new_op_dist_attr)
+
+
+register_distributed_operator_impl(
+    "check_finite_and_unscale",
+    DistributedCheckFiniteAndUnscaleImpl("check_finite_and_unscale"),
+)

python/paddle/distributed/auto_parallel/static/operators/dist_concat.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ..completion import get_phi_spmd_rule
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedConcat(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        axis_tensor = op_desc.input('AxisTensor')
+        assert len(axis_tensor) == 0, (
+            "Please use axis attr instead of AxisTensor"
+        )
+
+        input_arg_names = op_desc.input_arg_names()
+        output_arg_names = op_desc.output_arg_names()
+        axis = op_desc.attr('axis')
+
+        input_specs = []
+        for name in input_arg_names:
+            input_specs.append(get_dist_tensor_spec(dist_op, name))
+        output_spec = get_dist_tensor_spec(dist_op, output_arg_names[0], False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("concat")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(input_specs, axis)
+        bw_results = rule.infer_backward(input_specs, output_spec, axis)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            input_arg_names,
+            output_arg_names,
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedConcat("concat"))

python/paddle/distributed/auto_parallel/static/operators/dist_cross_entropy.py

@@ -0,0 +1,528 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import copy
+
+from paddle.common_ops_import import check_variable_and_dtype
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+
+from ..completion import get_phi_spmd_rule
+from ..dist_attribute import OperatorDistAttr
+from ..process_group import new_process_group
+from ..utils import (
+    _get_comm_group,
+    _get_corresponding_rank,
+    get_dist_tensor_spec,
+    is_dim_shard,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    ParallelMode,
+    copy_op_without_infer_shape,
+    naive_copy_op_dist_attr_for_program,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedCrossEntropy(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        logits_name = op_desc.input('Logits')[0]
+        label_name = op_desc.input('Label')[0]
+        loss_name = op_desc.output('Loss')[0]
+        softmax_name = op_desc.output('Softmax')[0]
+
+        soft_label = op_desc.attr('soft_label')
+        ignore_index = op_desc.attr('ignore_index')
+        numeric_stable_mode = op_desc.attr('numeric_stable_mode')
+        axis = op_desc.attr('axis')
+
+        logits_spec = get_dist_tensor_spec(dist_op, logits_name)
+        label_spec = get_dist_tensor_spec(dist_op, label_name)
+        loss_spec = get_dist_tensor_spec(dist_op, loss_name, False)
+        softmax_spec = get_dist_tensor_spec(dist_op, softmax_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("softmax_with_cross_entropy")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(
+            logits_spec,
+            label_spec,
+            soft_label,
+            True,
+            numeric_stable_mode,
+            ignore_index,
+            axis,
+        )
+        bw_results = rule.infer_backward(
+            logits_spec,
+            label_spec,
+            softmax_spec,
+            loss_spec,
+            soft_label,
+            True,
+            numeric_stable_mode,
+            ignore_index,
+            axis,
+        )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            [logits_name, label_name],
+            [softmax_name, loss_name],
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        op_dist_attr.impl_type = op_desc.type()
+
+        logits_name = op_desc.input('Logits')[0]
+
+        soft_label = op_desc.attr('soft_label')
+        axis = op_desc.attr('axis')
+
+        logits_dims_mapping = copy.deepcopy(
+            op_dist_attr.get_input_dims_mapping(logits_name)
+        )
+        logits_ndim = len(logits_dims_mapping)
+        axis = axis + logits_ndim if axis < 0 else axis
+
+        if is_dim_shard(logits_dims_mapping[axis]):
+            assert soft_label is False, (
+                "parallel_cross_entropy does not support soft_label now."
+            )
+            assert axis == logits_ndim - 1, (
+                "parallel_cross_entropy can only support shard on the last dim now."
+            )
+            op_dist_attr.impl_idx = 1
+        else:
+            op_dist_attr.impl_idx = 0
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedCrossEntropy("softmax_with_cross_entropy")
+)
+
+
+# The softmax_norm axis is not sharded
+class DistributedCrossEntropyImpl0(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        return True
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        """
+        kwargs: inputname_mapping & outputname_mapping
+        """
+
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None, (
+            f"forward op [{src_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        assert 'Logits' in kwargs, "input [Logits] is not given"
+        assert 'Label' in kwargs, "input [Label] is not given"
+        assert 'Loss' in kwargs, "output [Loss] is not given"
+        assert 'Softmax' in kwargs, "output [Softmax] is not given"
+
+        assert len(kwargs['Logits']) == 1, (
+            "input [Logits] take 1 variable but got {}".format(kwargs['Logits'])
+        )
+        assert len(kwargs['Label']) == 1, (
+            "input [Label] take 1 variable but got {}".format(kwargs['Label'])
+        )
+
+        logits_var = main_block._var_recursive(kwargs['Logits'][0])
+        label_var = main_block._var_recursive(kwargs['Label'][0])
+        loss_var = main_block._var_recursive(kwargs['Loss'][0])
+        softmax_var = main_block._var_recursive(kwargs['Softmax'][0])
+
+        # got dist attribute info
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
+        if rank_id not in process_mesh_group:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        check_variable_and_dtype(
+            logits_var,
+            'input',
+            ['bfloat16', 'float16', 'float32', 'float64'],
+            'cross_entropy_with_softmax',
+        )
+        check_variable_and_dtype(
+            label_var,
+            'input',
+            ['int32', 'int64', 'float32', 'float64'],
+            'cross_entropy_with_softmax',
+        )
+        check_variable_and_dtype(
+            loss_var,
+            'output',
+            ['bfloat16', 'float16', 'float32', 'float64'],
+            'cross_entropy_with_softmax',
+        )
+        check_variable_and_dtype(
+            softmax_var,
+            'output',
+            ['bfloat16', 'float16', 'float32', 'float64'],
+            'cross_entropy_with_softmax',
+        )
+        copy_op_without_infer_shape(src_op, main_block, ctx, kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        backward_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
+
+        assert op_dist_attr is not None, (
+            f"backward op [{backward_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        assert 'Softmax' in kwargs, "input [Logits] is not given"
+        assert 'Label' in kwargs, "input [Label] is not given"
+        assert 'Loss@GRAD' in kwargs, "input [Loss@GRAD] is not given"
+        assert 'Logits@GRAD' in kwargs, "output [Logits@GRAD] is not given"
+
+        assert len(kwargs['Softmax']) == 1, (
+            "input [Softmax] take 1 variable but got {}".format(
+                kwargs['Softmax']
+            )
+        )
+        assert len(kwargs['Label']) == 1, (
+            "input [Label] take 1 variable but got {}".format(kwargs['Label'])
+        )
+        assert len(kwargs['Loss@GRAD']) == 1, (
+            "input [Loss@GRAD] take 1 variable but got {}".format(kwargs['Out'])
+        )
+        assert len(kwargs['Logits@GRAD']) == 1, (
+            "output [Logits@GRAD] take 1 variable but got {}".format(
+                kwargs['Logits@GRAD']
+            )
+        )
+        # replicate op in dist program
+        copy_op_without_infer_shape(backward_op, main_block, ctx, kwargs)
+
+
+class DistributedCrossEntropyImpl1(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        return True
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        """
+        kwargs: inputname_mapping & outputname_mapping
+        """
+
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None, (
+            f"forward op [{src_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        assert 'Logits' in kwargs, "input [Logits] is not given"
+        assert 'Label' in kwargs, "input [Label] is not given"
+        assert 'Loss' in kwargs, "output [Loss] is not given"
+        assert 'Softmax' in kwargs, "output [Softmax] is not given"
+
+        assert len(kwargs['Logits']) == 1, (
+            "input [Logits] take 1 variable but got {}".format(kwargs['Logits'])
+        )
+        assert len(kwargs['Label']) == 1, (
+            "input [Label] take 1 variable but got {}".format(kwargs['Label'])
+        )
+
+        logits_var = main_block._var_recursive(kwargs['Logits'][0])
+        label_var = main_block._var_recursive(kwargs['Label'][0])
+        loss_var = main_block._var_recursive(kwargs['Loss'][0])
+        softmax_var = main_block._var_recursive(kwargs['Softmax'][0])
+
+        # got dist attribute info
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
+        if rank_id not in process_mesh_group:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        check_variable_and_dtype(
+            logits_var,
+            'input',
+            ['bfloat16', 'float16', 'float32', 'float64'],
+            'c_softmax_with_cross_entropy',
+        )
+        check_variable_and_dtype(
+            label_var,
+            'input',
+            ['int32', 'int64', 'float32', 'float64'],
+            'c_softmax_with_cross_entropy',
+        )
+        check_variable_and_dtype(
+            loss_var,
+            'output',
+            ['bfloat16', 'float16', 'float32', 'float64'],
+            'c_softmax_with_cross_entropy',
+        )
+        check_variable_and_dtype(
+            softmax_var,
+            'output',
+            ['bfloat16', 'float16', 'float32', 'float64'],
+            'c_softmax_with_cross_entropy',
+        )
+
+        # infer new var shape with op dist attr
+        # the dims mapping in dist_op may be different from that in tensor
+        # so we should
+        loss_dist_attr = ctx.get_tensor_dist_attr_for_program(loss_var)
+        assert loss_dist_attr is not None
+        softmax_dist_attr = ctx.get_tensor_dist_attr_for_program(softmax_var)
+        assert softmax_dist_attr is not None
+        op_dist_attr_loss = op_dist_attr.get_output_dist_attr(loss_var.name)
+        assert op_dist_attr_loss is not None
+        op_dist_attr_softmax = op_dist_attr.get_output_dist_attr(
+            softmax_var.name
+        )
+        assert op_dist_attr_softmax is not None
+
+        # TODO calculate ring id
+        softmax_axis = src_op.desc.attr('axis')
+        logits_dims_mapping = op_dist_attr.get_input_dims_mapping(
+            logits_var.name
+        )
+        parallel_axis = logits_dims_mapping[softmax_axis]
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        c_cross_entropy_op = main_block.append_op(
+            type='c_softmax_with_cross_entropy',
+            inputs={
+                'Logits': logits_var,
+                'Label': label_var,
+            },
+            outputs={
+                'Loss': loss_var,
+                'Softmax': softmax_var,
+            },
+            attrs={
+                'ring_id': group.id,
+                'rank': group.local_rank(rank_id),
+                'nranks': group.nranks,
+                'ignore_index': src_op.desc.attr('ignore_index'),
+                OP_ROLE_KEY: src_op.attr('op_role'),
+            },
+        )
+        naive_copy_op_dist_attr_for_program(c_cross_entropy_op, src_op, ctx)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        backward_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
+
+        assert op_dist_attr is not None, (
+            f"backward op [{backward_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        assert 'Softmax' in kwargs, "input [Softmax] is not given"
+        assert 'Label' in kwargs, "input [Label] is not given"
+        assert 'Loss@GRAD' in kwargs, "input [Loss@GRAD] is not given"
+        assert 'Logits@GRAD' in kwargs, "output [Logits@GRAD] is not given"
+
+        assert len(kwargs['Softmax']) == 1, (
+            "input [Softmax] take 1 variable but got {}".format(
+                kwargs['Softmax']
+            )
+        )
+        assert len(kwargs['Label']) == 1, (
+            "input [Label] take 1 variable but got {}".format(kwargs['Label'])
+        )
+        assert len(kwargs['Loss@GRAD']) == 1, (
+            "input [Loss@GRAD] take 1 variable but got {}".format(
+                kwargs['Loss@GRAD']
+            )
+        )
+        assert len(kwargs['Logits@GRAD']) == 1, (
+            "output [Logits@GRAD] take 1 variable but got {}".format(
+                kwargs['Logits@GRAD']
+            )
+        )
+
+        # got dist attribute info
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
+        if rank_id not in process_mesh_group:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        for op in main_block.ops:
+            # the output value of reduce_mean_grad is 1/numel, so when the
+            # tensor is sharded, we should insert a scale op to make the
+            # grad correct.
+            if (
+                op.type == "reduce_mean_grad"
+                and kwargs['Loss@GRAD'][0] in op.output_arg_names
+            ):
+                loss_grad_var = main_block._var_recursive(
+                    kwargs['Loss@GRAD'][0]
+                )
+                loss_grad_dims_mapping = op_dist_attr.get_input_dims_mapping(
+                    loss_grad_var.name
+                )
+                degree = 1.0
+                for i in range(len(loss_grad_dims_mapping) - 1):
+                    if loss_grad_dims_mapping[i] != -1:
+                        degree *= process_mesh_shape[loss_grad_dims_mapping[i]]
+                if degree > 1:
+                    scale_op = main_block.append_op(
+                        type='scale',
+                        inputs={'X': loss_grad_var},
+                        outputs={'Out': loss_grad_var},
+                        attrs={
+                            'scale': 1.0 / degree,
+                            OP_ROLE_KEY: OpRole.Backward,
+                        },
+                    )
+                    scale_op._set_attr(
+                        'op_namescope', '/' + ParallelMode.DataParallel
+                    )
+                    dims_mapping = op_dist_attr.get_input_dims_mapping(
+                        loss_grad_var.name
+                    )
+                    scale_op_attr = OperatorDistAttr()
+                    scale_op_attr.process_mesh = op_dist_attr.process_mesh
+                    scale_op_attr.chunk_id = op_dist_attr.chunk_id
+                    scale_op_attr.set_output_dims_mapping(
+                        loss_grad_var.name, dims_mapping
+                    )
+                    scale_op_attr.set_input_dims_mapping(
+                        loss_grad_var.name, dims_mapping
+                    )
+                    ctx.set_op_dist_attr_for_program(scale_op, scale_op_attr)
+
+        # TODO calculate ring id
+        softmax_axis = backward_op.desc.attr('axis')
+        # softmax_dims_mapping is the same as logits_dims_mapping
+        softmax_dims_mapping = op_dist_attr.get_input_dims_mapping(
+            kwargs['Softmax'][0]
+        )
+        parallel_axis = softmax_dims_mapping[softmax_axis]
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        cross_entropy_grad_op_desc = main_block.append_op(type='nop').desc
+        cross_entropy_grad_op_desc.set_type("c_softmax_with_cross_entropy_grad")
+        cross_entropy_grad_op_desc.set_input('Softmax', [kwargs['Softmax'][0]])
+        cross_entropy_grad_op_desc.set_input('Label', [kwargs['Label'][0]])
+        cross_entropy_grad_op_desc.set_input(
+            'Loss@GRAD', [kwargs['Loss@GRAD'][0]]
+        )
+        cross_entropy_grad_op_desc.set_output(
+            'Logits@GRAD', [kwargs['Logits@GRAD'][0]]
+        )
+
+        ignore_index = backward_op.desc.attr('ignore_index')
+        # the ignore_index attribute in c_cross_entropy_grad kernel
+        # is int64_t type, so we should set this attribute with
+        # _set_int64_attr. Otherwise ignore_index will be int32 type,
+        # causing an error.
+        cross_entropy_grad_op_desc._set_int64_attr('ignore_index', ignore_index)
+        cross_entropy_grad_op_desc._set_attr('ring_id', group.id)
+        cross_entropy_grad_op_desc._set_attr('rank', group.local_rank(rank_id))
+        cross_entropy_grad_op_desc._set_attr('nranks', group.nranks)
+        cross_entropy_grad_op_desc._set_attr(OP_ROLE_KEY, OpRole.Backward)
+
+        cross_entropy_grad_op = main_block.ops[-1]
+        naive_copy_op_dist_attr_for_program(
+            cross_entropy_grad_op, backward_op, ctx
+        )
+
+
+register_distributed_operator_impl(
+    "softmax_with_cross_entropy", DistributedCrossEntropyImpl0("cross_entropy")
+)
+register_distributed_operator_impl(
+    "softmax_with_cross_entropy",
+    DistributedCrossEntropyImpl1("c_cross_entropy"),
+)

python/paddle/distributed/auto_parallel/static/operators/dist_default.py

@@ -0,0 +1,681 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import paddle
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+
+from ..completion import contains_spmd_rule, get_phi_spmd_rule
+from ..cost import (
+    _g_op_cost_factory,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_dp_costs,
+)
+from ..dist_attribute import DistTensorSpec, OperatorDistAttr
+from ..process_group import new_process_group
+from ..utils import (
+    _get_comm_group,
+    _get_corresponding_rank,
+    compute_compatible_dim_mapping,
+    get_dist_tensor_spec,
+    is_prim_op,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    copy_op_without_infer_shape,
+    get_default_distributed_operator_impl,
+    gradient_synchronization,
+    is_parameter_related,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    set_comm_op_dist_attr_for_program,
+    update_op_dims_mapping,
+)
+
+__op_not_need_param_init__ = ["while", "cond"]
+__op_has_shape_attr__ = [
+    "fill_constant_batch_size_like",
+    "fill_constant",
+    "expand_v2",
+    "expand_as_v2",
+]
+
+
+def prim_operator_data_parallel_functor(ctx, src_op):
+    dist_op_context = ctx.dist_op_context
+    main_block = dist_op_context.work_block
+    startup_block = dist_op_context.startup_block
+
+    var_name = src_op.output_arg_names[0]
+    if var_name in ctx.grads_params:
+        assert var_name not in ctx.synced_gradient, (
+            f"in primitive mode, grad is already {var_name} synced"
+        )
+        ctx.synced_gradient.add(var_name)
+        sync_group = new_process_group(ctx.data_parallel_group)
+
+        allreduce_op = main_block.append_op(
+            type='all_reduce',
+            inputs={'x': [var_name]},
+            outputs={'out': [var_name]},
+            attrs={
+                'ring_id': sync_group.id,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
+                OP_ROLE_KEY: OpRole.Backward,
+            },
+        )
+
+        param = ctx.grads_params[var_name]
+        startup_block = dist_op_context.startup_block
+        new_op = startup_block.append_op(
+            type='broadcast',
+            inputs={'x': [param]},
+            outputs={'out': [param]},
+            attrs={
+                'ring_id': sync_group.id,
+                'root': 0,
+                OP_ROLE_KEY: OpRole.Forward,
+            },
+        )
+
+        grad_var = main_block._var_recursive(var_name)
+        dims_mapping = ctx.get_tensor_dist_attr_for_program(
+            grad_var
+        ).dims_mapping
+        dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        process_mesh = dist_attr.process_mesh
+        op_attr = OperatorDistAttr()
+        op_attr.process_mesh = process_mesh
+        op_attr.set_output_dims_mapping(grad_var.name, dims_mapping)
+        op_attr.set_input_dims_mapping(grad_var.name, dims_mapping)
+        ctx.set_op_dist_attr_for_program(allreduce_op, op_attr)
+
+
+class DistributedDefault(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+
+        op_desc = dist_op.serial_op.desc
+        input_arg_names = op_desc.input_arg_names()
+        output_arg_names = op_desc.output_arg_names()
+        main_block = dist_op.serial_op.block
+
+        num_inputs = len(input_arg_names)
+        input_specs = []
+        for i in range(num_inputs):
+            assert not is_parameter_related(input_arg_names[i], main_block), (
+                f"input {input_arg_names[i]} of op {dist_op.serial_op} is parameter, op should not use default rule."
+            )
+            input_specs.append(
+                get_dist_tensor_spec(dist_op, input_arg_names[i])
+            )
+        num_outputs = len(output_arg_names)
+        output_specs = []
+        for i in range(num_outputs):
+            assert not is_parameter_related(output_arg_names[i], main_block), (
+                f"output {output_arg_names[i]} of op {dist_op.serial_op} is parameter, op should not use default rule."
+            )
+            output_specs.append(
+                get_dist_tensor_spec(dist_op, output_arg_names[i], False)
+            )
+
+        # step2: infer spmd
+        if contains_spmd_rule(dist_op.serial_op.type):
+            # when some inputs are optional, the input_arg_names will be less than input_names
+            # and we can pass empty DistTensorSpec() as argument
+            if len(op_desc.input_names()) > len(op_desc.input_arg_names()):
+                for i in range(
+                    len(op_desc.input_names()) - len(op_desc.input_arg_names())
+                ):
+                    input_specs.append(DistTensorSpec())
+            rule = get_phi_spmd_rule(dist_op.serial_op.type)
+            fw_results = rule.infer_forward(*input_specs)
+            bw_results = rule.infer_backward(*input_specs, output_specs)
+        else:
+            rule = get_phi_spmd_rule('default_')
+            # tensor order following order in PHI definition
+            fw_results = rule.infer_forward(input_specs, output_specs)
+            bw_results = rule.infer_backward(input_specs, output_specs)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, input_arg_names, output_arg_names, fw_results, bw_results
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        # all op use default dist operator impl.
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedDefault("default"))
+
+
+# Replicated Default
+class DistributedDefaultImpl0(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        """Calculate the cost by the op role."""
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_op.dist_attr.process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            _g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
+        )
+        res_cost = [cost_mapping]
+
+        return res_cost
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        backward_op = dist_op.serial_op
+        op_type = backward_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            _g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and not is_parameter_related(
+                    varname, main_block
+                ):
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        need_gradient_allreduce = True
+                        break
+
+        if need_gradient_allreduce:
+            for input_name in backward_op.desc.input_names():
+                for varname in backward_op.desc.input(input_name):
+                    if "@GRAD" not in varname and is_parameter_related(
+                        varname, main_block
+                    ):
+                        var_dim_mapping = dist_attr.get_input_dims_mapping(
+                            varname
+                        )
+                        mesh_shape = process_mesh.shape
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+        return res
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        batch_dim_mappings = []
+        input_names = op_desc.input_names()
+        xshape_arg_names = []
+        if "XShape" in input_names:
+            xshape_arg_names = op_desc.input("XShape")
+        for arg_name in op_desc.input_arg_names():
+            serial_tensor = dist_op.get_serial_input(arg_name)
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if serial_tensor.is_parameter:
+                for mapping in dims_mapping:
+                    if mapping != -1:
+                        return False
+                continue
+            if arg_name not in xshape_arg_names:
+                if len(dims_mapping) > 1:
+                    for mapping in dims_mapping[1:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 1:
+                    batch_dim_mappings.append(dims_mapping[0])
+            else:
+                if dims_mapping[0] != -1:
+                    return False
+                if len(dims_mapping) > 2:
+                    for mapping in dims_mapping[2:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 2:
+                    batch_dim_mappings.append(dims_mapping[1])
+
+        if compute_compatible_dim_mapping(batch_dim_mappings) is None:
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        output_names = op_desc.output_names()
+        batch_dim_mappings = []
+        xshape_arg_names = []
+        if "XShape" in output_names:
+            xshape_arg_names = op_desc.output("XShape")
+        for arg_name in op_desc.output_arg_names():
+            serial_tensor = dist_op.get_serial_output(arg_name)
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if serial_tensor.is_parameter:
+                for mapping in dims_mapping:
+                    if mapping != -1:
+                        return False
+                continue
+            if arg_name not in xshape_arg_names:
+                if len(dims_mapping) > 1:
+                    for mapping in dims_mapping[1:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 1:
+                    batch_dim_mappings.append(dims_mapping[0])
+            else:
+                if dims_mapping[0] != -1:
+                    return False
+                if len(dims_mapping) > 2:
+                    for mapping in dims_mapping[2:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 2:
+                    batch_dim_mappings.append(dims_mapping[1])
+
+        if compute_compatible_dim_mapping(batch_dim_mappings) is None:
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        batch_dim_mappings = []
+        # Check input compatibility
+        input_names = op_desc.input_names()
+        xshape_arg_names = []
+        if "XShape" in input_names:
+            xshape_arg_names = op_desc.input("XShape")
+        for arg_name in op_desc.input_arg_names():
+            serial_tensor = dist_op.get_serial_input(arg_name)
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if serial_tensor is not None and serial_tensor.is_parameter:
+                for mapping in dims_mapping:
+                    if mapping != -1:
+                        return False
+                continue
+            if arg_name not in xshape_arg_names:
+                if len(dims_mapping) > 1:
+                    for mapping in dims_mapping[1:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 1:
+                    batch_dim_mappings.append(dims_mapping[0])
+            else:
+                if dims_mapping[0] != -1:
+                    return False
+                if len(dims_mapping) > 2:
+                    for mapping in dims_mapping[2:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 2:
+                    batch_dim_mappings.append(dims_mapping[1])
+
+        # Check output compatibility
+        output_names = op_desc.output_names()
+        xshape_arg_names = []
+        if "XShape" in output_names:
+            xshape_arg_names = op_desc.output("XShape")
+        for arg_name in op_desc.output_arg_names():
+            serial_tensor = dist_op.get_serial_output(arg_name)
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if serial_tensor is not None and serial_tensor.is_parameter:
+                for mapping in dims_mapping:
+                    if mapping != -1:
+                        return False
+                continue
+            if arg_name not in xshape_arg_names:
+                if len(dims_mapping) > 1:
+                    for mapping in dims_mapping[1:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 1:
+                    batch_dim_mappings.append(dims_mapping[0])
+            else:
+                if dims_mapping[0] != -1:
+                    return False
+                if len(dims_mapping) > 2:
+                    for mapping in dims_mapping[2:]:
+                        if mapping != -1:
+                            return False
+                if len(dims_mapping) >= 2:
+                    batch_dim_mappings.append(dims_mapping[1])
+
+        # Check batch dim mapping compatibility
+        if not all(
+            batch_dim_mappings[0] == dim_mapping
+            for dim_mapping in batch_dim_mappings
+        ):
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+
+        if op_desc.type() == "while":
+            return False
+
+        input_names = op_desc.input_names()
+        input_xshape_arg_names = []
+        if "XShape" in input_names:
+            input_xshape_arg_names = op_desc.input("XShape")
+
+        output_names = op_desc.output_names()
+        output_xshape_arg_names = []
+        if "XShape" in output_names:
+            output_xshape_arg_names = op_desc.output("XShape")
+
+        batch_dim_mappings = []
+        for arg_name in op_desc.input_arg_names():
+            serial_tensor = dist_op.get_serial_input(arg_name)
+            if serial_tensor.is_parameter:
+                continue
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if arg_name not in input_xshape_arg_names:
+                if len(dims_mapping) >= 1:
+                    batch_dim_mappings.append(dims_mapping[0])
+            else:
+                batch_dim_mappings.append(dims_mapping[1])
+        for arg_name in op_desc.output_arg_names():
+            if op_desc.type() == 'fill_any_like':
+                input_tensor = dist_op.get_serial_input(
+                    op_desc.input_arg_names()[0]
+                )
+                if input_tensor.is_parameter:
+                    continue
+            serial_tensor = dist_op.get_serial_output(arg_name)
+            if serial_tensor.is_parameter:
+                continue
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if arg_name not in output_xshape_arg_names:
+                if len(dims_mapping) >= 1:
+                    batch_dim_mappings.append(dims_mapping[0])
+            else:
+                batch_dim_mappings.append(dims_mapping[1])
+
+        if not batch_dim_mappings:
+            return changed
+
+        compatible_dim_mapping = compute_compatible_dim_mapping(
+            batch_dim_mappings
+        )
+        if compatible_dim_mapping is None:
+            return False
+
+        for arg_name in op_desc.input_arg_names():
+            serial_tensor = dist_op.get_serial_input(arg_name)
+            if serial_tensor.is_parameter:
+                continue
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if arg_name not in input_xshape_arg_names:
+                if (
+                    len(dims_mapping) >= 1
+                    and compatible_dim_mapping != dims_mapping[0]
+                ):
+                    dims_mapping[0] = compatible_dim_mapping
+                    op_dist_attr.set_input_dims_mapping(arg_name, dims_mapping)
+                    changed = True
+            else:
+                if (
+                    len(dims_mapping) >= 2
+                    and compatible_dim_mapping != dims_mapping[1]
+                ):
+                    dims_mapping[1] = compatible_dim_mapping
+                    op_dist_attr.set_input_dims_mapping(arg_name, dims_mapping)
+                    changed = True
+        for arg_name in op_desc.output_arg_names():
+            if op_desc.type() == 'fill_any_like':
+                input_tensor = dist_op.get_serial_input(
+                    op_desc.input_arg_names()[0]
+                )
+                if input_tensor.is_parameter:
+                    continue
+            if op_desc.type() in ["shape", "slice"]:
+                continue
+            serial_tensor = dist_op.get_serial_output(arg_name)
+            if serial_tensor.is_parameter:
+                continue
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if arg_name not in output_xshape_arg_names:
+                if (
+                    len(dims_mapping) >= 1
+                    and compatible_dim_mapping != dims_mapping[0]
+                ):
+                    dims_mapping[0] = compatible_dim_mapping
+                    op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
+                    changed = True
+            else:
+                if (
+                    len(dims_mapping) >= 2
+                    and compatible_dim_mapping != dims_mapping[1]
+                ):
+                    dims_mapping[1] = compatible_dim_mapping
+                    op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
+                    changed = True
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+
+        # check validation of inputs / outputs
+        for input_name in src_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                src_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in src_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                src_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        # replicate op in dist program
+        dst_op = copy_op_without_infer_shape(src_op, main_block, ctx, kwargs)
+
+        def get_shape_attr_name():
+            for name in ["shape", "target_shape"]:
+                if src_op.has_attr(name) and src_op.attr(name):
+                    return name
+            return None
+
+        shape_attr_name = get_shape_attr_name()
+        if shape_attr_name and src_op.type in __op_has_shape_attr__:
+            shape_list = src_op.attr(shape_attr_name)
+            Out_var = main_block._var_recursive(kwargs['Out'][0])
+            op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+            dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
+            process_mesh_shape = op_dist_attr.process_mesh.shape
+            assert len(shape_list) == len(dim_mapping)
+            # modify target shape
+            for idx, axis in enumerate(dim_mapping):
+                if axis >= 0:
+                    if len(shape_list) > idx:
+                        shape_list[idx] = (
+                            shape_list[idx] // process_mesh_shape[axis]
+                        )
+            dst_op.desc._set_attr(shape_attr_name, shape_list)
+
+        # data parallel synchronization for primitive operators
+        from paddle.incubate.autograd import prim_enabled
+
+        if prim_enabled():
+            assert is_prim_op(src_op)
+            prim_operator_data_parallel_functor(ctx, src_op)
+            return
+
+        # param initialization sync
+        if src_op.type in __op_not_need_param_init__:
+            return
+
+        for varname in dst_op.desc.input_arg_names():
+            if (
+                startup_block.has_var(varname)
+                and startup_block.var(varname).is_parameter
+                and varname not in dist_op_context.already_init_sync_vars
+            ):
+                dist_op_context.already_init_sync_vars.add(varname)
+                param = startup_block.var(varname)
+                param_dist_attr = ctx.get_tensor_dist_attr_for_program(param)
+                process_mesh = param_dist_attr.process_mesh
+                dims_mapping = param_dist_attr.dims_mapping
+
+                # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
+                if rank_id not in process_mesh.process_ids:
+                    rank_id = _get_corresponding_rank(
+                        ctx, process_mesh, rank_id
+                    )
+
+                # NOTE all not splited axis should be presented in mesh
+                for axis, size in enumerate(process_mesh.shape):
+                    if size <= 1 or axis in dims_mapping:
+                        pass
+                    else:
+                        group_ranks = _get_comm_group(
+                            process_mesh.process_ids,
+                            process_mesh.shape,
+                            axis,
+                            rank_id,
+                        )
+                        sync_group = new_process_group(group_ranks)
+
+                        new_op = startup_block.append_op(
+                            type='broadcast',
+                            inputs={'x': param},
+                            outputs={'out': param},
+                            attrs={
+                                'ring_id': sync_group.id,
+                                'root': 0,
+                                OP_ROLE_KEY: OpRole.Forward,
+                            },
+                        )
+                        set_comm_op_dist_attr_for_program(
+                            new_op,
+                            process_mesh,
+                            param_dist_attr,
+                            ctx,
+                        )
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        # by now the backward function only insert the gradient allreduce for dist op itself
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        backward_op = dist_op_context.cur_src_op
+        dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
+        assert dist_attr is not None, (
+            f"backward op [{backward_op}] don't have dist attribute !"
+        )
+        rank_id = dist_op_context.rank_id
+
+        # check validation of inputs / outputs
+        for input_name in backward_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                backward_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in backward_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                backward_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        # replicate op in dist program
+        copy_op_without_infer_shape(backward_op, main_block, ctx, kwargs)
+
+        # data parallel gradient synchronization
+        act_grad_names = []
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and not is_parameter_related(
+                    varname, main_block
+                ):
+                    act_grad_names.append(varname)
+
+        out_grad_names = []
+        for output_name in backward_op.desc.output_names():
+            for varname in backward_op.desc.output(output_name):
+                if varname in kwargs["grad_var_to_var"]:
+                    fwd_name = kwargs["grad_var_to_var"][varname]
+                    if not main_block._find_var_recursive(fwd_name):
+                        continue
+                    if is_parameter_related(fwd_name, main_block):
+                        out_grad_names.append(varname)
+
+        gradient_synchronization(
+            ctx, backward_op, act_grad_names, out_grad_names, rank_id
+        )
+
+
+register_distributed_operator_impl(
+    "default", DistributedDefaultImpl0("replicate_parallel")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_dropout.py

@@ -0,0 +1,238 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+import logging
+
+import paddle
+from paddle.base.log_helper import get_logger
+from paddle.framework import core
+from paddle.utils import unique_name
+
+from ...random import determinate_rng, is_enable_auto_rand_ctrl
+from ..completion import get_phi_spmd_rule
+from ..utils import (
+    get_dist_tensor_spec,
+    naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
+    set_var_dist_attr,
+)
+from .common import (
+    DistributedOperatorImplContainer,
+    merge_forward_backward_dims_mapping,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+from .dist_eltwise import DistributedDefaultImpl0, DistributedElementwiseImpl0
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class DistributedDropout(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        mask_name = op_desc.output('Mask')[0]
+        # seed_name = op_desc.input('Seed')[0]  // seed is a scalar and leave it to be unsharded
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        output_spec = get_dist_tensor_spec(dist_op, out_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("dropout")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec)
+        bw_results = rule.infer_backward(x_spec, output_spec)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, [x_name], [out_name], fw_results, bw_results
+        )
+
+        # step5: update mask and seed dropout special
+        if changed:
+            (
+                _,
+                inferred_output_dims_mappings,
+            ) = merge_forward_backward_dims_mapping(fw_results, bw_results)
+            dist_op.dist_attr.set_output_dims_mapping(
+                mask_name, inferred_output_dims_mappings[0]
+            )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        # all dropout op use Dropout with Random Control dist operator impl.
+        op_dist_attr = dist_op.dist_attr
+        op_dist_attr.impl_type = "dropout"
+        op_dist_attr.impl_idx = 0
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedDropout("dropout"))
+
+
+# Dist Dropout with Random Control
+# Dropout re-use the compatible and cost function of elementwise
+class DistributedDropoutImpl0(DistributedElementwiseImpl0):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None, (
+            f"forward op [{src_op}] don't have dist attribute !"
+        )
+
+        if is_enable_auto_rand_ctrl() and not op_dist_attr.is_recompute:
+            # check validation of inputs / outputs
+            assert 'X' in kwargs, "input [{}] is not given".format('X')
+            assert len(kwargs['X']) == 1, (
+                "input X should be only one tensor but got {}".format(
+                    kwargs['X']
+                )
+            )
+            assert 'Seed' in kwargs, "input [{}] is not given".format('Seed')
+
+            if (
+                src_op.has_attr("fix_seed")
+                and src_op.attr("fix_seed")
+                and src_op.has_attr("seed")
+                and src_op.attr("seed")
+            ):
+                _logger.info(
+                    f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
+                )
+            elif rank_id not in op_dist_attr.process_mesh.process_ids:
+                pass
+            # NOTE Adopt for recompute
+            # If user already set seed, We should not modify it. But if the seed is added by recompute pass, it should be under control.
+            # TODO  in future recompute pass should happen after parallel partition. and remove this at that time.
+            elif len(kwargs['Seed']) > 0 or len(src_op.input("Seed")) > 0:
+                seed_var_name = kwargs['Seed'][0]
+                if seed_var_name.startswith('rc_seed'):
+                    pre_op = main_block.ops[-1]
+                    assert (
+                        pre_op.type == "seed"
+                        and len(pre_op.attr("rng_name")) == 0
+                    ), f"found exception op {pre_op}"
+
+                    # determinate rng
+                    X_var = main_block._var_recursive(kwargs['X'][0])
+                    X_dims_mapping = op_dist_attr.get_input_dims_mapping(
+                        X_var.name
+                    )
+                    process_mesh = op_dist_attr.process_mesh
+                    rng_name = determinate_rng(
+                        rank_id, X_dims_mapping, process_mesh
+                    )
+                    # make recompute seed under control
+                    pre_op._set_attr("rng_name", rng_name)
+                    pre_op._set_attr("deterministic", True)
+                    pre_op._set_attr("force_cpu", True)
+                else:
+                    _logger.info(
+                        f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
+                    )
+            else:
+                # determinate rng
+                X_var = main_block._var_recursive(kwargs['X'][0])
+                X_dims_mapping = op_dist_attr.get_input_dims_mapping(X_var.name)
+                process_mesh = op_dist_attr.process_mesh
+
+                rng_name = determinate_rng(
+                    rank_id, X_dims_mapping, process_mesh
+                )
+                assert rng_name is not None and rng_name != ""
+
+                # insert seed op
+                seed_var = main_block.create_var(
+                    name=unique_name.generate_with_ignorable_key(
+                        ".".join(["tensor_parallel_seed", 'tmp'])
+                    ),
+                    dtype=paddle.int32,
+                    type=core.VarDesc.VarType.DENSE_TENSOR,
+                    persistable=False,
+                    stop_gradient=False,
+                )
+
+                # set new seed_var's dist_attr
+                seed_var_dims_mapping = [-1]
+                seed_var_dist_attr = set_var_dist_attr(
+                    ctx,
+                    seed_var,
+                    seed_var_dims_mapping,
+                    process_mesh,
+                    chunk_id=op_dist_attr.chunk_id,
+                )
+
+                # adopt for recompute
+                # force_cpu to reduce sync copy from CPU->GPU->CPU, and reduce pipeline hang
+                seed_op = main_block.append_op(
+                    type='seed',
+                    outputs={'Out': seed_var},
+                    attrs={
+                        'deterministic': True,
+                        'rng_name': rng_name,
+                        'force_cpu': True,
+                    },
+                )
+                seed_op._set_attr('op_namescope', 'auto_tensor_parallel_seed')
+                # set new seed op's dist_attr
+                naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
+                    seed_op,
+                    process_mesh,
+                    seed_var_dims_mapping,
+                    ctx,
+                    chunk_id=op_dist_attr.chunk_id,
+                )
+
+                # modify dropout op
+                src_op.desc.set_input("Seed", [seed_var.name])
+                src_op.desc._set_attr("fix_seed", False)
+                src_op.desc._set_attr("seed", 0)
+                op_dist_attr.set_input_dist_attr(
+                    seed_var.name, seed_var_dist_attr
+                )
+                kwargs['Seed'] = [seed_var.name]
+
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        # dropout backward is deterministic by mask, and not need for random state control
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "dropout", DistributedDropoutImpl0("random_control")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_eltwise.py

@@ -0,0 +1,400 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..completion import get_phi_spmd_rule
+from ..cost import (
+    _g_op_cost_factory,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_dp_costs,
+)
+from ..utils import (
+    compute_compatible_dim_mapping,
+    compute_compatible_dims_mapping,
+    get_dist_tensor_spec,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    is_elementwise_op,
+    is_parameter_related,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedElementwise(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+        assert len(op_desc.input_arg_names()) >= 1, (
+            f"elementwise op [{op_desc.type}] has [{len(op_desc.input_arg_names())}] inputs"
+        )
+        input_arg_names = op_desc.input_arg_names()
+        assert len(op_desc.output_arg_names()) == 1, (
+            f"elementwise op [{dist_op.serial_op}] has [{len(op_desc.output_arg_names())}] outputs"
+        )
+        output_arg_name = op_desc.output_arg_names()[0]
+        num_inputs = len(input_arg_names)
+
+        # TODO (zhangyichen) replace dist tensor specs by dist tensor in future.
+        input_specs = []
+        for i in range(num_inputs):
+            input_specs.append(
+                get_dist_tensor_spec(dist_op, input_arg_names[i])
+            )
+        output_spec = get_dist_tensor_spec(dist_op, output_arg_name, False)
+
+        # step2: infer spmd
+        # TODO revise me
+        op_type = op_desc.type()
+        rule = get_phi_spmd_rule(op_type)
+        fw_results = rule.infer_forward(*input_specs)
+        bw_results = rule.infer_backward(*input_specs, output_spec)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, input_arg_names, [output_arg_name], fw_results, bw_results
+        )
+
+        return changed
+
+    # NOTE this function will be remove once we use local reshard to replace distopimpls
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        # all elementwise op use default dist operator impl.
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedElementwise("elementwise")
+)
+
+
+# Replicated Elementwise
+class DistributedElementwiseImpl0(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = False
+        self._backward_implemented = False
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        """Calculate the cost by the op role."""
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_op.dist_attr.process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            _g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
+        )
+        res_cost = [cost_mapping]
+
+        return res_cost
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        backward_op = dist_op.serial_op
+        op_type = backward_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            _g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and not is_parameter_related(
+                    varname, main_block
+                ):
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        need_gradient_allreduce = True
+                        break
+
+        if need_gradient_allreduce:
+            for input_name in backward_op.desc.input_names():
+                for varname in backward_op.desc.input(input_name):
+                    if "@GRAD" not in varname and is_parameter_related(
+                        varname, main_block
+                    ):
+                        var_dim_mapping = dist_attr.get_input_dims_mapping(
+                            varname
+                        )
+                        mesh_shape = process_mesh.shape
+                        batch_size_axis = var_dim_mapping[0]
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+        return res
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        if not is_elementwise_op(op_desc.type()):
+            return False
+        op_dist_attr = dist_op.dist_attr
+        dims_mapping_list = []
+        input_arg_names = op_desc.input_arg_names()
+        max_dims_mapping_len = -1
+        for arg_name in input_arg_names:
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if max_dims_mapping_len < len(dims_mapping):
+                max_dims_mapping_len = len(dims_mapping)
+            dims_mapping_list.append(dims_mapping)
+
+        for idx in range(max_dims_mapping_len):
+            dim_mappings = []
+            for dims_mapping in dims_mapping_list:
+                if idx < len(dims_mapping):
+                    dim_mappings.append(dims_mapping[-(idx + 1)])
+            if compute_compatible_dim_mapping(dim_mappings) is None:
+                return False
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        if not is_elementwise_op(op_desc.type()):
+            return False
+        op_dist_attr = dist_op.dist_attr
+        dims_mapping_list = []
+        output_arg_names = op_desc.output_arg_names()
+        max_dims_mapping_len = -1
+        for arg_name in output_arg_names:
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if max_dims_mapping_len < len(dims_mapping):
+                max_dims_mapping_len = len(dims_mapping)
+            dims_mapping_list.append(dims_mapping)
+
+        for idx in range(max_dims_mapping_len):
+            dim_mappings = []
+            for dims_mapping in dims_mapping_list:
+                if idx < len(dims_mapping):
+                    dim_mappings.append(dims_mapping[-(idx + 1)])
+            if compute_compatible_dim_mapping(dim_mappings) is None:
+                return False
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        if not is_elementwise_op(op_desc.type()):
+            return False
+        op_dist_attr = dist_op.dist_attr
+        dims_mapping_list = []
+
+        input_arg_names = op_desc.input_arg_names()
+        input_max_dims_mapping_len = -1
+        for arg_name in input_arg_names:
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if input_max_dims_mapping_len < len(dims_mapping):
+                input_max_dims_mapping_len = len(dims_mapping)
+            dims_mapping_list.append(dims_mapping)
+
+        output_arg_names = op_desc.output_arg_names()
+        output_max_dims_mapping_len = -1
+        for arg_name in output_arg_names:
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if output_max_dims_mapping_len < len(dims_mapping):
+                output_max_dims_mapping_len = len(dims_mapping)
+            dims_mapping_list.append(dims_mapping)
+
+        assert input_max_dims_mapping_len == output_max_dims_mapping_len
+        max_dims_mapping_len = input_max_dims_mapping_len
+
+        for idx in range(max_dims_mapping_len):
+            dim_mappings = []
+            for dims_mapping in dims_mapping_list:
+                if idx < len(dims_mapping):
+                    dim_mappings.append(dims_mapping[-(idx + 1)])
+            if not all(
+                dim_mappings[0] == dim_mapping for dim_mapping in dim_mappings
+            ):
+                return False
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        dims_mapping_list = []
+
+        input_arg_names = op_desc.input_arg_names()
+        input_dims_mapping_dict = {}
+        input_dims_mapping_lens = {}
+        input_max_dims_mapping_len = -1
+        for arg_name in input_arg_names:
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if input_max_dims_mapping_len < len(dims_mapping):
+                input_max_dims_mapping_len = len(dims_mapping)
+            input_dims_mapping_dict[arg_name] = dims_mapping
+            input_dims_mapping_lens[arg_name] = len(dims_mapping)
+        for arg_name in input_arg_names:
+            if input_dims_mapping_lens[arg_name] < input_max_dims_mapping_len:
+                new_dims_mapping = [
+                    -1 for _ in range(input_max_dims_mapping_len)
+                ]
+                for i in range(input_dims_mapping_lens[arg_name]):
+                    new_idx = (
+                        input_max_dims_mapping_len
+                        - input_dims_mapping_lens[arg_name]
+                    ) + i
+                    new_dims_mapping[new_idx] = input_dims_mapping_dict[
+                        arg_name
+                    ][i]
+                dims_mapping_list.append(new_dims_mapping)
+            else:
+                dims_mapping_list.append(input_dims_mapping_dict[arg_name])
+
+        output_arg_names = op_desc.output_arg_names()
+        output_dims_mapping_dict = {}
+        output_dims_mapping_lens = {}
+        output_max_dims_mapping_len = -1
+        for arg_name in output_arg_names:
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if output_max_dims_mapping_len < len(dims_mapping):
+                output_max_dims_mapping_len = len(dims_mapping)
+            output_dims_mapping_dict[arg_name] = dims_mapping
+            output_dims_mapping_lens[arg_name] = len(dims_mapping)
+        for arg_name in output_arg_names:
+            if output_dims_mapping_lens[arg_name] < output_max_dims_mapping_len:
+                new_dims_mapping = [
+                    -1 for _ in range(output_max_dims_mapping_len)
+                ]
+                for i in range(output_dims_mapping_lens[arg_name]):
+                    new_idx = (
+                        output_max_dims_mapping_len
+                        - output_dims_mapping_lens[arg_name]
+                    ) + i
+                    new_dims_mapping[new_idx] = output_dims_mapping_dict[
+                        arg_name
+                    ][i]
+                dims_mapping_list.append(new_dims_mapping)
+            else:
+                dims_mapping_list.append(output_dims_mapping_dict[arg_name])
+
+        assert input_max_dims_mapping_len == output_max_dims_mapping_len
+        max_dims_mapping_len = input_max_dims_mapping_len
+        compatible_dims_mapping = compute_compatible_dims_mapping(
+            dims_mapping_list
+        )
+        if compatible_dims_mapping is None:
+            return False
+
+        for arg_name in input_arg_names:
+            if input_dims_mapping_lens[arg_name] < max_dims_mapping_len:
+                new_dims_mapping = [
+                    -1 for _ in range(input_dims_mapping_lens[arg_name])
+                ]
+                for i in range(input_dims_mapping_lens[arg_name]):
+                    new_idx = (
+                        max_dims_mapping_len - input_dims_mapping_lens[arg_name]
+                    ) + i
+                    new_dims_mapping[i] = compatible_dims_mapping[new_idx]
+                if new_dims_mapping != input_dims_mapping_dict[arg_name]:
+                    op_dist_attr.set_input_dims_mapping(
+                        arg_name, new_dims_mapping
+                    )
+                    changed = True
+            else:
+                if compatible_dims_mapping != input_dims_mapping_dict[arg_name]:
+                    op_dist_attr.set_input_dims_mapping(
+                        arg_name, compatible_dims_mapping
+                    )
+                    changed = True
+
+        for arg_name in output_arg_names:
+            if output_dims_mapping_lens[arg_name] < max_dims_mapping_len:
+                new_dims_mapping = [
+                    -1 for _ in range(output_dims_mapping_lens[arg_name])
+                ]
+                for i in range(output_dims_mapping_lens[arg_name]):
+                    new_idx = (
+                        max_dims_mapping_len
+                        - output_dims_mapping_lens[arg_name]
+                    ) + i
+                    new_dims_mapping[i] = compatible_dims_mapping[new_idx]
+                if new_dims_mapping != output_dims_mapping_dict[arg_name]:
+                    op_dist_attr.set_output_dims_mapping(
+                        arg_name, new_dims_mapping
+                    )
+                    changed = True
+            else:
+                if (
+                    compatible_dims_mapping
+                    != output_dims_mapping_dict[arg_name]
+                ):
+                    op_dist_attr.set_output_dims_mapping(
+                        arg_name, compatible_dims_mapping
+                    )
+                    changed = True
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "elementwise", DistributedElementwiseImpl0("replicate_parallel")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_embedding.py

@@ -0,0 +1,671 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import paddle
+from paddle.common_ops_import import check_variable_and_dtype
+from paddle.distributed.auto_parallel.static.cost.comm_op_cost import (
+    AllReduceOpCost,
+    IdentityOpCost,
+)
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+from paddle.framework import core
+from paddle.utils import unique_name
+
+from ..completion import get_phi_spmd_rule
+from ..cost import (
+    EmbeddingGradOpCost,
+    EmbeddingOpCost,
+    build_comm_costs_from_descs,
+    build_comm_desc_from_dist_op,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_dp_costs,
+)
+from ..dist_attribute import OperatorDistAttr
+from ..process_group import new_process_group
+from ..utils import (
+    _get_comm_group,
+    _get_corresponding_rank,
+    _get_idx_in_axis,
+    compute_compatible_and_update_dim_mapping,
+    get_dist_tensor_spec,
+    is_dim_replicate,
+    is_dim_shard,
+    set_var_dist_attr,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    ParallelMode,
+    get_default_distributed_operator_impl,
+    gradient_synchronization,
+    naive_copy_op_dist_attr_for_program,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    set_comm_op_dist_attr_for_program,
+    update_op_dims_mapping,
+)
+
+
+class DistributedEmbedding(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+        assert dist_op.serial_op.type == "lookup_table_v2", (
+            f"{dist_op.serial_op.type} is not supported by dist embedding yet."
+        )
+
+        x_name = op_desc.input('Ids')[0]
+        w_name = op_desc.input('W')[0]
+        out_name = op_desc.output('Out')[0]
+        padding_idx = op_desc.attr('padding_idx')
+        is_sparse = op_desc.attr('is_sparse')
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        w_spec = get_dist_tensor_spec(dist_op, w_name)
+        output_spec = get_dist_tensor_spec(dist_op, out_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("embedding")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec, w_spec, padding_idx, is_sparse)
+        bw_results = rule.infer_backward(
+            x_spec, w_spec, output_spec, padding_idx, is_sparse
+        )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, [x_name, w_name], [out_name], fw_results, bw_results
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        reverted = False
+        op_dist_attr = dist_op.dist_attr
+        op_desc = dist_op.serial_op.desc
+        out_name = op_desc.output('Out')[0]
+        out_dist_attr = op_dist_attr.get_output_dist_attr(out_name)
+
+        # vocab parallel embedding
+        if out_dist_attr._is_partial():
+            op_dist_attr.impl_type = op_desc.type()
+            op_dist_attr.impl_idx = 0
+        # data parallel or col parallel of weight
+        else:
+            default_impl = get_default_distributed_operator_impl()
+            op_dist_attr.impl_type = default_impl.type
+            op_dist_attr.impl_idx = default_impl.idx
+
+        return reverted
+
+
+register_distributed_operator_impl_container(
+    DistributedEmbedding("lookup_table_v2")
+)
+register_distributed_operator_impl_container(
+    DistributedEmbedding("c_embedding")
+)
+register_distributed_operator_impl_container(
+    DistributedEmbedding("lookup_table")
+)
+
+
+def adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var):
+    assert len(Ids_var.shape) == 3, (
+        f"input Ids to lookup_table should have 3 dimensions but got [{Ids_var.name}] with shape [{Ids_var.shape}]"
+    )
+    if not Ids_var.stop_gradient:
+        raise NotImplementedError(
+            'Requiring the gradient of Ids of lookup_table(v1) dist op is not currently supported. Please open an issue with details on your use case so that we can prioritize adding this (for instance, adversarial training for language model).'
+        )
+
+    target_shape = list(Ids_var.shape[:-1])
+    intermediate_var_0 = main_block.create_var(
+        name=unique_name.generate_with_ignorable_key(
+            ".".join(["dist_reshape", 'tmp'])
+        ),
+        dtype=Ids_var.dtype,
+        shape=target_shape,
+        type=core.VarDesc.VarType.DENSE_TENSOR,
+        persistable=False,
+        stop_gradient=True,
+    )
+
+    target_shape = [0, *list(Ids_var.shape[:-1])]
+    xshape_var = main_block.create_var(
+        name=unique_name.generate_with_ignorable_key(
+            ".".join(["dist_Xshape", 'tmp'])
+        ),
+        dtype=Ids_var.dtype,
+        shape=target_shape,
+        type=core.VarDesc.VarType.DENSE_TENSOR,
+        persistable=False,
+        stop_gradient=True,
+    )
+
+    # TODO use inplace reshape for memory saving
+    reshape_op = main_block.append_op(
+        type='reshape2',
+        inputs={'X': [Ids_var]},
+        outputs={'Out': [intermediate_var_0], 'XShape': [xshape_var]},
+        attrs={
+            "shape": [0, -1],
+        },
+    )
+
+    # set dist attr
+    op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+    Ids_var_dist_attr = op_dist_attr.get_input_dist_attr(Ids_var.name)
+    assert Ids_var_dist_attr is not None
+    intermediate_var_0_dist_attr = set_var_dist_attr(
+        ctx,
+        intermediate_var_0,
+        Ids_var_dist_attr.dims_mapping,
+        Ids_var_dist_attr.process_mesh,
+        chunk_id=Ids_var_dist_attr.chunk_id,
+    )
+    set_var_dist_attr(
+        ctx,
+        xshape_var,
+        [-1, *list(Ids_var_dist_attr.dims_mapping)],
+        Ids_var_dist_attr.process_mesh,
+        chunk_id=Ids_var_dist_attr.chunk_id,
+    )
+    # rename src_op's input
+    src_op._rename_input(Ids_var.name, intermediate_var_0.name)
+    op_dist_attr.del_input_dist_attr(Ids_var.name)
+    op_dist_attr.set_input_dist_attr(
+        intermediate_var_0.name, intermediate_var_0_dist_attr
+    )
+
+    new_op_dist_attr = OperatorDistAttr()
+    new_op_dist_attr.process_mesh = Ids_var_dist_attr.process_mesh
+    new_op_dist_attr.impl_type = "default"
+    new_op_dist_attr.impl_idx = 0
+    new_op_dist_attr.chunk_id = Ids_var_dist_attr.chunk_id
+    new_op_dist_attr.set_input_dims_mapping(
+        Ids_var.name, Ids_var_dist_attr.dims_mapping
+    )
+    new_op_dist_attr.set_output_dims_mapping(
+        intermediate_var_0.name, Ids_var_dist_attr.dims_mapping
+    )
+    new_op_dist_attr.set_output_dims_mapping(
+        xshape_var.name, [-1, *list(Ids_var_dist_attr.dims_mapping)]
+    )
+    ctx.set_op_dist_attr_for_program(reshape_op, new_op_dist_attr)
+
+    return intermediate_var_0
+
+
+# RowParallel
+class DistributedEmbeddingImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        """Calculate the cost by the op role."""
+        cost = None
+        if int(op_role) == int(OpRole.Forward):
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        elif int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_op.dist_attr.process_mesh.process_ids
+        # embedding need start_index
+        cost_mapping = build_comp_costs_from_descs(
+            EmbeddingOpCost, ctx, processes, desc_mapping, cluster
+        )
+
+        serial_op = dist_op.serial_op
+        parallel_axis = dist_op.dist_attr.get_input_dims_mapping(
+            serial_op.input("W")[0]
+        )[0]
+        attrs = {"use_calc_stream": True, "use_model_parallel": True}
+        var_names = serial_op.output("Out")
+        all_reduce_sum_desc_mapping = build_comm_desc_from_dist_op(
+            "all_reduce",
+            dist_op,
+            ctx,
+            var_names,
+            attrs=attrs,
+            parallel_axis=parallel_axis,
+        )
+
+        comm_op_cost_list = build_comm_costs_from_descs(
+            AllReduceOpCost,
+            ctx,
+            processes,
+            all_reduce_sum_desc_mapping,
+            cluster,
+        )
+
+        res_cost = [cost_mapping, comm_op_cost_list]
+
+        return res_cost
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # by now the backward function only insert the gradient allreduce for dist op itself
+        res = []
+        backward_op = dist_op.serial_op
+        main_block = backward_op.block
+        dist_attr = dist_op.dist_attr
+
+        embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
+            backward_op.input("W")[0]
+        )[0]
+        parallel_axis = embedding_row_dim_mapping
+        attrs = {"use_calc_stream": True, "use_model_parallel": True}
+        var_names = [backward_op.input("Out@GRAD")[0]]
+        c_identity_desc_mapping = build_comm_desc_from_dist_op(
+            "c_identity",
+            dist_op,
+            ctx,
+            var_names,
+            attrs=attrs,
+            parallel_axis=parallel_axis,
+        )
+
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        comm_op_cost_list = build_comm_costs_from_descs(
+            IdentityOpCost, ctx, processes, c_identity_desc_mapping, cluster
+        )
+        res.append(comm_op_cost_list)
+
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        cost_mapping = build_comp_costs_from_descs(
+            EmbeddingGradOpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        # need gradient allreduce
+        var_dim_mapping = dist_attr.get_input_dims_mapping(
+            backward_op.input("Ids")[0]
+        )
+        mesh_shape = process_mesh.shape
+        batch_size_axis = var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+        if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+            parallel_axis = batch_size_axis
+            attrs = {"use_calc_stream": True}
+            var_names = [backward_op.output('W@GRAD')[0]]
+            build_dp_costs(
+                res, dist_op, ctx, var_names, attrs, parallel_axis, cluster
+            )
+
+        return res
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        ids_name = op_desc.input('Ids')[0]
+        w_name = op_desc.input('W')[0]
+        ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
+        w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
+        if is_dim_replicate(w_dims_mapping[-2]) or is_dim_shard(
+            w_dims_mapping[-1]
+        ):
+            return False
+        # Other dimensions must be replicate except the batch dimension
+        for mapping in ids_dims_mapping[1:]:
+            if is_dim_shard(mapping):
+                return False
+
+        if is_dim_shard(ids_dims_mapping[0]) and is_dim_shard(
+            w_dims_mapping[-2]
+        ):
+            if ids_dims_mapping[0] == w_dims_mapping[-2]:
+                return False
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_name = op_desc.output('Out')[0]
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        # Other dimensions must be replicate except the batch dimension
+        for mapping in out_dims_mapping[1:]:
+            if is_dim_shard(mapping):
+                return False
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        ids_name = op_desc.input('Ids')[0]
+        w_name = op_desc.input('W')[0]
+        out_name = op_desc.output('Out')[0]
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
+        w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
+
+        if ids_dims_mapping != out_dims_mapping[: len(ids_dims_mapping)]:
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        ids_name = op_desc.input('Ids')[0]
+        w_name = op_desc.input('W')[0]
+        out_name = op_desc.output('Out')[0]
+        ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
+        w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        for i in range(len(ids_dims_mapping)):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [ids_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                changed = True
+
+        dim_changed = compute_compatible_and_update_dim_mapping(
+            [w_dims_mapping, out_dims_mapping], [-1, -1]
+        )
+        if dim_changed:
+            changed = True
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(ids_name, ids_dims_mapping)
+            op_dist_attr.set_input_dims_mapping(w_name, w_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        """
+        kwargs: inputname_mapping & outputname_mapping
+        """
+
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None, (
+            f"forward op [{src_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        assert 'Ids' in kwargs, "input [{}] is not given".format('Ids')
+        assert 'W' in kwargs, "input [{}] is not given".format('W')
+        assert 'Out' in kwargs, "output [{}] is not given".format('Out')
+
+        assert len(kwargs['Ids']) == 1, (
+            "row_parallel_embedding input Ids take 1 variable but got {}".format(
+                kwargs['Ids']
+            )
+        )
+        assert len(kwargs['W']) == 1, (
+            "row_parallel_embedding input W take 1 variable but got {}".format(
+                kwargs['W']
+            )
+        )
+        assert len(kwargs['Out']) == 1, (
+            "row_parallel_embedding output Out take 1 variable but got {}".format(
+                kwargs['Out']
+            )
+        )
+
+        Ids_var = main_block._var_recursive(kwargs['Ids'][0])
+        Weight_var = main_block._var_recursive(kwargs['W'][0])
+        Out_var = main_block._var_recursive(kwargs['Out'][0])
+
+        # support lookup_table_v1
+        if src_op.type == 'lookup_table':
+            Ids_var = adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var)
+
+        # got dist attribute info
+        embedding_row_dim_mapping = op_dist_attr.get_input_dims_mapping(
+            Weight_var.name
+        )[0]
+        assert embedding_row_dim_mapping >= 0, (
+            f"row_parallel_embedding's row should be divided by a specific mesh axis, but got [{embedding_row_dim_mapping}]"
+        )
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
+        if rank_id not in process_mesh_group:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        # A generalized method to calculate embedding offset using cartesian product
+        relative_idx = _get_idx_in_axis(
+            process_mesh_group,
+            process_mesh_shape,
+            embedding_row_dim_mapping,
+            rank_id,
+        )
+
+        per_part_size = Weight_var.shape[0]
+        relative_idx = relative_idx * per_part_size
+
+        # TODO calculate ring id
+        parallel_axis = embedding_row_dim_mapping
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        # append op
+        check_variable_and_dtype(
+            Ids_var, 'input', ['int32', 'int64'], 'c_embedding'
+        )
+
+        # infer new var shape with op dist attr
+        out_tensor_dist_attr = ctx.get_tensor_dist_attr_for_program(Out_var)
+        assert out_tensor_dist_attr is not None
+        out_var_dist_attr = op_dist_attr.get_output_dist_attr(Out_var.name)
+        assert out_var_dist_attr is not None
+
+        c_embedding_op_desc = main_block.append_op(type='nop').desc
+        c_embedding_op_desc.set_type("c_embedding")
+        c_embedding_op_desc.set_input('Ids', [Ids_var.name])
+        c_embedding_op_desc.set_input('W', [Weight_var.name])
+        c_embedding_op_desc.set_output('Out', [Out_var.name])
+        c_embedding_op_desc._set_attr('start_index', relative_idx)
+        c_embedding_op_desc._set_attr(OP_ROLE_KEY, src_op.attr('op_role'))
+        c_embedding_op = main_block.ops[-1]
+        assert c_embedding_op.type == "c_embedding"
+        naive_copy_op_dist_attr_for_program(c_embedding_op, src_op, ctx)
+
+        # use_model_parallel
+        all_reduce_sum_op = main_block.append_op(
+            type='all_reduce',
+            inputs={'x': [Out_var]},
+            outputs={'out': [Out_var]},
+            attrs={
+                'ring_id': group.id,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
+                'use_model_parallel': True,
+                OP_ROLE_KEY: src_op.attr('op_role'),
+            },
+        )
+        all_reduce_sum_op._set_attr(
+            'op_namescope', '/' + ParallelMode.TensorParallel
+        )
+        # allreduce
+        set_comm_op_dist_attr_for_program(
+            all_reduce_sum_op,
+            op_dist_attr.process_mesh,
+            out_var_dist_attr,
+            ctx,
+            chunk_id=op_dist_attr.chunk_id,
+        )
+
+        # param initialization sync
+        if Weight_var.is_parameter and not op_dist_attr.is_recompute:
+            if Weight_var.name in dist_op_context.already_init_sync_vars:
+                return
+            dist_op_context.already_init_sync_vars.add(Weight_var.name)
+            param = startup_block.var(Weight_var.name)
+            param_dist_attr = ctx.get_tensor_dist_attr_for_program(param)
+            process_mesh = param_dist_attr.process_mesh
+            dim_mapping = param_dist_attr.dims_mapping
+
+            # NOTE all not split axis should be presented in mesh
+            for axis, size in enumerate(process_mesh.shape):
+                if size <= 1 or axis in dim_mapping:
+                    pass
+                else:
+                    group_ranks = _get_comm_group(
+                        process_mesh.process_ids,
+                        process_mesh.shape,
+                        axis,
+                        rank_id,
+                    )
+                    sync_group = new_process_group(group_ranks)
+
+                    broadcast_op = startup_block.append_op(
+                        type='broadcast',
+                        inputs={'x': param},
+                        outputs={'out': param},
+                        attrs={
+                            'ring_id': sync_group.id,
+                            'root': 0,
+                            OP_ROLE_KEY: OpRole.Forward,
+                        },
+                    )
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        # by now the backward function only insert the gradient allreduce for dist op itself
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        backward_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
+        assert dist_attr is not None, (
+            f"backward op [{backward_op}] don't have dist attribute !"
+        )
+
+        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
+        if rank_id not in dist_attr.process_mesh.process_ids:
+            rank_id = _get_corresponding_rank(
+                ctx, dist_attr.process_mesh, rank_id
+            )
+
+        assert 'Ids' in kwargs, "input [{}] is not given".format('Ids')
+        assert 'W' in kwargs, "input [{}] is not given".format('W')
+        assert 'Out@GRAD' in kwargs, "input [{}] is not given".format('Out')
+        assert 'W@GRAD' in kwargs, "output [{}] is not given".format('W@GRAD')
+
+        assert len(kwargs['Ids']) == 1, (
+            "row_parallel_embedding input Ids take 1 variable but got {}".format(
+                kwargs['Ids']
+            )
+        )
+        assert len(kwargs['W']) == 1, (
+            "row_parallel_embedding input Ids take 1 variable but got {}".format(
+                kwargs['W']
+            )
+        )
+        assert len(kwargs['Out@GRAD']) == 1, (
+            "row_parallel_embedding input Ids take 1 variable but got {}".format(
+                kwargs['Out']
+            )
+        )
+        assert len(kwargs['W@GRAD']) == 1, (
+            "row_parallel_embedding output Ids take 1 variable but got {}".format(
+                kwargs['W@GRAD']
+            )
+        )
+
+        Ids_var = main_block._var_recursive(kwargs['Ids'][0])
+        Weight_var = main_block._var_recursive(kwargs['W'][0])
+        Out_grad = main_block._var_recursive(kwargs['Out@GRAD'][0])
+        Weight_grad = main_block._var_recursive(kwargs['W@GRAD'][0])
+
+        embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
+            Weight_var.name
+        )[0]
+        assert embedding_row_dim_mapping >= 0, (
+            f"row_parallel_embedding's row should be divided by a specific mesh axis, but got [{embedding_row_dim_mapping}]"
+        )
+        process_mesh_shape = dist_attr.process_mesh.shape
+        process_mesh_group = dist_attr.process_mesh.process_ids
+
+        # A generalized method to calculate embedding offset using cartesian product
+        relative_idx = _get_idx_in_axis(
+            process_mesh_group,
+            process_mesh_shape,
+            embedding_row_dim_mapping,
+            rank_id,
+        )
+        per_part_size = Weight_var.shape[0]
+        relative_idx = relative_idx * per_part_size
+
+        c_embedding_grad_op_desc = main_block.append_op(type='nop').desc
+        c_embedding_grad_op_desc.set_type("c_embedding_grad")
+        c_embedding_grad_op_desc.set_input('Ids', [Ids_var.name])
+        c_embedding_grad_op_desc.set_input('W', [Weight_var.name])
+        c_embedding_grad_op_desc.set_input('Out@GRAD', [Out_grad.name])
+        c_embedding_grad_op_desc.set_output('W@GRAD', [Weight_grad.name])
+        c_embedding_grad_op_desc._set_attr('start_index', relative_idx)
+        c_embedding_grad_op_desc._set_attr(OP_ROLE_KEY, OpRole.Backward)
+
+        c_embedding_grad_op = main_block.ops[-1]
+        assert c_embedding_grad_op.type == "c_embedding_grad"
+        naive_copy_op_dist_attr_for_program(
+            c_embedding_grad_op, backward_op, ctx
+        )
+
+        # data parallel gradient synchronization
+        act_grad_names = [Ids_var.name]
+        out_grad_names = [kwargs['W@GRAD'][0]]
+
+        gradient_synchronization(
+            ctx, backward_op, act_grad_names, out_grad_names, rank_id
+        )
+
+
+register_distributed_operator_impl(
+    "lookup_table_v2", DistributedEmbeddingImpl("row_parallel")
+)
+register_distributed_operator_impl(
+    "c_embedding", DistributedEmbeddingImpl("row_parallel")
+)
+register_distributed_operator_impl(
+    "lookup_table", DistributedEmbeddingImpl("row_parallel")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_expand_as.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ..completion import get_phi_spmd_rule
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedExpandAs(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        input_arg_names = op_desc.input_arg_names()
+        output_arg_names = op_desc.output_arg_names()
+        target_shape = op_desc.attr('target_shape')
+
+        input_specs = []
+        for name in input_arg_names:
+            input_specs.append(get_dist_tensor_spec(dist_op, name))
+
+        assert len(input_specs) == 2
+
+        output_spec = get_dist_tensor_spec(dist_op, output_arg_names[0], False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("expand_as")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(
+            input_specs[0], input_specs[1], target_shape
+        )
+        bw_results = rule.infer_backward(
+            input_specs[0], input_specs[1], output_spec, target_shape
+        )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            input_arg_names,
+            output_arg_names,
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedExpandAs("expand_as_v2")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_fill_constant_batch_size_like.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..cost import (
+    FillConstantBatchSizeLikeOpCost,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+)
+from ..utils import compute_compatible_and_update_dim_mapping
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedFillConstantBatchSizeLike(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(
+    DistributedFillConstantBatchSizeLike("fill_constant_batch_size_like")
+)
+
+
+class DistributedFillConstantBatchSizeLikeImpl0(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            raise ValueError(
+                "The fill_constant_batch_size_like has no grad op."
+            )
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_op.dist_attr.process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            FillConstantBatchSizeLikeOpCost,
+            ctx,
+            processes,
+            desc_mapping,
+            cluster,
+        )
+
+        res_cost = [cost_mapping]
+        return res_cost
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_name = op_desc.output('Out')[0]
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        shape_list = op_desc.attr("shape")
+
+        if len(shape_list) != len(out_dims_mapping):
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_name = op_desc.output('Out')[0]
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        in_name = op_desc.input('Input')[0]
+        in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
+
+        # the dim_mapping of batch dimension should be the same
+        return out_dims_mapping[0] == in_dims_mapping[0]
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('Input')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        # only the batch size dimension of input and output are relative.
+        dim_changed = compute_compatible_and_update_dim_mapping(
+            [x_dims_mapping, out_dims_mapping], [0, 0]
+        )
+        if dim_changed:
+            changed = True
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        """
+        kwargs: inputname_mapping & outputname_mapping
+        """
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "fill_constant_batch_size_like",
+    DistributedFillConstantBatchSizeLikeImpl0("fill_by_shape"),
+)

python/paddle/distributed/auto_parallel/static/operators/dist_flash_attn.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ...random import determinate_rng, is_enable_auto_rand_ctrl
+from .common import (
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_eltwise import DistributedElementwiseImpl0
+
+
+class DistributedFlashAttn(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(DistributedFlashAttn("flash_attn"))
+
+
+# Dist FlashAttn with Random Control
+# NOTE(zhiqiu): trick implementation, copy dist_attr of q,k,v to out
+class DistributedFlashAttnImpl0(DistributedElementwiseImpl0):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        return True
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+
+        if (
+            is_enable_auto_rand_ctrl()
+            and not op_dist_attr.is_recompute
+            and rank_id in op_dist_attr.process_mesh.process_ids
+        ):
+            assert op_dist_attr is not None, (
+                f"forward op [{src_op}] don't have dist attribute !"
+            )
+
+            if (
+                len(kwargs.get('fixed_seed_offset', [])) > 0
+                or len(src_op.input("fixed_seed_offset")) > 0
+            ):
+                # TODO(kuizhiqing) recompute should go here
+                pass
+            else:
+                # determinate rng
+                q_var = main_block._var_recursive(kwargs['q'][0])
+                k_var = main_block._var_recursive(kwargs['k'][0])
+                q_dims_mapping = op_dist_attr.get_input_dims_mapping(q_var.name)
+                k_dims_mapping = op_dist_attr.get_input_dims_mapping(k_var.name)
+                process_mesh = op_dist_attr.process_mesh
+                dims_mapping = [*q_dims_mapping[:3], q_dims_mapping[2]]
+
+                rng_name = determinate_rng(rank_id, dims_mapping, process_mesh)
+                assert rng_name is not None and rng_name != ""
+
+                src_op._set_attr('rng_name', rng_name)
+
+        DistributedElementwiseImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        # dropout backward is deterministic by mask, and not need for random state control
+        DistributedElementwiseImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "flash_attn", DistributedFlashAttnImpl0("random_control")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_fused_attention.py

@@ -0,0 +1,235 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..process_group import new_process_group
+from ..utils import (
+    _get_comm_group,
+    _get_corresponding_rank,
+    compute_compatible_and_update_dim_mapping,
+    is_dim_replicate,
+    is_dim_shard,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedFusedAttention(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(
+    DistributedFusedAttention("fused_attention")
+)
+
+
+class DistributedFusedAttentionImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        qkv_w = op_desc.input('QKVW')[0]
+        qkv_bias = op_desc.input('QKVBias')[0]
+        out_w = op_desc.input('OutLinearW')[0]
+        out_bias = op_desc.input('OutLinearBias')[0]
+
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        qkv_w_dims_mapping = op_dist_attr.get_input_dims_mapping(qkv_w)
+        qkv_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(qkv_bias)
+        out_w_dims_mapping = op_dist_attr.get_input_dims_mapping(out_w)
+        out_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(out_bias)
+
+        head_axis = 1
+        for mapping in x_dims_mapping[1:-1]:
+            if is_dim_shard(mapping):
+                return False
+        if len(qkv_w_dims_mapping) != 4 or is_dim_replicate(
+            qkv_w_dims_mapping[head_axis]
+        ):
+            return False
+        if len(qkv_bias_dims_mapping) != 3 or is_dim_replicate(
+            qkv_bias_dims_mapping[head_axis]
+        ):
+            return False
+        if is_dim_replicate(out_w_dims_mapping[0]):
+            return False
+        if is_dim_shard(out_bias_dims_mapping[-1]):
+            return False
+
+        replicated_dims = [
+            qkv_w_dims_mapping[0],
+            qkv_w_dims_mapping[-2],
+            qkv_w_dims_mapping[-1],
+            qkv_bias_dims_mapping[0],
+            qkv_bias_dims_mapping[-1],
+            out_w_dims_mapping[-1],
+            out_bias_dims_mapping[-1],
+        ]
+        for mapping in replicated_dims:
+            if is_dim_shard(mapping):
+                return False
+        if qkv_bias_dims_mapping[head_axis] != qkv_w_dims_mapping[head_axis]:
+            return False
+        if qkv_bias_dims_mapping[head_axis] != out_w_dims_mapping[0]:
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+
+        # none of output should be sharded
+        for out_name in op_desc.output_names():
+            out = op_desc.output(out_name)[0]
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out)
+            for mapping in out_dims_mapping[1:-1]:
+                if is_dim_shard(mapping):
+                    return False
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_names = op_desc.output('Y')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        for out_name in out_names:
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+            if x_dims_mapping != out_dims_mapping:
+                return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_names = op_desc.output('Y')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+
+        for out_name in out_names:
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+            for i in range(len(x_dims_mapping)):
+                dim_changed = compute_compatible_and_update_dim_mapping(
+                    [x_dims_mapping, out_dims_mapping], [i, i]
+                )
+                if dim_changed:
+                    changed = True
+                    op_dist_attr.set_output_dims_mapping(
+                        out_name, out_dims_mapping
+                    )
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+
+        if rank_id not in op_dist_attr.process_mesh.process_ids:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        # infer logic comm presentation
+        head_axis = 1
+        qkv_w = src_op.input('QKVW')[0]
+        qkv_w_col_dim_mapping = op_dist_attr.get_input_dims_mapping(qkv_w)[
+            head_axis
+        ]
+        assert qkv_w_col_dim_mapping >= 0, (
+            f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{qkv_w_col_dim_mapping}]"
+        )
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        parallel_axis = qkv_w_col_dim_mapping
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        # insert op
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+        # setting comm id
+        new_op = main_block.ops[-1]
+        assert new_op.type == "fused_attention"
+        new_op._set_attr("ring_id", int(group.id))
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+
+        if rank_id not in op_dist_attr.process_mesh.process_ids:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        # infer logic comm presentation
+        out_w = src_op.input('OutLinearW')[0]
+        out_w_col_dim_mapping = op_dist_attr.get_input_dims_mapping(out_w)[-1]
+        assert out_w_col_dim_mapping >= 0, (
+            f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{out_w_col_dim_mapping}]"
+        )
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        parallel_axis = out_w_col_dim_mapping
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        # insert op
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+        # setting comm id
+        new_op = main_block.ops[-1]
+        assert new_op.type == "fused_attention_grad"
+        new_op._set_attr("ring_id", int(group.id))
+
+
+register_distributed_operator_impl(
+    "fused_attention", DistributedFusedAttentionImpl("tensor_parallel")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_fused_dropout_add.py

@@ -0,0 +1,195 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+import logging
+
+import paddle
+from paddle.base.log_helper import get_logger
+from paddle.framework import core
+from paddle.utils import unique_name
+
+from ...random import determinate_rng, is_enable_auto_rand_ctrl
+from ..utils import (
+    naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
+    set_var_dist_attr,
+)
+from .common import (
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_eltwise import DistributedDefaultImpl0, DistributedElementwiseImpl0
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class DistributedDropout(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(
+    DistributedDropout("fused_dropout_add")
+)
+
+
+# Dist Dropout with Random Control
+# Dropout re-use the compatible and cost function of elementwise
+class DistributedDropoutImpl0(DistributedElementwiseImpl0):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        return True
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+
+        if is_enable_auto_rand_ctrl() and not op_dist_attr.is_recompute:
+            assert op_dist_attr is not None, (
+                f"forward op [{src_op}] don't have dist attribute !"
+            )
+
+            assert 'seed_tensor' in kwargs, "input [{}] is not given".format(
+                'seed_tensor'
+            )
+
+            if (
+                src_op.has_attr("fix_seed")
+                and src_op.attr("fix_seed")
+                and src_op.has_attr("seed")
+                and src_op.attr("seed")
+            ):
+                _logger.info(
+                    f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
+                )
+            elif rank_id not in op_dist_attr.process_mesh.process_ids:
+                pass
+            elif (
+                len(kwargs['seed_tensor']) > 0
+                or len(src_op.input("seed_tensor")) > 0
+            ):
+                seed_var_name = kwargs['seed_tensor'][0]
+                if seed_var_name.startswith('rc_seed'):
+                    pre_op = main_block.ops[-1]
+                    assert (
+                        pre_op.type == "seed"
+                        and len(pre_op.attr("rng_name")) == 0
+                    ), f"found exception op {pre_op}"
+
+                    # determinate rng
+                    X_var = main_block._var_recursive(kwargs['x'][0])
+                    X_dims_mapping = op_dist_attr.get_input_dims_mapping(
+                        X_var.name
+                    )
+                    process_mesh = op_dist_attr.process_mesh
+                    rng_name = determinate_rng(
+                        rank_id, X_dims_mapping, process_mesh
+                    )
+                    # make recompute seed under control
+                    pre_op._set_attr("rng_name", rng_name)
+                    pre_op._set_attr("deterministic", True)
+                    pre_op._set_attr("force_cpu", True)
+                else:
+                    _logger.info(
+                        f"Auto Parallel Random Control Skipped Since manual seed is set by user: {src_op}"
+                    )
+            else:
+                # determinate rng
+                X_var = main_block._var_recursive(kwargs['x'][0])
+                X_dims_mapping = op_dist_attr.get_input_dims_mapping(X_var.name)
+                process_mesh = op_dist_attr.process_mesh
+
+                rng_name = determinate_rng(
+                    rank_id, X_dims_mapping, process_mesh
+                )
+                assert rng_name is not None and rng_name != ""
+
+                # insert seed op
+                seed_var = main_block.create_var(
+                    name=unique_name.generate_with_ignorable_key(
+                        ".".join(["tensor_parallel_seed", 'tmp'])
+                    ),
+                    dtype=paddle.int32,
+                    type=core.VarDesc.VarType.DENSE_TENSOR,
+                    persistable=False,
+                    stop_gradient=False,
+                )
+
+                # set new seed_var's dist_attr
+                seed_var_dims_mapping = [-1]
+                seed_var_dist_attr = set_var_dist_attr(
+                    ctx,
+                    seed_var,
+                    seed_var_dims_mapping,
+                    process_mesh,
+                    chunk_id=op_dist_attr.chunk_id,
+                )
+
+                # adopt for recompute
+                # force_cpu to reduce sync copy from CPU->GPU->CPU, and reduce pipeline hang
+                seed_op = main_block.append_op(
+                    type='seed',
+                    outputs={'Out': seed_var},
+                    attrs={
+                        'deterministic': True,
+                        'rng_name': rng_name,
+                        'force_cpu': True,
+                    },
+                )
+                seed_op._set_attr('op_namescope', 'auto_tensor_parallel_seed')
+                # set new seed op's dist_attr
+                naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
+                    seed_op,
+                    process_mesh,
+                    seed_var_dims_mapping,
+                    ctx,
+                    chunk_id=op_dist_attr.chunk_id,
+                )
+
+                # modify dropout op
+                src_op.desc.set_input("seed_tensor", [seed_var.name])
+                src_op._remove_attr("fix_seed")
+                src_op._remove_attr("seed")
+                op_dist_attr.set_input_dist_attr(
+                    seed_var.name, seed_var_dist_attr
+                )
+                kwargs['seed_tensor'] = [seed_var.name]
+
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        # dropout backward is deterministic by mask, and not need for random state control
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "fused_dropout_add", DistributedDropoutImpl0("random_control")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_fused_feedforward.py

@@ -0,0 +1,228 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..process_group import new_process_group
+from ..utils import (
+    _get_comm_group,
+    _get_corresponding_rank,
+    compute_compatible_and_update_dim_mapping,
+    is_dim_replicate,
+    is_dim_shard,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedFusedFeedForward(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(
+    DistributedFusedFeedForward("fused_feedforward")
+)
+
+
+class DistributedFusedFeedForwardImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        linear1_weight = op_desc.input('Linear1Weight')[0]
+        linear1_bias = op_desc.input('Linear1Bias')[0]
+        linear2_weight = op_desc.input('Linear2Weight')[0]
+        linear2_bias = op_desc.input('Linear2Bias')[0]
+
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        linear1_weight_dims_mapping = op_dist_attr.get_input_dims_mapping(
+            linear1_weight
+        )
+        linear1_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(
+            linear1_bias
+        )
+        linear2_weight_dims_mapping = op_dist_attr.get_input_dims_mapping(
+            linear2_weight
+        )
+        linear2_bias_dims_mapping = op_dist_attr.get_input_dims_mapping(
+            linear2_bias
+        )
+
+        for mapping in x_dims_mapping[1:-1]:
+            if is_dim_shard(mapping):
+                return False
+        if is_dim_shard(linear1_weight_dims_mapping[-2]) or is_dim_replicate(
+            linear1_weight_dims_mapping[-1]
+        ):
+            return False
+        if is_dim_replicate(linear1_bias_dims_mapping[-1]):
+            return False
+        if is_dim_replicate(linear2_weight_dims_mapping[-2]) or is_dim_shard(
+            linear2_weight_dims_mapping[-1]
+        ):
+            return False
+        if is_dim_shard(linear2_bias_dims_mapping[-1]):
+            return False
+        if linear1_weight_dims_mapping[-1] != linear2_weight_dims_mapping[-2]:
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+
+        # none of output should be sharded
+        for out_name in op_desc.output_names():
+            out = op_desc.output(out_name)[0]
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out)
+            for mapping in out_dims_mapping[1:-1]:
+                if is_dim_shard(mapping):
+                    return False
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_names = op_desc.output('Out')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        for out_name in out_names:
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+            if x_dims_mapping != out_dims_mapping:
+                return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_names = op_desc.output('Out')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+
+        for out_name in out_names:
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+            for i in range(len(x_dims_mapping)):
+                dim_changed = compute_compatible_and_update_dim_mapping(
+                    [x_dims_mapping, out_dims_mapping], [i, i]
+                )
+                if dim_changed:
+                    changed = True
+                    op_dist_attr.set_output_dims_mapping(
+                        out_name, out_dims_mapping
+                    )
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+
+        if rank_id not in op_dist_attr.process_mesh.process_ids:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        # infer logic comm presentation
+        linear1_weight = src_op.input('Linear1Weight')[0]
+        linear1_weight_col_dim_mapping = op_dist_attr.get_input_dims_mapping(
+            linear1_weight
+        )[-1]
+        assert linear1_weight_col_dim_mapping >= 0, (
+            f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{linear1_weight_col_dim_mapping}]"
+        )
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        parallel_axis = linear1_weight_col_dim_mapping
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        # insert op
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+        # setting comm id
+        new_op = main_block.ops[-1]
+        assert new_op.type == "fused_feedforward"
+        new_op._set_attr("ring_id", int(group.id))
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+
+        if rank_id not in op_dist_attr.process_mesh.process_ids:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        # infer logic comm presentation
+        linear2_weight = src_op.input('Linear2Weight')[0]
+        linear2_weight_col_dim_mapping = op_dist_attr.get_input_dims_mapping(
+            linear2_weight
+        )[-1]
+        assert linear2_weight_col_dim_mapping >= 0, (
+            f"col_parallel_matmul's row should be divided by a specific mesh axis, but got [{linear2_weight_col_dim_mapping}]"
+        )
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+
+        parallel_axis = linear2_weight_col_dim_mapping
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        # insert op
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+        # setting comm id
+        new_op = main_block.ops[-1]
+        assert new_op.type == "fused_feedforward_grad"
+        new_op._set_attr("ring_id", int(group.id))
+
+
+register_distributed_operator_impl(
+    "fused_feedforward", DistributedFusedFeedForwardImpl("tensor_parallel")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_fused_rms_norm.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import logging
+
+from paddle.base.log_helper import get_logger
+
+from ..completion import get_phi_spmd_rule
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class DistributedLayerNorm(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        x_name = op_desc.input('x')[0]
+        scale_name = op_desc.input('scale')[0]
+        y_name = op_desc.output('y')[0]
+        invvar_name = op_desc.output('invvar')[0]
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        scale_spec = get_dist_tensor_spec(dist_op, scale_name)
+
+        y_spec = get_dist_tensor_spec(dist_op, y_name, is_input=False)
+        invvar_spec = get_dist_tensor_spec(dist_op, invvar_name, is_input=False)
+
+        epsilon = op_desc.attr('epsilon')
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("fused_rms_norm")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec, scale_spec, epsilon)
+        bw_results = rule.infer_backward(
+            x_spec,
+            scale_spec,
+            y_spec,
+            invvar_spec,
+            epsilon,
+        )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            [x_name, scale_name],
+            [y_name, invvar_name],
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+
+        # default impl
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedLayerNorm("fused_rms_norm")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_fused_rope.py

@@ -0,0 +1,189 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ..completion import get_phi_spmd_rule
+from ..dist_attribute import DistTensorSpec
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedFusedRope(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args), build fake spec for optional args
+
+        op_desc = dist_op.serial_op.desc
+        input_parameters = op_desc.input_names()
+        output_parameters = op_desc.output_names()
+        is_input_arg_exist = lambda parameter: (
+            parameter in input_parameters and op_desc.input(parameter)
+        )
+        is_output_arg_exist = lambda parameter: (
+            parameter in output_parameters and op_desc.output(parameter)
+        )
+
+        q = op_desc.input('q')[0]
+        k = op_desc.input('k')[0] if is_input_arg_exist('k') else None
+        v = op_desc.input('v')[0] if is_input_arg_exist('v') else None
+        sin = op_desc.input('sin')[0] if is_input_arg_exist('sin') else None
+        cos = op_desc.input('cos')[0] if is_input_arg_exist('cos') else None
+        position_ids = (
+            op_desc.input('position_ids')[0]
+            if is_input_arg_exist('position_ids')
+            else None
+        )
+        out_q = op_desc.output('out_q')[0]
+        out_k = (
+            op_desc.output('out_k')[0] if is_output_arg_exist('out_k') else None
+        )
+        out_v = (
+            op_desc.output('out_v')[0] if is_output_arg_exist('out_v') else None
+        )
+
+        q_spec = get_dist_tensor_spec(dist_op, q)
+        k_spec = (
+            get_dist_tensor_spec(dist_op, k)
+            if k is not None
+            else DistTensorSpec()
+        )
+        v_spec = (
+            get_dist_tensor_spec(dist_op, v)
+            if v is not None
+            else DistTensorSpec()
+        )
+        sin_spec = (
+            get_dist_tensor_spec(dist_op, sin)
+            if sin is not None
+            else DistTensorSpec()
+        )
+        cos_spec = (
+            get_dist_tensor_spec(dist_op, cos)
+            if cos is not None
+            else DistTensorSpec()
+        )
+        position_ids_spec = (
+            get_dist_tensor_spec(dist_op, position_ids)
+            if position_ids is not None
+            else DistTensorSpec()
+        )
+        out_q_spec = get_dist_tensor_spec(dist_op, out_q, is_input=False)
+        out_k_spec = (
+            get_dist_tensor_spec(dist_op, out_k, is_input=False)
+            if out_k is not None
+            else DistTensorSpec()
+        )
+        out_v_spec = (
+            get_dist_tensor_spec(dist_op, out_v, is_input=False)
+            if out_v is not None
+            else DistTensorSpec()
+        )
+
+        use_neox_rotary_style = op_desc.attr("use_neox_rotary_style")
+        time_major = op_desc.attr("time_major")
+        rotary_emb_base = op_desc.attr("rotary_emb_base")
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("fused_rotary_position_embedding")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(
+            q_spec,
+            k_spec,
+            v_spec,
+            sin_spec,
+            cos_spec,
+            position_ids_spec,
+            use_neox_rotary_style,
+            time_major,
+            rotary_emb_base,
+        )
+        bw_results = rule.infer_backward(
+            q_spec,
+            k_spec,
+            v_spec,
+            sin_spec,
+            cos_spec,
+            position_ids_spec,
+            out_q_spec,
+            out_k_spec,
+            out_v_spec,
+            use_neox_rotary_style,
+            time_major,
+            rotary_emb_base,
+        )
+
+        # remove optional args in spmd results
+        input_args = [q, k, v, sin, cos, position_ids]
+        output_args = [out_q, out_k, out_v]
+        fw_and_bw_results_without_optional_arg = []
+        for results in [fw_results, bw_results]:
+            input_results = results[0]
+            output_results = results[1]
+            input_results_without_optional_arg = []
+            output_results_without_optional_arg = []
+            for idx, input_arg in enumerate(input_args):
+                if input_arg is not None:
+                    input_results_without_optional_arg.append(
+                        input_results[idx]
+                    )
+            for idx, output_arg in enumerate(output_args):
+                if output_arg is not None:
+                    output_results_without_optional_arg.append(
+                        output_results[idx]
+                    )
+            fw_and_bw_results_without_optional_arg.append(
+                [
+                    input_results_without_optional_arg,
+                    output_results_without_optional_arg,
+                ]
+            )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            input_arg_names=[
+                input_arg for input_arg in input_args if input_arg is not None
+            ],
+            output_arg_names=[
+                output_arg
+                for output_arg in output_args
+                if output_arg is not None
+            ],
+            fw_results=fw_and_bw_results_without_optional_arg[0],
+            bw_results=fw_and_bw_results_without_optional_arg[1],
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedFusedRope("fused_rotary_position_embedding")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_gather_nd.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ..completion import get_phi_spmd_rule
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedGatherNd(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        x_name = op_desc.input('X')[0]
+        index_name = op_desc.input('Index')[0]
+        out_name = op_desc.output('Out')[0]
+
+        x_specs = get_dist_tensor_spec(dist_op, x_name)
+        index_specs = get_dist_tensor_spec(dist_op, index_name)
+        output_spec = get_dist_tensor_spec(dist_op, out_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("gather_nd")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_specs, index_specs)
+        bw_results = rule.infer_backward(x_specs, index_specs, output_spec)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            [x_name, index_name],
+            [out_name],
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedGatherNd("gather_nd"))

python/paddle/distributed/auto_parallel/static/operators/dist_layer_norm.py

@@ -0,0 +1,151 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import copy
+import logging
+
+from paddle.base.log_helper import get_logger
+
+from ..completion import get_phi_spmd_rule
+from ..dist_attribute import DistTensorSpec, TensorDistAttr
+from ..utils import get_dist_tensor_spec, is_dim_shard
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class DistributedLayerNorm(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        x_name = op_desc.input('X')[0]
+        scale_name = (
+            op_desc.input('Scale')[0]
+            if len(op_desc.input('Scale')) > 0
+            else None
+        )
+        bias_name = (
+            op_desc.input('Bias')[0] if len(op_desc.input('Bias')) > 0 else None
+        )
+        y_name = op_desc.output('Y')[0]
+        var_name = op_desc.output('Variance')[0]
+        mean_name = op_desc.output('Mean')[0]
+        begin_norm_axis = op_desc.attr('begin_norm_axis')
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        scale_spec = (
+            DistTensorSpec([0], TensorDistAttr())
+            if scale_name is None
+            else get_dist_tensor_spec(dist_op, scale_name)
+        )
+        bias_spec = (
+            DistTensorSpec([0], TensorDistAttr())
+            if bias_name is None
+            else get_dist_tensor_spec(dist_op, bias_name)
+        )
+        y_spec = get_dist_tensor_spec(dist_op, y_name, False)
+        var_spec = get_dist_tensor_spec(dist_op, var_name, False)
+        mean_spec = get_dist_tensor_spec(dist_op, mean_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("layer_norm")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(
+            x_spec, scale_spec, bias_spec, 1.0, begin_norm_axis
+        )
+        bw_results = rule.infer_backward(
+            x_spec,
+            scale_spec,
+            bias_spec,
+            y_spec,
+            var_spec,
+            mean_spec,
+            1.0,
+            begin_norm_axis,
+        )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        input_arg_names = [x_name]
+        if scale_name is not None:
+            input_arg_names.append(scale_name)
+        if bias_name is not None:
+            input_arg_names.append(bias_name)
+        changed = update_op_dims_mapping(
+            dist_op,
+            input_arg_names,
+            [y_name, var_name, mean_name],
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        begin_norm_axis = op_desc.attr('begin_norm_axis')
+
+        # sharded on begin_norm_axis
+        x_name = op_desc.input('X')[0]
+        x_dims_mapping = copy.deepcopy(
+            op_dist_attr.get_input_dims_mapping(x_name)
+        )
+        if (begin_norm_axis > 0) and is_dim_shard(
+            x_dims_mapping[begin_norm_axis]
+        ):
+            # TODO (ljz) support sharding on `begin_norm_axis`
+            _logger.info(
+                "sharding on `begin_norm_axis` is not supported yet, we resharded it as replicated"
+            )
+            x_dims_mapping[begin_norm_axis] = -1
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+
+            param_names = [op_desc.input('Scale')[0], op_desc.input('Bias')[0]]
+            for p_name in param_names:
+                p_dims_mapping = copy.deepcopy(
+                    op_dist_attr.get_input_dims_mapping(p_name)
+                )
+                p_dims_mapping[begin_norm_axis] = -1
+                op_dist_attr.set_input_dims_mapping(p_name, p_dims_mapping)
+
+            y_name = op_desc.output('Y')[0]
+            y_dims_mapping = copy.deepcopy(
+                op_dist_attr.get_output_dims_mapping(y_name)
+            )
+            y_dims_mapping[begin_norm_axis] = -1
+            op_dist_attr.set_input_dims_mapping(y_name, y_dims_mapping)
+
+        # default impl
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedLayerNorm("layer_norm"))

python/paddle/distributed/auto_parallel/static/operators/dist_pnorm.py

@@ -0,0 +1,387 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from paddle.common_ops_import import check_dtype, check_variable_and_dtype
+from paddle.distributed.utils.stream_utils import ExecutionStreamType
+from paddle.framework import core
+from paddle.static import Operator
+
+from ..dist_attribute import OperatorDistAttr, TensorDistAttr
+from ..process_group import new_process_group
+from ..utils import (
+    _get_comm_group,
+    _get_corresponding_rank,
+    compute_compatible_dim_mapping,
+    is_dim_replicate,
+    is_dim_shard,
+    set_dist_op_desc_original_id,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+
+
+class DistributedPNorm(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(DistributedPNorm("p_norm"))
+
+
+# Data Parallel
+class DistributedPNormImpl0(DistributedOperatorImpl):
+    """
+    TODO: p_norm scene
+
+    1. axis == None, isinstance(p, (int, float)), asvector = True
+        1.1 x_dims_mapping == [0, -1, -1]
+            allgather input if it is split by dp group
+        1.2 x_dims_mapping == [-1, 0, -1]
+            allgather, split and concat input if it is split by mp group
+    2. isinstance(axis, int), asvector = False
+        1.1 axis == 0 and x_dims_mapping == [0, -1, -1]
+            allgather input if it's input[0] is splited by dp group.
+        1.2 axis == 1 and x_dims_mapping == [-1, 0, -1]
+            allgather, split and concat input if it's input[1] is split by mp group
+    """
+
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        axis = op_desc.attr('axis')
+        asvector = op_desc.attr('asvector')
+        x_name = op_desc.input('X')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        if is_dim_replicate(x_dims_mapping[0]):
+            return False
+        # Other dimensions must be replicate except the batch dimension
+        for mapping in x_dims_mapping[1:]:
+            if is_dim_shard(mapping):
+                return False
+        if not (axis == -1 and asvector) and not (axis == 0 and not asvector):
+            return False
+        return True
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (
+            (not self.is_input_compatible(dist_op))
+            or (not self.is_output_compatible(dist_op))
+            or (not self.is_compatible(dist_op))
+        ):
+            return False
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        axis = op_desc.attr('axis')
+        keepdim = op_desc.attr('keepdim')
+
+        batch_dim_mappings = []
+        for arg_name in op_desc.input_arg_names():
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if len(dims_mapping) >= 1:
+                batch_dim_mappings.append(dims_mapping[0])
+        for arg_name in op_desc.output_arg_names():
+            dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+            if len(dims_mapping) >= 1:
+                batch_dim_mappings.append(dims_mapping[0])
+
+        compatible_dim_mapping = compute_compatible_dim_mapping(
+            batch_dim_mappings
+        )
+        if compatible_dim_mapping is None:
+            return False
+
+        for arg_name in op_desc.input_arg_names():
+            dims_mapping = op_dist_attr.get_input_dims_mapping(arg_name)
+            if (
+                len(dims_mapping) >= 1
+                and compatible_dim_mapping != dims_mapping[0]
+            ):
+                dims_mapping[0] = compatible_dim_mapping
+                op_dist_attr.set_input_dims_mapping(arg_name, dims_mapping)
+                changed = True
+
+        if axis == 0 and not keepdim:
+            for arg_name in op_desc.output_arg_names():
+                dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+                if len(dims_mapping) >= 1 and dims_mapping[0] != -1:
+                    dims_mapping[0] = -1
+                    op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
+                    changed = True
+        else:
+            for arg_name in op_desc.output_arg_names():
+                dims_mapping = op_dist_attr.get_output_dims_mapping(arg_name)
+                if (
+                    len(dims_mapping) >= 1
+                    and compatible_dim_mapping != dims_mapping[0]
+                ):
+                    dims_mapping[0] = compatible_dim_mapping
+                    op_dist_attr.set_output_dims_mapping(arg_name, dims_mapping)
+                    changed = True
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None
+
+        # check validation of inputs / outputs
+        for input_name in src_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                src_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in src_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                src_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        if rank_id not in op_dist_attr.process_mesh.process_ids:
+            rank_id = _get_corresponding_rank(
+                ctx, op_dist_attr.process_mesh, rank_id
+            )
+
+        X_var = main_block._var_recursive(kwargs['X'][0])
+        in_dims_mapping = op_dist_attr.get_input_dims_mapping(X_var.name)
+        for axis in range(len(in_dims_mapping)):
+            if in_dims_mapping[axis] != -1:
+                break
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+        group_ranks = _get_comm_group(
+            process_mesh_group, process_mesh_shape, axis, rank_id
+        )
+        group = new_process_group(group_ranks)
+
+        check_variable_and_dtype(
+            X_var, 'x', ['float16', 'float32', 'float64'], 'norm'
+        )
+        check_dtype(
+            X_var.dtype, 'dtype', ['float16', 'float32', 'float64'], 'norm'
+        )
+
+        # 2. insert all_gather op
+        # create all_gather output var
+        allgather_out = main_block.create_var(
+            name=".".join(["all_gather", X_var.name]),
+            dtype=X_var.dtype,
+            shape=X_var.shape,
+            type=core.VarDesc.VarType.DENSE_TENSOR,
+            persistable=False,
+            stop_gradient=X_var.stop_gradient,
+        )
+        # set allgather_out tensor dist_attr
+        allgather_out_dist_attr = TensorDistAttr()
+        allgather_out_dist_attr.process_mesh = op_dist_attr.process_mesh
+        allgather_out_dist_attr.chunk_id = op_dist_attr.chunk_id
+        allgather_out_dist_attr.dims_mapping = [
+            -1 for i in range(len(allgather_out.shape))
+        ]
+        ctx.set_tensor_dist_attr_for_program(
+            allgather_out, allgather_out_dist_attr
+        )
+        all_gather_op = main_block.append_op(
+            type='all_gather',
+            inputs={'x': [X_var]},
+            outputs={'out': [allgather_out]},
+            attrs={
+                'ring_id': group.id,
+                'use_calc_stream': True,
+                'nranks': group.nranks,
+                'op_role': src_op.attr('op_role'),
+            },
+        )
+        # set all_gather op dist_attr
+        allgather_op_dist_attr = OperatorDistAttr()
+        allgather_op_dist_attr.process_mesh = op_dist_attr.process_mesh
+        allgather_op_dist_attr.chunk_id = op_dist_attr.chunk_id
+        allgather_op_dist_attr.set_input_dims_mapping(
+            X_var.name, in_dims_mapping
+        )
+        allgather_op_dist_attr.set_output_dims_mapping(
+            allgather_out.name, allgather_out_dist_attr.dims_mapping
+        )
+        allgather_op_dist_attr.execution_stream = (
+            ExecutionStreamType.DefaultStream.value
+        )
+        ctx.set_op_dist_attr_for_program(all_gather_op, allgather_op_dist_attr)
+
+        # 3. copy p_norm op desc and reset input name
+        # rename input
+        kwargs['X'] = [allgather_out.name]
+        # replicate op in dist program
+        dist_op = main_block.append_op(type='nop')
+        dist_op_desc = dist_op.desc
+        dist_op_desc.copy_from(src_op.desc)
+        set_dist_op_desc_original_id(dist_op_desc, src_op.desc, ctx)
+        for input_name in src_op.desc.input_names():
+            dist_op_desc.set_input(input_name, kwargs[input_name])
+        for output_name in src_op.desc.output_names():
+            dist_op_desc.set_output(output_name, kwargs[output_name])
+        pnorm_op = Operator(main_block, dist_op_desc)
+        op_dist_attr.set_input_dims_mapping(
+            allgather_out.name, allgather_out_dist_attr.dims_mapping
+        )
+        # Remove the unrelated dist attr
+        op_dist_attr.del_input_dist_attr(X_var.name)
+        ctx.set_op_dist_attr_for_program(pnorm_op, op_dist_attr)
+        # TODO: should we add a new dist attr for the new op here?
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        backward_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
+        assert op_dist_attr is not None
+
+        # check validation of inputs / outputs
+        for input_name in backward_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                backward_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in backward_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                backward_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        X_var = main_block._var_recursive(kwargs['X'][0])
+        X_grad_var = main_block._var_recursive(kwargs['X@GRAD'][0])
+
+        # 1. copy p_norm_grad op and reset input name and output name
+        new_kwargs = copy.deepcopy(kwargs)
+        new_kwargs['X'] = [".".join(["all_gather", X_var.name])]
+        new_X_var = main_block._var_recursive(new_kwargs['X'][0])
+        new_X_grad = main_block.create_var(
+            name=".".join(["all_gather", X_grad_var.name]),
+            dtype=X_grad_var.dtype,
+            shape=new_X_var.shape,
+            type=core.VarDesc.VarType.DENSE_TENSOR,
+            persistable=False,
+            stop_gradient=X_grad_var.stop_gradient,
+        )
+        new_kwargs['X@GRAD'] = [new_X_grad.name]
+        new_X_var_dist_attr = ctx.get_tensor_dist_attr_for_program(new_X_var)
+        ctx.set_tensor_dist_attr_for_program(new_X_grad, new_X_var_dist_attr)
+        # replicate op in dist program with new kwargs
+        dist_op = main_block.append_op(type='nop')
+        dist_op_desc = dist_op.desc
+        dist_op_desc.copy_from(backward_op.desc)
+        # Refer to the related dist op
+        set_dist_op_desc_original_id(dist_op_desc, backward_op.desc, ctx)
+        for input_name in backward_op.desc.input_names():
+            dist_op_desc.set_input(input_name, new_kwargs[input_name])
+        for output_name in backward_op.desc.output_names():
+            dist_op_desc.set_output(output_name, new_kwargs[output_name])
+        p_norm_grad_op = Operator(main_block, dist_op_desc)
+        op_dist_attr.set_input_dims_mapping(
+            new_X_var.name, new_X_var_dist_attr.dims_mapping
+        )
+        # Store X_grad_var dims_mapping for later use
+        X_grad_var_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            X_grad_var.name
+        )
+        # Remove the unrelated dist attr
+        op_dist_attr.del_input_dist_attr(X_var.name)
+        op_dist_attr.set_output_dims_mapping(
+            new_X_grad.name, new_X_var_dist_attr.dims_mapping
+        )
+        # Remove the unrelated dist attr
+        op_dist_attr.del_output_dist_attr(X_grad_var.name)
+        ctx.set_op_dist_attr_for_program(p_norm_grad_op, op_dist_attr)
+        # TODO: should we add a new dist attr for the new op here?
+
+        # 2. insert slice op
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+        process_mesh_group = op_dist_attr.process_mesh.process_ids
+        dims_mapping = [0] + [-1 for _ in range(len(new_X_grad.shape) - 1)]
+        from ..reshard import Resharder
+
+        partition_idx = Resharder.compute_partition_index(
+            rank_id,
+            new_X_grad.shape,
+            dims_mapping,
+            process_mesh_shape,
+            process_mesh_group,
+        )
+        slice_starts = []
+        slice_ends = []
+        slices_axes = []
+        for idx, item in enumerate(partition_idx):
+            slice_starts.append(item[0])
+            slice_ends.append(item[1])
+            slices_axes.append(idx)
+
+        infer_flags = [1 for i in range(len(slices_axes))]
+        attrs = {
+            "axes": slices_axes,
+            "starts": slice_starts,
+            "ends": slice_ends,
+            "infer_flags": infer_flags,
+            "op_role": backward_op.attr('op_role'),
+        }
+        slice_op = main_block.append_op(
+            type='slice',
+            inputs={'Input': [new_X_grad]},
+            outputs={'Out': [X_grad_var]},
+            attrs=attrs,
+        )
+        slice_op_dist_attr = OperatorDistAttr()
+        slice_op_dist_attr.process_mesh = op_dist_attr.process_mesh
+        slice_op_dist_attr.chunk_id = op_dist_attr.chunk_id
+        slice_op_dist_attr.set_input_dims_mapping(
+            new_X_grad.name, new_X_var_dist_attr.dims_mapping
+        )
+        slice_op_dist_attr.set_output_dims_mapping(
+            X_grad_var.name, X_grad_var_dims_mapping
+        )
+        ctx.set_op_dist_attr_for_program(slice_op, slice_op_dist_attr)
+
+
+register_distributed_operator_impl(
+    "p_norm", DistributedPNormImpl0("data_parallel")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_reduce_sum_p.py

@@ -0,0 +1,240 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+import copy
+
+import paddle
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+
+from ..completion import get_phi_spmd_rule
+from ..dist_attribute import OperatorDistAttr
+from ..process_group import new_process_group
+from ..utils import (
+    get_dist_tensor_spec,
+    is_dim_shard,
+    set_dist_op_desc_original_id,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedReduceSum(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+
+        op_desc = dist_op.serial_op.desc
+        assert len(op_desc.input_arg_names()) == 1, (
+            f"reduce_sum op [{op_desc.type}] has [{len(op_desc.input_arg_names())}] inputs"
+        )
+        input_arg_name = op_desc.input_arg_names()[0]
+        assert len(op_desc.output_arg_names()) == 1, (
+            f"reduce_sum op [{op_desc.type}] has [{len(op_desc.output_arg_names())}] outputs"
+        )
+        output_arg_name = op_desc.output_arg_names()[0]
+        keep_dim = op_desc.attr('keep_dim')
+        dims = op_desc.attr('dim')
+
+        # TODO (zhangyichen) replace dist tensor spec by dist tensor in future.
+        input_spec = get_dist_tensor_spec(dist_op, input_arg_name)
+        output_spec = get_dist_tensor_spec(dist_op, output_arg_name, False)
+        # len(dims) == 0 means reduce_all
+        if len(dims) == 0:
+            dims = list(range(len(input_spec.shape)))
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("reduce_sum")
+        fw_results = rule.infer_forward(input_spec, dims, keep_dim)
+        bw_results = rule.infer_backward(
+            input_spec, output_spec, dims, keep_dim
+        )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, [input_arg_name], [output_arg_name], fw_results, bw_results
+        )
+
+        return changed
+
+    # NOTE this function will be remove once we use local reshard to replace distopimpls
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        op_desc = dist_op.serial_op.desc
+        input_name = op_desc.input_arg_names()[0]
+        input_dims_mapping = copy.deepcopy(
+            op_dist_attr.get_input_dims_mapping(input_name)
+        )
+        axes = op_desc.attr('dim')
+
+        op_dist_attr = dist_op.dist_attr
+        reverted = False
+
+        def is_partial_reduce(axes, dims_mapping):
+            # FIXME(ljz) Hack for performance:
+            # if the reduce result is a scalar, it is the loss reduce in GPT case,
+            # and if any axis of reduce input is sharded, the result loss would be partial.
+            # BUT we keep the loss as partial instead of allreduce it for performance, since it would effect the backward.
+            # we should use an optimization pass for the Hack in future.
+            if len(axes) != 0 and (len(axes) < len(dims_mapping)):
+                for axis in axes:
+                    if is_dim_shard(dims_mapping[axis]):
+                        return True  # reverted
+            return False
+
+        # if reduce_axis is sharded, the output is partial and need to be allreduce
+        if is_partial_reduce(axes, input_dims_mapping):
+            # TODO (ljz) support reduce where the reduce_axis is sharded
+            dist_op.dist_attr = original_op_dist_attr
+            reverted = True
+        # if reduce_axis is unsharded, NO extra operator need.
+        else:
+            default_impl = get_default_distributed_operator_impl()
+            op_dist_attr.impl_type = default_impl.type
+            op_dist_attr.impl_idx = default_impl.idx
+        return reverted
+
+
+register_distributed_operator_impl_container(DistributedReduceSum("reduce_sum"))
+
+
+class DistributedReduceSumPrimitive(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(
+    DistributedReduceSumPrimitive("reduce_sum_p")
+)
+
+
+# Batch Dimension ReduceSum Primitive
+class DistributedReduceSumPrimitiveImpl0(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+
+        return len(op_desc.input_arg_names()) == 1
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        outputs = op_desc.output_arg_names()
+
+        if len(outputs) != 1:
+            return False
+
+        output_name = outputs[0]
+        output_var = dist_op.serial_op.block._var_recursive(output_name)
+        if output_var.shape != ():
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+
+        return self.is_input_compatible(dist_op) and self.is_output_compatible(
+            dist_op
+        )
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        startup_block = dist_op_context.startup_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+
+        # check validation of inputs / outputs
+        for input_name in src_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                src_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in src_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                src_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        # replicate op in dist program
+        dist_op = main_block.append_op(type='nop')
+        dist_op_desc = dist_op.desc
+        dist_op_desc.copy_from(src_op.desc)
+        set_dist_op_desc_original_id(dist_op_desc, src_op.desc, ctx)
+        for input_name in src_op.desc.input_names():
+            dist_op_desc.set_input(input_name, kwargs[input_name])
+        for output_name in src_op.desc.output_names():
+            dist_op_desc.set_output(output_name, kwargs[output_name])
+        # TODO: should we add a new dist attr for the new op here?
+
+        # batch dimension synchronization
+        var_name = src_op.output_arg_names[0]
+        sync_group = new_process_group(ctx.data_parallel_group)
+        allreduce_op = main_block.append_op(
+            type='all_reduce',
+            inputs={'x': [var_name]},
+            outputs={'out': [var_name]},
+            attrs={
+                'ring_id': sync_group.id,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
+                OP_ROLE_KEY: OpRole.Forward,
+            },
+        )
+
+        # dist attr
+        var = main_block._var_recursive(var_name)
+        tensor_dist_attr = ctx.get_tensor_dist_attr_for_program(var)
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        new_op_attr = OperatorDistAttr()
+        new_op_attr.process_mesh = op_dist_attr.process_mesh
+        new_op_attr.set_output_dims_mapping(
+            var.name, tensor_dist_attr.dims_mapping
+        )
+        new_op_attr.set_input_dims_mapping(
+            var.name, tensor_dist_attr.dims_mapping
+        )
+        ctx.set_op_dist_attr_for_program(allreduce_op, new_op_attr)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        raise RuntimeError("primitive operator does NOT have backward function")
+
+
+register_distributed_operator_impl(
+    "reduce_sum_p",
+    DistributedReduceSumPrimitiveImpl0("batch_dimension_reduce_sum_p"),
+)

python/paddle/distributed/auto_parallel/static/operators/dist_reshape.py

@@ -0,0 +1,866 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..completion import get_phi_spmd_rule
+from ..cost import (
+    Reshape2GradOpCost,
+    Reshape2OpCost,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_dp_costs,
+)
+from ..utils import (
+    compute_compatible_and_update_dim_mapping,
+    get_dist_tensor_spec,
+    is_dim_shard,
+    set_dist_op_desc_original_id,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    is_parameter_related,
+    merge_forward_backward_dims_mapping,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedReshape2(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+        assert dist_op.serial_op.type == "reshape2", (
+            f"{dist_op.serial_op.type} is not supported by dist reshape yet."
+        )
+
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        xshape_name = op_desc.output('XShape')[0]
+        shape = op_desc.attr('shape')
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        output_spec = get_dist_tensor_spec(dist_op, out_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("reshape")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec, shape)
+        bw_results = rule.infer_backward(x_spec, output_spec, shape)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, [x_name], [out_name], fw_results, bw_results
+        )
+
+        # step4: update xshape
+        inferred_input_dims_mappings, _ = merge_forward_backward_dims_mapping(
+            fw_results, bw_results
+        )
+        dist_op.dist_attr.set_output_dims_mapping(
+            xshape_name, [-1] + inferred_input_dims_mappings[0]
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        reverted = False
+        op_dist_attr = dist_op.dist_attr
+
+        # all reshape mapping to impl0
+        op_dist_attr.impl_type = "reshape2"
+        op_dist_attr.impl_idx = 0
+
+        return reverted
+
+
+register_distributed_operator_impl_container(DistributedReshape2("reshape2"))
+
+
+class DistributedReshapeImpl0(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = False
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        res = []
+        op = dist_op.serial_op
+        dist_attr = dist_op.dist_attr
+
+        shape_list = op.desc.attr("shape")
+        # got dist attribute info
+        dim_mapping = dist_attr.get_output_dims_mapping(op.output("Out")[0])
+        process_mesh_shape = dist_attr.process_mesh.shape
+
+        # modify target shape
+        for idx, axis in enumerate(dim_mapping):
+            if axis >= 0:
+                if len(shape_list) > idx:
+                    shape_list[idx] = (
+                        shape_list[idx] // process_mesh_shape[axis]
+                    )
+
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_attr.process_mesh.process_ids
+        for key in desc_mapping:
+            desc_mapping[key]["shape"] = shape_list
+
+        cost_mapping = build_comp_costs_from_descs(
+            Reshape2OpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        return res
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+
+        cost_mapping = build_comp_costs_from_descs(
+            Reshape2GradOpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        backward_op = dist_op.serial_op
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and is_parameter_related(
+                    varname, main_block
+                ):
+                    # NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+
+        return res
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        if len(x_dims_mapping) != len(out_dims_mapping) - 1:
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        if len(x_dims_mapping) != len(out_dims_mapping) - 1:
+            return False
+
+        if is_dim_shard(out_dims_mapping[-1]):
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        for idx, dim_mapping in enumerate(out_dims_mapping[:-1]):
+            if x_dims_mapping[idx] != dim_mapping:
+                return False
+
+        if x_shape_dims_mapping[0] != -1:
+            return False
+
+        if x_shape_dims_mapping[1:] != x_dims_mapping[:]:
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+
+        for i in range(len(x_dims_mapping)):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [x_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                changed = True
+
+        for i in range(len(x_dims_mapping)):
+            x_shape_dims_mapping[i + 1] = x_dims_mapping[i]
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(
+                x_shape_name, x_shape_dims_mapping
+            )
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        """
+        kwargs: inputname_mapping & outputname_mapping
+        """
+
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None, (
+            f"backward op [{src_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        for input_name in src_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                src_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in src_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                src_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        X_var = main_block._var_recursive(kwargs['X'][0])
+        Out_var = main_block._var_recursive(kwargs['Out'][0])
+        XShape_var = main_block._var_recursive(kwargs['XShape'][0])
+        shape_list = src_op.desc.attr("shape")
+        ShapeTensor_var_list = []
+        for name in kwargs['ShapeTensor']:
+            ShapeTensor_var_list.append(name)
+        Shape_var_list = []
+        for name in kwargs['Shape']:
+            Shape_var_list.append(name)
+
+        # got dist attribute info
+        dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+
+        # modify target shape
+        for idx, axis in enumerate(dim_mapping):
+            if axis >= 0:
+                if len(shape_list) > idx:
+                    shape_list[idx] = (
+                        shape_list[idx] // process_mesh_shape[axis]
+                    )
+
+        # create op
+        new_op = main_block.append_op(type='nop')
+        new_op_desc = new_op.desc
+        new_op_desc.copy_from(src_op.desc)
+        set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
+        new_op_desc.set_input('ShapeTensor', ShapeTensor_var_list)
+        new_op_desc.set_input('Shape', Shape_var_list)
+        new_op_desc.set_input('X', [X_var.name])
+        new_op_desc.set_output('XShape', [XShape_var.name])
+        new_op_desc.set_output('Out', [Out_var.name])
+        new_op_desc._set_attr('shape', shape_list)
+        # TODO: should we add a new dist attr for the new op here?
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+class DistributedReshapeImpl1(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = False
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        res = []
+        op = dist_op.serial_op
+        dist_attr = dist_op.dist_attr
+
+        shape_list = op.desc.attr("shape")
+        # got dist attribute info
+        dim_mapping = dist_attr.get_output_dims_mapping(op.output("Out")[0])
+        process_mesh_shape = dist_attr.process_mesh.shape
+
+        # modify target shape
+        for idx, axis in enumerate(dim_mapping):
+            if axis >= 0:
+                if len(shape_list) > idx:
+                    shape_list[idx] = (
+                        shape_list[idx] // process_mesh_shape[axis]
+                    )
+
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_attr.process_mesh.process_ids
+        for key in desc_mapping:
+            desc_mapping[key]["shape"] = shape_list
+
+        cost_mapping = build_comp_costs_from_descs(
+            Reshape2OpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        return res
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+
+        cost_mapping = build_comp_costs_from_descs(
+            Reshape2GradOpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        backward_op = dist_op.serial_op
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and not is_parameter_related(
+                    varname, main_block
+                ):
+                    # NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+
+        return res
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        if len(x_dims_mapping) != len(out_dims_mapping) + 1:
+            return False
+
+        if is_dim_shard(x_dims_mapping[-1]):
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        if len(x_dims_mapping) != len(out_dims_mapping) + 1:
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+
+        if is_dim_shard(x_dims_mapping[-1]):
+            return False
+
+        for idx, item in enumerate(x_dims_mapping[:-1]):
+            if out_dims_mapping[idx] != item:
+                return False
+
+        if x_shape_dims_mapping[0] != -1:
+            return False
+
+        if x_shape_dims_mapping[1:] != x_dims_mapping[:]:
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+
+        for i in range(len(out_dims_mapping)):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [x_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                changed = True
+
+        for i in range(len(x_dims_mapping)):
+            x_shape_dims_mapping[i + 1] = x_dims_mapping[i]
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(
+                x_shape_name, x_shape_dims_mapping
+            )
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        """
+        kwargs: inputname_mapping & outputname_mapping
+        """
+
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        src_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None, (
+            f"backward op [{src_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        for input_name in src_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                src_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in src_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                src_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        X_var = main_block._var_recursive(kwargs['X'][0])
+        Out_var = main_block._var_recursive(kwargs['Out'][0])
+        XShape_var = main_block._var_recursive(kwargs['XShape'][0])
+        shape_list = src_op.desc.attr("shape")
+        ShapeTensor_var_list = []
+        for name in kwargs['ShapeTensor']:
+            ShapeTensor_var_list.append(name)
+        Shape_var_list = []
+        for name in kwargs['Shape']:
+            Shape_var_list.append(name)
+
+        # got dist attribute info
+        dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+
+        # modify target shape
+        for idx, axis in enumerate(dim_mapping):
+            if axis >= 0:
+                if len(shape_list) > idx:
+                    shape_list[idx] = (
+                        shape_list[idx] // process_mesh_shape[axis]
+                    )
+
+        # create op
+        new_op = main_block.append_op(type='nop')
+        new_op_desc = new_op.desc
+        new_op_desc.copy_from(src_op.desc)
+        set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
+        new_op_desc.set_input('ShapeTensor', ShapeTensor_var_list)
+        new_op_desc.set_input('Shape', Shape_var_list)
+        new_op_desc.set_input('X', [X_var.name])
+        new_op_desc.set_output('XShape', [XShape_var.name])
+        new_op_desc.set_output('Out', [Out_var.name])
+        new_op_desc._set_attr('shape', shape_list)
+        # TODO: should we add a new dist attr for the new op here?
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+class DistributedReshapeImpl2(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = False
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        res = []
+        op = dist_op.serial_op
+        dist_attr = dist_op.dist_attr
+
+        shape_list = op.desc.attr("shape")
+        # got dist attribute info
+        dim_mapping = dist_attr.get_output_dims_mapping(op.output("Out")[0])
+        process_mesh_shape = dist_attr.process_mesh.shape
+
+        # modify target shape
+        for idx, axis in enumerate(dim_mapping):
+            if axis >= 0:
+                if len(shape_list) > idx:
+                    shape_list[idx] = (
+                        shape_list[idx] // process_mesh_shape[axis]
+                    )
+
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_attr.process_mesh.process_ids
+        for key in desc_mapping:
+            desc_mapping[key]["shape"] = shape_list
+
+        cost_mapping = build_comp_costs_from_descs(
+            Reshape2OpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        return res
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+
+        cost_mapping = build_comp_costs_from_descs(
+            Reshape2GradOpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        backward_op = dist_op.serial_op
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and not is_parameter_related(
+                    varname, main_block
+                ):
+                    # NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+
+        return res
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        if len(x_dims_mapping) != len(out_dims_mapping):
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_name = op_desc.output('Out')[0]
+        x_name = op_desc.input('X')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        if len(x_dims_mapping) != len(out_dims_mapping):
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+
+        for idx, item in enumerate(x_dims_mapping[:-1]):
+            if out_dims_mapping[idx] != item:
+                return False
+
+        if x_shape_dims_mapping[0] != -1:
+            return False
+
+        if x_shape_dims_mapping[1:] != out_dims_mapping[:]:
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+
+        for i in range(len(out_dims_mapping) - 1):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [x_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                changed = True
+
+        for i in range(len(out_dims_mapping)):
+            x_shape_dims_mapping[i + 1] = out_dims_mapping[i]
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(
+                x_shape_name, x_shape_dims_mapping
+            )
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        """
+        kwargs: inputname_mapping & outputname_mapping
+        """
+
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.work_block
+        src_op = dist_op_context.cur_src_op
+        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
+        assert op_dist_attr is not None, (
+            f"backward op [{src_op}] don't have dist attribute !"
+        )
+
+        # check validation of inputs / outputs
+        for input_name in src_op.desc.input_names():
+            assert input_name in kwargs, f"input [{input_name}] is not given"
+            assert len(kwargs[input_name]) == len(
+                src_op.desc.input(input_name)
+            ), f"number of tensor for input [{input_name}] is not match"
+        for output_name in src_op.desc.output_names():
+            assert output_name in kwargs, f"input [{output_name}] is not given"
+            assert len(kwargs[output_name]) == len(
+                src_op.desc.output(output_name)
+            ), f"number of tensor for input [{output_name}] is not match"
+
+        X_var = main_block._var_recursive(kwargs['X'][0])
+        Out_var = main_block._var_recursive(kwargs['Out'][0])
+        XShape_var = main_block._var_recursive(kwargs['XShape'][0])
+        shape_list = src_op.desc.attr("shape")
+        ShapeTensor_var_list = []
+        for name in kwargs['ShapeTensor']:
+            ShapeTensor_var_list.append(name)
+        Shape_var_list = []
+        for name in kwargs['Shape']:
+            Shape_var_list.append(name)
+
+        # got dist attribute info
+        out_dim_mapping = op_dist_attr.get_output_dims_mapping(Out_var.name)
+        process_mesh_shape = op_dist_attr.process_mesh.shape
+
+        # modify target shape
+        for idx, axis in enumerate(out_dim_mapping):
+            if axis >= 0:
+                if len(shape_list) > idx:
+                    shape_list[idx] = (
+                        shape_list[idx] // process_mesh_shape[axis]
+                    )
+
+        # create op
+        new_op = main_block.append_op(type='nop')
+        new_op_desc = new_op.desc
+        new_op_desc.copy_from(src_op.desc)
+        set_dist_op_desc_original_id(new_op_desc, src_op.desc, ctx)
+        new_op_desc.set_input('ShapeTensor', ShapeTensor_var_list)
+        new_op_desc.set_input('Shape', Shape_var_list)
+        new_op_desc.set_input('X', [X_var.name])
+        new_op_desc.set_output('XShape', [XShape_var.name])
+        new_op_desc.set_output('Out', [Out_var.name])
+        new_op_desc._set_attr('shape', shape_list)
+        # TODO: should we add a new dist attr for the new op here?
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "reshape2", DistributedReshapeImpl0("add_one_dim_back")
+)
+register_distributed_operator_impl(
+    "reshape2", DistributedReshapeImpl1("remove_one_dim_back")
+)
+register_distributed_operator_impl(
+    "reshape2", DistributedReshapeImpl2("same_dim_shape")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_scale.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..cost import (
+    _g_op_cost_factory,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_dp_costs,
+)
+from ..utils import compute_compatible_and_update_dim_mapping
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    is_parameter_related,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedScale(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+# TODO remove assign dist op
+# register_distributed_operator_impl_container(DistributedScale("scale"))
+# register_distributed_operator_impl_container(DistributedScale("fill_any_like"))
+# register_distributed_operator_impl_container(DistributedScale("where"))
+# register_distributed_operator_impl_container(DistributedScale("tanh"))
+
+
+class DistributedScaleImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        """Calculate the cost by the op role."""
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_op.dist_attr.process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            _g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
+        )
+        res_cost = [cost_mapping]
+
+        return res_cost
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        backward_op = dist_op.serial_op
+        op_type = backward_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            _g_op_cost_factory[op_type], ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and not is_parameter_related(
+                    varname, main_block
+                ):
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        need_gradient_allreduce = True
+                        break
+
+        if need_gradient_allreduce:
+            for input_name in backward_op.desc.input_names():
+                for varname in backward_op.desc.input(input_name):
+                    if "@GRAD" not in varname and is_parameter_related(
+                        varname, main_block
+                    ):
+                        var_dim_mapping = dist_attr.get_input_dims_mapping(
+                            varname
+                        )
+                        mesh_shape = process_mesh.shape
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+        return res
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_name = op_desc.output('Out')[0]
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        in_dims_mappings = []
+        for in_name in op_desc.input_arg_names():
+            in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
+            in_dims_mappings.append(in_dims_mapping)
+
+        for x_dims_mapping in in_dims_mappings:
+            if x_dims_mapping != out_dims_mapping:
+                return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        for i in range(len(x_dims_mapping)):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [x_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+                op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+                changed = True
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+# register_distributed_operator_impl("scale", DistributedScaleImpl("scale"))
+# register_distributed_operator_impl(
+#     "fill_any_like", DistributedScaleImpl("fill_any_like")
+# )
+# register_distributed_operator_impl("where", DistributedScaleImpl("where"))
+# register_distributed_operator_impl("tanh", DistributedScaleImpl("tanh"))

python/paddle/distributed/auto_parallel/static/operators/dist_shape.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..utils import is_dim_shard
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedShape(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(DistributedShape("shape"))
+
+
+class DistributedShapeImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_name = op_desc.output('Out')[0]
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        assert len(out_dims_mapping) == 1
+        if is_dim_shard(out_dims_mapping[0]):
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        return False
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl("shape", DistributedShapeImpl("shape"))

python/paddle/distributed/auto_parallel/static/operators/dist_slice.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ..utils import compute_compatible_dim_mapping, is_dim_shard
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedSlice(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(DistributedSlice("slice"))
+
+
+class DistributedSliceImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        in_name = op_desc.input('Input')[0]
+        out_name = op_desc.output('Out')[0]
+        in_var = dist_op.serial_op.block._var_recursive(in_name)
+        out_var = dist_op.serial_op.block._var_recursive(out_name)
+        axes = op_desc.attr('axes')
+        in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
+        for axis in axes:
+            if (
+                is_dim_shard(in_dims_mapping[axis])
+                and in_var.shape[axis] != out_var.shape[axis]
+            ):
+                return False
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        in_name = op_desc.input('Input')[0]
+        out_name = op_desc.output('Out')[0]
+        in_var = dist_op.serial_op.block._var_recursive(in_name)
+        out_var = dist_op.serial_op.block._var_recursive(out_name)
+        axes = op_desc.attr('axes')
+        decrease_axis = op_desc.attr('decrease_axis')
+        in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        ref_indices = []
+        for i in range(len(in_dims_mapping)):
+            if i not in decrease_axis:
+                ref_indices.append(i)
+        if ref_indices == []:
+            assert len(out_dims_mapping) == 0
+        else:
+            for i in range(len(out_dims_mapping)):
+                ref_index = ref_indices[i]
+                if (
+                    ref_index in axes
+                    and is_dim_shard(out_dims_mapping[i])
+                    and in_var.shape[ref_index] != out_var.shape[ref_index]
+                ):
+                    return False
+
+        return True
+
+    def is_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        in_name = op_desc.input('Input')[0]
+        out_name = op_desc.output('Out')[0]
+        decrease_axis = op_desc.attr('decrease_axis')
+        in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        if len(in_dims_mapping) - len(decrease_axis) != 0 and len(
+            out_dims_mapping
+        ) != len(in_dims_mapping) - len(decrease_axis):
+            return False
+
+        new_out_dims_mapping = []
+        for i in range(len(in_dims_mapping)):
+            if i not in decrease_axis:
+                new_out_dims_mapping.append(in_dims_mapping[i])
+        if new_out_dims_mapping == []:
+            new_out_dims_mapping = [-1]
+        if new_out_dims_mapping != out_dims_mapping:
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (
+            (not self.is_input_compatible(dist_op))
+            or (not self.is_output_compatible(dist_op))
+            or (not self.is_compatible(dist_op))
+        ):
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        in_name = op_desc.input('Input')[0]
+        out_name = op_desc.output('Out')[0]
+        decrease_axis = op_desc.attr('decrease_axis')
+        in_dims_mapping = op_dist_attr.get_input_dims_mapping(in_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        ref_dims_mapping = []
+        ref_indices = []
+        for i in range(len(in_dims_mapping)):
+            if i not in decrease_axis:
+                ref_dims_mapping.append(in_dims_mapping[i])
+                ref_indices.append(i)
+
+        if ref_dims_mapping == []:
+            assert len(ref_dims_mapping) == len(out_dims_mapping)
+            changed = False
+        else:
+            assert len(ref_dims_mapping) == len(out_dims_mapping)
+            for i in range(len(out_dims_mapping)):
+                compatible_dim_mapping = compute_compatible_dim_mapping(
+                    [out_dims_mapping[i], ref_dims_mapping[i]]
+                )
+                if compatible_dim_mapping is None:
+                    continue
+                if ref_dims_mapping[i] != compatible_dim_mapping:
+                    in_dims_mapping[ref_indices[i]] = compatible_dim_mapping
+                    changed = True
+                if out_dims_mapping[i] != compatible_dim_mapping:
+                    out_dims_mapping[i] = compatible_dim_mapping
+                    changed = True
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(in_name, in_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "slice", DistributedSliceImpl("decrease_in_axis")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_softmax.py

@@ -0,0 +1,200 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..cost import (
+    SoftmaxGradOpCost,
+    SoftmaxOpCost,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_dp_costs,
+)
+from ..utils import compute_compatible_and_update_dim_mapping, is_dim_shard
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    is_parameter_related,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedSoftmax(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(DistributedSoftmax("softmax"))
+
+
+class DistributedSoftmaxImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = False
+        self._backward_implemented = False
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_op.dist_attr.process_mesh.process_ids
+        cost_mapping = build_comp_costs_from_descs(
+            SoftmaxOpCost, ctx, processes, desc_mapping, cluster
+        )
+
+        res_cost = [cost_mapping]
+        return res_cost
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        cost_mapping = build_comp_costs_from_descs(
+            SoftmaxGradOpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        backward_op = dist_op.serial_op
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and is_parameter_related(
+                    varname, main_block
+                ):
+                    # NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+
+        return res
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        axis = op_desc.attr('axis')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+
+        # if axis != -1 and axis != len(x_dims_mapping) - 1:
+        #     return False
+
+        if is_dim_shard(x_dims_mapping[axis]):
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_name = op_desc.output('Out')[0]
+        axis = op_desc.attr('axis')
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        # if axis != -1 and axis != len(out_dims_mapping) - 1:
+        #     return False
+
+        if is_dim_shard(out_dims_mapping[axis]):
+            return False
+
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        axis = op_desc.attr('axis')
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        # if axis != -1 and axis != len(x_dims_mapping) - 1:
+        #     return False
+
+        if x_dims_mapping != out_dims_mapping:
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+
+        for i in range(len(x_dims_mapping)):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [x_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                changed = True
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+
+        return changed
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "softmax", DistributedSoftmaxImpl("replicate_last_axis")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_split.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..completion import get_phi_spmd_rule
+from ..utils import (
+    compute_compatible_and_update_dim_mapping,
+    get_dist_tensor_spec,
+    is_dim_shard,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedSplit(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        x_name = op_desc.input('X')[0]
+        assert len(op_desc.input('AxisTensor')) == 0, (
+            "Attribute AxisTensor is not supported by dist split."
+        )
+        assert len(op_desc.input('SectionsTensorList')) == 0, (
+            "Attribute SectionsTensorList is not supported by dist split."
+        )
+        output_arg_names = op_desc.output('Out')
+
+        num = op_desc.attr('num')
+        sections = op_desc.attr('sections')
+        if num:
+            assert (sections is None) or (len(sections) == 0), (
+                f"Both Attributes of num: {num} and sections: {sections} are specified."
+            )
+            first_attr = num
+            rule_type = "split_with_num"
+        else:
+            assert not num, (
+                f"Both Attributes of num: {num} and sections: {sections} are specified."
+            )
+            first_attr = sections
+            rule_type = "split"
+        axis = op_desc.attr('axis')
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        num_outputs = len(output_arg_names)
+        output_specs = []
+        for i in range(num_outputs):
+            output_specs.append(
+                get_dist_tensor_spec(dist_op, output_arg_names[i], False)
+            )
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule(rule_type)
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec, first_attr, axis)
+        bw_results = rule.infer_backward(x_spec, output_specs, first_attr, axis)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, [x_name], output_arg_names, fw_results, bw_results
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        # all split op use default dist operator impl.
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedSplit("split"))
+register_distributed_operator_impl_container(DistributedSplit("split_with_num"))
+
+
+class DistributedSplitImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = True
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        axis = op_desc.attr('axis')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+
+        if is_dim_shard(x_dims_mapping[axis]):
+            return False
+
+        return True
+
+    def is_output_compatible(self, dist_op):
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        out_names = op_desc.output('Out')
+        axis = op_desc.attr('axis')
+        for out_name in out_names:
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+            if is_dim_shard(out_dims_mapping[axis]):
+                return False
+
+        return True
+
+    def is_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        axis = op_desc.attr('axis')
+        out_names = op_desc.output('Out')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        for out_name in out_names:
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+            if x_dims_mapping != out_dims_mapping:
+                return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_names = op_desc.output('Out')
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+
+        for out_name in out_names:
+            out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+            for i in range(len(x_dims_mapping)):
+                dim_changed = compute_compatible_and_update_dim_mapping(
+                    [x_dims_mapping, out_dims_mapping], [i, i]
+                )
+                if dim_changed:
+                    changed = True
+                    op_dist_attr.set_output_dims_mapping(
+                        out_name, out_dims_mapping
+                    )
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+        return changed
+
+    def is_auto_compatible(self, dist_op):
+        if (
+            (not self.is_input_compatible(dist_op))
+            or (not self.is_output_compatible(dist_op))
+            or (not self.is_compatible(dist_op))
+        ):
+            return False
+
+        return True
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "split", DistributedSplitImpl("replicate_in_axis")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_stack.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ..completion import get_phi_spmd_rule
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedStack(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        input_arg_names = op_desc.input_arg_names()
+        output_arg_names = op_desc.output_arg_names()
+        axis = op_desc.attr('axis')
+
+        input_specs = []
+        for name in input_arg_names:
+            input_specs.append(get_dist_tensor_spec(dist_op, name))
+        output_spec = get_dist_tensor_spec(dist_op, output_arg_names[0], False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("stack")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(input_specs, axis)
+        bw_results = rule.infer_backward(input_specs, output_spec, axis)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            input_arg_names,
+            output_arg_names,
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedStack("stack"))

python/paddle/distributed/auto_parallel/static/operators/dist_strided_slice.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..completion import get_phi_spmd_rule
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedStridedSlice(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        x_name = op_desc.input('Input')[0]
+        y_name = op_desc.output('Out')[0]
+        axes = op_desc.attr('axes')
+        starts = op_desc.attr('starts')
+        ends = op_desc.attr('ends')
+        strides = op_desc.attr('strides')
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        y_spec = get_dist_tensor_spec(dist_op, y_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("strided_slice")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec, axes, starts, ends, strides)
+        bw_results = rule.infer_backward(
+            x_spec,
+            y_spec,
+            axes,
+            starts,
+            ends,
+            strides,
+        )
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            [x_name],
+            [y_name],
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedStridedSlice("strided_slice")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_tile.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+
+from ..completion import get_phi_spmd_rule
+from ..utils import (
+    get_dist_tensor_spec,
+)
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedTile(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+        assert dist_op.serial_op.type == "tile", (
+            f"{dist_op.serial_op.type} is not supported by dist transpose yet."
+        )
+
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        repeat_times = op_desc.attr('repeat_times')
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        output_spec = get_dist_tensor_spec(dist_op, out_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("tile")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec, repeat_times)
+        bw_results = rule.infer_backward(x_spec, output_spec, repeat_times)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, [x_name], [out_name], fw_results, bw_results
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        # all elementwise op use default dist operator impl.
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(DistributedTile("tile"))

python/paddle/distributed/auto_parallel/static/operators/dist_transpose.py

@@ -0,0 +1,270 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..completion import get_phi_spmd_rule
+from ..cost import (
+    Transpose2GradOpCost,
+    Transpose2OpCost,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_dp_costs,
+)
+from ..utils import (
+    compute_compatible_and_update_dim_mapping,
+    get_dist_tensor_spec,
+)
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    is_parameter_related,
+    merge_forward_backward_dims_mapping,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+from .dist_default import DistributedDefaultImpl0
+
+
+class DistributedTranspose2(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+        assert dist_op.serial_op.type == "transpose2", (
+            f"{dist_op.serial_op.type} is not supported by dist transpose yet."
+        )
+
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        xshape_name = op_desc.output('XShape')[0]
+        axes = op_desc.attr('axis')
+
+        x_spec = get_dist_tensor_spec(dist_op, x_name)
+        output_spec = get_dist_tensor_spec(dist_op, out_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("transpose")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(x_spec, axes)
+        bw_results = rule.infer_backward(x_spec, output_spec, axes)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op, [x_name], [out_name], fw_results, bw_results
+        )
+
+        # step4: update xshape
+        inferred_input_dims_mappings, _ = merge_forward_backward_dims_mapping(
+            fw_results, bw_results
+        )
+        dist_op.dist_attr.set_output_dims_mapping(
+            xshape_name, [-1] + inferred_input_dims_mappings[0]
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        # all elementwise op use default dist operator impl.
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedTranspose2("transpose2")
+)
+
+
+class DistributedTranspose2Impl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._forward_implemented = False
+        self._backward_implemented = False
+
+    def is_input_compatible(self, dist_op):
+        return True
+
+    def is_output_compatible(self, dist_op):
+        return True
+
+    def is_auto_compatible(self, dist_op):
+        if (not self.is_input_compatible(dist_op)) or (
+            not self.is_output_compatible(dist_op)
+        ):
+            return False
+
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        perm = op_desc.attr('axis')
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        new_dims_mapping = [-1 for i in range(len(x_dims_mapping))]
+        for i in range(len(x_dims_mapping)):
+            new_dims_mapping[i] = x_dims_mapping[perm[i]]
+
+        if len(x_dims_mapping) != len(out_dims_mapping):
+            return False
+
+        if new_dims_mapping != out_dims_mapping:
+            return False
+
+        if x_shape_dims_mapping[0] != -1:
+            return False
+
+        if x_shape_dims_mapping[1:] != x_dims_mapping[:]:
+            return False
+
+        return True
+
+    def update_dims_mapping(self, dist_op):
+        changed = False
+        op_desc = dist_op.serial_op.desc
+        op_dist_attr = dist_op.dist_attr
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        x_shape_name = op_desc.output('XShape')[0]
+        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
+        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
+        x_shape_dims_mapping = op_dist_attr.get_output_dims_mapping(
+            x_shape_name
+        )
+        perm = op_desc.attr('axis')
+
+        assert len(x_dims_mapping) == len(perm)
+
+        new_dims_mapping = [-1 for i in range(len(x_dims_mapping))]
+        for i in range(len(x_dims_mapping)):
+            new_dims_mapping[i] = x_dims_mapping[perm[i]]
+
+        for i in range(len(out_dims_mapping)):
+            dim_changed = compute_compatible_and_update_dim_mapping(
+                [new_dims_mapping, out_dims_mapping], [i, i]
+            )
+            if dim_changed:
+                changed = True
+
+        for i in range(len(x_dims_mapping)):
+            if x_dims_mapping[perm[i]] != new_dims_mapping[i]:
+                x_dims_mapping[perm[i]] = new_dims_mapping[i]
+                changed = True
+
+        for i in range(len(x_dims_mapping)):
+            x_shape_dims_mapping[i + 1] = x_dims_mapping[i]
+
+        if changed:
+            op_dist_attr.set_input_dims_mapping(x_name, x_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)
+            op_dist_attr.set_output_dims_mapping(
+                x_shape_name, x_shape_dims_mapping
+            )
+
+        return changed
+
+    def calc_cost(self, op_role, dist_op, ctx, cluster):
+        cost = None
+        if int(op_role) == int(OpRole.Backward):
+            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
+        else:
+            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
+        assert cost is not None
+        return cost
+
+    def calc_fwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        processes = dist_op.dist_attr.process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            Transpose2OpCost, ctx, processes, desc_mapping, cluster
+        )
+
+        res_cost = [cost_mapping]
+        return res_cost
+
+    def calc_bwd_cost(self, dist_op, ctx, cluster):
+        # calc comp op cost
+        res = []
+        desc_mapping = build_comp_desc_from_dist_op(
+            dist_op=dist_op, dist_context=ctx
+        )
+        dist_attr = dist_op.dist_attr
+        process_mesh = dist_attr.process_mesh
+        processes = process_mesh.process_ids
+        op_type = dist_op.serial_op.type
+        cost_mapping = build_comp_costs_from_descs(
+            Transpose2GradOpCost, ctx, processes, desc_mapping, cluster
+        )
+        res.append(cost_mapping)
+
+        backward_op = dist_op.serial_op
+        main_block = backward_op.block
+        need_gradient_allreduce = False
+        for input_name in backward_op.desc.input_names():
+            for varname in backward_op.desc.input(input_name):
+                if "@GRAD" not in varname and is_parameter_related(
+                    varname, main_block
+                ):
+                    # NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
+                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)
+
+                    mesh_shape = process_mesh.shape
+                    batch_size_axis = (
+                        var_dim_mapping[0] if len(var_dim_mapping) > 0 else -1
+                    )
+                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
+                        parallel_axis = batch_size_axis
+                        attrs = {"use_calc_stream": True}
+                        var_names = [varname + "@GRAD"]
+                        build_dp_costs(
+                            res,
+                            dist_op,
+                            ctx,
+                            var_names,
+                            attrs,
+                            parallel_axis,
+                            cluster,
+                        )
+        return res
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
+
+
+register_distributed_operator_impl(
+    "transpose2", DistributedTranspose2Impl("same_mapping_transpose")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_unsqueeze2.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ..completion import get_phi_spmd_rule
+from ..utils import get_dist_tensor_spec
+from .common import (
+    DistributedOperatorImplContainer,
+    get_default_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+    update_op_dims_mapping,
+)
+
+
+class DistributedUnSqueeze2(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+    @staticmethod
+    def update_dims_mapping(dist_op):
+        # step1: prepare inputs need for rule (order args as PHI definition and filter out unnecessary args)
+        op_desc = dist_op.serial_op.desc
+
+        axes_tensor = op_desc.input('AxesTensor')
+        axes_tensor_list = op_desc.input('AxesTensorList')
+        assert len(axes_tensor) == 0 and len(axes_tensor_list) == 0
+
+        x_name = op_desc.input('X')[0]
+        out_name = op_desc.output('Out')[0]
+        axes = op_desc.attr('axes')
+
+        input_spec = get_dist_tensor_spec(dist_op, x_name)
+        output_spec = get_dist_tensor_spec(dist_op, out_name, False)
+
+        # step2: infer spmd
+        rule = get_phi_spmd_rule("unsqueeze2")
+        # tensor order following order in PHI definition
+        fw_results = rule.infer_forward(input_spec, axes)
+        bw_results = rule.infer_backward(input_spec, output_spec, axes)
+
+        # step3: update dist_attr
+        # tensor order following order in PHI definition
+        changed = update_op_dims_mapping(
+            dist_op,
+            [x_name],
+            [out_name],
+            fw_results,
+            bw_results,
+        )
+
+        return changed
+
+    @staticmethod
+    def mapping_to_dist_operator_impl(dist_op, original_op_dist_attr):
+        op_dist_attr = dist_op.dist_attr
+        default_impl = get_default_distributed_operator_impl()
+        op_dist_attr.impl_type = default_impl.type
+        op_dist_attr.impl_idx = default_impl.idx
+
+        return False
+
+
+register_distributed_operator_impl_container(
+    DistributedUnSqueeze2("unsqueeze2")
+)

python/paddle/distributed/auto_parallel/static/operators/dist_update_loss_scaling.py

@@ -0,0 +1,171 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from ..utils import set_dist_op_desc_original_id
+from .common import (
+    DistributedOperatorImpl,
+    DistributedOperatorImplContainer,
+    register_distributed_operator_impl,
+    register_distributed_operator_impl_container,
+)
+
+
+class DistributedUpdateLossScaling(DistributedOperatorImplContainer):
+    def __init__(self, op_type):
+        super().__init__(op_type)
+
+
+register_distributed_operator_impl_container(
+    DistributedUpdateLossScaling("update_loss_scaling")
+)
+
+
+class DistributedUpdateLossScalingImpl(DistributedOperatorImpl):
+    def __init__(self, name):
+        super().__init__(name)
+        self._name = name
+        self._forward_implemented = False
+        self._backward_implemented = True
+
+    def is_input_compatible(self, dist_op):
+        raise RuntimeError(
+            "DistributedUpdateLossScalingImpl's is_input_compatible should not be called !"
+        )
+
+    def is_output_compatible(self, dist_op):
+        raise RuntimeError(
+            "DistributedUpdateLossScalingImpl's is_output_compatible should not be called !"
+        )
+
+    def is_auto_compatible(self, dist_op):
+        raise RuntimeError(
+            "DistributedUpdateLossScalingImpl's is_auto_compatible should not be called !"
+        )
+
+    def update_dims_mapping(self, dist_op):
+        raise RuntimeError(
+            "DistributedUpdateLossScalingImpl's update_dims_mapping should not be called !"
+        )
+
+    @staticmethod
+    def forward(ctx, *args, **kwargs):
+        raise RuntimeError(
+            "DistributedUpdateLossScalingImpl's forward should not be called !"
+        )
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        # the backward function only filter the gradient with current rank id
+        dist_op_context = ctx.dist_op_context
+        main_block = dist_op_context.main_block
+        backward_op = dist_op_context.cur_src_op
+        rank_id = dist_op_context.rank_id
+        dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
+        assert dist_attr is not None, (
+            f"backward op [{backward_op}] don't have dist attribute !"
+        )
+
+        assert rank_id in dist_attr.process_mesh.process_ids
+
+        assert 'X' in kwargs, "input [{}] is not given".format('X')
+        assert 'FoundInfinite' in kwargs, "input [{}] is not given".format(
+            'FoundInfinite'
+        )
+        assert 'PrevLossScaling' in kwargs, "input [{}] is not given".format(
+            'PrevLossScaling'
+        )
+        assert 'InGoodSteps' in kwargs, "input [{}] is not given".format(
+            'InGoodSteps'
+        )
+        assert 'InBadSteps' in kwargs, "input [{}] is not given".format(
+            'InBadSteps'
+        )
+
+        assert 'Out' in kwargs, "output [{}] is not given".format('Out')
+        assert 'LossScaling' in kwargs, "output [{}] is not given".format(
+            'LossScaling'
+        )
+        assert 'OutGoodSteps' in kwargs, "output [{}] is not given".format(
+            'OutGoodSteps'
+        )
+        assert 'OutBadSteps' in kwargs, "output [{}] is not given".format(
+            'OutBadSteps'
+        )
+
+        assert len(kwargs['FoundInfinite']) == 1, (
+            "update_loss_scaling input FoundInfinite take 1 variable but got {}".format(
+                kwargs['FoundInfinite']
+            )
+        )
+        assert len(kwargs['PrevLossScaling']) == 1, (
+            "update_loss_scaling input PrevLossScaling take 1 variable but got {}".format(
+                kwargs['PrevLossScaling']
+            )
+        )
+        assert len(kwargs['InGoodSteps']) == 1, (
+            "update_loss_scaling input InGoodSteps take 1 variable but got {}".format(
+                kwargs['InGoodSteps']
+            )
+        )
+        assert len(kwargs['InBadSteps']) == 1, (
+            "update_loss_scaling input InBadSteps take 1 variable but got {}".format(
+                kwargs['InBadSteps']
+            )
+        )
+        assert len(kwargs['LossScaling']) == 1, (
+            "update_loss_scaling output LossScaling take 1 variable but got {}".format(
+                kwargs['LossScaling']
+            )
+        )
+        assert len(kwargs['OutGoodSteps']) == 1, (
+            "update_loss_scaling output OutGoodSteps take 1 variable but got {}".format(
+                kwargs['OutGoodSteps']
+            )
+        )
+        assert len(kwargs['OutBadSteps']) == 1, (
+            "update_loss_scaling output OutBadSteps take 1 variable but got {}".format(
+                kwargs['OutBadSteps']
+            )
+        )
+
+        assert len(kwargs['X']) == len(kwargs['Out']), (
+            "update_loss_scaling got [{}] X and [{}] Out, which are supposed to be equal".format(
+                len(kwargs['X']), len(kwargs['Out'])
+            )
+        )
+
+        filter_vars = []
+        for varname in kwargs['X']:
+            if (
+                rank_id
+                in ctx.get_tensor_dist_attr_for_program(
+                    main_block._var_recursive(varname)
+                ).process_mesh.process_ids
+            ):
+                filter_vars.append(varname)
+
+        # replicate op in dist program
+        dist_op = main_block.append_op(type='nop')
+        dist_op_desc = dist_op.desc
+        dist_op_desc.copy_from(backward_op.desc)
+        set_dist_op_desc_original_id(dist_op_desc, backward_op.desc, ctx)
+        dist_op_desc.set_input('X', filter_vars)
+        dist_op_desc.set_output('Out', filter_vars)
+        # TODO: should we add a new dist attr for the new op here?
+
+
+register_distributed_operator_impl(
+    "update_loss_scaling",
+    DistributedUpdateLossScalingImpl("update_loss_scaling"),
+)