chore: import upstream snapshot with attribution

python/paddle/distributed/__init__.py

@@ -0,0 +1,253 @@
+# Copyright (c) 2019 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 atexit  # noqa: F401
+
+from .value_patch import monkey_patch_value_in_dist
+
+monkey_patch_value_in_dist()
+from paddle.base.core import Placement, ProcessGroup, ReduceType
+from paddle.distributed.fleet.base.topology import (
+    ParallelMode,
+    create_nccl_config,
+)
+from paddle.distributed.fleet.dataset import InMemoryDataset, QueueDataset
+
+from . import (
+    cloud_utils,  # noqa: F401
+    io,
+    rpc,  # noqa: F401
+)
+from .auto_parallel import shard_op  # noqa: F401
+from .auto_parallel.api import (
+    DistAttr,
+    DistModel,
+    ShardingStage1,
+    ShardingStage2,
+    ShardingStage3,
+    Strategy,
+    dtensor_from_fn,
+    enable_auto_dp,  # noqa: F401
+    in_auto_parallel_align_mode,  # noqa: F401
+    reshard,
+    shard_dataloader,
+    shard_layer,
+    shard_optimizer,
+    shard_scaler,
+    shard_tensor,
+    to_static,
+    unshard_dtensor,
+)
+from .auto_parallel.high_level_api import to_distributed
+from .auto_parallel.interface import get_mesh, set_mesh
+from .auto_parallel.intermediate.context_parallel import (
+    ContextParallel,
+    PrepareContextParallel,
+)
+from .auto_parallel.intermediate.parallelize import parallelize
+from .auto_parallel.intermediate.pipeline_parallel import SplitPoint
+from .auto_parallel.intermediate.tensor_parallel import (
+    ColWiseParallel,
+    ConvParallel,
+    PrepareLayerInput,
+    PrepareLayerOutput,
+    RowWiseParallel,
+    SequenceParallelBegin,
+    SequenceParallelDisable,
+    SequenceParallelEnable,
+    SequenceParallelEnd,
+)
+from .auto_parallel.local_layer import LocalLayer
+from .auto_parallel.local_map import local_map
+from .auto_parallel.placement_type import (
+    Partial,
+    Replicate,
+    Shard,
+)
+from .auto_parallel.process_mesh import ProcessMesh
+from .collective import (
+    is_available,
+    new_group,
+    restart_process_group,
+    shutdown_process_group,
+    split,
+)
+from .communication import (  # noqa: F401
+    P2POp,
+    ReduceOp,
+    all_gather,
+    all_gather_object,
+    all_reduce,
+    alltoall,
+    alltoall_single,
+    barrier,
+    batch_isend_irecv,
+    broadcast,
+    broadcast_object_list,
+    destroy_process_group,
+    gather,
+    get_backend,
+    get_group,
+    irecv,
+    is_initialized,
+    isend,
+    recv,
+    recv_object_list,
+    reduce,
+    reduce_scatter,
+    scatter,
+    scatter_object_list,
+    send,
+    send_object_list,
+    stream,
+    wait,
+)
+
+# Import the namespace class directly from the submodule so it does not
+# shadow ``communication.group`` (the submodule) inside the package.
+from .communication.group import _DistGroupNamespace as group
+from .entry_attr import (
+    CountFilterEntry,
+    ProbabilityEntry,
+    ShowClickEntry,
+)
+from .flex_checkpoint.dcp.load_state_dict import (
+    load_merged_state_dict,
+    load_state_dict,
+)
+from .flex_checkpoint.dcp.save_state_dict import save_state_dict
+from .flex_checkpoint.dcp.sharded_weight import (
+    ShardedStateDict,
+    ShardedWeight,
+    build_sharded_state_dict,
+    shard_weight,
+)
+from .launch.main import launch
+from .parallel import (  # noqa: F401
+    DataParallel,
+    ParallelEnv,
+    get_rank,
+    get_world_size,
+    init_parallel_env,
+    init_process_group,
+)
+from .parallel_with_gloo import (
+    gloo_barrier,
+    gloo_init_parallel_env,
+    gloo_release,
+)
+from .sharding import (  # noqa: F401
+    group_sharded_parallel,
+    save_group_sharded_model,
+)
+from .spawn import spawn
+
+__all__ = [
+    "io",
+    "spawn",
+    "launch",
+    "scatter",
+    "gather",
+    "scatter_object_list",
+    "broadcast",
+    "broadcast_object_list",
+    "ParallelEnv",
+    "new_group",
+    "shutdown_process_group",
+    "restart_process_group",
+    "init_parallel_env",
+    "init_process_group",
+    "group",
+    "ProcessGroup",
+    "gloo_init_parallel_env",
+    "gloo_barrier",
+    "gloo_release",
+    "QueueDataset",
+    "split",
+    "CountFilterEntry",
+    "ShowClickEntry",
+    "get_world_size",
+    "get_group",
+    "all_gather",
+    "all_gather_object",
+    "InMemoryDataset",
+    "barrier",
+    "all_reduce",
+    "alltoall",
+    "alltoall_single",
+    "send",
+    "reduce",
+    "recv",
+    "ReduceOp",
+    "wait",
+    "get_rank",
+    "ProbabilityEntry",
+    "ParallelMode",
+    "is_initialized",
+    "destroy_process_group",
+    "isend",
+    "irecv",
+    "send_object_list",
+    "recv_object_list",
+    "reduce_scatter",
+    "is_available",
+    "get_backend",
+    "ProcessMesh",
+    "DistAttr",
+    "shard_tensor",
+    "dtensor_from_fn",
+    "reshard",
+    "shard_layer",
+    "shard_dataloader",
+    "ReduceType",
+    "Placement",
+    "Shard",
+    "Replicate",
+    "Partial",
+    "save_state_dict",
+    "load_state_dict",
+    "load_merged_state_dict",
+    "shard_optimizer",
+    "shard_scaler",
+    "ShardingStage1",
+    "ShardingStage2",
+    "ShardingStage3",
+    "to_static",
+    "Strategy",
+    "DistModel",
+    "LocalLayer",
+    "local_map",
+    "unshard_dtensor",
+    "parallelize",
+    "SequenceParallelEnd",
+    "SequenceParallelBegin",
+    "SequenceParallelEnable",
+    "SequenceParallelDisable",
+    "ColWiseParallel",
+    "RowWiseParallel",
+    "PrepareLayerOutput",
+    "PrepareLayerInput",
+    "SplitPoint",
+    "set_mesh",
+    "get_mesh",
+    "to_distributed",
+    "ConvParallel",
+    "ContextParallel",
+    "PrepareContextParallel",
+    "create_nccl_config",
+    "ShardedWeight",
+    "ShardedStateDict",
+    "shard_weight",
+    "build_sharded_state_dict",
+]

python/paddle/distributed/auto_parallel/__init__.py

@@ -0,0 +1,32 @@
+#   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 .interface import (  # noqa: F401
+    create_mesh,
+    exclude_ops_in_recompute,
+    fetch,
+    get_mesh,
+    recompute,
+    set_mesh,
+    shard_op,
+    shard_tensor,
+)
+from .process_mesh import ProcessMesh  # noqa: F401
+from .random import parallel_manual_seed  # noqa: F401
+from .ring_attention import RingFlashAttention  # noqa: F401
+from .ring_conv import RingConv2d  # noqa: F401
+from .static.engine import Engine  # noqa: F401
+from .strategy import Strategy  # noqa: F401
+
+__all__ = []

python/paddle/distributed/auto_parallel/_utils.py

@@ -0,0 +1,98 @@
+# Copyright (c) 2025 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 functools import wraps
+
+import paddle
+
+
+# NOTE(zhengtianyu): align ClipGradByGlobalNorm in auto_parallel_align_mode.
+# In old dygraph semi-auto parallel, each rank has parameter and gradient information
+# from other ranks. To align with this behavior, this decorator ensures auto_hybrid_pp
+# uses the same logic as old dygraph semi-auto parallel for ClipGradByGlobalNorm in align mode.
+# Pay attention to the auto_hybrid_pp's default logic matches dynamic manual-parallel,
+# Refer to NOTE: Fix grad_clip in auto_hybrid_pp mode
+def _patch_grads_for_step(
+    amp_master_grad=False,
+):
+    """
+    Only for auto parallel align mode, use this decorator to handle None gradients in optimizer step.
+
+    This decorator is applied to optimizer step methods to handle cases where parameters
+    have None gradients. It creates zero gradients for parameters that need gradients
+    but currently have None gradients.
+
+    Args:
+        amp_master_grad (bool, optional): Whether to use master gradient mode.
+            If True, gradients will be created as float32 regardless of parameter dtype.
+            If False, gradients will be created with the same dtype as the parameter.
+            Default is False.
+
+    Returns:
+        function: Decorated step method that handles None gradients.
+
+    Example:
+        .. code-block:: pycon
+
+            >>> from __future__ import annotations
+            >>> import paddle.distributed as dist
+            >>> import types
+            >>> from paddle.distributed.auto_parallel._utils import _patch_grads_for_step
+
+            >>> opt = paddle.optimizer.AdamW(
+            ...     learning_rate=0.001,
+            ...     parameters=self.model.parameters(),
+            ...     grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+            ... )
+            >>> if dist.in_auto_parallel_align_mode():
+            >>>     orig_step = (
+            ...         opt.step.__func__ if hasattr(opt.step, "__func__") else opt.step
+            ...     )
+            >>>     decorator = (
+            ...         _patch_grads_for_step(
+            ...             amp_master_grad=True
+            ...         )
+            ...     )
+            >>>     new_step = decorator(orig_step)
+            >>>     opt.step = types.MethodType(new_step, opt)
+
+    """
+
+    def decorator(step_method):
+        @wraps(step_method)
+        def wrapper(self, *args, **kwargs):
+            # Helper function to set gradient for a parameter
+            def set_param_grad(param):
+                if param.stop_gradient or param.grad is not None:
+                    return
+
+                if hasattr(param, "main_grad"):
+                    param.main_grad = paddle.zeros_like(
+                        param, dtype=paddle.float32
+                    )
+                else:
+                    dtype = paddle.float32 if amp_master_grad else param.dtype
+                    param.grad = paddle.zeros_like(param, dtype=dtype)
+
+            if not isinstance(self._parameter_list[0], dict):
+                for param in self._parameter_list:
+                    set_param_grad(param)
+            else:
+                for param_group in self._param_groups:
+                    for param in param_group['params']:
+                        set_param_grad(param)
+            return step_method(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator

python/paddle/distributed/auto_parallel/auto_dp_utils.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2025 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
+import paddle.distributed as dist
+
+_enable_auto_dp_mode = False
+
+
+def _fake_replicate_grad_to_partial(grad, partial_axis):
+    new_placements = grad.placements
+    assert new_placements[partial_axis] == dist.Replicate(), (
+        "when reshard fake replicated grad to partial, the partial axis of grad should be Replicate"
+    )
+
+    new_placements[partial_axis] = dist.Partial(dist.ReduceType.kRedSum)
+
+    grad_mesh = grad.process_mesh
+    grad = dist.auto_parallel.api.dtensor_to_local(
+        grad, grad_mesh, grad.placements
+    )
+    grad = dist.auto_parallel.api.dtensor_from_local(
+        grad, grad_mesh, new_placements
+    )
+    return grad
+
+
+def _convert_fake_replicate_grad_to_partial(params_grads):
+    # skip non-parallel cases
+    world_size = paddle.distributed.get_world_size()
+    if world_size == 1:
+        return
+
+    if isinstance(params_grads, list):
+        for idx in range(len(params_grads)):
+            param, grad = params_grads[idx][0], params_grads[idx][1]
+            if grad.is_dist():
+                grad_placements = grad.placements
+                if not isinstance(grad_placements[0], dist.Partial):
+                    grad = _fake_replicate_grad_to_partial(grad, 0)
+            else:
+                default_grad_placements = [
+                    dist.Partial(dist.ReduceType.kRedSum)
+                ]
+                default_grad_mesh = dist.ProcessMesh(
+                    list(range(0, world_size)), dim_names=["dp"]
+                )
+                grad = dist.auto_parallel.api.dtensor_from_local(
+                    grad, default_grad_mesh, default_grad_placements
+                )
+            params_grads[idx] = (param, grad)
+    else:
+        for idx in range(len(params_grads['params'])):
+            grad = params_grads['params'][idx][1]
+            if grad.is_dist():
+                grad_placements = grad.placements
+                if not isinstance(grad_placements[0], dist.Partial):
+                    grad = _fake_replicate_grad_to_partial(grad, 0)
+            else:
+                default_grad_placements = [
+                    dist.Partial(dist.ReduceType.kRedSum)
+                ]
+                default_grad_mesh = dist.ProcessMesh(
+                    list(range(0, world_size)), dim_names=["dp"]
+                )
+                grad = dist.auto_parallel.api.dtensor_from_local(
+                    grad, default_grad_mesh, default_grad_placements
+                )
+            params_grads['params'][idx] = (params_grads['params'][idx][0], grad)
+
+
+def in_auto_dp_mode():
+    world_size = paddle.distributed.get_world_size()
+    if world_size <= 1:
+        return False
+
+    global _enable_auto_dp_mode
+    return _enable_auto_dp_mode
+
+
+def _enable_auto_dp():
+    global _enable_auto_dp_mode
+    _enable_auto_dp_mode = True

python/paddle/distributed/auto_parallel/constants.py

@@ -0,0 +1,370 @@
+# 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 __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING, TypedDict
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing.dtype_like import _DTypeLiteral
+
+# _g_default_config[category][field] = default_value
+_g_default_config = defaultdict(dict)
+
+
+def get_category_default_config(category):
+    return _g_default_config[category]
+
+
+def set_category_default_config(category, default_value):
+    _g_default_config[category] = default_value
+
+
+def get_field_default_config(category, field):
+    return _g_default_config[category][field]
+
+
+def set_field_default_config(category, field, default_value):
+    _g_default_config[category][field] = default_value
+
+
+NOT_FOUND = "not_found"
+
+#########################################
+# base configuration
+#########################################
+BASE = "base"
+set_field_default_config(BASE, "auto_mode", "semi")
+set_field_default_config(BASE, "gradient_scale", True)
+set_field_default_config(BASE, "gradient_scale_using_allreduce_avg", False)
+set_field_default_config(BASE, "use_cache", True)
+set_field_default_config(BASE, "return_numpy", True)
+set_field_default_config(BASE, "all_ranks", False)
+set_field_default_config(BASE, "split_data", True)
+set_field_default_config(BASE, "seed", None)
+set_field_default_config(BASE, "reinit", False)  # Only for debug
+
+if TYPE_CHECKING:
+
+    class _BaseConfig(TypedDict, total=False):  # noqa: PYI049
+        auto_mode: str
+        gradient_scale: bool
+        gradient_scale_using_allreduce_avg: bool
+        use_cache: bool
+        return_numpy: bool
+        all_ranks: bool
+        split_data: bool
+        seed: int | None
+        reinit: bool
+
+
+#########################################
+# recompute configuration
+#########################################
+RECOMPUTE = "recompute"
+set_field_default_config(RECOMPUTE, "enable", False)
+set_field_default_config(RECOMPUTE, "checkpoints", [])
+set_field_default_config(RECOMPUTE, "no_recompute_segments", [])
+set_field_default_config(RECOMPUTE, "sr", 0)
+set_field_default_config(RECOMPUTE, "refined_ops_patterns", [])  # List[Dict]
+set_field_default_config(RECOMPUTE, "enable_tuning", False)
+
+if TYPE_CHECKING:
+
+    class _RefinedOpsPatterns(TypedDict, total=False):
+        main_ops: list[str]
+        num: int
+        pre_ops: list[str]
+        suf_ops: list[str]
+
+    class _RecomputeConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        checkpoints: list[Tensor]
+        no_recompute_segments: list[int]
+        sr: int
+        refined_ops_patterns: list[_RefinedOpsPatterns]
+        enable_tuning: bool
+
+
+#########################################
+# AMP configuration
+#########################################
+AMP = "amp"
+set_field_default_config(AMP, "enable", False)
+set_field_default_config(AMP, "dtype", "float16")
+set_field_default_config(AMP, "level", "o1")
+set_field_default_config(AMP, "init_loss_scaling", 32768.0)
+set_field_default_config(AMP, "incr_every_n_steps", 1000)
+set_field_default_config(AMP, "decr_every_n_nan_or_inf", 2)
+set_field_default_config(AMP, "incr_ratio", 2.0)
+set_field_default_config(AMP, "decr_ratio", 0.8)
+set_field_default_config(AMP, "use_dynamic_loss_scaling", True)
+set_field_default_config(AMP, "custom_white_list", [])
+set_field_default_config(AMP, "custom_black_list", [])
+set_field_default_config(AMP, "custom_black_varnames", [])
+set_field_default_config(AMP, "use_fp16_guard", False)
+set_field_default_config(AMP, "use_bf16_guard", False)
+set_field_default_config(AMP, "use_master_grad", False)
+set_field_default_config(AMP, "use_promote", True)
+
+if TYPE_CHECKING:
+
+    class _AMPConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        dtype: _DTypeLiteral
+        level: str
+        init_loss_scaling: float
+        incr_every_n_steps: int
+        decr_every_n_nan_or_inf: int
+        incr_ratio: float
+        decr_ratio: float
+        use_dynamic_loss_scaling: bool
+        custom_white_list: list[str]
+        custom_black_list: list[str]
+        custom_black_varnames: list[str]
+        use_fp16_guard: bool
+        use_bf16_guard: bool
+        use_master_grad: bool
+        use_promote: bool
+
+
+#########################################
+# sharding configuration
+#########################################
+SHARDING = "sharding"
+set_field_default_config(SHARDING, "enable", False)
+set_field_default_config(SHARDING, "stage", 1)
+set_field_default_config(SHARDING, "degree", 8)
+set_field_default_config(SHARDING, "enable_overlap", False)
+set_field_default_config(SHARDING, "param_comm_stream_num", 1)
+set_field_default_config(SHARDING, "grad_comm_stream_num", 1)
+set_field_default_config(SHARDING, "param_bucket_size_numel", 1)
+set_field_default_config(SHARDING, "grad_bucket_size_numel", 1)
+set_field_default_config(SHARDING, "enable_hierarchical_comm", False)
+set_field_default_config(SHARDING, "partition_algor", "greedy_even")
+set_field_default_config(SHARDING, "enable_tuning", False)
+set_field_default_config(SHARDING, "tuning_range", [])
+set_field_default_config(SHARDING, "release_gradients", False)
+set_field_default_config(SHARDING, "comm_buffer_size_MB", 256)
+set_field_default_config(SHARDING, "enable_tensor_fusion", False)
+set_field_default_config(SHARDING, "save_unbalanced_param", True)
+
+
+if TYPE_CHECKING:
+
+    class _ShardingConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        stage: int
+        degree: int
+        enable_overlap: bool
+        param_comm_stream_num: int
+        grad_comm_stream_num: int
+        param_bucket_size_numel: int
+        grad_bucket_size_numel: int
+        enable_hierarchical_comm: bool
+        partition_algor: str
+        enable_tuning: bool
+        tuning_range: list[int] | tuple[int, int]
+
+
+#########################################
+# gradient merge configuration
+#########################################
+GRADIENT_MERGE = "gradient_merge"
+set_field_default_config(GRADIENT_MERGE, "enable", False)
+set_field_default_config(GRADIENT_MERGE, "k_steps", 1)
+set_field_default_config(GRADIENT_MERGE, "avg", True)
+
+if TYPE_CHECKING:
+
+    class _GradientMergeConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        k_steps: int
+        avg: bool
+
+
+#########################################
+# pipeline configuration
+#########################################
+PIPELINE = "pipeline"
+set_field_default_config(PIPELINE, "enable", False)
+set_field_default_config(PIPELINE, "schedule_mode", "1F1B")
+set_field_default_config(PIPELINE, "pp_degree", 1)
+set_field_default_config(PIPELINE, "vpp_degree", 1)
+set_field_default_config(PIPELINE, "vpp_seg_method", "")
+set_field_default_config(PIPELINE, "micro_batch_size", 1)
+set_field_default_config(PIPELINE, "accumulate_steps", 1)
+set_field_default_config(PIPELINE, "generation_batch_size", 1)
+set_field_default_config(PIPELINE, "enable_send_recv_overlap", False)
+set_field_default_config(PIPELINE, "job_schedule_profiler_start", -1)
+set_field_default_config(PIPELINE, "job_schedule_profiler_stop", -1)
+set_field_default_config(PIPELINE, "program_runtimes", [61, 72, 71, 34, 3])
+set_field_default_config(PIPELINE, "memory_limit_times", -1)
+set_field_default_config(PIPELINE, "split_backward", False)
+set_field_default_config(PIPELINE, "auto_parallel_sync_shared_params", False)
+
+if TYPE_CHECKING:
+
+    class _PipelineConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        schedule_mode: str
+        pp_degree: int
+        vpp_degree: int
+        vpp_seg_method: str
+        micro_batch_size: int
+        accumulate_steps: int
+        generation_batch_size: int
+        enable_send_recv_overlap: bool
+        job_schedule_profiler_start: int
+        job_schedule_profiler_stop: int
+        split_backward: bool
+        auto_parallel_sync_shared_params: bool
+
+
+#########################################
+# quantization configuration
+#########################################
+QAT = "qat"
+set_field_default_config(QAT, "enable", False)
+set_field_default_config(QAT, "channel_wise_abs_max", True)
+set_field_default_config(QAT, "weight_bits", 8)
+set_field_default_config(QAT, "activation_bits", 8)
+set_field_default_config(QAT, "not_quant_pattern", ['skip_quant'])
+set_field_default_config(QAT, "algo", None)
+set_field_default_config(QAT, "onnx_format", True)
+
+if TYPE_CHECKING:
+
+    class _QATConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        channel_wise_abs_max: bool
+        weight_bits: int
+        activation_bits: int
+        not_quant_pattern: list[str]
+        algo: str | None
+        onnx_format: bool
+
+
+#########################################
+# auto tuning configuration
+#########################################
+TUNING = "tuning"
+set_field_default_config(TUNING, "enable", False)
+set_field_default_config(TUNING, "profile_start_step", 1)
+set_field_default_config(TUNING, "profile_end_step", 1)
+set_field_default_config(TUNING, "run_after_tuning", True)
+set_field_default_config(TUNING, "debug", False)
+
+if TYPE_CHECKING:
+
+    class _TuningConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        profile_start_step: int
+        profile_end_step: int
+        run_after_tuning: bool
+        debug: bool
+
+
+#########################################
+# dataset configuration
+#########################################
+DATASET = "dataset"
+set_field_default_config(DATASET, "enable", False)
+set_field_default_config(DATASET, "num_shards", 1)
+
+if TYPE_CHECKING:
+
+    class _DatasetConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        num_shards: int
+
+
+# #########################################
+# # offload configuration
+# #########################################
+FUSEDLINEARPROMOTION = "fused_linear_promotion"
+set_field_default_config(FUSEDLINEARPROMOTION, "enable", False)
+
+if TYPE_CHECKING:
+
+    class _FusedLinearPromotionConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+
+
+#########################################
+# fused passes configuration
+#########################################
+FUSED_PASSES = "fused_passes"
+set_field_default_config(FUSED_PASSES, "enable", False)
+set_field_default_config(FUSED_PASSES, "fused_passes_list", [])
+
+if TYPE_CHECKING:
+
+    class _FusedPassesConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        fused_passes_list: list[str]
+
+
+#########################################
+# data parallel configuration
+#########################################
+DP_OPTIMIZATION = "dp_optimization"
+set_field_default_config(DP_OPTIMIZATION, "enable", False)
+set_field_default_config(DP_OPTIMIZATION, "fuse_all_reduce_ops", True)
+set_field_default_config(DP_OPTIMIZATION, "fuse_grad_size_in_MB", 32)
+set_field_default_config(DP_OPTIMIZATION, "overlap_comm_cacl", True)
+set_field_default_config(
+    DP_OPTIMIZATION, "gradient_sync_after_accumulate", False
+)
+
+if TYPE_CHECKING:
+
+    class _DPOptimizationConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool
+        fuse_all_reduce_ops: bool
+        fuse_grad_size_in_MB: int
+        overlap_comm_cacl: bool
+        gradient_sync_after_accumulate: bool
+
+
+#########################################
+# model parallel configuration
+#########################################
+MP_OPTIMIZATION = "mp_optimization"
+set_field_default_config(
+    MP_OPTIMIZATION, "allreduce_matmul_grad_overlapping", False
+)
+set_field_default_config(MP_OPTIMIZATION, "replace_with_c_embedding", False)
+
+set_field_default_config(
+    MP_OPTIMIZATION, "replace_with_parallel_cross_entropy", False
+)
+if TYPE_CHECKING:
+
+    class _MPOptimizationConfig(TypedDict, total=False):  # noqa: PYI049
+        allreduce_matmul_grad_overlapping: bool
+
+
+#########################################
+# sequence parallel configuration
+#########################################
+SP_OPTIMIZATION = "sp_optimization"
+set_field_default_config(SP_OPTIMIZATION, "enable", True)
+
+if TYPE_CHECKING:
+
+    class _SPOptimizationConfig(TypedDict, total=False):  # noqa: PYI049
+        enable: bool

python/paddle/distributed/auto_parallel/dygraph/__init__.py

@@ -0,0 +1,13 @@
+#   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.

python/paddle/distributed/auto_parallel/fully_shard.py

@@ -0,0 +1,199 @@
+# Copyright (c) 2025 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 os
+from types import MethodType
+
+import paddle
+import paddle.distributed as dist
+from paddle.autograd import PyLayer
+
+from .auto_dp_utils import in_auto_dp_mode
+from .fully_shard_fusion import FullyShardFusion
+
+
+def shard_accumulators(parameters_and_grads, optimizer, target_block):
+    if getattr(optimizer, "_has_sharded_accumulators", False):
+        return
+    optimizer._has_sharded_accumulators = True
+    for param, _ in parameters_and_grads:
+        optimizer._create_accumulators(
+            target_block,
+            [param],
+        )
+        target_name = param.name
+        if param.name in optimizer._master_weights.keys():
+            master_weight = optimizer._master_weights[param.name]
+            target_name = master_weight.name
+        for key in optimizer._accumulators.keys():
+            accumulator = optimizer._accumulators[key][target_name]
+            if accumulator.is_dist():
+                continue
+            origin_accumulator_name = accumulator.name
+
+            if 'beta' not in key:
+                placements = copy.deepcopy(param.placements)
+            else:
+                placements = [
+                    dist.Replicate()
+                    for _ in range(len(param.process_mesh.shape))
+                ]
+            optimizer._accumulators[key][target_name] = dist.shard_tensor(
+                accumulator,
+                mesh=param.process_mesh,
+                placements=placements,
+            )
+            optimizer._accumulators[key][
+                target_name
+            ].name = origin_accumulator_name
+
+    def _finish_update_impl(self, block, parameters_and_grads):
+        if not isinstance(parameters_and_grads, list):
+            parameters_and_grads = parameters_and_grads['params']
+        for param, _ in parameters_and_grads:
+            param.main_grad = None
+
+    optimizer._finish_update = MethodType(_finish_update_impl, optimizer)
+
+
+class FullyShardAuto:
+    def __init__(
+        self,
+        model,
+        mesh,
+        enable_tensor_fusion_and_overlap=True,
+        fsdp_unit_layers=None,
+        moe_layers_name=None,
+    ):
+        if enable_tensor_fusion_and_overlap:
+            FullyShardFusion(model, mesh, fsdp_unit_layers, moe_layers_name)
+        else:
+            self.model = model
+            self.mesh = mesh
+            # use first dims as sharding axis
+            self._shard_fn = dist.ShardingStage3(0, self.mesh)
+            for param in self.model.parameters():
+                param._need_shard_auto = True
+                self._shard_fn._shard_parameter(param)
+                if not in_auto_dp_mode():
+                    self._shard_fn._register_hook_for_param_grad(param)
+            if in_auto_dp_mode():
+                self._register_comm_hook(self.model)
+            os.environ["skip_sharding3_output_reshard"] = "1"
+
+    def _register_comm_hook(self, model):
+        def _pre_forward_hook(sublayers):
+            @paddle.autograd.no_grad()
+            def gather_comm(*_):
+                dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
+                for key, param in sublayers._parameters.items():
+                    if param.placements[dp_axis] != dist.Replicate():
+                        new_placements = copy.deepcopy(param.placements)
+                        new_placements[dp_axis] = dist.Replicate()
+                        replicate_param = dist.reshard(
+                            param, param.process_mesh, new_placements
+                        )
+                        param.get_tensor()._share_data_with(
+                            replicate_param.get_tensor()
+                        )
+
+            return gather_comm
+
+        def _post_forward_hook(sublayers):
+            @paddle.autograd.no_grad()
+            def shard_comm(*_):
+                dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
+                for key, param in sublayers._parameters.items():
+                    if (
+                        param.trainable
+                        and param.placements[dp_axis] == dist.Replicate()
+                    ):
+                        new_placements = copy.deepcopy(param.placements)
+                        new_placements[dp_axis] = dist.Shard(dp_axis)
+                        shard_param = dist.reshard(
+                            param, param.process_mesh, new_placements
+                        )
+                        param.get_tensor()._share_data_with(
+                            shard_param.get_tensor()
+                        )
+
+            return shard_comm
+
+        def _post_backward_hook(param):
+            def shard_comm(grad):
+                dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
+                if param.placements[dp_axis] == dist.Replicate():
+                    new_placements = copy.deepcopy(param.placements)
+                    new_placements[dp_axis] = dist.Shard(dp_axis)
+                    shard_param = dist.reshard(
+                        param, param.process_mesh, new_placements
+                    )
+                    param.get_tensor()._share_data_with(
+                        shard_param.get_tensor()
+                    )
+                return grad
+
+            param.register_hook(shard_comm)
+
+        # register forward hooks
+        for name, sublayers in model.named_sublayers(include_self=True):
+            sublayers.register_forward_pre_hook(_pre_forward_hook(sublayers))
+            sublayers.register_forward_post_hook(_post_forward_hook(sublayers))
+
+        # register backward hooks
+        for param in model.parameters():
+            if param.trainable:
+                _post_backward_hook(param)
+
+        # register layer hooks for param sync in tie weights
+        self._register_layer_hooks(model)
+
+    def _register_layer_hooks(self, layer, name="last_layer"):
+        def _forward_post_hook(layer, inputs, outputs):
+            return LayerHook.apply(
+                outputs,
+                layer=layer,
+            )
+
+        if layer.parameters(include_sublayers=False):
+            layer.register_forward_post_hook(_forward_post_hook)
+        for name, sub_layer in layer.named_children():
+            self._register_layer_hooks(sub_layer, name)
+
+
+class LayerHook(PyLayer):
+    @staticmethod
+    def forward(ctx, inputs, layer):
+        ctx.layer = layer
+        return inputs
+
+    @staticmethod
+    def backward(ctx, *args):
+        layer = ctx.layer
+        dp_axis = dist.auto_parallel.get_mesh().dim_names.index('dp')
+        for param in layer.parameters(include_sublayers=False):
+            if (
+                param.trainable
+                and param.placements[dp_axis] != dist.Replicate()
+            ):
+                new_placements = copy.deepcopy(param.placements)
+                new_placements[dp_axis] = dist.Replicate()
+                replicate_param = dist.reshard(
+                    param, param.process_mesh, new_placements
+                )
+                param.get_tensor()._share_data_with(
+                    replicate_param.get_tensor()
+                )
+        return args

python/paddle/distributed/auto_parallel/fully_shard_fusion.py

@@ -0,0 +1,809 @@
+# Copyright (c) 2026 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 collections import OrderedDict
+from dataclasses import dataclass, field
+from enum import Enum
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.autograd import PyLayer
+from paddle.distributed.fleet.utils.tensor_fusion_helper import (
+    align,
+    alignment,
+    get_current_device_type,
+)
+
+# Global registry for fsdp_context
+_g_fsdp_context = None
+
+
+def register_fsdp_context(context):
+    global _g_fsdp_context
+    _g_fsdp_context = context
+
+
+def get_fsdp_context():
+    return _g_fsdp_context
+
+
+class BufferState(Enum):
+    # Buffer status for lazy double buffer mechanism
+    #
+    # State transitions:
+    #     FREED ──all_gather──> USING ──computation done──> READY ──release──> FREED
+    #                             ^                            │
+    #                             │         (reuse)            │
+    #                             └────────────────────────────┘
+
+    FREED = 1  # Released, buffer data is sharded, tmp_buffer not allocated
+    USING = 2  # Unsharded and actively in use
+    READY = 3  # Unsharded, marked for lazy release, can be reused
+    SYNCING = 4  # Communication in progress
+
+
+@dataclass
+class BufferGroup:
+    params: list = field(default_factory=list)
+    dtype: object = None
+    trainable: bool = None
+    fsdp_unit_id: int = None
+    is_tie: bool = False
+    is_expert_param: bool = False
+    fsdp_group: object = None
+    params_buffer: 'TensorFusionBuffer' = None
+    grads_buffer: 'TensorFusionBuffer' = None
+    params_use_sum: int = 0
+    params_use_cnt: int = 0
+    grads_use_sum: int = 0
+    grads_use_cnt: int = 0
+
+
+def _dtensor_from_local(local_tensor, mesh, placements):
+    global_dims = list(local_tensor.shape)
+    for idx, placement in enumerate(placements):
+        if placement.is_shard():
+            global_dims[placement.get_dim()] = (
+                global_dims[placement.get_dim()] * mesh.shape[idx]
+            )
+    place = paddle.framework._current_expected_place()
+    place = paddle.framework._get_paddle_place(place)
+
+    return paddle.Tensor(
+        local_tensor,
+        dims=global_dims,
+        process_mesh=mesh,
+        placements=placements,
+        place=place,
+    )
+
+
+class TensorFusionBuffer:
+    def __init__(
+        self,
+        group_id,
+        params,
+        fsdp_degree,
+        dtype,
+        is_params=False,
+        main_grad_dtype=None,
+    ):
+        # Calculate total buffer size needed (with padding)
+        self.group_id = group_id
+        self.fsdp_degree = fsdp_degree
+        self.dtype = dtype
+        self.main_grad_dtype = (
+            main_grad_dtype if main_grad_dtype is not None else dtype
+        )
+        self.total_buffer_size = 0
+        self.param_offsets = {}
+        self.tmp_data_buffer = None
+        self.comm_task = None
+        self.trainable = params[0].trainable
+
+        for param in params:
+            self.param_offsets[param.name] = self.total_buffer_size
+            self.total_buffer_size += self.get_padded_size(param)
+
+        if is_params:
+            # Create fused params_buffer
+            # TODO(lizhenxing): Build full params_buffer on CPU and only move shards to GPU to minimize mem peaks
+            self.data_buffer = paddle.zeros(
+                shape=[self.total_buffer_size],
+                dtype=dtype,
+            )
+            # Use BufferState enum instead of is_shard boolean, initial state is FREED (sharded)
+            self.status = BufferState.FREED
+
+            for param in params:
+                offset = self.param_offsets[param.name]
+                stop_gradient = param.stop_gradient
+                local_shape = param._local_shape
+                param.stop_gradient = True
+                param._local_value().flatten_()
+                paddle.assign(
+                    param._local_value(),
+                    self.data_buffer._slice(
+                        offset,
+                        offset + param._numel(),
+                    ),
+                )
+
+                param._clear_data()
+                param.stop_gradient = stop_gradient
+                param._local_value().get_tensor()._set_dims(local_shape)
+            paddle.device.cuda.empty_cache()
+
+            mesh = dist.auto_parallel.get_mesh()
+            curr_global_rank = paddle.distributed.get_rank()
+            if curr_global_rank in mesh.process_ids:
+                total_nums = self.data_buffer.shape[0]
+                num_of_pieces = mesh.shape[0]
+                piece_len = (total_nums + num_of_pieces - 1) // num_of_pieces
+                rank_relative = mesh.process_ids.index(curr_global_rank)
+                start = rank_relative * piece_len
+                end = min(start + piece_len, total_nums)
+                self.data_buffer = paddle.slice(
+                    self.data_buffer, [0], [start], [end]
+                ).clone()
+
+            # Init params_buffer attr
+            self.data_buffer.name = "fuse_params_" + str(group_id)
+            self.data_buffer.stop_gradient = params[0].stop_gradient
+            self.data_buffer.optimize_attr = params[0].optimize_attr
+        else:
+            # Create fused grads_buffer with shard
+            self.data_buffer = paddle.zeros(
+                shape=[self.total_buffer_size // self.fsdp_degree],
+                dtype=self.main_grad_dtype,
+            )
+
+            # Register get_main_grad method for each param, returns view_slice of grad_buffer
+            for param in params:
+                if param.trainable:
+                    param._fusion_buffer = self
+                    param._param_offsets = self.param_offsets
+
+                    def get_grad_from_tmp_buf(param):
+                        tmp_buffer = param._fusion_buffer.get_tmp_buffer()
+                        offset = param._param_offsets[param.name]
+                        main_grad = paddle._C_ops.view_slice(
+                            tmp_buffer,
+                            offset,
+                            offset + param._numel(),
+                        )
+                        return main_grad
+
+                    param.get_main_grad = get_grad_from_tmp_buf.__get__(param)
+
+    def get_padded_size(self, param):
+        size = np.prod(param.shape)
+        align_size = (
+            alignment[get_current_device_type()]
+            // align[param.dtype]
+            * self.fsdp_degree
+        )
+        return ((size + align_size - 1) // align_size) * align_size
+
+    def get_tmp_buffer(self):
+        # Reuse tmp_buffer if exists, else create
+        if self.tmp_data_buffer is None:
+            self.tmp_data_buffer = paddle.zeros(
+                shape=[self.total_buffer_size], dtype=self.dtype
+            )
+        return self.tmp_data_buffer
+
+    def clear_tmp_buffer(self):
+        if self.tmp_data_buffer is not None:
+            self.tmp_data_buffer._clear_data()
+            self.tmp_data_buffer = None
+            # paddle.device.cuda.empty_cache()
+
+
+class FSDPBufferManager:
+    def __init__(
+        self, model, mesh, fsdp_unit_layers=None, moe_layers_name=None
+    ):
+        self.model = model
+        self._fsdp_group = mesh.get_group("dp")
+        self.main_grad_dtype = paddle.float32
+
+        # Get EP group if "ep" dimension exists in mesh
+        if "ep" in mesh.dim_names:
+            self._ep_fsdp_group = mesh.get_group("ep")
+        else:
+            self._ep_fsdp_group = self._fsdp_group
+
+        topk = None
+        if hasattr(self.model, 'config') and hasattr(
+            self.model.config, 'num_experts_per_tok'
+        ):
+            topk = self.model.config.num_experts_per_tok
+
+        # Layer types to wrap as FSDP sharding layers
+        # Note: 'Qwen3VLTextDecoderLayer' is temporary; fleet models all use 'TransformerLayer'
+        self.fsdp_unit_layers = fsdp_unit_layers or [
+            'TransformerLayer',
+            'Qwen3VLTextDecoderLayer',
+            'Qwen3MoeDecoderLayer',
+        ]
+        # Layer types to identify MoE expert layers
+        self.moe_layers_name = moe_layers_name or [
+            'StandardMLPExpert',
+        ]
+
+        # Get tie_param_name if using tie_weights
+        self.tie_param_name = None
+        # Note: need add get_input_embeddings in fleet modeling
+        # if hasattr(self.model, "get_input_embeddings"):
+        #     self.tie_param_name = self.model.get_input_embeddings().weight.name
+
+        # Create buffer_groups
+        grouped_params, group_is_expert = self._build_groups()
+        self.buffer_groups = []
+        self.param_to_buffer_id = {}
+
+        # Create params_buffer, grads_buffer with groups
+        for gid, params in grouped_params.items():
+            is_expert = group_is_expert.get(gid, False)
+            # Use EP group for expert params, DP group for regular params
+            fsdp_group = self._ep_fsdp_group if is_expert else self._fsdp_group
+
+            params_buffer = TensorFusionBuffer(
+                gid,
+                params,
+                fsdp_group.nranks,
+                params[0].dtype,
+                is_params=True,
+            )
+
+            if not params[0].stop_gradient:
+                grads_buffer = TensorFusionBuffer(
+                    gid,
+                    params,
+                    fsdp_group.nranks,
+                    params[0].dtype,
+                    main_grad_dtype=self.main_grad_dtype,
+                )
+            else:
+                grads_buffer = None
+
+            if is_expert:
+                _params_use_sum = topk
+                _grads_use_sum = topk
+            else:
+                _params_use_sum = len(params)
+                _grads_use_sum = len(params)
+            self.buffer_groups.append(
+                BufferGroup(
+                    params=params,
+                    dtype=params[0].dtype,
+                    trainable=params[0].trainable,
+                    is_expert_param=is_expert,
+                    fsdp_group=fsdp_group,
+                    params_buffer=params_buffer,
+                    grads_buffer=grads_buffer,
+                    params_use_sum=_params_use_sum,
+                    params_use_cnt=0,
+                    grads_use_sum=_grads_use_sum,
+                    grads_use_cnt=0,
+                )
+            )
+
+            for param in params:
+                self.param_to_buffer_id[param.name] = gid
+
+    def _build_groups(self):
+        parameters = self.model.parameters()
+        grouped_params = OrderedDict()
+        group_is_expert = {}
+        curr_gid = 0
+
+        param_to_unit_id = {}
+        for unit_id, module in enumerate(self.model.modules()):
+            if type(module).__name__ in self.fsdp_unit_layers:
+                for param in module.parameters():
+                    param_to_unit_id[param.name] = unit_id
+            if type(module).__name__ in self.moe_layers_name:
+                for param in module.parameters():
+                    param.is_moe_param = True
+
+        temp_groups = []
+        for param in parameters:
+            name = param.name
+            is_expert = getattr(param, "is_moe_param", False)
+            if is_expert:
+                continue
+            is_tie = (
+                self.tie_param_name is not None and name == self.tie_param_name
+            )
+
+            param_attrs = {
+                "dtype": param.dtype,
+                "trainable": param.trainable,
+                "fsdp_unit_id": param_to_unit_id.get(name),
+                "is_tie": is_tie,
+                "is_expert_param": is_expert,
+            }
+
+            found_group = False
+            for param_group in temp_groups:
+                if (
+                    param_group.dtype == param_attrs["dtype"]
+                    and param_group.trainable == param_attrs["trainable"]
+                    and param_group.fsdp_unit_id == param_attrs["fsdp_unit_id"]
+                    and param_group.is_tie == param_attrs["is_tie"]
+                    and param_group.is_expert_param
+                    == param_attrs["is_expert_param"]
+                ):
+                    param_group.params.append(param)
+                    found_group = True
+                    break
+
+            # Create new group if no matching
+            if not found_group:
+                temp_groups.append(BufferGroup(params=[param], **param_attrs))
+
+        def group_sort_key(group):
+            priority = 0 if group.is_tie else (1 if not group.trainable else 2)
+            return (
+                priority,
+                group.fsdp_unit_id
+                if group.fsdp_unit_id is not None
+                else float('inf'),
+            )
+
+        sorted_groups = sorted(temp_groups, key=group_sort_key)
+
+        # For each sorted parameter group, buffer them by execution order
+        for param_group in sorted_groups:
+            cur_params = param_group.params
+            if len(cur_params) == 0:
+                continue
+            for p in cur_params:
+                grouped_params.setdefault(curr_gid, []).append(p)
+            group_is_expert[curr_gid] = param_group.is_expert_param
+            curr_gid += 1
+
+        return grouped_params, group_is_expert
+
+
+class FSDPCommManager:
+    def __init__(
+        self,
+        buffer_manager,
+        enable_overlap=True,
+        double_buffer_limit=2,
+    ):
+        self.buffer_manager = buffer_manager
+        self.enable_overlap = enable_overlap
+        self.grad_reduce_queue = []
+
+        # for double buffer mechanism config
+        self.double_buffer_limit = double_buffer_limit
+        self.buffer_cnt_in_using = 0
+        self._need_zero_grads = True
+
+    def _release_one_buffer_if_needed(self):
+        # Release a buffer with the READY status if needed
+        while self.buffer_cnt_in_using >= self.double_buffer_limit:
+            found = False
+            for group in self.buffer_manager.buffer_groups:
+                if group.params_buffer.status == BufferState.READY:
+                    group.params_buffer.status = BufferState.FREED
+                    group.params_buffer.clear_tmp_buffer()
+                    self.buffer_cnt_in_using -= 1
+                    found = True
+                    break
+            if not found:
+                break
+
+    def _next_buffer_id(self, gid, is_backward):
+        # Get next buffer id for prefetch
+        if is_backward:
+            next_gid = gid - 1
+            # Search backward for trainable buffer_groups
+            while (
+                next_gid >= 0
+                and not self.buffer_manager.buffer_groups[
+                    next_gid
+                ].params_buffer.trainable
+            ):
+                next_gid -= 1
+            return max(next_gid, 0)
+        else:
+            return min(gid + 1, len(self.buffer_manager.buffer_groups) - 1)
+
+    def all_gather_params(self, params, is_backward=False):
+        if len(params) == 0:
+            return
+        for param in params:
+            if hasattr(param, "is_moe_param"):
+                continue
+            gid = self.buffer_manager.param_to_buffer_id[param.name]
+            group = self.buffer_manager.buffer_groups[gid]
+            group.params_use_cnt += 1
+            params_buffer = group.params_buffer
+            # Use group-specific fsdp_group
+            fsdp_group = group.fsdp_group or self.buffer_manager._fsdp_group
+
+            # Double buffer: reuse buffer if status is READY
+            if params_buffer.status == BufferState.READY:
+                # Reuse: READY -> USING, no need to all_gather again
+                params_buffer.status = BufferState.USING
+
+            # Overlap prefetch comm
+            if self.enable_overlap:
+                next_gid = self._next_buffer_id(gid, is_backward)
+                next_group = self.buffer_manager.buffer_groups[next_gid]
+                next_params_buffer = next_group.params_buffer
+                next_fsdp_group = (
+                    next_group.fsdp_group or self.buffer_manager._fsdp_group
+                )
+                if next_params_buffer.status == BufferState.FREED:
+                    # Check double_buffer_limit before prefetch
+                    self._release_one_buffer_if_needed()
+                    next_params_buffer.status = BufferState.SYNCING
+                    tmp_buffer_prefetch = next_params_buffer.get_tmp_buffer()
+                    next_params_buffer.comm_task = (
+                        paddle.distributed.all_gather(
+                            tmp_buffer_prefetch,
+                            next_params_buffer.data_buffer,
+                            group=next_fsdp_group,
+                            sync_op=False,
+                        )
+                    )
+                    self.buffer_cnt_in_using += 1
+
+            # Wait for async comm to complete: SYNCING -> USING
+            if params_buffer.status == BufferState.SYNCING:
+                params_buffer.status = BufferState.USING
+                params_buffer.comm_task.wait()
+                params_buffer.comm_task = None
+
+            tmp_buffer = params_buffer.get_tmp_buffer()
+            # Do all_gather in sync: FREED -> USING
+            if params_buffer.status == BufferState.FREED:
+                fsdp_group.process_group.all_gather(
+                    params_buffer.data_buffer, tmp_buffer
+                ).wait()
+                params_buffer.status = BufferState.USING
+                self.buffer_cnt_in_using += 1
+
+            # Bind the unsharded param to the real param
+            offset = params_buffer.param_offsets[param.name]
+            tmp_param = paddle._C_ops.view_slice(
+                tmp_buffer,
+                offset,
+                offset + param._numel(),
+            )
+            tmp_param.get_tensor()._set_dims(param.shape)
+            tmp_param = _dtensor_from_local(
+                tmp_param,
+                param.process_mesh,
+                param.placements,
+            )
+            param.get_tensor()._share_data_with(tmp_param.get_tensor())
+
+    def shard_params(self, params, is_backward=False):
+        affected_gids = set()
+        for param in params:
+            if hasattr(param, "is_moe_param"):
+                continue
+            gid = self.buffer_manager.param_to_buffer_id.get(param.name)
+
+            group = self.buffer_manager.buffer_groups[gid]
+            stop_gradient = param.stop_gradient
+            local_shape = param._local_shape
+            param._clear_data()
+            param.stop_gradient = stop_gradient
+            param._local_value().get_tensor()._set_dims(local_shape)
+
+            affected_gids.add(gid)
+
+        for gid in affected_gids:
+            group = self.buffer_manager.buffer_groups[gid]
+            if group.params_buffer.status == BufferState.USING:
+                group.params_buffer.status = BufferState.READY
+
+    def reduce_scatter_grads(self, param):
+        if self._need_zero_grads:
+            self._need_zero_grads = False
+            for group in self.buffer_manager.buffer_groups:
+                if group.grads_buffer is not None:
+                    group.grads_buffer.data_buffer.zero_()
+        gid = self.buffer_manager.param_to_buffer_id.get(param.name)
+        group = self.buffer_manager.buffer_groups[gid]
+        group.grads_use_cnt += 1
+        fsdp_group = group.fsdp_group or self.buffer_manager._fsdp_group
+        param.main_grad = None
+
+        if group.grads_use_cnt == group.grads_use_sum:
+            group.grads_use_cnt = 0
+
+            # reduce_scatter from tmp_grad_buffer into grads_buffer
+            grads_buffer = group.grads_buffer
+
+            # Grad queue mechanism: wait and release completed reduce_scatter async tasks
+            self._wait_for_grad_comm()
+
+            tmp_buffer = grads_buffer.get_tmp_buffer()
+            shard_size = grads_buffer.data_buffer.shape[0]
+            grad_buffer_shard = tmp_buffer._slice(0, shard_size)
+            if self.enable_overlap:
+                # Comm grads async and check all comm_task before optimizer update
+                grads_buffer.comm_task = paddle.distributed.reduce_scatter(
+                    grad_buffer_shard,
+                    tmp_buffer,
+                    op=paddle.distributed.ReduceOp.SUM,
+                    group=fsdp_group,
+                    sync_op=False,
+                )
+
+                # Add async task to queue
+                self.grad_reduce_queue.append(grads_buffer)
+            else:
+                paddle.distributed.reduce_scatter(
+                    grad_buffer_shard,
+                    tmp_buffer,
+                    op=paddle.distributed.ReduceOp.SUM,
+                    group=fsdp_group,
+                    sync_op=False,
+                ).wait()
+                grads_buffer.data_buffer.add_(grad_buffer_shard)
+                grads_buffer.clear_tmp_buffer()
+
+    def _wait_for_grad_comm(self, queue_limit=2):
+        # Wait for async reduce_scatter tasks to complete and release resources
+        # queue_limit: max queue size, default use 2, 0 means wait for all
+        while len(self.grad_reduce_queue) > queue_limit:
+            grads_buffer = self.grad_reduce_queue.pop(0)
+            if grads_buffer.comm_task is not None:
+                grads_buffer.comm_task.wait()
+                grads_buffer.comm_task = None
+                tmp_buffer = grads_buffer.get_tmp_buffer()
+                shard_size = grads_buffer.data_buffer.shape[0]
+                grad_buffer_shard = tmp_buffer._slice(0, shard_size)
+                grads_buffer.data_buffer.add_(grad_buffer_shard)
+            grads_buffer.clear_tmp_buffer()
+
+    def _finish_grads_sync(self):
+        # Wait for all async reduce_scatter tasks, call before optimizer.step()
+        self._wait_for_grad_comm(queue_limit=0)
+
+    def _reset_params_buffer_status(self):
+        for group in self.buffer_manager.buffer_groups:
+            params_buffer = group.params_buffer
+            if params_buffer.status in (BufferState.READY, BufferState.USING):
+                # Clear stale tmp_buffer to force re-all_gather with updated data_buffer
+                params_buffer.clear_tmp_buffer()
+                params_buffer.status = BufferState.FREED
+                if self.buffer_cnt_in_using > 0:
+                    self.buffer_cnt_in_using -= 1
+
+
+class FusionBackwardHook(PyLayer):
+    @staticmethod
+    def forward(ctx, *inputs, layer, comm_manager, recursive=False):
+        ctx.layer = layer
+        ctx.comm_manager = comm_manager
+        ctx.recursive = recursive
+        return inputs if len(inputs) > 1 else inputs[0]
+
+    @staticmethod
+    def backward(ctx, *args):
+        trainable_params = []
+
+        for param in ctx.layer.parameters(include_sublayers=ctx.recursive):
+            if param.trainable:
+                trainable_params.append(param)
+
+        ctx.comm_manager.all_gather_params(trainable_params, is_backward=True)
+        return args
+
+
+class FusionForwardHook(PyLayer):
+    @staticmethod
+    def forward(ctx, *inputs, layer, comm_manager, recursive=False):
+        ctx.layer = layer
+        ctx.comm_manager = comm_manager
+        ctx.recursive = recursive
+        return inputs
+
+    @staticmethod
+    def backward(ctx, *args):
+        ctx.comm_manager.shard_params(
+            ctx.layer.parameters(include_sublayers=ctx.recursive),
+            is_backward=True,
+        )
+        return args
+
+
+class FullyShardFusion:
+    def __init__(
+        self, model, mesh, fsdp_unit_layers=None, moe_layers_name=None
+    ):
+        self.model = model
+        self.mesh = self._check_mesh(mesh)
+        self._shard_all_params()
+        self.buffer_manager = FSDPBufferManager(
+            self.model, self.mesh, fsdp_unit_layers, moe_layers_name
+        )
+        self.comm_manager = FSDPCommManager(self.buffer_manager)
+        self.register_tensor_fusion_hooks(self.model)
+        register_fsdp_context(self)
+
+    def _check_mesh(self, mesh, pp_idx=0):
+        if "pp" in mesh.dim_names:
+            mesh = mesh.get_mesh_with_dim("pp", pp_idx)
+        return mesh
+
+    def _shard_all_params(self):
+        def shard_layer_param(layer):
+            for param_name in list(layer._parameters.keys()):
+                param = getattr(layer, param_name)
+                if param is not None:
+                    param_placements = [
+                        dist.Replicate() for _ in range(len(self.mesh.shape))
+                    ]
+                    if not param.is_dist():
+                        param = dist.shard_tensor(
+                            param, self.mesh, param_placements
+                        )
+                        setattr(layer, param_name, param)
+
+        for name, layer in self.model.named_sublayers(include_self=True):
+            shard_layer_param(layer)
+
+    def comm_sync_and_reset_status(self):
+        self.comm_manager._finish_grads_sync()
+        self.comm_manager._reset_params_buffer_status()
+        self.comm_manager._need_zero_grads = True
+        # Reset main_grad for all trainable parameters
+        for param in self.model.parameters():
+            if param.trainable:
+                param.main_grad = None
+
+    def register_tensor_fusion_hooks(self, model):
+        def _pre_forward_hook(sublayers, recursive=False):
+            comm_manager = self.comm_manager
+
+            @paddle.autograd.no_grad()
+            def all_gather_comm(*_):
+                comm_manager.all_gather_params(
+                    sublayers.parameters(include_sublayers=recursive)
+                )
+
+            return all_gather_comm
+
+        def _post_forward_hook(sublayers, recursive=False):
+            comm_manager = self.comm_manager
+
+            @paddle.autograd.no_grad()
+            def shard_comm(*_):
+                comm_manager.shard_params(
+                    sublayers.parameters(include_sublayers=recursive)
+                )
+
+            return shard_comm
+
+        def _update_main_grad_hook(param):
+            comm_manager = self.comm_manager
+
+            @paddle.autograd.no_grad()
+            def comm_hook(grad):
+                if grad is not None and grad._is_initialized():
+                    # Share mem with grads_tmp_buffer
+                    _main_grad = param.get_main_grad()
+                    _main_grad.get_tensor()._set_dims(grad._local_shape)
+                    param.main_grad = _dtensor_from_local(
+                        _main_grad,
+                        grad.process_mesh,
+                        grad.placements,
+                    )
+                    param.main_grad._local_value().copy_(grad._local_value())
+                    grad._clear_data()
+                comm_manager.shard_params([param], is_backward=True)
+                comm_manager.reduce_scatter_grads(param)
+
+            return comm_hook
+
+        def _post_backward_hook(param):
+            param.main_grad = None
+            if hasattr(param, "get_main_grad"):
+                param._register_grad_hook(_update_main_grad_hook(param))
+
+        for param in model.parameters():
+            if param.trainable:
+                _post_backward_hook(param)
+
+        def _register_recursive(layer):
+            is_unit = (
+                type(layer).__name__ in self.buffer_manager.fsdp_unit_layers
+            )
+
+            if is_unit:
+                # For FSDP Unit, register recursive hooks and stop recursion
+                layer.register_forward_pre_hook(
+                    _pre_forward_hook(layer, recursive=True)
+                )
+                layer.register_forward_post_hook(
+                    _post_forward_hook(layer, recursive=True)
+                )
+                self._register_fusion_layer_hooks(layer, recursive=True)
+                return
+
+            if layer.parameters(include_sublayers=False):
+                layer.register_forward_pre_hook(
+                    _pre_forward_hook(layer, recursive=False)
+                )
+                layer.register_forward_post_hook(
+                    _post_forward_hook(layer, recursive=False)
+                )
+                self._register_fusion_layer_hooks(layer, recursive=False)
+
+            for child in layer.children():
+                _register_recursive(child)
+
+        _register_recursive(model)
+
+    def _register_fusion_layer_hooks(self, layer, recursive=False):
+        def _forward_post_hook(layer, inputs, outputs):
+            if isinstance(outputs, dict):
+                for key, value in outputs.items():
+                    if (
+                        isinstance(value, paddle.Tensor)
+                        and not value.stop_gradient
+                    ):
+                        outputs[key] = FusionBackwardHook.apply(
+                            value,
+                            layer=layer,
+                            comm_manager=self.comm_manager,
+                            recursive=recursive,
+                        )
+                return outputs
+            elif isinstance(outputs, tuple):
+                result = FusionBackwardHook.apply(
+                    *outputs,
+                    layer=layer,
+                    comm_manager=self.comm_manager,
+                    recursive=recursive,
+                )
+                if not isinstance(result, tuple):
+                    result = (result,)
+                return result
+            else:
+                return FusionBackwardHook.apply(
+                    outputs,
+                    layer=layer,
+                    comm_manager=self.comm_manager,
+                    recursive=recursive,
+                )
+
+        def _forward_pre_hook(layer, inputs):
+            return FusionForwardHook.apply(
+                *inputs,
+                layer=layer,
+                comm_manager=self.comm_manager,
+                recursive=recursive,
+            )
+
+        layer.register_forward_post_hook(_forward_post_hook)
+
+        # Register an additional hook for tie_weights shard_params
+        for param in layer.parameters(include_sublayers=False):
+            if param.name == self.comm_manager.buffer_manager.tie_param_name:
+                layer.register_forward_pre_hook(_forward_pre_hook)

python/paddle/distributed/auto_parallel/high_level_api.py

@@ -0,0 +1,933 @@
+# 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 __future__ import annotations
+
+import logging
+import math
+import warnings
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.framework import in_dygraph_mode
+
+logger = logging.getLogger(__name__)
+
+
+class ToDistributedConfig:
+    def __init__(self):
+        self.input_spec = None
+        self.sequence_parallel = False
+
+
+def cost_model(matched_programs, device_num, node_num):
+    # TODO(jeff41404): multi-node will be supported later
+    assert node_num == 1, (
+        "we only support single node now, multi-node will be supported later"
+    )
+
+    # TODO(jeff41404): will evaluate the best combination of parallel strategies
+    # based on cost_model and return global_mesh, currently using pre-defined parallel strategy
+    if device_num % 2 == 0:
+        if device_num == 8:
+            return dist.ProcessMesh(
+                np.arange(device_num).reshape(2, 2, 2).tolist(),
+                dim_names=["pp", "dp", "mp"],
+            )
+        elif device_num == 6:
+            return dist.ProcessMesh(
+                np.arange(device_num).reshape(3, 2).tolist(),
+                dim_names=["dp", "mp"],
+            )
+        elif device_num == 4:
+            return dist.ProcessMesh(
+                np.arange(device_num).reshape(2, 2).tolist(),
+                dim_names=["dp", "mp"],
+            )
+        elif device_num == 2:
+            return dist.ProcessMesh(list(range(device_num)), dim_names=["dp"])
+        else:
+            raise ValueError(
+                f"device_num must be an even number to be able to use at least 2 parallel strategies, but got: {device_num}"
+            )
+    else:
+        logger.debug(
+            f'device_num must be an even number to be able to use at least 2 parallel strategies, but got: {device_num}, only use data parallel.'
+        )
+        return dist.ProcessMesh(list(range(device_num)), dim_names=["dp"])
+
+
+def record_program_ops_pre_hook(layer, inputs):
+    """
+    A pre-hook to mark op numbers before enter layer.forward.
+    """
+    if not in_dygraph_mode():
+        # Because ir_guard._switch_to_pir() will change default_main_program in python/paddle/__init__.py.
+        # In order to avoid errors, we import default_main_program until this hook running.
+        # After fully switching to pir, can move this import to the beginning of the file.
+        from paddle.base import default_main_program
+
+        if layer._op_recorder.start < 0:
+            layer._op_recorder.start = len(
+                default_main_program().global_block().ops
+            )
+            layer._op_recorder.is_valid = True
+        else:
+            layer._op_recorder.is_valid = False
+            warnings.warn(
+                f"{layer._full_name} has recorded the op information before. Please check whether you call this layer twice."
+            )
+
+
+def transpose_reshard_embedding_layer_output(layer, inputs, outputs):
+    if hasattr(layer, "current_mesh"):
+        current_mesh = layer.__getattr__("current_mesh")
+        new_output = paddle.transpose(outputs, [1, 0, 2])
+        new_output = dist.reshard(
+            new_output, current_mesh, [dist.Shard(1), dist.Shard(0)]
+        )
+        return new_output
+
+
+def reshard_transpose_attention_layer_input(layer, inputs):
+    new_inputs = list(inputs)
+    x = new_inputs[0]
+    if hasattr(layer, "current_mesh"):
+        current_mesh = layer.__getattr__("current_mesh")
+        new_x = dist.reshard(x, current_mesh, [dist.Shard(1), dist.Replicate()])
+        new_x = paddle.transpose(new_x, [1, 0, 2])
+        new_inputs[0] = new_x
+        return tuple(new_inputs)
+
+
+def transpose_reshard_attention_layer_output(layer, inputs, outputs):
+    attn_out = outputs
+    if hasattr(layer, "current_mesh"):
+        current_mesh = layer.__getattr__("current_mesh")
+        new_attn_out = paddle.transpose(attn_out, [1, 0, 2])
+        new_attn_out = dist.reshard(
+            new_attn_out, current_mesh, [dist.Shard(1), dist.Shard(0)]
+        )
+        return new_attn_out
+
+
+def reshard_mlp_layer_input(layer, inputs):
+    new_inputs = list(inputs)
+    mlp_input = new_inputs[0]
+    if hasattr(layer, "current_mesh"):
+        current_mesh = layer.__getattr__("current_mesh")
+        new_mlp_input = dist.reshard(
+            mlp_input, current_mesh, [dist.Shard(1), dist.Replicate()]
+        )
+        new_inputs[0] = new_mlp_input
+        return tuple(new_inputs)
+
+
+def reshard_mlp_layer_output(layer, inputs, outputs):
+    mlp_out = outputs
+    if hasattr(layer, "current_mesh"):
+        current_mesh = layer.__getattr__("current_mesh")
+        new_mlp_out = dist.reshard(
+            mlp_out, current_mesh, [dist.Shard(1), dist.Shard(0)]
+        )
+        return new_mlp_out
+
+
+def reshard_transpose_rms_norm_layer_output(layer, inputs, outputs):
+    if hasattr(layer, "current_mesh"):
+        current_mesh = layer.__getattr__("current_mesh")
+        new_output = dist.reshard(
+            outputs, current_mesh, [dist.Shard(1), dist.Replicate()]
+        )
+        new_output = paddle.transpose(new_output, [1, 0, 2])
+        return new_output
+
+
+def reshard_all_inputs(layer, inputs):
+    if hasattr(layer, "current_mesh"):
+        current_mesh = layer.__getattr__("current_mesh")
+        if type(inputs) is tuple:
+            new_inputs = []
+            for input in inputs:
+                if paddle.is_tensor(input):
+                    if input.is_dist():
+                        new_input = dist.reshard(
+                            input,
+                            current_mesh,
+                            input.placements,
+                        )
+                    else:
+                        new_input = dist.shard_tensor(
+                            input,
+                            current_mesh,
+                            [dist.Shard(0), dist.Replicate()],
+                        )
+                    new_inputs.append(new_input)
+                else:
+                    new_inputs.append(input)
+            return tuple(new_inputs)
+        else:
+            if input.is_dist():
+                new_input = dist.reshard(
+                    input, current_mesh, [dist.Shard(0), dist.Replicate()]
+                )
+            else:
+                new_input = dist.shard_tensor(
+                    input, current_mesh, [dist.Shard(0), dist.Replicate()]
+                )
+            return new_input
+
+
+def reshard_all_outputs(layer, inputs, outputs):
+    if hasattr(layer, "next_mesh"):
+        next_mesh = layer.__getattr__("next_mesh")
+        if type(outputs) is tuple:
+            new_outputs = []
+            for output in outputs:
+                if paddle.is_tensor(output):
+                    new_output = dist.reshard(
+                        output, next_mesh, [dist.Shard(0), dist.Replicate()]
+                    )
+                    new_outputs.append(new_output)
+                else:
+                    new_outputs.append(output)
+            return new_outputs
+        else:
+            new_output = dist.reshard(
+                outputs, next_mesh, [dist.Shard(0), dist.Replicate()]
+            )
+            return new_output
+
+
+def record_program_ops_post_hook(layer, inputs, outputs):
+    """
+    A post-hook to mark op numbers after enter layer.forward, and record corresponding ops of the layer.
+    """
+    if not in_dygraph_mode():
+        # Because ir_guard._switch_to_pir() will change default_main_program in python/paddle/__init__.py.
+        # In order to avoid errors, we import default_main_program until this hook running.
+        # After fully switching to pir, can move this import to the beginning of the file.
+        from paddle.base import default_main_program
+
+        assert (
+            layer._op_recorder.start >= 0
+            and layer._op_recorder.is_valid is True
+        ), (
+            f"{layer._full_name} has not recorded the start of the corresponding ops before"
+        )
+        end = len(default_main_program().global_block().ops)
+        # some layers, such as rotary_embedding, will not add new ops to program
+        # assert end > layer._op_recorder.start, f"{layer._full_name} has not added new ops to the program"
+        ops = []
+        if end > layer._op_recorder.start:
+            layer._op_recorder.end = end
+            ops = (
+                default_main_program()
+                .global_block()
+                .ops[layer._op_recorder.start : layer._op_recorder.end]
+            )
+        logger.debug(
+            f'start: {layer._op_recorder.start}, end: {layer._op_recorder.end}, ops: {ops}'
+        )
+        layer._op_recorder.ops = ops
+
+
+def get_layer_pp_info(mesh, num_hidden_layers, layer_index):
+    if "pp" in mesh.dim_names:
+        pp_degree = mesh.get_dim_size("pp")
+        layer_per_stage = math.ceil(num_hidden_layers / pp_degree)
+        return layer_index // layer_per_stage
+    else:
+        # return None, False
+        return None
+
+
+def to_distributed(
+    model: paddle.nn.Layer,
+    optimizer: paddle.optimizer.Optimizer,
+    dataloader: paddle.io.DataLoader,
+    device_num: int,
+    node_num: int | None = 1,
+    config: ToDistributedConfig | None = None,
+) -> tuple[
+    paddle.nn.Layer,
+    paddle.optimizer.Optimizer,
+    paddle.distributed.auto_parallel.ShardDataloader,
+]:
+    """
+    `to_distributed` can automatically convert neural networks, optimizer, and dataloader
+    that do not contain any distributed code into neural networks, optimizers, and dataloader
+    that are suitable for distributed training and ensure their correctness.
+    At the same time, during the transformation process, the optimal distributed strategy
+    will be automatically selected based on `node_num` and `device_num` to maximize performance.
+
+    Args:
+        model(paddle.nn.Layer): The model in dygraph mode, whose parameters
+            are ordinary tensors, do not contain any distributed code.
+            If one device has sufficient memory, it can train directly.
+        optimizer(paddle.optimizer.Optimizer): The optimizer for training.
+            one instance of a regular optimizer, e.g. `paddle.optimizer.Adam` etc.
+        dataloader(paddle.io.DataLoader): The dataloader used in dygraph mode,
+            It is instantiated through regular `paddle.io.Dataset` and `paddle.io.Sampler`,
+            not `paddle.io.DistributedBatchSampler`.
+        device_num(int): the number of devices on each node or machine.
+        node_num(int|None, optional): the number of nodes or machines.
+        config(ToDistributedConfig| None = None): Configs for input_spec and sequence_parallel.
+            The custom input specs specify the most likely shape, dtype, and name information
+            of each model inputs. If it is not None, the input specs and
+            will be inferred from the custom input specs. If it is None, will use default with
+            shape of [BATCH_SIZE=4, SEQ_LENGTH=1024], The custom
+            input specs should be a list of `paddle.static.InputSpec`. Default: None.
+            sequence_parallel indicates whether to use sequence parallel. Default: False.
+
+    Returns:
+        model. The model in dygraph mode but contain distributed attributes.
+
+        optimizer. The optimizer for training and may be sharded states.
+
+        dataloader. The dataloader can be used in distributed training.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('run in distributed env')
+            >>> import math
+            >>> import numpy as np
+            >>> import paddle
+            >>> import paddle.nn.functional as F
+            >>> from paddle import nn
+            >>> from paddle.distributed import to_distributed
+            >>> from paddle.distributed.auto_parallel.high_level_api import ToDistributedConfig
+
+            >>> EPOCHS = 1
+            >>> VOCAB_SIZE = 8000
+            >>> BATCH_NUM = 2
+            >>> BATCH_SIZE = 4
+            >>> HIDDEN_SIZE = 2048
+            >>> INTERMEDIATE_SIZE = 4096
+            >>> SEQ_LENGTH = 1024
+            >>> N_HEAD = 32
+            >>> NUM_HIDDEN_LAYERS = 4
+            >>> class RandomDataset(paddle.io.Dataset):  # type: ignore[type-arg]
+            ...     def __init__(self, inputs, labels, num_samples):
+            ...         self.inputs = inputs
+            ...         self.labels = labels
+            ...         self.num_samples = num_samples
+            ...
+            ...     def __getitem__(self, idx):
+            ...         return self.inputs[idx], self.labels[idx]
+            ...
+            ...     def __len__(self):
+            ...         return self.num_samples
+
+            >>> class RotaryEmbedding(nn.Layer):
+            ...     def __init__(self, dim, max_position_embeddings=2048, base=10000):
+            ...         super().__init__()
+            ...         self.dim = dim
+            ...         self.max_position_embeddings = max_position_embeddings
+            ...         self.base = base
+            ...         self.inv_freq = 1.0 / (self.base ** (paddle.cast(paddle.arange(0, self.dim, 2), dtype="float32") / self.dim))
+            ...         self._set_cos_sin_cache(seq_len=max_position_embeddings)
+
+            ...     def _set_cos_sin_cache(self, seq_len):
+            ...         self.max_seq_len_cached = seq_len
+            ...         t = paddle.arange(seq_len, dtype="float32")
+            ...         freqs = paddle.einsum("i,j->ij", t, self.inv_freq)
+            ...         emb = paddle.concat([freqs, freqs], axis=-1)
+            ...         self.cos_cached = emb.cos()[None, :, None, :]
+            ...         self.sin_cached = emb.sin()[None, :, None, :]
+
+            ...     def forward(self, x, seq_len=None):
+            ...         cos = self.cos_cached[:, :seq_len, :, :]
+            ...         sin = self.sin_cached[:, :seq_len, :, :]
+            ...         return (
+            ...             cos.cast(x.dtype) if cos.dtype != x.dtype else cos,
+            ...             sin.cast(x.dtype) if sin.dtype != x.dtype else sin,
+            ...         )
+
+            >>> def rotate_half(x):
+            ...     x1 = x[..., : x.shape[-1] // 2]
+            ...     x2 = x[..., x.shape[-1] // 2 :]
+            ...     return paddle.concat([-x2, x1], axis=-1)
+
+            >>> def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+            ...     if position_ids is None:
+            ...         cos = cos[:, : q.shape[1], :, :]
+            ...         sin = sin[:, : q.shape[1], :, :]
+            ...     else:
+            ...         cos = cos.squeeze(axis=[0, 2])
+            ...         sin = sin.squeeze(axis=[0, 2])
+            ...         cos = cos[position_ids].unsqueeze(2)
+            ...         sin = sin[position_ids].unsqueeze(2)
+            ...     q_embed = (q * cos) + (rotate_half(q) * sin)
+            ...     k_embed = (k * cos) + (rotate_half(k) * sin)
+            ...     return q_embed, k_embed
+
+            >>> def scaled_dot_product_attention(
+            ...     query_states,
+            ...     key_states,
+            ...     value_states,
+            ...     attention_mask,
+            ... ):
+            ...     bsz, q_len, num_heads, head_dim = query_states.shape
+            ...     _, kv_seq_len, _, _ = value_states.shape
+            ...     query_states = paddle.transpose(query_states, [0, 2, 1, 3])
+            ...     key_states = paddle.transpose(key_states, [0, 2, 1, 3])
+            ...     value_states = paddle.transpose(value_states, [0, 2, 1, 3])
+            ...     attn_weights = paddle.matmul(query_states / math.sqrt(head_dim), key_states.transpose([0, 1, 3, 2]))
+            ...     attention_mask = attention_mask.reshape([bsz, 1, q_len, kv_seq_len])
+            ...     attn_weights = attn_weights + attention_mask
+            ...     if not paddle.in_dynamic_mode():
+            ...         attn_weights = F.softmax(attn_weights, axis=-1, dtype="float32").astype(query_states.dtype)
+            ...     else:
+            ...         with paddle.amp.auto_cast(False):
+            ...             attn_weights = F.softmax(attn_weights, axis=-1, dtype="float32").astype(query_states.dtype)
+            ...     attn_output = paddle.matmul(attn_weights, value_states)
+            ...     attn_output = attn_output.transpose([0, 2, 1, 3])
+            ...     attn_output = attn_output.reshape([bsz, q_len, head_dim * num_heads])
+            ...     return attn_output
+
+            >>> class Attention(nn.Layer):
+            ...     def __init__(self, hidden_size=HIDDEN_SIZE, n_head=N_HEAD):
+            ...         super().__init__()
+            ...         self.hidden_size = hidden_size
+            ...         self.num_heads = n_head
+            ...         self.head_dim = hidden_size // n_head
+            ...         self.q_proj = nn.Linear(hidden_size, hidden_size, bias_attr=False)
+            ...         self.k_proj = nn.Linear(hidden_size, hidden_size, bias_attr=False)
+            ...         self.v_proj = nn.Linear(hidden_size, hidden_size, bias_attr=False)
+            ...         self.o_proj = nn.Linear(hidden_size, hidden_size, bias_attr=False)
+            ...         self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=SEQ_LENGTH, base=10000)
+
+            ...     def forward(
+            ...         self,
+            ...         hidden_states,
+            ...         position_ids=None,
+            ...         attention_mask=None,
+            ...     ):
+            ...         query_states = self.q_proj(hidden_states)
+            ...         key_states = self.k_proj(hidden_states)
+            ...         value_states = self.v_proj(hidden_states)
+            ...         target_query_shape = [0, 0, self.num_heads, self.head_dim]
+            ...         target_key_value_shape = [0, 0, self.num_heads, self.head_dim]
+            ...         query_states = query_states.reshape(shape=target_query_shape)
+            ...         key_states = key_states.reshape(shape=target_key_value_shape)
+            ...         value_states = value_states.reshape(shape=target_key_value_shape)
+            ...         kv_seq_len = key_states.shape[-3]
+            ...         cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+            ...         query_states, key_states = apply_rotary_pos_emb(
+            ...             query_states, key_states, cos, sin, position_ids
+            ...         )
+            ...         output = scaled_dot_product_attention(
+            ...             query_states,
+            ...             key_states,
+            ...             value_states,
+            ...             attention_mask,
+            ...         )
+            ...         attn_output = output
+            ...         attn_output = self.o_proj(attn_output)
+            ...         return attn_output
+
+            >>> class Mlp(nn.Layer):
+            ...     def __init__(
+            ...         self,
+            ...         hidden_size=HIDDEN_SIZE,
+            ...         intermediate_size=INTERMEDIATE_SIZE,
+            ...     ):
+            ...         super().__init__()
+            ...         self.hidden_size = hidden_size
+            ...         self.intermediate_size = intermediate_size
+            ...         self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias_attr=False)
+            ...         self.up_proj = nn.Linear(hidden_size, intermediate_size, bias_attr=False)
+            ...         self.down_proj = nn.Linear(intermediate_size, hidden_size, bias_attr=False)
+
+            ...     def forward(self, x):
+            ...         x = paddle.nn.functional.swiglu(
+            ...             self.gate_proj(x), self.up_proj(x)
+            ...         )
+            ...         out = self.down_proj(x)
+            ...         return out
+
+            >>> class RMSNorm(nn.Layer):
+            ...     def __init__(self, hidden_size=HIDDEN_SIZE):
+            ...         super().__init__()
+            ...         self.hidden_size = hidden_size
+            ...         self.weight = paddle.create_parameter(
+            ...             shape=[self.hidden_size],
+            ...             dtype=paddle.get_default_dtype(),
+            ...             default_initializer=nn.initializer.Constant(1.0),
+            ...         )
+            ...         self.variance_epsilon = 1.0
+
+            ...     def forward(self, hidden_states):
+            ...         with paddle.amp.auto_cast(False):
+            ...             hidden_states = hidden_states.astype("float32")
+            ...             variance = hidden_states.pow(2).mean(-1, keepdim=True)
+            ...             hidden_states = (
+            ...                 paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
+            ...             )
+            ...         if self.weight.dtype in [paddle.float16, paddle.bfloat16]:
+            ...             hidden_states = paddle.cast(hidden_states, self.weight.dtype)
+            ...         return hidden_states * self.weight
+
+            >>> class DecoderLayer(nn.Layer):
+            ...     def __init__(
+            ...         self,
+            ...         hidden_size=HIDDEN_SIZE,
+            ...         intermediate_size=INTERMEDIATE_SIZE,
+            ...     ):
+            ...         super().__init__()
+            ...         self.hidden_size = hidden_size
+            ...         self.intermediate_size = intermediate_size
+            ...         self.self_attn = Attention(hidden_size)
+            ...         self.mlp = Mlp()
+            ...         self.input_layernorm = RMSNorm(hidden_size)
+            ...         self.post_attn_layernorm = RMSNorm(hidden_size)
+
+            ...     def forward(
+            ...         self,
+            ...         hidden_states,
+            ...         position_ids=None,
+            ...         attention_mask=None,
+            ...     ):
+            ...         residual = hidden_states
+            ...         hidden_states = self.input_layernorm(hidden_states)
+            ...         hidden_states = self.self_attn(
+            ...             hidden_states, position_ids, attention_mask
+            ...         )
+            ...         hidden_states = residual + hidden_states
+            ...         residual = hidden_states
+            ...         hidden_states = self.post_attn_layernorm(hidden_states)
+            ...         hidden_states = self.mlp(hidden_states)
+            ...         hidden_states = residual + hidden_states
+            ...         return hidden_states
+
+            >>> def _prepare_decoder_attention_mask(attention_mask, input_shape, dtype):
+            ...     batch_size, src_length = attention_mask.shape[0], attention_mask.shape[-1]
+            ...     batch_size, target_length = input_shape
+            ...     attention_mask = attention_mask[:, None, None, :].astype("bool")
+            ...     attention_mask.stop_gradient = True
+            ...     expanded_attn_mask = attention_mask.expand([batch_size, 1, target_length, src_length])
+            ...     mask = paddle.tril(paddle.ones((target_length, target_length), dtype="bool"))
+            ...     combined_attention_mask = mask[None, None, :, :].expand([batch_size, 1, target_length, target_length])
+            ...     expanded_attn_mask = expanded_attn_mask & combined_attention_mask
+            ...     expanded_attn_mask = paddle.where(expanded_attn_mask, 0.0, paddle.finfo(dtype).min).astype(dtype)
+            ...     return expanded_attn_mask
+
+            >>> class Model(nn.Layer):
+            ...     def __init__(
+            ...         self,
+            ...         vocab_size=VOCAB_SIZE,
+            ...         hidden_size=HIDDEN_SIZE,
+            ...         intermediate_size=INTERMEDIATE_SIZE,
+            ...     ):
+            ...         super().__init__()
+            ...         self.vocab_size = vocab_size
+            ...         self.hidden_size = hidden_size
+            ...         self.intermediate_size = intermediate_size
+            ...         self.embed_tokens = nn.Embedding(
+            ...             vocab_size,
+            ...             hidden_size,
+            ...         )
+            ...         self.layers = nn.LayerList([DecoderLayer() for i in range(NUM_HIDDEN_LAYERS)])
+            ...         self.norm = RMSNorm(hidden_size)
+            ...         self.weight = self.create_parameter(
+            ...             shape=[hidden_size, vocab_size],
+            ...             dtype=paddle.get_default_dtype(),
+            ...         )
+            ...         self.ignore_index = -100
+            ...         self.loss_func = paddle.nn.CrossEntropyLoss(reduction="none", ignore_index=self.ignore_index)
+
+            ...     def forward(
+            ...         self,
+            ...         input_ids=None,
+            ...         position_ids=None,
+            ...         attention_mask=None,
+            ...         labels=None,
+            ...     ):
+            ...         batch_size, seq_length = input_ids.shape
+            ...         inputs_embeds = self.embed_tokens(input_ids)
+            ...         attention_mask = paddle.ones(
+            ...             (batch_size, seq_length), dtype=paddle.bool
+            ...         )
+            ...         if position_ids is None:
+            ...             position_ids = paddle.arange(seq_length, dtype="int64").expand(
+            ...                 (batch_size, seq_length)
+            ...             )
+            ...         attention_mask = _prepare_decoder_attention_mask(
+            ...             attention_mask,
+            ...             (batch_size, seq_length),
+            ...             inputs_embeds.dtype,
+            ...         )
+            ...         hidden_states = inputs_embeds
+            ...         for idx, (decoder_layer) in enumerate(self.layers):
+            ...             layer_outputs = decoder_layer(
+            ...                 hidden_states,
+            ...                 position_ids,
+            ...                 attention_mask,
+            ...             )
+            ...             hidden_states = layer_outputs
+            ...         hidden_states = self.norm(hidden_states)
+            ...         logits = paddle.matmul(hidden_states, self.weight)
+            ...         loss = None
+            ...         if labels is not None:
+            ...             masked_lm_loss = self.loss_func(
+            ...                 logits.astype("float32"),
+            ...                 labels.unsqueeze(2),
+            ...             )
+            ...             binary_sequence = paddle.where(
+            ...                 masked_lm_loss > 0,
+            ...                 paddle.ones_like(masked_lm_loss),
+            ...                 paddle.zeros_like(masked_lm_loss),
+            ...             )
+            ...             count = paddle.sum(binary_sequence)
+            ...             if count == 0:
+            ...                 loss = paddle.sum(masked_lm_loss * binary_sequence)
+            ...             else:
+            ...                 loss = paddle.sum(masked_lm_loss * binary_sequence) / count
+            ...         return (loss, logits)
+
+            >>> model = Model()  # There is no distributed code or markup in Model
+            >>> input_seqs = np.random.randint(low=0, high=1024, size=(BATCH_SIZE * BATCH_NUM, SEQ_LENGTH)).astype("int64")
+            >>> labels = np.random.randint(low=0, high=1024, size=(BATCH_SIZE * BATCH_NUM, SEQ_LENGTH)).astype("int64")
+            >>> dataset = RandomDataset(input_seqs, labels, BATCH_SIZE * BATCH_NUM)
+            >>> sampler = paddle.io.BatchSampler(
+            ...     dataset,
+            ...     batch_size=BATCH_SIZE,
+            ...     shuffle=False,
+            ...     drop_last=True,
+            ... )
+            >>> loader = paddle.io.DataLoader(dataset, batch_sampler=sampler)
+            >>> opt = paddle.optimizer.SGD(learning_rate=0.1, parameters=model.parameters())
+            >>> input_seq_spec = paddle.static.InputSpec([BATCH_SIZE, SEQ_LENGTH], 'float32', 'input_seq', True)
+            >>> dist_config = ToDistributedConfig()
+            >>> dist_config.sequence_parallel = True
+
+            >>> # wrap model, opt, dataloader by using **to_distributed**
+            >>> dist_model, dist_opt, dist_loader = to_distributed(
+            ...     model,
+            ...     opt,
+            ...     loader,
+            ...     device_num=8,
+            ...     node_num=1,
+            ...     config=dist_config,
+            ... )
+
+            >>> for epoch in range(EPOCHS):
+            ...     dist_model.train()
+            ...     for i, data in enumerate(dist_loader()):
+            ...         inputs, labels = data
+            ...         loss, _ = dist_model(inputs, labels=labels)
+            ...         print(f"epoch {epoch}, step {i}: loss {loss}")
+            ...         loss.backward()
+            ...         dist_opt.step()
+            ...         dist_opt.clear_grad()
+            >>> # This case need to be executed in multi-card environment
+            >>> # python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 {test_case}.py
+    """
+    # Because some API(`paddle.randn` etc.) will be used when building pattern,
+    # In order to avoid circle import, we import get_pattern until function running.
+    from .static.tuner.to_distributed_api_patterns import (
+        clear_used_patterns,
+        get_pattern,
+        match_all_patterns,
+        register_used_patterns,
+    )
+
+    logger.debug(f'input model: {model}')
+    # paddle.distributed.init_parallel_env()
+
+    # step 1: identifying network structure and pattern recogincation
+    # step 1.1: register pre-hooks and post-hooks, thus recording corresponding static ops in following paddle.jit.to_static
+    for layer in model.sublayers():
+        pre_hook_helper = layer.register_forward_pre_hook(
+            record_program_ops_pre_hook
+        )
+        post_hook_helper = layer.register_forward_post_hook(
+            record_program_ops_post_hook
+        )
+        layer._op_recorder.hooks.append(pre_hook_helper)
+        layer._op_recorder.hooks.append(post_hook_helper)
+
+    # step 1.2: call @to_static, get program, and corresponding static ops of each layer
+    custom_input_spec = (
+        config.input_spec
+        if config.input_spec
+        else [paddle.static.InputSpec([4, 1024], 'float32', 'input_seq', True)]
+    )
+    static_func = paddle.jit.to_static(
+        model.forward, input_spec=custom_input_spec, full_graph=True
+    )
+    program = static_func.concrete_program.main_program
+    # currently, paddle.jit.to_static has side effects that will affect model.
+    # After fixing it, one line of code below can be dropped
+    static_func.rollback()
+    logger.debug(
+        f'Converted model to pir program: {program}, for pattern matching'
+    )
+
+    # step 1.3: get the mapping [dynamic-layers : static ops]
+    op_to_id = {}
+    for idx, op in enumerate(program.global_block().ops):
+        op_to_id[op] = idx
+
+    ops_id_to_layer = {}
+    op_id_to_layer = {}
+    for layer in model.sublayers():
+        layer_ops = layer._op_recorder.ops
+        logger.debug(
+            f'layer name: {layer.__class__.__name__}, layer_ops: {layer_ops}'
+        )
+        ops_id = []
+        for op in layer_ops:
+            assert op in op_to_id.keys(), f"{op.name} is not in program"
+            op_id = op_to_id[op]
+            op_id_to_layer[op_id] = layer
+            ops_id.append(op_id)
+        ops_id_to_layer[tuple(ops_id)] = layer
+    logger.debug(f'ops_id_to_layer is: {ops_id_to_layer}')
+
+    # step 1.4: pattern recogincation
+    DECODER_LAYER_NAME = 'decoder_layer'
+    register_used_patterns(DECODER_LAYER_NAME)
+    results = match_all_patterns(program)
+    logger.debug(f'Matched decoder layer patterns are: {results}')
+
+    matched_programs = {}
+    for pattern_name, matched_patterns in results.items():
+        # process one pattern
+        pattern_ops_dist_infos = get_pattern(pattern_name).ops_dist_infos
+        assert pattern_ops_dist_infos is not None, (
+            f"{pattern_name} does not contain ops_dist_infos, cannot reshard, please check"
+        )
+        processed_patterns = []
+        for matched_pattern in matched_patterns:
+            # convert pattern_ops_dist_infos to program_ops_dist_infos
+            program_ops_dist_infos = {}
+            for pattern_ops_id, op_dist_info in pattern_ops_dist_infos.items():
+                program_ops_id = []
+                for pattern_op_id in pattern_ops_id:
+                    assert pattern_op_id in matched_pattern.keys(), (
+                        f"please check ops_dist_infos of {pattern_name}, {pattern_op_id} not in matched_pattern: {matched_pattern.keys()}"
+                    )
+                    program_op_id = matched_pattern[pattern_op_id]
+                    program_ops_id.append(program_op_id)
+                program_ops_dist_infos[tuple(program_ops_id)] = op_dist_info
+            processed_patterns.append(program_ops_dist_infos)
+        matched_programs[pattern_name] = processed_patterns
+    logger.debug(f'Matched decoder layer patterns are: {matched_programs}')
+
+    # step 2: calculate the optimal parallel strategies based on the network structure
+    mesh = cost_model(matched_programs, device_num, node_num)
+    logger.debug(f'mesh: {mesh}')
+
+    with_pp = True if "pp" in mesh.dim_names else False
+    with_mp = True if "mp" in mesh.dim_names else False
+    with_dp = True if "dp" in mesh.dim_names else False
+    with_sp = (
+        True if "mp" in mesh.dim_names and config.sequence_parallel else False
+    )
+
+    # step 3: processing tensor parallel if necessary, according to the optimal parallel strategies shard weight tensors in decoder blocks
+    if with_mp:
+        num_hidden_layers = len(matched_programs[DECODER_LAYER_NAME])
+        for pattern_name, processed_patterns in matched_programs.items():
+            assert len(processed_patterns) == num_hidden_layers, (
+                "transformer patterns matched are incomplete"
+            )
+            for idx, processed_pattern in enumerate(processed_patterns):
+                local_mesh = mesh
+                if with_pp:
+                    pp_stage_id = get_layer_pp_info(
+                        mesh, num_hidden_layers, idx
+                    )
+                    local_mesh = mesh.get_mesh_with_dim("pp", pp_stage_id)
+
+                for program_ops_id, dist_infos in processed_pattern.items():
+                    assert program_ops_id in ops_id_to_layer.keys(), (
+                        f"program_ops: {program_ops_id} is not corresponding to a dynamic layer"
+                    )
+                    dynamic_layer = ops_id_to_layer[program_ops_id]
+                    mesh_num_dims = len(local_mesh.shape)
+                    sharding_info = dist_infos.get_dist_info(mesh_num_dims)
+                    dynamic_layer.weight = dist.shard_tensor(
+                        dynamic_layer.weight, local_mesh, sharding_info[0]
+                    )
+                    if dynamic_layer.bias is not None:
+                        dynamic_layer.bias = dist.shard_tensor(
+                            dynamic_layer.bias, local_mesh, sharding_info[1]
+                        )
+    logger.debug(f'after tensor parallel, model: {model}')
+
+    # step 4: processing pipeline parallel if necessary, reshard inputs of decoder blocks to next pp mesh b when switching from pp stage a to pp stage b
+    if with_pp:
+        decoder_layers = []
+        for pattern_name, matched_all_patterns in results.items():
+            if pattern_name == DECODER_LAYER_NAME:
+                for matched_pattern in matched_all_patterns:
+                    program_ops_id = []
+                    for a, b in matched_pattern.items():
+                        program_ops_id.append(b)
+                    if tuple(sorted(program_ops_id)) in ops_id_to_layer.keys():
+                        decoder_layers.append(
+                            ops_id_to_layer[tuple(sorted(program_ops_id))]
+                        )
+
+        if decoder_layers is not None:
+            num_decoder_blocks = len(decoder_layers)
+            assert num_decoder_blocks == num_hidden_layers, (
+                f"decoder pattern layers matched are incomplete, num_decoder_blocks: {num_decoder_blocks} should be equal to num_hidden_layers: {num_hidden_layers}"
+            )
+
+            pp_degree = mesh.get_dim_size("pp")
+            num_blocks_per_stage = num_decoder_blocks // pp_degree
+            for i in range(num_decoder_blocks):
+                pp_stage_id = get_layer_pp_info(mesh, num_decoder_blocks, i)
+                current_mesh = mesh.get_mesh_with_dim("pp", pp_stage_id)
+                decoder_layer = decoder_layers[i]
+                decoder_layer.__setattr__("current_mesh", current_mesh)
+                pre_hook_helper = decoder_layer.register_forward_pre_hook(
+                    reshard_all_inputs
+                )
+    logger.debug(f'after pipeline parallel, model: {model}')
+
+    # step 5: processing sequence parallel if necessary, reshard or transpose sequence dims for inputs of attention/mlp inputs
+    if with_sp:
+        clear_used_patterns()
+        EMBEDDING_LAYER_NAME = "embedding"
+        ATTENTION_LAYER_NAME = "attention"
+        MLP_LAYER_NAME = "mlp_3_with_swiglu"
+        RMS_NORM_LAYER_NAME = "rmsnorm"
+        used_patterns = [
+            EMBEDDING_LAYER_NAME,
+            ATTENTION_LAYER_NAME,
+            MLP_LAYER_NAME,
+            RMS_NORM_LAYER_NAME,
+        ]
+        register_used_patterns(used_patterns)
+        results = match_all_patterns(program)
+
+        matched_layers = {}
+        for pattern_name, matched_all_patterns in results.items():
+            if pattern_name in used_patterns:
+                for matched_pattern in matched_all_patterns:
+                    program_ops_id = []
+                    for a, b in matched_pattern.items():
+                        program_ops_id.append(b)
+                    if tuple(sorted(program_ops_id)) in ops_id_to_layer.keys():
+                        if pattern_name in matched_layers.keys():
+                            matched_layers[pattern_name].append(
+                                ops_id_to_layer[tuple(sorted(program_ops_id))]
+                            )
+                        else:
+                            matched_layers[pattern_name] = [
+                                ops_id_to_layer[tuple(sorted(program_ops_id))]
+                            ]
+
+        logger.debug(f'Matched attention/mlp layers are: {matched_layers}')
+        # init mesh
+        GLOBAL_MESH = []
+        if with_pp:
+            pp_degree = mesh.get_dim_size("pp")
+            for i in range(pp_degree):
+                local_mesh = mesh.get_mesh_with_dim("pp", i)
+                GLOBAL_MESH.append(local_mesh)
+        else:
+            GLOBAL_MESH.append(mesh)
+
+        # embedding: from [b/dp_degree, s, h] reshard+transpose to [s/mp_degree, b/dp_degree, h]
+        embedding_layer = matched_layers[EMBEDDING_LAYER_NAME][0]
+        embedding_layer_mesh = GLOBAL_MESH[0]
+        embedding_layer.__setattr__("current_mesh", embedding_layer_mesh)
+        post_hook_helper = embedding_layer.register_forward_post_hook(
+            transpose_reshard_embedding_layer_output
+        )
+
+        # attention: input from [s/mp_degree, b/dp_degree, h] to [b/dp_degree, s, h], output from [b/dp_degree, s, h] to [s/mp_degree, b/dp_degree, h]
+        attention_layers = matched_layers[ATTENTION_LAYER_NAME]
+        num_attention_layers = len(attention_layers)
+        if attention_layers is not None:
+            for i in range(num_attention_layers):
+                current_mesh = GLOBAL_MESH[0]
+                if with_pp:
+                    pp_stage_id = get_layer_pp_info(
+                        mesh, num_attention_layers, i
+                    )
+                    current_mesh = GLOBAL_MESH[pp_stage_id]
+                attention_layer = attention_layers[i]
+                attention_layer.__setattr__("current_mesh", current_mesh)
+                pre_hook_helper = attention_layer.register_forward_pre_hook(
+                    reshard_transpose_attention_layer_input
+                )
+                post_hook_helper = attention_layer.register_forward_post_hook(
+                    transpose_reshard_attention_layer_output
+                )
+
+        # mlp: input from [s/mp_degree, b/dp_degree, h] to [s, b/dp_degree, h], output from [s, b/dp_degree, h] to [s/mp_degree, b/dp_degree, h]
+        mlp_layers = matched_layers[MLP_LAYER_NAME]
+        num_mlp_layers = len(mlp_layers)
+        if mlp_layers is not None:
+            for i in range(num_mlp_layers):
+                current_mesh = GLOBAL_MESH[0]
+                if with_pp:
+                    pp_stage_id = get_layer_pp_info(
+                        mesh, num_attention_layers, i
+                    )
+                    current_mesh = GLOBAL_MESH[pp_stage_id]
+                mlp_layer = mlp_layers[i]
+                mlp_layer.__setattr__("current_mesh", current_mesh)
+                pre_hook_helper = mlp_layer.register_forward_pre_hook(
+                    reshard_mlp_layer_input
+                )
+                post_hook_helper = mlp_layer.register_forward_post_hook(
+                    reshard_mlp_layer_output
+                )
+
+        # rms norm: for the last rms norm (after decoder blocks), input from [s/mp_degree, b/dp_degree, h] to [b, s, h]
+        rms_norm_layers = matched_layers[RMS_NORM_LAYER_NAME]
+        if rms_norm_layers is not None:
+            last_rms_norm_layer = rms_norm_layers[-1]
+            current_mesh = GLOBAL_MESH[-1]
+            last_rms_norm_layer.__setattr__("current_mesh", current_mesh)
+            post_hook_helper = last_rms_norm_layer.register_forward_post_hook(
+                reshard_transpose_rms_norm_layer_output
+            )
+
+    # step 6: processing data parallel if necessary, shard dataloader
+    # TODO(jeff41404): shard optimizer
+    if with_dp:
+        if with_pp:
+            first_stage_mesh = mesh.get_mesh_with_dim("pp", 0)
+            last_stage_mesh = mesh.get_mesh_with_dim("pp", 1)
+            loader = dist.shard_dataloader(
+                dataloader,
+                meshes=[first_stage_mesh, last_stage_mesh],
+                shard_dims="dp",
+            )
+        else:
+            loader = dist.shard_dataloader(
+                dataloader, meshes=[mesh], shard_dims="dp"
+            )
+    else:
+        loader = dist.shard_dataloader(
+            dataloader, meshes=[mesh], shard_dims=None
+        )
+
+    # step 7: clean layer_op recorder hooks
+    for layer in model.sublayers():
+        for hook_helper in layer._op_recorder.hooks:
+            hook_helper.remove()
+
+    return model, optimizer, loader

python/paddle/distributed/auto_parallel/interface.py

@@ -0,0 +1,402 @@
+#   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 __future__ import annotations
+
+from functools import reduce
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+
+from .process_mesh import ProcessMesh, get_current_process_mesh
+from .static.dist_context import get_default_distributed_context
+from .static.dist_op import DistributedOperatorHelper
+from .static.dist_tensor import DistributedTensor
+from .static.utils import (
+    __no_shape_var_type__,
+    convert_to_dims_mapping,
+    verify_shard_spec,
+)
+
+
+def shard_tensor(x, process_mesh=None, shard_spec=None):
+    """
+    Shard a tensor on a process mesh according to the shard specification.
+
+    Args:
+        x (Tensor): the tensor to be sharded.
+        process_mesh (ProcessMesh, optional): An instance of ProcessMesh describes a mesh
+            topology of the used logical processes where the tensor is sharded. If it is None,
+            the found current process mesh will be used. And an error will be raised if the
+            current process mesh cannot be found. Default: None.
+        shard_spec (list, optional): a list to describe the sharding mapping between `x` and `process_mesh`,
+            which means the dimension `i` of `x` is split across the dimension `shard_spec[i]` of `process_mesh`,
+            where `None` means that tensor dimension is not split. For example, given a tensor with
+            the shape [6, 12] and a process mesh with the shape [2, 3] and the dimension names ["x", "y"]:
+                If `shard_spec=["x", "y"]`, each shard of the tensor will have a shape [3, 4];
+                If `shard_spec=["y", "x"]`, each shard of the tensor will have a shape [2, 6];
+                If `shard_spec=["x", None]`, each shard of the tensor will have a shape [3, 12];
+                If `shard_spec=[None, "x"]`, each shard of the tensor will have a shape [6, 4];
+                If `shard_spec=["y", None]`, each shard of the tensor will have a shape [2, 12];
+                If `shard_spec=[None, "y"]`, each shard of the tensor will have a shape [6, 4];
+                If `shard_spec=[None, None]`, each shard of the tensor will have a shape [6, 12];
+        If the `shard_spec` is None, the tensor will be replicated across all the processes of `process_mesh`.
+        In the above example, the `shard_spec=None` is same as 'shard_spec=[None, None]'. Defaults: None.
+
+    Returns:
+        Tensor: the tensor `x` annotated with sharding information.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle
+            >>> from paddle.distributed.fleet import auto
+
+            >>> mesh = auto.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+            >>> x = paddle.ones([4, 6])
+            >>> shard_spec = ["x", "y"]
+            >>> auto.shard_tensor(x, mesh, shard_spec)
+
+    """
+
+    if process_mesh is not None:
+        assert isinstance(process_mesh, core.ProcessMesh), (
+            f"Argument process_mesh {process_mesh} is not an instance of ProcessMesh"
+        )
+    else:
+        process_mesh = get_current_process_mesh()
+        assert process_mesh is not None, (
+            "Specify the process mesh argument or use ProcessMesh context manager first."
+        )
+    assert isinstance(shard_spec, list), (
+        f"Argument shard_spec {shard_spec} is not an instance of list"
+    )
+    if isinstance(x, str):
+        x = (
+            paddle.static.default_main_program()
+            .global_block()
+            ._var_recursive(x)
+        )
+        dist_tensor = DistributedTensor(x)
+    else:
+        dist_tensor = DistributedTensor(x)
+    serial_tensor = dist_tensor.serial_tensor
+    dist_tensor.dist_attr.process_mesh = process_mesh
+    if serial_tensor.type in __no_shape_var_type__:
+        tensor_shape = []
+    else:
+        tensor_shape = serial_tensor.shape
+    if shard_spec is not None:
+        valid_dims = (
+            process_mesh.get_dim_names()
+            if hasattr(process_mesh, "get_dim_names")
+            else process_mesh.dim_names
+        )
+        for i, dim in enumerate(shard_spec):
+            if dim is not None and (
+                not isinstance(dim, str) or dim not in valid_dims
+            ):
+                raise ValueError(
+                    f"Invalid shard_spec at index {i}: '{dim}' "
+                    f"is not a valid dimension name in process_mesh {valid_dims}."
+                )
+        assert verify_shard_spec(shard_spec, tensor_shape, process_mesh), (
+            f"For tensor {serial_tensor.name}, shard_spec {shard_spec} is invalid with tensor_shape {tensor_shape} and process_mesh {process_mesh}."
+        )
+        dist_tensor.dist_attr.dims_mapping = convert_to_dims_mapping(
+            shard_spec, process_mesh
+        )
+    if process_mesh is not None:
+        dist_tensor.dist_attr.mark_annotated("process_mesh")
+    if shard_spec is not None:
+        dist_tensor.dist_attr.mark_annotated("dims_mapping")
+    default_dist_ctx = get_default_distributed_context()
+    default_dist_ctx.add_dist_tensor_for_program(dist_tensor)
+    dist_tensor = default_dist_ctx.get_dist_tensor_for_program(x)
+    default_dist_ctx.add_process_mesh(process_mesh)
+    return x
+
+
+def shard_op(
+    op, process_mesh=None, in_shard_specs=None, out_shard_specs=None, **kwargs
+):
+    """
+    Shard an operation on a process mesh according to its input and output shard specification.
+
+    Args:
+        op (Callable): a callable operator or module to be sharded.
+        process_mesh (ProcessMesh, optional): An instance of ProcessMesh describes a mesh
+            topology of the used logical processes where the op is sharded. All of its inputs and
+            outputs are sharded by this process mesh. If it is None, the found current process mesh
+            will be used. And an error will be raised if the current process mesh cannot be found.
+            Default: None.
+        in_shard_specs (list of list, optional): a list of list to describe the sharding specifications
+            for the inputs. Each item of `in_shard_specs` is a `shard_spec` between the corresponding input
+            and `process_mesh`. If one item is None, the corresponding input is replicated across all processes
+            If it is None, all inputs are replicated across all processes. Note that the length of the
+            `in_shard_specs` should be equal to the actual number of inputs when calling this operation.
+            Default: None.
+        out_shard_specs (list of list, optional): a list of list to describe the sharding specifications
+            for the outputs. Each item of `out_shard_specs` is a `shard_spec` between the corresponding output
+            and `process_mesh`. If one item is None, the corresponding output is replicated across all processes
+            If it is None, all outputs are replicated across all processes. Note that the length of the
+            `in_shard_specs` should be equal to the actual number of inputs when calling this operation.
+            Default: None. Default: None.
+
+    Returns:
+        Outputs of `op`, each of which is annotated with sharding information.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.distributed.fleet import auto
+
+            >>> x = paddle.ones([4, 6])
+            >>> y = paddle.zeros([4, 6])
+            >>> mesh = auto.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+            >>> dist_add = auto.shard_op(
+            ...     paddle.add,
+            ...     mesh,
+            ...     in_shard_specs=[["x", "y"], ["y", None]],
+            ...     out_shard_specs=[[None, "x"]],
+            ... )
+            >>> dist_add(x, y)
+
+    """
+
+    if process_mesh is not None:
+        assert isinstance(process_mesh, ProcessMesh), (
+            f"Argument process_mesh {process_mesh} is not an instance of ProcessMesh"
+        )
+    else:
+        process_mesh = get_current_process_mesh()
+        assert process_mesh is not None, (
+            "Specify the process mesh argument or use ProcessMesh context manager first."
+        )
+    in_dims_mappings = []
+    if in_shard_specs is not None:
+        assert all(
+            (isinstance(shard_spec, list) or shard_spec is None)
+            for shard_spec in in_shard_specs
+        ), f"in_shard_spec {in_shard_specs} is not a list of list or None"
+        for shard_spec in in_shard_specs:
+            if shard_spec is not None:
+                in_dims_mappings.append(
+                    convert_to_dims_mapping(shard_spec, process_mesh)
+                )
+            else:
+                in_dims_mappings.append(None)
+    out_dims_mappings = []
+    if out_shard_specs is not None:
+        assert all(
+            (isinstance(shard_spec, list) or shard_spec is None)
+            for shard_spec in out_shard_specs
+        ), f"out_shard_spec {out_shard_specs} is not a list of list or None"
+        for shard_spec in out_shard_specs:
+            if shard_spec is not None:
+                out_dims_mappings.append(
+                    convert_to_dims_mapping(shard_spec, process_mesh)
+                )
+            else:
+                out_dims_mappings.append(None)
+    op = DistributedOperatorHelper(
+        op, process_mesh, in_dims_mappings, out_dims_mappings, kwargs
+    )
+    return op
+
+
+_g_recompute_idx = -1
+
+
+def recompute(op):
+    global _g_recompute_idx
+    _g_recompute_idx += 1
+
+    class RecomputeOperator:
+        def __init__(self, op):
+            self._op = op
+
+        def __call__(self, *args, **kwargs):
+            block = paddle.static.default_main_program().global_block()
+            rc_begin_id = len(block.ops)
+
+            with paddle.static.name_scope(
+                f'/auto_parallel/rc_{_g_recompute_idx}'
+            ):
+                if paddle.base.dygraph.base.in_to_static_mode():
+                    output = (
+                        paddle.jit.dy2static.convert_call_func.convert_call(
+                            self._op
+                        )(*args, **kwargs)
+                    )
+                else:
+                    output = self._op(*args, **kwargs)
+
+            if paddle.framework.in_pir_mode():
+                block = paddle.static.default_main_program().global_block()
+                rc_end_id = len(block.ops)
+                for idx in range(rc_begin_id, rc_end_id):
+                    rc_op = block.ops[idx]
+                    rc_op.set_int_attr("fwd_recompute_id", _g_recompute_idx)
+
+            return output
+
+    return RecomputeOperator(op)
+
+
+def exclude_ops_in_recompute(run_function):
+    """
+    Exclude some operators in recompute segments.
+        Args:
+        run_function (callable): The callable function to be excluded.
+
+    Returns:
+        ExcludeOperator: The callable object.
+
+    """
+
+    class ExcludeOperator:
+        def __init__(self, run_function):
+            self._run_function = run_function
+
+        def __call__(self, *args, **kwargs):
+            with paddle.static.name_scope('/exclude_rc'):
+                if paddle.base.dygraph.base.in_to_static_mode():
+                    output = (
+                        paddle.jit.dy2static.convert_call_func.convert_call(
+                            self._run_function
+                        )(*args, **kwargs)
+                    )
+                else:
+                    output = self._run_function(*args, **kwargs)
+
+            return output
+
+    return ExcludeOperator(run_function)
+
+
+_g_collections = {}
+
+
+class CollectionNames:
+    FETCHES = "fetches"
+    LOGGING = "logging"
+
+
+def get_collection(name):
+    collection = _g_collections.get(name, None)
+    if collection is None:
+        collection = []
+        _g_collections[name] = collection
+    return _g_collections[name]
+
+
+def add_to_collection(collection_name, value, name=None):
+    if collection_name not in _g_collections:
+        _g_collections[collection_name] = []
+    if name is not None:
+        for _, v in _g_collections[collection_name]:
+            if v == value:
+                return
+        _g_collections[collection_name].append((name, value))
+    else:
+        for _, v in _g_collections[collection_name]:
+            if v == value:
+                return
+        _g_collections[collection_name].append((None, value))
+
+
+def fetch(tensor, name=None, logging=False):
+    if isinstance(tensor, paddle.static.Variable):
+        tensor = tensor.name
+    elif isinstance(tensor, str):
+        tensor = tensor
+    else:
+        raise TypeError(
+            f"Only support fetch `Variable` or `str`[`Variable`'s name], but got `{type(tensor)}`"
+        )
+    add_to_collection(CollectionNames.FETCHES, tensor, name)
+    if logging:
+        add_to_collection(CollectionNames.LOGGING, tensor, name)
+
+
+_g_mesh = None
+
+
+def get_mesh() -> paddle.distributed.ProcessMesh:
+    """
+    Get the global mesh set by set_mesh.
+
+    Returns:
+        mesh (paddle.distributed.ProcessMesh): the global mesh.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> mesh = dist.ProcessMesh([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=["dp", "mp", "pp"])
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> dist.auto_parallel.set_mesh(mesh)
+            >>> mesh = dist.auto_parallel.get_mesh()
+            >>> # This case need to be executed in multi-card environment
+            >>> # python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 {test_case}.py
+    """
+    global _g_mesh
+    return _g_mesh
+
+
+def set_mesh(mesh: paddle.distributed.ProcessMesh) -> None:
+    """
+    Set the global mesh.
+
+    Args:
+        mesh (paddle.distributed.ProcessMesh): global mesh to be set.
+
+    Returns:
+        None
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> mesh = dist.ProcessMesh([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=["dp", "mp", "pp"])
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> dist.auto_parallel.set_mesh(mesh)
+            >>> # This case need to be executed in multi-card environment
+            >>> # python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 {test_case}.py
+    """
+    global _g_mesh
+    _g_mesh = mesh
+
+
+def create_mesh(mesh_dims: list[tuple[str, int]]):
+    """
+    Create a global process_mesh for auto parallel.
+
+    Args:
+        mesh_dims (list[tuple[str, int]]): A list of tuple, each element is (dim_name, dim_degree).
+    """
+    global _g_mesh
+    dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+    mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+    mesh_arr = np.arange(0, reduce(lambda x, y: x * y, mesh_shape, 1)).reshape(
+        mesh_shape
+    )
+    _g_mesh = ProcessMesh(mesh_arr, dim_names)
+    return _g_mesh

python/paddle/distributed/auto_parallel/intermediate/__init__.py

@@ -0,0 +1,15 @@
+#   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.
+
+__all__ = []

python/paddle/distributed/auto_parallel/intermediate/context_parallel.py

@@ -0,0 +1,391 @@
+#   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 __future__ import annotations
+
+import logging
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.ring_attention import (
+    shard_seq_load_balance,
+)
+
+from .tensor_parallel import PlanBase
+
+
+class PrepareContextParallel(PlanBase):
+    """
+
+    Prepare Input for context parallel optimizations.
+
+    This will work for Layer that calls like whole-llama Layer which is the first layer in the network.
+
+    Users can set backend='p2p/all2all' for different context parallel strategys.
+
+    backend='p2p' will use Ring FlashAttention strategy which segments input with balance in the sequence dimension before whole-llama Layer.
+    backend='all2all' will use Deepspeed Ulysses strategy(Paddle SegmentParallel strategy) which segments input in the sequence dimension before whole-llama Layer.
+
+    Args:
+        backend (string): select strategy for context parallel, now support 'p2p' and 'all2all'.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class SDPALayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...
+            ...     def forward(self, q, k, v):
+            ...         return paddle.nn.functional.scaled_dot_product_attention(q, k, v)
+            >>>
+            >>> class AttentionLayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.hidden_size = 64
+            ...         self.num_key_value_heads = 10
+            ...         self.head_dim = 64
+            ...         self.sdpa = SDPALayer()
+            ...         self.q = paddle.nn.Linear(
+            ...             self.hidden_size,
+            ...             self.hidden_size,
+            ...             bias_attr=False,
+            ...         )
+            ...         self.k = paddle.nn.Linear(
+            ...             self.hidden_size,
+            ...             self.num_key_value_heads * self.head_dim,
+            ...             bias_attr=False,
+            ...         )
+            ...         self.v = paddle.nn.Linear(
+            ...             self.hidden_size,
+            ...             self.num_key_value_heads * self.head_dim,
+            ...             bias_attr=False,
+            ...         )
+            ...
+            ...     def forward(self, input):
+            ...         q = self.q(input)
+            ...         k = self.k(input)
+            ...         v = self.v(input)
+            ...         return self.sdpa(q, k, v)
+            >>>
+            >>> class LlamaLayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.attention = AttentionLayer()
+            ...
+            ...     def forward(self, input, label):
+            ...         return self.attention(input)
+            >>>
+            >>> class LlamaForCausalLayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.llama = LlamaLayer()
+            ...         self.weight = self.create_parameter(shape=[64, 1024])
+            ...         self.loss_func = paddle.nn.CrossEntropyLoss()
+            ...
+            ...     def forward(self, input, label):
+            ...         out = self.llama(input, label)
+            ...         logits = paddle.matmul(out, self.weight)
+            ...         loss = self.loss_func(logits, label)
+            ...         return logits
+            >>>
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = LlamaForCausalLayer()
+            >>> mp_config = {
+            ...     'llama': dist.PrepareContextParallel('p2p'),
+            ...     'sdpa': dist.ContextParallel('p2p'),
+            ... }
+    """
+
+    def __init__(self, backend: str = 'p2p') -> None:
+        super().__init__()
+        self.backend = backend
+        assert self.backend in [
+            'p2p',
+            'all2all',
+        ], f"backend must be 'p2p' or 'all2all', but got {self.backend}"
+
+    def all2all_split_input_pre_hook(self, process_mesh):
+        def shard_tensor(input_tensor, seq_dim):
+            cp_index = process_mesh.dim_names.index('sep')
+            placements = input_tensor.placements
+            if placements is None:
+                placements = [
+                    dist.Replicate() for _ in range(len(process_mesh.shape))
+                ]
+            # split sequence dim
+            placements[cp_index] = dist.Shard(seq_dim)
+            reshard_input = dist.reshard(input_tensor, process_mesh, placements)
+            return reshard_input
+
+        def all2all_split_input(layer, args):
+            cp_index = process_mesh.dim_names.index('sep')
+            cp_degree = process_mesh.shape[cp_index]
+            # check input_ids
+            if isinstance(args, (list, tuple)):
+                all_args = []
+                for input_tensor in args:
+                    assert input_tensor.is_dist(), (
+                        "Input tensor must be a distributed tensor."
+                    )
+                    assert len(input_tensor.shape) == 2, (
+                        f"input_ids should be [batch_size, seq_len], but got {input_tensor.shape}"
+                    )
+                    _, seq_len = input_tensor.shape
+                    assert seq_len % cp_degree == 0, (
+                        f"sequence length {seq_len} must be divisible by cp degree {cp_degree}"
+                    )
+                    reshard_input = shard_tensor(input_tensor, 1)
+                    all_args.append(reshard_input)
+                new_args = tuple(all_args)
+                return new_args
+            elif isinstance(args, paddle.Tensor):
+                reshard_input = shard_tensor(args, 1)
+                return reshard_input
+            else:
+                raise ValueError(
+                    f"Unsupported argument type: {type(args)}. Expected list of tensors or single tensor."
+                )
+
+        return all2all_split_input
+
+    def p2p_split_input_pre_hook(self, process_mesh):
+        def p2p_split_input(layer, args):
+            cp_index = process_mesh.dim_names.index('sep')
+            cp_degree = process_mesh.shape[cp_index]
+            if isinstance(args, (list, tuple)):
+                all_args = []
+                for input_tensor in args:
+                    # check input_ids
+                    assert input_tensor.is_dist(), (
+                        "Input tensor must be a distributed tensor."
+                    )
+                    assert len(input_tensor.shape) == 2, (
+                        f"input_ids should be [batch_size, seq_len], but got {input_tensor.shape}"
+                    )
+                    placements = input_tensor.placements
+                    if placements is None:
+                        placements = [
+                            dist.Replicate()
+                            for _ in range(len(process_mesh.shape))
+                        ]
+                    assert placements[cp_index] == dist.Replicate(), (
+                        "Input tensor must be a replicated tensor in cp mesh."
+                    )
+                    reshard_input = shard_seq_load_balance(input_tensor, 1)
+                    all_args.append(reshard_input)
+                new_args = tuple(all_args)
+                return new_args
+            elif isinstance(args, paddle.Tensor):
+                reshard_input = shard_seq_load_balance(input_tensor, 1)
+                return reshard_input
+            else:
+                raise ValueError(
+                    f"Unsupported argument type: {type(args)}. Expected list of tensors or single tensor."
+                )
+
+        return p2p_split_input
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        if self.backend == 'all2all':
+            # Deepspeed Ulysses
+            layer.register_forward_pre_hook(
+                self.all2all_split_input_pre_hook(process_mesh)
+            )
+        elif self.backend == 'p2p':
+            # Ring FlashAttention
+            layer.register_forward_pre_hook(
+                self.p2p_split_input_pre_hook(process_mesh)
+            )
+        else:
+            logging.warning(
+                f'{self.backend} is not supported backend for context parallel'
+            )
+
+
+class ContextParallel(PlanBase):
+    """
+
+    Applies context parallel optimizations to the attention layer.
+
+    This will work for Layer that calls paddle.nn.functional.scaled_dot_product_attention).
+
+    Users can set backend='p2p/all2all' for different context parallel strategys.
+
+    backend='p2p' will use Ring FlashAttention strategy which segments q/k/v in the sequence dimension and communicates k/v between ranks.
+    backend='all2all' will use Deepspeed Ulysses strategy(Paddle SegmentParallel strategy) which inserts all2all before and after sdpa compute.
+
+    Note:
+
+
+    Args:
+        backend (string): select strategy for context parallel, now support 'p2p' and 'all2all'.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class SDPALayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...
+            ...     def forward(self, q, k, v):
+            ...         return paddle.nn.functional.scaled_dot_product_attention(q, k, v)
+            >>>
+            >>> class AttentionLayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.hidden_size = 64
+            ...         self.num_key_value_heads = 10
+            ...         self.head_dim = 64
+            ...         self.sdpa = SDPALayer()
+            ...         self.q = paddle.nn.Linear(
+            ...             self.hidden_size,
+            ...             self.hidden_size,
+            ...             bias_attr=False,
+            ...         )
+            ...         self.k = paddle.nn.Linear(
+            ...             self.hidden_size,
+            ...             self.num_key_value_heads * self.head_dim,
+            ...             bias_attr=False,
+            ...         )
+            ...         self.v = paddle.nn.Linear(
+            ...             self.hidden_size,
+            ...             self.num_key_value_heads * self.head_dim,
+            ...             bias_attr=False,
+            ...         )
+            ...
+            ...     def forward(self, input):
+            ...         q = self.q(input)
+            ...         k = self.k(input)
+            ...         v = self.v(input)
+            ...         return self.sdpa(q, k, v)
+            >>>
+            >>> class LlamaLayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.attention = AttentionLayer()
+            ...
+            ...     def forward(self, input, label):
+            ...         return self.attention(input)
+            >>>
+            >>> class LlamaForCausalLayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.llama = LlamaLayer()
+            ...         self.weight = self.create_parameter(shape=[64, 1024])
+            ...         self.loss_func = paddle.nn.CrossEntropyLoss()
+            ...
+            ...     def forward(self, input, label):
+            ...         out = self.llama(input, label)
+            ...         logits = paddle.matmul(out, self.weight)
+            ...         loss = self.loss_func(logits, label)
+            ...         return logits
+            >>>
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = LlamaForCausalLayer()
+            >>> mp_config = {
+            ...     'llama': dist.PrepareContextParallel('p2p'),
+            ...     'sdpa': dist.ContextParallel('p2p'),
+            ... }
+    """
+
+    def __init__(self, backend: str = 'p2p') -> None:
+        super().__init__()
+        self.backend = backend
+
+    def all2all_reshard_pre_hook(self, process_mesh):
+        def all2all_reshard_hook(layer, args):
+            cp_index = process_mesh.dim_names.index('sep')
+            cp_degree = process_mesh.shape[cp_index]
+            all_args = []
+            for arg in args:
+                # check q k v
+                assert arg.is_dist(), f"arg {arg} must be a distributed tensor."
+                assert len(arg.shape) == 3 or len(arg.shape) == 4
+                placements = arg.placements
+                assert placements[cp_index] == dist.Shard(1), (
+                    f"arg {arg} must be sharded in sequence dimension."
+                )
+                # reshard [batch_size，seq_len/sep，num_head，head_dim] -> [batch_size，seq_len，num_head/sep，head_dim]
+                placements[cp_index] = dist.Shard(2)
+                target_arg = dist.reshard(arg, process_mesh, placements)
+                all_args.append(target_arg)
+            new_args = tuple(all_args)
+            return new_args
+
+        return all2all_reshard_hook
+
+    def all2all_reshard_post_hook(self, process_mesh):
+        def all2all_reshard_hook(layer, input, output):
+            cp_index = process_mesh.dim_names.index('sep')
+            cp_degree = process_mesh.shape[cp_index]
+            placements = output.placements
+            assert output.is_dist(), (
+                f"output {output} must be a distributed tensor."
+            )
+            assert len(output.shape) == 4 or len(output.shape) == 3
+            assert placements[cp_index] == dist.Shard(2), (
+                f"output {output} must be Shard(2) in sequence dimension."
+            )
+            # reshard [batch_size，seq_len，num_head/seq，head_dim]  ->  [batch_size，seq_len/sep，num_head，head_dim]
+            placements[cp_index] = dist.Shard(1)
+            target_output = dist.reshard(output, process_mesh, placements)
+            return target_output
+
+        return all2all_reshard_hook
+
+    def p2p_reshard_pre_hook(self, process_mesh):
+        def input_hook(layer, args, kwargs):
+            cp_index = process_mesh.dim_names.index('sep')
+            cp_degree = process_mesh.shape[cp_index]
+            for arg in args:
+                # check q k v
+                assert arg.is_dist(), (
+                    "Input tensor must be a distributed tensor."
+                )
+                assert len(arg.shape) == 3 or len(arg.shape) == 4
+                placements = arg.placements
+                assert placements[cp_index] == dist.Shard(1), (
+                    f"arg {arg} must be Shard(1) in sequence dimension."
+                )
+            # edit kwarg backend to 'p2p'
+            new_kwargs = kwargs
+            new_kwargs['backend'] = 'p2p'
+            return args, new_kwargs
+
+        return input_hook
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        if self.backend == 'all2all':
+            # Deepspeed Ulysses
+            layer.register_forward_pre_hook(
+                self.all2all_reshard_pre_hook(process_mesh)
+            )
+            layer.register_forward_post_hook(
+                self.all2all_reshard_post_hook(process_mesh)
+            )
+        elif self.backend == 'p2p':
+            # Ring FlashAttention
+            layer.register_forward_pre_hook(
+                self.p2p_reshard_pre_hook(process_mesh), with_kwargs=True
+            )
+        else:
+            logging.warning(
+                f'{self.backend} is not supported backend for context parallel'
+            )

python/paddle/distributed/auto_parallel/intermediate/parallel_base.py

@@ -0,0 +1,294 @@
+#   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
+
+import paddle
+import paddle.distributed as dist
+from paddle import pir
+from paddle.base.framework import (
+    in_dygraph_mode,
+    in_pir_mode,
+)
+from paddle.distributed import fleet
+from paddle.nn import Layer
+from paddle.optimizer import Optimizer
+
+
+def is_tensor(tensor):
+    if in_dygraph_mode():
+        return isinstance(tensor, paddle.Tensor)
+    elif in_pir_mode():
+        return isinstance(tensor, pir.Value)
+    else:
+        raise RuntimeError(
+            "PipelineParallel are only supported in dynamic or pir mode."
+        )
+
+
+class ParallelOptimizer:
+    def __init__(
+        self,
+        optimizer,
+        level=None,
+        sharding_mesh_dim=None,
+    ):
+        self.level = None
+        self.sharding_mesh_dim = None
+        self.optimizer = None
+
+        if isinstance(optimizer, ParallelOptimizer):
+            self.optimizer = optimizer.optimizer
+            if level is None:
+                self.level = optimizer.level
+                self.sharding_mesh_dim = optimizer.sharding_mesh_dim
+            else:
+                if isinstance(level, int):
+                    level = str(level)
+                assert level in ("0", "1", "2", "3", None)
+                if optimizer.level is not None:
+                    assert level == optimizer.level, (
+                        f"The level passed in is not identical with previous level. Current level is {level}, previous level is {optimizer.level}"
+                    )
+                self.level = level
+                self.sharding_mesh_dim = sharding_mesh_dim
+        else:
+            assert isinstance(optimizer, Optimizer)
+            self.optimizer = optimizer
+            if isinstance(level, int):
+                level = str(level)
+            assert level in ("0", "1", "2", "3", None)
+            # level=0 and level=None are all mean pure dp
+            self.level = level
+            self.sharding_mesh_dim = sharding_mesh_dim
+
+        self.is_initialized = False
+
+    def parallelize(self):
+        assert self.optimizer is not None
+        if self.is_initialized:
+            return self.optimizer
+
+        mesh = fleet.auto.get_mesh()
+        if self.level == "1":
+            self.optimizer = dist.shard_optimizer(
+                self.optimizer,
+                dist.ShardingStage1(self.sharding_mesh_dim, mesh),
+            )
+        elif self.level == "2":
+            self.optimizer = dist.shard_optimizer(
+                self.optimizer,
+                dist.ShardingStage2(self.sharding_mesh_dim, mesh),
+            )
+        elif self.level == "3":
+            self.optimizer = dist.shard_optimizer(
+                self.optimizer,
+                dist.ShardingStage3(self.sharding_mesh_dim, mesh),
+            )
+        else:
+            self.optimizer = dist.shard_optimizer(self.optimizer, None)
+        self.is_initialized = True
+
+        return self.optimizer
+
+    def update_param_list(self, parallelized_parameters):
+        self.optimizer._parameter_list = parallelized_parameters
+        if isinstance(parallelized_parameters[0], dict):
+            self.optimizer._param_groups = []
+            for param_group in self.parallelized_parameters:
+                self.optimizer._add_param_group(param_group.copy())
+        else:
+            self.optimizer._param_groups = self.optimizer._parameter_list
+
+
+class ParallelModel:
+    def __init__(self, model):
+        super().__init__()
+        self.pp_parallelizer = None
+        self.tp_parallelizer = None
+        self.sharding_parallelizer = None
+        self.model = None
+        self.share_param_list = {}
+        self.layer_param_placements = {}
+        if isinstance(model, ParallelModel):
+            self.pp_parallelizer = model.pp_parallelizer
+            self.tp_parallelizer = model.tp_parallelizer
+            self.sharding_parallelizer = model.sharding_parallelizer
+            self.model = model.model
+        else:
+            assert isinstance(model, Layer)
+            self.model = model
+
+        self.is_parallelized = False
+
+    def get_mesh(self, pp_idx=0):
+        mesh = fleet.auto.get_mesh()
+        if "pp" in mesh.dim_names:
+            mesh = mesh.get_mesh_with_dim("pp", pp_idx)
+        return mesh
+
+    def parallelize_model(self):
+        assert self.model is not None
+        if self.is_parallelized:
+            return self.model
+
+        if self.pp_parallelizer is not None:
+            assert callable(self.pp_parallelizer)
+            self.model = self.pp_parallelizer(self.model)
+
+        if self.tp_parallelizer is not None:
+            assert callable(self.tp_parallelizer)
+            self.model, self.layer_param_placements = self.tp_parallelizer(
+                self.model
+            )
+        if self.sharding_parallelizer is not None:
+            assert callable(self.sharding_parallelizer)
+            self.model = self.sharding_parallelizer(self.model)
+        self._shard_all_param(self.model)
+        self.is_parallelized = True
+
+        return self.model
+
+    def _process_share_weight_layer(
+        self, layer, origin_weight, param_name, param_placements
+    ):
+        ipp = (
+            layer.pipeline_stage_index
+            if hasattr(layer, "pipeline_stage_index")
+            else 0
+        )
+
+        def create_pre_hook(origin_weight, param_name):
+            def forward_pre_hook(layer, input):
+                setattr(
+                    layer,
+                    param_name,
+                    None,
+                )
+                delattr(layer, param_name)
+                mesh = self.get_mesh(ipp)
+                share_weight = dist.reshard(
+                    origin_weight,
+                    mesh,
+                    param_placements,
+                )
+                setattr(
+                    layer,
+                    param_name,
+                    share_weight,
+                )
+
+            return forward_pre_hook
+
+        def create_post_hook(origin_weight, param_name):
+            def forward_post_hook(layer, input, output):
+                setattr(
+                    layer,
+                    param_name,
+                    origin_weight,
+                )
+
+            return forward_post_hook
+
+        layer.register_forward_pre_hook(
+            create_pre_hook(origin_weight, param_name)
+        )
+        layer.register_forward_post_hook(
+            create_post_hook(origin_weight, param_name)
+        )
+
+    def _shard_all_param(self, model):
+        param_name_to_shard_param = {}
+        param_name_to_pp_stage = {}
+
+        def shard_layer_param(layer):
+            if self.pp_parallelizer is not None:
+                assert hasattr(layer, "pipeline_stage_index")
+            for param_name in list(layer._parameters.keys()):
+                param = getattr(layer, param_name)
+                if param is not None:
+                    param_full_name = param.name
+                    ipp = (
+                        layer.pipeline_stage_index
+                        if hasattr(layer, "pipeline_stage_index")
+                        else 0
+                    )
+                    mesh = self.get_mesh(ipp)
+                    param_placements = [
+                        dist.Replicate() for _ in range(len(mesh._shape))
+                    ]
+                    if layer in self.layer_param_placements:
+                        if param_name in self.layer_param_placements[layer]:
+                            param_placements = (
+                                self.layer_param_placements[layer][param_name]
+                                if self.layer_param_placements[layer][
+                                    param_name
+                                ]
+                                is not None
+                                else param_placements
+                            )
+                    if not param.is_dist():
+                        if param_full_name in param_name_to_shard_param:
+                            setattr(
+                                layer,
+                                param_name,
+                                param_name_to_shard_param[param_full_name],
+                            )
+                            if ipp != param_name_to_pp_stage[param_full_name]:
+                                self._process_share_weight_layer(
+                                    layer,
+                                    param_name_to_shard_param[param_full_name],
+                                    param_name,
+                                    param_placements,
+                                )
+                        else:
+                            param = dist.shard_tensor(
+                                param, mesh, param_placements
+                            )
+                            param_name_to_shard_param[param_full_name] = param
+                            param_name_to_pp_stage[param_full_name] = ipp
+                            setattr(layer, param_name, param)
+                    else:
+                        if (
+                            param_full_name in param_name_to_shard_param
+                            and ipp != param_name_to_pp_stage[param_full_name]
+                        ):
+                            self._process_share_weight_layer(
+                                layer,
+                                param_name_to_shard_param[param_full_name],
+                                param_name,
+                                param_placements,
+                            )
+                        elif param_full_name not in param_name_to_shard_param:
+                            param_name_to_shard_param[param_full_name] = param
+                            param_name_to_pp_stage[param_full_name] = ipp
+
+        for name, layer in model.named_sublayers():
+            shard_layer_param(layer)
+
+
+def parallelize_model_and_optimizer(model, optimizer=None):
+    if not isinstance(model, ParallelModel):
+        assert not isinstance(optimizer, ParallelOptimizer)
+        logging.warning(
+            "The method `parallelize_model_and_optimizer` won't do anything since the model is not parallelized."
+        )
+        return model, optimizer
+    parallelized_model = model.parallelize_model()
+    parallelized_optimizer = None
+    if optimizer is not None:
+        assert isinstance(optimizer, ParallelOptimizer)
+        optimizer.update_param_list(parallelized_model.parameters())
+        parallelized_optimizer = optimizer.parallelize()
+
+    return parallelized_model, parallelized_optimizer

python/paddle/distributed/auto_parallel/intermediate/parallelize.py

@@ -0,0 +1,385 @@
+#   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 __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, TypedDict
+
+from typing_extensions import NotRequired
+
+from paddle.distributed import fleet
+from paddle.framework import core
+
+from .parallel_base import ParallelOptimizer, parallelize_model_and_optimizer
+from .pipeline_parallel import pipeline_parallel
+from .sharded_data_parallel import sharded_data_parallel
+from .tensor_parallel import tensor_parallel
+
+if TYPE_CHECKING:
+    import paddle
+
+    from .pipeline_parallel import SplitPoint
+    from .tensor_parallel import PlanBase
+
+    class _DPConfig(TypedDict):
+        sharding_level: str | int
+
+    class _MPConfig(TypedDict):
+        parallelize_plan: dict[str, PlanBase | list[PlanBase]]
+
+    class _PPConfig(TypedDict):
+        split_spec: str | dict[str, SplitPoint]
+        global_spec: NotRequired[str]
+
+    class _ParallelizeConfig(TypedDict):
+        dp_config: NotRequired[_DPConfig]
+        mp_config: NotRequired[_MPConfig]
+        pp_config: NotRequired[_PPConfig]
+
+
+def parallelize(
+    model: paddle.nn.Layer,
+    optimizer: paddle.optimizer.Optimizer | None = None,
+    mesh: paddle.distributed.ProcessMesh | None = None,
+    config: _ParallelizeConfig | None = None,
+) -> tuple[paddle.nn.Layer, paddle.optimizer.Optimizer]:
+    """
+
+    Parallelize the model and optimizer from a single card version to a distributed version.
+
+    Args:
+        model (paddle.nn.Layer): the model to be parallelized.
+        optimizer (paddle.optimizer.Optimizer, optional): the optimizer to be parallelized.
+            Could be `None` if no optimizer to be parallelized.
+        mesh (paddle.distributed.ProcessMesh, optional): the process mesh for parallelize the model and the optimizer.
+            Best practice: calling `dist.auto_parallel.set_mesh` to set the global mesh ahead of calling `parallelize`
+            and keep the `mesh` parameter as `None.
+            If the `mesh` is not None, the mesh passed to `parallelize` will overwrite the mesh set by `set_mesh`.
+        config (dict, optional): a dict contains the parallel config.
+            The keys of the dict can be chosen from `dp_config`, `mp_config` and `pp_config` which will be used to
+            determine the parallel method for data parallel, tensor parallel and pipeline parallel separately.
+            A valid config can be like this: {"dp_config": for more information refer the `dp_config` section of
+            this doc, "mp_config": for more information refer the `mp_config` section of this doc, "pp_config":
+            for more information refer the `pp_config` section of this doc}.
+
+            dp_config (dict): a dict specifying the data parallel config. The keys of `dp_config` is `sharding_level`.
+                The value of `sharding_level` can be chosen from 0/1/2/3, which means pure data parallel, sharding
+                parallel stage 1, sharding parallel stage 2 and sharding parallel stage 3  separately. A valid
+                dp_config can be like this: {"sharding_level": 2}.
+
+            mp_config (dict): a dict specifying the tensor parallel config. The keys of `mp_config` is
+                `parallelize_plan`. The value of `parallelize_plan` is another dict, mapping a layer name or a param
+                name to a specific parallel plan. Note that the layer name could be written in regular format. If
+                mapping a param name to a specific plan, the name of the param must be ended with `weight` or `bias`.
+                And all valid parallel plan is `ColWiseParallel`, `RowWiseParallel`, `SequenceParallelBegin,
+                `SequenceParallelDisable`, `SequenceParallelEnable`, `SequenceParallelEnd`, `PrepareLayerInput` and
+                `PrepareLayerOutput`. A valid mp_config can be like this: {"llama.embed_tokens": dist.ColWiseParallel(),
+                "llama.norm": dist.SequenceParallelEnable(), "lm_head.weight": dist.ColWiseParallel()}.
+
+            pp_config (dict): a dict specifying the pipeline parallel config. The keys of `pp_config` is `split_spec`
+                and `global_spec`. The `split_spec` can be a dict or a string. If the `split_spec` is a dict, it maps
+                a layer name to a `SplitPoint`, note that the layer name could be written in regular format. The
+                pipeline parallel will exactly split the model at the point indicated by the map. If the `split_spec`
+                is a string, it contains the prefix of a set of layers. The pipeline parallel will automatically split
+                the model evenly at target layer. The `global_spec` is a string indicating a layer that contains global
+                tensors, which will be duplicated through all stages of the pipeline parallel. Some valid pp_config
+                can be list these:  {"split_spec": "llama.layers", "global_spec": "llama.global_layer"}
+                or {"split_spec": {"llama.layers.1": SplitPoint.END}}.
+
+            cp_config (dict): a dict specifying the context parallel config. The keys of `cp_config` is
+                `parallelize_plan`. The value of `parallelize_plan` is another dict, mapping a layer name or a param
+                name to a specific parallel plan. All valid parallel plan is `ContextParallel` and `PrepareContextParallel`.
+                A valid cp_config can be like this: {"llama": dist.PrepareContextParallel('p2p'),
+                "llama.sdpa": dist.ContextParallel('p2p')}.
+
+    Note:
+        If the mesh is `None` or neither of `dp_config`, `mp_config`, `pp_config` and `cp_config` is in the config, this
+        api will do nothing but return the model and optimizer passed in.
+
+    Returns:
+        model, optimizer: the model and the optimizer after parallelize
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class ModelConfig:
+            ...     def __init__(self):
+            ...         self.vocab_size = 10
+            ...         self.hidden_size = 20
+            ...         self.intermediate_size = 20
+            ...         self.num_layers = 2
+
+            >>> model_config = ModelConfig()
+
+            >>> class LlamaRMSNorm(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.weight = paddle.create_parameter(
+            ...             shape=[model_config.hidden_size],
+            ...             dtype=paddle.get_default_dtype(),
+            ...         )
+            ...
+            ...     def forward(self, input):
+            ...         pass
+
+            >>> class LlamaAttention(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...
+            ...         self.qkv_proj = paddle.nn.Linear(
+            ...             model_config.hidden_size,
+            ...             model_config.hidden_size * 3,
+            ...             bias_attr=False,
+            ...         )
+            ...
+            ...         self.o_proj = paddle.nn.Linear(
+            ...             model_config.hidden_size,
+            ...             model_config.hidden_size,
+            ...             bias_attr=False,
+            ...         )
+            ...
+            ...     def forward(self, input):
+            ...         pass
+
+            >>> class LlamaMLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.gate_up_proj = paddle.nn.Linear(
+            ...             model_config.hidden_size,
+            ...             model_config.intermediate_size * 2,
+            ...             bias_attr=False,
+            ...         )
+            ...
+            ...         self.down_proj = paddle.nn.Linear(model_config.intermediate_size, model_config.hidden_size, bias_attr=False)
+            ...
+            ...     def forward(self, input):
+            ...         pass
+
+            >>> class LlamaDecoderLayer(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.self_attn = LlamaAttention()
+            ...         self.mlp = LlamaMLP()
+            ...         self.input_layernorm = LlamaRMSNorm()
+            ...         self.post_attention_layernorm = LlamaRMSNorm()
+            ...
+            ...     def forward(self, input):
+            ...         pass
+
+            >>> class LlamaModel(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.embedding = paddle.nn.Embedding(model_config.vocab_size, model_config.hidden_size)
+            ...         decoder_layers = []
+            ...         for _ in range(model_config.num_layers):
+            ...             decoder_layers.append(LlamaDecoderLayer())
+            ...
+            ...         self.layers = paddle.nn.LayerList(decoder_layers)
+            ...         self.norm = LlamaRMSNorm()
+            ...
+            ...     def forward(self, input):
+            ...         pass
+
+            >>> class LlamaLMHead(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.weight = self.create_parameter(
+            ...             shape=[model_config.hidden_size, model_config.vocab_size],
+            ...             dtype=paddle.get_default_dtype(),
+            ...         )
+            ...
+            ...     def forward(self, input):
+            ...         pass
+
+            >>> class LlamaForCausalLM(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.llama = LlamaModel()
+            ...         self.lm_head = LlamaLMHead()
+            ...
+            ...     def forward(self, input):
+            ...         pass
+
+            >>> mesh = dist.ProcessMesh([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=["dp", "mp", "pp"])
+            >>> dist.auto_parallel.set_mesh(mesh)
+            >>> parallel_config = {
+            ...     "dp_config": {'sharding_level': 1},
+            ...     "mp_config": {
+            ...         "parallelize_plan": {
+            ...             "llama.embed_tokens": [
+            ...                 dist.ColWiseParallel(),
+            ...                 dist.SequenceParallelBegin(),
+            ...             ],
+            ...             "llama.position_embedding": [
+            ...                 dist.ColWiseParallel(),
+            ...                 dist.SequenceParallelBegin(),
+            ...             ],
+            ...             "llama.layers.*.self_attn.qkv_proj": dist.ColWiseParallel(),
+            ...             "llama.layers.*.self_attn.o_proj": dist.RowWiseParallel(),
+            ...             "llama.layers.*.self_attn": dist.SequenceParallelDisable(),
+            ...             "llama.layers.*.mlp.gate_up_proj": dist.ColWiseParallel(),
+            ...             "llama.layers.*.mlp.down_proj": dist.RowWiseParallel(),
+            ...             "llama.layers.*.mlp": dist.SequenceParallelDisable(need_transpose=False),
+            ...             "lm_head.weight": dist.ColWiseParallel(),
+            ...             "lm_head": dist.SequenceParallelEnd(),
+            ...         }
+            ...     },
+            ...     "pp_config": {'split_spec': "llama.layers"},
+            ... }
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> model = LlamaForCausalLM()
+            >>> optimizer = paddle.optimizer.AdamW(parameters=model.parameters())
+            >>> dist_model, dist_optimizer = dist.parallelize(model, optimizer, config=parallel_config)  # type: ignore[arg-type]
+            >>> # This case need to be executed in multi-card environment
+            >>> # python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 {test_case}.py
+
+    """
+    if config is None:
+        warnings.warn(
+            "The `parallelize will do nothing since the config is `None`."
+        )
+        return model, optimizer
+    assert isinstance(config, dict)
+    if mesh is not None:
+        assert isinstance(mesh, core.ProcessMesh), (
+            "The mesh must be an instance of paddle.distributed.ProcessMesh."
+        )
+        g_mesh = fleet.auto.get_mesh()
+        if g_mesh is not None and g_mesh != mesh:
+            warnings.warn(
+                "The mesh set by `fleet.auto.set_mesh` is different with the mesh pass to "
+                "`parallelize`. Will overwrite the previous mesh"
+            )
+        fleet.auto.set_mesh(mesh)
+    pp_config = config.get('pp_config')
+    mp_config = config.get('mp_config')
+    dp_config = config.get('dp_config')
+    cp_config = config.get('cp_config')
+    if pp_config is not None:
+        assert isinstance(pp_config, dict)
+        model, optimizer = pipeline_parallel(
+            model,
+            optimizer,
+            pp_config,
+        )
+    if mp_config is not None:
+        assert isinstance(mp_config, dict)
+        if cp_config is not None:
+            assert isinstance(cp_config, dict)
+            assert "parallelize_plan" in cp_config.keys()
+            assert "parallelize_plan" in mp_config.keys()
+            mp_config['parallelize_plan'].update(cp_config['parallelize_plan'])
+        model, optimizer = tensor_parallel(model, optimizer, mp_config)
+    elif cp_config is not None:
+        assert isinstance(cp_config, dict)
+        model, optimizer = tensor_parallel(
+            model,
+            optimizer,
+            cp_config,
+        )
+    if dp_config is not None:
+        assert isinstance(dp_config, dict)
+        if 'sharding_level' not in dp_config.keys():
+            warnings.warn(
+                "The dp_config doesn't contain sharding_level, will run under dp."
+            )
+        model, optimizer = sharded_data_parallel(
+            model,
+            optimizer,
+            config=dp_config,
+        )
+    model, optimizer = parallelize_model_and_optimizer(model, optimizer)
+    return model, optimizer
+
+
+has_parallelized_model = False
+
+
+def parallelize_model(model, mesh=None, config=None):
+    if config is None:
+        warnings.warn(
+            "The `parallelize_model will do nothing since the config is `None`."
+        )
+        return model
+    assert isinstance(config, dict)
+    if mesh is not None:
+        assert isinstance(mesh, core.ProcessMesh), (
+            "The mesh must be an instance of paddle.distributed.ProcessMesh."
+        )
+        g_mesh = fleet.auto.get_mesh()
+        if g_mesh is not None and g_mesh != mesh:
+            warnings.warn(
+                "The mesh set by `fleet.auto.set_mesh` is different with the mesh pass to "
+                "`parallelize_model`. Will overwrite the previous mesh"
+            )
+        fleet.auto.set_mesh(mesh)
+    global has_parallelized_model
+    has_parallelized_model = True
+    model, _ = parallelize(model, None, mesh, config)
+    return model
+
+
+def parallelize_optimizer(optimizer, mesh=None, config=None):
+    if config is None:
+        warnings.warn(
+            "The `parallelize_optimizer will do nothing since the config is `None`."
+        )
+        return optimizer
+    assert isinstance(config, dict)
+    if mesh is not None:
+        assert isinstance(mesh, core.ProcessMesh), (
+            "The mesh must be an instance of paddle.distributed.ProcessMesh."
+        )
+        g_mesh = fleet.auto.get_mesh()
+        if g_mesh is not None and g_mesh != mesh:
+            warnings.warn(
+                "The mesh set by `fleet.auto.set_mesh` is different with the mesh pass to "
+                "`parallelize_optimizer`. Will overwrite the previous mesh"
+            )
+        fleet.auto.set_mesh(mesh)
+
+    global has_parallelized_model
+    assert has_parallelized_model, (
+        "Please parallelize the model before parallelize optimizer."
+    )
+    param_list = optimizer._parameter_list
+    if isinstance(param_list[0], dict):
+        for param_group in param_list:
+            for param in param_group['params']:
+                assert param.is_dist(), (
+                    "Please use model after parallelize to create optimizer."
+                )
+    else:
+        for param in param_list:
+            assert param.is_dist(), (
+                "Please use model after parallelize to create optimizer."
+            )
+
+    dp_config = config.get('dp_config')
+    level = None
+    sharding_mesh_dim = None
+    if dp_config is not None:
+        if 'sharding_level' not in dp_config.keys():
+            warnings.warn(
+                "The dp_config doesn't contain sharding_level, will run under dp."
+            )
+        level = dp_config.get('sharding_level')
+        sharding_mesh_dim = dp_config.get('sharding_mesh_dim', "dp")
+    optimizer = ParallelOptimizer(optimizer, level, sharding_mesh_dim)
+    optimizer = optimizer.parallelize()
+    return optimizer

python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py

@@ -0,0 +1,419 @@
+#   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 itertools
+import logging
+import re
+from collections import OrderedDict
+from enum import Enum
+
+import paddle.distributed as dist
+from paddle.distributed import fleet
+from paddle.distributed.utils.log_utils import get_logger
+
+from .parallel_base import ParallelModel, ParallelOptimizer, is_tensor
+
+logger = get_logger("INFO", __name__)
+
+
+class SplitPoint(Enum):
+    """
+    Marking the position of the split.
+    BEGINNING: will split the model before the specified layer.
+    END: will split the model after the specified layer.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> pp_config = {
+            ...     'fc1': dist.SplitPoint.END,
+            ... }
+    """
+
+    BEGINNING = 0
+    END = 1
+
+
+class PipelineParallel(ParallelModel):
+    def __init__(self, model, split_spec, global_spec, pipeline_layers=None):
+        super().__init__(model)
+        self.split_spec = split_spec
+        self.global_spec = global_spec
+        self.pipeline_layers = pipeline_layers
+        self.pp_parallelizer = self.pipeline_parallel_fn
+        self.name_to_layer = {}
+        for layer_name, layer in model.named_sublayers():
+            self.name_to_layer[layer_name] = layer
+
+    def get_layer_by_name(self, name):
+        assert name in self.name_to_layer, (
+            f"layer name:{name} not in the model, please check the split_spec"
+        )
+        return self.name_to_layer[name]
+
+    def pipeline_parallel_fn(self, model):
+        mesh = fleet.auto.get_mesh()
+        pipeline_stage_num = mesh.get_dim_size("pp")
+        assert len(self.split_spec) == pipeline_stage_num - 1
+
+        def forward_post_hook(layer, input, output):
+            pipeline_stage_index = layer.pipeline_stage_index
+            split_point = layer.split_point
+            assert split_point == SplitPoint.END
+            # reshard to next pipeline stage
+            if isinstance(output, (dict, OrderedDict)):
+                for key, tensor in output.items():
+                    assert is_tensor(tensor)
+                    output[key] = dist.reshard(
+                        tensor,
+                        self.get_mesh(pipeline_stage_index + 1),
+                        tensor.placements,
+                    )
+            elif isinstance(output, list):
+                for i in range(len(output)):
+                    assert is_tensor(output[i])
+                    output[i] = dist.reshard(
+                        output[i],
+                        self.get_mesh(pipeline_stage_index + 1),
+                        output[i].placements,
+                    )
+            elif isinstance(output, tuple):
+                output = list(output)
+                for i in range(len(output)):
+                    assert is_tensor(output[i])
+                    output[i] = dist.reshard(
+                        output[i],
+                        self.get_mesh(pipeline_stage_index + 1),
+                        output[i].placements,
+                    )
+                output = tuple(output)
+            elif is_tensor(output):
+                output = dist.reshard(
+                    output,
+                    self.get_mesh(pipeline_stage_index + 1),
+                    output.placements,
+                )
+            else:
+                raise ValueError(
+                    f"output between pp stages should be a dict of tensors or list of tensors or tuple of tensors or tensor, but {type(output)}"
+                )
+            return output
+
+        def forward_pre_hook(layer, input):
+            split_point = layer.split_point
+            assert split_point == SplitPoint.BEGINNING
+            # TODO(deepllz): support in the future
+            return input
+
+        # step1: set every layer's own pipeline_stage_index
+        split_layer_names = list(self.split_spec.keys())
+        sublayer_names = [name for name, _ in model.named_sublayers()]
+        # Mark which layer is the next pipeline stage
+        pipeline_layer_mark = [0 for _ in range(len(sublayer_names))]
+        for split_layer_name in split_layer_names:
+            split_point = self.split_spec[split_layer_name]
+            index = sublayer_names.index(split_layer_name)
+            if split_point == SplitPoint.END:
+                is_valid = False
+                for i in range(index + 1, len(sublayer_names)):
+                    if not sublayer_names[i].startswith(split_layer_name):
+                        pipeline_layer_mark[i] = 1
+                        is_valid = True
+                        break
+                assert is_valid, (
+                    f"the last layer:{split_layer_name} must not be SplitPoint.END, please check the split_spec"
+                )
+            else:
+                raise NotImplementedError(
+                    "SplitPoint.BEGINNING is not supported currently"
+                )
+                pipeline_layer_mark[index] = 1
+        # the inclusiveSum of pipeline_layer_mark is the pipeline stage index
+        pipeline_stage_index = list(itertools.accumulate(pipeline_layer_mark))
+        for index, (name, layer) in enumerate(model.named_sublayers()):
+            layer.pipeline_stage_index = pipeline_stage_index[index]
+
+        # step2: insert reshard
+        for name in split_layer_names:
+            layer = self.get_layer_by_name(name)
+            split_point = self.split_spec[name]
+            layer.split_point = split_point
+            if split_point == SplitPoint.END:
+                layer.register_forward_post_hook(forward_post_hook)
+            else:
+                raise NotImplementedError(
+                    "SplitPoint.BEGINNING is not supported currently"
+                )
+                layer.register_forward_pre_hook(forward_pre_hook)
+
+        if self.global_spec:
+            self.process_global_mesh_layers()
+
+        return model
+
+    def process_global_mesh_layers(self):
+        g_mesh = fleet.auto.get_mesh()
+        g_mesh = g_mesh.get_mesh_with_dim("pp")
+
+        def forward_post_hook(layer, input, output):
+            if isinstance(output, (list, tuple)):
+                global_output = list(output)
+                for ind in range(len(global_output)):
+                    output_i = global_output[ind]
+                    if is_tensor(output_i):
+                        if output_i.is_dist():
+                            global_output[ind] = dist.reshard(
+                                output_i,
+                                g_mesh,
+                                [
+                                    dist.Replicate()
+                                    for _ in range(len(g_mesh._shape))
+                                ],
+                            )
+                        else:
+                            global_output[ind] = dist.shard_tensor(
+                                output_i,
+                                g_mesh,
+                                [
+                                    dist.Replicate()
+                                    for _ in range(len(g_mesh._shape))
+                                ],
+                            )
+
+                if isinstance(output, tuple):
+                    global_output = tuple(global_output)
+                return global_output
+            elif is_tensor(output):
+                if output.is_dist():
+                    return dist.reshard(
+                        output,
+                        g_mesh,
+                        [dist.Replicate() for _ in range(len(g_mesh._shape))],
+                    )
+                else:
+                    return dist.shard_tensor(
+                        output,
+                        g_mesh,
+                        [dist.Replicate() for _ in range(len(g_mesh._shape))],
+                    )
+            else:
+                raise TypeError(
+                    "layer output can only be tensor or list/tuple of tensor"
+                )
+
+        def forward_pre_hook(layer, args, kwargs):
+            pp_idx = getattr(layer, "pipeline_stage_index", 0)
+            new_args = []
+            new_kwargs = {}
+
+            def reshard_not_mesh_match_tensor(arg):
+                cur_pp_mesh = self.get_mesh(pp_idx)
+                if (
+                    arg is not None
+                    and is_tensor(arg)
+                    and arg.is_dist()
+                    and arg.process_mesh != cur_pp_mesh
+                ):
+                    return dist.reshard(
+                        arg,
+                        cur_pp_mesh,
+                        [dist.Replicate(), dist.Replicate()],
+                    )
+                return arg
+
+            for arg in args:
+                new_args.append(reshard_not_mesh_match_tensor(arg))
+
+            for key, arg in kwargs.items():
+                new_kwargs[key] = reshard_not_mesh_match_tensor(arg)
+
+            return (tuple(new_args), new_kwargs)
+
+        # wa because of pir in vpp mode send receive bug
+        for layer_name in self.global_spec:
+            layer = self.get_layer_by_name(layer_name)
+            layer.register_forward_post_hook(forward_post_hook)
+
+        if self.pipeline_layers is not None:
+            for layer_name in self.pipeline_layers:
+                layer = self.get_layer_by_name(layer_name)
+                layer.register_forward_pre_hook(
+                    forward_pre_hook, with_kwargs=True
+                )
+        else:
+            for layer in self.name_to_layer.values():
+                layer.register_forward_pre_hook(
+                    forward_pre_hook, with_kwargs=True
+                )
+
+
+def pipeline_parallel(model, optimizer=None, config=None):
+    """
+    pipeline_parallel converts model and optimizer to pipelined distributed model
+
+    Args:
+        model (paddle.nn.Layer): A single card model to be distributed
+        optimizer (paddle.optimizer.Optimizer): An optimizer to be distributed
+        config (dict): {
+            "split_spec": OrderedDict|dict|str|list(str), The pipeline parallel split point.
+                if split_spec is a string or list, such as "llama.layer" or ["llama.layerA", "llama.layerB"], Then the layer with same prefix a will be divided equally according to the size of pipeline degree.
+                if split_spec is a OrderedDict|dict, key is the layer name, and the value is the split position that can be SplitPoint.BEGINNING or SplitPoint.END, the order of the keys is the order of the pipeline stage.
+                NOTE: dict is also ordered after python3.7, so use dict at this time.
+            "global_spec": str|list(str), make the output tensor of specific layers on global mesh.
+        }
+
+    Returns:
+        PipelineParallel: a distributed model
+        ParallelOptimizer: a distributed optimizer
+    """
+
+    split_spec = config.get("split_spec")
+    if split_spec is None:
+        logging.warning("No split_spec, pipeline parallel won't do anything.")
+        return model, optimizer
+
+    mesh = fleet.auto.get_mesh()
+    assert mesh is not None, (
+        "global mesh must not be None, please call fleet.auto.set_mesh(global_mesh) firstly"
+    )
+    assert "pp" in mesh.dim_names, (
+        "pp must in the mesh dim_names when use pipeline_parallel"
+    )
+
+    global_spec = config.get("global_spec")
+    if isinstance(split_spec, str):
+        split_spec = [split_spec]
+
+    matched_layer_name = None
+    if isinstance(split_spec, (list, tuple)):
+        # match layer_name with split_spec following by a dot and numbers and no other characters
+        # such as split_spec = ["llama.layer"], then llama.layer.0 is matched, llama.layer.0.mlp is not matched
+        patterns = [rf"{prefix}\.\d+$" for prefix in split_spec]
+
+        def is_match(layer_name):
+            for pattern in patterns:
+                if re.match(pattern, layer_name) or layer_name in split_spec:
+                    return True
+            return False
+
+        def filter_matched_layer(matched_layer_name):
+            # remove the base name if it has a numbered suffix
+            string_set = set(matched_layer_name)
+            to_remove = set()
+
+            numbered_pattern = re.compile(r'^(.+)\.\d+$')
+            for s in matched_layer_name:
+                match = numbered_pattern.match(s)
+                if match:
+                    base_name = match.group(1)
+                    if base_name in string_set:
+                        to_remove.add(base_name)
+
+            res = []
+            for s in matched_layer_name:
+                if s not in to_remove:
+                    res.append(s)
+            return res
+
+        matched_layer_name = [
+            name for name, _ in model.named_sublayers() if is_match(name)
+        ]
+        matched_layer_name = filter_matched_layer(matched_layer_name)
+        pp_size = mesh.get_dim_size("pp")
+        layer_num = len(matched_layer_name)
+        assert layer_num > 0, (
+            "No layer match the split_spec, please check its correctness"
+        )
+        assert layer_num >= pp_size, (
+            "The number of layers must not be less than the pp size"
+        )
+        if layer_num % pp_size != 0:
+            logger.warning(
+                f"The number of layers({layer_num}) must be divisible by the pp size({pp_size}), but got {layer_num} and {pp_size}"
+            )
+
+            def divide_list_indices(n, k):
+                base_size = n // k
+                extra = n % k
+
+                indices = []
+                current_index = -1
+
+                for i in range(k - 1):
+                    current_index += base_size
+                    if i < extra:
+                        current_index += 1
+                    indices.append(current_index)
+                return indices
+
+            indices = divide_list_indices(layer_num, pp_size)
+            split_spec_dict = OrderedDict(
+                [
+                    (matched_layer_name[indices[i]], SplitPoint.END)
+                    for i in range(pp_size - 1)
+                ]
+            )
+        else:
+            layers_per_rank = layer_num // pp_size
+            split_spec_dict = OrderedDict(
+                [
+                    (
+                        matched_layer_name[i * layers_per_rank - 1],
+                        SplitPoint.END,
+                    )
+                    for i in range(1, pp_size)
+                ]
+            )
+    else:
+        sublayer_names = [name for name, _ in model.named_sublayers()]
+        split_spec_dict = split_spec
+        for key, value in split_spec_dict.items():
+            assert key in sublayer_names, (
+                f"wrong split layer, expected one of {sublayer_names}"
+            )
+            assert value is SplitPoint.END, "not supported split point at now."
+
+    if global_spec:
+        if isinstance(global_spec, str):
+            global_spec = [global_spec]
+        else:
+            assert isinstance(global_spec, (list, tuple)), (
+                f"global_spec can only be list or list(str), but got:{type(global_spec)}"
+            )
+
+    logger.info(
+        f"split_spec_dict: {split_spec_dict}, global_spec: {global_spec}, matched_layer_name: {matched_layer_name}"
+    )
+
+    model = PipelineParallel(
+        model, split_spec_dict, global_spec, matched_layer_name
+    )
+    if optimizer is not None:
+        optimizer = ParallelOptimizer(optimizer)
+
+    return model, optimizer

python/paddle/distributed/auto_parallel/intermediate/sharded_data_parallel.py

@@ -0,0 +1,88 @@
+#   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 paddle.distributed import fleet
+
+from .parallel_base import ParallelModel, ParallelOptimizer
+
+
+class ShardedDataParallel(ParallelModel):
+    """
+    ShardedDataParallel converts a single card model to a distributed data parallel model
+
+    Args:
+        model (paddle.nn.Layer): A single card model to be distributed.
+        optimizer (paddle.optimizer.Optimizer): an optimizer to be distributed.
+        level (str): Zero stage, can be the following values:
+            0: no sharding (pure dp)
+            1: Zero Stage1
+            2: Zero Stage2
+            3: Zero Stage3
+            Default: None, which means optimizer is replicated among all process.
+        offload (bool): whether enable cpu offload strategy, not implemented currently.
+        exclude_layer (list): Specify which layers do not use the zero stage strategy, not implemented currently.
+    """
+
+    def __init__(
+        self,
+        model,
+        offload=False,
+        exclude_layer=None,
+    ):
+        super().__init__(model)
+        assert offload is False
+        assert exclude_layer is None
+
+        self.sharding_parallelizer = self.sharding_parallelizer_func
+
+    def sharding_parallelizer_func(self, model):
+        return model
+
+
+def sharded_data_parallel(model, optimizer=None, config=None):
+    """
+    sharded_data_parallel converts model and optimizer to distributed and supports set zero stage1/2/3
+
+    Args:
+        model (paddle.nn.Layer): A single card model to be distributed
+        optimizer (paddle.optimizer.Optimizer): an optimizer to be distributed
+        config (dict): {
+            "sharding_level": 0,
+            "offload": False,
+            "exclude_layer": None,
+            "sharding_mesh_dim": "dp",
+        }
+
+    Returns:
+        ShardedDataParallel: a distributed model
+        ParallelOptimizer: a distributed optimizer
+    """
+    sdp_model = ShardedDataParallel(
+        model, bool(config.get('offload')), config.get('exclude_layer')
+    )
+    if optimizer is not None:
+        level = config.get('sharding_level')
+        sharding_mesh_dim = config.get('sharding_mesh_dim', "dp")
+        optimizer = ParallelOptimizer(optimizer, level, sharding_mesh_dim)
+
+    # check global_mesh
+    mesh = fleet.auto.get_mesh()
+    assert mesh is not None, (
+        "global mesh must not be None, please call fleet.auto.set_mesh(global_mesh) firstly"
+    )
+    assert "dp" in mesh.dim_names, (
+        "dp must in the mesh dim_names when use sharded_data_parallel"
+    )
+    return sdp_model, optimizer

python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py

@@ -0,0 +1,954 @@
+#   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 __future__ import annotations
+
+import logging
+import re
+from typing import TYPE_CHECKING
+
+import paddle
+import paddle.distributed as dist
+
+from .parallel_base import ParallelModel, ParallelOptimizer, is_tensor
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
+    from paddle.distributed import ProcessMesh
+    from paddle.nn import Layer
+
+
+def c_split(x, process_mesh, need_transpose, split_type="sp"):
+    mp_index = process_mesh.dim_names.index('mp')  # get the axis for the split
+    dp_index = process_mesh.dim_names.index('dp')
+    if isinstance(x, tuple):
+        target_x = x[0]
+    else:
+        target_x = x
+    assert is_tensor(target_x)
+    assert len(target_x.shape) == 3
+    if need_transpose:
+        target_x = paddle.transpose(target_x, perm=[1, 0, 2])
+    placements = target_x.placements
+    if placements is None:
+        placements = [dist.Replicate() for _ in range(len(process_mesh.shape))]
+    if split_type == "sp":
+        if placements[dp_index] == dist.Shard(0):
+            # NOTE(zhangwl):if shard(0) , input shape should be [b,s,h]
+            split_dims = dist.Shard(1)
+        elif placements[dp_index] == dist.Shard(1):
+            # NOTE(zhangwl):if shard(1) , input shape should be [s,b,h]
+            split_dims = dist.Shard(0)
+        else:
+            logging.warning(
+                f"parallel api don't know {target_x.shape} which dimension is batch, default is to cut to the 0th dimension"
+            )
+            split_dims = dist.Shard(0)
+    elif split_type == "mp":
+        split_dims = dist.Shard(2)  # split h [b,s,h]
+    else:
+        raise ValueError(f"Unsupported split type {split_type}")
+
+    placements[mp_index] = split_dims
+    target_x = dist.reshard(target_x, process_mesh, placements)
+    if isinstance(x, tuple):
+        x = list(x)
+        x[0] = target_x
+        x = tuple(x)
+    else:
+        x = target_x
+
+    return x
+
+
+def c_concat(x, process_mesh, need_transpose):
+    index = process_mesh.dim_names.index('mp')  # get the axis for the split
+    if isinstance(x, tuple):
+        target_x = x[0]
+    else:
+        target_x = x
+    assert is_tensor(target_x)
+    assert len(target_x.shape) == 3
+    placements = target_x.placements
+    if placements is None:
+        placements = [dist.Replicate() for _ in range(len(process_mesh.shape))]
+    placements[index] = dist.Replicate()
+    target_x = dist.reshard(target_x, process_mesh, placements)
+    if need_transpose:
+        target_x = paddle.transpose(target_x, perm=[1, 0, 2])
+    if isinstance(x, tuple):
+        x = list(x)
+        x[0] = target_x
+        x = tuple(x)
+    else:
+        x = target_x
+
+    return x
+
+
+class PlanBase:
+    def __init__(self):
+        self.share_param_list = {}
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        raise NotImplementedError("Don't call the PlanBase directly.")
+
+
+class ColWiseParallel(PlanBase):
+    """
+    Col wise parallel plan for mp config.
+    Will try to split weight on the second dim and the bias on the first dim.
+    This api is designed for paddle.nn.Linear or paddle.nn.Embedding.
+    If any other instance of paddle.nn.Layer is passed,
+    this plan will try to split `layer.weight` and `layer.bias` if it has.
+
+    Note:
+        1. `layer.weight` should have two dims.
+        2. `layer.bias` should have one dim.
+
+    Args:
+        gather_output (bool): Whether gather the output to change it from a local tensor to a global tensor.
+            If gather the local tensor to global, an extra communication will be called.
+            The default value is `False`, which means keeping the output as a local tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.ColWiseParallel(),
+            ... }
+
+    """
+
+    def __init__(self, gather_output: bool = False) -> None:
+        super().__init__()
+        self.gather_output = gather_output
+
+    def gather_output_hook(self, process_mesh):
+        def gather_hook(layer, input, output):
+            assert output is not None
+            return c_concat(output, process_mesh, False)
+
+        return gather_hook
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        index = process_mesh.dim_names.index('mp')  # get the axis for the split
+        size = len(process_mesh.shape)
+        placement = [dist.Replicate() for _ in range(size)]
+        param_placements = {}
+        assert isinstance(layer, paddle.nn.Layer)
+        if not isinstance(layer, (paddle.nn.Linear, paddle.nn.Embedding)):
+            logging.warning(
+                f"ColWiseParallel is designed to handle Linear and Embedding. "
+                f"But got {layer.__class__.__name__}. "
+                f"Will try to shard weight and bias if the layer contains one."
+            )
+        shard_param_list = set(shard_param_list)
+        if len(shard_param_list) == 0:
+            shard_param_list.add("weight")
+            shard_param_list.add("bias")
+
+        def shard_param(param_name):
+            if (
+                hasattr(layer, param_name)
+                and getattr(layer, param_name) is not None
+            ):
+                layer_param = getattr(layer, param_name)
+
+                if layer_param.is_dist():
+                    return
+
+                if len(layer_param.shape) == 2:
+                    placement[index] = dist.Shard(1)
+                elif len(layer_param.shape) == 1:
+                    placement[index] = dist.Shard(0)
+                else:
+                    raise ValueError(f"{layer_param} should have 1 or 2 dims.")
+                # NOTE(zhangweilong):for share parameter, the parameter should be handled uniformly in the end
+                if (
+                    self.share_param_list is not None
+                    and layer_param.name in self.share_param_list
+                    and self.share_param_list[layer_param.name] > 1
+                ):
+                    param_placements.update({param_name: placement})
+                else:
+                    layer_param = dist.shard_tensor(
+                        layer_param,
+                        process_mesh,
+                        placement,
+                    )
+                    setattr(layer, param_name, layer_param)
+
+        for param_name in shard_param_list:
+            shard_param(param_name)
+        if self.gather_output:
+            layer.register_forward_post_hook(
+                self.gather_output_hook(process_mesh)
+            )
+        return param_placements
+
+
+class RowWiseParallel(PlanBase):
+    """
+    Row wise parallel plan for mp config.
+    Will try to split weight on the first dim.
+    This api is designed for paddle.nn.Linear or paddle.nn.Embedding.
+    If any other instance of paddle.nn.Layer is passed, this plan will try to split `layer.weight` if it has.
+
+    Note:
+        `layer.weight` should have two dims.
+
+    Args:
+        is_input_parallel (bool): Whether the input is a local tensor or a global tensor. If the input is a
+            global tensor, an extra split will be called. The default value is `True`,
+            which means the input is a local tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.RowWiseParallel(),
+            ... }
+    """
+
+    def __init__(self, is_input_parallel: bool = True) -> None:
+        super().__init__()
+        self.is_input_parallel = is_input_parallel
+
+    def split_input_hook(self, process_mesh):
+        def split_hook(layer, input):
+            return c_split(input, process_mesh, False, split_type="mp")
+
+        return split_hook
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        index = process_mesh.dim_names.index('mp')  # get the axis for the split
+        size = len(process_mesh.shape)
+        placement = [dist.Replicate() for _ in range(size)]
+        placement[index] = dist.Shard(0)
+        param_placements = {}
+        assert isinstance(layer, paddle.nn.Layer)
+        if not isinstance(layer, (paddle.nn.Linear, paddle.nn.Embedding)):
+            logging.warning(
+                f"RowWiseParallel is designed to handle Linear and Embedding. "
+                f"But got {layer.__class__.__name__}. "
+                f"Will try to shard weight if the layer contains one."
+            )
+        shard_param_list = set(shard_param_list)
+        shard_param_list.discard("bias")
+        if len(shard_param_list) == 0:
+            shard_param_list.add("weight")
+
+        def shard_param(param_name):
+            if (
+                hasattr(layer, param_name)
+                and getattr(layer, param_name) is not None
+            ):
+                layer_param = getattr(layer, param_name)
+                if layer_param.is_dist():
+                    return
+                if len(layer_param.shape) != 2:
+                    raise ValueError(f"{layer_param} should have 2 dims.")
+                # NOTE(zhangweilong):for share parameter, the parameter should be handled uniformly in the end
+                if (
+                    self.share_param_list is not None
+                    and layer_param.name in self.share_param_list
+                    and self.share_param_list[layer_param.name] > 1
+                ):
+                    param_placements.update({param_name: placement})
+                else:
+                    layer_param = dist.shard_tensor(
+                        layer_param,
+                        process_mesh,
+                        placement,
+                    )
+                    setattr(layer, param_name, layer_param)
+
+        for param_name in shard_param_list:
+            shard_param(param_name)
+        if not self.is_input_parallel:
+            layer.register_forward_pre_hook(self.split_input_hook(process_mesh))
+        return param_placements
+
+
+class PrepareLayerInput(PlanBase):
+    """
+    Prepare the input of specific layer. User should provide one callable function.
+
+    Args:
+        fn (callable): A function that prepare the layer input. The function should take exactly
+            one parameter named `process_mesh` and return the pre hook.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> def layer_input_hook(process_mesh):
+            ...     def hook(layer, input, output):
+            ...         return input
+            ...
+            ...     return hook
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.PrepareLayerOutput(layer_input_hook),
+            ... }
+    """
+
+    def __init__(
+        self,
+        fn: (
+            Callable[
+                [ProcessMesh],
+                Callable[
+                    [Layer, tuple[Tensor], tuple[Tensor]], [tuple[Tensor]]
+                ],
+            ]
+            | None
+        ) = None,
+    ) -> None:
+        super().__init__()
+        assert callable(fn)
+        self.fn = fn
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        layer.register_forward_pre_hook(self.fn(process_mesh=process_mesh))
+
+
+class PrepareLayerOutput(PlanBase):
+    """
+    Prepare the output of specific layer. User should provide one callable function.
+
+    Args:
+        fn (callable): A function that prepare the layer input. The function should take exactly
+            one parameter named `process_mesh` and return the post hook.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> def layer_output_hook(process_mesh):
+            ...     def hook(layer, input, output):
+            ...         return output
+            ...
+            ...     return hook
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.PrepareLayerOutput(layer_output_hook),
+            ... }
+    """
+
+    def __init__(
+        self,
+        fn: (
+            Callable[
+                [ProcessMesh],
+                Callable[
+                    [Layer, tuple[Tensor], tuple[Tensor]], [tuple[Tensor]]
+                ],
+            ]
+            | None
+        ) = None,
+    ) -> None:
+        super().__init__()
+        assert callable(fn)
+        self.fn = fn
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        layer.register_forward_post_hook(self.fn(process_mesh=process_mesh))
+
+
+class SequenceParallelBegin(PlanBase):
+    """
+    Sequence parallel plan for mp config.
+    This plan marks the beginning of the sp and should be added to the LAST layer before the sp range.
+
+    Note:
+        DON'T mark any layer in the sp range.
+
+    Args:
+        need_transpose (bool): the default value is `True`. With `need_transpose=True`, this plan will transfer
+            the output from [b, s, h] to [s/mp, b, h].  With `need_transpose=False`, this plan will transfer
+            the output from [s, b, h] to [s/mp, b, h].
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.SequenceParallelBegin(),
+            ... }
+    """
+
+    def __init__(self, need_transpose: bool = True) -> None:
+        super().__init__()
+        self.need_transpose = need_transpose
+
+    def sequence_parallel_begin(self, process_mesh):
+        def begin(layer, input, output):
+            assert output is not None
+            return c_split(output, process_mesh, self.need_transpose)
+
+        return begin
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        layer.register_forward_post_hook(
+            self.sequence_parallel_begin(process_mesh)
+        )
+
+
+class SequenceParallelEnd(PlanBase):
+    """
+    Sequence parallel plan for mp config.
+    This plan marks the ending of the sp and should be added to the FIRST layer after the sp range.
+
+    Note:
+        DON'T mark any layer in the sp range.
+
+    Args:
+        need_transpose (bool): the default value is `True`. With `need_transpose=True`, this plan will transfer
+            the input from [s/mp, b, h] to [b, s, h]. With `need_transpose=False`, this plan will transfer the
+            input from [s/mp, b, h] to [s, b, h].
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.SequenceParallelEnd(),
+            ... }
+    """
+
+    def __init__(self, need_transpose: bool = True) -> None:
+        super().__init__()
+        self.need_transpose = need_transpose
+
+    def sequence_parallel_end(self, process_mesh):
+        def end(layer, input, output=None):
+            assert input is not None
+            return c_concat(input, process_mesh, self.need_transpose)
+
+        return end
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        layer.register_forward_pre_hook(
+            self.sequence_parallel_end(process_mesh)
+        )
+
+
+class SequenceParallelEnable(PlanBase):
+    """
+    Sequence parallel plan for mp config.
+    Do sequence parallel on the layer. Note the input should be in [b, s, h] format.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.SequenceParallelEnable(),
+            ... }
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def sequence_parallel_begin(self, process_mesh):
+        def begin(layer, input, output=None):
+            assert input is not None
+            return c_split(input, process_mesh, True)
+
+        return begin
+
+    def sequence_parallel_end(self, process_mesh):
+        def end(layer, input, output):
+            assert output is not None
+            return c_concat(output, process_mesh, True)
+
+        return end
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        logging.warning(
+            "Sequence parallel with the usage of SequenceParallel may not reach the best throughput. "
+            "Try to use SequenceParallelBegin/End to achieve better performance"
+        )
+        layer.register_forward_pre_hook(
+            self.sequence_parallel_begin(process_mesh)
+        )
+        layer.register_forward_post_hook(
+            self.sequence_parallel_end(process_mesh)
+        )
+
+
+class SequenceParallelDisable(PlanBase):
+    """
+    Sequence parallel plan for mp config.
+    Disable sequence parallel on the layer.
+
+    Args:
+        need_transpose (bool): the default value is `True`. If the need_transpose is `True`: this plan will transfer
+            the input from  [s/mp, b, h] to [b, s, h] and then transfer the output from [b, s, h] to [s/mp, b, h].
+            If the need_transpose is `False`: this plan will transfer the input from  [s/mp, b, h] to [s, b, h] and
+            then transfer the output from [s, b, h] to [s/mp, b, h].
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> mp_config = {
+            ...     'fc1': dist.SequenceParallelDisable(),
+            ... }
+    """
+
+    def __init__(self, need_transpose: bool = True) -> None:
+        super().__init__()
+        self.need_transpose = need_transpose
+
+    def sequence_parallel_begin(self, process_mesh):
+        def begin(layer, input, output=None):
+            return c_split(output, process_mesh, self.need_transpose)
+
+        return begin
+
+    def sequence_parallel_end(self, process_mesh):
+        def end(layer, input, output=None):
+            return c_concat(input, process_mesh, self.need_transpose)
+
+        return end
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        layer.register_forward_pre_hook(
+            self.sequence_parallel_end(process_mesh)
+        )
+
+        layer.register_forward_post_hook(
+            self.sequence_parallel_begin(process_mesh)
+        )
+
+
+class ConvParallel(PlanBase):
+    """
+    A strategy for enabling spatial parallelism on ``paddle.nn.Conv2D`` layers
+    by sharding the input tensor along its Width (W) dimension.
+
+    When this ``ConvParallel`` configuration is applied to a ``Conv2D`` layer,
+    the layer's input tensor will have its width dimension split across devices
+    in the model parallel group. This can help reduce memory usage from activations,
+    especially when dealing with inputs that have a large width.
+
+    To enable width-wise input sharding correctly, make sure your `Conv2D` layer
+    satisfies the following conditions along the width dimension:
+
+    - **Dilation** must be set to `1`.
+    - **If no width padding is used:**
+        - The input width must be evenly divisible by the stride width.
+        - The stride width must be equal to the kernel width.
+    - **If width padding is used:**
+        - The stride width must be `1`.
+        - The total input width must be at least half the kernel width.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.nn as nn
+            >>> import paddle.distributed as dist
+
+            >>> class SimpleConvNet(nn.Layer):
+            ...     def __init__(self, data_format="NCHW"):
+            ...         super().__init__()
+            ...         self.conv1 = nn.Conv2D(
+            ...             3,
+            ...             8,
+            ...             kernel_size=3,
+            ...             padding=1,
+            ...             data_format=data_format,
+            ...         )
+            ...         self.relu = nn.ReLU()
+            ...
+            ...     def forward(self, x):
+            ...         x = self.conv1(x)
+            ...         return self.relu(x)
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> model = SimpleConvNet(data_format="NCHW")
+            >>> mp_config = {
+            ...     "parallelize_plan": {
+            ...         "conv1": dist.ConvParallel(),
+            ...     },
+            ... }
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @staticmethod
+    def _is_supported(
+        input_size,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        data_format,
+        mp_group_size,
+    ):
+        idx_w_input = -1
+        idx_w_kernel = -1
+
+        if data_format == "NCHW":
+            idx_w_input = 3
+            idx_w_kernel = 3
+        elif data_format == "NHWC":
+            idx_w_input = 2
+            idx_w_kernel = 3
+        else:
+            return False
+
+        if input_size[idx_w_input] % mp_group_size != 0:
+            return False
+
+        dilation_w = dilation[1]
+        padding_w = padding[1]
+        stride_w = stride[1]
+
+        input_w = input_size[idx_w_input]
+        kernel_w = kernel_size[idx_w_kernel]
+
+        if dilation_w != 1:
+            # RingConv2d only supports dilation=1.
+            # Larger dilation would require enlarged halo regions and more complex communication.
+            return False
+
+        if padding_w == 0:
+            # To avoid halo exchange when padding=0, we require:
+            # - input_w must be divisible by stride_w, so partitions align evenly across ranks.
+            # - stride_w == kernel_w, so each kernel operates on disjoint local regions.
+            if input_w % stride_w != 0:
+                return False
+            if stride_w != kernel_w:
+                return False
+
+        else:
+            # When padding > 0, halo exchange is needed.
+            # To simplify halo logic, we require:
+            # - stride_w == 1: ensures each output element is computed from overlapping input,
+            #   and no input region is skipped, simplifying halo construction.
+            # - kernel_w // 2 <= input_w: prevents the kernel from exceeding local input.
+            if stride_w != 1:
+                return False
+            if kernel_w // 2 > input_w:
+                return False
+
+        return True
+
+    def conv_parallel_start(self, process_mesh, data_format):
+        def start(layer, input, output=None):
+            if data_format == "NCHW":
+                shard_w_dim = 3
+            elif data_format == "NHWC":
+                shard_w_dim = 2
+            else:
+                raise ValueError(
+                    f"Unsupported data_format: {data_format}. "
+                    "Only NCHW and NHWC are supported."
+                )
+
+            if isinstance(input, tuple):
+                x = input[0]
+            else:
+                x = input
+
+            placements = x.placements
+            mp_index = process_mesh.dim_names.index('mp')
+            mp_group_size = process_mesh.get_dim_size('mp')
+
+            # Note(luchang): for intermediate api, when this ConvLayer is
+            # not supported, we just skip apply parallelization.
+            if not ConvParallel._is_supported(
+                x.shape,
+                layer.weight.shape,
+                layer._stride,
+                layer._updated_padding,
+                layer._dilation,
+                data_format,
+                mp_group_size,
+            ):
+                return input
+
+            if placements is None:
+                placements = [
+                    dist.Replicate() for _ in range(len(process_mesh.shape))
+                ]
+            if placements[mp_index] == dist.Shard(shard_w_dim):
+                return input
+
+            placements[mp_index] = dist.Shard(shard_w_dim)
+
+            if not x.is_dist():
+                x = dist.shard_tensor(x, process_mesh, placements)
+            else:
+                x = dist.reshard(x, process_mesh, placements)
+
+            if isinstance(input, tuple):
+                input = list(input)
+                input[0] = x
+                input = tuple(input)
+            else:
+                input = x
+            return input
+
+        return start
+
+    def apply(self, layer, process_mesh, shard_param_list):
+        layer.register_forward_pre_hook(
+            self.conv_parallel_start(process_mesh, layer._data_format)
+        )
+
+
+class TensorParallel(ParallelModel):
+    def __init__(self, model, parallelize_plan=None):
+        super().__init__(model)
+        if parallelize_plan is not None:
+            assert isinstance(parallelize_plan, dict)
+            for key, plan in parallelize_plan.items():
+                assert isinstance(key, str), (
+                    "The key of the parallelize plan should be a string."
+                )
+                if not isinstance(plan, list):
+                    plan = [plan]
+                for p in plan:
+                    assert isinstance(p, PlanBase), (
+                        "The value the the parallelize plan should be a instance of PlanBase or a list of PlanBase."
+                    )
+
+            self.global_mesh = dist.auto_parallel.get_mesh()
+            self.parallelize_plan = parallelize_plan
+            self.tp_parallelizer = self.tensor_parallelizer_fn
+
+    def match_layer(self, layer, name):
+        # Match the layer to a plan.
+        # Will return the plan if the layer hits one, otherwise return None.
+        plans = []
+        for key, plan in self.parallelize_plan.items():
+            attr_name = key.split('.')[-1]
+            shard_param_list = []
+            # Find some plan for specific parameter, such as
+            # "lm_head.weight": ColWiseParallel()
+            # "qkv_proj.lora_A" ColWiseParallel()
+            # if there is no plan for specific parameter, layer will be sharded by default: layer.weight and layer.bias
+            if key.endswith(f".{attr_name}"):
+                if hasattr(layer, attr_name) and is_tensor(
+                    getattr(layer, attr_name)
+                ):
+                    key = key.replace(f".{attr_name}", "")
+                    shard_param_list.append(attr_name)
+            re_find = re.match(key, name)
+            if key == name or (
+                re_find is not None
+                and int(re_find.end()) - int(re_find.start()) == len(name)
+            ):
+                if isinstance(plan, PlanBase):
+                    plan = [plan]
+                plans.append([plan, shard_param_list])
+        return plans
+
+    def tensor_parallelizer_fn(self, model):
+        if self.parallelize_plan is None:
+            return
+        layer_param_placements = {}
+        share_param_list = {}
+        for name, layer in model.named_sublayers():
+            for param_name in list(layer._parameters.keys()):
+                param = getattr(layer, param_name)
+                if param.name not in share_param_list:
+                    share_param_list[param.name] = 1
+                    continue
+                share_param_list[param.name] += 1
+        for name, layer in model.named_sublayers():
+            plans = self.match_layer(layer, name)
+            layer_param_placements[layer] = {}
+            if len(plans) > 0:
+                pp_idx = getattr(layer, "pipeline_stage_index", 0)
+                for plan in plans:
+                    real_plan, shard_param_list = plan
+                    for p in real_plan:
+                        p.share_param_list = share_param_list
+                        param_placements = p.apply(
+                            layer, self.get_mesh(pp_idx), shard_param_list
+                        )
+                        if param_placements is not None and param_placements:
+                            layer_param_placements[layer].update(
+                                param_placements
+                            )
+        return model, layer_param_placements
+
+
+def tensor_parallel(model, optimizer=None, config=None):
+    """
+    Tensor parallel.
+    Args:
+        model (paddle.nn.Layer): the model to be shard into tensor parallel.
+        optimizer (paddle.optimizer.Optimizer): the optimizer.
+        config (dict): {
+            "parallelize_plan": dict, the plan to shard the layer.
+        }
+    Returns:
+        model: model after tp
+        optimizer: optimizer after tp
+
+    NOTE: the plan should be a dict maps layer name or parameter name to a split_plan,
+    which will be used to split the layer or the parameter. The name can be written in regular format.
+
+    An example for the plan is:
+    ```
+    plan = {
+        "llama.embed_tokens": ColWiseParallel(),
+        "llama.layers.*.self_attn.q_proj": ColWiseParallel(),
+        "llama.layers.*.self_attn.k_proj": ColWiseParallel(),
+        "llama.layers.*.self_attn.v_proj": ColWiseParallel(),
+        "llama.layers.*.self_attn.o_proj": RowWiseParallel(),
+        "llama.layers.*.mlp.gate_proj": ColWiseParallel(),
+        "llama.layers.*.mlp.up_proj": ColWiseParallel(),
+        "llama.layers.*.mlp.down_proj": RowWiseParallel(),
+        "lm_head.weight": ColWiseParallel(),
+    }
+    ```
+    """
+    parallelize_plan = config.get("parallelize_plan")
+    if parallelize_plan is None:
+        # Do nothing if no plan.
+        logging.warning(
+            "No parallelize plan, tensor parallel won't do anything."
+        )
+        return model, optimizer
+
+    global_mesh = dist.auto_parallel.get_mesh()
+
+    assert global_mesh is not None, (
+        "global mesh must not be None, please call fleet.auto.set_mesh(global_mesh) firstly"
+    )
+    assert "mp" in global_mesh.dim_names, (
+        "mp must in the mesh dim_names when use tensor_parallel"
+    )
+
+    model = TensorParallel(model, parallelize_plan)
+    if optimizer is not None:
+        optimizer = ParallelOptimizer(optimizer)
+
+    return model, optimizer

python/paddle/distributed/auto_parallel/local_layer.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2025 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 __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import paddle
+import paddle.distributed as dist
+from paddle.nn import Layer
+
+if TYPE_CHECKING:
+    from paddle.distributed import Placement, ProcessMesh
+
+
+class LocalLayer(Layer):
+    """
+    The `LocalLayer` class is a specialized `Layer` for managing distributed tensors during
+    forward and backward passes in a parallelized training environment. It converts distributed tensors
+    to local tensors for computation and then back to distributed tensors as output, ensuring seamless
+    integration with distributed parallelism frameworks.
+
+    Args:
+        out_dist_attrs (list[tuple[ProcessMesh, list[Placement]]]):
+            A list where each entry is a tuple containing the `ProcessMesh` and the list of `Placement`
+            attributes for the corresponding output tensors. These attributes define the distribution
+            strategy for the outputs.
+        grad_dist_attrs (list[tuple[ProcessMesh, list[Placement]]]):
+            Similar to `out_dist_attrs` but for gradient tensors. The tuple in the list can be None, indicating that the dist_attr of the gradient tensor is same as the corresponding input tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from __future__ import annotations
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> from paddle import Tensor
+            >>> from paddle.distributed import ProcessMesh
+
+            >>> class CustomLayer(dist.LocalLayer):
+            ...     def __init__(self, out_dist_attrs, grad_dist_attrs):
+            ...         super().__init__(out_dist_attrs, grad_dist_attrs)
+            ...         self.local_result = paddle.to_tensor(0.0)
+
+            ...     def forward(self, x):
+            ...         mask = paddle.zeros_like(x)
+            ...         if dist.get_rank() == 0:
+            ...             mask[1:3] = 1
+            ...         else:
+            ...             mask[4:7] = 1
+
+            ...         x = x * mask
+            ...         mask_sum = paddle.sum(x)
+            ...         mask_sum = mask_sum / mask.sum()
+            ...         self.local_result = mask_sum
+            ...         return mask_sum
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> dist.init_parallel_env()
+            >>> mesh = ProcessMesh([0, 1], dim_names=["x"])
+            >>> dist_attrs = [
+            ...     (mesh, [dist.Partial(dist.ReduceType.kRedSum)]),
+            ... ]
+            >>> local_input = paddle.arange(0, 10, dtype="float32")
+            >>> local_input = local_input + dist.get_rank()
+            >>> input_dist = dist.auto_parallel.api.dtensor_from_local(
+            ...     local_input,
+            ...     mesh,
+            ...     [dist.Shard(0)],
+            ... )
+            >>> custom_layer = CustomLayer(dist_attrs, dist_attrs)
+            >>> output_dist = custom_layer(input_dist)
+
+            >>> local_value = custom_layer.local_result
+            >>> gathered_values: list[Tensor] = []
+            >>> dist.all_gather(gathered_values, local_value)
+
+            >>> print(f"[Rank 0] local_loss={gathered_values[0]}")
+            [Rank 0] local_loss=1.5
+            >>> print(f"[Rank 1] local_loss={gathered_values[1]}")
+            [Rank 1] local_loss=6.0
+            >>> print(f"global_loss (distributed)={output_dist}")
+            global_loss (distributed)=7.5
+
+            >>> # This case needs to be executed in a multi-card environment
+            >>> # export CUDA_VISIBLE_DEVICES=0,1
+            >>> # python -m paddle.distributed.launch {test_case}.py
+    """
+
+    def __init__(
+        self,
+        out_dist_attrs: list[tuple[ProcessMesh, list[Placement]]],
+        grad_dist_attrs: list[tuple[ProcessMesh, list[Placement]]],
+    ) -> None:
+        super().__init__()
+        self.out_dist_attrs = out_dist_attrs
+        self.grad_dist_attrs = grad_dist_attrs
+
+    def __call__(self, *inputs: Any, **kwargs: Any) -> Any:
+        """
+        Overrides the base `Layer`'s `__call__` method. Transforms distributed tensors to local tensors
+        before computation, invokes the parent class's `__call__` method, and then transforms the
+        outputs back to distributed tensors based on the specified distribution attributes.
+        """
+        inputs = list(inputs)
+        assert len(inputs) == len(self.grad_dist_attrs), (
+            f"The number of inputs ({len(inputs)}) does not match the number of grad_dist_attrs ({len(self.grad_dist_attrs)})."
+        )
+        for idx in range(len(inputs)):
+            if inputs[idx].is_dist():
+                if self.grad_dist_attrs[idx] is None:
+                    if paddle.in_dynamic_mode():
+                        mesh, placement = (
+                            inputs[idx].process_mesh,
+                            inputs[idx].placements,
+                        )
+                    else:
+                        mesh, placement = (
+                            inputs[idx].dist_attr().process_mesh,
+                            inputs[idx].dist_attr().placements,
+                        )
+                else:
+                    mesh, placement = (
+                        self.grad_dist_attrs[idx][0],
+                        self.grad_dist_attrs[idx][1],
+                    )
+
+                inputs[idx] = dist.auto_parallel.api.dtensor_to_local(
+                    inputs[idx], mesh, placement
+                )
+
+        outputs = Layer.__call__(self, *inputs, **kwargs)
+        list_outs = paddle.utils.flatten(outputs)
+        assert len(list_outs) == len(self.out_dist_attrs), (
+            f"The number of outputs ({len(list_outs)}) does not match the number of distribution attributes ({len(self.out_dist_attrs)})."
+        )
+
+        dist_outs = []
+        for idx in range(len(list_outs)):
+            dist_outs.append(
+                dist.auto_parallel.api.dtensor_from_local(
+                    list_outs[idx],
+                    self.out_dist_attrs[idx][0],
+                    self.out_dist_attrs[idx][1],
+                )
+            )
+        return paddle.utils.pack_sequence_as(outputs, dist_outs)

python/paddle/distributed/auto_parallel/local_map.py

@@ -0,0 +1,256 @@
+# Copyright (c) 2025 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 __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING, Any
+
+import paddle
+import paddle.distributed as dist
+from paddle.utils import flatten, pack_sequence_as
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle.distributed import ProcessMesh
+
+
+def local_map(
+    func: Callable[..., Any],
+    out_placements: list[list[dist.Placement]],
+    in_placements: list[list[dist.Placement]] | None = None,
+    process_mesh: ProcessMesh | None = None,
+    reshard_inputs: bool = False,
+) -> Callable[..., Any]:
+    """
+    The `local_map` API allows users to pass dist_tensors to a function that is written
+    to be applied on ``paddle.Tensor`` s. It works by extracting the local components
+    of dist_tensors, calling the function, and wrapping the outputs as dist_tensors
+    according to the ``out_placements``.
+
+    Args:
+        func (Callable): The function to be applied on each local shard of dist_tensors.
+
+        out_placements (list[list[dist.Placement]]):
+            The desired placements for each output tensor. Must be a list where each element
+            is a list of Placement objects specifying the distribution strategy for that
+            output tensor. The length of the outer list must match the number of outputs
+            from ``func``. For non-tensor outputs, the corresponding placement must be None.
+            When there are no dist_tensor inputs, process_mesh must be specified to use
+            non-None placements.
+
+        in_placements (Optional[list[list[dist.Placement]]], optional):
+            The required placements for each input tensor. If specified, must be a list
+            where each element is a list of Placement objects defining the distribution
+            strategy for that input tensor. The length of the outer list must match the
+            number of input tensors.
+            Default: None
+
+        process_mesh (ProcessMesh, optional):
+            The process mesh that all dist_tensors are placed on. If not specified,
+            this will be inferred from the input dist_tensors' process mesh.
+            local_map requires all dist_tensors to be placed on the same process mesh.
+            Must be specified when there are no dist_tensor inputs but out_placements
+            contains non-None values.
+            Default: None
+
+        reshard_inputs (bool, optional):
+            the bool value indicating whether to reshard the input :dist_tensors when
+            their placements are different from the required input placements. If this
+            value is ``False`` and some :dist_tensor input has a different placement,
+            an exception will be raised. Default: False.
+
+    Returns:
+        Callable: A function that applies func to local shards of input dist_tensors and returns dist_tensors or original values.
+
+    Example:
+        .. code-block:: pycon
+
+            >>> from __future__ import annotations
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> from paddle import Tensor
+            >>> from paddle.distributed import ProcessMesh
+
+            >>> def custom_function(x):
+            ...     mask = paddle.zeros_like(x)
+            ...     if dist.get_rank() == 0:
+            ...         mask[1:3] = 1
+            ...     else:
+            ...         mask[4:7] = 1
+            ...     x = x * mask
+            ...     mask_sum = paddle.sum(x)
+            ...     mask_sum = mask_sum / mask.sum()
+            ...     return mask_sum
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> dist.init_parallel_env()
+            >>> mesh = ProcessMesh([0, 1], dim_names=["x"])
+            >>> local_input = paddle.arange(0, 10, dtype="float32")
+            >>> local_input = local_input + dist.get_rank()
+            >>> input_dist = dist.auto_parallel.api.dtensor_from_local(local_input, mesh, [dist.Shard(0)])
+            >>> wrapped_func = dist.local_map(
+            ...     custom_function,
+            ...     out_placements=[[dist.Partial(dist.ReduceType.kRedSum)]],
+            ...     in_placements=[[dist.Shard(0)]],
+            ...     process_mesh=mesh,
+            ... )
+            >>> output_dist = wrapped_func(input_dist)
+
+            >>> local_value = output_dist._local_value()
+            >>> gathered_values: list[Tensor] = []
+            >>> dist.all_gather(gathered_values, local_value)
+
+            >>> print(f"[Rank 0] local_value={gathered_values[0].item()}")
+            [Rank 0] local_value=1.5
+            >>> print(f"[Rank 1] local_value={gathered_values[1].item()}")
+            [Rank 1] local_value=6.0
+            >>> print(f"global_value (distributed)={output_dist.item()}")
+            global_value (distributed)=7.5
+
+            >>> # This case needs to be executed in a multi-card environment
+            >>> # export CUDA_VISIBLE_DEVICES=0,1
+            >>> # python -m paddle.distributed.launch {test_case}.py
+    """
+
+    def wrapped(process_mesh: ProcessMesh | None, *args, **kwargs):
+        # Process input arguments
+        flat_dist_args = flatten(args)
+        if in_placements is not None:
+            assert len(in_placements) == len(flat_dist_args), (
+                f"in_placements length {len(in_placements)} does not match "
+                f"number of input args {len(flat_dist_args)}!"
+            )
+
+        flat_local_args = []
+        seen_dist_tensor = False
+
+        for idx, arg in enumerate(flat_dist_args):
+            if dist.auto_parallel.api.is_dist_tensor(arg):
+                dist_tensor = arg
+                if process_mesh is None:
+                    if paddle.in_dynamic_mode():
+                        process_mesh = dist_tensor.process_mesh
+                    else:
+                        process_mesh = dist_tensor.dist_attr().process_mesh
+
+                seen_dist_tensor = True
+
+                if in_placements is not None:
+                    in_placement = in_placements[idx]
+                    if in_placement is None:
+                        if paddle.in_dynamic_mode():
+                            in_placement = dist_tensor.placements
+                        else:
+                            in_placement = dist_tensor.dist_attr().placements
+                    else:
+                        if paddle.in_dynamic_mode():
+                            if in_placement != dist_tensor.placements:
+                                if reshard_inputs:
+                                    dist_tensor = dist.reshard(
+                                        dist_tensor, process_mesh, in_placement
+                                    )
+                                else:
+                                    raise ValueError(
+                                        f"in_placement {in_placement} does not match dist_tensor.placements {dist_tensor.placements}"
+                                    )
+
+                        else:
+                            if (
+                                in_placement
+                                != dist_tensor.dist_attr().placements
+                            ):
+                                if reshard_inputs:
+                                    dist_tensor = dist.reshard(
+                                        dist_tensor, process_mesh, in_placement
+                                    )
+                                else:
+                                    raise ValueError(
+                                        f"in_placement {in_placement} does not match dist_tensor.dist_attr().placements {dist_tensor.dist_attr().placements}"
+                                        "If reshard_inputs is wanted, set "
+                                        "reshard_inputs=True to local_map."
+                                    )
+                local_tensor = dist.auto_parallel.api.dtensor_to_local(
+                    dist_tensor, process_mesh, in_placement
+                )
+                flat_local_args.append(local_tensor)
+            else:
+                flat_local_args.append(arg)
+
+        local_args = pack_sequence_as(args, flat_local_args)
+        out = func(*local_args, **kwargs)
+        original_out = out
+        if seen_dist_tensor:
+            flat_out = flatten(out)
+            assert len(flat_out) == len(out_placements), (
+                "local_map requires one PlacementType for each output value, "
+                f"got {len(out_placements)} placements but expected "
+                f"{len(flat_out)}!"
+            )
+
+            flat_dist_and_arg_out = []
+            for out, out_placement in zip(flat_out, out_placements):
+                if paddle.in_dynamic_mode():
+                    if isinstance(out, paddle.Tensor):
+                        assert not dist.auto_parallel.api.is_dist_tensor(out), (
+                            f"Expected dense tensor output but got {type(out)}: {out}"
+                        )
+
+                        flat_dist_and_arg_out.append(
+                            dist.auto_parallel.api.dtensor_from_local(
+                                out, process_mesh, out_placement
+                            )
+                        )
+                    else:
+                        assert out_placement is None, (
+                            f"Expected None placements for non-tensor output {out} "
+                            f"but got {out_placement}!"
+                        )
+                        flat_dist_and_arg_out.append(out)
+                else:
+                    if isinstance(out, paddle.base.libpaddle.pir.Value):
+                        assert not dist.auto_parallel.api.is_dist_tensor(out), (
+                            f"Expected dense tensor output but got {type(out)}: {out}"
+                        )
+
+                        flat_dist_and_arg_out.append(
+                            dist.auto_parallel.api.dtensor_from_local(
+                                out, process_mesh, out_placement
+                            )
+                        )
+                    else:
+                        assert out_placement is None, (
+                            f"Expected None placements for non-tensor output {out} "
+                            f"but got {out_placement}!"
+                        )
+                        flat_dist_and_arg_out.append(out)
+            return pack_sequence_as(original_out, flat_dist_and_arg_out)
+        else:
+            flat_out = flatten(out)
+            flat_dist_and_arg_out = []
+            for out, out_placement in zip(flat_out, out_placements):
+                if out_placement is not None:
+                    assert process_mesh is not None, (
+                        "process_mesh must be specified when out_placements is not None"
+                    )
+                    flat_dist_and_arg_out.append(
+                        dist.auto_parallel.api.dtensor_from_local(
+                            out, process_mesh, out_placement
+                        )
+                    )
+                else:
+                    flat_dist_and_arg_out.append(out)
+            return pack_sequence_as(original_out, flat_dist_and_arg_out)
+
+    return functools.partial(wrapped, process_mesh)

python/paddle/distributed/auto_parallel/moe_utils.py

@@ -0,0 +1,468 @@
+# 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 __future__ import annotations
+
+import copy
+import os
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import Tensor
+from paddle.autograd import PyLayer
+
+from .placement_type import check_placements_equal, to_dim_map
+from .static.reshard_funcs.base_reshard_func import choose_reshard_func
+from .static.reshard_funcs.nd_mesh_reshard_func import get_1D_sub_process_mesh
+from .static.utils import split_mesh
+
+if TYPE_CHECKING:
+    from paddle.distributed import Placement
+    from paddle.distributed.auto_parallel.process_mesh import ProcessMesh
+
+
+def _specific_alltoall_dim(
+    dist_tensor: Tensor, mesh: ProcessMesh, placements: list[Placement]
+):
+    """
+    Get the specific dimension for alltoall communication in nd_mesh reshard.
+    """
+    if not os.getenv("FLAGS_enable_moe_utils") == "true":
+        return None
+
+    mesh_dim = None
+    if paddle.in_dynamic_mode():
+        src_mesh = dist_tensor.process_mesh
+        src_placements = dist_tensor.placements
+    elif paddle.framework.in_pir_mode():
+        src_mesh = dist_tensor.process_mesh
+        src_placements = dist_tensor.dist_attr().placements_attr
+
+    if src_mesh != mesh or src_mesh.ndim == 1:
+        return None
+
+    if any(p.is_partial() for p in src_placements):
+        return None
+    if any(p.is_partial() for p in placements):
+        return None
+
+    for i in range(min(len(src_placements), len(placements))):
+        src_p = src_placements[i]
+        dst_p = placements[i]
+        if src_p.is_shard() and dst_p.is_shard() and src_p != dst_p:
+            # reshard from shard to shard, needs alltoall
+            # now only supports reshard on one dimension
+            src_dim = src_p.get_dim()
+            dst_dim = dst_p.get_dim()
+            if mesh_dim is not None or abs(src_dim - dst_dim) != 1:
+                return None
+            else:
+                mesh_dim = i
+
+    return mesh_dim
+
+
+def _dtensor_from_local(
+    local_tensor, mesh, placements, local_tensor_shape=None
+):
+    # assume the each rank has the same tensor shape for now, just use the local shape to calculate the global shape
+    global_dims = list(local_tensor.shape)
+    if local_tensor_shape is not None:
+        global_dims = local_tensor_shape
+    for idx, placement in enumerate(placements):
+        if placement.is_shard():
+            shard_dim = placement.get_dim()
+            local_dim_size = global_dims[shard_dim]
+            global_dims[shard_dim] = local_dim_size * mesh.shape[idx]
+
+    if paddle.in_dynamic_mode():
+        place = paddle.framework._current_expected_place()
+        place = paddle.framework._get_paddle_place(place)
+
+        return paddle.Tensor(
+            local_tensor,
+            dims=global_dims,
+            process_mesh=mesh,
+            placements=placements,
+            place=place,
+        )
+
+    # TODO Adopt Mix2Dist Pass to allow the program could be executed actually.
+    elif paddle.framework.in_pir_mode():
+        assert isinstance(local_tensor, (type(None), paddle.pir.Value)), (
+            "input tensor is not pir value."
+        )
+        assert local_tensor.is_dense_tensor_type(), (
+            "dtensor_from_local() are only supported dense tensor type right."
+        )
+        sharding_specs = (
+            paddle.distributed.auto_parallel.placement_type.get_shard_spec(
+                mesh, placements, local_tensor.ndim
+            )
+        )
+        dims_mapping = paddle.distributed.auto_parallel.static.utils.convert_to_dims_mapping(
+            sharding_specs, mesh
+        )
+        local_shape = local_tensor.shape
+        global_tensor_type = paddle.pir.create_shaped_type(
+            local_tensor.type(), global_dims
+        )
+        dist_dense_tensor_type = paddle.base.libpaddle.pir.create_dist_dense_tensor_type_by_dense_tensor(
+            global_tensor_type, local_shape, mesh, dims_mapping
+        )
+        local_tensor.set_type(dist_dense_tensor_type)
+        return local_tensor
+    else:
+        raise RuntimeError(
+            "dtensor_from_local() are only supported in dynamic or pir mode."
+        )
+
+
+def _pir_nd_mesh_all2all(src_value, dst_type, mesh, placements, dim):
+    """
+    Use all to all communication in nd_mesh reshard.
+    """
+    # create value on sub 1D mesh
+    sub_value = paddle._C_ops.share_data(src_value)
+    sub_mesh = get_1D_sub_process_mesh(mesh, dim)
+    sub_placements = [src_value.dist_attr().placements_attr[dim]]
+    sub_value_shape = dist.auto_parallel.api._cal_global_shape(
+        src_value._local_shape, sub_mesh, sub_placements
+    )
+    sub_value_type = paddle.pir.create_shaped_type(
+        sub_value.type(), sub_value_shape
+    )
+    sub_dims_mapping, partial_status = to_dim_map(
+        sub_placements, len(sub_value_shape)
+    )
+    sub_value_dist_attr = (
+        paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+            sub_mesh, sub_dims_mapping, partial_status
+        )
+    )
+    sub_value_dist_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+        sub_value_type, sub_value_dist_attr
+    )
+    sub_value.set_type(sub_value_dist_type)
+
+    # 1D mesh reshard
+    dst_placements = [placements[dim]]
+    sub_dst_dims_mapping, partial_status = to_dim_map(
+        dst_placements, len(sub_value_shape)
+    )
+    sub_dst_type = paddle.pir.create_shaped_type(
+        sub_value.type(), sub_value_shape
+    )
+    sub_dst_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+        sub_mesh, sub_dst_dims_mapping, partial_status
+    )
+    sub_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+        sub_dst_type, sub_dst_dist_attr
+    )
+    reshard_func = choose_reshard_func(sub_value_dist_attr, sub_dst_dist_attr)
+    out = reshard_func.reshard(
+        sub_value_dist_attr, sub_dst_dist_attr, sub_value, sub_dst_type
+    )
+
+    # set the type of the output value with global mesh
+    if out is not None:
+        out.set_type(dst_type)
+
+    return out
+
+
+class _NdMeshAlltoAll(PyLayer):
+    @staticmethod
+    def forward(
+        ctx,
+        dist_tensor: Tensor,
+        mesh: ProcessMesh,
+        placements: list[Placement],
+        dim: int,
+    ):
+        sub_mesh = get_1D_sub_process_mesh(mesh, dim)
+        ctx.alltoall_dim = dim
+        ctx.x_mesh = copy.deepcopy(dist_tensor.process_mesh)
+        ctx.x_placements = copy.deepcopy(dist_tensor.placements)
+        ctx.out_mesh = copy.deepcopy(mesh)
+        ctx.out_placements = copy.deepcopy(placements)
+
+        local_shape = _cal_local_shape(
+            dist_tensor.shape, sub_mesh, [dist_tensor.placements[dim]]
+        )
+        out = _dtensor_from_local(
+            dist_tensor._local_value(),
+            sub_mesh,
+            [dist_tensor.placements[dim]],
+            local_shape,
+        )
+        out = dist.reshard(out, sub_mesh, [placements[dim]])
+        local_shape = _cal_local_shape(out.shape, sub_mesh, out.placements)
+        out = _dtensor_from_local(
+            out._local_value(), mesh, placements, local_shape
+        )
+        out.stop_gradient = dist_tensor.stop_gradient
+        return out
+
+    @staticmethod
+    def backward(ctx, out_grad):
+        if not check_placements_equal(ctx.out_placements, out_grad.placements):
+            out = dist.reshard(out_grad, ctx.out_mesh, ctx.out_placements)
+        out = _NdMeshAlltoAll.apply(
+            out_grad, ctx.x_mesh, ctx.x_placements, ctx.alltoall_dim
+        )
+        return out
+
+
+def _cal_local_shape(global_shape, mesh, placements):
+    local_shape = list(global_shape)
+    for idx, placement in enumerate(placements):
+        if placement.is_shard():
+            shard_dim = placement.get_dim()
+            local_shape[shard_dim] = local_shape[shard_dim] // mesh.shape[idx]
+    return local_shape
+
+
+def infer_positive_shape(src_shape, tgt_shape):
+    if isinstance(tgt_shape, (list, tuple)):
+        ret_shape = np.array(tgt_shape)
+    else:
+        ret_shape = tgt_shape.copy()
+
+    minus_one_idx = np.where(ret_shape == -1)[0]
+    if minus_one_idx.size > 0:
+        assert minus_one_idx.size <= 1, (
+            "At most one -1 is allowed in target shape."
+        )
+
+        nelem = np.prod(src_shape)
+        ret_shape[minus_one_idx[0]] = 1
+        ret_shape[minus_one_idx[0]] = nelem // np.prod(ret_shape)
+
+    return list(ret_shape)
+
+
+class _local_reshape(PyLayer):
+    @staticmethod
+    def forward(
+        ctx,
+        dist_tensor: Tensor,
+        global_shape: list,
+        local_shape: list,
+        mesh: ProcessMesh,
+        placements: list[Placement],
+    ):
+        place = paddle.framework._current_expected_place()
+        place = paddle.framework._get_paddle_place(place)
+
+        if dist_tensor._local_value()._is_initialized():
+            local_tensor = dist_tensor._local_value().clone()
+        else:
+            local_tensor = dist_tensor._local_value()
+        ctx.x_global_shape = copy.deepcopy(dist_tensor.shape)
+        ctx.x_local_shape = copy.deepcopy(local_tensor.shape)
+        ctx.x_mesh = copy.deepcopy(dist_tensor.process_mesh)
+        ctx.x_placements = copy.deepcopy(dist_tensor.placements)
+
+        local_tensor = local_tensor.reshape(local_shape)
+        out = paddle.Tensor(
+            local_tensor,
+            dims=global_shape,
+            process_mesh=mesh,
+            placements=placements,
+            place=place,
+        )
+        out.stop_gradient = dist_tensor.stop_gradient
+        return out
+
+    @staticmethod
+    def backward(ctx, out_grad):
+        place = paddle.framework._current_expected_place()
+        place = paddle.framework._get_paddle_place(place)
+
+        if out_grad._local_value()._is_initialized():
+            local_grad = out_grad._local_value().clone()
+            x_local_shape = ctx.x_local_shape
+        else:
+            local_grad = out_grad._local_value()
+            x_local_shape = [0]
+        local_grad = local_grad.reshape(x_local_shape)
+        ret = paddle.Tensor(
+            local_grad,
+            dims=ctx.x_global_shape,
+            process_mesh=ctx.x_mesh,
+            placements=ctx.x_placements,
+            place=place,
+        )
+        return ret
+
+
+def _dist_reshape(
+    dist_tensor: Tensor,
+    global_shape: list,
+    mesh: ProcessMesh,
+    placements: list[Placement],
+):
+    """
+    Reshape the local tensors of the dist tensor on each rank,
+    and manually set the process_mesh and placements of the output.
+    """
+    tgt_global_shape = infer_positive_shape(dist_tensor.shape, global_shape)
+    tgt_local_shape = _cal_local_shape(tgt_global_shape, mesh, placements)
+    if paddle.in_dynamic_mode():
+        src_local_shape = dist_tensor._local_value().shape
+        if not dist_tensor._local_value()._is_initialized():
+            tgt_local_shape = dist_tensor._local_value().shape
+    elif paddle.framework.in_pir_mode():
+        # src_local_shape = dist_tensor._local_shape
+        src_local_shape = _cal_local_shape(
+            dist_tensor.shape,
+            dist_tensor.dist_attr().process_mesh,
+            dist_tensor.dist_attr().placements_attr,
+        )
+    else:
+        raise NotImplementedError(
+            "dist_reshape is only supported in dynamic and pir mode."
+        )
+
+    assert np.prod(tgt_local_shape) == np.prod(src_local_shape), (
+        f"The local shapes {src_local_shape} and {tgt_local_shape} are mismatched."
+    )
+
+    if paddle.in_dynamic_mode():
+        return _local_reshape.apply(
+            dist_tensor, tgt_global_shape, tgt_local_shape, mesh, placements
+        )
+    elif paddle.framework.in_pir_mode():
+        return paddle._C_ops.dist_reshape(
+            dist_tensor,
+            dist_tensor.placements,
+            tgt_global_shape,
+            tgt_local_shape,
+            mesh,
+            placements,
+        )
+
+
+def shard_submesh_and_slice(mesh, tensor_slice, tensor_dim, mesh_dim):
+    new_sub_meshes = split_mesh(mesh, mesh_dim)
+    num_shards = len(new_sub_meshes)
+
+    total_size = tensor_slice[tensor_dim][1] - tensor_slice[tensor_dim][0]
+    shard_size = (total_size + num_shards - 1) // num_shards
+    effective_size = shard_size * (num_shards - 1)
+    last_shard_size = total_size - effective_size
+
+    new_slices = []
+    for i in range(num_shards):
+        start = tensor_slice[tensor_dim][0] + i * shard_size
+        if i == num_shards - 1:
+            end = min(start + last_shard_size, tensor_slice[tensor_dim][1])
+        else:
+            end = min(start + shard_size, tensor_slice[tensor_dim][1])
+        new_slice = list(tensor_slice)
+        new_slice[tensor_dim] = (start, end)
+        new_slices.append(new_slice)
+    return new_sub_meshes, new_slices
+
+
+def get_rank2tensor_indices(sub_mesh_indices_info, sub_mesh_partial_info):
+    rank2tensor_indices = {}
+    for sub_mesh, slice_info in sub_mesh_indices_info.items():
+        for rank in sub_mesh.process_ids:
+            rank2tensor_indices[rank] = {
+                'slice': slice_info,
+                'partial': sub_mesh_partial_info,
+            }
+    return rank2tensor_indices
+
+
+def get_local_slices(tensor, mesh, placements):
+    # TODO(nieyuntao): Temporarily disable this check to bypass certain special cases (shard one tensor dim by many mesh dim)
+    # if len(mesh.shape) < len(placements):
+    #     raise ValueError(
+    #         f"placements length ({len(placements)}) must be smaller or equal to mesh_shape({len(mesh.shape)})"
+    #     )
+    if len(placements) < len(mesh.shape):
+        for _ in range(len(mesh.shape) - len(placements)):
+            placements.append(dist.Replicate())
+
+    sub_mesh_indices_info = {mesh: [(0, s) for s in tensor.shape]}
+    sub_mesh_partial_info = {}
+    for mesh_dim, placement in enumerate(placements):
+        if placement.is_shard():
+            tensor_dim = placement.get_dim()
+            tmp = {}
+            while sub_mesh_indices_info:
+                sub_mesh, slice_info = sub_mesh_indices_info.popitem()
+                new_sub_meshes, new_slices = shard_submesh_and_slice(
+                    sub_mesh, slice_info, tensor_dim, mesh_dim
+                )
+                tmp.update(dict(zip(new_sub_meshes, new_slices)))
+            sub_mesh_indices_info.update(tmp)
+
+        if hasattr(placement, 'is_partial') and placement.is_partial():
+            sub_mesh_partial_info[mesh_dim] = placement.reduce_type()
+
+    return get_rank2tensor_indices(sub_mesh_indices_info, sub_mesh_partial_info)
+
+
+def _only_reshard_mesh_shape(
+    dist_tensor: Tensor, mesh: ProcessMesh, placements: list[Placement]
+):
+    if not os.getenv("FLAGS_enable_moe_utils") == "true":
+        return False
+
+    if paddle.in_dynamic_mode():
+        src_placements = dist_tensor.placements
+        src_mesh = dist_tensor.process_mesh
+    elif paddle.framework.in_pir_mode():
+        src_placements = dist_tensor.dist_attr().placements_attr
+        src_mesh = dist_tensor.dist_attr().process_mesh
+    else:
+        raise NotImplementedError(
+            "_only_reshard_mesh_shape is only supported in dynamic and pir mode."
+        )
+    if src_mesh == mesh or src_mesh.process_ids != mesh.process_ids:
+        return False
+    src_rank2tensor_indices = get_local_slices(
+        dist_tensor, src_mesh, src_placements
+    )
+    dst_rank2tensor_indices = get_local_slices(dist_tensor, mesh, placements)
+    if src_rank2tensor_indices != dst_rank2tensor_indices:
+        return False
+    return True
+
+
+def _reshard_mesh_shape(
+    dist_tensor: Tensor, mesh: ProcessMesh, placements: list[Placement]
+):
+    if not os.getenv("FLAGS_enable_moe_utils") == "true":
+        return False
+
+    src_mesh = dist_tensor.process_mesh
+    if src_mesh == mesh or src_mesh.process_ids != mesh.process_ids:
+        return False
+
+    # only the mesh shapes are different,
+    # if the placements are all replicate,
+    # then we can reshard the mesh shapes
+    if not all(p.is_replicated() for p in dist_tensor.placements + placements):
+        return False
+
+    return True

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

@@ -0,0 +1,101 @@
+# 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 logging
+
+from ...utils.log_utils import get_logger
+
+_logger = get_logger(logging.INFO)
+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 DistributedDefaultImpl0, 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
+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)
+
+        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(
+    "flash_attn", DistributedFlashAttnImpl0("random_control")
+)

python/paddle/distributed/auto_parallel/pipelining/__init__.py

@@ -0,0 +1,15 @@
+#   Copyright (c) 2025 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.
+
+__all__ = []

python/paddle/distributed/auto_parallel/pipelining/_backward.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2025 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 __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+import paddle
+
+from .utils import _map_debug_info
+
+logger = logging.getLogger(__name__)
+
+
+def stage_backward_input(
+    stage_outputs_or_loss: list[paddle.Tensor],
+    output_grads: list[paddle.Tensor] | None,
+    input_values: list[paddle.Tensor],
+    weights: Iterator[paddle.Tensor],
+) -> tuple[tuple[paddle.Tensor | None, ...], list[dict[str, Any]]]:
+    raise NotImplementedError("stage_backward_input is not implemented yet")
+
+
+def stage_backward_weight(
+    weights: Iterator[paddle.Tensor],
+    param_groups: list[dict[str, Any]],
+    retain_graph=False,
+) -> tuple[paddle.Tensor | None, ...]:
+    raise NotImplementedError("stage_backward_weight is not implemented yet")
+
+
+def stage_backward(
+    stage_output,
+    output_grads,
+    input_values,
+) -> tuple[paddle.Tensor | None, ...]:
+    """
+    This is a helper function to:
+    1. compute the gradients for the stage inputs, and
+    2. accumulate gradients for the stage module's parameters.
+
+    Given the input value(s) and the corresponding gradient for the output
+    value(s), compute and accumulate gradients for all parameter values (leaves
+    in the autograd trace) as well as return a list of the gradients for the
+    input values
+
+    """
+
+    try:
+        # stage_output may be a composite datatype like dict. Extract all individual
+        # tensor values here
+        stage_output_tensors: list[paddle.Tensor] = []
+        output_grad_tensors: list[paddle.Tensor | None] = []
+
+        def extract_tensors_with_grads(
+            output_val,
+            grad_val,
+            extract_tensors_with_grads,
+        ):
+            if isinstance(output_val, paddle.Tensor):
+                if output_val.stop_gradient and output_val.grad_fn is None:
+                    return
+                assert isinstance(grad_val, (paddle.Tensor, type(None))), (
+                    f"Expected Tensor or None gradient but got {type(grad_val)}"
+                )
+                stage_output_tensors.append(output_val)
+                output_grad_tensors.append(grad_val)
+            elif isinstance(output_val, (tuple, list)):
+                if grad_val is None:
+                    return
+                assert isinstance(grad_val, (tuple, list)), (
+                    f"grad_value expected to have type {type(output_val)} but got {type(grad_val)}"
+                )
+                assert len(output_val) == len(grad_val)
+                for ov, gv in zip(output_val, grad_val):
+                    extract_tensors_with_grads(
+                        ov,
+                        gv,
+                        extract_tensors_with_grads,
+                    )
+            elif isinstance(output_val, dict):
+                if grad_val is None:
+                    return
+                assert isinstance(grad_val, dict)
+                assert set(output_val.keys()) == set(grad_val.keys())
+                for k in output_val.keys():
+                    extract_tensors_with_grads(
+                        output_val[k], grad_val[k], extract_tensors_with_grads
+                    )
+            else:
+                # Output is a non-tensor type; just ignore it
+                pass
+
+        # Note: ref cycle
+        # break a ref cycle that would keep tensors alive until GC runs
+        # 1. extract_tensors_with_grads refers to a cell that holds refs to any vars defined in stage_backward
+        #    and used in extract_tensors_with_grads
+        # 2. extract_tensors_with_grads referred to both stage_output_tensors, output_grad_tensors,
+        #    and to itself (extract_tensors_with_grads) since it makes a recursive call
+        # 3. stage_output_tensors was kept alive by the above refcycle, and it holds activation tensors, which is bad
+        # fix -> explicitly pass in the ref to the fn, so there is no gc cycle anymore
+        extract_tensors_with_grads(
+            stage_output, output_grads, extract_tensors_with_grads
+        )
+        # Deactivate auto mixed precision context in the backward phase
+        with paddle.amp.auto_cast(enable=False):
+            paddle.autograd.backward(
+                stage_output_tensors,
+                grad_tensors=output_grad_tensors,
+            )
+
+        # Extract gradients wrt the input values
+        grad_inputs: list[paddle.Tensor | None] = []
+        for val in input_values:
+            if isinstance(val, paddle.Tensor):
+                grad_inputs.append(val.grad)
+            else:
+                grad_inputs.append(None)
+
+    except Exception as e:
+        exc_msg = f"""
+        Failed to run stage backward:
+        Stage output: {_map_debug_info(stage_output)}
+        Output gradient: {_map_debug_info(output_grads)}
+        Input: {_map_debug_info(input_values)}
+        """
+        raise RuntimeError(exc_msg) from e
+
+    return tuple(grad_inputs)

python/paddle/distributed/auto_parallel/pipelining/microbatch.py

@@ -0,0 +1,326 @@
+# Copyright (c) 2025 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 __future__ import annotations
+
+import logging
+from typing import Any
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import Replicate, Shard
+from paddle.distributed.auto_parallel.api import (
+    dtensor_from_local,
+    dtensor_to_local,
+)
+from paddle.utils import flatten, map_structure, pack_sequence_as
+
+logger = logging.getLogger(__name__)
+
+# Default chunking dimension is 0. This is used for the case where the user did
+# not specify a chunking dimension.
+DEFAULT_CHUNK_DIM = 0
+
+
+def _split_tensor(x, num_chunks, split_axis=0):
+    if not x.is_dist():
+        chunk_tensors = paddle.tensor_split(x, num_chunks, split_axis)
+    # dp_degree > 1 , placements of model input is [S(0), R, ...]
+    else:
+        if dist.in_auto_parallel_align_mode():
+
+            def _reorder_data_for_align():
+                nonlocal x
+                assert x.placements[0] == dist.Shard(0), (
+                    "inputs should be placed on S(0)."
+                )
+
+                shardings = x.process_mesh.shape[0]
+
+                rows_per_shard = x.shape[0] // shardings
+                new_indices = []
+                for s_id in range(shardings):
+                    for row_in_shard in range(rows_per_shard):
+                        new_indices.append(s_id + row_in_shard * shardings)
+                tmp = x[new_indices]
+                x = dist.reshard(tmp, x.process_mesh, x.placements)
+
+            _reorder_data_for_align()
+        mesh = x.process_mesh
+        placements = x.placements
+        dense_x = dtensor_to_local(x, mesh, placements)
+        chunk_tensors = paddle.tensor_split(dense_x, num_chunks, split_axis)
+        for i in range(num_chunks):
+            chunk_tensors[i] = dtensor_from_local(
+                chunk_tensors[i], mesh, placements
+            )
+    return chunk_tensors
+
+
+def _concat_tensor(chunk_tensors, axis=0):
+    chunk0 = chunk_tensors[0]
+    if not chunk0.is_dist():
+        out = paddle.concat(chunk_tensors, axis)
+
+    else:
+        # loss_fun(out, labels), placements of labels is [S(0), R, ...]
+        mesh = chunk0.process_mesh
+        placements = [Replicate() for _ in range(mesh.ndim)]
+        dp_index = mesh.dim_names.index("dp") if "dp" in mesh.dim_names else 0
+        placements[dp_index] = Shard(0)
+
+        for i in range(len(chunk_tensors)):
+            chunk_tensors[i] = dist.reshard(chunk_tensors[i], mesh, placements)
+
+            chunk_tensors[i] = dtensor_to_local(
+                chunk_tensors[i], mesh, placements
+            )
+        out = paddle.concat(chunk_tensors, axis)
+        out = dtensor_from_local(out, mesh, placements)
+    return out
+
+
+class TensorChunkSpec:
+    """
+    Class used to specify chunking of inputs
+    """
+
+    def __init__(self, split_axis):
+        self.split_axis = split_axis
+
+    split_axis: int
+
+    def __repr__(self):
+        return f"{self.__class__.__module__}.{self.__class__.__name__}({self.split_axis})"
+
+    def __str__(self):
+        return f"TensorChunkSpec({self.split_axis})"
+
+
+def _split_args_helper(
+    args_dict,
+    args_chunk_spec,
+    num_chunks,
+):
+    """
+    A helper function of split_args_kwargs_into_chunks.
+    """
+    assert len(args_dict) == len(args_chunk_spec), (
+        f"args_dict.keys() = {list(args_dict.keys())} args_chunk_spec.keys() = {list(args_chunk_spec.keys())}"
+    )
+
+    shared_args_dict_flat = {}
+    # handle args one by one
+    for arg_key, arg in args_dict.items():
+        arg_flat = flatten(arg)
+
+        chunk_spec = args_chunk_spec[arg_key]
+        assert chunk_spec is not None
+
+        chunk_spec_flat = flatten(chunk_spec)
+        assert len(chunk_spec_flat) == len(arg_flat), (
+            f"{arg_key} {len(arg_flat)} != {len(chunk_spec_flat)}"
+        )
+
+        shard_arg_flat = []
+
+        for v, chunk_v in zip(arg_flat, chunk_spec_flat):
+            if not isinstance(v, paddle.Tensor):
+                shard_arg_flat.append([v] * num_chunks)
+            elif isinstance(chunk_v, TensorChunkSpec):
+                v_split_axis_size = v.shape[chunk_v.split_axis]
+
+                if v_split_axis_size < num_chunks:
+                    raise ValueError(
+                        f"Arg {arg_key} on chunking dimension has a size of {v_split_axis_size}, "
+                        f"smaller than the number of chunks {num_chunks}. "
+                        "Please adjust your num_chunks setting."
+                    )
+                # split tensor v
+                chunk_tensors = _split_tensor(v, num_chunks, chunk_v.split_axis)
+
+                shard_arg_flat.append(chunk_tensors)
+            else:
+                raise TypeError(f"Unrecognized chunk spec: {chunk_v}")
+
+        shared_args_dict_flat[arg_key] = shard_arg_flat
+
+    # the structure of each element in args_split is the same as the original args_dict
+    args_split = []
+    for idx in range(num_chunks):
+        chunk_args = {}
+        for key, arg in shared_args_dict_flat.items():
+            last_arg = None if not arg else arg[0][idx]
+            arg_of_curr_chunk = (
+                [v[idx] for v in arg] if len(arg) > 1 else last_arg
+            )
+            chunk_args[key] = arg_of_curr_chunk
+
+        # flatten chunk_args first, and then pack chunk_args as the origin args_dict
+        flatten_chunk_args = [x for x in flatten(chunk_args) if x is not None]
+        chunk_args = pack_sequence_as(args_dict, flatten_chunk_args)
+        args_split.append(chunk_args)
+    return args_split
+
+
+def split_args_kwargs_into_chunks(
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+    chunks: int,
+    args_chunk_spec: (
+        tuple[
+            tuple[TensorChunkSpec, ...]
+            | list[TensorChunkSpec, ...]
+            | TensorChunkSpec,
+            ...,
+        ]
+        | None
+    ) = None,
+    kwargs_chunk_spec: (
+        dict[
+            str,
+            tuple[TensorChunkSpec, ...]
+            | list[TensorChunkSpec, ...]
+            | TensorChunkSpec,
+        ]
+        | None
+    ) = None,
+) -> tuple[list[tuple], list[dict]]:
+    """
+    Given a sequence of args and kwargs, split them into a number of chunks
+    according to  their respective chunking specs.
+
+    Args:
+        args: tuple of args
+        kwargs: dict of kwargs
+        chunks: Number of chunks to split the args and kwargs into
+        args_chunk_spec: chunking specs for args, in same shape as args
+        kwargs_chunk_spec: chunking specs for kwargs, in same shape as kwargs
+
+    Returns:
+        args_split: list of sharded args
+        kwargs_split: list of sharded kwargs
+    """
+
+    if kwargs is None:
+        kwargs = {}
+
+    if args_chunk_spec is None:
+        args_chunk_spec = map_structure(
+            lambda _: TensorChunkSpec(DEFAULT_CHUNK_DIM), args
+        )
+
+    if kwargs_chunk_spec is None:
+        kwargs_chunk_spec = map_structure(
+            lambda _: TensorChunkSpec(DEFAULT_CHUNK_DIM), kwargs
+        )
+
+    args_split_dict = _split_args_helper(
+        dict(enumerate(args)),
+        dict(enumerate(args_chunk_spec)),
+        chunks,
+    )
+    kwargs_split = _split_args_helper(
+        kwargs,
+        kwargs_chunk_spec,
+        chunks,
+    )
+
+    assert len(args_split_dict) == len(kwargs_split), (
+        "args and kwargs are split into difference number of chunks: "
+        f"{len(args_split_dict)}, {len(kwargs_split)}"
+    )
+
+    # the form of each args_chunk should be tuple
+    args_split = [
+        tuple(args_chunk[i] for i in range(len(args_chunk)))
+        for args_chunk in args_split_dict
+    ]
+
+    return args_split, kwargs_split
+
+
+def merge_chunks(
+    chunks: list[Any],
+    chunk_spec,
+):
+    """
+    Given a list of chunks, merge them into a single chunk according to
+    the chunk spec.
+
+    Args:
+        chunks: list of chunks
+        chunk_spec: Chunking spec for the chunks
+
+    Returns:
+        chunk: chunks merged value
+    """
+    if len(chunks) == 0:
+        logger.warning("No chunks to merge.")
+        return chunks
+
+    if chunk_spec is None:
+        chunk_spec = map_structure(
+            lambda _: TensorChunkSpec(DEFAULT_CHUNK_DIM), chunks[0]
+        )
+
+    chunks_flat = []
+    # flatten chunk_spec first
+    chunk_spec = flatten(chunk_spec)
+    for chunk in chunks:
+        chunk_flat = flatten(chunk)
+        assert len(chunk_flat) == len(chunk_spec), (
+            f"Chunk {chunk} did not match chunk spec {chunk_spec}"
+        )
+        chunks_flat.append(chunk_flat)
+
+    def _merge_non_tensor_type_arg(chunks, idx, chunk_spec_of_arg=None):
+        # use the first chunk's value as the merged result
+        arg_0 = chunks[0][idx]
+        for chunk_idx in range(1, len(chunks)):
+            assert chunks[chunk_idx][idx] == arg_0, (
+                f"Cannot merge chunks with index 0 and {idx} with different values,"
+                f"When the arg's TensorChunkSpec is {chunk_spec_of_arg}"
+            )
+        return arg_0
+
+    args_flat = []
+    for arg_idx, chunk_spec_of_arg in enumerate(chunk_spec):
+        if isinstance(chunk_spec_of_arg, TensorChunkSpec):
+            if isinstance(chunks_flat[0][arg_idx], paddle.Tensor):
+                arg_chunks_to_merge = [
+                    chunks_flat[chunk_idx][arg_idx]
+                    for chunk_idx in range(len(chunks_flat))
+                ]
+                merged_arg = _concat_tensor(
+                    arg_chunks_to_merge, axis=chunk_spec_of_arg.split_axis
+                )
+            else:
+                logger.warning(
+                    f"Cannot merge chunks with TensorChunkSpec {chunk_spec_of_arg}."
+                    "The TensorChunkSpec only supports paddle.Tensor type."
+                )
+
+                merged_arg = _merge_non_tensor_type_arg(
+                    chunks_flat, arg_idx, chunk_spec_of_arg
+                )
+        else:
+            merged_arg = _merge_non_tensor_type_arg(
+                chunks_flat, arg_idx, chunk_spec_of_arg
+            )
+
+        args_flat.append(merged_arg)
+
+    # pack args_flat as the input chunks[0]
+    return pack_sequence_as(chunks[0], args_flat)

python/paddle/distributed/auto_parallel/pipelining/utils.py

@@ -0,0 +1,177 @@
+# Copyright (c) 2025 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 __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+import paddle
+from paddle.distributed import fleet
+from paddle.distributed.auto_parallel.api import (
+    dtensor_from_local,
+)
+from paddle.utils import map_structure
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+logger = logging.getLogger(__name__)
+
+
+def _detach_and_requires_grad(x):
+    o = x.detach()
+    o.stop_gradient = False
+    return o
+
+
+def _detach_and_keep_grad(x):
+    o = x.detach_()
+    o.stop_gradient = x.stop_gradient
+    return o
+
+
+def _zero_initialize_with_meta(meta, mesh):
+    assert isinstance(meta, TensorMeta)
+    x = paddle.zeros(
+        meta._local_shape if meta._local_shape else meta.shape, dtype=meta.dtype
+    )
+    if meta.placements:
+        x = dtensor_from_local(x, mesh, meta.placements)
+    return x
+
+
+def _flatten_args(args):
+    """
+    Flatten the args into a list form.
+    """
+    flat_args = []
+
+    def extract_tensor_args(a):
+        nonlocal flat_args
+        if isinstance(a, paddle.Tensor):
+            flat_args.append(a)
+        return a
+
+    paddle.utils.map_structure(
+        extract_tensor_args,
+        args,
+    )
+
+    return flat_args
+
+
+class PipeliningShapeError(RuntimeError):
+    """Shape mismatch between configured and runtime values."""
+
+
+def _validate_tensor_metadata(desc, expected, given):
+    if not expected.shape == given.shape:
+        raise PipeliningShapeError(
+            f"{desc} has a shape mismatch: expected {expected.shape} actual {given.shape}"
+        )
+    if not expected.dtype == given.dtype:
+        raise PipeliningShapeError(
+            f"{desc} has a dtype mismatch: expected {expected.dtype} actual {given.dtype}"
+        )
+
+
+def _validate_tensors_metadata(
+    desc,
+    expected_tensors: list[paddle.Tensor] | tuple[paddle.Tensor, ...],
+    actual_tensors: list[paddle.Tensor] | tuple[paddle.Tensor, ...],
+):
+    if len(expected_tensors) != len(actual_tensors):
+        raise PipeliningShapeError(
+            f"{desc}: Number of values ({len(actual_tensors)}) does not match expected number ({len(expected_tensors)})"
+        )
+    for i in range(len(expected_tensors)):
+        _validate_tensor_metadata(
+            f"{desc}: value {i}", expected_tensors[i], actual_tensors[i]
+        )
+
+
+NestedStruct = list[Any] | tuple[Any, ...] | dict[Any, Any]
+
+
+def _map_structure_only(
+    type_: Any, fn: Callable[[Any], Any], structure: NestedStruct
+) -> NestedStruct:
+    """
+    Apply `fn` to each entry which matches `type_` in `structure` and return a new structure with the same shape.
+    """
+    return map_structure(
+        lambda x: fn(x) if isinstance(x, type_) else x, structure
+    )
+
+
+class TensorMeta:
+    def __init__(self, tensor: paddle.Tensor):
+        if tensor.is_dist():
+            self.shape = tensor.shape
+            self._local_shape = tensor._local_shape
+        else:
+            self.shape = tensor.shape
+            self._local_shape = None
+        self.dtype = tensor.dtype
+        self.placements = None if not tensor.is_dist() else tensor.placements
+        self.stop_gradient = tensor.stop_gradient
+
+    def __repr__(self):
+        return f"TensorMeta(global_shape={self.shape},local_shape={self._local_shape}, dtype={self.dtype}, placements={self.placements})"
+
+
+def _get_pp_mesh(pp_idx=0, pp_dim_names="pp"):
+    """
+    Get the mesh of the {pp_idx}th PipelineStage.
+    """
+    mesh = fleet.auto.get_mesh()
+    assert mesh is not None, (
+        "the mesh is None, please call fleet.auto.set_mesh first."
+    )
+    if "pp" in mesh.dim_names:
+        mesh = mesh.get_mesh_with_dim("pp", pp_idx)
+    else:
+        logger.warning(
+            f"The dim name of pp {pp_dim_names} not exist in global mesh {mesh}"
+        )
+    return mesh
+
+
+def _get_stage_mesh(stage_index, pp_group_size, style=None):
+    if style == "v":
+        raise NotImplementedError
+    if style is not None:
+        raise ValueError(f"Unknown style: {style}, style can be None, v.")
+    else:
+        pp_idx = stage_index % pp_group_size
+        return _get_pp_mesh(pp_idx)
+
+
+def _friendly_debug_info(v):
+    """
+    Helper function to print out debug info in a friendly way.
+    """
+    if isinstance(v, paddle.Tensor):
+        return f"Tensor({v.shape}, stop_gradient={v.stop_gradient}, dtype={v.dtype})"
+    else:
+        return str(v)
+
+
+def _map_debug_info(a):
+    """
+    Helper function to apply `friendly_debug_info` to items in `a`.
+    `a` may be a list, tuple, or dict.
+    """
+    return map_structure(_friendly_debug_info, a)

python/paddle/distributed/auto_parallel/placement_type.py

@@ -0,0 +1,214 @@
+#   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 typing import cast
+
+import paddle
+from paddle.base.core import Partial, Replicate, Shard
+
+
+def dims_mapping_to_placements(dim_map, mesh, partial_idx=[], split_factor={}):
+    """
+    convert dim_map to placements.
+
+    Args:
+        dim_map(List[int]): a list of integer that represents sharding on each tensor dimension.
+        mesh(paddle.distributed.ProcessMesh): The `ProcessMesh` object describes the Cartesian topology of the used processes.
+        partial_idx(List[int], Optional): a list of integer that represents the DTensor have pending sum on which device mesh dimension
+
+    Returns:
+        List[Placement]: a list contains some `paddle.distributed.Placement`.
+    """
+    placements = [Replicate() for _ in mesh.shape]
+
+    for s in partial_idx:
+        placements[s] = Partial()
+
+    for tensor_dim, mesh_dims in enumerate(dim_map):
+        if len(mesh_dims) <= 0:
+            continue
+        is_co_shard = len(mesh_dims) > 1
+        for shard_order, mesh_dim in enumerate(mesh_dims):
+            p = placements[mesh_dim]
+            if p.is_shard():
+                raise Exception(
+                    f"ProcessMesh dimension can not be mapped to two dimension of same tensor: {tensor_dim} and {p.get_dim()}."
+                )
+            elif p.is_partial():
+                raise Exception(
+                    f"ProcessMesh dimension {mesh_dim} can not be both shard and partial!"
+                )
+
+            shard = (
+                Shard(tensor_dim, co_shard_order=shard_order)
+                if is_co_shard
+                else Shard(tensor_dim)
+            )
+            placements[mesh_dim] = shard
+
+    if len(split_factor) > 1:
+        raise RuntimeError("At now only support to rearrange at one mesh dim.")
+
+    for k, v in split_factor.items():
+        placements[k].set_split_factor(v)
+
+    return placements
+
+
+def to_placements(dim_map, mesh, partial_idx=[]):
+    """
+    convert dim_map to placements.
+    Args:
+        dim_map(List[int]): a list of integer that represents sharding on each tensor dimension.
+        mesh(paddle.distributed.ProcessMesh): The `ProcessMesh` object describes the Cartesian topology of the used processes.
+        partial_idx(List[int], Optional): a list of integer that represents the DTensor have pending sum on which device mesh dimension
+    Returns:
+        List[Placement]: a list contains some `paddle.distributed.Placement`.
+    """
+    if isinstance(mesh, paddle.base.libpaddle.ProcessMesh):
+        shape = mesh.shape
+    else:
+        shape = mesh.mesh.shape
+    placements = [Replicate() for _ in range(len(shape))]
+
+    for s in partial_idx:
+        placements[s] = Partial()
+
+    for i, m in enumerate(dim_map):
+        if m >= 0:
+            p = placements[m]
+            if p.is_shard():
+                p = cast("Shard", p)
+                raise Exception(
+                    f"ProcessMesh dimension can not be mapped to two dimension of same tensor: {i} and {p.get_dim()}."
+                )
+            elif p.is_partial():
+                raise Exception(
+                    f"ProcessMesh dimension {m} can not be both shard and partial!"
+                )
+            placements[m] = Shard(i)
+
+    return placements
+
+
+def check_placements_equal(this, that):
+    assert isinstance(this, list) and isinstance(that, list)
+    small_placements = this if len(this) < len(that) else that
+    large_placements = that if len(this) < len(that) else this
+    for i in range(len(large_placements)):
+        if i < len(small_placements):
+            if small_placements[i] != large_placements[i]:
+                return False
+        else:
+            if large_placements[i] != Replicate():
+                return False
+    return True
+
+
+def placemetns_to_dist_status(
+    placements, tensor_dims, return_split_factor=False
+):
+    """
+    convert placements to dim_map.
+
+    Args:
+        placements(List[Placement]): a list contains some `paddle.distributed.Placement`.
+        tensor_dims(int): the dimension of dist_tensor.
+
+    Returns:
+        List[int]: a list of integer that represents sharding on each tensor dimension.
+    """
+    output_list = []
+    dim_map = [[] for _ in range(tensor_dims)]
+    partial_status = {}
+    split_factor_map = {}
+    for i, placement in enumerate(placements):
+        if placement.is_shard():
+            shard_dim = cast("Shard", placement).get_dim()
+            dim_map[shard_dim].append(i)
+            if cast("Shard", placement).get_split_factor() > 1:
+                split_factor_map[i] = cast(
+                    "Shard", placement
+                ).get_split_factor()
+                assert len(split_factor_map) == 1, (
+                    "only support to rerrange at one mesh dim."
+                )
+        if placement.is_partial():
+            partial_status[i] = cast("Partial", placement).reduce_type()
+
+    for shard_dim in dim_map:
+        if len(shard_dim) > 0:
+            shard_dim.sort(key=lambda idx: placements[idx].get_co_shard_order())
+
+    output_list.append(dim_map)
+    output_list.append(partial_status)
+
+    if return_split_factor:
+        output_list.append(split_factor_map)
+
+    return output_list
+
+
+def to_dim_map(placements, tensor_dims):
+    """
+    convert placements to dim_map.
+
+    Args:
+        placements(List[Placement]): a list contains some `paddle.distributed.Placement`.
+        tensor_dims(int): the dimension of dist_tensor.
+
+    Returns:
+        List[int]: a list of integer that represents sharding on each tensor dimension.
+    """
+    dim_map = [-1] * tensor_dims
+    partial_status = {}
+    for i, placement in enumerate(placements):
+        if placement.is_shard():
+            shard_dim = cast("Shard", placement).get_dim()
+            if dim_map[shard_dim] > -1:
+                import logging
+
+                logging.warning(
+                    f"Tensor dim {shard_dim} is already sharded on mesh dim {dim_map[shard_dim]}."
+                )
+
+            dim_map[shard_dim] = i
+        if placement.is_partial():
+            partial_status[i] = cast("Partial", placement).reduce_type()
+
+    return dim_map, partial_status
+
+
+# TODO(lfw): delete it in future.
+def get_shard_spec(mesh, placements, tensor_dims):
+    """to get shard_spec for construct DistAttr for static API."""
+    dim_map = [-1] * tensor_dims
+    for i, placement in enumerate(placements):
+        if placement.is_shard():
+            shard_dim = cast("Shard", placement).get_dim()
+            if dim_map[shard_dim] > -1:
+                import logging
+
+                logging.warning(
+                    f"Tensor dim {shard_dim} is already sharded on mesh dim {dim_map[shard_dim]}."
+                )
+
+            dim_map[shard_dim] = i
+
+    mesh_dim_names = mesh.dim_names
+    shard_spec = [None] * len(dim_map)
+    for i, d in enumerate(dim_map):
+        if d > -1:
+            shard_spec[i] = mesh_dim_names[d]
+
+    return shard_spec

python/paddle/distributed/auto_parallel/process_mesh.py

@@ -0,0 +1,603 @@
+#   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 __future__ import annotations
+
+import copy
+import logging
+from typing import TYPE_CHECKING, Any, SupportsIndex
+
+import numpy as np
+
+import paddle
+from paddle.distributed import fleet
+from paddle.distributed.collective import _get_group_map
+from paddle.distributed.communication.group import is_initialized
+from paddle.framework import core
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Sequence
+    from types import TracebackType
+
+    import numpy.typing as npt
+
+    from paddle._typing import NestedNumericSequence
+
+    _NumpyShapeLike = SupportsIndex | Sequence[SupportsIndex]
+
+
+# Use to store the previous and current process mesh
+_g_previous_process_mesh = None
+_g_current_process_mesh = None
+# {shape_process_ids : unique_id}
+_g_unique_process_mesh_map = {}
+_g_group_map = {}
+
+
+def get_current_process_mesh():
+    global _g_current_process_mesh
+    return _g_current_process_mesh
+
+
+def set_current_process_mesh(process_mesh):
+    global _g_previous_process_mesh
+    global _g_current_process_mesh
+    _g_previous_process_mesh = _g_current_process_mesh
+    _g_current_process_mesh = process_mesh
+
+
+def reset_current_process_mesh():
+    global _g_previous_process_mesh
+    global _g_current_process_mesh
+    _g_current_process_mesh = _g_previous_process_mesh
+
+
+def get_unique_id_for_process_mesh(shape, process_ids):
+    key = f"shape {shape}, process_ids {process_ids}"
+    global _g_unique_process_mesh_map
+    if key in _g_unique_process_mesh_map:
+        unique_id = _g_unique_process_mesh_map[key]
+    else:
+        unique_id = len(_g_unique_process_mesh_map) + 1
+        _g_unique_process_mesh_map[key] = unique_id
+
+    return unique_id
+
+
+def retrieve_unique_id_for_process_mesh(shape, process_ids):
+    key = f"shape {shape}, process_ids {process_ids}"
+    global _g_unique_process_mesh_map
+    assert key in _g_unique_process_mesh_map
+    return _g_unique_process_mesh_map[key]
+
+
+def get_unique_process_mesh_map():
+    global _g_unique_process_mesh_map
+    return _g_unique_process_mesh_map
+
+
+def init_group_by_process_mesh(dim_names):
+    global _g_group_map
+    if dim_names is None:
+        dim_names = []
+    assert isinstance(dim_names, list), "dim_names must be a list."
+    for dim_name in dim_names:
+        if dim_name in _g_group_map:
+            continue
+        _g_group_map[dim_name] = {}
+
+
+def get_group_map_by_dim_name(dim_name):
+    global _g_group_map
+    if dim_name not in _g_group_map:
+        raise RuntimeError(f'No group found for dim_name {dim_name}')
+    return _g_group_map[dim_name]
+
+
+class ProcessMesh(core.ProcessMesh):
+    """
+    The `ProcessMesh` object describes the Cartesian topology of the used processes.
+
+    Args:
+        mesh (list|numpy.array): an n-dimensional array describes the topology
+            of the processes.
+        dim_names (list, optional): the i-th element of this list gives the name of the
+            i-th dimension of the mesh.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]], dim_names=["x", "y"])
+            >>> assert mesh.shape == [2, 3]
+            >>> assert mesh.process_ids == [2, 4, 5, 0, 1, 3]
+
+    """
+
+    shape: list[int]
+    process_ids: list[int]
+
+    def __init__(
+        self,
+        mesh: npt.NDArray[Any] | NestedNumericSequence | None = None,
+        dim_names: list[str] | None = None,
+        shape: _NumpyShapeLike | None = None,
+        process_ids: Iterable[Any] | None = None,
+    ) -> None:
+        paddle.base.framework.global_var._in_auto_parallel_ = True
+
+        # Use shape and process_ids just for compatibility
+        # Users should not use these directly
+        if mesh is None:
+            assert shape is not None
+            assert process_ids is not None
+            mesh = np.array(process_ids).reshape(shape)
+
+        if not isinstance(mesh, list) and not isinstance(mesh, np.ndarray):
+            raise ValueError(
+                'The mesh must be an instance of list or np.ndarray.'
+            )
+        if isinstance(mesh, list):
+            mesh = np.array(mesh)
+
+        if dim_names is not None and not isinstance(dim_names, list):
+            raise ValueError('The dim_names must be an instance of list.')
+
+        self._mesh = mesh
+        self._shape = list(self._mesh.shape)
+        self._process_ids = self._mesh.flatten().tolist()
+
+        assert all(isinstance(p, int) for p in self._process_ids), (
+            "All elements of the mesh must be integer"
+        )
+        assert min(self._process_ids) >= 0, (
+            'All elements of the mesh must be >= 0.'
+        )
+        unique_process_ids = set(self._process_ids)
+        assert len(unique_process_ids) == len(self._process_ids), (
+            'All elements of the mesh must be unique.'
+        )
+
+        if dim_names is not None:
+            assert len(dim_names) == len(self._shape), (
+                "The length of dims_names must be same as the shape of the mesh."
+            )
+            self._dim_names = copy.deepcopy(dim_names)
+        else:
+            self._dim_names = ["d" + str(i) for i in range(len(self._shape))]
+        unique_dim_names = set(self._dim_names)
+        assert len(unique_dim_names) == len(self._dim_names), (
+            f'All dim_names {dim_names} must be unique.'
+        )
+
+        # Follow the requirement for using pybind11
+        core.ProcessMesh.__init__(
+            self, self._shape, self._process_ids, self._dim_names
+        )
+
+        # Store all process meshes
+        from .static.dist_context import get_default_distributed_context
+
+        default_dist_cxt = get_default_distributed_context()
+        default_dist_cxt.add_process_mesh(self)
+        # Add new processes to process group 0
+        from .static.process_group import get_process_group
+
+        pg0 = get_process_group(0)
+        pg0.add_ranks(self.process_ids)
+
+        # Unique Mesh Id
+        self._unique_id = get_unique_id_for_process_mesh(
+            self._shape, self._process_ids
+        )
+        init_group_by_process_mesh(self._dim_names)
+
+    @property
+    def mesh(self) -> npt.NDArray[Any]:
+        """
+        Get the underlying mesh of ProcessMesh.
+        """
+        return self._mesh
+
+    @property
+    def dim_names(self) -> list[str]:
+        """
+        Get the underlying dimension names of ProcessMesh.
+        """
+        return self._dim_names
+
+    @property
+    def unique_id(self) -> int:
+        """
+        Get the unique id of ProcessMesh.
+        NOTE
+        Unique id only take process_ids and shape into account.
+        Different ProcessMesh with same process_ids and shape have same unique id.
+        """
+        return self._unique_id
+
+    def __getitem__(
+        self, index: slice | tuple[slice, ...] | str | SupportsIndex
+    ) -> ProcessMesh:
+        if isinstance(index, tuple):
+            new_dim_names = []
+            for i, item in enumerate(index):
+                if isinstance(item, slice):
+                    new_dim_names.append(self._dim_names[i])
+            new_mesh = self._mesh[index]
+            if new_mesh.shape:
+                return ProcessMesh(new_mesh, new_dim_names)
+            else:
+                # Wrap a scalar into a list but without dim_names
+                return ProcessMesh([new_mesh])
+        elif isinstance(index, slice):
+            new_mesh = self._mesh[index]
+            new_dim_names = self._dim_names
+            return ProcessMesh(new_mesh, new_dim_names)
+        elif isinstance(index, str):
+            return self.get_submesh_with_dim(index)
+        else:
+            new_mesh = self._mesh[index]
+            new_dim_names = self._dim_names[1:]
+            if new_mesh.shape:
+                return ProcessMesh(new_mesh, new_dim_names)
+            else:
+                return ProcessMesh([new_mesh])
+
+    def get_rank_by_dim_and_process_id(
+        self, dim: str | int, process_id: int
+    ) -> int:
+        # do some check
+        if process_id not in self._process_ids:
+            # -1 means invalid rank
+            return -1
+
+        if dim is None:
+            # if dim is None, all process's rank is 0
+            return 0
+
+        if isinstance(dim, int):
+            dim_name = self._dim_names[dim]
+        elif isinstance(dim, str):
+            dim_name = dim
+        else:
+            raise ValueError("dim must be a string or an integer.")
+        dim_name_index = self._dim_names.index(dim_name)
+        rank_index = np.where(self._mesh == process_id)[dim_name_index]
+        return int(rank_index.item())
+
+    def get_dim_size(self, dim: str | int) -> int:
+        if dim is None:
+            return 1
+
+        if isinstance(dim, int):
+            dim_name = self._dim_names[dim]
+        elif isinstance(dim, str):
+            dim_name = dim
+        else:
+            raise ValueError("dim must be a string or an integer.")
+        assert dim_name in self._dim_names
+        return self._shape[self._dim_names.index(dim_name)]
+
+    def get_mesh_with_dim(
+        self,
+        dim_name: str,
+        index: slice | tuple[slice, ...] | SupportsIndex | None = None,
+    ) -> ProcessMesh:
+        assert dim_name in self._dim_names, (
+            f'{dim_name} is not a valid dim name.'
+        )
+        index_axis = self._dim_names.index(dim_name)
+        new_order = [index_axis] + [
+            i for i in range(len(self._dim_names)) if i != index_axis
+        ]
+        new_dim_names = [dim_name] + [
+            dim for dim in self._dim_names if dim != dim_name
+        ]
+        new_mesh = self._mesh.transpose(new_order)
+
+        if index is not None:
+            if len(new_dim_names[1:]) > 0:
+                return ProcessMesh(new_mesh[index], new_dim_names[1:])
+            # satisfy the single dimension mesh case
+            else:
+                return ProcessMesh([new_mesh[index]], new_dim_names)
+        return ProcessMesh(new_mesh, new_dim_names)
+
+    def get_submesh_with_dim(
+        self,
+        dim_name: str,
+    ) -> ProcessMesh:
+        """
+        Slice the current ProcessMesh based on the dim_name given to create a submesh with single dimension remained.
+
+        Args:
+            dim_name (str): the name of the mesh dimension of the ProcessMesh to create the submesh for.
+        Returns:
+            A :class:`ProcessMesh` object
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> import paddle.distributed as dist
+
+                >>> dist.init_parallel_env()
+                >>> mesh_2d = dist.ProcessMesh([[0, 1, 2, 3], [4, 5, 6, 7]], dim_names=["dp", "tp"])
+
+                >>> dp_mesh = mesh_2d.get_submesh_with_dim("dp")
+                >>> # ProcessMesh:([0, 4]) on rank 0, 4
+                >>> # ProcessMesh:([1, 5]) on rank 1, 5
+                >>> # ProcessMesh:([2, 6]) on rank 2, 6
+                >>> # ProcessMesh:([3, 7]) on rank 3, 7
+
+                >>> tp_mesh = mesh_2d.get_submesh_with_dim("tp")
+                >>> # ProcessMesh:([0, 1, 2, 3]) on rank 0, 1, 2, 3
+                >>> # ProcessMesh:([4, 5, 6, 7]) on rank 4, 5, 6, 7
+
+                >>> mesh_3d = dist.ProcessMesh([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=["pp", "dp", "tp"])
+
+                >>> pp_mesh = mesh_3d.get_submesh_with_dim("pp")
+                >>> # ProcessMesh:([0, 4]) on rank 0, 4
+                >>> # ProcessMesh:([1, 5]) on rank 1, 5
+                >>> # ProcessMesh:([2, 6]) on rank 2, 6
+                >>> # ProcessMesh:([3, 7]) on rank 3, 7
+
+                >>> dp_mesh = mesh_3d.get_submesh_with_dim("dp")
+                >>> # ProcessMesh:([0, 2]) on rank 0, 2
+                >>> # ProcessMesh:([1, 3]) on rank 1, 3
+                >>> # ProcessMesh:([4, 6]) on rank 4, 6
+                >>> # ProcessMesh:([5, 7]) on rank 5, 7
+
+                >>> tp_mesh = mesh_3d.get_submesh_with_dim("tp")
+                >>> # ProcessMesh:([0, 1]) on rank 0, 1
+                >>> # ProcessMesh:([2, 3]) on rank 2, 3
+                >>> # ProcessMesh:([4, 5]) on rank 4, 5
+                >>> # ProcessMesh:([6, 7]) on rank 6, 7
+        """
+
+        reorder_mesh = self.get_mesh_with_dim(dim_name)._mesh.reshape(
+            self.get_dim_size(dim_name), -1
+        )
+        curr_rank = paddle.distributed.get_rank()
+        if curr_rank not in self._process_ids:
+            logger.warning(
+                f"Rank {curr_rank} is not in the process mesh, just return None"
+            )
+            return None
+        # find curr_rank in reorder_mesh, get the column index
+        col_idx = np.argmax(reorder_mesh == curr_rank) % reorder_mesh.shape[-1]
+        sub_mesh = ProcessMesh(reorder_mesh[:, col_idx], [dim_name])
+        return sub_mesh
+
+    def _get_group(
+        self,
+        dim_name: str | None = None,
+    ) -> paddle.distributed.communication.group.Group:
+        """ """
+        assert is_initialized(), (
+            "When you want to get a group from the ProcessMesh."
+            " Call paddle.distributed.init_parallel_env first "
+            "to initialize the distributed environment."
+        )
+        if len(self._dim_names) > 1 and dim_name is None:
+            raise ValueError(
+                "You should specify the dim_name when the ProcessMesh has more than one dimensions."
+            )
+        reorder_mesh = self.get_mesh_with_dim(dim_name)._mesh.reshape(
+            self.get_dim_size(dim_name), -1
+        )
+        curr_rank = paddle.distributed.get_rank()
+        groups = get_group_map_by_dim_name(dim_name)
+
+        for rank in self._process_ids:
+            col_idx = np.argmax(reorder_mesh == rank) % reorder_mesh.shape[-1]
+            if col_idx in groups:
+                continue
+            pg = paddle.distributed.new_group(reorder_mesh[:, col_idx])
+            groups[col_idx] = pg
+
+        cur_col_idx = (
+            np.argmax(reorder_mesh == curr_rank) % reorder_mesh.shape[-1]
+        )
+        return groups[cur_col_idx]
+
+    def get_group(
+        self,
+        dim_name: str | None = None,
+    ) -> paddle.distributed.communication.group.Group:
+        """
+        Convert single dimension ProcessMesh to the corresponding Group.
+
+        Args:
+            dim_name (str, optional): it can be the name of the mesh dimension. Default is None.
+
+        Returns:
+            A :class:`Group` object.
+        """
+
+        # check parallel environment whether ready or not
+        assert is_initialized(), (
+            "When you want to get a group from the ProcessMesh."
+            " Call paddle.distributed.init_parallel_env first "
+            "to initialize the distributed environment."
+        )
+        if len(self._dim_names) > 1 and dim_name is None:
+            raise ValueError(
+                "You should specify the dim_name when the ProcessMesh has more than one dimensions."
+            )
+        if len(self._dim_names) == 1:
+            if dim_name is not None and dim_name not in self._dim_names:
+                raise ValueError(
+                    f"{dim_name} not in the dimension names {self._dim_names}"
+                )
+            else:
+                if hasattr(fleet.fleet, "_hcg"):
+                    hcg = fleet.get_hybrid_communicate_group()
+                    if hcg is not None:
+                        parallel_group_map = {
+                            "pp": hcg.get_pipe_parallel_group,
+                            "dp": hcg.get_data_parallel_group,
+                            "mp": hcg.get_model_parallel_group,
+                            "sep": hcg.get_sep_parallel_group,
+                            "sharding": hcg.get_sharding_parallel_group,
+                        }
+
+                        if dim_name not in parallel_group_map:
+                            raise ValueError(
+                                f"{dim_name} is not a valid dim name."
+                            )
+
+                        return parallel_group_map[dim_name]()
+                group_map = _get_group_map()
+                for group in group_map.values():
+                    if set(group.ranks) == set(self._process_ids):
+                        return group
+                return paddle.distributed.new_group(self._process_ids)
+        else:
+            if dim_name not in self._dim_names:
+                raise ValueError(
+                    f"{dim_name} not in the dimension names {self._dim_names}"
+                )
+            sub_mesh = self.get_submesh_with_dim(dim_name)
+            return sub_mesh.get_group(dim_name)
+
+    def __enter__(self) -> None:
+        set_current_process_mesh(self)
+        default_prog = paddle.static.default_main_program()
+        cur_block = default_prog.current_block()
+        self._old_var_names = list(cur_block.vars.keys())
+        self._old_op_size = len(cur_block.ops)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        from .static.dist_op import DistributedOperator
+        from .static.dist_tensor import DistributedTensor
+
+        default_prog = paddle.static.default_main_program()
+        cur_block = default_prog.current_block()
+        new_var_names = list(cur_block.vars.keys())
+        new_op_size = len(cur_block.ops)
+        from .static.dist_context import get_default_distributed_context
+
+        default_dist_ctx = get_default_distributed_context()
+        for name in new_var_names:
+            if name not in self._old_var_names:
+                tensor = cur_block.vars[name]
+                dist_tensor = default_dist_ctx.get_dist_tensor_for_program(
+                    tensor
+                )
+                if dist_tensor is None:
+                    dist_tensor = DistributedTensor(cur_block.vars[name])
+                    dist_tensor.dist_attr.process_mesh = self
+                    dist_tensor.dist_attr.mark_annotated("process_mesh")
+                    default_dist_ctx.add_dist_tensor_for_program(dist_tensor)
+                else:
+                    if dist_tensor.dist_attr.process_mesh is None:
+                        dist_tensor.dist_attr.process_mesh = self
+                        dist_tensor.dist_attr.mark_annotated("process_mesh")
+
+        for idx in range(self._old_op_size, new_op_size):
+            op = cur_block.ops[idx]
+            dist_op = default_dist_ctx.get_dist_op_for_program(op)
+            if dist_op is None:
+                dist_op = DistributedOperator(op)
+                dist_op.dist_attr.process_mesh = self
+                dist_op.dist_attr.mark_annotated("process_mesh")
+                default_dist_ctx.add_dist_op_for_program(dist_op)
+            else:
+                if dist_op.dist_attr.process_mesh is None:
+                    dist_op.dist_attr.process_mesh = self
+                    dist_op.dist_attr.mark_annotated("process_mesh")
+        reset_current_process_mesh()
+
+    def __deepcopy__(self, memo: Any) -> ProcessMesh:
+        if id(self) in memo:
+            return memo[id(self)]
+        new_process_mesh = ProcessMesh(np.array(self.mesh), self.dim_names)
+        memo[id(self)] = new_process_mesh
+        return new_process_mesh
+
+    def __eq__(self, other: ProcessMesh | core.ProcessMesh) -> bool:
+        if not isinstance(other, (ProcessMesh, core.ProcessMesh)):
+            return False
+        if self.shape != other.shape or self.process_ids != other.process_ids:
+            return False
+        return True
+
+    def __ne__(self, other: ProcessMesh | core.ProcessMesh) -> None:
+        return not self.__eq__(other)
+
+    def __str__(self) -> str:
+        str = f"shape {self.shape}, process_ids {self.process_ids}, dim_names {self.dim_names}"
+        return str
+
+    def __hash__(self) -> int:
+        return super().__hash__()
+
+
+def compute_compatible_process_mesh(process_mesh_list):
+    """Compute the compatible process mesh given a list of process meshes."""
+    if not process_mesh_list:
+        return None
+
+    def _compute_compatible_process_mesh_of_two(pm1, pm2):
+        if pm1 is None:
+            return True, pm2
+        if pm2 is None:
+            return True, pm1
+        if pm1 == pm2:
+            return True, pm1
+        if pm1.process_ids == pm2.process_ids:
+            if len(pm1.shape) >= len(pm2.shape):
+                return True, pm1
+            else:
+                return True, pm2
+        process_set1 = set(pm1.process_ids)
+        process_set2 = set(pm2.process_ids)
+        if process_set1.issubset(process_set2):
+            return True, pm2
+        if process_set2.issubset(process_set1):
+            return True, pm1
+        return False, None
+
+    compatible_result = None
+    for process_mesh in process_mesh_list:
+        compatible, compatible_result = _compute_compatible_process_mesh_of_two(
+            compatible_result, process_mesh
+        )
+        if not compatible:
+            return None
+    return copy.deepcopy(compatible_result)
+
+
+def merge_process_meshes(process_meshes):
+    """Merge a list of process meshes."""
+    merged_process_mesh = None
+    merged_process_ids = set()
+    for process_mesh in process_meshes:
+        if process_mesh is not None:
+            process_ids = set(process_mesh.process_ids)
+            merged_process_ids = merged_process_ids.union(process_ids)
+    if len(merged_process_ids) != 0:
+        merged_process_mesh = ProcessMesh(list(merged_process_ids))
+    return merged_process_mesh

python/paddle/distributed/auto_parallel/random.py

@@ -0,0 +1,176 @@
+# 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 __future__ import annotations
+
+import contextlib
+import logging
+
+import paddle
+
+from ..utils.log_utils import get_logger
+from .process_mesh import retrieve_unique_id_for_process_mesh
+from .static.utils import _get_idx_in_axis
+
+_logger = get_logger(logging.INFO)
+
+_rng_name_to_seed = {}
+_rng_name_to_states = {}
+_inited_rng_name_to_seed = {}
+_enable_random_control = False
+_basic_seed = 42
+_basic_name = ""
+
+# use Prime number as offset to avoid conflict
+_mesh_offset = 173
+_dim_offsets = [11, 23, 37, 73]
+
+
+def is_enable_auto_rand_ctrl():
+    global _enable_random_control
+    return _enable_random_control
+
+
+def enable_auto_rand_ctrl():
+    global _enable_random_control
+    _enable_random_control = True
+
+
+def parallel_manual_seed(seed: int, name: str = "") -> None:
+    """Enable auto parallel random control.
+    Random control maintain the randomness when tensor is distributed across devices on a Mesh(any order).
+        * Independency: If tensor is **Sharded** on a Mesh dimension, Devices along that Mesh dimension should have Different randomness.
+
+        * Consistency:  Meanwhile if the tensor is **Replicated** on another Mesh dimension, randomness of Devices along that Mesh dimension should be Consistent.
+
+    For instance: rank0 ~ rank7 consist a Mesh of shape of [2, 4]; A 2D tensor is distributed in that Mesh using dims_mapping [-1, 1].
+    Randomness for rank0-rank1-rank2-rank3 (rank4-rank5-rank6-rank7) should be Independent;
+    Randomness for rank0 and rank4 (rank1 and rank5, ...) should be Consistent.
+
+    This function should be called only once before auto parallel compiles the computation graph (e.g. auto_parallel.engine.prepare() or fit()).
+
+    This seed only affects how randomness-relative **operators** (dropout, fuse op with dropout inside, etc) are execute among mesh, and would NOT affect other process like Parameter initialization.
+
+    Examples:
+        # seed relative to training step
+        auto_parallel_random_seed((step + 13) * 257)
+        ...
+        engine.prepare()
+    """
+
+    enable_auto_rand_ctrl()
+    global _basic_seed
+    _basic_seed = seed
+    global _basic_name
+    _basic_name = name
+
+
+def determinate_rng(
+    rank, dims_mapping=None, process_mesh=None, placements=None
+):
+    assert process_mesh is not None, "Must provide process mesh"
+    assert dims_mapping is not None or placements is not None, (
+        "Must provide one of dims mapping or placements."
+    )
+    assert not (dims_mapping is not None and placements is not None), (
+        "Cannot provide dims mapping and placements at same time."
+    )
+    # TODO(JZ-LIANG) Support Mesh with any high rank
+    # use a string to unique integer hashing algorithm for seed computation.
+    # instead of using offsets to coordinate seed across devices.
+    if len(process_mesh.shape) > 4:
+        raise NotImplementedError(
+            f"Auto Parallel Random Control for Mesh's rank > 4 is NOT supported! Got {process_mesh}"
+        )
+    global _basic_seed
+    seed_ = _basic_seed
+    global _basic_name
+    name_ = _basic_name
+
+    if name_:
+        name_ += "_"
+
+    # FIXME
+    # unique_id = process_mesh.unique_id
+    unique_id = retrieve_unique_id_for_process_mesh(
+        process_mesh.shape, process_mesh.process_ids
+    )
+    sharding_expr = name_ + f'mesh:{unique_id}'
+    seed_ += _mesh_offset * (unique_id + 1)
+
+    for i in range(len(process_mesh.shape)):
+        if (dims_mapping is not None and i not in dims_mapping) or (
+            placements is not None and not placements[i].is_shard()
+        ):
+            relative_idx = -1
+        else:
+            relative_idx = _get_idx_in_axis(
+                process_mesh.process_ids,
+                process_mesh.shape,
+                i,
+                rank,
+            )
+
+        sharding_expr += f"_dim{i}:{relative_idx}"
+        seed_ += _dim_offsets[i] * (relative_idx + 1)
+
+    global _rng_name_to_seed
+    global _rng_name_to_states
+    if sharding_expr in _rng_name_to_seed:
+        assert _rng_name_to_seed[sharding_expr] == seed_
+    else:
+        assert seed_ not in _rng_name_to_seed.values(), (
+            f"Seed Conflict! current seed: {seed_}, current sharding expr: {sharding_expr}, generated seed: {_rng_name_to_seed}"
+        )
+        _rng_name_to_seed[sharding_expr] = seed_
+        if paddle.in_dynamic_mode():
+            # for dygraph, just init the seed when meeting a new seed
+            orig_rng_state = paddle.get_rng_state()
+            paddle.seed(seed_)
+            _rng_name_to_states[sharding_expr] = paddle.get_rng_state()
+            paddle.set_rng_state(orig_rng_state)
+    return sharding_expr
+
+
+@contextlib.contextmanager
+def rng_state(name):
+    global _rng_name_to_states
+    assert name in _rng_name_to_states, (
+        f"The rng state name {name} haven't been init. "
+    )
+    orig_rng_state = paddle.get_rng_state()
+    paddle.set_rng_state(_rng_name_to_states[name])
+    try:
+        yield
+    finally:
+        _rng_name_to_states[name] = paddle.get_rng_state()
+        paddle.set_rng_state(orig_rng_state)
+
+
+def init_auto_parallel_rng():
+    if not is_enable_auto_rand_ctrl():
+        return
+
+    global _rng_name_to_seed
+    # NOTE init rng maybe call multiple times, avoid init same rng twice
+    global _inited_rng_name_to_seed
+
+    for rng_name, seed in _rng_name_to_seed.items():
+        if rng_name in _inited_rng_name_to_seed:
+            assert _inited_rng_name_to_seed[rng_name] == seed
+        else:
+            _logger.info(
+                f"Init Auto Parallel RNG: {rng_name}, with seed {seed}"
+            )
+            paddle.framework.random.set_random_seed_generator(rng_name, seed)
+            _inited_rng_name_to_seed[rng_name] = seed

python/paddle/distributed/auto_parallel/ring_attention.py

@@ -0,0 +1,542 @@
+# Copyright (c) 2025 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
+import paddle.distributed as dist
+import paddle.nn.functional as F
+from paddle import _C_ops
+
+
+def shard_seq_load_balance(tensor, seq_dim):
+    # dtensor Replicate() -> reorder -> Shard(seq_dim)
+    placements = tensor.placements
+    process_mesh = tensor.process_mesh
+    cp_index = process_mesh.dim_names.index('sep')
+    cp_degree = process_mesh.shape[cp_index]
+    if cp_degree > 1:
+        # split
+        sliced_datas = paddle.split(
+            tensor, num_or_sections=cp_degree * 2, axis=seq_dim
+        )
+        # resort [q0,q1,q2,q3] -> [q0,q3,q1,q2]
+        indices = []
+        for i in range(cp_degree):
+            indices.append(i)
+            indices.append(cp_degree * 2 - 1 - i)
+        reorder_indices = indices
+        reordered = [sliced_datas[i] for i in reorder_indices]
+        reordered_tensor = paddle.concat(reordered, axis=seq_dim)
+        # reshard q/k/v -> Shard(seq_dim)
+        placements[cp_index] = paddle.distributed.Shard(seq_dim)  # seq_dim:1
+        tensor = paddle.distributed.reshard(
+            reordered_tensor, process_mesh, placements
+        )
+    return tensor
+
+
+def unshard_seq_load_balance(tensor, seq_dim):
+    # dtensor Shard(seq_dim) -> Replicate() -> reorder
+    placements = tensor.placements
+    process_mesh = tensor.process_mesh
+    cp_index = process_mesh.dim_names.index('sep')
+    cp_degree = process_mesh.shape[cp_index]
+    all_tensor = dist.reshard(tensor, process_mesh, [dist.Replicate()])
+    sliced_datas = paddle.split(
+        all_tensor, num_or_sections=cp_degree * 2, axis=seq_dim
+    )
+    reorder_indices = []
+    for i in range(cp_degree):
+        reorder_indices.append(i)
+        reorder_indices.append(cp_degree * 2 - 1 - i)
+    inverse_indices = [0] * len(reorder_indices)
+    for idx, v in enumerate(reorder_indices):
+        inverse_indices[v] = idx
+    restored = [sliced_datas[i] for i in inverse_indices]
+    return paddle.concat(restored, axis=seq_dim)
+
+
+class RingCommunicator:
+    def __init__(self, group, local_key, local_value):
+        self._k_buffer = [
+            local_key.clone().contiguous(),
+            local_key.clone().contiguous(),
+        ]
+        self._v_buffer = [
+            local_value.clone().contiguous(),
+            local_value.clone().contiguous(),
+        ]
+
+        self._next_buffer_idx = 0
+
+        self.group = group
+        mesh = dist.auto_parallel.get_mesh()
+        process_id = dist.get_rank()
+        self.group_rank = mesh.get_rank_by_dim_and_process_id("sep", process_id)
+        self.cp_size = mesh.get_dim_size("sep")
+        cp_index = mesh.dim_names.index("sep")
+
+        self.send_rank = self.group.ranks[
+            (self.group_rank + 1) % self.cp_size
+        ]  # 1%2=1
+        self.recv_rank = self.group.ranks[(self.group_rank - 1) % self.cp_size]
+
+        self._reqs = []
+
+    def wait(self):
+        paddle.device.synchronize()
+
+    def add_to_buffers(self, key, value):
+        if key.shape != self._k_buffer[self._next_buffer_idx].shape:
+            self._k_buffer[self._next_buffer_idx][:, : key.shape[1], :, :].add_(
+                key
+            )
+            self._v_buffer[self._next_buffer_idx][:, : key.shape[1], :, :].add_(
+                value
+            )
+        else:
+            self._k_buffer[self._next_buffer_idx].add_(key)
+            self._v_buffer[self._next_buffer_idx].add_(value)
+
+    def get_buffers(self):
+        return (
+            self._k_buffer[self._next_buffer_idx],
+            self._v_buffer[self._next_buffer_idx],
+        )
+
+    def send_recv(self):
+        send_k_op = dist.P2POp(
+            dist.isend,
+            self._k_buffer[self._next_buffer_idx].contiguous(),
+            self.send_rank,
+            self.group,
+        )
+        send_v_op = dist.P2POp(
+            dist.isend,
+            self._v_buffer[self._next_buffer_idx].contiguous(),
+            self.send_rank,
+            self.group,
+        )
+        recv_k_op = dist.P2POp(
+            dist.irecv,
+            self._k_buffer[(self._next_buffer_idx + 1) % 2],
+            self.recv_rank,
+            self.group,
+        )
+        recv_v_op = dist.P2POp(
+            dist.irecv,
+            self._v_buffer[(self._next_buffer_idx + 1) % 2],
+            self.recv_rank,
+            self.group,
+        )
+
+        self._next_buffer_idx = (self._next_buffer_idx + 1) % 2
+
+        ops = [send_k_op, send_v_op, recv_k_op, recv_v_op]
+
+        self._reqs = dist.batch_isend_irecv(ops)
+
+
+def update_out_and_lse(
+    old_out, old_lse, block_out, block_lse, second_chunk_only=False
+):
+    if second_chunk_only:
+        second_chunk_out = old_out[:, old_out.shape[1] // 2 :, :, :]
+        second_chunk_lse = old_lse[:, old_lse.shape[1] // 2 :, :, :]
+        second_chunk_out, second_chunk_lse = update_out_and_lse(
+            second_chunk_out, second_chunk_lse, block_out, block_lse
+        )
+        old_out[:, old_out.shape[1] // 2 :, :, :] = second_chunk_out
+        old_lse[:, old_lse.shape[1] // 2 :, :, :] = second_chunk_lse
+        return old_out, old_lse
+    else:
+        block_out, block_lse = (
+            paddle.cast(block_out, "float32"),
+            paddle.cast(block_lse, "float32"),
+        )
+        with paddle.amp.auto_cast(enable=False):
+            return old_out - (old_out - block_out) * F.sigmoid(
+                block_lse - old_lse
+            ), old_lse - F.log_sigmoid(old_lse - block_lse)
+
+
+def get_chunk_id(rank, cp_size):
+    return rank, (2 * cp_size - 1 - rank)
+
+
+def concat_masks(attn_masks_list, rank, cp_size):
+    assert len(attn_masks_list) == 2 * cp_size
+    first_chunk_id, second_chunk_id = get_chunk_id(rank, cp_size)
+    return paddle.concat(
+        [attn_masks_list[first_chunk_id], attn_masks_list[second_chunk_id]],
+        axis=3,
+    )
+
+
+def ring_flash_attention_forward_func(
+    group,
+    local_query,
+    local_key,
+    local_value,
+    attn_mask=None,
+    dropout=0.0,
+    is_causal=False,
+    fixed_seed_offset=None,
+    training=True,
+):
+    cp_size = group.world_size
+    group_rank = group.rank
+
+    comm_buffer = RingCommunicator(group, local_key, local_value)
+    local_q_seq_len = local_query.shape[1]
+    if attn_mask is not None:
+        attn_masks_list = paddle.split(
+            attn_mask, num_or_sections=cp_size * 2, axis=3
+        )
+    if is_causal:
+        local_query_second_chunk = local_query[
+            :, local_q_seq_len // 2 :, :, :
+        ].contiguous()
+    for step in range(cp_size):
+        block_k, block_v = comm_buffer.get_buffers()
+        if step != cp_size - 1:
+            comm_buffer.send_recv()
+        if not is_causal:
+            # out [bs, seq, nhead, headdim]
+            # lse [bs, nhead, seq]
+            block_out, _, block_lse, _ = _C_ops.flash_attn(
+                local_query,
+                block_k,
+                block_v,
+                fixed_seed_offset,
+                (
+                    None
+                    if attn_mask is None
+                    else concat_masks(
+                        attn_masks_list, (group_rank - step) % cp_size, cp_size
+                    )
+                ),
+                dropout,
+                False,
+                False,
+                not training,
+                "",
+            )
+            paddle.unsqueeze_(paddle.transpose_(block_lse, [0, 2, 1]), axis=-1)
+
+            if step == 0:
+                out, lse = block_out, block_lse
+            else:
+                out, lse = update_out_and_lse(out, lse, block_out, block_lse)
+        else:
+            if step == 0:
+                block_out, _, block_lse, _ = _C_ops.flash_attn(
+                    local_query,
+                    block_k,
+                    block_v,
+                    fixed_seed_offset,
+                    None,
+                    dropout,
+                    True,
+                    False,
+                    not training,
+                    "",
+                )
+                paddle.unsqueeze_(
+                    paddle.transpose_(block_lse, [0, 2, 1]), axis=-1
+                )
+                out, lse = block_out, block_lse
+            elif step > group_rank:
+                block_out, _, block_lse, _ = _C_ops.flash_attn(
+                    local_query_second_chunk,
+                    block_k,
+                    block_v,
+                    fixed_seed_offset,
+                    None,
+                    dropout,
+                    False,
+                    False,
+                    not training,
+                    "",
+                )
+                block_lse = block_lse[:, :, 0 : (local_q_seq_len // 2)]
+                paddle.unsqueeze_(
+                    paddle.transpose_(block_lse, [0, 2, 1]), axis=-1
+                )
+                out, lse = update_out_and_lse(
+                    out, lse, block_out, block_lse, True
+                )
+            else:
+                block_out, _, block_lse, _ = _C_ops.flash_attn(
+                    local_query,
+                    block_k[:, : local_q_seq_len // 2, :, :],
+                    block_v[:, : local_q_seq_len // 2, :, :],
+                    fixed_seed_offset,
+                    None,
+                    dropout,
+                    False,
+                    False,
+                    not training,
+                    "",
+                )
+                paddle.unsqueeze_(
+                    paddle.transpose_(block_lse, [0, 2, 1]), axis=-1
+                )
+                out, lse = update_out_and_lse(out, lse, block_out, block_lse)
+
+        paddle.device.synchronize()
+    out = paddle.cast(out, local_query.dtype)
+    lse = paddle.transpose_(paddle.squeeze(lse, axis=-1), [0, 2, 1])
+    return out, lse
+
+
+def ring_flash_attention_backward_func(
+    group,
+    local_out_grad,
+    local_query,
+    local_key,
+    local_value,
+    local_out,
+    lse,
+    attn_mask,
+    dropout=0.0,
+    is_causal=False,
+    fixed_seed_offset=None,
+):
+    cp_size = group.world_size
+    group_rank = group.rank
+
+    lse = lse.contiguous()
+
+    local_q_seq_len = local_query.shape[1]
+    query_grad_buffer = paddle.zeros_like(local_query)
+    key_grad_buffer = paddle.zeros_like(local_key)
+    value_grad_buffer = paddle.zeros_like(local_value)
+
+    kv_comm_buffer = RingCommunicator(group, local_key, local_value)
+    grad_comm_buffer = RingCommunicator(
+        group, key_grad_buffer, value_grad_buffer
+    )
+    if is_causal:
+        local_query_second_chunk = local_query[:, local_q_seq_len // 2 :, :, :]
+        local_out_second_chunk = local_out[:, local_q_seq_len // 2 :, :, :]
+        lse_second_chunk = lse[:, :, local_q_seq_len // 2 :].contiguous()
+        out_grad_second_chunk = local_out_grad[:, local_q_seq_len // 2 :, :, :]
+
+    if attn_mask is not None:
+        attn_masks_list = paddle.split(
+            attn_mask, num_or_sections=cp_size * 2, axis=3
+        )
+    for step in range(cp_size):
+        block_k, block_v = kv_comm_buffer.get_buffers()
+        if step != cp_size - 1:
+            kv_comm_buffer.send_recv()
+
+        if not is_causal:
+            block_q_grad, block_k_grad, block_v_grad = _C_ops.flash_attn_grad(
+                local_query,
+                block_k,
+                block_v,
+                local_out,
+                lse,
+                fixed_seed_offset,
+                (
+                    None
+                    if attn_mask is None
+                    else concat_masks(
+                        attn_masks_list, (group_rank - step) % cp_size, cp_size
+                    )
+                ),
+                local_out_grad,
+                dropout,
+                False,
+            )
+            query_grad_buffer.add_(block_q_grad)
+        else:
+            if step == 0:
+                block_q_grad, block_k_grad, block_v_grad = (
+                    _C_ops.flash_attn_grad(
+                        local_query,
+                        block_k,
+                        block_v,
+                        local_out,
+                        lse,
+                        fixed_seed_offset,
+                        None,
+                        local_out_grad,
+                        dropout,
+                        True,
+                    )
+                )
+                query_grad_buffer.add_(block_q_grad)
+            elif step > group_rank:
+                block_q_grad, block_k_grad, block_v_grad = (
+                    _C_ops.flash_attn_grad(
+                        local_query_second_chunk,
+                        block_k,
+                        block_v,
+                        local_out_second_chunk,
+                        lse_second_chunk,
+                        fixed_seed_offset,
+                        None,
+                        out_grad_second_chunk,
+                        dropout,
+                        False,
+                    )
+                )
+                query_grad_buffer[:, local_q_seq_len // 2 :, :, :].add_(
+                    block_q_grad
+                )
+            else:
+                block_q_grad, block_k_grad, block_v_grad = (
+                    _C_ops.flash_attn_grad(
+                        local_query,
+                        block_k[:, : local_q_seq_len // 2, :, :],
+                        block_v[:, : local_q_seq_len // 2, :, :],
+                        local_out,
+                        lse,
+                        fixed_seed_offset,
+                        None,
+                        local_out_grad,
+                        dropout,
+                        False,
+                    )
+                )
+                query_grad_buffer.add_(block_q_grad)
+        paddle.device.synchronize()
+
+        grad_comm_buffer.add_to_buffers(
+            block_k_grad.contiguous(), block_v_grad.contiguous()
+        )
+        grad_comm_buffer.send_recv()
+
+        grad_comm_buffer.wait()
+    key_grad_buffer, value_grad_buffer = grad_comm_buffer.get_buffers()
+
+    return query_grad_buffer, key_grad_buffer, value_grad_buffer
+
+
+class RingFlashAttention(paddle.autograd.PyLayer):
+    @staticmethod
+    def forward(
+        ctx,
+        query,
+        key,
+        value,
+        attn_mask=None,
+        dropout=0.0,
+        is_causal=False,
+        fixed_seed_offset=None,
+        training=True,
+    ):
+        if dropout > 0.0:
+            raise NotImplementedError(
+                "Dropout is not supported in ring attention yet."
+            )
+        mesh = dist.auto_parallel.get_mesh()
+        cp_index = mesh.dim_names.index('sep')
+        process_id = dist.get_rank()
+        rank = mesh.get_rank_by_dim_and_process_id("sep", process_id)
+        dist.init_parallel_env()
+
+        group = mesh._get_group("sep")
+        local_query = dist.auto_parallel.api.dtensor_to_local(
+            query, query.process_mesh, query.placements
+        )
+        local_key = dist.auto_parallel.api.dtensor_to_local(
+            key, key.process_mesh, key.placements
+        )
+        local_value = dist.auto_parallel.api.dtensor_to_local(
+            value, value.process_mesh, value.placements
+        )
+        if attn_mask is not None:
+            is_causal = False
+
+        out, lse = ring_flash_attention_forward_func(
+            group,
+            local_query,
+            local_key,
+            local_value,
+            attn_mask,
+            dropout,
+            is_causal,
+            fixed_seed_offset,
+            training,
+        )
+        ctx.save_for_backward(group, query, key, value, out, lse, attn_mask)
+        ctx.fixed_seed_offset = fixed_seed_offset
+        ctx.dropout = dropout
+        ctx.is_causal = is_causal
+        out_dtensor = dist.auto_parallel.api.dtensor_from_local(
+            out, query.process_mesh, query.placements
+        )
+        return out_dtensor.contiguous()
+
+    @staticmethod
+    def backward(ctx, out_grad):
+        mesh = dist.auto_parallel.get_mesh()
+        cp_index = mesh.dim_names.index('sep')
+        group, query, key, value, out, lse, attn_mask = ctx.saved_tensor()
+        fixed_seed_offset = ctx.fixed_seed_offset
+        dropout = ctx.dropout
+        is_causal = ctx.is_causal
+
+        if fixed_seed_offset is None:
+            fixed_seed_offset = paddle.to_tensor(
+                [0, 0], place=paddle.CPUPlace(), dtype=paddle.int64
+            )
+        local_query = dist.auto_parallel.api.dtensor_to_local(
+            query, query.process_mesh, query.placements
+        )
+        local_key = dist.auto_parallel.api.dtensor_to_local(
+            key, key.process_mesh, key.placements
+        )
+        local_value = dist.auto_parallel.api.dtensor_to_local(
+            value, value.process_mesh, value.placements
+        )
+        local_out_grad = dist.auto_parallel.api.dtensor_to_local(
+            out_grad, out_grad.process_mesh, out_grad.placements
+        )
+        query_grad, key_grad, value_grad = ring_flash_attention_backward_func(
+            group,
+            local_out_grad,
+            local_query,
+            local_key,
+            local_value,
+            out,
+            lse,
+            attn_mask,
+            dropout,
+            is_causal,
+            fixed_seed_offset,
+        )
+        query_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
+            query_grad, query.process_mesh, query.placements
+        )
+        key_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
+            key_grad, key.process_mesh, key.placements
+        )
+        value_grad_dtensor = dist.auto_parallel.api.dtensor_from_local(
+            value_grad, value.process_mesh, value.placements
+        )
+
+        if attn_mask is not None and not attn_mask.stop_gradient:
+            return (
+                query_grad_dtensor,
+                key_grad_dtensor,
+                value_grad_dtensor,
+                None,
+            )
+        else:
+            return query_grad_dtensor, key_grad_dtensor, value_grad_dtensor

python/paddle/distributed/auto_parallel/ring_conv.py

@@ -0,0 +1,740 @@
+# Copyright (c) 2025 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 itertools
+
+import paddle
+import paddle.distributed as dist
+import paddle.nn.functional as F
+
+
+def _get_comm_group_by_dim(mesh, dim):
+    dim_names = mesh.dim_names
+    assert dim in dim_names, f"dim '{dim}' not in mesh.dim_names {dim_names}"
+
+    shape = mesh.shape
+    dim_idx = dim_names.index(dim)
+    ids = mesh.process_ids
+
+    def nest(flat, shape):
+        if not shape:
+            return flat[0]
+        step = int(len(flat) // shape[0])
+        return [
+            nest(flat[i * step : (i + 1) * step], shape[1:])
+            for i in range(shape[0])
+        ]
+
+    mesh_nd = nest(ids, shape)
+
+    other_axes = [i for i in range(len(shape)) if i != dim_idx]
+    other_ranges = [range(shape[i]) for i in other_axes]
+
+    comm_groups = []
+    for index in itertools.product(*other_ranges):
+        group = []
+        for i in range(shape[dim_idx]):
+            idx = list(index)
+            idx.insert(dim_idx, i)
+            val = mesh_nd
+            for j in idx:
+                val = val[j]
+            group.append(val)
+        comm_groups.append(group)
+
+    return comm_groups
+
+
+def _get_conv_tp_group(x_mesh, x_placements, data_format):
+    if data_format == "NCHW":
+        shard_axis = 3
+    else:
+        shard_axis = 2
+
+    axis_name = None
+    for i, placement in enumerate(x_placements):
+        if placement == dist.Shard(shard_axis):
+            axis_name = x_mesh.dim_names[i]
+            break
+
+    if not axis_name:
+        raise ValueError(
+            f"Input tensor placements {x_placements} do not contain a Shard on W axis:{shard_axis}."
+        )
+
+    tp_groups = _get_comm_group_by_dim(x_mesh, axis_name)
+    rank = dist.get_rank()
+
+    for group in tp_groups:
+        if rank in group:
+            return axis_name, group
+
+    raise RuntimeError(
+        f"Rank {rank} not found in any tensor parallel group for mesh {x_mesh}."
+    )
+
+
+def _ring_conv_halo_exchange(
+    local_input_tensor,
+    halo_width_to_receive_from_left,
+    halo_width_to_receive_from_right,
+    left_neighbor_rank,
+    right_neighbor_rank,
+    current_rank,
+    conv_tp_group,
+    data_format,
+):
+    if len(conv_tp_group) == 1:
+        return local_input_tensor
+
+    if not (
+        len(local_input_tensor.shape) == 4
+    ):  # Assuming 4D tensors like NCHW/NHWC
+        raise ValueError(
+            f"Input tensor is expected to be 4D for NCHW/NHWC formats, "
+            f"but got {len(local_input_tensor.shape)}D."
+        )
+
+    if data_format == "NCHW":
+        width_dim_idx = 3
+    elif data_format == "NHWC":
+        width_dim_idx = 2
+    else:
+        raise ValueError(
+            f"Unsupported data_format: {data_format}. Must be 'NCHW' or 'NHWC'."
+        )
+
+    # Segment to send to the right neighbor (right_neighbor_rank)
+    slices_for_send_right = [slice(None)] * 4
+    slices_for_send_right[width_dim_idx] = slice(
+        -halo_width_to_receive_from_left, None
+    )
+    segment_to_send_right = local_input_tensor[
+        tuple(slices_for_send_right)
+    ].contiguous()
+
+    # Segment to send to the left neighbor (left_neighbor_rank)
+    slices_for_send_left = [slice(None)] * 4
+    slices_for_send_left[width_dim_idx] = slice(
+        None, halo_width_to_receive_from_right
+    )
+    segment_to_send_left = local_input_tensor[
+        tuple(slices_for_send_left)
+    ].contiguous()
+
+    buffer_for_halo_from_right = paddle.zeros_like(segment_to_send_left)
+    buffer_for_halo_from_left = paddle.zeros_like(segment_to_send_right)
+
+    op_isend_to_right = dist.P2POp(
+        dist.isend, segment_to_send_right, right_neighbor_rank
+    )
+    op_isend_to_left = dist.P2POp(
+        dist.isend, segment_to_send_left, left_neighbor_rank
+    )
+
+    op_irecv_from_right = dist.P2POp(
+        dist.irecv, buffer_for_halo_from_right, right_neighbor_rank
+    )
+    op_irecv_from_left = dist.P2POp(
+        dist.irecv, buffer_for_halo_from_left, left_neighbor_rank
+    )
+
+    p2p_requests = dist.batch_isend_irecv(
+        [
+            op_isend_to_right,
+            op_isend_to_left,
+            op_irecv_from_left,
+            op_irecv_from_right,
+        ]
+    )
+    for req in p2p_requests:
+        req.wait()
+
+    # Concatenate received halo regions with the local tensor
+    if current_rank == conv_tp_group[0]:
+        # First rank: original tensor || halo_from_right
+        reconstructed_tensor = paddle.concat(
+            [local_input_tensor, buffer_for_halo_from_right], axis=width_dim_idx
+        )
+    elif current_rank == conv_tp_group[-1]:
+        # Last rank: halo_from_left || original tensor
+        reconstructed_tensor = paddle.concat(
+            [buffer_for_halo_from_left, local_input_tensor], axis=width_dim_idx
+        )
+    else:
+        # Middle ranks: halo_from_left || original tensor || halo_from_right
+        reconstructed_tensor = paddle.concat(
+            [
+                buffer_for_halo_from_left,
+                local_input_tensor,
+                buffer_for_halo_from_right,
+            ],
+            axis=width_dim_idx,
+        )
+
+    return reconstructed_tensor.contiguous()
+
+
+def _ring_conv_halo_aggregate(
+    local_gradient_tensor,
+    halo_width_send_left,
+    halo_width_send_right,
+    left_neighbor_rank,
+    right_neighbor_rank,
+    current_process_rank,
+    conv_tp_group,
+    data_format,
+):
+    if len(conv_tp_group) == 1:
+        return local_gradient_tensor
+
+    if data_format == "NCHW":
+        width_dim_idx = 3
+    elif data_format == "NHWC":
+        width_dim_idx = 2
+    else:
+        raise ValueError(
+            f"Unsupported data_format: {data_format}. Must be 'NCHW' or 'NHWC'."
+        )
+
+    # Prepare gradient segments to send
+    slices_for_send_right = [slice(None)] * 4
+    slices_for_send_right[width_dim_idx] = slice(
+        -halo_width_send_right, None
+    )  # Send the rightmost part
+    segment_to_send_right = local_gradient_tensor[
+        tuple(slices_for_send_right)
+    ].contiguous()
+
+    slices_for_send_left = [slice(None)] * 4
+    slices_for_send_left[width_dim_idx] = slice(
+        None, halo_width_send_left
+    )  # Send the leftmost part
+    segment_to_send_left = local_gradient_tensor[
+        tuple(slices_for_send_left)
+    ].contiguous()
+
+    # Buffers for receiving gradients
+    buffer_for_gradient_from_left = paddle.zeros_like(segment_to_send_right)
+    buffer_for_gradient_from_right = paddle.zeros_like(segment_to_send_left)
+
+    op_isend_to_right = dist.P2POp(
+        dist.isend, segment_to_send_right, right_neighbor_rank
+    )
+    op_isend_to_left = dist.P2POp(
+        dist.isend, segment_to_send_left, left_neighbor_rank
+    )
+    op_irecv_from_right = dist.P2POp(
+        dist.irecv, buffer_for_gradient_from_right, right_neighbor_rank
+    )
+    op_irecv_from_left = dist.P2POp(
+        dist.irecv, buffer_for_gradient_from_left, left_neighbor_rank
+    )
+
+    p2p_requests = dist.batch_isend_irecv(
+        [
+            op_isend_to_right,
+            op_isend_to_left,
+            op_irecv_from_left,
+            op_irecv_from_right,
+        ]
+    )
+    for req in p2p_requests:
+        req.wait()
+
+    processed_gradient_tensor = local_gradient_tensor
+    # Crop local tensor and aggregate received gradients
+    if current_process_rank == conv_tp_group[0]:
+        # Crop the part sent to the right neighbor
+        crop_slices = [slice(None)] * 4
+        crop_slices[width_dim_idx] = slice(None, -halo_width_send_right)
+        processed_gradient_tensor = processed_gradient_tensor[
+            tuple(crop_slices)
+        ]
+
+        # Aggregate gradient received from the right neighbor
+        # This is added to the new rightmost part of the processed_gradient_tensor
+        agg_slices = [slice(None)] * 4
+        agg_slices[width_dim_idx] = slice(-halo_width_send_left, None)
+
+        target_slice = processed_gradient_tensor[tuple(agg_slices)]
+        target_slice.add_(buffer_for_gradient_from_right)
+
+    elif current_process_rank == conv_tp_group[-1]:
+        # Crop the part sent to the left neighbor
+        crop_slices = [slice(None)] * 4
+        crop_slices[width_dim_idx] = slice(halo_width_send_left, None)
+        processed_gradient_tensor = processed_gradient_tensor[
+            tuple(crop_slices)
+        ]
+
+        # Aggregate gradient received from the left neighbor
+        agg_slices = [slice(None)] * 4
+        agg_slices[width_dim_idx] = slice(None, halo_width_send_right)
+
+        target_slice = processed_gradient_tensor[tuple(agg_slices)]
+        target_slice.add_(buffer_for_gradient_from_left)
+
+    else:
+        # Crop parts sent to both left and right neighbors
+        crop_slices = [slice(None)] * 4
+        crop_slices[width_dim_idx] = slice(
+            halo_width_send_left, -halo_width_send_right
+        )
+        processed_gradient_tensor = processed_gradient_tensor[
+            tuple(crop_slices)
+        ]
+
+        # Aggregate gradient received from the right neighbor
+        agg_slices_right_edge = [slice(None)] * 4
+        agg_slices_right_edge[width_dim_idx] = slice(
+            -halo_width_send_left, None
+        )
+        target_slice_right = processed_gradient_tensor[
+            tuple(agg_slices_right_edge)
+        ]
+        target_slice_right.add_(buffer_for_gradient_from_right)
+
+        # Aggregate gradient received from the left neighbor
+        agg_slices_left_edge = [slice(None)] * 4
+        agg_slices_left_edge[width_dim_idx] = slice(None, halo_width_send_right)
+        target_slice_left = processed_gradient_tensor[
+            tuple(agg_slices_left_edge)
+        ]
+        target_slice_left.add_(buffer_for_gradient_from_left)
+
+    return processed_gradient_tensor.contiguous()
+
+
+class RingConv2d(paddle.autograd.PyLayer):
+    @staticmethod
+    def _is_supported(
+        input_size, kernel_size, stride, padding, dilation, data_format="NCHW"
+    ):
+        idx_w_input = -1
+        idx_w_kernel = -1
+
+        if data_format == "NCHW":
+            # input_size: (N, C, H, W)
+            # kernel_size: (OutChannels, InChannels/Groups, KernelH, KernelW)
+            idx_w_input = 3
+            idx_w_kernel = 3
+        elif data_format == "NHWC":
+            # input_size: (N, H, W, C)
+            # kernel_size: (OutChannels, InChannels/Groups, KernelH, KernelW)
+            idx_w_input = 2
+            idx_w_kernel = 3
+        else:
+            raise ValueError(
+                f"Unsupported data_format '{data_format}'. Expected 'NCHW' or 'NHWC'."
+            )
+
+        dilation_w = dilation[1]
+        padding_w = padding[1]
+        stride_w = stride[1]
+
+        input_w = input_size[idx_w_input]
+        kernel_w = kernel_size[idx_w_kernel]
+
+        if dilation_w != 1:
+            # RingConv2d only supports dilation=1.
+            # Larger dilation would require enlarged halo regions and more complex communication.
+            raise RuntimeError(
+                f"Only dilation=1 on the W-dimension is supported for tensor-parallel convolution. "
+                f"Got dilation_w={dilation_w} (data_format='{data_format}')."
+            )
+
+        if padding_w == 0:
+            # To avoid halo exchange when padding=0, we require:
+            # - input_w must be divisible by stride_w, so partitions align evenly across ranks.
+            # - stride_w == kernel_w, so each kernel operates on disjoint local regions.
+            if input_w % stride_w != 0:
+                raise RuntimeError(
+                    f"When padding_w=0, input_W={input_w} must be divisible by stride_W={stride_w} "
+                    f"for tensor-parallel convolution (data_format='{data_format}')."
+                )
+            if stride_w != kernel_w:
+                raise RuntimeError(
+                    f"When padding_w=0, stride_W={stride_w} must equal kernel_W={kernel_w} "
+                    f"to avoid halo exchange (data_format='{data_format}')."
+                )
+
+        else:
+            # When padding > 0, halo exchange is needed.
+            # To simplify halo logic, we require:
+            # - stride_w == 1: ensures each output element is computed from overlapping input,
+            #   and no input region is skipped, simplifying halo construction.
+            # - kernel_w // 2 <= input_w: prevents the kernel from exceeding local input.
+            if stride_w != 1:
+                raise RuntimeError(
+                    f"When padding_w={padding_w}, stride_W must be 1 for tensor-parallel convolution. "
+                    f"Got stride_W={stride_w} (data_format='{data_format}')."
+                )
+            if kernel_w // 2 > input_w:
+                raise RuntimeError(
+                    f"Half of kernel_W ({kernel_w // 2}) must not exceed input_W={input_w} "
+                    f"to ensure halo region fits (data_format='{data_format}')."
+                )
+
+        return True
+
+    @staticmethod
+    def forward(
+        ctx,
+        x,
+        weight,
+        bias=None,
+        stride=1,
+        padding=0,
+        padding_algorithm=None,
+        dilation=1,
+        groups=1,
+        data_format="NCHW",
+        channel_dim=1,
+    ):
+        rank = dist.get_rank()
+
+        assert RingConv2d._is_supported(
+            x.shape, weight.shape, stride, padding, dilation, data_format
+        )
+        assert x.is_dist(), "Input tensor `x` must be a distributed tensor."
+
+        if not weight.is_dist():
+            weight_placements = [
+                dist.Replicate() for _ in range(len(x.placements))
+            ]
+            weight = dist.auto_parallel.api.dtensor_from_local(
+                weight, x.process_mesh, weight_placements
+            )
+
+        if bias is not None and not bias.is_dist():
+            bias_placements = [
+                dist.Replicate() for _ in range(len(x.placements))
+            ]
+            bias = dist.auto_parallel.api.dtensor_from_local(
+                bias, x.process_mesh, bias_placements
+            )
+
+        ctx.save_for_backward(x, weight, bias)
+
+        x_mesh = x.process_mesh
+        x_placements = x.placements
+        x = dist.auto_parallel.api.dtensor_to_local(x, x_mesh, x_placements)
+
+        weight = dist.auto_parallel.api.dtensor_to_local(
+            weight, weight.process_mesh, weight.placements
+        )
+        if bias is not None:
+            bias = dist.auto_parallel.api.dtensor_to_local(
+                bias, bias.process_mesh, bias.placements
+            )
+
+        ctx.attrs = (
+            stride,
+            padding,
+            padding_algorithm,
+            dilation,
+            groups,
+            data_format,
+        )
+
+        mesh_axis_name, conv_tp_group = _get_conv_tp_group(
+            x_mesh, x_placements, data_format
+        )
+        if padding[1] == 0 or len(conv_tp_group) <= 1:
+            final_local_results = paddle._C_ops.conv2d(
+                x,
+                weight,
+                stride,
+                padding,
+                padding_algorithm,
+                dilation,
+                groups,
+                data_format,
+            )
+        else:
+            # step 0: calculate the required overlap (halo) pixels for the input tensor
+            if data_format == "NCHW":
+                kernel_width_dim_idx = 3
+                output_width_dim_idx = 3
+            elif data_format == "NHWC":
+                kernel_width_dim_idx = 3
+                output_width_dim_idx = 2
+            else:
+                raise ValueError(
+                    f"Unsupported data_format: {data_format}. Must be 'NCHW' or 'NHWC'."
+                )
+
+            kernel_width = weight.shape[kernel_width_dim_idx]
+            kernel_total_halo_span = kernel_width - 1
+            left_halo_width = kernel_total_halo_span // 2
+            right_halo_width = kernel_total_halo_span - left_halo_width
+            assert left_halo_width + right_halo_width == kernel_total_halo_span
+
+            ctx.mesh_axis_name = mesh_axis_name
+            rank_idx = conv_tp_group.index(rank)
+            next_rank = conv_tp_group[(rank_idx + 1) % len(conv_tp_group)]
+            prev_rank = conv_tp_group[(rank_idx - 1) % len(conv_tp_group)]
+
+            # step 1: reconstruct the local input tensor including halo regions via ring communication
+            # `x` is updated here, now including halo data received from neighboring ranks.
+            x = _ring_conv_halo_exchange(
+                x,
+                left_halo_width,
+                right_halo_width,
+                prev_rank,
+                next_rank,
+                rank,
+                conv_tp_group,
+                data_format,
+            )
+
+            # step 2: feed the reconstructed local input tensor to the actual computation (op_call)
+            local_results_with_halo = paddle._C_ops.conv2d(
+                x,
+                weight,
+                stride,
+                padding,
+                padding_algorithm,
+                dilation,
+                groups,
+                data_format,
+            )
+
+            # step 3: remove extra output portions from the results, generated from processing halo regions
+            # `padding[1]` (from outer scope) is assumed here to be the width of the halo/overlap
+            # that needs to be trimmed from each side of the output.
+            output_halo_trim_width = padding[1]
+            width_before_trimming = local_results_with_halo.shape[
+                output_width_dim_idx
+            ]
+
+            if data_format == "NCHW":
+                if rank == conv_tp_group[0]:
+                    final_local_results = local_results_with_halo[
+                        :,
+                        :,
+                        :,
+                        : width_before_trimming - output_halo_trim_width,
+                    ]
+                elif rank == conv_tp_group[-1]:
+                    final_local_results = local_results_with_halo[
+                        :, :, :, output_halo_trim_width:
+                    ]
+                else:
+                    final_local_results = local_results_with_halo[
+                        :,
+                        :,
+                        :,
+                        output_halo_trim_width : width_before_trimming
+                        - output_halo_trim_width,
+                    ]
+            else:
+                if rank == conv_tp_group[0]:
+                    final_local_results = local_results_with_halo[
+                        :,
+                        :,
+                        : width_before_trimming - output_halo_trim_width,
+                        :,
+                    ]
+                elif rank == conv_tp_group[-1]:
+                    final_local_results = local_results_with_halo[
+                        :, :, output_halo_trim_width:, :
+                    ]
+                else:
+                    final_local_results = local_results_with_halo[
+                        :,
+                        :,
+                        output_halo_trim_width : width_before_trimming
+                        - output_halo_trim_width,
+                        :,
+                    ]
+
+            ctx.left_halo_width = left_halo_width
+            ctx.right_halo_width = right_halo_width
+            ctx.output_halo_trim_width = output_halo_trim_width
+            ctx.output_width_dim_idx = output_width_dim_idx
+
+        final_local_results = dist.auto_parallel.api.dtensor_from_local(
+            final_local_results, x_mesh, x_placements
+        )
+
+        return final_local_results.contiguous()
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        current_rank = dist.get_rank()
+        x, weight, bias = ctx.saved_tensor()
+
+        x_stop_gradient = x.stop_gradient
+        weight_stop_gradient = weight.stop_gradient
+        bias_stop_gradient = bias.stop_gradient if bias is not None else True
+
+        x_mesh = x.process_mesh
+        x_placements = x.placements
+        x = dist.auto_parallel.api.dtensor_to_local(x, x_mesh, x_placements)
+
+        weight_mesh = weight.process_mesh
+        weight_placements = weight.placements
+        weight = dist.auto_parallel.api.dtensor_to_local(
+            weight, weight_mesh, weight_placements
+        )
+
+        grad_out = dist.auto_parallel.api.dtensor_to_local(
+            grad_out, grad_out.process_mesh, grad_out.placements
+        )
+
+        if bias is not None:
+            bias_mesh = bias.process_mesh
+            bias_placements = bias.placements
+            bias = dist.auto_parallel.api.dtensor_to_local(
+                bias, bias_mesh, bias_placements
+            )
+
+        conv_attrs = ctx.attrs
+        data_format = conv_attrs[-1]
+        padding = conv_attrs[1]
+
+        grad_x = None
+        grad_weight = None
+        grad_bias = None
+
+        _, conv_tp_group = _get_conv_tp_group(x_mesh, x_placements, data_format)
+
+        if padding[1] == 0 or len(conv_tp_group) <= 1:
+            grad_x, grad_weight = paddle._C_ops.conv2d_grad(
+                x, weight, grad_out, *conv_attrs
+            )
+        else:
+            rank_idx = conv_tp_group.index(current_rank)
+            next_rank = conv_tp_group[(rank_idx + 1) % len(conv_tp_group)]
+            prev_rank = conv_tp_group[(rank_idx - 1) % len(conv_tp_group)]
+
+            left_halo_width = ctx.left_halo_width
+            right_halo_width = ctx.right_halo_width
+            output_halo_trim_width = ctx.output_halo_trim_width
+            output_width_dim_idx = ctx.output_width_dim_idx
+
+            # Step 1: Reconstruct `in_tensor_augmented` (original input to local conv in forward)
+            in_tensor_augmented = _ring_conv_halo_exchange(
+                x,
+                left_halo_width,
+                right_halo_width,
+                prev_rank,
+                next_rank,
+                current_rank,
+                conv_tp_group,
+                data_format,
+            )
+
+            # Step 2: Pad `grad_out` to match the output shape of conv on augmented input
+            padding_w = padding[1]
+            if data_format == "NCHW":
+                if current_rank == conv_tp_group[0]:
+                    padding_list = [0, padding_w]
+                elif current_rank == conv_tp_group[-1]:
+                    padding_list = [padding_w, 0]
+                else:
+                    padding_list = [padding_w, padding_w]
+            else:
+                if current_rank == conv_tp_group[0]:
+                    padding_list = [0, padding_w, 0, 0]
+                elif current_rank == conv_tp_group[-1]:
+                    padding_list = [padding_w, 0, 0, 0]
+                else:
+                    padding_list = [padding_w, padding_w, 0, 0]
+
+            grad_out_padded = F.pad(
+                grad_out,
+                padding_list,
+                mode="constant",
+                value=0.0,
+                data_format=data_format,
+            )
+
+            # Step 3: Local backward computation using augmented/padded tensors
+            # `padding` here is the original conv padding from forward.
+            grad_x_augmented, grad_weight = paddle._C_ops.conv2d_grad(
+                in_tensor_augmented, weight, grad_out_padded, *conv_attrs
+            )
+
+            # Step 4: Aggregate "halo" regions for grad_input
+            if not x_stop_gradient:
+                grad_x = _ring_conv_halo_aggregate(
+                    grad_x_augmented,
+                    left_halo_width,
+                    right_halo_width,
+                    prev_rank,
+                    next_rank,
+                    current_rank,
+                    conv_tp_group,
+                    data_format,
+                )
+
+        if bias is not None:
+            sum_axes = [0, 2, 3] if data_format == "NCHW" else [0, 1, 2]
+            grad_bias = paddle.sum(grad_out, axis=sum_axes, keepdim=True)
+            grad_bias = grad_bias.reshape(bias.shape)
+
+        if grad_x is not None:
+            grad_x = dist.auto_parallel.api.dtensor_from_local(
+                grad_x, x_mesh, x_placements
+            )
+
+        # Note(luchang): With input X sharded along tp_axis_name and weight W replicated,
+        # the locally computed grad_weight is only a partial sum for the full dL/dW,
+        # as dL/dW depends on contributions from all input shards.
+        # Aggregation across TP ranks is therefore necessary. Partial(ReduceSum)
+        # declares this averaging intent, and reshard to Replicate() executes
+        # the AllReduce-average, making the correct averaged grad_weight available
+        # and replicated on all TP ranks.
+        tp_axis_name, _ = _get_conv_tp_group(x_mesh, x_placements, data_format)
+        for idx, axis_name in enumerate(weight_mesh.dim_names):
+            if axis_name == tp_axis_name:
+                weight_placements[idx] = dist.Partial(dist.ReduceType.kRedSum)
+                if bias is not None:
+                    bias_placements[idx] = dist.Partial(dist.ReduceType.kRedSum)
+
+        grad_weight = dist.auto_parallel.api.dtensor_from_local(
+            grad_weight, weight_mesh, weight_placements
+        )
+        # do allreduce to get right grad_weight
+        grad_weight = dist.reshard(
+            grad_weight,
+            weight_mesh,
+            [dist.Replicate() for _ in range(len(weight_placements))],
+        )
+
+        if bias is not None:
+            grad_bias = dist.auto_parallel.api.dtensor_from_local(
+                grad_bias, bias_mesh, bias_placements
+            )
+            # do allreduce to get right grad_bias
+            grad_bias = dist.reshard(
+                grad_bias,
+                weight_mesh,
+                [dist.Replicate() for _ in range(len(bias_placements))],
+            )
+
+        if x_stop_gradient:
+            grad_x = None
+        if weight_stop_gradient:
+            grad_weight = None
+        if bias_stop_gradient:
+            grad_bias = None
+
+        if bias is not None:
+            return grad_x, grad_weight, grad_bias
+
+        return grad_x, grad_weight

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

@@ -0,0 +1,13 @@
+#   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.

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

@@ -0,0 +1,532 @@
+# 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 collections
+import copy
+import os
+import pickle
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.base import core
+from paddle.base.framework import Program
+from paddle.distributed.auto_parallel.static.converter import Converter
+from paddle.distributed.auto_parallel.static.dist_context import (
+    get_default_distributed_context,
+)
+from paddle.distributed.auto_parallel.static.utils import (
+    is_backward_op,
+    is_forward_op,
+    is_loss_op,
+)
+from paddle.static.io import deserialize_program
+
+_valid_types = [
+    core.VarDesc.VarType.DENSE_TENSOR,
+    core.VarDesc.VarType.SELECTED_ROWS,
+    core.VarDesc.VarType.DENSE_TENSOR_ARRAY,
+]
+
+paddle.enable_static()
+
+
+class AutoAlignTool:
+    """
+    This is an automatic parallel precision alignment tool。
+    """
+
+    def __init__(self, program: Program, step=1, fetch_list=None):
+        """Set some initialization information of the tool.
+        step: Step when returning a specific variable name。
+        fetch_list: initialization fetch_list.When a specific step is not reached, return this.
+                 It can combine with Engine class。
+                 example:in Engine.fit function,like this
+                         try:
+                             fetch_list = []
+                             align_tool = AutoAlignTool(self.main_program, 0, fetch_names)
+                             level = 0
+                             fetch_list = align_tool.get_var(level, step)
+                             outs = self._executor.run(
+                                 self.main_program,
+                                 fetch_list=fetch_list,
+                                 use_program_cache=self._strategy.use_cache,
+                                 return_numpy=self._strategy.return_numpy,
+                             )
+                             if fetch_list != fetch_names:
+                                 align_tool.save(dir_path, outs, fetch_list, self._dist_contexts["train"], self.serial)
+                                 exit(0)
+                         except core.EOFException:
+                             break
+        """
+        assert isinstance(program, Program)
+        self._program = program
+        self._blocks = program.blocks
+        self._step = step
+        self._fetch_list = fetch_list
+        assert self._blocks is not None
+
+    def set_step(self, step):
+        self._step = step
+
+    def get_var(self, level, step):
+        """
+        level must be in [0,1,2,3,4,5].
+        """
+        if step != self._step or step == -1:
+            return self._fetch_list
+        if level == 0:
+            return self.get_loss_lr_var()
+        elif level == 1:
+            return self.get_data_var()
+        elif level == 2:
+            return self.get_param_var()
+        elif level == 3:
+            return self.get_param_grad_var()
+        elif level == 4:
+            return self.get_forward_tmp_var()
+        elif level == 5:
+            return self.get_backward_tmp_var()
+        else:
+            raise ValueError
+
+    def set_program(self, program: Program):
+        assert isinstance(program, Program)
+        self._program = program
+        self._blocks = program.blocks
+        assert self._blocks is not None
+
+    def get_loss_lr_var(self):
+        """
+        Returns the variable name of learning rate and loss
+        """
+        fetch_set = set()
+        loss_ops = []
+        for block in self._blocks:
+            for op in block.ops:
+                if is_loss_op(op):
+                    assert len(op.desc.output_arg_names()) == 1, (
+                        "loss op should only output loss var"
+                    )
+                    loss_ops.append(op)
+
+        for block in self._blocks:
+            for varname in block.vars:
+                var = block._find_var_recursive(varname)
+
+                if var is None or var.type not in _valid_types:
+                    continue
+
+                if "learning_rate" in var.name:
+                    fetch_set.add(var.name)
+
+        for loss_op in loss_ops:
+            fetch_set.add(loss_op.output_arg_names[0])
+
+        return list(fetch_set)
+
+    def get_data_var(self):
+        """
+        Returns the variable name of data.
+        """
+        fetch_set = set()
+        for block in self._blocks:
+            for varname in block.vars:
+                var = block._find_var_recursive(varname)
+
+                if var is None or var.type not in _valid_types:
+                    continue
+
+                if var.is_data:
+                    fetch_set.add(var.name)
+        return list(fetch_set)
+
+    def get_param_var(self):
+        """
+        Returns the variable name of parameters.
+        """
+        fetch_set = set()
+        for block in self._blocks:
+            for op in block.ops:
+                if is_backward_op(op):
+                    break
+                for varname in op.input_arg_names + op.output_arg_names:
+                    var = block._find_var_recursive(varname)
+                    if var is None or var.type not in _valid_types:
+                        continue
+                    if var.is_parameter:
+                        fetch_set.add(varname)
+
+        return list(fetch_set)
+
+    def get_param_grad_var(self):
+        """
+        Returns the variable name of parameters' gradient.
+        """
+        fetch_set = set()
+        for block in self._blocks:
+            for op in block.ops:
+                if is_forward_op(op):
+                    continue
+                for varname in op.input_arg_names + op.output_arg_names:
+                    if "@GRAD" not in varname:
+                        continue
+                    fwd_varname = varname.split("@GRAD")[0]
+                    fwd_var = block._find_var_recursive(fwd_varname)
+                    if fwd_var is None or fwd_var.type not in _valid_types:
+                        continue
+                    if fwd_var.is_parameter is False:
+                        continue
+                    var = block._find_var_recursive(varname)
+                    if var is None or var.type not in _valid_types:
+                        continue
+                    fetch_set.add(varname)
+
+        return list(fetch_set)
+
+    def get_forward_tmp_var(self):
+        """
+        Returns the name of the temporary variable in the forward propagation
+        """
+        fetch_set = set()
+        loss_lr_list = self.get_loss_lr_var()
+        for block in self._blocks:
+            for op in block.ops:
+                if is_backward_op(op):
+                    break
+                for varname in op.input_arg_names + op.output_arg_names:
+                    if varname in loss_lr_list:
+                        continue
+                    var = block._find_var_recursive(varname)
+                    if var is None or var.type not in _valid_types:
+                        continue
+                    if var.is_data or var.is_parameter:
+                        continue
+                    fetch_set.add(varname)
+
+        return list(fetch_set)
+
+    def get_backward_tmp_var(self):
+        """
+        Returns the name of a temporary variable in back-propagation
+        """
+        fetch_set = set()
+        loss_lr_list = self.get_loss_lr_var()
+        forward_tmp_list = self.get_forward_tmp_var()
+        for block in self._blocks:
+            for op in block.ops:
+                if is_backward_op(op):
+                    for varname in op.input_arg_names + op.output_arg_names:
+                        if (
+                            varname in loss_lr_list
+                            or varname in forward_tmp_list
+                        ):
+                            continue
+                        if "@GRAD" in varname:
+                            fwd_varname = varname.split("@GRAD")[0]
+                            fwd_var = block._find_var_recursive(fwd_varname)
+                            if (
+                                fwd_var is not None
+                                and fwd_var.type in _valid_types
+                            ):
+                                if fwd_var.is_parameter:
+                                    continue
+                        var = block._find_var_recursive(varname)
+                        if var is None or var.type not in _valid_types:
+                            continue
+                        if var.is_data or var.is_parameter:
+                            continue
+                        fetch_set.add(varname)
+
+        return list(fetch_set)
+
+    def save(self, save_dir, vars, fetch_list, dist_context=None):
+        """
+        save fetch variables, distributed properties of variables and program.
+        """
+        if os.path.exists(save_dir) is False:
+            os.mkdir(save_dir)
+        if dist_context is None:
+            dist_context = get_default_distributed_context()
+        assert os.path.exists(save_dir)
+        if dist.get_world_size() == 1:
+            vars_path = os.path.join(save_dir, "vars.pkl")
+            program_path = os.path.join(save_dir, "program.pdmodel")
+            dist_attr_path = os.path.join(save_dir, "dist_attr.pkl")
+        else:
+            vars_path = os.path.join(
+                save_dir, f"vars_rank{dist.get_rank()}.pkl"
+            )
+            program_path = os.path.join(
+                save_dir, f"program_rank{dist.get_rank()}.pdmodel"
+            )
+            dist_attr_path = os.path.join(
+                save_dir, f"dist_attr_rank{dist.get_rank()}.pkl"
+            )
+        if vars is not None:
+            vars_dict = {}
+            assert len(fetch_list) == len(vars)
+            for i in range(len(fetch_list)):
+                if vars[i] is None:
+                    continue
+                vars_dict[fetch_list[i]] = vars[i]
+            with open(vars_path, "wb") as f:
+                pickle.dump(vars_dict, f)
+            dist_attr = {}
+            for var in self._program.list_vars():
+                if var.name not in fetch_list:
+                    continue
+                tensor_dist_attr = (
+                    dist_context.get_tensor_dist_attr_for_program(var)
+                )
+                if tensor_dist_attr is None:
+                    continue
+                process_mesh = tensor_dist_attr.process_mesh
+                dims_mapping = tensor_dist_attr.dims_mapping
+                dist_attr[var.name] = {
+                    "process_shape": process_mesh.shape,
+                    "process_group": process_mesh.process_ids,
+                    "dims_mapping": dims_mapping,
+                }
+            if len(dist_attr) > 0:
+                with open(dist_attr_path, "wb") as f:
+                    pickle.dump(dist_attr, f)
+        if self._program is not None:
+            with open(program_path, "wb") as f:
+                f.write(self._program.desc.serialize_to_string())
+
+    @staticmethod
+    def load(save_dir):
+        assert os.path.exists(save_dir)
+        filename_list = sorted(os.listdir(save_dir))
+        vars_list = []
+        program_list = []
+        dist_attr_list = []
+        for filename in filename_list:
+            filepath = os.path.join(save_dir, filename)
+            assert os.path.isfile(filepath)
+            if "vars" in filename:
+                assert filename.endswith("pkl")
+                with open(filepath, "rb") as f:
+                    from paddle.framework.restricted_unpickler import (
+                        safe_load_pickle,
+                    )
+
+                    vars_list.append(safe_load_pickle(f))
+            elif "program" in filename:
+                assert filename.endswith("pdmodel")
+                with open(filepath, "rb") as f:
+                    program_string = f.read()
+                program_list.append(deserialize_program(program_string))
+            elif "dist_attr" in filename:
+                assert filename.endswith("pkl")
+                with open(filepath, "rb") as f:
+                    from paddle.framework.restricted_unpickler import (
+                        safe_load_pickle,
+                    )
+
+                    dist_attr_list.append(safe_load_pickle(f))
+
+        dist_attr_map = {}
+        for dist_attrs in dist_attr_list:
+            for dist_attr_name in dist_attrs.keys():
+                if dist_attr_name not in dist_attr_map:
+                    dist_attr_map[dist_attr_name] = dist_attrs[dist_attr_name]
+        assert len(vars_list) == len(program_list)
+        return vars_list, program_list, dist_attr_map
+
+    @staticmethod
+    def convert_src_tensor_2_dst_tensor(vars_list, src_attr_map, dst_attr_map):
+        """
+        Converter is a class object for auto parallel to convert tensors from
+        one parallel strategy to another one. Tensors will merge and slice value
+        with their strategy when strategies are different.
+        But like dp to pp or dp to serial is not supported.
+        """
+        assert len(vars_list) >= 1
+        # if dist_attr_map is None or len(dist_attr_map) == 0 or len(vars_list) == 1:
+        if src_attr_map is None or len(src_attr_map) == 0:
+            return vars_list[0]
+
+        dst_strategies = {}
+        src_strategies = {}
+        tensors_dict = {}
+
+        convert_tensor_dict = None
+        for var_name in src_attr_map.keys():
+            assert var_name not in dst_strategies
+            dist_vars = []
+            for vars in vars_list:
+                if var_name in vars.keys():
+                    dist_vars.append(vars[var_name])
+            if len(dist_vars) == 0:
+                continue
+
+            if var_name in dst_attr_map and var_name in src_attr_map:
+                dst_strategies[var_name] = copy.deepcopy(dst_attr_map[var_name])
+                src_strategies[var_name] = copy.deepcopy(src_attr_map[var_name])
+                tensors_dict[var_name] = dist_vars
+
+        if src_attr_map == dst_attr_map:
+            return tensors_dict
+        converter = Converter(tensors_dict, src_strategies, dst_strategies)
+        convert_tensor_dict = converter.convert()
+
+        return convert_tensor_dict
+
+    @staticmethod
+    def find_diff_vars(fixed_vars_map, query_vars_map):
+        """
+        Found two variable names with different variable lists
+        """
+        diff_var_name_list = set()
+        for var_name in fixed_vars_map.keys():
+            if var_name in query_vars_map:
+                fixed_vars = fixed_vars_map[var_name]
+                query_vars = query_vars_map[var_name]
+                if isinstance(fixed_vars, np.ndarray):
+                    fixed_vars = [fixed_vars]
+                if isinstance(query_vars, np.ndarray):
+                    query_vars = [query_vars]
+
+                length = min(len(fixed_vars), len(query_vars))
+                if len(fixed_vars) != len(query_vars):
+                    print()
+                for i in range(length):
+                    if not np.allclose(fixed_vars[i], query_vars[i]):
+                        diff_var_name_list.add(var_name)
+        return diff_var_name_list
+
+    @staticmethod
+    def diff_information(right_dir, wrong_dir):
+        """
+        Find the corresponding operator according to the variable name.
+        """
+        (
+            right_vars_list,
+            right_program_list,
+            right_dist_attr_map,
+        ) = AutoAlignTool.load(right_dir)
+        (
+            wrong_vars_list,
+            wrong_program_list,
+            wrong_dist_attr_map,
+        ) = AutoAlignTool.load(wrong_dir)
+        right_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
+            right_vars_list, right_dist_attr_map, right_dist_attr_map
+        )
+        wrong_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
+            wrong_vars_list, wrong_dist_attr_map, right_dist_attr_map
+        )
+
+        diff_var_name_list = AutoAlignTool.find_diff_vars(
+            right_tensors_dict, wrong_tensors_dict
+        )
+
+        diff_ops_varname_dict = collections.OrderedDict()
+
+        for program in wrong_program_list:
+            for block in program.blocks:
+                for op in block.ops:
+                    for varname in op.input_arg_names + op.output_arg_names:
+                        if varname in diff_var_name_list:
+                            if len(diff_ops_varname_dict) == 0:
+                                print(
+                                    "first different op:\n",
+                                    op,
+                                    f"\ndifferent varname is:{varname}",
+                                )
+                            if op not in diff_ops_varname_dict:
+                                diff_ops_varname_dict[op] = [varname]
+                            else:
+                                diff_ops_varname_dict[op].append(varname)
+
+        return diff_ops_varname_dict
+
+    @staticmethod
+    def diff_information_from_dirs(right_dirs, wrong_dirs):
+        right_vars_list = []
+        right_program_list = []
+        right_dist_attr_map = {}
+        for right_dir in right_dirs:
+            (
+                tmp_vars_list,
+                right_program_list,
+                tmp_dist_attr_map,
+            ) = AutoAlignTool.load(right_dir)
+            if len(right_vars_list) == 0:
+                right_vars_list = tmp_vars_list
+            else:
+                for i in range(len(tmp_vars_list)):
+                    vars_list = tmp_vars_list[i]
+                    for key in vars_list.keys():
+                        if key not in right_vars_list[i].keys():
+                            right_vars_list[i][key] = vars_list[key]
+
+            for key in tmp_dist_attr_map.keys():
+                if key not in right_dist_attr_map:
+                    right_dist_attr_map[key] = tmp_dist_attr_map[key]
+
+        wrong_vars_list = []
+        wrong_program_list = []
+        wrong_dist_attr_map = {}
+        for wrong_dir in wrong_dirs:
+            (
+                tmp_vars_list,
+                wrong_program_list,
+                tmp_dist_attr_map,
+            ) = AutoAlignTool.load(wrong_dir)
+            if len(wrong_vars_list) == 0:
+                wrong_vars_list = tmp_vars_list
+            else:
+                for i in range(len(tmp_vars_list)):
+                    vars_list = tmp_vars_list[i]
+                    for key in vars_list.keys():
+                        if key not in wrong_vars_list[i].keys():
+                            wrong_vars_list[i][key] = vars_list[key]
+
+            for key in tmp_dist_attr_map.keys():
+                if key not in wrong_dist_attr_map:
+                    wrong_dist_attr_map[key] = tmp_dist_attr_map[key]
+
+        right_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
+            right_vars_list, right_dist_attr_map, right_dist_attr_map
+        )
+        wrong_tensors_dict = AutoAlignTool.convert_src_tensor_2_dst_tensor(
+            wrong_vars_list, wrong_dist_attr_map, right_dist_attr_map
+        )
+        diff_var_name_list = AutoAlignTool.find_diff_vars(
+            right_tensors_dict, wrong_tensors_dict
+        )
+
+        diff_ops_varname_dict = collections.OrderedDict()
+
+        for program in wrong_program_list:
+            for block in program.blocks:
+                for op in block.ops:
+                    for varname in op.input_arg_names + op.output_arg_names:
+                        if varname in diff_var_name_list:
+                            if len(diff_ops_varname_dict) == 0:
+                                print(
+                                    "first different op:\n",
+                                    op,
+                                    f"\ndifferent varname is:{varname}",
+                                )
+                            if op not in diff_ops_varname_dict:
+                                diff_ops_varname_dict[op] = [varname]
+                            else:
+                                diff_ops_varname_dict[op].append(varname)
+
+        return diff_ops_varname_dict

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

@@ -0,0 +1,244 @@
+# 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 os
+import time
+
+import paddle
+from paddle.hapi.callbacks import (
+    Callback,
+    CallbackList,
+    LRScheduler,
+    ModelCheckpoint,
+    ProgBarLogger,
+)
+
+from ..interface import CollectionNames, get_collection
+
+
+def config_callbacks(
+    callbacks=None,
+    engine=None,
+    batch_size=None,
+    epochs=None,
+    steps=None,
+    log_freq=2,
+    verbose=2,
+    save_freq=1,
+    save_dir=None,
+    metrics=None,
+    acc_step=1,
+    mode='train',
+):
+    cbks = callbacks or []
+    cbks = cbks if isinstance(cbks, (list, tuple)) else [cbks]
+
+    if not any(isinstance(k, ProgBarLogger) for k in cbks) and verbose:
+        cbks = [ProgBarLoggerAuto(log_freq, verbose=verbose), *cbks]
+
+    if not any(isinstance(k, LRScheduler) for k in cbks):
+        cbks = [LRSchedulerAuto(), *cbks]
+
+    if not any(isinstance(k, ModelCheckpoint) for k in cbks):
+        cbks = [*cbks, ModelCheckpointAuto(save_freq, save_dir)]
+
+    if not any(isinstance(k, Profiler) for k in cbks) and verbose == 3:
+        cbks = [*cbks, Profiler(timer_only=True)]
+
+    if not any(isinstance(k, History) for k in cbks):
+        cbks = [*cbks, History()]
+
+    for i, k in enumerate(cbks):
+        if isinstance(k, ProgBarLogger):
+            cbks[i] = ProgBarLoggerAuto(k.log_freq, k.verbose)
+        if isinstance(k, LRScheduler):
+            cbks[i] = LRSchedulerAuto(k.by_step, k.by_epoch)
+        if isinstance(k, ModelCheckpoint):
+            cbks[i] = ModelCheckpointAuto(k.save_freq, k.save_dir)
+
+    cbk_list = CallbackList(cbks)
+    cbk_list.set_model(engine)
+    metrics = metrics or [] if mode != 'test' else []
+    params = {
+        'batch_size': batch_size,
+        'epochs': epochs,
+        'steps': steps,
+        'verbose': verbose,
+        'metrics': metrics,
+        'acc_step': acc_step,
+    }
+    cbk_list.set_params(params)
+    return cbk_list
+
+
+class ProgBarLoggerAuto(ProgBarLogger):
+    def __init__(self, log_freq=1, verbose=2):
+        super().__init__(log_freq, verbose)
+
+    def _is_print(self):
+        return True
+
+    def _updates(self, logs, mode):
+        values = []
+        metrics = getattr(self, f'{mode}_metrics')
+        progbar = getattr(self, f'{mode}_progbar')
+        steps = getattr(self, f'{mode}_step')
+
+        for k in metrics:
+            if k in logs:
+                values.append((k, logs[k]))
+
+        if 'lr' in logs:
+            values.append(('lr', logs['lr']))
+
+        fetches_logs = logs.get('fetches', {})
+        collect_logging = get_collection(CollectionNames.LOGGING)
+        for name, var in collect_logging:
+            k = name or var.name
+            if k in fetches_logs:
+                values.append((k, fetches_logs[k]))
+
+        out_logs = logs.get('outputs', {})
+        for k in out_logs:
+            values.append((k, out_logs[k]))
+
+        if self.verbose == 3 and hasattr(self, f'_{mode}_timer'):
+            timer = getattr(self, f'_{mode}_timer')
+            cnt = timer['count'] if timer['count'] > 0 else 1.0
+            samples = timer['samples'] if timer['samples'] > 0 else 1.0
+            values.append(
+                ('avg_reader_cost', "%.5f sec" % (timer['data_time'] / cnt))
+            )
+            values.append(
+                ('avg_batch_cost', "%.5f sec" % (timer['batch_time'] / cnt))
+            )
+            values.append(
+                (
+                    'ips',
+                    "%.5f samples/sec"
+                    % (samples / (timer['data_time'] + timer['batch_time'])),
+                )
+            )
+            timer['count'] = 0
+            timer['samples'] = 0
+            timer['data_time'] = 0.0
+            timer['batch_time'] = 0.0
+
+        progbar.update(steps, values)
+
+    def on_eval_batch_end(self, step, logs=None):
+        logs = logs or {}
+        self.eval_step += 1
+        samples = self.params['batch_size']
+        self.evaled_samples += samples
+
+        self._eval_timer['batch_time'] += (
+            time.time() - self._eval_timer['batch_data_end_time']
+        )
+        self._eval_timer['count'] += 1
+        samples = self.params['batch_size']
+        self._eval_timer['samples'] += samples
+
+        if self._is_print() and self.eval_step % self.log_freq == 0:
+            if self.eval_steps is None or self.eval_step < self.eval_steps:
+                self._updates(logs, 'eval')
+
+        self._eval_timer['batch_start_time'] = time.time()
+
+
+class LRSchedulerAuto(LRScheduler):
+    def __init__(self, by_step=True, by_epoch=False):
+        super().__init__(by_step, by_epoch)
+
+    def on_epoch_begin(self, epoch=None, logs=None):
+        self.acc_step = self.params["acc_step"]
+        self.epoch = epoch
+        self.train_step = 0
+
+    def on_train_batch_end(self, step, logs=None):
+        self.train_step += 1
+
+        if self.by_step and self.train_step % self.acc_step == 0:
+            if (
+                self.model.optimizer
+                and hasattr(self.model.optimizer, '_learning_rate')
+                and isinstance(
+                    self.model.optimizer._learning_rate,
+                    paddle.optimizer.lr.LRScheduler,
+                )
+            ):
+                self.model.optimizer._learning_rate.step()
+
+
+class History(Callback):
+    def __init__(self):
+        self.history = {}
+
+    def on_train_begin(self, logs=None):
+        self.epoch = []
+
+    def on_epoch_end(self, epoch, logs=None):
+        logs = logs or {}
+        self.epoch.append(epoch)
+        for k, v in logs.items():
+            self.history.setdefault(k, []).append(v)
+
+        self.model.history = self
+
+
+class Profiler(Callback):
+    def __init__(self, *args, **kwargs):
+        self.prof = paddle.profiler.Profiler(*args, **kwargs)
+
+    def on_epoch_begin(self, epoch=None, logs=None):
+        self.epoch = epoch
+        self.train_step = 0
+        self.batch_size = self.params["batch_size"]
+        self.steps = self.params['steps']
+
+    def on_train_begin(self, logs=None):
+        self.prof.start()
+
+    def on_train_batch_end(self, step, logs=None):
+        self.train_step += 1
+        self.prof.step(num_samples=self.batch_size)
+        print(
+            "step {}:{}".format(
+                self.train_step, self.prof.step_info(unit='samples')
+            )
+        )
+
+    def on_train_end(self, logs=None):
+        self.prof.stop()
+        self.prof.summary()
+
+
+class ModelCheckpointAuto(ModelCheckpoint):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _is_save(self):
+        return self.model and self.save_dir
+
+    def on_epoch_end(self, epoch, logs=None):
+        if self._is_save() and (self.epoch + 1) % self.save_freq == 0:
+            path = f'{self.save_dir}/epoch{epoch}'
+            print(f'save checkpoint at {os.path.abspath(path)}')
+            self.model.save(path)
+
+    def on_train_end(self, logs=None):
+        if self._is_save():
+            path = f'{self.save_dir}/final'
+            print(f'save checkpoint at {os.path.abspath(path)}')
+            self.model.save(path)

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

@@ -0,0 +1,129 @@
+# 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 enum import IntEnum, unique
+
+import numpy as np
+
+from paddle.framework import core
+
+
+@unique
+class DeviceType(IntEnum):
+    UNKNOWN = 0
+    CPU = 1
+    GPU = 2
+    XPU = 3
+    DCU = 5
+    NIC = 6
+
+
+@unique
+class LinkType(IntEnum):
+    UNKNOWN = 0
+    LOC = 1
+    SYS = 2
+    PHB = 3
+    PIX = 4
+    PIB = 5
+    NVL = 6
+    NVB = 7
+    NET = 8
+
+
+class DeviceMesh(core.DeviceMesh):
+    r"""
+    The class `DeviceMesh` describes the topology of physical devices.
+
+    Args:
+        mesh (list|numpy.array): an N-dimensional array describes the topology
+            of logical processes.
+        dim_names (list, optional): the i-th element of this list gives the name of the
+            i-th dimension.
+
+    Returns:
+        None
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> paddle.enable_static()
+
+            >>> mesh = dist.DeviceMesh([[2, 4, 5], [0, 1, 3]])
+            >>> assert mesh.shape == [2, 3]
+            >>> assert mesh.device_ids == [2, 4, 5, 0, 1, 3]
+
+    """
+
+    def __init__(self, name, mesh, dim_names=None):
+        self._name = name
+
+        if not isinstance(mesh, list) and not isinstance(mesh, np.ndarray):
+            raise ValueError(
+                'The mesh must be an instance of list or np.ndarray.'
+            )
+        if isinstance(mesh, list):
+            mesh = np.array(mesh)
+
+        self._mesh = mesh
+
+        self._shape = list(self._mesh.shape)
+
+        self._device_ids = self._mesh.flatten().tolist()
+        assert all(isinstance(p, int) for p in self._device_ids), (
+            "All elements of the mesh be integer"
+        )
+        assert min(self._device_ids) >= 0, (
+            'All elements of the mesh must be >= 0.'
+        )
+        unique_device_ids = set(self._device_ids)
+        assert len(unique_device_ids) == len(self._device_ids), (
+            'All elements of the mesh must be unique.'
+        )
+
+        if dim_names is not None:
+            assert len(dim_names) == len(self._shape), (
+                "The length of dims_names must be same as the shape of the mesh."
+            )
+            self._dim_names = dim_names
+        else:
+            self._dim_names = ["d" + str(i) for i in range(len(self._shape))]
+
+        # Follow the requirement for using pybind11
+        core.DeviceMesh.__init__(
+            self, self._name, self._shape, self._device_ids, self._dim_names
+        )
+
+    @property
+    def mesh(self):
+        return self._mesh
+
+
+# class Cluster:
+#     """
+#     The cluster represents the hardware resource.
+#     """
+
+#     def __init__(self):
+#         self._device_meshes = {}
+
+#     def device_mesh(self, device_mesh_name):
+#         return self._device_meshes[device_mesh_name]
+
+#     def add_device_mesh(self, device_mesh):
+#         self._device_meshes[device_mesh.name] = device_mesh

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

@@ -0,0 +1,543 @@
+# 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 logging
+import warnings
+
+import numpy as np
+
+import paddle
+
+from ...utils.log_utils import get_logger
+
+
+class Converter:
+    """
+    Converter is a class object for auto parallel to convert tensors from
+    one parallel strategy to another one. Tensors will merge and slice value
+    with their strategy when strategies are different.
+    """
+
+    def __init__(self, tensors_dict, pre_strategy, cur_strategy):
+        """
+        Args:
+            tensors_dict(dict): tensors' value of all ranks that to be converted.
+                key is tensor's name(str), value is all ranks' data(list(numpy.ndarray))
+            pre_strategy(dict): tensors' distributed attribute of last training process.
+                key is tensor's name(str), value is tensor's distributed attribute in last
+                training process.
+            cur_strategy(dict): tensors' distributed attribute of current rank.
+                key is tensor's name(str), value is tensor's distributed attribute in current
+                rank.
+        """
+        self._tensors_dict = self._check_tensor_dict(tensors_dict)
+        self._pre_strategy = self._check_pre_strategy(pre_strategy)
+        self._cur_strategy = self._check_cur_strategy(cur_strategy)
+        self._logger = get_logger(logging.INFO)
+
+    def _check_tensor_dict(self, tensors_dict):
+        if not tensors_dict:
+            raise ValueError(
+                "'tensors_dict' is None, "
+                "the tensors to be converted cannot be None."
+            )
+        if not isinstance(tensors_dict, dict):
+            raise TypeError(
+                f"The type of 'tensors_dict' should be 'dict', but got '{type(tensors_dict)}'."
+            )
+        return tensors_dict
+
+    def _check_pre_strategy(self, pre_strategy):
+        if not pre_strategy:
+            raise ValueError(
+                "'pre_strategy' is None, there are not tensors in pre process."
+            )
+        if not isinstance(pre_strategy, dict):
+            raise TypeError(
+                "The type of 'pre_strategy' should be 'dict', "
+                f"but got '{type(pre_strategy)}'."
+            )
+        return pre_strategy
+
+    def _check_cur_strategy(self, cur_strategy):
+        if not cur_strategy:
+            warnings.warn(
+                "'cur_strategy' is None, there are not tensors in cur process"
+            )
+        if not isinstance(cur_strategy, dict):
+            raise TypeError(
+                "The type of 'cur_strategy' should be 'dict', "
+                f"but got '{type(cur_strategy)}'."
+            )
+        return cur_strategy
+
+    def convert(self, strict=True):
+        """
+        Convert tensors
+
+        Args:
+            strict(bool): whether to strict convert tensor with tensor's name. If False, it will
+            convert tensors by prefix matching. Otherwise, tensors will be converted with
+            their name strictly.
+
+        Returns:
+            converted tensors(dict)
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                >>> import numpy as np
+                >>> from paddle.distributed.auto_parallel.static.converter import Converter
+                >>> complete_tensors = np.arange(4).reshape([2, 2])
+                >>> partial_tensors = np.split(complete_tensors, 2, axis=0)
+                >>> name = "tmp_0"
+                >>> tensors_dict = {name: partial_tensors}
+                >>> strategy_1 = {
+                ...     name: {
+                ...         "process_shape": [2],
+                ...         "process_group": [0, 1],
+                ...         "dims_mapping": [0, -1],
+                ...     },
+                ... }
+                >>> strategy_2 = {
+                ...     name: {
+                ...         "process_shape": [2],
+                ...         "process_group": [0, 1],
+                ...         "dims_mapping": [-1, -1],
+                ...     },
+                ... }
+                >>> converter = Converter(tensors_dict, strategy_1, strategy_2)
+                >>> result = converter.convert()
+                >>> # the result's value is equal to `complete_tensors`
+        """
+        tensors_dict = {}
+        # the name which is in cur_process but not in pre_process
+        tensor_not_in_pre = []
+        # the name which is in pre_process but not in cur_process
+        tensor_not_in_cur = []
+        # the name which is in strategy but not in ckpt files
+        tensor_not_in_ckpt = []
+        self._logger.info("Start to convert tensors.")
+        for tensor_name in self._cur_strategy:
+            if tensor_name not in self._pre_strategy:
+                tensor_not_in_pre.append(tensor_name)
+                continue
+            if tensor_name not in self._tensors_dict:
+                tensor_not_in_ckpt.append(tensor_name)
+                continue
+            self._pre_name = tensor_name
+            self._cur_name = tensor_name
+            tensor_list = self._tensors_dict[tensor_name]
+            pre_dist_attr = self._pre_strategy[tensor_name]
+            cur_dist_attr = self._cur_strategy[tensor_name]
+            try:
+                tensors_dict[tensor_name] = Converter.merge_and_slice(
+                    tensor_list, pre_dist_attr, cur_dist_attr
+                )
+            except ValueError as err:
+                raise ValueError(
+                    f"Fail to convert tensor '{tensor_name}'. {err}"
+                )
+
+        for tensor_name in self._pre_strategy:
+            if tensor_name not in self._cur_strategy:
+                tensor_not_in_cur.append(tensor_name)
+
+        if not strict:
+            (
+                tensors_dict,
+                tensor_match_with_pre,
+                tensor_match_with_cur,
+            ) = self.convert_with_prefix_match(
+                tensors_dict, tensor_not_in_pre, tensor_not_in_cur
+            )
+        else:
+            tensors_dict, tensor_match_with_pre, tensor_match_with_cur = (
+                tensors_dict,
+                [],
+                [],
+            )
+
+        tensor_not_in_pre = set(tensor_not_in_pre) - set(tensor_match_with_pre)
+        tensor_not_in_cur = set(tensor_not_in_cur) - set(tensor_match_with_cur)
+        if tensor_not_in_pre:
+            warnings.warn(
+                f"tensors [{tensor_not_in_pre}] are not found in last training strategy."
+            )
+        if tensor_not_in_cur:
+            warnings.warn(
+                f"tensors [{tensor_not_in_cur}] are not found in current training strategy."
+            )
+        if tensor_not_in_ckpt:
+            warnings.warn(
+                f"tensors [{tensor_not_in_ckpt}] are found in pre_strategy, but are not found"
+                "in checkpoint files, please check your checkpoint files."
+            )
+
+        return tensors_dict
+
+    def convert_with_prefix_match(
+        self, tensors_dict, tensor_not_in_pre, tensor_not_in_cur
+    ):
+        # the name which in cur_process and can match with pre_process
+        tensor_match_with_pre = []
+        # the name which in pre_process and can match with cur_process
+        tensor_match_with_cur = []
+        for cur_name in tensor_not_in_pre:
+            prefix_name = cur_name
+            while prefix_name.find("_") != -1:
+                prefix_name = prefix_name[: prefix_name.rfind("_")]
+                for pre_name in tensor_not_in_cur:
+                    if prefix_name in pre_name:
+                        # 'cur_name' of cur_process can match with 'pre_name' of pre_process
+                        self._pre_name = pre_name
+                        self._cur_name = cur_name
+                        pre_tensor_list = self._tensors_dict[pre_name]
+                        pre_dist_attr = self._pre_strategy[pre_name]
+                        cur_dist_attr = self._cur_strategy[cur_name]
+                        try:
+                            tensors_dict[cur_name] = Converter.merge_and_slice(
+                                pre_tensor_list, pre_dist_attr, cur_dist_attr
+                            )
+                        except ValueError as err:
+                            raise ValueError(
+                                f"Fail to convert tensor '{cur_name}' by '{pre_name}'. {err}"
+                            )
+                        self._logger.info(
+                            f"tensor [{cur_name}] is matched with tensor [{pre_name}]"
+                        )
+                        tensor_match_with_pre.append(cur_name)
+                        tensor_match_with_cur.append(pre_name)
+                        break
+                break
+
+        return tensors_dict, tensor_match_with_pre, tensor_match_with_cur
+
+    @staticmethod
+    def merge_and_slice(tensor_list, pre_dist_attr, cur_dist_attr):
+        """
+        Merge tensors with previous dist_attr and slice tensors with current dist_attr
+
+        Returns:
+            tensor(numpy.narray): a tensor's value of current rank.
+        """
+        assert isinstance(tensor_list, list)
+        assert all(isinstance(p, np.ndarray) for p in tensor_list)
+
+        if pre_dist_attr == cur_dist_attr:
+            # skip merge and slice tensor
+            rank_id = paddle.distributed.get_rank()
+            index = cur_dist_attr["process_group"].index(rank_id)
+            tensor = tensor_list[index]
+        else:
+            pre_dims_mapping = pre_dist_attr["dims_mapping"]
+            cur_dims_mapping = cur_dist_attr["dims_mapping"]
+
+            if len(pre_dims_mapping) and (
+                len(set(pre_dims_mapping)) > 1 or -1 not in pre_dims_mapping
+            ):
+                # merge tensor
+                tensor = Converter.merge_with_dist_attr(
+                    tensor_list, pre_dist_attr
+                )
+            else:
+                # skip merge tensor
+                tensor = tensor_list[0]
+
+            if len(cur_dims_mapping) and (
+                len(set(cur_dims_mapping)) > 1 or -1 not in cur_dims_mapping
+            ):
+                # slice tensor
+                tensor = Converter.slice_with_dist_attr(tensor, cur_dist_attr)
+
+        return tensor
+
+    @staticmethod
+    def merge_with_dist_attr(tensor_list, dist_attr):
+        """Merge tensor with distributed attribute"""
+        from .reshard import Resharder
+
+        dims_mapping = dist_attr["dims_mapping"]
+        process_shape = dist_attr["process_shape"]
+        process_group = dist_attr["process_group"]
+        # get the complete shape of the tensor
+        complete_shape = Resharder.compute_complete_shape(
+            tensor_list[0].shape, process_shape, dims_mapping
+        )
+        # merge the tensor with dist_attr
+        partition_tensor_list = []
+        merged_partition = []
+        for process in process_group:
+            partition_index = Resharder.compute_partition_index(
+                process,
+                complete_shape,
+                dims_mapping,
+                process_shape,
+                process_group,
+            )
+            index = process_group.index(process)
+            if partition_index not in merged_partition:
+                merged_partition.append(partition_index)
+                Converter.merge(
+                    partition_tensor_list,
+                    tensor_list[index],
+                    partition_index,
+                    complete_shape,
+                )
+
+        if len(partition_tensor_list) != 1:
+            raise ValueError(
+                f"Fail to merge tensor with dist_attr '{dist_attr}'."
+            )
+        complete_tensor = partition_tensor_list[0][0]
+        return complete_tensor
+
+    @staticmethod
+    def slice_with_dist_attr(tensor, dist_attr):
+        """Slice tensor with distributed attribute"""
+        dims_mapping = dist_attr["dims_mapping"]
+        if len(dims_mapping) == 0:
+            # NOTE: scalar tensor no need to split
+            return tensor
+        process_shape = dist_attr["process_shape"]
+        process_group = dist_attr["process_group"]
+        # slice the tensor with dist_attr
+        partition_index_list = Converter._get_split_indices(
+            tensor.shape, dims_mapping, process_shape, process_group
+        )
+        sliced_tensor_list = Converter.split(
+            tensor, partition_index_list, len(partition_index_list)
+        )
+        # get the current tensor's index in sliced_tensor_list
+        rank_id = paddle.distributed.get_rank()
+        sliced_tensor_index = Converter._get_sliced_index(
+            rank_id, tensor.shape, dims_mapping, process_shape, process_group
+        )
+        if sliced_tensor_index not in range(len(sliced_tensor_list)):
+            raise ValueError(
+                f"Fail to slice tensor with dist_attr '{dist_attr}'."
+            )
+        sliced_tensor = sliced_tensor_list[sliced_tensor_index]
+        return sliced_tensor
+
+    @staticmethod
+    def merge(partition_tensor_list, tensor, partition_index, complete_shape):
+        """
+        Merge partial tensors to a complete.
+
+        Returns:
+            None
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                >>> import numpy as np
+                >>> import paddle
+                >>> from paddle.distributed.auto_parallel.static.converter import Converter
+                >>> partition_tensor_list = [(np.array([[[1.11, 1.12]]]), [[0, 1], [0, 1], [0, 2]])]
+                >>> tensor = np.array([[[1.13, 1.14]]])
+                >>> partition_index = [[0, 1], [0, 1], [2, 4]]
+                >>> complete_shape = [3, 2]
+
+                >>> Converter.merge(partition_tensor_list, tensor, partition_index, complete_shape)
+                >>> print(partition_tensor_list)
+                [(array([[[1.11, 1.12, 1.13, 1.14]]]), [[0, 1], [0, 1], [0, 4]])]
+        """
+        from .reshard import Resharder
+
+        if len(partition_tensor_list) == 1:
+            is_complete_data = True
+            for idx, item in enumerate(partition_tensor_list[0][1]):
+                if item[0] != 0 or item[1] != complete_shape[idx]:
+                    is_complete_data = False
+                    break
+            if is_complete_data:
+                return
+
+        if not partition_tensor_list:
+            partition_tensor_list.append((tensor, partition_index))
+        else:
+            i = 0
+            while i < len(partition_tensor_list):
+                (
+                    concat_axis,
+                    first_order,
+                    new_partition,
+                ) = Resharder.compute_concat_info(
+                    partition_tensor_list[i][1], partition_index
+                )
+                if concat_axis != -1:
+                    if first_order == 0:
+                        new_tensor = np.concatenate(
+                            (partition_tensor_list[i][0], tensor),
+                            axis=concat_axis,
+                        )
+                    else:
+                        new_tensor = np.concatenate(
+                            (tensor, partition_tensor_list[i][0]),
+                            axis=concat_axis,
+                        )
+
+                    partition_tensor_list.pop(i)
+                    Converter.merge(
+                        partition_tensor_list,
+                        new_tensor,
+                        new_partition,
+                        complete_shape,
+                    )
+                    break
+                i += 1
+
+    @staticmethod
+    def split(complete_tensor, partition_index_list, length):
+        """
+        Slice a complete tensor.
+
+        Returns:
+            sliced_tensor_list(list): sliced tensors with 'partition_index_list'
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                >>> import numpy as np
+                >>> from paddle.distributed.auto_parallel.static.converter import Converter
+                >>> complete_tensor = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
+                >>> rank = 2
+                >>> complete_shape = [1, 1, 6]
+                >>> dims_mapping = [-1, -1, 0]
+                >>> process_shape = [3]
+                >>> process_group = [0, 1, 2]
+
+                >>> sliced_tensor_list = Converter.split(complete_tensor, [[], [], [2, 4]], 3)
+                >>> print(sliced_tensor_list)
+                [array([[[1.11, 1.12]]]), array([[[1.13, 1.14]]]), array([[[1.15, 1.16]]])]
+        """
+        sliced_tensor_list = []
+        axis = len(complete_tensor.shape) - length
+        sliced_tensor = np.split(
+            complete_tensor, partition_index_list[axis], axis=axis
+        )
+        if length == 1:
+            return sliced_tensor
+        for tensor in sliced_tensor:
+            sliced_tensor_list.extend(
+                Converter.split(tensor, partition_index_list, length - 1)
+            )
+        return sliced_tensor_list
+
+    @staticmethod
+    def _get_split_indices(
+        complete_shape, dims_mapping, process_shape, process_group
+    ):
+        """
+        Get split indices of every dimension.
+
+        Returns:
+            split_indices_list(list): the split indices of every dimension of the tensor
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                >>> import numpy as np
+                >>> from paddle.distributed.auto_parallel.static.utils import _get_split_indices
+                >>> complete_tensor = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
+                >>> complete_shape = [1, 1, 6]
+                >>> dims_mapping = [-1, -1, 0]
+                >>> process_shape = [3]
+                >>> process_group = [0, 1, 2]
+
+                >>> index = _get_split_indices(complete_shape, dims_mapping, process_shape, process_group)
+                >>> print(index)
+                [[], [], [2, 4]]
+        """
+        from .reshard import Resharder
+
+        split_indices_list = []
+        for process in process_group:
+            partition_index = Resharder.compute_partition_index(
+                process,
+                complete_shape,
+                dims_mapping,
+                process_shape,
+                process_group,
+            )
+            if split_indices_list:
+                for dim in range(len(partition_index)):
+                    split_indices_list[dim].extend(partition_index[dim])
+            else:
+                split_indices_list = partition_index
+        split_indices_list = list(
+            map(
+                lambda x, y: list(set(x) - {y} - {0}),
+                split_indices_list,
+                complete_shape,
+            )
+        )
+        split_indices_list = [sorted(x) for x in split_indices_list]
+        return split_indices_list
+
+    @staticmethod
+    def _get_sliced_index(
+        rank_id, complete_shape, dims_mapping, process_shape, process_group
+    ):
+        """
+        Get sliced_tensor's index of current rank in all sliced tensors list.
+
+        Returns:
+            sliced_tensor_index(int): the index of sliced tensor in sliced_tensor_list
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+                >>> import numpy as np
+                >>> from paddle.distributed.auto_parallel.static.converter import Converter
+                >>> complete_tensor = np.array([[[1.11, 1.12, 1.13, 1.14, 1.15, 1.16]]])
+                >>> rank = 2
+                >>> complete_shape = [1, 1, 6]
+                >>> dims_mapping = [-1, -1, 0]
+                >>> process_shape = [3]
+                >>> process_group = [0, 1, 2]
+
+                >>> index = Converter._get_sliced_index(
+                ...     rank,
+                ...     complete_shape,
+                ...     dims_mapping,
+                ...     process_shape,
+                ...     process_group,
+                ... )
+                >>> print(index)
+                2
+        """
+        from .reshard import Resharder
+
+        partition_index = Resharder.compute_partition_index(
+            rank_id, complete_shape, dims_mapping, process_shape, process_group
+        )
+        sliced_index = 0
+        for i, shape in enumerate(complete_shape):
+            if dims_mapping[i] == -1:
+                slice_shape = shape
+            else:
+                slice_shape = shape // process_shape[dims_mapping[i]]
+            if slice_shape == 1:
+                index = partition_index[i][0]
+            else:
+                index = (partition_index[i][0] + 1) // slice_shape
+            sliced_index = sliced_index * (shape // slice_shape) + index
+        return sliced_index

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

@@ -0,0 +1,62 @@
+# 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 .base_cost import (  # noqa: F401
+    CommContext,
+    Cost,
+    _g_op_cost_factory,
+    build_comm_costs_from_descs,
+    build_comm_desc,
+    build_comm_desc_from_dist_op,
+    build_comp_costs_from_descs,
+    build_comp_desc_from_dist_op,
+    build_comp_desc_str_for_predict,
+    build_dp_costs,
+    calc_time_by_cost_model,
+)
+from .comm_op_cost import (  # noqa: F401
+    AllgatherOpCost,
+    AllReduceOpCost,
+    AllreduceSumOpCost,
+    BroadcastOpCost,
+    IdentityOpCost,
+    RecvOpCost,
+    SendOpCost,
+)
+from .comp_op_cost import (  # noqa: F401
+    ConcatOpCost,
+    EmbeddingGradOpCost,
+    EmbeddingOpCost,
+    FillConstantBatchSizeLikeOpCost,
+    MatmulGradOpCost,
+    MatmulOpCost,
+    MatmulV2GradOpCost,
+    MatmulV2OpCost,
+    MulGradOpCost,
+    MulOpCost,
+    Reshape2GradOpCost,
+    Reshape2OpCost,
+    SliceOpCost,
+    SoftmaxGradOpCost,
+    SoftmaxOpCost,
+    SplitOpCost,
+    Transpose2GradOpCost,
+    Transpose2OpCost,
+)
+from .estimate_cost import CostEstimator  # noqa: F401
+from .op_runtime_cost import (  # noqa: F401
+    check_if_op_supports_runtime_profiling,
+    measure_program_real_op_cost,
+)
+from .tensor_cost import TensorCost  # noqa: F401

python/paddle/distributed/auto_parallel/static/cost/comm_op_cost.py

@@ -0,0 +1,316 @@
+# 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 math
+
+import numpy as np
+
+import paddle
+
+from .base_cost import CommOpCost, register_op_cost
+
+
+@register_op_cost
+class AllreduceSumOpCost(CommOpCost):
+    OP_TYPE = "c_allreduce_sum"
+
+    def __init__(self, op=None, op_desc=None, comm_context=None):
+        super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
+
+    def calc_time(self):
+        # use tree if cross machine and use ring if in a single machine
+        time = None
+        cluster = self.comm_context.cluster
+        if not cluster.cross_machine(self.group_ranks):
+            time = self.calc_time_ring()
+        else:
+            time = self.calc_time_tree()
+
+        return time
+
+    def calc_time_ring(self):
+        alpha = self.comm_context.base_ring
+        alpha += (
+            2
+            * (self.rank_count - self.machine_count)
+            * self.comm_context.intra_ring
+        )
+        alpha += (
+            2
+            * (self.machine_count - 1)
+            * (
+                self.comm_context.inter_ring
+                + self.hops * self.comm_context.switch
+            )
+        )
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+        time = (
+            alpha
+            + 2
+            * (self.rank_count - 1)
+            / self.rank_count
+            * self.comm_count
+            * beta
+        )
+
+        return time
+
+    def calc_time_tree(self):
+        alpha = self.comm_context.base_tree
+        alpha += (
+            2
+            * (self.rank_count / self.machine_count - 1)
+            * self.comm_context.intra_tree
+        )
+        alpha += math.log2(self.machine_count) * (
+            self.comm_context.inter_tree + self.hops * self.comm_context.switch
+        )
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+
+        time = alpha + 2 * self.comm_count * beta
+
+        return time
+
+
+@register_op_cost
+class AllReduceOpCost(CommOpCost):
+    OP_TYPE = "all_reduce"
+
+    def __init__(self, op=None, op_desc=None, comm_context=None):
+        super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
+
+    def calc_time(self):
+        # use tree if cross machine and use ring if in a single machine
+        time = None
+        cluster = self.comm_context.cluster
+        if not cluster.cross_machine(self.group_ranks):
+            time = self.calc_time_ring()
+        else:
+            time = self.calc_time_tree()
+
+        return time
+
+    def calc_time_ring(self):
+        alpha = self.comm_context.base_ring
+        alpha += (
+            2
+            * (self.rank_count - self.machine_count)
+            * self.comm_context.intra_ring
+        )
+        alpha += (
+            2
+            * (self.machine_count - 1)
+            * (
+                self.comm_context.inter_ring
+                + self.hops * self.comm_context.switch
+            )
+        )
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+        time = (
+            alpha
+            + 2
+            * (self.rank_count - 1)
+            / self.rank_count
+            * self.comm_count
+            * beta
+        )
+
+        return time
+
+    def calc_time_tree(self):
+        alpha = self.comm_context.base_tree
+        alpha += (
+            2
+            * (self.rank_count / self.machine_count - 1)
+            * self.comm_context.intra_tree
+        )
+        alpha += math.log2(self.machine_count) * (
+            self.comm_context.inter_tree + self.hops * self.comm_context.switch
+        )
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+
+        time = alpha + 2 * self.comm_count * beta
+
+        return time
+
+    @property
+    def comm_count(self):
+        from ..reshard import get_var_with_recursion
+
+        if self._comm_count is None:
+            dtype = None
+            shape = None
+            if self.op is not None:
+                vars = self.op.block.vars
+                try:
+                    var_name = self.op.input("x")[0]
+                except:
+                    var_name = self.op.output("out")[0]
+                var = get_var_with_recursion(
+                    var_name, self.op.block, self.op.block.program
+                )
+                dtype = var.dtype
+                shape = var.shape
+            elif self.op_desc is not None:
+                dtype = self.op_desc["inputs"]["x"][0][0]
+                shape = self.op_desc["inputs"]["x"][0][1]
+
+            factor = None
+            if dtype == paddle.float32 or dtype == paddle.int32:
+                factor = 4
+            else:
+                raise ValueError(f"Unsupported comm dtype {dtype}")
+            comm_count = int(np.prod(shape)) * factor
+            self._comm_count = comm_count
+
+        return self._comm_count
+
+
+@register_op_cost
+class AllgatherOpCost(CommOpCost):
+    OP_TYPE = "all_gather"
+
+    def __init__(self, op=None, op_desc=None, comm_context=None):
+        super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
+
+    def calc_time(self):
+        time = self.calc_time_ring()
+        return time
+
+    def calc_time_ring(self):
+        alpha = self.comm_context.base_ring
+        alpha += (
+            self.rank_count - self.machine_count
+        ) * self.comm_context.intra_ring
+        alpha += (self.machine_count - 1) * (
+            self.comm_context.inter_ring + self.hops * self.comm_context.switch
+        )
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+        time = (
+            alpha
+            + (self.rank_count - 1) / self.rank_count * self.comm_count * beta
+        )
+        return time
+
+    @property
+    def comm_count(self):
+        from ..reshard import get_var_with_recursion
+
+        if self._comm_count is None:
+            dtype = None
+            shape = None
+            if self.op is not None:
+                vars = self.op.block.vars
+                try:
+                    var_name = self.op.input("x")[0]
+                except:
+                    var_name = self.op.output("out")[0]
+                var = get_var_with_recursion(
+                    var_name, self.op.block, self.op.block.program
+                )
+                dtype = var.dtype
+                shape = var.shape
+            elif self.op_desc is not None:
+                dtype = self.op_desc["inputs"]["X"][0][0]
+                shape = self.op_desc["inputs"]["X"][0][1]
+
+            factor = None
+            if dtype == paddle.float32 or dtype == paddle.int32:
+                factor = 4
+            else:
+                raise ValueError(f"Unsupported comm dtype {dtype}")
+            comm_count = int(np.prod(shape)) * factor
+            self._comm_count = comm_count
+
+        return self._comm_count
+
+
+@register_op_cost
+class BroadcastOpCost(CommOpCost):
+    OP_TYPE = "broadcast"
+
+    def __init__(self, op=None, op_desc=None, comm_context=None):
+        super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
+
+    def calc_time(self):
+        time = self.calc_time_ring()
+        return time
+
+    def calc_time_ring(self):
+        alpha = self.comm_context.base_ring
+        if self.machine_count > 1:
+            alpha += (
+                self.comm_context.inter_ring
+                + self.hops * self.comm_context.switch
+            )
+        else:
+            alpha += self.comm_context.intra_ring
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+        time = alpha + self.comm_count * beta
+
+        return time
+
+
+@register_op_cost
+class IdentityOpCost(CommOpCost):
+    OP_TYPE = "c_identity"
+
+    def __init__(self, op=None, op_desc=None, comm_context=None):
+        super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
+
+    def calc_time(self):
+        return self.comm_count * 1 / (144 * 1e3)
+
+
+@register_op_cost
+class RecvOpCost(CommOpCost):
+    OP_TYPE = "recv_v2"
+
+    def __init__(self, op=None, op_desc=None, comm_context=None):
+        super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
+
+    def calc_time(self):
+        alpha = self.comm_context.base_ring
+        if self.machine_count > 1:
+            alpha += (
+                self.comm_context.inter_ring
+                + self.hops * self.comm_context.switch
+            )
+        else:
+            alpha += self.comm_context.intra_ring
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+        time = alpha + self.comm_count * beta
+        return time
+
+
+@register_op_cost
+class SendOpCost(CommOpCost):
+    OP_TYPE = "send_v2"
+
+    def __init__(self, op=None, op_desc=None, comm_context=None):
+        super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
+
+    def calc_time(self):
+        alpha = self.comm_context.base_ring
+        if self.machine_count > 1:
+            alpha += (
+                self.comm_context.inter_ring
+                + self.hops * self.comm_context.switch
+            )
+        else:
+            alpha += self.comm_context.intra_ring
+        beta = self.comm_context.get_max_beta(self.group_ranks)
+        time = alpha + self.comm_count * beta
+
+        return time

python/paddle/distributed/auto_parallel/static/cost/comp_op_cost.py

@@ -0,0 +1,591 @@
+# 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 .base_cost import CompOpCost, register_op_cost
+
+
+@register_op_cost
+class AdamOpCost(CompOpCost):
+    OP_TYPE = "adam"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ArgsortOpCost(CompOpCost):
+    OP_TYPE = "argsort"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class AssignOpCost(CompOpCost):
+    OP_TYPE = "assign"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class AssignValueOpCost(CompOpCost):
+    OP_TYPE = "assign_value"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class BeamSearchOpCost(CompOpCost):
+    OP_TYPE = "beam_search"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class BeamSearchDecodeOpCost(CompOpCost):
+    OP_TYPE = "beam_search_decode"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class CastOpCost(CompOpCost):
+    OP_TYPE = "cast"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ConcatOpCost(CompOpCost):
+    OP_TYPE = "concat"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class DropoutOpCost(CompOpCost):
+    OP_TYPE = "dropout"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class DropoutGradOpCost(CompOpCost):
+    OP_TYPE = "dropout_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseAddOpCost(CompOpCost):
+    OP_TYPE = "elementwise_add"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseAddGradOpCost(CompOpCost):
+    OP_TYPE = "elementwise_add_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseDivOpCost(CompOpCost):
+    OP_TYPE = "elementwise_div"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseDivGradOpCost(CompOpCost):
+    OP_TYPE = "elementwise_div_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseMulOpCost(CompOpCost):
+    OP_TYPE = "elementwise_mul"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseMulGradOpCost(CompOpCost):
+    OP_TYPE = "elementwise_mul_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseSubOpCost(CompOpCost):
+    OP_TYPE = "elementwise_sub"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ElementwiseSubGradOpCost(CompOpCost):
+    OP_TYPE = "elementwise_sub_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class EqualOpCost(CompOpCost):
+    OP_TYPE = "equal"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class EmbeddingOpCost(CompOpCost):
+    OP_TYPE = "c_embedding"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class EmbeddingGradOpCost(CompOpCost):
+    OP_TYPE = "c_embedding_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class FillConstantOpCost(CompOpCost):
+    OP_TYPE = "fill_constant"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class FillConstantBatchSizeLikeOpCost(CompOpCost):
+    OP_TYPE = "fill_constant_batch_size_like"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class FusedSoftmaxMaskUpperTriangleOpCost(CompOpCost):
+    OP_TYPE = "fused_softmax_mask_upper_triangle"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class FusedSoftmaxMaskUpperTriangleGradOpCost(CompOpCost):
+    OP_TYPE = "fused_softmax_mask_upper_triangle_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class GatherOpCost(CompOpCost):
+    OP_TYPE = "gather"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class GeluOpCost(CompOpCost):
+    OP_TYPE = "gelu"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class GeluGradOpCost(CompOpCost):
+    OP_TYPE = "gelu_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class GreaterEqualOpCost(CompOpCost):
+    OP_TYPE = "greater_equal"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class IncrementOpCost(CompOpCost):
+    OP_TYPE = "increment"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class IsEmptyOpCost(CompOpCost):
+    OP_TYPE = "is_empty"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LayerNormOpCost(CompOpCost):
+    OP_TYPE = "layer_norm"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LayerNormGradOpCost(CompOpCost):
+    OP_TYPE = "layer_norm_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LessThanOpCost(CompOpCost):
+    OP_TYPE = "less_than"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LogicalNotOpCost(CompOpCost):
+    OP_TYPE = "logical_not"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LogicalAndOpCost(CompOpCost):
+    OP_TYPE = "logical_and"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LodResetOpCost(CompOpCost):
+    OP_TYPE = "lod_reset"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LogOpCost(CompOpCost):
+    OP_TYPE = "log"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LookupTableV2OpCost(CompOpCost):
+    OP_TYPE = "lookup_table_v2"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class LookupTableV2GradOpCost(CompOpCost):
+    OP_TYPE = "lookup_table_v2_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class MatmulOpCost(CompOpCost):
+    OP_TYPE = "matmul"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class MatmulGradOpCost(CompOpCost):
+    OP_TYPE = "matmul_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class MatmulV2OpCost(CompOpCost):
+    OP_TYPE = "matmul_v2"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class MatmulV2GradOpCost(CompOpCost):
+    OP_TYPE = "matmul_v2_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class MemcpyOpCost(CompOpCost):
+    OP_TYPE = "memcpy"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class MulOpCost(CompOpCost):
+    OP_TYPE = "mul"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class MulGradOpCost(CompOpCost):
+    OP_TYPE = "mul_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class OneHotOpCost(CompOpCost):
+    OP_TYPE = "one_hot"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ReadFromArrayOpCost(CompOpCost):
+    OP_TYPE = "read_from_array"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ReduceSumOpCost(CompOpCost):
+    OP_TYPE = "reduce_sum"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ReduceSumGradOpCost(CompOpCost):
+    OP_TYPE = "reduce_sum_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class Reshape2OpCost(CompOpCost):
+    OP_TYPE = "reshape2"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class Reshape2GradOpCost(CompOpCost):
+    OP_TYPE = "reshape2_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ReduceMeanOpCost(CompOpCost):
+    OP_TYPE = "reduce_mean"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ReduceMeanGradOpCost(CompOpCost):
+    OP_TYPE = "reduce_mean_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ScaleOpCost(CompOpCost):
+    OP_TYPE = "scale"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class ShapeOpCost(CompOpCost):
+    OP_TYPE = "shape"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SliceOpCost(CompOpCost):
+    OP_TYPE = "slice"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SoftmaxOpCost(CompOpCost):
+    OP_TYPE = "softmax"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SoftmaxGradOpCost(CompOpCost):
+    OP_TYPE = "softmax_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SoftmaxWithCrossEntropyOpCost(CompOpCost):
+    OP_TYPE = "softmax_with_cross_entropy"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SoftmaxWithCrossEntropyGradOpCost(CompOpCost):
+    OP_TYPE = "softmax_with_cross_entropy_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SplitOpCost(CompOpCost):
+    OP_TYPE = "split"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class Squeeze2OpCost(CompOpCost):
+    OP_TYPE = "squeeze2"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SquareOpCost(CompOpCost):
+    OP_TYPE = "square"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SquareGradOpCost(CompOpCost):
+    OP_TYPE = "square_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class SumOpCost(CompOpCost):
+    OP_TYPE = "sum"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class TopKOpCost(CompOpCost):
+    OP_TYPE = "top_k"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class Transpose2OpCost(CompOpCost):
+    OP_TYPE = "transpose2"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class Transpose2GradOpCost(CompOpCost):
+    OP_TYPE = "transpose2_grad"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class Unsqueeze2OpCost(CompOpCost):
+    OP_TYPE = "unsqueeze2"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
+
+
+@register_op_cost
+class WriteToArrayOpCost(CompOpCost):
+    OP_TYPE = "write_to_array"
+
+    def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
+        super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)

python/paddle/distributed/auto_parallel/static/cost/estimate_cost.py

@@ -0,0 +1,671 @@
+# 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 collections import OrderedDict
+from functools import reduce
+
+import paddle
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+from ..dist_tensor import DistributedTensor
+from ..operators.common import get_distributed_operator_impl_container
+from .base_cost import Cost
+
+
+class CostEstimator:
+    _special_op_type = ["fused_attention", "fused_feedforward"]
+
+    def __init__(
+        self, program, cluster, mode="modeling", rank=None, loop_count=10
+    ):
+        self._program = program
+        self._cluster = cluster
+        self._check_mode(mode)
+        self._mode = mode
+        self._rank = rank if rank is not None else paddle.distributed.get_rank()
+        self._loop_count = loop_count
+        self._global_cost = Cost()
+        self._local_cost_mapping = {}
+        self._detailed_cost = OrderedDict()  # {`op_id`: {"reshard": [], "dist_op": [], "local_cost": local_cost}}}
+        self._bubble_time_mapping = {}
+        self._ordered_ops = []
+        self.max_memories = {}
+        self.max_memory = None
+
+    @property
+    def loop_count(self):
+        return self._loop_count
+
+    @property
+    def detailed_cost(self):
+        return self._detailed_cost
+
+    @property
+    def program(self):
+        return self._program
+
+    @property
+    def rank(self):
+        return self._rank
+
+    @property
+    def dist_context(self):
+        return self._dist_context
+
+    @property
+    def cluster(self):
+        return self._cluster
+
+    @property
+    def mode(self):
+        return self._mode
+
+    @property
+    def global_cost(self):
+        max_time = 0
+        memory = 0
+        flops = 0
+        for rank in self._local_cost_mapping:
+            cost = self._local_cost_mapping[rank]
+            if cost.time > max_time:
+                max_time = cost.time
+            memory += cost.memory
+            flops += cost.flops
+        self._global_cost.time = max_time
+        self._global_cost.memory = memory
+        self._global_cost.flops = flops
+        return self._global_cost
+
+    def local_cost(self, rank=None):
+        rank = self.rank if rank is None else rank
+        if rank not in self._local_cost_mapping:
+            self._local_cost_mapping[rank] = Cost()
+
+        return self._local_cost_mapping[rank]
+
+    def local_bubble_time(self, rank=None):
+        rank = self.rank if rank is None else rank
+        return self._bubble_time_mapping[rank]
+
+    def _check_mode(self, mode):
+        if mode not in ["modeling", "profiling"]:
+            raise ValueError(
+                f"Just support modeling and profiling, but got {mode}"
+            )
+
+    def _is_special_var_name(self, var_name):
+        special_var_name = ["lod_tensor_blocking_queue_0"]
+        if var_name in special_var_name:
+            return True
+        return False
+
+    def _estimate_core(self, dist_context, resharder, block):
+        from ..reshard import get_var_with_recursion
+
+        ops = block.ops
+        loop_count = None
+        if block.desc.id != self.program.global_block().desc.id:
+            loop_count = self.loop_count
+        else:
+            loop_count = 1
+        for i in range(loop_count):
+            for op in ops:
+                self._detailed_cost[op.desc.id()] = OrderedDict()
+                # If in the while sub block, the detail of cost is the last cost
+                detail = self._detailed_cost[op.desc.id()]
+                detail["reshard_cost"] = OrderedDict()  #
+                detail["dist_op_cost"] = []
+                if int(op.attr('op_role')) == int(OpRole.Optimize):
+                    continue
+                if op.type in [
+                    "create_py_reader",
+                    "create_double_buffer_reader",
+                    "read",
+                ]:
+                    continue
+
+                # NOTE: It does not support nested loop and just supports while op when op has sub block now.
+                if op.type == "while":
+                    while_block = self.program.blocks[op.attr("sub_block").id]
+                    self._estimate_core(dist_context, resharder, while_block)
+                    continue
+
+                for var_name in op.input_arg_names:
+                    if self._is_special_var_name(var_name):
+                        continue
+                    var = get_var_with_recursion(var_name, block, self.program)
+                    reshard_cost = resharder.get_cost(op, var, self.cluster)
+
+                    # Calc reshard cost
+                    if reshard_cost is not None:
+                        detail["reshard_cost"][var_name] = reshard_cost
+
+                        comm_costs = reshard_cost[0]
+                        local_comp_cost = reshard_cost[1]
+                        for comm_cost in comm_costs:
+                            # Time is cumulative in global cost and local cost, but memory and flops just are cumulative in global cost.
+                            # Comm sync
+                            for item in comm_cost:
+                                group_ranks, cost = item
+                                max_time = None
+                                cost_time = {}
+                                for rank in group_ranks:
+                                    rank_cost = self.local_cost(rank)
+                                    cost_time[rank] = rank_cost.time
+                                    if max_time is None:
+                                        max_time = rank_cost.time
+                                    else:
+                                        if max_time < rank_cost.time:
+                                            max_time = rank_cost.time
+
+                                for rank in group_ranks:
+                                    self.local_cost(rank).time = (
+                                        max_time + cost.time
+                                    )
+
+                                    if rank not in self._bubble_time_mapping:
+                                        self._bubble_time_mapping[rank] = 0
+
+                                    self._bubble_time_mapping[rank] += (
+                                        max_time - cost_time[rank]
+                                    )
+
+                        for rank in local_comp_cost:
+                            for comp_cost in local_comp_cost[rank]:
+                                self.local_cost(rank).time += comp_cost.time
+
+                # Calc dist op cost
+                dist_op = dist_context.get_dist_op_for_program(op)
+                if not dist_op:
+                    continue
+
+                op_dist_attr = dist_op.dist_attr
+                processes = op_dist_attr.process_mesh.process_ids
+
+                container = get_distributed_operator_impl_container(
+                    op_dist_attr.impl_type
+                )
+                dist_impl = container.impls[op_dist_attr.impl_idx]
+
+                dist_op_cost = dist_impl.calc_cost(
+                    op.attr('op_role'), dist_op, dist_context, self.cluster
+                )
+                detail["dist_op_cost"] = dist_op_cost
+
+                if dist_op_cost is None:
+                    assert (
+                        dist_op.serial_op.type in CostEstimator._special_op_type
+                    )
+                    continue
+                for item in dist_op_cost:
+                    if isinstance(item, list):
+                        # Comm sync
+                        for comm_op_cost in item:
+                            max_time = None
+                            cost_time = {}
+                            group_ranks = comm_op_cost.group_ranks
+                            for rank in comm_op_cost.group_ranks:
+                                rank_cost = self.local_cost(rank)
+                                cost_time[rank] = rank_cost.time
+                                if max_time is None:
+                                    max_time = rank_cost.time
+                                else:
+                                    if max_time < rank_cost.time:
+                                        max_time = rank_cost.time
+                            for rank in group_ranks:
+                                self.local_cost(rank).time = (
+                                    max_time + comm_op_cost.time
+                                    if op.attr('op_role') != OpRole.Backward
+                                    else max_time + 0.9 * comm_op_cost.time
+                                )
+                                if rank not in self._bubble_time_mapping:
+                                    self._bubble_time_mapping[rank] = 0
+                                self._bubble_time_mapping[rank] += (
+                                    max_time - cost_time[rank]
+                                )
+                    elif isinstance(item, dict):
+                        # Op just one
+                        for rank in processes:
+                            # DP+PP+MP
+                            if rank not in item:
+                                continue
+                            self.local_cost(rank).time += item[rank].time
+
+    def prepare(self):
+        self._global_cost = Cost()
+        self._local_cost_mapping = {}
+        self._detailed_cost = OrderedDict()
+        self._bubble_time_mapping = {}
+
+    def _calculate_bytes(self, sizes, dtype):
+        if sizes:
+            total_count = reduce(lambda x, y: x * y, sizes, 1)
+        else:
+            total_count = 0
+
+        if dtype == paddle.float64 or dtype == paddle.int64:
+            dtype_factor = 8
+        elif dtype == paddle.float32 or dtype == paddle.int32:
+            dtype_factor = 4
+        elif (
+            dtype == paddle.float16
+            or dtype == paddle.bfloat16
+            or dtype == paddle.int16
+        ):
+            dtype_factor = 2
+        elif dtype == paddle.int8 or dtype == paddle.uint8:
+            dtype_factor = 1
+        else:
+            dtype_factor = 8
+
+        memory = total_count * dtype_factor
+        return memory
+
+    def _estimate_max_memory_by_dist_op(self, dist_context):
+        # This estimation will be improved, now reshard and inplace are not considered.
+        # Persist var is not free.
+        def _convert_pm_and_dm_to_str(process_mesh, dims_mapping):
+            processes = ",".join([str(x) for x in process_mesh.process_ids])
+            topology = ",".join([str(x) for x in process_mesh.shape])
+            dims_mapping = ",".join([str(x) for x in dims_mapping])
+            result = processes + topology + dims_mapping
+            return result
+
+        memories = {}
+        self.max_memories = {}
+        var_info = {}  # var_name: [[process_mesh, dims_mapping], [id]], [[process_mesh, dims_mapping], [id]]}
+
+        for block in self.program.blocks:
+            for op in block.ops:
+                self._ordered_ops.append([op.desc.id(), op])
+        self._ordered_ops.sort(key=lambda x: x[0])
+
+        parameters = set()
+        for op_id, op in self._ordered_ops:
+            if op.type in [
+                "create_py_reader",
+                "create_double_buffer_reader",
+                "read",
+            ]:
+                continue
+            dist_op = dist_context.get_dist_op_for_program(op)
+            if not dist_op:
+                continue
+            process_mesh = dist_op.dist_attr.process_mesh
+            for var_name in op.input_arg_names:
+                input_dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
+                    var_name
+                )
+
+                if var_name not in var_info:
+                    var_info[var_name] = {}
+                key = _convert_pm_and_dm_to_str(
+                    process_mesh, input_dims_mapping
+                )
+                if key not in var_info[var_name]:
+                    var_info[var_name][key] = {}
+                # It is even partition now
+                if "position" not in var_info[var_name][key]:
+                    var_info[var_name][key]["position"] = []
+                var_info[var_name][key]["position"].append(op_id)
+
+                if "memory" not in var_info[var_name][key]:
+                    var = dist_op.get_serial_input(var_name)
+                    global_sizes = var.shape
+                    dtype = var.dtype
+                    sizes = DistributedTensor.get_local_sizes(
+                        global_sizes,
+                        input_dims_mapping,
+                        process_mesh.shape,
+                        process_mesh.process_ids,
+                    )
+                    var_info[var_name][key]["memory"] = self._calculate_bytes(
+                        sizes, dtype
+                    )
+                    if var.persistable:
+                        name = var_name + key
+                        if name not in parameters:
+                            parameters.add(name)
+                            for process in process_mesh.process_ids:
+                                if process not in memories:
+                                    memories[process] = 0
+                                memories[process] += var_info[var_name][key][
+                                    "memory"
+                                ]
+
+            for var_name in op.output_arg_names:
+                output_dims_mapping = dist_op.dist_attr.get_output_dims_mapping(
+                    var_name
+                )
+                if var_name not in var_info:
+                    var_info[var_name] = {}
+                key = _convert_pm_and_dm_to_str(
+                    process_mesh, output_dims_mapping
+                )
+                if key not in var_info[var_name]:
+                    var_info[var_name][key] = {}
+                if "position" not in var_info[var_name][key]:
+                    var_info[var_name][key]["position"] = []
+                var_info[var_name][key]["position"].append(op_id)
+
+                if "memory" not in var_info[var_name][key]:
+                    var = dist_op.get_serial_output(var_name)
+                    global_sizes = var.shape
+                    dtype = var.dtype
+                    sizes = DistributedTensor.get_local_sizes(
+                        global_sizes,
+                        output_dims_mapping,
+                        process_mesh.shape,
+                        process_mesh.process_ids,
+                    )
+                    var_info[var_name][key]["memory"] = self._calculate_bytes(
+                        sizes, dtype
+                    )
+                    if var.persistable:
+                        name = var_name + key
+                        if name not in parameters:
+                            parameters.add(name)
+                            for process in process_mesh.process_ids:
+                                if process not in memories:
+                                    memories[process] = 0
+                                memories[process] += var_info[var_name][key][
+                                    "memory"
+                                ]
+
+        has_used_vars = set()
+        not_calc_vars = set()
+        for op_id, op in self._ordered_ops:
+            if op.type in [
+                "create_py_reader",
+                "create_double_buffer_reader",
+                "read",
+            ]:
+                continue
+            can_free_memories = {}
+            can_free_vars = set()
+            dist_op = dist_context.get_dist_op_for_program(op)
+            if not dist_op:
+                continue
+            process_mesh = dist_op.dist_attr.process_mesh
+            for var_name in op.input_arg_names:
+                input_dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
+                    var_name
+                )
+                key = _convert_pm_and_dm_to_str(
+                    process_mesh, input_dims_mapping
+                )
+                has_used_var = var_name + key
+                var = dist_op.get_serial_input(var_name)
+                # Not used
+                if (
+                    has_used_var not in has_used_vars
+                    and has_used_var not in parameters
+                ):
+                    if has_used_var in not_calc_vars:
+                        continue
+                    has_used_vars.add(has_used_var)
+                    for process in process_mesh.process_ids:
+                        if process not in memories:
+                            memories[process] = 0
+                        memories[process] += var_info[var_name][key]["memory"]
+                # Used
+                if op_id == var_info[var_name][key]["position"][-1]:
+                    if (
+                        has_used_var not in can_free_vars
+                        and not var.persistable
+                    ):
+                        can_free_vars.add(has_used_var)
+                        for process in process_mesh.process_ids:
+                            if process not in can_free_memories:
+                                can_free_memories[process] = 0
+                            can_free_memories[process] += var_info[var_name][
+                                key
+                            ]["memory"]
+
+            for var_name in op.output_arg_names:
+                output_dims_mapping = dist_op.dist_attr.get_output_dims_mapping(
+                    var_name
+                )
+                key = _convert_pm_and_dm_to_str(
+                    process_mesh, output_dims_mapping
+                )
+                has_used_var = var_name + key
+                var = dist_op.get_serial_output(var_name)
+                if (
+                    op.type == "reshape2"
+                    or op.type == "transpose2"
+                    or op.type == "elementwise_add"
+                ):
+                    not_calc_vars.add(has_used_var)
+                    continue
+                # Not used
+                if (
+                    has_used_var not in has_used_vars
+                    and has_used_var not in parameters
+                ):
+                    has_used_vars.add(has_used_var)
+                    for process in process_mesh.process_ids:
+                        if process not in memories:
+                            memories[process] = 0
+                        memories[process] += var_info[var_name][key]["memory"]
+                # Used
+                if op_id == var_info[var_name][key]["position"][-1]:
+                    if (
+                        has_used_var not in can_free_vars
+                        and not var.persistable
+                    ):
+                        can_free_vars.add(has_used_var)
+                        for process in process_mesh.process_ids:
+                            if process not in can_free_memories:
+                                can_free_memories[process] = 0
+                            can_free_memories[process] += var_info[var_name][
+                                key
+                            ]["memory"]
+
+            # Calc peak memory
+            for process in memories:
+                if process not in self.max_memories:
+                    self.max_memories[process] = memories[process]
+                else:
+                    if memories[process] > self.max_memories[process]:
+                        self.max_memories[process] = memories[process]
+            # Free memory
+            for process in can_free_memories:
+                if process in memories:
+                    memories[process] -= can_free_memories[process]
+
+        # Calculate the max memory in all ranks
+        max_memory = max(self.max_memories.values())
+        self.max_memory = max_memory
+
+        return max_memory
+
+    def estimate(self, dist_context, resharder=None):
+        self.prepare()
+        from ..reshard import Resharder
+
+        resharder = (
+            Resharder(self.program, None, self.rank, dist_context, [])
+            if resharder is None
+            else resharder
+        )
+
+        block = self.program.global_block()
+        self._estimate_core(dist_context, resharder, block)
+
+        return self.global_cost
+
+    def _print_tag(self, max_len, length):
+        tag = "+" + "-" * max_len
+        for i in range(length):
+            print(tag, end="")
+            if i == length - 1:
+                print("+")
+
+    def _print_vals(self, vals, max_len):
+        for idx, val in enumerate(vals):
+            s = "|" + str(val).center(max_len)
+            print(s, end="")
+            if idx == len(vals) - 1:
+                print("|")
+
+    def _pretty_print_memory_cost(self):
+        """Print memory of every rank prettily."""
+        if not self.max_memories or not self.max_memory:
+            raise ValueError("Please calculate memory cost before print.")
+
+        # Padding automatically
+        max_len = 0
+        header = ["Rank", "Memory(MiB)"]
+        memories = [
+            int(item // 1e6) for item in list(self.max_memories.values())
+        ]
+        for memory in memories + header:
+            if len(str(memory)) > max_len:
+                max_len = len(str(memory))
+        max_len += 4  # for pretty print of center
+
+        # Print tag
+        self._print_tag(max_len, len(header))
+
+        # Print header
+        self._print_vals(header, max_len)
+
+        # Print tag
+        self._print_tag(max_len, len(header))
+
+        # Print rank and its memory
+        for i in range(len(self.max_memories)):
+            memory = memories[i]
+            vals = [i, memory]
+            self._print_vals(vals, max_len)
+            self._print_tag(max_len, len(header))
+
+    def _pretty_print_global(self):
+        """Print global execution time and max memory prettily."""
+        if not self.max_memories or not self.max_memory:
+            raise ValueError("Please calculate cost before print.")
+
+        # Padding automatically
+        max_len = 0
+        header = ["Execution Time(us)", "Max Memory(MiB)"]
+        vals = [round(self.global_cost.time, 3), int(self.max_memory // 1e6)]
+        for memory in vals + header:
+            if len(str(memory)) > max_len:
+                max_len = len(str(memory))
+        max_len += 4  # for pretty print of center
+
+        # Print tag
+        self._print_tag(max_len, len(header))
+
+        # Print header
+        self._print_vals(header, max_len)
+
+        # Print tag
+        self._print_tag(max_len, len(header))
+
+        # Print exec time and max memory
+        self._print_vals(vals, max_len)
+
+        # Print tag
+        self._print_tag(max_len, len(header))
+
+    def pretty_print_cost(self):
+        """Print cost prettily."""
+        print("The global execution time and max memory are as follows:")
+        self._pretty_print_global()
+        print("The memory of every rank is as follows:")
+        self._pretty_print_memory_cost()
+
+
+def get_cost_from_engine(engine, mode):
+    import copy
+
+    from ..utils import to_list
+
+    # Construct cost estimator by original main program
+    serial_main_prog = (
+        engine._fwd_main_progs[mode].clone()
+        if mode in engine._fwd_main_progs
+        else engine._orig_main_prog.clone()
+    )
+
+    serial_startup_prog = (
+        engine._fwd_dist_contexts[mode]._original_serial_main_program.clone()
+        if mode in engine._fwd_dist_contexts
+        else engine._orig_startup_prog.clone()
+    )
+    losses = (
+        to_list(engine._loss)
+        if (
+            not isinstance(engine._loss, paddle.nn.Layer)
+            and not callable(engine._loss)
+        )
+        else engine._losses
+    )
+    serial_optimizer = copy.deepcopy(engine._orig_optimizer)
+    if mode in engine._fwd_dist_contexts:
+        dist_context = copy.deepcopy(engine._fwd_dist_contexts[mode])
+    else:
+        from ..dist_context import DistributedContext
+
+        dist_context = DistributedContext(
+            serial_main_prog,
+            serial_startup_prog,
+            serial_optimizer,
+            losses,
+            {},
+            {"loss": losses},
+            engine._cluster,
+            engine._strategy,
+        )
+    from ..completion import Completer
+
+    completer = Completer(dist_context)
+    completer.complete_forward_annotation()
+    dist_context.block_state.parse_forward_blocks(
+        dist_context.serial_main_program
+    )
+
+    if mode == "eval" or mode == "predict":
+        cost_estimator = CostEstimator(serial_main_prog, engine._cluster)
+    elif mode == "train":
+        from ..parallelizer_v2 import Parallelizer
+
+        # Get serial main program with backward
+        parallelizer = Parallelizer(mode, completer, dist_context)
+        # Generate backward
+        loss_name = dist_context.serial_loss.name
+        serial_loss = serial_main_prog.global_block()._var_recursive(loss_name)
+        params_grads = parallelizer._generate_backward(
+            serial_main_prog, serial_startup_prog, serial_loss
+        )
+
+        # Generate optimizer
+        optimizer_ops = parallelizer._generate_optimizer(
+            serial_main_prog,
+            serial_startup_prog,
+            serial_optimizer,
+            params_grads,
+        )
+        cost_estimator = CostEstimator(serial_main_prog, engine._cluster)
+
+    # Estimate global_cost and  max memory
+    global_cost = cost_estimator.estimate(dist_context)
+    max_memory = cost_estimator._estimate_max_memory_by_dist_op(dist_context)
+
+    # Print the cost
+    cost_estimator.pretty_print_cost()
+
+    return global_cost, max_memory

python/paddle/distributed/auto_parallel/static/cost/op_runtime_cost.py

@@ -0,0 +1,320 @@
+# 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 logging
+import warnings
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+from paddle.base.data_feeder import convert_dtype
+from paddle.base.executor import (
+    _as_lodtensor,
+    _StandaloneExecutor,
+    check_feed_shape_type,
+)
+from paddle.base.framework import Operator, Program
+from paddle.distributed.auto_parallel.static.utils import get_logger, is_comm_op
+
+
+def check_if_op_supports_runtime_profiling(op):
+    return not is_comm_op(op)
+
+
+def _measure_program_real_op_cost_multipass(program, place, run_iters, verbose):
+    '''
+    Run op profiling for a single pass. Internal function, do not call this directly.
+    '''
+
+    # clone the program to avoid accidental change made to the vanilla program.
+    cloned_program = program.clone()
+    cloned_main_block = cloned_program.global_block()
+
+    # We will run the executor in a newly created scope, so that our
+    # executor will not pollute the global scope when running. Since
+    # we created a brand new scope, we need to manually create input
+    # tensors and network parameters and feed fake data into them.
+    scope = core.Scope()
+
+    logger = get_logger(log_level=logging.INFO)
+
+    def _analyze_graph_and_collect_all_vars_with_zero_in_degree():
+        var_in_degree = {}
+
+        def _collect_op_input_var_names(op: Operator):
+            input_var_names = []
+            for input_name in op.input_names:
+                input_var_names += op.input(input_name)
+            return input_var_names
+
+        def _collect_op_output_var_names(op: Operator):
+            output_var_names = []
+            for output_name in op.output_names:
+                output_var_names += op.output(output_name)
+            return output_var_names
+
+        def _record_op_output_vars_in_degree(in_var_names, out_var_names):
+            for out_var_name in out_var_names:
+                if out_var_name in in_var_names:
+                    # NOTE (liuchenghao): if an op's input var is its output var,
+                    # this means this var forms an in-place connection to itself,
+                    # in this situation we need to ignore this variable, this way
+                    # we can ensure that vars with zero in-degree are dangling vars
+                    # and they should be created manually before program executes.
+                    continue
+                var_in_degree[out_var_name] += 1
+
+        def _filter_vars_with_zero_in_degree_and_ignore_feed_fetch_vars():
+            filtered_vars = []
+            for var_name in var_in_degree:
+                if var_name in ['feed', 'fetch']:
+                    continue
+                if var_in_degree[var_name] == 0:
+                    filtered_vars.append(var_name)
+            return filtered_vars
+
+        for op in cloned_main_block.ops:
+            op: Operator
+            if is_comm_op(op):
+                # ignore communication op from graph, because sometimes we want to profile a sub-graph
+                # and these dangling operators will not work (no graph to communicate to/from)
+                continue
+            input_var_names, output_var_names = (
+                _collect_op_input_var_names(op),
+                _collect_op_output_var_names(op),
+            )
+            for var_name in input_var_names + output_var_names:
+                if var_name not in var_in_degree:
+                    var_in_degree[var_name] = 0
+            _record_op_output_vars_in_degree(input_var_names, output_var_names)
+        return _filter_vars_with_zero_in_degree_and_ignore_feed_fetch_vars()
+
+    def _alloc_and_fill_var(var_name):
+        supported_var_dtypes = [
+            "paddle.float16",
+            "paddle.float32",
+            "paddle.float64",
+            "paddle.int8",
+            "paddle.int16",
+            "paddle.int32",
+            "paddle.int64",
+            "paddle.bool",
+        ]
+        var = cloned_main_block.var(var_name)
+        var_shape = var.shape
+        var_dtype = var.dtype
+        assert str(var_dtype) in supported_var_dtypes, (
+            'Found unsupported variable dtype: "{}", current supported '
+            'dtype(s) is/are: [{}]. '.format(
+                str(var_dtype), ", ".join(supported_var_dtypes)
+            )
+        )
+        (
+            logger.info(
+                f'[+] var: "{var_name}", shape={var_shape}, dtype="{var_dtype}".\n'
+            )
+            if verbose
+            else None
+        )
+        np_dtype = (
+            convert_dtype(var_dtype)
+            if isinstance(var_dtype, core.VarDesc.VarType)
+            else var_dtype
+        )
+        if str(var_dtype).find('int') != -1:
+            # target variable's type is int* (uint*, int*), it is highly possible that
+            # the target variable contains indices (such as lookup_table op's input var)
+            # for safety we need to fill it with all one instead of random numbers
+            # NOTE (liuchenghao): filling with zero will generate "division by zero" error
+            # in mod ops, so filling with one seems to be the simplest way to make it work,
+            # although it is possible that for array with only one element, index "1" is
+            # invalid, that situation is very rare and we don't need to care about it now.
+            new_tensor = np.array(np.ones(var_shape)).astype(np_dtype)
+        else:
+            # target variable's type is float*, we treat it as an ordinary tensor, fill it
+            # with random gaussian numbers
+            new_tensor = np.array(np.random.randn(*var_shape)).astype(np_dtype)
+        new_tensor = _as_lodtensor(new_tensor, place, var_dtype)
+        check_feed_shape_type(var, new_tensor)
+        core.set_variable(scope, new_tensor, var_name)
+        return new_tensor
+
+    def _configure_feed_ops_and_return_feed_names():
+        """
+        configure feed op,
+        1. alloc feed op output var storage
+        2. fill feed op's input var
+        return feed var names
+        """
+
+        feed_names = []
+        has_feed_op = False
+        for op in cloned_main_block.ops:
+            if op.type == "feed":
+                has_feed_op = True
+                out_var_name = op.desc.output('Out')[0]
+                in_var_name = op.desc.input('X')[0]  # this is usually "feed"
+                input_index = op.desc.attr('col')
+                new_tensor = _alloc_and_fill_var(out_var_name)
+                core.set_feed_variable(
+                    scope, new_tensor, in_var_name, input_index
+                )
+                feed_names.append(out_var_name)
+        if not has_feed_op:
+            (
+                logger.info("WARNING: program does not have any feed op.\n")
+                if verbose
+                else None
+            )
+        return feed_names
+
+    for var_name in _analyze_graph_and_collect_all_vars_with_zero_in_degree():
+        _alloc_and_fill_var(var_name)
+    feed_names = _configure_feed_ops_and_return_feed_names()
+
+    # build a simple plan from program and run profiling
+    plan = core.Plan([core.Job("default")], {"default": cloned_program.desc})
+    exe = _StandaloneExecutor(place, plan, scope)
+
+    num_ops = len(cloned_main_block.ops)
+    prof_results = [[None for _ in range(run_iters)] for _ in range(num_ops)]
+
+    for iter_id in range(run_iters):
+        # for each iteration, run profiling and retrieve modified version of program desc
+        program_desc = exe.run_profile(feed_names)
+
+        # rebuild program object from the new program desc
+        temp_program = cloned_program.clone()
+        temp_program._rebuild_from_desc(program_desc)
+        temp_main_block = temp_program.global_block()
+
+        # collect profiling result
+        for op_id, temp_op in zip(
+            range(len(temp_main_block.ops)), temp_main_block.ops
+        ):
+            run_time_us = temp_op.dist_attr.run_time_us
+            prof_results[op_id][iter_id] = (
+                run_time_us
+                if check_if_op_supports_runtime_profiling(temp_op)
+                and run_time_us >= 0.0
+                else None
+            )
+    return prof_results
+
+
+def measure_program_real_op_cost(
+    program: paddle.static.Program,
+    run_iters: int = 8,
+    place=paddle.base.framework._current_expected_place(),
+    verbose_level: int = 0,
+):
+    '''
+    Description
+    -----------
+    Measuring real op run time (us) with respect to the given "program" and "place".
+
+    Parameters
+    -----------
+    @param program: paddle.static.Program
+        The program object waiting to be executed.
+    @param run_iters: int
+        Specify how many iterations will be run during profiling. Larger value tends
+        to give more accurate profiling result but requires more time.
+    @param place: paddle.CPUPlace | paddle.CUDAPlace
+        Where the program is going to be executed.
+    @param verbose_level: int
+        Set up verbose level during profiling. Can be set to one of the following:
+        0 = turn off, don't output anything,
+        1 = output profiling messages only,
+        2 = output profiling and debug messages.
+
+    Returns
+    -----------
+    Nothing to return. This API will write op run time directly into program object.
+    For example, to retrieve the run time for the first op in program, use:
+    >>> program.global_block().ops[0].dist_attr.run_time_us
+
+    Note
+    -----------
+    Not all ops support runtime profiling. Currently communication ops do not support
+    runtime profiling feature since their execution times rely on other ops. To check
+    if an op supports runtime profiling, use:
+    >>> check_if_op_supports_runtime_profiling(op)
+    where "op" is an instance of "paddle.base.framework.Operator".
+
+    Example
+    -----------
+    * Profiling a simple program from scratch:
+    >>> from paddle.distributed.auto_parallel.static.utils import (
+    ...     measure_program_real_op_cost,
+    ... )
+    >>> program = ...  # build your own program object here.
+    >>> measure_program_real_op_cost(
+    >>>     program, verbose_level=1
+    >>> )
+    * Profiling a program which is already embedded into an Executor or some other class instance:
+    >>> import paddle
+    >>> from paddle.distributed.auto_parallel.static.utils import (
+    ...     measure_program_real_op_cost,
+    ... )
+    >>> place: str = paddle.device.get_device()  # here we assume place = "cuda:x"
+    >>> place = paddle.CUDAPlace(int(place.split(':')[1]))
+    >>> # here "program" is an inner object that has already been built before
+    >>> measure_program_real_op_cost(program, verbose_level=1)
+    '''
+
+    assert isinstance(program, Program), (
+        f'"program" should be a instance of "paddle.base.framework.Program" but got type "{type(program).__name__}".'
+    )
+    supported_places = [
+        paddle.CUDAPlace,
+    ]
+    assert any(
+        isinstance(place, supported_place)
+        for supported_place in supported_places
+    ), (
+        f'Current place ({place}) does not support runtime profiling. "place" should be one of the following: {supported_places}.'
+    )
+    assert isinstance(run_iters, int) and run_iters >= 1, (
+        'Invalid parameter run_iters set. run_iters should be an integer >= 1.'
+    )
+    if run_iters == 1:
+        warnings.warn(
+            'run_iters was set to 1, profiling results might be inaccurate due to outliers.'
+        )
+
+    logger = get_logger(log_level=logging.INFO)
+
+    # run profiling multiple times and record op run time of each run
+    prof_results = _measure_program_real_op_cost_multipass(
+        program, place, run_iters, verbose=(verbose_level >= 2)
+    )
+    op_num = len(prof_results)
+    for op_id, op in zip(range(op_num), program.global_block().ops):
+        op_runtime_us_final = None
+        if prof_results[op_id][0] is not None:
+            op_runtime_us_final = np.median(prof_results[op_id])
+        if (
+            op_runtime_us_final is not None
+            and check_if_op_supports_runtime_profiling(op)
+        ):
+            op.dist_attr.run_time_us = op_runtime_us_final
+        (
+            logger.info(
+                f"{op_id!s:>4} {op.type!s:>32}  {op_runtime_us_final:.1f} us"
+            )
+            if verbose_level >= 1
+            else None
+        )

python/paddle/distributed/auto_parallel/static/cost/tensor_cost.py

@@ -0,0 +1,110 @@
+#   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 functools import reduce
+
+import paddle
+from paddle.distributed.auto_parallel.static.dist_tensor import (
+    DistributedTensor,
+)
+from paddle.static import Variable
+
+from .base_cost import Cost
+
+
+class TensorCost:
+    def __init__(self, tensor=None, dist_tensor=None, shape=None, dtype=None):
+        self._check_args(tensor, dist_tensor, shape, dtype)
+        self._tensor = tensor
+        self._dist_tensor = dist_tensor
+        self._shape = shape
+        self._dtype = dtype
+        self._cost = self.calc_cost()
+
+    @property
+    def tensor(self):
+        return self._tensor
+
+    @property
+    def dist_tensor(self):
+        return self._dist_tensor
+
+    @property
+    def shape(self):
+        return self._shape
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    def _check_args(self, tensor, dist_tensor, shape, dtype):
+        if tensor is not None:
+            assert shape is None and dist_tensor is None and dtype is None
+
+            if not isinstance(tensor, Variable):
+                raise TypeError(
+                    f"Please check tensor type is Variable, but got {type(tensor)}"
+                )
+
+        elif dist_tensor is not None:
+            assert tensor is None and shape is None
+            if not isinstance(dist_tensor, DistributedTensor):
+                raise TypeError(
+                    f"Please check dist_tensor type is DistributedTensor, but got {type(dist_tensor)}"
+                )
+
+        elif shape is not None:
+            assert tensor is None and dist_tensor is None and dtype is not None
+            if not isinstance(shape, (list, set)):
+                raise TypeError(
+                    f"Please check shape type is list or set, but got {type(shape)}"
+                )
+
+        elif dtype is not None:
+            assert tensor is None and dist_tensor is None and shape is not None
+
+    @property
+    def cost(self):
+        return self._cost
+
+    def calc_cost(self):
+        dtype = None
+        shape = None
+
+        if self.dist_tensor:
+            shape = self.dist_tensor.local_sizes()
+            dtype = self.dist_tensor.serial_tensor.dtype
+        elif self.tensor:
+            shape = self.tensor.shape
+            dtype = self.tensor.dtype
+        elif self.shape and self.dtype:
+            shape = self.shape
+            dtype = self.dtype
+
+        total_count = reduce(lambda x, y: x * y, shape, 1)
+
+        if dtype == paddle.float32 or dtype == paddle.int32:
+            dtype_factor = 4
+        elif dtype == paddle.int64:
+            dtype_factor = 8
+        elif dtype == paddle.uint8:
+            dtype_factor = 1
+        else:
+            dtype_factor = 2
+
+        memory = total_count * dtype_factor
+        assert memory >= 0
+        cost = Cost(memory=memory)
+
+        return cost

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

@@ -0,0 +1,855 @@
+#   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 queue
+from enum import Enum
+
+import numpy as np
+
+import paddle
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+from paddle.framework import core
+
+SUCC = 0  # successor
+PRED = 1  # predecessor
+
+
+class CostNodeType(Enum):
+    DEFAULT = 0
+    COMPUTATION = 1
+    COMMUNICATION = 2
+    VARIABLE = 3
+    MERGED = 4
+    NOP = 5
+
+
+class Cost:
+    def __init__(self):
+        self.runtime = None
+        self.static_mem = None
+        self.peak_mem = None
+
+
+class CostModelMode(Enum):
+    DEFAULT = 0
+    BENCHMARKING = 1  # costs based on trial runs
+    ANALYSIS = 2  # costs based on analysis
+    MIXED = 3
+
+
+class CostNode:
+    def __init__(self, node, node_type, id=None):
+        self.id = id
+        self.node = node
+        self.type = node_type
+        self._cost = 0
+        self.is_optim = False
+        self.is_bwd = False
+
+    @property
+    def cost(self):
+        return self._cost
+
+    @cost.setter
+    def cost(self, cost):
+        if cost < 0:
+            raise ValueError('Cost must be above 0.')
+        self._cost = cost
+
+
+class MergedOpsCostNode(CostNode):
+    def __init__(self, node_type, id=None, base_node_list=None, is_bwd=False):
+        super().__init__(None, node_type, id)
+        self.node_list = base_node_list
+        self.is_bwd = is_bwd
+
+
+class CommOpCostNode(CostNode):
+    def __init__(
+        self, node, node_type, id=None, comm_node_list=None, is_bwd=False
+    ):
+        super().__init__(node, node_type, id)
+        self.node_list = comm_node_list
+        self.ranks = []
+        self.comm_type = node.type
+        self.is_bwd = is_bwd
+
+    def set_ranks(self, ranks):
+        self.ranks = ranks
+
+    def set_shapes(self, input_shape, output_shape):
+        self.input_shape = input_shape
+        self.output_shape = output_shape
+
+    def init_comm_cost(self, cluster=None):
+        # ref: https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md
+        # should get from `cluster`
+        BANDWIDTH = 32 * 1024 / 1000  # MB/ms, V100 PCIe
+        num_ranks = len(self.ranks)
+        comm_volume = np.prod(self.input_shape) * 4
+
+        if 'allreduce' in self.comm_type:
+            self._cost = comm_volume / (
+                BANDWIDTH * num_ranks / (2 * (num_ranks - 1))
+            )
+        elif 'gather' in self.comm_type:
+            self._cost = comm_volume / (BANDWIDTH * num_ranks / (num_ranks - 1))
+        elif 'broadcast' in self.comm_type:
+            self._cost = comm_volume / BANDWIDTH
+        elif 'send' in self.comm_type or 'recv' in self.comm_type:
+            self._cost = comm_volume / BANDWIDTH
+        else:
+            self._cost = 0
+
+
+class TensorCostNode(CostNode):
+    def __init__(
+        self,
+        node,
+        node_type,
+        id=None,
+        base_node_list=None,
+        batch_size=None,
+        shared_node_id=None,
+    ):
+        super().__init__(node, node_type, id)
+        if node.name == "create_py_reader_0" or node.name == "double_buffer_0":
+            self.shape = [2, 2]
+            self.dtype = paddle.float32
+        else:
+            self.shape = node.shape
+            self.dtype = node.dtype
+        self.dtype_factor = 1
+        self.persistable = None
+        self.shared_node_id = shared_node_id
+        if self.dtype == paddle.float32 or node.dtype == paddle.int32:
+            self.dtype_factor *= 4
+        elif node.dtype == paddle.int64:
+            self.dtype_factor *= 8
+        elif node.dtype == paddle.uint8:
+            self.dtype_factor = 1
+        else:
+            self.dtype_factor = 2
+            # raise NotImplementedError("{} not counted".format(node.dtype))
+        self.batch_size = None
+        if batch_size is not None:
+            self.batch_size = batch_size
+
+    def get_size(self):
+        p = 1
+        for i in self.node.shape:
+            if i == -1:  # deal with placeholder
+                assert self.batch_size is not None, "Batch size not decided."
+                i = self.batch_size
+            p *= i
+        return p
+
+
+class CompOpCostNode(CostNode):
+    def __init__(self, node, node_type, id=None, is_bwd=False, is_optim=False):
+        super().__init__(node, node_type, id)
+        self.is_bwd = is_bwd
+        self.is_optim = is_optim
+
+    def init_comp_cost(self, cost_data):
+        # TODO: improve base.CostModel for more specific cost_data
+        op_id = self.node.desc.id()
+        if op_id in cost_data.keys():
+            self.cost = cost_data[op_id]
+        else:
+            self.cost = 0.0
+
+
+class PipeEvent:
+    def __init__(self, stage_id, event_name, duration, start_time=-1):
+        self.stage_id = stage_id
+        self.name = event_name
+        self.duration = duration
+        self.s_time = start_time
+        self.e_time = -1
+
+
+class CostModel:
+    def __init__(
+        self,
+        mode=CostModelMode.BENCHMARKING,
+        cluster=None,
+        batch_size=1,
+        microbatch_num=1,
+        opcall_overhead=0,
+        standalone_cost_data=None,
+        pipeline_config=None,
+    ):
+        self.mode = mode
+
+        # parameters
+        self.opcall_overhead = opcall_overhead
+        self.batch_size = batch_size
+        self.microbatch_num = microbatch_num
+
+        self.nodes = {}  # name -> node
+
+        self.origin_graph = {}  # original graph
+        self.op_graph = {}  # op graph (no variables nodes)
+        self.runtime_graph = {}  # runtime graph, for simulation
+
+        self.cluster = cluster
+        self.cost_data = standalone_cost_data
+        self.pp2rank = pipeline_config
+        if self.pp2rank is not None:
+            self.rank2pp = {}
+            for stage_idx, ranks in enumerate(self.pp2rank):
+                for rank in ranks:
+                    self.rank2pp[rank] = stage_idx
+        else:
+            self.rank2pp = None
+
+        self.ring2rank = {}
+
+        self.fwd_time = []
+        self.bwd_time = []
+        self.optim_time = []
+
+    def _parse_sub_program(self, program, nodes, graph, cost_data, sub_idx):
+        assert len(program.blocks) == 1, (
+            "Program more than 1 block not supported."
+        )
+        block = program.blocks[0]
+
+        var_id = "lod_tensor_blocking_queue_0"
+        new_var = program.global_block().create_var(
+            name=var_id,
+            dtype=paddle.float32,
+            type=core.VarDesc.VarType.DENSE_TENSOR,
+        )
+        nodes[var_id] = TensorCostNode(
+            new_var, CostNodeType.VARIABLE, "lod_tensor_blocking_queue_0"
+        )
+        for var in block.vars.values():
+            var_id = var.name
+            # if var.name == "create_py_reader_0" or var.name == "double_buffer_0":
+            #     continue
+            nodes[var_id] = TensorCostNode(var, CostNodeType.VARIABLE, var_id)
+            graph[var_id] = [[], []]
+
+        for op in block.ops:
+            op_id = op.type + "_" + str(op.idx)
+            if (
+                op.type.startswith('c_')
+                or op.type.startswith('send')
+                or op.type.startswith('recv')
+            ):
+                is_bwd = False
+                if (
+                    op.type.startswith('c_')
+                    and op.type != "c_sync_calc_stream"
+                    and not op.type.startswith('c_embedding')
+                ):
+                    ring_id = op.attr('ring_id')
+                    if ring_id not in self.ring2rank:
+                        self.ring2rank[ring_id] = set()
+                    self.ring2rank[ring_id].add(sub_idx)
+                    is_bwd = '@GRAD' in op.output('Out')[0]
+                elif op.type.startswith('recv'):
+                    is_bwd = '@GRAD' in op.output('Out')[0]
+                elif op.type.startswith('send'):
+                    is_bwd = '@GRAD' in op.input('X')[0]
+                op_node = CommOpCostNode(
+                    op, CostNodeType.COMMUNICATION, op_id, is_bwd
+                )
+            else:
+                is_bwd = (
+                    int(op.attr('op_role')) == int(OpRole.Backward)
+                ) or "@GRAD" in op.input_arg_names
+                is_optim = 'LearningRate' in op.input_names
+                op_node = CompOpCostNode(
+                    op, CostNodeType.COMPUTATION, op_id, is_bwd, is_optim
+                )
+                op_node.init_comp_cost(cost_data)
+
+            nodes[op_id] = op_node
+            graph[op_id] = [[], []]
+
+            comm_input_shape = [0]
+            comm_output_shape = [0]
+            for i in range(len(op.input_names)):
+                try:
+                    var_id = op.input(op.input_names[i])[0]
+                    var_node = nodes[var_id]
+                    graph[op_id][PRED].append(var_node.id)
+                    graph[var_id][SUCC].append(op_node.id)
+                    comm_input_shape = var_node.shape
+                except:
+                    continue
+
+            for i in range(len(op.output_names)):
+                try:
+                    var_id = op.output(op.output_names[i])[0]
+                    var_node = nodes[var_id]
+                    graph[op_id][SUCC].append(var_node.id)
+                    graph[var_id][PRED].append(op_node.id)
+                    comm_output_shape = var_node.shape
+                except:
+                    continue
+            if op_node.type == CostNodeType.COMMUNICATION:
+                op_node.set_shapes(comm_input_shape, comm_output_shape)
+
+        # resolve hazard: rename the r/w hazard variable nodes to ensure self.origin_graph is a DAG
+        new_var_dict = {}
+        for node_id, node in nodes.items():
+            if node.type == CostNodeType.VARIABLE and node.node.persistable:
+                write_op_cnt = 0
+                for pred_id in graph[node_id][PRED]:
+                    pred = nodes[pred_id]
+                    if pred.type == CostNodeType.COMPUTATION and (
+                        pred_id in graph[node_id][SUCC]
+                    ):
+                        graph[pred_id][SUCC].remove(node_id)
+                        graph[node_id][PRED].remove(pred_id)
+
+                        write_op_cnt += 1
+                        new_var_id = node_id + f'_write_{write_op_cnt}'
+                        new_var = TensorCostNode(
+                            node.node,
+                            CostNodeType.VARIABLE,
+                            new_var_id,
+                            shared_node_id=node_id,
+                        )
+
+                        graph[new_var_id] = [[], []]
+                        graph[pred_id][SUCC].append(new_var_id)
+                        graph[new_var_id][PRED].append(pred_id)
+
+                        new_var_dict[new_var_id] = new_var
+        for k, v in new_var_dict.items():
+            nodes[k] = v
+        return nodes
+
+    def parse_program(self, distributed_program):
+        self.distributed_program = distributed_program
+        self.total_rank = len(self.distributed_program)
+        sub_prog_cnt = len(distributed_program)
+        self.nodes = [] * sub_prog_cnt
+        self.origin_graph = [] * sub_prog_cnt  # original graph
+        self.op_graph = [] * sub_prog_cnt  # op graph (no variables nodes)
+        self.runtime_graph = [] * sub_prog_cnt  # runtime graph, for simulation
+
+        for sub_idx, sub_prog in enumerate(distributed_program):
+            self.nodes.append({})
+            self.origin_graph.append({})
+            self.op_graph.append({})
+            self.runtime_graph.append({})
+            self._parse_sub_program(
+                sub_prog,
+                self.nodes[sub_idx],
+                self.origin_graph[sub_idx],
+                self.cost_data[
+                    0 if self.rank2pp is None else self.rank2pp[sub_idx]
+                ],
+                sub_idx,
+            )
+        return self.nodes
+
+    def _find_succ_op(self, node_id, sub_idx=0):
+        succ_ops_id = []
+        for succ_id in self.origin_graph[sub_idx][node_id][SUCC]:
+            succ = self.nodes[sub_idx][succ_id]
+            if (
+                succ.type == CostNodeType.COMMUNICATION
+                or succ.type == CostNodeType.COMPUTATION
+            ):
+                succ_ops_id.append(succ_id)
+            elif succ.type == CostNodeType.VARIABLE:
+                succ_ops_id = succ_ops_id + self._find_succ_op(succ_id, sub_idx)
+            else:
+                raise NotImplementedError(
+                    f'This type of node not supported yet:{succ.type}'
+                )
+        return succ_ops_id
+
+    def build_op_graph(self):
+        for sub_idx in range(self.total_rank):
+            op_nodes_id = []
+            for node_id, node in self.nodes[sub_idx].items():
+                if node.type == CostNodeType.VARIABLE:
+                    continue
+                self.op_graph[sub_idx][node_id] = [[], []]
+                op_nodes_id.append(node_id)
+            for op_id in op_nodes_id:
+                succ_nodes_id = self._find_succ_op(op_id, sub_idx)
+
+                self.op_graph[sub_idx][op_id][SUCC] = succ_nodes_id
+                for succ_id in succ_nodes_id:
+                    self.op_graph[sub_idx][succ_id][PRED].append(op_id)
+
+    def build_runtime_graph(self):
+        self.runtime_graph = copy.deepcopy(self.op_graph)
+
+    def eliminate_multi_edges(self, graph=None):
+        for node_id, edges in graph.items():
+            graph[node_id][PRED] = list(set(edges[PRED]))
+            graph[node_id][SUCC] = list(set(edges[SUCC]))
+
+    def merge_comm(self):
+        for sub_idx in range(self.total_rank):
+            for node_id, edges in self.op_graph[sub_idx].items():
+                node = self.nodes[sub_idx][node_id]
+                if (
+                    node_id.startswith('c_')
+                    and not node.id.startswith("c_sync_calc_stream")
+                    and not node.id.startswith('c_embedding')
+                ):
+                    ring_id = node.node.attr('ring_id')
+                    node.set_ranks(list(self.ring2rank[ring_id]))
+                    node.init_comm_cost(self.cluster)
+                elif node_id.startswith('send') or node_id.startswith('recv'):
+                    peer_rank = node.node.attr('peer')
+                    node.set_ranks([sub_idx, peer_rank])
+                    node.init_comm_cost(self.cluster)
+                else:
+                    pass  # Not communication op
+
+    def _merge_node(self, to_merge_node_list, merge_type='linear', nodes=None):
+        nodes_list = []
+        node_cost = 0
+        for node in to_merge_node_list:
+            if isinstance(node, MergedOpsCostNode):
+                nodes_list += node.node_list
+            else:
+                nodes_list.append(node.id)
+            if merge_type == 'linear':
+                node_cost += node.cost
+            elif merge_type == 'branch':
+                node_cost = max(node_cost, node.cost)
+            else:
+                raise NotImplementedError(
+                    f'This type of merging is not supported:{merge_type}'
+                )
+        merged_node_id = 'merged_' + str(len(nodes))
+        is_bwd = to_merge_node_list[0].is_bwd
+        merged_node = MergedOpsCostNode(
+            CostNodeType.MERGED,
+            id=merged_node_id,
+            base_node_list=nodes_list,
+            is_bwd=is_bwd,
+        )
+        merged_node.cost = node_cost
+        return merged_node_id, merged_node
+
+    def merge_linear(self):
+        r'''
+        This method does the following:
+        If X depends on Y only, they must be run sequentially.
+            [ e.g. A ->- C ->- D   D and E depends on C only.]
+            [      B ->-/ \->- E   C depends on A and B.     ]
+        We merge X and Y into a new node and sum up their cost time.
+        '''
+        cnt = 0
+        for sub_idx in range(self.total_rank):
+            cnt += self._merge_linear(
+                self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=False
+            )
+            cnt += self._merge_linear(
+                self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=True
+            )
+        return cnt
+
+    def merge_branch(self):
+        r'''
+        This method does the following:
+        If a node has more than one successor, there is *branch*.
+            [ e.g. A ->- B ->- D                                       ]
+            [       \->- C ->- / , B and C can be run at the same time ]
+            case 1: if B or C is null (or D is directly dependent on A),
+                    it's equivalent to A->C->D or A->B->D, fall back to self.merge_linear
+            case 2: if both B and C are some op,
+                    merged_cost = max(cost(B), cost(C))
+        '''
+        cnt = 0
+        for sub_idx in range(self.total_rank):
+            cnt += self._merge_branch(
+                self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=False
+            )
+            cnt += self._merge_branch(
+                self.nodes[sub_idx], self.runtime_graph[sub_idx], is_bwd=True
+            )
+        return cnt
+
+    def _merge_linear(self, nodes, runtime_graph, is_bwd=False):
+        reduct_cnt = 0
+        rt_nodes_id = list(runtime_graph.keys())
+        for node_id in rt_nodes_id:
+            if node_id not in runtime_graph.keys():
+                continue
+            node = nodes[node_id]
+            if not is_bwd == node.is_bwd or node.is_optim:
+                continue
+            edges = runtime_graph[node_id]
+            ind = len(edges[PRED])  # in_degree
+            if ind == 1:  # only depend on one node
+                pred_id = edges[PRED][0]
+                pred = nodes[pred_id]
+                merged_node_id, merged_node = self._merge_node(
+                    [node, pred], merge_type='linear', nodes=nodes
+                )
+                nodes[merged_node_id] = merged_node
+                runtime_graph[merged_node_id] = [[], []]
+
+                # delete edges and add new edges
+                succ = None
+                try:
+                    runtime_graph[merged_node_id][SUCC] = copy.deepcopy(
+                        edges[SUCC]
+                    )
+
+                    if len(runtime_graph[pred_id][SUCC]) > 1:
+                        # predecessor has more than 1 successor
+                        # the merged_node is to inherit the rest of its successors
+                        succ = runtime_graph[pred_id][SUCC]
+                        succ.remove(node_id)
+                        runtime_graph[merged_node_id][SUCC] += succ
+                    runtime_graph[merged_node_id][PRED] = runtime_graph[
+                        pred_id
+                    ][PRED]
+                except:
+                    pass
+                try:
+                    for i in runtime_graph[pred_id][PRED]:
+                        try:
+                            runtime_graph[i][SUCC].remove(pred_id)
+                        except:
+                            continue
+                        runtime_graph[i][SUCC].append(merged_node_id)
+                except:
+                    pass
+
+                try:
+                    for i in edges[SUCC]:
+                        runtime_graph[i][PRED].remove(node_id)
+                        runtime_graph[i][PRED].append(merged_node_id)
+                except:
+                    pass
+                if succ is not None:
+                    for i in succ:
+                        try:
+                            runtime_graph[i][PRED].remove(pred_id)
+                        except:
+                            continue
+                        runtime_graph[i][PRED].append(merged_node_id)
+
+                runtime_graph.pop(node_id)
+                try:
+                    runtime_graph.pop(pred_id)
+                except:
+                    continue
+                reduct_cnt += 1
+                self.eliminate_multi_edges(runtime_graph)
+                break
+        return reduct_cnt  # the number of nodes that have been reduced
+
+    def _merge_branch(self, nodes, runtime_graph, is_bwd=False):
+        reduct_cnt = 0
+        rt_nodes_id = list(runtime_graph.keys())
+        for node_id in rt_nodes_id:
+            node = nodes[node_id]
+            if not is_bwd == node.is_bwd or node.is_optim:
+                continue
+            edges = runtime_graph[node_id]
+            outd = len(edges[SUCC])  # out_degree
+            if outd > 1:  # branch out
+                succ_nodes_id = edges[SUCC]
+
+                succ_to_elim = []
+                for succ_id in succ_nodes_id:
+                    for succ_2_id in succ_nodes_id:
+                        try:
+                            tmp = runtime_graph[succ_2_id][SUCC]
+                        except:
+                            continue
+                        if succ_id in tmp:
+                            succ_to_elim.append(succ_id)
+                            break
+                for id in succ_to_elim:
+                    edges[SUCC].remove(id)
+                    runtime_graph[id][PRED].remove(node_id)
+                    reduct_cnt += 1
+
+                to_merge = True
+                try:
+                    if (
+                        len(edges[SUCC]) < 1
+                        or len(runtime_graph[edges[SUCC][0]][SUCC]) < 1
+                    ):
+                        continue
+                except:
+                    continue
+                end_node_id = runtime_graph[edges[SUCC][0]][SUCC][0]
+                for i in succ_nodes_id:
+                    try:
+                        if (
+                            len(runtime_graph[i][SUCC]) != 1
+                            or runtime_graph[i][SUCC][0] != end_node_id
+                        ):
+                            to_merge = False  # if branches has different end node, we don't merge them
+                            break
+                    except:
+                        continue
+                if to_merge and len(succ_nodes_id) > 1:
+                    to_merge_node_list = [nodes[i] for i in succ_nodes_id]
+                    merged_node_id, merged_node = self._merge_node(
+                        to_merge_node_list, merge_type='branch', nodes=nodes
+                    )
+                    nodes[merged_node_id] = merged_node
+                    runtime_graph[merged_node_id] = [[], []]
+
+                    # delete edges and add new edges
+                    runtime_graph[merged_node_id][SUCC] = [end_node_id]
+                    runtime_graph[merged_node_id][PRED] = edges[PRED]
+
+                    runtime_graph[end_node_id][PRED] = [merged_node_id]
+                    runtime_graph[node_id][SUCC] = [merged_node_id]
+
+                    try:
+                        for i in succ_nodes_id:
+                            runtime_graph.pop(i)
+                        reduct_cnt += len(to_merge_node_list) - 1
+                        break
+                    except:
+                        pass
+        return reduct_cnt
+
+    def get_runtime_cost(self):
+        def get_node_cost(node):
+            node_cost = node.cost + self.opcall_overhead
+            if isinstance(node, MergedOpsCostNode):
+                for it in node.node_list:
+                    node_cost += self.opcall_overhead
+            return node_cost
+
+        for sub_idx in range(self.total_rank):
+            fwd_cost = 0
+            bwd_cost = 0
+            optim_cost = 0
+            for node_id in self.runtime_graph[sub_idx].keys():
+                node = self.nodes[sub_idx][node_id]
+                if node.is_optim:
+                    optim_cost += get_node_cost(node)
+                elif node.is_bwd:
+                    bwd_cost += get_node_cost(node)
+                else:
+                    fwd_cost += get_node_cost(node)
+            self.fwd_time.append(fwd_cost)
+            self.bwd_time.append(bwd_cost)
+            self.optim_time.append(optim_cost)
+        return self.fwd_time, self.bwd_time, self.optim_time
+
+    def get_mem(self):
+        static_list = []
+        top_list = []
+        for sub_idx in range(self.total_rank):
+            static_mem, cur_mem, top_mem = self._simulate_mem(
+                self.nodes[sub_idx], self.origin_graph[sub_idx]
+            )
+            static_list.append(static_mem)
+            top_list.append(top_mem)
+        return static_list, top_list
+
+    def _simulate_mem(self, nodes, origin_graph):
+        q = queue.Queue(1024)
+        sim_graph = copy.deepcopy(origin_graph)
+        for node_id, node in nodes.items():
+            if len(sim_graph[node_id][PRED]) == 0:
+                q.put(node_id)
+
+        q.put('nop')
+        cur_mem = 0
+        top_mem = -1
+        static_mem = 0
+        while not q.empty():
+            node_id = q.get()
+            node = None
+            size = 0
+            if node_id == 'nop':
+                top_mem = max(cur_mem, top_mem)
+                if q.empty():
+                    break
+                else:
+                    q.put(node_id)
+                    continue
+            else:
+                node = nodes[node_id]
+            if node.type == CostNodeType.VARIABLE:
+                size = node.get_size()
+                if node.node.persistable:
+                    static_mem += size
+                cur_mem += size
+            edges = sim_graph[node_id]
+            if not (
+                node.type == CostNodeType.VARIABLE and node.node.persistable
+            ):
+                for succ_id in edges[SUCC]:
+                    sim_graph[succ_id][PRED].remove(node_id)
+                    if len(sim_graph[succ_id][PRED]) == 0:
+                        q.put(succ_id)
+            for pred_id in edges[PRED]:
+                pred = nodes
+                if pred.type == CostNodeType.VARIABLE:
+                    sim_graph[pred_id][SUCC].remove(node_id)
+                    if (
+                        len(sim_graph[pred_id][SUCC]) == 0
+                        and not pred.node.persistable
+                    ):
+                        cur_mem -= pred.get_size()
+        return static_mem, cur_mem, top_mem
+
+    def get_pipeline_time(self):
+        if self.pp2rank is None:
+            return self.fwd_time[0] + self.bwd_time[0] + self.optim_time[0]
+        else:
+            return self._simulate_pipeline()
+
+    def _simulate_pipeline(self):
+        stage_num = len(self.pp2rank)
+        event_list = []
+        global_time = [0] * stage_num
+        total_time = 0
+        fwd_cnt = list(range(stage_num, 0, -1))
+        bwd_cnt = [self.microbatch_num] * stage_num
+        q = queue.Queue(1024)
+
+        for i in range(self.microbatch_num):
+            q.put(PipeEvent(0, 'fwd', self.fwd_time[0]))
+
+        while not q.empty():
+            e = q.get()
+            stid = e.stage_id
+            if e.name == 'fwd':
+                if fwd_cnt[stid] > 0:
+                    e.s_time = max(global_time[stid], e.s_time)
+                    e.e_time = e.s_time + e.duration
+                    event_list.append(e)
+                    if stid != stage_num - 1:
+                        q.put(
+                            PipeEvent(
+                                stid + 1,
+                                'fwd',
+                                self.fwd_time[stid + 1],
+                                start_time=e.e_time,
+                            )
+                        )
+                    else:
+                        q.put(
+                            PipeEvent(
+                                stid,
+                                'bwd',
+                                self.bwd_time[stid],
+                                start_time=e.e_time,
+                            )
+                        )
+                    fwd_cnt[stid] -= 1
+                    global_time[stid] = e.e_time
+                else:
+                    q.put(e)
+            elif e.name == 'bwd':
+                e.s_time = max(global_time[stid], e.s_time)
+                e.e_time = e.s_time + e.duration
+                event_list.append(e)
+                if stid != 0:
+                    q.put(
+                        PipeEvent(
+                            stid - 1,
+                            'bwd',
+                            self.bwd_time[stid - 1],
+                            start_time=e.e_time,
+                        )
+                    )
+                fwd_cnt[stid] += 1
+                bwd_cnt[stid] -= 1
+                if bwd_cnt[stid] == 0:
+                    q.put(
+                        PipeEvent(
+                            stid,
+                            'optim',
+                            self.optim_time[stid],
+                            start_time=e.e_time,
+                        )
+                    )
+                global_time[stid] = e.e_time
+            elif e.name == 'optim':
+                e.s_time = max(global_time[stid], e.s_time)
+                e.e_time = e.s_time + e.duration
+                event_list.append(e)
+                global_time[stid] = e.e_time
+            else:
+                raise NotImplementedError(
+                    f'This type of pipe event is not supported yet.{e.name}'
+                )
+
+        for t in global_time:
+            total_time = max(total_time, t)
+        return total_time
+
+    def get_cost(self):
+        cost = Cost()
+        static_mem, peak_mem = self.get_mem()
+        cost.static_mem = static_mem
+        cost.peak_mem = peak_mem
+        self.merge_comm()
+        while True:
+            cnt = 0
+            cnt += self.merge_linear()
+            cnt += self.merge_branch()
+            if cnt == 0:  # can't be further merged
+                break
+        self.get_runtime_cost()
+        cost.runtime = self.get_pipeline_time()
+        return cost
+
+    def init(self, distributed_program):
+        self.parse_program(distributed_program)
+        self.build_op_graph()
+        for sub_idx in range(self.total_rank):
+            self.eliminate_multi_edges(self.op_graph[sub_idx])
+        self.build_runtime_graph()
+
+
+def estimate_cost(
+    distributed_program,
+    cluster,
+    pipeline_config,
+    standalone_cost_data,
+    batch_size,
+):
+    """
+    Estimated cost from distributed program, cluster model and distributed settings.
+
+    Args:
+        distributed_program(list): list of paddle programs
+        cluster(Cluster): cluster model
+        standalone_cost_data(CostData): cost data given by paddle.core
+        batch_size(int): batch size of the training workload
+        pipeline_config(list): configuration of pipeline stage allocation
+    """
+    # the following line is left for now, cluster model will be involved in the future
+    assert cluster is None, "For now, cluster remains None"
+    cm_ctx = CostModel(
+        cluster=cluster,
+        batch_size=batch_size,
+        standalone_cost_data=standalone_cost_data,
+        pipeline_config=pipeline_config,
+    )
+    cm_ctx.init(distributed_program)
+    cost = cm_ctx.get_cost()
+    return cost

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

@@ -0,0 +1,19 @@
+#   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.base.core import (  # noqa: F401
+    DistTensorSpec,
+    OperatorDistAttr,
+    TensorDistAttr,
+)

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

@@ -0,0 +1,62 @@
+#   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.static import InputSpec
+
+from ..placement_type import get_shard_spec
+from .utils import convert_to_dims_mapping
+
+
+class DistributedInputSpec(InputSpec):
+    def __init__(
+        self,
+        shape,
+        dtype='float32',
+        name=None,
+        stop_gradient=False,
+        mesh=None,
+        placements=None,
+        local_shape=None,
+    ):
+        super().__init__(shape, dtype, name, stop_gradient)
+        self.mesh = copy.deepcopy(mesh)
+        sharding_specs = get_shard_spec(mesh, placements, len(self.shape))
+        self.dims_mapping = convert_to_dims_mapping(sharding_specs, mesh)
+        self.local_shape = local_shape
+
+    @classmethod
+    def from_dtensor(cls, dtensor, name=None, shape=None):
+        """
+        Generates a DistributedInputSpec based on dist tensor.
+
+        Args:
+            dtensor: the dist tensor.
+
+        Returns:
+            A DistributedInputSpec instance generated from dtensor.
+        """
+        return cls(
+            shape=dtensor.shape if shape is None else shape,
+            dtype=dtensor.dtype,
+            name=name,
+            stop_gradient=dtensor.stop_gradient,
+            mesh=dtensor.process_mesh,
+            placements=dtensor.placements,
+            local_shape=dtensor._local_value().shape,
+        )
+
+    def __repr__(self):
+        return f"{super().__repr__()}, mesh:{self.mesh}, placements:{self.dims_mapping}"

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

@@ -0,0 +1,290 @@
+# 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 abc
+
+import numpy as np
+
+import paddle
+from paddle.io import BatchSampler, IterableDataset
+from paddle.io.dataloader.batch_sampler import (
+    DistributedBatchSampler,
+    _InfiniteIterableSampler,
+)
+from paddle.io.dataloader.dataloader_iter import (
+    _DatasetKind,
+    default_collate_fn,
+    default_convert_fn,
+)
+
+
+class DistributedDataLoaderBase(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def __iter__(self):
+        raise NotImplementedError
+
+
+class DistributedDataLoaderFromGenerator(DistributedDataLoaderBase):
+    def __init__(
+        self,
+        dataset,
+        feed_list=None,
+        capacity=None,
+        use_double_buffer=True,
+        iterable=True,
+        return_list=False,
+        use_multiprocess=False,
+        drop_last=True,
+        places=None,
+        batch_size=1,
+        epochs=1,
+        steps_per_epoch=None,
+        collate_fn=None,
+        split_data=True,
+        data_parallel_world_size=[],
+        data_parallel_rank=[],
+        acc_steps=1,
+    ):
+        self.dataset = dataset
+        self.feed_list = feed_list
+        self.capacity = capacity
+        self.use_double_buffer = use_double_buffer
+        self.iterable = iterable
+        self.return_list = return_list
+        self.use_multiprocess = use_multiprocess
+        self.drop_last = drop_last
+        self.places = places
+        self.batch_size = batch_size
+        self.epochs = epochs
+        self.steps_per_epoch = steps_per_epoch
+        self.collate_fn = collate_fn
+        self.split_data = split_data
+        assert len(data_parallel_world_size) == len(feed_list)
+        assert len(data_parallel_rank) == len(feed_list)
+        self.dp_world_sizes = data_parallel_world_size
+        self.dp_ranks = data_parallel_rank
+        self.acc_steps = acc_steps
+
+        if isinstance(dataset, IterableDataset):
+            self.dataset_kind = _DatasetKind.ITER
+        else:
+            self.dataset_kind = _DatasetKind.MAP
+
+        if self.batch_size is None:
+            self.batch_sampler = None
+        else:
+            if isinstance(dataset, IterableDataset):
+                self.batch_sampler = _InfiniteIterableSampler(
+                    dataset, batch_size
+                )
+            else:
+                self.batch_sampler = BatchSampler(
+                    dataset,
+                    batch_size=batch_size,
+                    shuffle=False,
+                    drop_last=drop_last,
+                )
+
+        self.auto_collate_batch = self.batch_sampler is not None
+        self.sampler_iter = iter(self.index_sampler)
+
+        if self.auto_collate_batch:
+            self.collate_fn = collate_fn or default_collate_fn
+        else:
+            self.collate_fn = collate_fn or default_convert_fn
+
+        self.dataset_fetcher = _DatasetKind.create_fetcher(
+            self.dataset_kind,
+            self.dataset,
+            self.auto_collate_batch,
+            self.collate_fn,
+            self.drop_last,
+        )
+
+        self._steps = self._infer_steps()
+        self._inner_dataloader = self._create_inner_dataloader()
+
+    def __iter__(self):
+        self._cur_step = 0
+        self._inner_dataloader.start()
+        return self
+
+    def __next__(self):
+        if not self._steps:
+            self._cur_step += 1
+            return None
+        elif self._cur_step < self._steps:
+            self._cur_step += 1
+            return None
+        else:
+            self._inner_dataloader.reset()
+            self.sampler_iter = iter(self.index_sampler)
+            raise StopIteration
+
+    def _infer_steps(self):
+        if isinstance(self.steps_per_epoch, int) and self.steps_per_epoch > 0:
+            return self.steps_per_epoch
+        try:
+            if isinstance(self.dataset, IterableDataset):
+                steps_per_epoch = None
+            elif self.batch_size is None:
+                steps_per_epoch = len(self.dataset) // self.acc_steps
+            else:
+                steps_per_epoch = (
+                    len(self.dataset) // self.batch_size // self.acc_steps
+                )
+        except:
+            raise ValueError(
+                "Please set `steps_per_epoch` or implement `__len__` method in dataset class."
+            )
+        return steps_per_epoch
+
+    @property
+    def index_sampler(self):
+        if self.auto_collate_batch:
+            return self.batch_sampler
+        else:
+            if self.dataset_kind == _DatasetKind.MAP:
+                return list(range(len(self.dataset)))
+            else:
+                return _InfiniteIterableSampler(self.dataset, 1)
+
+    def _create_inner_dataloader(self):
+        def data_generator():
+            while True:
+                try:
+                    indices = next(self.sampler_iter)
+                    batch = self.dataset_fetcher.fetch(indices)
+                    if batch is None:
+                        break
+                except StopIteration:
+                    self.dataset_fetcher = _DatasetKind.create_fetcher(
+                        self.dataset_kind,
+                        self.dataset,
+                        self.auto_collate_batch,
+                        self.collate_fn,
+                        self.drop_last,
+                    )
+                    break
+
+                partial_data = []
+                for i, d in enumerate(batch):
+                    array = np.array(d)
+                    if not self.split_data:
+                        partial_data.append(array)
+                        continue
+
+                    batch_size = array.shape[0]
+                    assert batch_size % self.dp_world_sizes[i] == 0, (
+                        f"batch_size [{batch_size}] is not divisible by dp_world_size [{self.dp_world_sizes[i]}]"
+                    )
+                    partial_data.append(
+                        np.split(array, self.dp_world_sizes[i])[
+                            self.dp_ranks[i]
+                        ]
+                    )
+
+                yield partial_data
+
+        dataloader = paddle.base.io.DataLoader.from_generator(
+            feed_list=self.feed_list,
+            capacity=self.capacity,
+            use_double_buffer=self.use_double_buffer,
+            # iterable=self.iterable,
+            iterable=False,
+            return_list=self.return_list,
+            use_multiprocess=self.use_multiprocess,
+            drop_last=self.drop_last,
+        )
+        dataloader.set_batch_generator(data_generator, self.places)
+
+        return dataloader
+
+
+class DistributedDataLoader(DistributedDataLoaderBase):
+    def __init__(
+        self,
+        dataset,
+        feed_list=None,
+        places=None,
+        return_list=True,
+        batch_size=1,
+        shuffle=False,
+        drop_last=False,
+        collate_fn=None,
+        num_workers=0,
+        use_buffer_reader=True,
+        use_shared_memory=True,
+        timeout=0,
+        worker_init_fn=None,
+        epochs=1,
+        steps_per_epoch=None,
+        split_data=True,
+        data_parallel_world_size=[],
+        data_parallel_rank=[],
+    ):
+        self.dataset = dataset
+        self.feed_list = feed_list
+        self.return_list = return_list
+        self.places = places
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+        self.collate_fn = collate_fn
+        self.num_workers = num_workers
+        self.use_buffer_reader = use_buffer_reader
+        self.use_shared_memory = use_shared_memory
+        self.timeout = timeout
+        self.worker_init_fn = worker_init_fn
+        self.epochs = epochs
+        self.steps_per_epoch = steps_per_epoch
+        self.dp_world_sizes = data_parallel_world_size
+        self.dp_ranks = data_parallel_rank
+        self.split_data = split_data
+
+        if self.batch_size is None:
+            self.batch_sampler = None
+        else:
+            self.batch_sampler = DistributedBatchSampler(
+                dataset=self.dataset,
+                batch_size=self.batch_size,
+                num_replicas=self.dp_world_sizes[0],
+                rank=self.dp_ranks[0],
+                shuffle=self.shuffle,
+                drop_last=self.drop_last,
+            )
+
+        self._dataloader = paddle.io.DataLoader(
+            self.dataset,
+            feed_list=self.feed_list,
+            places=self.places,
+            return_list=self.return_list,
+            batch_sampler=self.batch_sampler,
+            batch_size=1 if self.batch_sampler else self.batch_size,
+            collate_fn=self.collate_fn,
+            num_workers=self.num_workers,
+            use_buffer_reader=self.use_buffer_reader,
+            use_shared_memory=self.use_shared_memory,
+            timeout=self.timeout,
+            worker_init_fn=self.worker_init_fn,
+        )
+
+    def __len__(self):
+        return len(self._dataloader)
+
+    def __iter__(self):
+        return self._dataloader.__iter__()
+
+    def __call__(self):
+        return self._dataloader.__iter__()

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

@@ -0,0 +1,323 @@
+#   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.static import Variable
+
+from .dist_attribute import OperatorDistAttr
+from .utils import (
+    __no_shape_var_type__,
+    convert_to_shard_spec,
+    verify_shard_spec,
+)
+
+
+class DistributedOperator:
+    def __init__(self, serial_op, dist_attr=None):
+        self._serial_op = serial_op
+        if dist_attr is not None and isinstance(dist_attr, OperatorDistAttr):
+            # TODO: remove this deepcopy after we fix the issue
+            self._dist_attr = copy.deepcopy(dist_attr)
+            # self._dist_attr = dist_attr
+            # TODO: Do we really need to write back to serial op？
+            self._serial_op.dist_attr = dist_attr
+        else:
+            assert dist_attr is None, f"{dist_attr}"
+            # Use the dist attr of serial_op to do the initialization
+            self._dist_attr = self._serial_op.dist_attr
+        self._serial_inputs = {}
+        self._serial_outputs = {}
+
+    @property
+    def serial_op(self):
+        return self._serial_op
+
+    @property
+    def dist_attr(self):
+        return self._dist_attr
+
+    @dist_attr.setter
+    def dist_attr(self, dist_attr):
+        self._dist_attr = dist_attr
+        # TODO: Do we really need to write back to serial op？
+        self._serial_op.dist_attr = dist_attr
+
+    def get_serial_input(self, name):
+        if self._serial_op.type == "create_py_reader":
+            tensor = None
+        elif self._serial_op.block._find_var_recursive(name) is not None:
+            tensor = self._serial_op.block._var_recursive(name)
+        else:
+            tensor = None
+        return tensor
+
+    def get_serial_output(self, name):
+        tensor = self._serial_op.block._var_recursive(name)
+        return tensor
+
+    def validate_dist_attr(self):
+        if "read" in self.serial_op.type or "while" == self.serial_op.type:
+            return True
+        for name in self.serial_op.input_arg_names:
+            input_dist_attr = self.dist_attr.get_input_dist_attr(name)
+            dims_mapping = input_dist_attr.dims_mapping
+            if self.get_serial_input(name).type in __no_shape_var_type__:
+                shape = []
+            else:
+                shape = self.get_serial_input(name).shape
+            if len(shape) != len(dims_mapping):
+                return False
+            for i in range(len(dims_mapping)):
+                if dims_mapping[i] < -1 or dims_mapping[i] >= len(
+                    self.dist_attr.process_mesh.shape
+                ):
+                    return False
+            for i in range(len(self.dist_attr.process_mesh.shape)):
+                if dims_mapping.count(i) > 1:
+                    return False
+            if self.dist_attr.process_mesh != input_dist_attr.process_mesh:
+                return False
+
+        for name in self.serial_op.output_arg_names:
+            output_dist_attr = self.dist_attr.get_output_dist_attr(name)
+            dims_mapping = output_dist_attr.dims_mapping
+            if self.get_serial_output(name).type in __no_shape_var_type__:
+                shape = []
+            else:
+                shape = self.get_serial_output(name).shape
+            if len(shape) != len(dims_mapping):
+                return False
+            for i in range(len(dims_mapping)):
+                if dims_mapping[i] < -1 or dims_mapping[i] >= len(
+                    self.dist_attr.process_mesh.shape
+                ):
+                    return False
+            for i in range(len(self.dist_attr.process_mesh.shape)):
+                if dims_mapping.count(i) > 1:
+                    return False
+            if self.dist_attr.process_mesh != output_dist_attr.process_mesh:
+                return False
+        return True
+
+    def __str__(self):
+        str = f"{{op type: {self.serial_op.desc.type()}, op id: {self.serial_op.desc.id()}, op original_id: {self.serial_op.desc.original_id()}"
+
+        # str += ", {}".format(self.dist_attr)
+        # return str
+
+        if self.dist_attr.is_annotated("process_mesh"):
+            annotated_str = "annotated"
+        else:
+            annotated_str = "non-annotated"
+        str += (
+            f", process_mesh ({annotated_str}): {self.dist_attr.process_mesh}"
+        )
+
+        str += f" , execution_stream: {self.dist_attr.execution_stream}"
+
+        for arg_name in self.serial_op.desc.input_arg_names():
+            try:
+                dims_mapping = self.dist_attr.get_input_dims_mapping(arg_name)
+            except IndexError:
+                raise IndexError(
+                    f"There is not input var '{arg_name}''s dist_attr in current op '{self.serial_op.desc.type()}'"
+                )
+            if self.dist_attr.is_annotated_input_dims_mapping(arg_name):
+                annotated_str = "annotated"
+            else:
+                annotated_str = "non-annotated"
+            if self.get_serial_input(arg_name) is not None:
+                if self.get_serial_input(arg_name).is_parameter:
+                    is_parameter_str = "parameter"
+                else:
+                    is_parameter_str = "non-parameter"
+            else:
+                is_parameter_str = "non-parameter"
+
+            # partial
+            input_dist_attr = self.dist_attr.get_input_dist_attr(arg_name)
+            partial_dims = sorted(input_dist_attr._partial_dims())
+
+            str += f"; {arg_name}'s dims_mapping (input, {annotated_str}, {is_parameter_str}): {dims_mapping}, partial on dims: {partial_dims}"
+
+        for arg_name in self.serial_op.desc.output_arg_names():
+            try:
+                dims_mapping = self.dist_attr.get_output_dims_mapping(arg_name)
+            except IndexError:
+                raise IndexError(
+                    f"There is not output var '{arg_name}''s dist_attr in current op '{self.serial_op.desc.type()}'"
+                )
+            if self.dist_attr.is_annotated_output_dims_mapping(arg_name):
+                annotated_str = "annotated"
+            else:
+                annotated_str = "non-annotated"
+            if self.get_serial_output(arg_name) is not None:
+                if self.get_serial_output(arg_name).is_parameter:
+                    is_parameter_str = "parameter"
+                else:
+                    is_parameter_str = "non-parameter"
+            else:
+                is_parameter_str = "non-parameter"
+
+            # partial
+            output_dist_attr = self.dist_attr.get_output_dist_attr(arg_name)
+            partial_dims = sorted(output_dist_attr._partial_dims())
+
+            str += f"; {arg_name}'s dims_mapping (output, {annotated_str}, {is_parameter_str}): {dims_mapping}, partial on dims: {partial_dims}"
+
+        str += f", dist_impl idx: {self.dist_attr.impl_idx} , dist_impl type: {self.dist_attr.impl_type}, chunk_id: {self.dist_attr.chunk_id} }}"
+
+        return str
+
+    def __deepcopy__(self, memo):
+        cls = self.__class__
+        result = cls.__new__(cls)
+        memo[id(self)] = result
+        for k, v in self.__dict__.items():
+            if (
+                k == "_serial_op"
+                or k == "_serial_inputs"
+                or k == "_serial_outputs"
+            ):
+                setattr(result, k, v)
+            else:
+                setattr(result, k, copy.deepcopy(v, memo))
+        return result
+
+
+class DistributedOperatorHelper:
+    def __init__(
+        self,
+        serial_op,
+        process_mesh,
+        in_dims_mappings,
+        out_dims_mappings,
+        kwargs,
+    ):
+        self._serial_op = serial_op
+        self._process_mesh = process_mesh
+        self._in_dims_mappings = in_dims_mappings
+        self._out_dims_mappings = out_dims_mappings
+        self._chunk_id = kwargs["chunk_id"] if "chunk_id" in kwargs else 0
+
+    def __call__(self, *args, **kwargs):
+        tensor_to_dims_mapping = {}
+        index = 0
+        if self._in_dims_mappings:
+            assert len(args) + len(kwargs) == len(self._in_dims_mappings), (
+                f"The length of dims_mapping {len(self._in_dims_mappings)} does not matching the length output {len(args) + len(kwargs)}."
+            )
+        for arg in args:
+            if isinstance(arg, Variable) and self._in_dims_mappings:
+                tensor_to_dims_mapping[arg.name] = self._in_dims_mappings[index]
+            index += 1
+        for arg in kwargs.values() and self._in_dims_mappings:
+            if isinstance(arg, Variable):
+                tensor_to_dims_mapping[arg.name] = self._in_dims_mappings[index]
+            index += 1
+
+        default_prog = paddle.static.default_main_program()
+        cur_block = default_prog.current_block()
+        op_size = len(cur_block.ops)
+        if paddle.base.dygraph.base.in_to_static_mode():
+            output = paddle.jit.dy2static.convert_call_func.convert_call(
+                self._serial_op
+            )(*args, **kwargs)
+        else:
+            output = self._serial_op(*args, **kwargs)
+        new_op_size = len(cur_block.ops)
+
+        if isinstance(output, (tuple, list)):
+            new_output = list(output)
+        elif isinstance(output, Variable):
+            new_output = [output]
+        else:
+            raise ValueError("Unrecognized output.")
+
+        if self._out_dims_mappings:
+            assert len(new_output) == len(self._out_dims_mappings), (
+                f"The length of dims_mapping {len(self._out_dims_mappings)} does not matching the length output {len(new_output)}."
+            )
+        for i, item in enumerate(new_output):
+            if isinstance(item, Variable) and self._out_dims_mappings:
+                tensor_to_dims_mapping[item.name] = self._out_dims_mappings[i]
+
+        from .dist_context import get_default_distributed_context
+
+        default_dist_ctx = get_default_distributed_context()
+        for idx in range(op_size, new_op_size):
+            op = cur_block.ops[idx]
+            dist_op = DistributedOperator(op)
+            for name in dist_op.serial_op.input_arg_names:
+                if name in tensor_to_dims_mapping.keys():
+                    tensor = dist_op.get_serial_input(name)
+                    tensor_dist_attr = dist_op.dist_attr.get_input_dist_attr(
+                        name
+                    )
+                    dims_mapping = tensor_to_dims_mapping[name]
+                    if tensor is None:
+                        tensor_shape = []
+                    else:
+                        if tensor.type in __no_shape_var_type__:
+                            tensor_shape = []
+                        else:
+                            tensor_shape = tensor.shape
+                    if dims_mapping is not None:
+                        dims_mapping = tensor_to_dims_mapping[name]
+                        shard_spec = convert_to_shard_spec(
+                            dims_mapping, self._process_mesh
+                        )
+                        assert verify_shard_spec(
+                            shard_spec, tensor_shape, self._process_mesh
+                        ), (
+                            f"For tensor {name}, shard_spec {shard_spec} is invalid with tensor_shape {tensor_shape} and process_mesh {self._process_mesh}."
+                        )
+                        tensor_dist_attr.dims_mapping = dims_mapping
+                        tensor_dist_attr.mark_annotated("dims_mapping")
+            for name in dist_op.serial_op.output_arg_names:
+                if name in tensor_to_dims_mapping.keys():
+                    tensor = dist_op.get_serial_output(name)
+                    tensor_dist_attr = dist_op.dist_attr.get_output_dist_attr(
+                        name
+                    )
+                    dims_mapping = tensor_to_dims_mapping[name]
+                    if tensor is None:
+                        tensor_shape = []
+                    else:
+                        if tensor.type in __no_shape_var_type__:
+                            tensor_shape = []
+                        else:
+                            tensor_shape = tensor.shape
+                    if dims_mapping is not None:
+                        dims_mapping = tensor_to_dims_mapping[name]
+                        shard_spec = convert_to_shard_spec(
+                            dims_mapping, self._process_mesh
+                        )
+                        assert verify_shard_spec(
+                            shard_spec, tensor_shape, self._process_mesh
+                        ), (
+                            f"For tensor {name}, shard_spec {shard_spec} is invalid with tensor_shape {tensor_shape} and process_mesh {self._process_mesh}."
+                        )
+                        tensor_dist_attr.dims_mapping = dims_mapping
+                        tensor_dist_attr.mark_annotated("dims_mapping")
+            dist_op.dist_attr.process_mesh = self._process_mesh
+            dist_op.dist_attr.chunk_id = self._chunk_id
+            if self._process_mesh is not None:
+                dist_op.dist_attr.mark_annotated("process_mesh")
+            default_dist_ctx.add_dist_op_for_program(dist_op)
+            default_dist_ctx.add_process_mesh(self._process_mesh)
+
+        return output

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

@@ -0,0 +1,261 @@
+# 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 errno
+import logging
+import os
+import pickle
+import re
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+
+from ...utils.log_utils import get_logger
+from .process_group import _g_process_group_map
+from .utils import get_dist_attr
+
+
+def check_filename(re_exp, filename):
+    if re.search(re_exp, filename):
+        return True
+    else:
+        return False
+
+
+def _process_path(path):
+    filename = os.path.basename(path)
+    if filename == "":
+        raise ValueError(
+            "path should be of 'dirname/filename' format, but received filename is empty string"
+        )
+    try:
+        dirname = os.path.dirname(path)
+        os.makedirs(dirname)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+    return dirname, filename
+
+
+class DistributedSaver:
+    def __init__(self):
+        self._logger = get_logger(logging.INFO)
+
+    def save(self, path, serial_program, dist_main_program, dist_context):
+        def _save_state(program, path, mode="param"):
+            state = {
+                k: np.array(v) for k, v in program.state_dict(mode).items()
+            }
+            with open(path, "wb") as f:
+                pickle.dump(state, f)
+
+        dirname, filename = _process_path(path)
+
+        rank_id = paddle.distributed.get_rank()
+        # save serial program when rank id is 0
+        if rank_id == 0:
+            self._save_rank_mapping(dirname)
+            serial_model_filename = filename + "_serial.pdmodel"
+            serial_model_path = os.path.join(dirname, serial_model_filename)
+            with open(serial_model_path, "wb") as f:
+                f.write(serial_program.desc.serialize_to_string())
+
+        # save distributed main program
+        dist_model_filename = filename + "_dist" + str(rank_id) + ".pdmodel"
+        dist_model_path = os.path.join(dirname, dist_model_filename)
+        with open(dist_model_path, "wb") as f:
+            f.write(dist_main_program.desc.serialize_to_string())
+
+        # save distributed attribute
+        dist_attr_filename = filename + "_dist" + str(rank_id) + ".pdattr"
+        dist_attr_path = os.path.join(dirname, dist_attr_filename)
+        dist_attrs = get_dist_attr(dist_main_program, dist_context)
+        with open(dist_attr_path, "wb") as f:
+            pickle.dump(dist_attrs, f)
+
+        # save distributed params
+        dist_param_filename = filename + "_dist" + str(rank_id) + ".pdparams"
+        dist_param_path = os.path.join(dirname, dist_param_filename)
+        _save_state(dist_main_program, dist_param_path)
+
+        # save distributed opt states
+        dist_opt_filename = filename + "_dist" + str(rank_id) + ".pdopt"
+        dist_opt_path = os.path.join(dirname, dist_opt_filename)
+        _save_state(dist_main_program, dist_opt_path, "opt")
+
+        # TODO:save cluster.json
+
+    def load(self, path, load_optimizer=True):
+        # TODO: if `program` is None, load `path.pdmodel`.
+        def _load_file(filename, dirname, suffix="pdparams"):
+            file_list = []
+            for file in os.listdir(dirname):
+                if check_filename(f'{filename}(.*)_dist(.*).{suffix}', file):
+                    file_list.append(os.path.join(dirname, file))
+            file_list.sort()
+            return file_list
+
+        def _load_state(filename, dirname, suffix="pdparams"):
+            file_list = _load_file(filename, dirname, suffix)
+            state_dict = {}
+            for file in file_list:
+                with open(file, 'rb') as f:
+                    from paddle.framework.restricted_unpickler import (
+                        safe_load_pickle,
+                    )
+
+                    state_dict_info = safe_load_pickle(f, encoding='latin1')
+                for name, value in state_dict_info.items():
+                    if name in state_dict:
+                        state_dict[name].append(np.array(value))
+                    else:
+                        state_dict[name] = [np.array(value)]
+            self._logger.info(f"Load param file: {file_list}")
+            return state_dict
+
+        filename = os.path.basename(path)
+        if filename == "":
+            raise ValueError(
+                "path should be of 'dirname/filename' format, but received filename is empty string"
+            )
+        dirname = os.path.dirname(path)
+
+        # load path.pdparam and path.pdopt
+        param_state_dict = _load_state(filename, dirname)
+        opt_state_dict = (
+            _load_state(filename, dirname, "pdopt") if load_optimizer else {}
+        )
+        state_dict = dict(param_state_dict, **opt_state_dict)
+
+        # load path.pdattr
+        dist_attr_file_list = _load_file(filename, dirname, "pdattr")
+        self._logger.info(
+            f"Load distributed attribute file: {dist_attr_file_list}"
+        )
+        dist_attr = {}
+        for dist_attr_file in dist_attr_file_list:
+            with open(dist_attr_file, 'rb') as f:
+                from paddle.framework.restricted_unpickler import (
+                    safe_load_pickle,
+                )
+
+                dist_attr_info = safe_load_pickle(f, encoding='latin1')
+            for name, attr in dist_attr_info.items():
+                if name not in dist_attr:
+                    dist_attr[name] = attr
+
+        return state_dict, dist_attr
+
+    def save_inference_model(self, path, feed_vars, fetch_vars, exe, **kwargs):
+        dirname, filename = _process_path(path)
+
+        # save distributed inference program
+        rank_id = paddle.distributed.get_rank()
+        if rank_id == 0:
+            self._save_rank_mapping(dirname)
+        op_role_key = core.op_proto_and_checker_maker.kOpRoleAttrName()
+        op_role_forward = int(core.op_proto_and_checker_maker.OpRole.Forward)
+
+        dist_main_prog = kwargs.get('program', None)
+        if not dist_main_prog:
+            dist_main_prog = paddle.static.default_main_program()
+        global_block = dist_main_prog.global_block()
+
+        ops = global_block.ops
+        feed_vars_names = [x.name for x in feed_vars]
+        fetch_vars_names = [x.name for x in fetch_vars]
+
+        last_idx = -1
+        for idx, op in enumerate(ops):
+            if op.attr(op_role_key) != op_role_forward:
+                continue
+            if op.type == "read" or op.type == "feed" or op.type == 'recv_v2':
+                feed_vars_names += op.output("Out")
+            if op.type == "send_v2":
+                fetch_vars_names += op.input("X")
+                last_idx = max(idx, last_idx)
+            for out_name in op.output_arg_names:
+                if out_name in fetch_vars_names:
+                    last_idx = max(idx, last_idx)
+
+        used_inputs = []
+        used_outputs = []
+        for idx, op in enumerate(ops):
+            if idx > last_idx:
+                break
+            used_inputs += op.input_arg_names
+            used_outputs += op.output_arg_names
+
+        # delete duplicated elements and keep order
+        feed_vars_names = list({}.fromkeys(feed_vars_names).keys())
+        used_inputs = list({}.fromkeys(used_inputs).keys())
+        fetch_vars_names = list({}.fromkeys(fetch_vars_names).keys())
+        used_outputs = list({}.fromkeys(used_outputs).keys())
+
+        dist_feed_vars_names = [
+            var_name for var_name in feed_vars_names if var_name in used_inputs
+        ]
+        dist_fetch_vars_names = [
+            var_name
+            for var_name in fetch_vars_names
+            if var_name in used_outputs
+        ]
+
+        dist_feed_vars = list(
+            reversed([global_block.vars[name] for name in dist_feed_vars_names])
+        )
+        dist_fetch_vars = [
+            global_block.vars[name] for name in dist_fetch_vars_names
+        ]
+
+        dist_filename = filename + "_dist" + str(rank_id)
+        dist_path = os.path.join(dirname, dist_filename)
+        legacy_format = kwargs.get("legacy_format", False)
+        paddle.static.save_inference_model(
+            dist_path,
+            dist_feed_vars,
+            dist_fetch_vars,
+            exe,
+            program=dist_main_prog,
+            legacy_format=legacy_format,
+        )
+
+    def _save_rank_mapping(self, dirname):
+        path = os.path.join(dirname, 'rank_mapping.csv')
+        f = open(path, 'w')
+        f.write('[ring_id -> ranks]\n')
+        for process_group in _g_process_group_map.values():
+            ring_id = process_group._group_id
+            ranks = [str(rank) for rank in process_group._ranks]
+            id_to_rank = str(ring_id) + "," + ",".join(ranks) + '\n'
+            f.write(id_to_rank)
+            id_to_rank = ""
+        f.write('[rank -> ring_ids]\n')
+        rank_to_id_dict = {}
+        for process_group in _g_process_group_map.values():
+            ring_id = process_group._group_id
+            for rank in process_group._ranks:
+                if rank in rank_to_id_dict:
+                    rank_to_id_dict[rank].append(str(ring_id))
+                else:
+                    rank_to_id_dict[rank] = [str(ring_id)]
+        rank_to_id = ""
+        for item, val in rank_to_id_dict.items():
+            rank_to_id += str(item) + ","
+            rank_to_id += ",".join(val) + "\n"
+            f.write(rank_to_id)
+            rank_to_id = ""
+        f.close()

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

@@ -0,0 +1,412 @@
+#   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 inspect
+
+import paddle
+from paddle.framework import Block
+from paddle.static import Parameter, Variable
+
+from .dist_attribute import TensorDistAttr
+from .utils import __no_shape_var_type__, _linear_idx2coordinate
+
+
+class DistributedTensor:
+    """
+    DistributedTensor represents the distribution of tensor on the process group and
+    local tensors can be created by DistributedTensor.
+    Only support even sharding now and uneven sharding will be supported in the future.
+    Local tensor information can be obtained from the DistributedTensor instance object,
+    or obtained by the static methods provided by DistributedTensor,
+    including shard (i.e. the index in the serial tensor), offsets, and sizes.
+    """
+
+    @staticmethod
+    def _validate_sizes_and_dist_attr(
+        sizes, dims_mapping, topology, processes, rank=None, shard_sizes=None
+    ):
+        if not (
+            isinstance(sizes, (list, tuple))
+            and all(isinstance(x, int) and x >= 0 for x in sizes)
+        ):
+            raise ValueError(
+                f"The sizes must be list or tuple and item in sizes must be non-negative integer, but got {sizes}"
+            )
+        if not (
+            isinstance(dims_mapping, (list, tuple))
+            and all(isinstance(x, int) and x >= -1 for x in dims_mapping)
+        ):
+            raise ValueError(
+                f"The dims_mapping must be list or tuple and item in dims_mapping must >= -1, but got {dims_mapping}"
+            )
+        if not (
+            isinstance(processes, (list, tuple))
+            and all(isinstance(x, int) and x >= 0 for x in processes)
+        ):
+            raise ValueError(
+                f"The processes must be list or tuple and item in processes must be integer, but got {processes}"
+            )
+        if not (
+            isinstance(topology, (list, tuple))
+            and all(isinstance(x, int) and x > 0 for x in topology)
+        ):
+            raise ValueError(
+                f"The topology must be list or tuple and item in topology must be non-negative integer, but got {topology}"
+            )
+        if rank is not None and not (isinstance(rank, int) and rank >= 0):
+            raise ValueError(f"The rank must >= 0, but got {rank}")
+
+        # # NOTE: Only support even sharding now
+        # if shard_sizes is not None:
+        #     raise ValueError("Only support even sharding now.")
+
+    @staticmethod
+    def get_local_sizes(
+        global_sizes,
+        dims_mapping,
+        topology,
+        processes,
+        rank=None,
+        shard_sizes=None,
+    ):
+        DistributedTensor._validate_sizes_and_dist_attr(
+            global_sizes, dims_mapping, topology, processes, rank, shard_sizes
+        )
+
+        local_sizes = []
+        # for even sharding, the local sizes of every rank are equal
+
+        for idx, item in enumerate(global_sizes):
+            # This is a trick to avoid dims_mapping is []
+            val = dims_mapping[idx] if idx < len(dims_mapping) else -1
+            if val == -1:
+                local_sizes.append(item)
+            else:
+                local_sizes.append(item // topology[dims_mapping[idx]])
+
+        return local_sizes
+
+    @staticmethod
+    def get_local_offsets(
+        global_sizes, dims_mapping, topology, processes, rank, shard_sizes=None
+    ):
+        local_sizes = DistributedTensor.get_local_sizes(
+            global_sizes, dims_mapping, topology, processes, rank, shard_sizes
+        )
+        local_offsets = []
+        rank_relative = processes.index(rank)
+        coordinate = _linear_idx2coordinate(topology, rank_relative)
+
+        for i in range(len(global_sizes)):
+            if dims_mapping[i] == -1:
+                local_offsets.append(0)
+            else:
+                local_offsets.append(
+                    coordinate[dims_mapping[i]] * local_sizes[i]
+                )
+        return local_offsets
+
+    @staticmethod
+    def get_global_sizes(
+        local_sizes,
+        dims_mapping,
+        topology,
+        processes,
+        rank=None,
+        shard_sizes=None,
+    ):
+        DistributedTensor._validate_sizes_and_dist_attr(
+            local_sizes, dims_mapping, topology, processes, rank, shard_sizes
+        )
+        global_sizes = []
+        for idx, item in enumerate(local_sizes):
+            if dims_mapping[idx] == -1:
+                global_sizes.append(item)
+            else:
+                global_sizes.append(item * topology[dims_mapping[idx]])
+        return global_sizes
+
+    @staticmethod
+    def get_local_shard(
+        global_sizes, dims_mapping, topology, processes, rank, shard_sizes=None
+    ):
+        local_offsets = DistributedTensor.get_local_offsets(
+            global_sizes, dims_mapping, topology, processes, rank, shard_sizes
+        )
+        local_sizes = DistributedTensor.get_local_sizes(
+            global_sizes, dims_mapping, topology, processes, rank, shard_sizes
+        )
+        assert len(local_sizes) == len(local_offsets), (
+            f"The length of local_sizes must be equal to local_offsets, but got {len(local_sizes)} and {len(local_offsets)}."
+        )
+
+        local_end_offsets = [
+            x[0] + x[1] for x in zip(local_offsets, local_sizes)
+        ]
+        local_shard = list(zip(local_offsets, local_end_offsets))
+        return local_shard
+
+    def __init__(self, serial_tensor, dist_attr=None, dist_context=None):
+        self._serial_tensor = serial_tensor
+        if dist_attr is not None and isinstance(dist_attr, TensorDistAttr):
+            # TODO: remove this deepcopy after we fix the issue
+            self._dist_attr = copy.deepcopy(dist_attr)
+            # self._dist_attr = dist_attr
+            # TODO: Do we really need to write dist_attr back to serial_tensor？
+            self._serial_tensor.dist_attr = dist_attr
+        else:
+            assert dist_attr is None, f"{dist_attr}"
+            # Use the dist attr of serial_tensor to do the initialization
+            self._dist_attr = self._serial_tensor.dist_attr
+
+        self._batch_dim = 0
+        self._local_offsets_map = {}
+        self._local_shard_map = {}
+        self._local_tensor_map = {}
+
+        from .dist_context import get_default_distributed_context
+
+        self._dist_context = (
+            dist_context
+            if dist_context is not None
+            else get_default_distributed_context()
+        )
+        # TODO: Add Automatically to dist_context after initialized and it will be adapted in the future.
+        # self._dist_context.add_dist_tensor_for_program(self)
+
+    @property
+    def serial_tensor(self):
+        return self._serial_tensor
+
+    @property
+    def dist_attr(self):
+        return self._dist_attr
+
+    @dist_attr.setter
+    def dist_attr(self, dist_attr):
+        self._dist_attr = dist_attr
+        # TODO: Do we really need to write back dist_attr to serial_tensor？
+        self._serial_tensor.dist_attr = dist_attr
+
+    @property
+    def dist_context(self):
+        return self._dist_context
+
+    # def _init_default_dist_attr(self):
+    #     if self._dist_attr.dims_mapping is None:
+    #         if self.serial_tensor.type in __no_shape_var_type__:
+    #             tensor_shape = []
+    #         else:
+    #             tensor_shape = self._serial_tensor.shape
+    #         tensor_dims_mapping = [-1 for _ in range(len(tensor_shape))]
+    #         self._dist_attr.dims_mapping = tensor_dims_mapping
+
+    def validate_dist_attr(self):
+        if self.serial_tensor.type in __no_shape_var_type__:
+            return True
+        tensor_shape = self.serial_tensor.shape
+        if len(tensor_shape) != len(self.dist_attr.dims_mapping):
+            return False
+        for i in range(len(self.dist_attr.dims_mapping)):
+            if self.dist_attr.dims_mapping[
+                i
+            ] < -1 or self.dist_attr.dims_mapping[i] >= len(
+                self.dist_attr.process_mesh.shape
+            ):
+                return False
+        for i in range(len(self.dist_attr.process_mesh.shape)):
+            if self.dist_attr.dims_mapping.count(i) > 1:
+                return False
+        return True
+
+    def local_sizes(self, rank=None):
+        """Get local sizes of the given rank."""
+        rank = paddle.distributed.get_rank() if rank is None else rank
+        global_sizes = self.serial_tensor.shape
+        dims_mapping = self.dist_attr.dims_mapping
+        # shard_sizes = self.dist_attr.shard_sizes
+        processes = self.dist_attr.process_mesh.process_ids
+        topology = self.dist_attr.process_mesh.shape
+        local_sizes = DistributedTensor.get_local_sizes(
+            global_sizes, dims_mapping, topology, processes, rank
+        )
+
+        return local_sizes
+
+    def local_offsets(self, rank=None):
+        rank = paddle.distributed.get_rank() if rank is None else rank
+        local_offsets = None
+        if rank in self._local_offsets_map.keys():
+            local_offsets = self._local_offsets_map[rank]
+        else:
+            global_sizes = self.serial_tensor.shape
+            dims_mapping = self.dist_attr.dims_mapping
+            # shard_sizes = self.dist_attr.shard_sizes
+            processes = self.dist_attr.process_mesh.process_ids
+            topology = self.dist_attr.process_mesh.shape
+            local_offsets = DistributedTensor.get_local_offsets(
+                global_sizes, dims_mapping, topology, processes, rank
+            )
+            self._local_offsets_map[rank] = local_offsets
+
+        return local_offsets
+
+    def global_sizes(self):
+        return self.serial_tensor.shape
+
+    def local_shard(self, rank=None):
+        rank = paddle.distributed.get_rank() if rank is None else rank
+        local_shard = None
+        if rank in self._local_shard_map.keys():
+            local_shard = self._local_shard_map[rank]
+        else:
+            global_sizes = self.serial_tensor.shape
+            dims_mapping = self.dist_attr.dims_mapping
+            # shard_sizes = self.dist_attr.shard_sizes
+            processes = self.dist_attr.process_mesh.process_ids
+            topology = self.dist_attr.process_mesh.shape
+            local_shard = DistributedTensor.get_local_shard(
+                global_sizes, dims_mapping, topology, processes, rank
+            )
+            self._local_shard_map[rank] = local_shard
+
+        return local_shard
+
+    def new_local_tensor(self, block=None, rank=None, name=None):
+        """
+        Create a new local tensor of serial tensor corresponding to rank.
+        Args:
+            block (Block): The block contains the new tensor. Default value is recommend and it will be created in the block of dist main program corresponding to the serial tensor block id. Default: None.
+            rank (int): The rank id. Default value is recommend and it will be the current rank. Default: None.
+        """
+
+        def _copy_kwargs(serial_tensor):
+            kwargs = {}
+            no_need_copy_args = ["self", "block", "shape", "name"]
+            arg_spec = inspect.getfullargspec(Variable.__init__)
+
+            for key in arg_spec.args:
+                # TODO: Check the copied attribute from serial tensor whether valid
+                if key in no_need_copy_args:
+                    continue
+                elif key not in kwargs:
+                    if key == "type":
+                        kwargs[key] = serial_tensor.desc.type()
+                    elif key == "dtype":
+                        kwargs[key] = serial_tensor.desc.dtype()
+                    elif key == "lod_level":
+                        kwargs[key] = serial_tensor.desc.lod_level()
+                    elif key == "persistable":
+                        kwargs[key] = serial_tensor.desc.persistable()
+                    elif key == "stop_gradient":
+                        kwargs[key] = serial_tensor.desc.stop_gradient()
+                    elif key == "need_check_feed":
+                        kwargs[key] = serial_tensor.desc.need_check_feed()
+                    # TODO: Get capacity by framework
+                    elif key == "capacity":
+                        continue
+                    else:
+                        kwargs[key] = self.serial_tensor.__dict__[key]
+
+            if isinstance(serial_tensor, Parameter):
+                kwargs["trainable"] = serial_tensor.trainable
+                kwargs["optimize_attr"] = serial_tensor.trainable
+                kwargs["regularizer"] = serial_tensor.regularizer
+                kwargs["do_model_average"] = serial_tensor.do_model_average
+                kwargs["need_clip"] = serial_tensor.need_clip
+                kwargs["is_distributed"] = serial_tensor.is_distributed
+                kwargs["is_parameter"] = serial_tensor.is_parameter
+
+            return kwargs
+
+        if rank is not None and not (isinstance(rank, int) and rank >= 0):
+            raise ValueError(f"The rank must >= 0, but got {rank}")
+        if block is not None and not isinstance(block, Block):
+            raise TypeError(f"The block must be Block, but got {type(block)}.")
+        rank = paddle.distributed.get_rank() if rank is None else rank
+
+        if block is None:
+            block_id = self.serial_tensor.block.idx
+            block = self.dist_context.dist_main_programs[rank].block(block_id)
+
+        # copy serial tensor attribute
+        kwargs = _copy_kwargs(self.serial_tensor)
+        kwargs["name"] = name
+        kwargs["shape"] = self.local_sizes(rank)
+
+        if isinstance(self.serial_tensor, Parameter):
+            kwargs.pop("persistable")
+            local_tensor = Parameter(block=block, **kwargs)
+        else:
+            local_tensor = block.create_var(**kwargs)
+
+        # TODO: Set original id when set original_id is approved
+        local_tensor.desc.set_original_id(self.serial_tensor.desc.id())
+        self._local_tensor_map[rank] = local_tensor
+        return local_tensor
+
+    def local_tensor(self, rank=None):
+        rank = paddle.distributed.get_rank() if rank is None else rank
+        assert rank in self._local_tensor_map, (
+            f"The rank {rank} local tensor has not been created."
+        )
+        return self._local_tensor_map[rank]
+
+    def __deepcopy__(self, memo):
+        cls = self.__class__
+        result = cls.__new__(cls)
+        memo[id(self)] = result
+        for k, v in self.__dict__.items():
+            if k == "_serial_tensor" or k == "_local_tensor_map":
+                setattr(result, k, v)
+            else:
+                setattr(result, k, copy.deepcopy(v, memo))
+        return result
+
+    def __str__(self):
+        str = f"{{tensor name: {self.serial_tensor.desc.name()}, tensor id: {self.serial_tensor.desc.id()}, tensor original_id {self.serial_tensor.desc.original_id()}"
+
+        # str += ", {}".format(self.dist_attr)
+        # return str
+
+        if self.dist_attr.is_annotated("process_mesh"):
+            annotated_str = "annotated"
+        else:
+            annotated_str = "non-annotated"
+        str += (
+            f", process_mesh ({annotated_str}): {self.dist_attr.process_mesh}"
+        )
+
+        str += f", is_parameter: {self.serial_tensor.is_parameter}"
+        str += f", chunk_id: {self.dist_attr.chunk_id}"
+
+        if self.dist_attr.is_annotated("dims_mapping"):
+            annotated_str = "annotated"
+        else:
+            annotated_str = "non-annotated"
+        str += f", dims_mapping ({annotated_str}): {self.dist_attr.dims_mapping} }}"
+
+        # if self.dist_attr.is_annotated("shard_mask"):
+        #     annotated_str = "annotated"
+        # else:
+        #     annotated_str = "non-annotated"
+        # str += ", shard_mask ({}): {}".format(annotated_str, None)
+
+        # if self.dist_attr.is_annotated("offload_device"):
+        #     annotated_str = "annotated"
+        # else:
+        #     annotated_str = "non-annotated"
+        # str += ", offload_device ({}): {} }}".format(annotated_str, None)
+        return str

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

@@ -0,0 +1,186 @@
+#   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 collections import OrderedDict
+
+
+class Node:
+    def __init__(self, id, **attrs):
+        # Each node must has a unique id
+        self._id = id
+        # Attributes for Node
+        self._attrs = {}
+        self._attrs.update(attrs)
+
+    @property
+    def id(self):
+        return self._id
+
+    @property
+    def attrs(self):
+        return self._attrs
+
+    def __getitem__(self, attr_name):
+        return self._attrs[attr_name]
+
+    def __setitem__(self, attr_name, attr_value):
+        self._attrs[attr_name] = attr_value
+
+    def __contains__(self, attr_name):
+        try:
+            return attr_name in self._attrs
+        except TypeError:
+            return False
+
+    def __str__(self):
+        str = f"(id: {self.id}, attrs: {self.attrs})"
+        return str
+
+
+class Edge:
+    def __init__(self, src_id, tgt_id, **attrs):
+        # The id of source node in an Edge
+        self._src_id = src_id
+        # The id of target node in an Edge
+        self._tgt_id = tgt_id
+        # Attributes for Edge
+        self._attrs = {}
+        self._attrs.update(attrs)
+
+    @property
+    def src_id(self):
+        return self._src_id
+
+    @property
+    def tgt_id(self):
+        return self._tgt_id
+
+    @property
+    def attrs(self):
+        return self._attrs
+
+    def __getitem__(self, attr_name):
+        return self._attrs[attr_name]
+
+    def __setitem__(self, attr_name, attr_value):
+        self._attrs[attr_name] = attr_value
+
+    def __contains__(self, attr_name):
+        try:
+            return attr_name in self._attrs
+        except TypeError:
+            return False
+
+    def __str__(self):
+        str = ""
+        str += f"(src_id: {self.src_id}, tgt_id: {self.tgt_id}, attrs: {self._attrs})"
+        return str
+
+
+class Graph:
+    def __init__(self, **attrs):
+        # _nodes is dict for storing the nodes of the graph.
+        # The key of this dict is the node id.
+        self._nodes = {}
+        # _adjs is a dict of dict for storing the adjacency of the graph.
+        # The key of the outer dict is the node id of the source node and
+        # the key of the inner dict is the node id of the target node.
+        self._adjs = {}
+        # Attributes for Graph
+        self._attrs = {}
+        self._attrs.update(attrs)
+        self._reverse_adjs = {}
+        self._attr_to_nodes = {}
+
+    @property
+    def nodes(self):
+        return self._nodes
+
+    @property
+    def attrs(self):
+        return self._attrs
+
+    @property
+    def adjs(self):
+        return self._adjs
+
+    def add_node(self, node_id, **attrs):
+        if node_id is None:
+            raise ValueError("None cannot be a node")
+        if node_id not in self._nodes:
+            node = Node(node_id, **attrs)
+            self._nodes[node_id] = node
+            self._adjs[node_id] = {}
+            self._reverse_adjs[node_id] = []
+        else:
+            self._nodes[node_id].attrs.update(attrs)
+
+        return self._nodes[node_id]
+
+    def add_edge(self, src_id, tgt_id, **attrs):
+        # add nodes
+        if src_id is None:
+            raise ValueError("None cannot be a node")
+        if tgt_id is None:
+            raise ValueError("None cannot be a node")
+        if src_id not in self._nodes:
+            src_node = Node(src_id)
+            self._nodes[src_id] = src_node
+            # for one tensor to multiple ops
+            self._adjs[src_id] = OrderedDict()
+            self._reverse_adjs[src_id] = []
+        if tgt_id not in self._nodes:
+            tgt_node = Node(tgt_id)
+            self._nodes[tgt_id] = tgt_node
+            # for one tensor to multiple ops
+            self._adjs[tgt_id] = OrderedDict()
+            self._reverse_adjs[tgt_id] = []
+        # add the edge
+        edge = Edge(src_id, tgt_id, **attrs)
+        self._adjs[src_id][tgt_id] = edge
+
+        # add the reverse adj
+        self._reverse_adjs[tgt_id].append(self.nodes[src_id])
+        return edge
+
+    def __len__(self):
+        return len(self._nodes)
+
+    def __iter__(self):
+        return iter(self._nodes.values())
+
+    def __getitem__(self, node_id):
+        # Return the adjacency of a node
+        return self._adjs[node_id]
+
+    def __contains__(self, node_id):
+        # Check whether a node in the graph
+        try:
+            return node_id in self._nodes
+        except TypeError:
+            return False
+
+    def __str__(self):
+        str = ""
+        str += "**************Nodes**************\n"
+        for node_id in self.nodes:
+            str += f"{self.nodes[node_id]}\n"
+
+        str += "**************Edges**************\n"
+        for src_id in self.adjs:
+            str += f"--------------{src_id}--------------\n"
+            for idx, tgt_id in enumerate(self.adjs[src_id]):
+                str += f"{self.adjs[src_id][tgt_id]}\n"
+
+        return str

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

@@ -0,0 +1,673 @@
+# 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
+import inspect
+import logging
+from collections import defaultdict
+
+import paddle
+from paddle import core
+from paddle.jit import not_to_static, to_static
+from paddle.jit.dy2static.program_translator import (
+    ProgramTranslator,
+    StaticFunction,
+)
+from paddle.jit.dy2static.utils import as_not_paddle_func
+from paddle.nn import Layer
+from paddle.static import Parameter, global_scope, program_guard
+from paddle.static.amp.fp16_utils import (
+    DEFAULT_AMP_OPTIONS,
+    prepare_op_amp_options,
+)
+
+from .converter import Converter
+from .dist_attribute import TensorDistAttr
+from .process_group import get_world_process_group
+from .utils import get_logger, to_list
+
+
+class ProxyLayer(Layer):
+    """
+    ProxyLayer implements all logic for converting dygraph model into
+    static Program IR. Meanwhile, it provides conventional interfaces for
+    auto parallel to visit feed/fetch/loss/metric variables.
+    """
+
+    def __init__(self, layer, loss_func, metrics):
+        super().__init__()
+        # NOTE: All verify logics are finished in Engine.Prepare
+        self.inner_layer = layer
+        self.loss_func = loss_func
+        self.metrics = metrics
+        # train / eval / predict
+        self.mode = None
+
+        # generated program vars
+        self._input_vars = defaultdict(list)
+        self._label_vars = defaultdict(list)
+        self._output_vars = defaultdict(list)
+        self._loss_vars = defaultdict(list)
+        self._loss_names = defaultdict(list)
+        self._metric_vars = defaultdict(list)
+
+        # Consider ProxyLayer as not Paddle inner function because it contains
+        # user-defined layer.
+        for fn_name in [
+            "_train",
+            "_eval",
+            "_predict",
+            "call_loss",
+            "call_metrics",
+        ]:
+            as_not_paddle_func(
+                f"{inspect.getmodule(ProxyLayer).__name__}.ProxyLayer.{fn_name}"
+            )
+
+    @paddle.jit.not_to_static
+    def append_loss_to_shadow_output(self, mode):
+        name = paddle.utils.unique_name.generate('loss')
+        paddle._C_ops.set_persistable_value(self._loss_vars[mode], name)
+        self._loss_names[mode] = name
+
+    def _train(self, inputs, labels):
+        """
+        Train process of inner_layer with forward/loss/metric logic.
+        """
+        # step 1. save feed variables of Program
+        mode = 'train'
+        self._input_vars[mode] = inputs
+        self._label_vars[mode] = labels
+
+        # step 2. call inner_layer.forward
+        self._output_vars[mode] = self.inner_layer(*inputs)
+
+        # step 3. calculate loss if needed
+        new_inputs = self._prepare(self.output_vars, labels)
+        self._loss_vars[mode] = self.call_loss(new_inputs)
+        if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
+            "FLAGS_enable_pir_api"
+        ]:
+            self.append_loss_to_shadow_output(mode)
+
+        # step 4. calculate metrics if needed
+        self._metric_vars[mode] = self.call_metrics(new_inputs)
+
+    def _eval(self, inputs, labels):
+        """
+        Evaluate process of inner_layer with forward/loss/metric logic.
+        """
+        # TODO(dev): we can reuse codes with self._train after making
+        # sure if they can.
+
+        # step 1. save feed variables of Program
+        mode = 'eval'
+        self._input_vars[mode] = inputs
+        self._label_vars[mode] = labels
+
+        # step 2. call inner_layer.forward
+        self._output_vars[mode] = self.inner_layer(*inputs)
+
+        # step 3. calculate loss if needed
+        new_inputs = self._prepare(self.output_vars, labels)
+        self._loss_vars[mode] = self.call_loss(new_inputs)
+        if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
+            "FLAGS_enable_pir_api"
+        ]:
+            self.append_loss_to_shadow_output(mode)
+
+        # step 4. calculate metrics if needed
+        self._metric_vars[mode] = self.call_metrics(new_inputs)
+
+    def _predict(self, inputs, labels):
+        """
+        Predict process of inner_layer with forward logic.
+        """
+        # step 1. save feed variables of Program
+        mode = 'predict'
+        self._input_vars[mode] = inputs
+        self._label_vars[mode] = labels
+
+        # step 2. call inner_layer.forward
+        self._output_vars[mode] = self.inner_layer(*inputs)
+
+    @not_to_static
+    def _prepare(self, outputs, labels):
+        """
+        Concat outputs and labels as a single list
+
+        NOTE(dev): We use @not_to_static to avoid AST Analysis.
+        """
+        return to_list(outputs) + to_list(labels)
+
+    def call_loss(self, inputs):
+        """
+        Apply Loss Function on outputs and labels.
+
+        Args:
+            inputs: List[Variable]
+
+        Returns: List[Variable]
+        """
+        res = []
+        if self.loss_func is not None:
+            res = self.loss_func(*inputs)
+        return res
+
+    def call_metrics(self, inputs):
+        """
+        Apply Metrics Function on outputs and labels.
+
+        Args:
+            inputs: List[Variable]
+
+        Returns: List[Variable]
+        """
+        outs = []
+        for metric in self.metrics:
+            outs.append(to_list(metric.compute(*inputs)))
+
+        return outs
+
+    def set_mode(self, mode):
+        self.mode = mode
+        self.training = mode == 'train'
+
+    def clone(self):
+        return ProxyLayer(self.inner_layer, self.loss_func, self.metrics)
+
+    @property
+    def input_vars(self):
+        return self._input_vars[self.mode]
+
+    @property
+    def label_vars(self):
+        return self._label_vars[self.mode]
+
+    @property
+    def output_vars(self):
+        return self._output_vars[self.mode]
+
+    @property
+    def loss_vars(self):
+        return self._loss_vars[self.mode]
+
+    @property
+    def loss_names(self):
+        return self._loss_names[self.mode]
+
+    @property
+    def metric_vars(self):
+        return self._metric_vars[self.mode]
+
+    @property
+    def startup_program(self):
+        return self.inner_layer._startup_program()
+
+
+class BuildInfo:
+    def __init__(self):
+        self.clear()
+
+    def has_cache(self, mode, update=False):
+        is_cache = self.states[mode]
+        if update:
+            self.cache(mode)
+        return is_cache
+
+    def cache(self, mode):
+        self.states[mode] = True
+
+    def clear(self):
+        self.states = defaultdict(bool)
+
+
+class ProgramHelper:
+    """
+    A Helper class for Engine to provides different Program IR according specified 'mode'.
+    """
+
+    def __init__(self, layer, loss_func, metrics, inputs_spec, labels_spec):
+        # original model config information
+        # TODO(Aurelius84): Implement append_backward and optimizer in ProxyLayer
+        # after distribute engine satisfy basic condition.
+        self.proxy_layer = ProxyLayer(layer, loss_func, metrics)
+        self.inputs_spec = inputs_spec
+        self.labels_spec = labels_spec
+
+        self.build_info = BuildInfo()
+        self._logger = get_logger(logging.INFO)
+        self.lazy_init = False
+        self._all_params_dist_attr = {}
+
+    def reset(self):
+        """
+        Reset all state of current Object.
+        """
+        self.build_info.clear()
+        self.proxy_layer = self.proxy_layer.clone()
+
+    def build_program(self, mode):
+        """
+        Convert dygraph model into static Program IR.
+        """
+        assert mode in ['train', 'eval', 'predict']
+        self.proxy_layer.set_mode(mode)
+        # skip if we has already built program.
+        if self.build_info.has_cache(mode, True):
+            self._logger.info(
+                f"Already build program with mode = {mode}, use cached program."
+            )
+            return
+
+        self._logger.info(f"start to build program for mode = {mode}.")
+        input_spec = [self.inputs_spec, self.labels_spec]
+        static_func = to_static(
+            self.static_func(), input_spec=input_spec, full_graph=True
+        )
+
+        func_name = '_' + mode
+        setattr(self.proxy_layer, func_name, static_func)
+
+        # NOTE(dev): Because @to_static is a Lazy mechanism, so we explicitly call this to trigger
+        # generating Program IR immediately.
+        concrete_program = getattr(self.proxy_layer, func_name).concrete_program
+
+        # TODO(zhiqiu): prepare_op_amp_options is not supported for PIR program
+        # It will to use dynamic-static unified amp in pir program, and there is
+        # no need to fit for prepare_op_amp_options
+        if not paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
+            "FLAGS_enable_pir_api"
+        ]:
+            prepare_op_amp_options(
+                concrete_program.main_program,
+                ProgramTranslator.get_instance()._amp_records,
+                DEFAULT_AMP_OPTIONS,
+            )
+        self._build_startup_program()
+
+    def _build_startup_program(self):
+        """
+        Create and Sync parameters into startup program.
+        """
+        startup_program = self.startup_program
+        if len(startup_program.global_block().ops) > 1:
+            self.lazy_init = True
+            return
+
+        for param in self.concrete_program.parameters:
+            Parameter(
+                name=param.name,
+                desc=param,
+                type=param.type,
+                shape=param.shape,
+                dtype=param.dtype,
+                stop_gradient=param.stop_gradient,
+                block=startup_program.global_block(),
+            )
+
+    def apply_optimizer(self, optimizer):
+        """
+        Append backward and generate optimizer operations.
+        """
+        self._verify_optimizer(optimizer)
+        self._logger.info(
+            "start to apply optimizer: %s ", type(optimizer).__name__
+        )
+        # clear optimizer parameters
+        original_params = optimizer._parameter_list
+        optimizer._parameter_list = None
+        with program_guard(self.main_program, self.startup_program):
+            res = optimizer.minimize(self.loss_vars[0])
+
+        # restore optimizer parameters
+        optimizer._parameter_list = original_params
+        return res
+
+    def _verify_optimizer(self, optimizer):
+        assert optimizer is not None
+        assert hasattr(optimizer, "minimize"), (
+            "Optimizer must have minimize() method."
+        )
+        assert self.proxy_layer.mode == 'train', (
+            f"Required mode == 'train', but received '{self.proxy_layer.mode}'"
+        )
+        assert len(self.loss_vars) == 1, (
+            f"Required len(loss_vars) == 1, but received len(loss_vars) = {len(self.loss_vars)}"
+        )
+
+    def to(self, mode):
+        """
+        Switch underly proxy layer mode into target mode.
+        """
+        assert mode in ['train', 'eval', 'predict']
+        func = getattr(self.proxy_layer, '_' + mode)
+        assert isinstance(func, StaticFunction), (
+            "Please call build_program(mode) firstly."
+        )
+        self.proxy_layer.set_mode(mode)
+
+    def static_func(self):
+        """
+        Return StaticFunction instance with underly target mode.
+        """
+        assert self.proxy_layer.mode in [
+            'train',
+            'eval',
+            'predict',
+        ], "Please call build_program(mode) firstly."
+        func_name = '_' + self.proxy_layer.mode
+        return getattr(self.proxy_layer, func_name)
+
+    def init_pir(self, main_program, place):
+        # collect all params in current dist program
+        param_values = main_program.global_block().all_parameters()
+        value_name_to_value = {}
+        dy_param_name_to_pir_param_name = {}
+        for value in param_values:
+            value_name_to_value[value.name] = value
+
+        dy_params = self.concrete_program.parameters[0]
+        pir_param = self.concrete_program.parameters[1]
+
+        for i in range(len(pir_param)):
+            if pir_param[i].name in value_name_to_value:
+                dy_param_name_to_pir_param_name[dy_params[i].name] = pir_param[
+                    i
+                ].name
+
+        is_comm = False
+        for param in dy_params:
+            if param.is_dist():
+                process_mesh, dims_mapping = self._all_params_dist_attr[
+                    param.name
+                ]
+                var_dist_attr = TensorDistAttr()
+                var_dist_attr.process_mesh = process_mesh
+                var_dist_attr.dims_mapping = dims_mapping
+                is_comm = True
+                with paddle.no_grad():
+                    tmp = paddle.base.core.reshard(param, var_dist_attr)
+                if tmp._is_initialized():
+                    param.get_tensor()._share_data_with(tmp.get_tensor())
+                else:
+                    # Only setting the "param" to "None" can't release the memory
+                    param.get_tensor()._clear()
+                    param = None
+
+            # create var in scope and share parameters to scope
+            if param is None:
+                continue
+            if param.name not in dy_param_name_to_pir_param_name:
+                # Release the redundant params
+                param.get_tensor()._clear()
+                continue
+            if not param._is_initialized():
+                continue
+            if param.is_dense():
+                value_name = dy_param_name_to_pir_param_name[param.name]
+                value = value_name_to_value[value_name]
+                # get param_var's dist_attr
+                assert value.is_dist_dense_tensor_type(), (
+                    f"param [{value.name}] is not dist tensor type"
+                )
+                dist_attr = {
+                    "dims_mapping": value.dist_attr().dims_mapping,
+                    "process_shape": value.dist_attr().process_mesh.shape,
+                    "process_group": value.dist_attr().process_mesh.process_ids,
+                }
+                # slice param_value with dist_attr
+                # share sliced_param_value with param_tensor in global_scope
+                pir_scope_param = global_scope().var(value_name).get_tensor()
+                sliced_param = Converter.slice_with_dist_attr(
+                    param.numpy(), dist_attr
+                )
+                pir_scope_param.set(sliced_param, place)
+                param.get_tensor()._clear()
+
+            elif param.is_dist():
+                value_name = dy_param_name_to_pir_param_name[param.name]
+                value = value_name_to_value[value_name]
+                # assert value.is_dist_dense_tensor_type(), "param [{}] is not dist tensor type".format(value.name)
+                pir_scope_param = global_scope().var(value_name).get_tensor()
+                pir_scope_param._share_data_with(
+                    param.get_tensor().get_tensor()
+                )
+                param.get_tensor()._clear()
+
+        world_group = get_world_process_group()
+        if (
+            is_comm
+            and world_group.nranks > 1
+            and paddle.distributed.get_world_size() > 1
+        ):
+            paddle.disable_static()
+            barrier_tensor = paddle.full([1], 1, dtype="int32")
+            # barrier is not available in xpu for now
+            if not paddle.framework.core.is_compiled_with_xpu():
+                paddle._legacy_C_ops.barrier(
+                    barrier_tensor, barrier_tensor, 'ring_id', 0
+                )
+            paddle.enable_static()
+
+    def init(self, main_program, place, dist_context):
+        if self.lazy_init:
+            return
+
+        amp_strategy = dist_context.strategy.amp
+        amp_config = copy.deepcopy(amp_strategy.to_dict())
+        need_cast_parameter = amp_strategy.enable and amp_config["level"] in [
+            "o2",
+            "o3",
+        ]
+        is_comm = False
+        for param in self.concrete_program.parameters:
+            if param.is_dist():
+                serial_main_program = self.concrete_program.main_program
+                var = serial_main_program.global_block().vars[param.name]
+                var_dist_attr = dist_context.get_tensor_dist_attr_for_program(
+                    var
+                )
+                is_comm = True
+                # No need to construct backward.
+                with paddle.no_grad():
+                    tmp = paddle.base.core.reshard(param, var_dist_attr)
+                if tmp._is_initialized():
+                    param.get_tensor()._share_data_with(tmp.get_tensor())
+                else:
+                    # Only setting the "param" to "None" can't release the memory
+                    param.get_tensor()._clear()
+                    param = None
+            paddle.device.synchronize()
+
+            # create var in scope and share parameters to scope
+            if param is None:
+                continue
+            if param.name not in main_program.global_block().vars:
+                # Release the redundant params
+                param.get_tensor()._clear()
+                continue
+            if not param._is_initialized():
+                continue
+            if param.is_dense():
+                # get param_var's dist_attr
+                var = main_program.global_block().vars[param.name]
+                var_dist_attr = dist_context.get_tensor_dist_attr_for_program(
+                    var
+                )
+                dist_attr = {
+                    "dims_mapping": var_dist_attr.dims_mapping,
+                    "process_shape": var_dist_attr.process_mesh.shape,
+                    "process_group": var_dist_attr.process_mesh.process_ids,
+                }
+                # slice param_value with dist_attr
+                # share sliced_param_value with param_tensor in global_scope
+                param_tensor = global_scope().var(param.name).get_tensor()
+                sliced_param = Converter.slice_with_dist_attr(
+                    param.numpy(), dist_attr
+                )
+                param_tensor.set(sliced_param, place)
+                if not need_cast_parameter:
+                    param.get_tensor()._clear()
+            elif param.is_dist():
+                dense_tensor = global_scope().var(param.name).get_tensor()
+                dense_tensor._share_data_with(param.get_tensor().get_tensor())
+
+        # transform the parameter in eager mode for amp.
+        if need_cast_parameter:
+            for param in self.concrete_program.parameters:
+                amp_dtype = amp_config["dtype"]
+                scope_var = global_scope().find_var(param.name)
+                # The parameter is not in this rank.
+                if not scope_var:
+                    continue
+                # The parameter do not need to transform
+                if param.dtype in [paddle.float16, paddle.bfloat16]:
+                    continue
+                scope_tensor = global_scope().var(param.name).get_tensor()
+                assert scope_var and scope_tensor._is_initialized(), (
+                    f"Parameter: {param.name} is not put into global_scope or not initialized."
+                )
+                param_used = param
+                # For the params without dist_attr.
+                # NOTE(lizhiyu): In principle, each param should have dist_attr.
+                if param.is_dense():
+                    # get param_var's dist_attr
+                    var = main_program.global_block().vars[param.name]
+                    var_dist_attr = (
+                        dist_context.get_tensor_dist_attr_for_program(var)
+                    )
+                    dist_attr = {
+                        "dims_mapping": var_dist_attr.dims_mapping,
+                        "process_shape": var_dist_attr.process_mesh.shape,
+                        "process_group": var_dist_attr.process_mesh.process_ids,
+                    }
+                    # slice param_value with dist_attr
+                    sliced_param = Converter.slice_with_dist_attr(
+                        param.numpy(), dist_attr
+                    )
+                    with paddle.base.dygraph.guard():
+                        param_used = paddle.to_tensor(
+                            sliced_param, place=param.place
+                        )
+                    param.get_tensor()._clear()
+                with paddle.base.dygraph.guard():
+                    if amp_dtype == "float16":
+                        with (
+                            paddle.no_grad(),
+                            paddle.base.framework._dygraph_place_guard(
+                                place=place
+                            ),
+                        ):
+                            t_casted = param_used.cast(
+                                dtype=core.VarDesc.VarType.FP16
+                            )
+                    elif amp_dtype == "bfloat16":
+                        with (
+                            paddle.no_grad(),
+                            paddle.base.framework._dygraph_place_guard(
+                                place=place
+                            ),
+                        ):
+                            t_casted = param_used.cast(
+                                dtype=core.VarDesc.VarType.BF16
+                            )
+                    # NOTE(lizhiyu): Clear the origin param. Don't use `param_used.get_tensor().get_tensor()._clear()` to
+                    #                clear the `DistTensor`, because it can't clear the `_holder`,
+                    #                which `param_used.get_tensor().get_tensor()` will copy one `DenseTensor`.
+                    param_used.get_tensor()._clear()
+                    if t_casted.is_dist():
+                        scope_tensor._share_data_with(
+                            t_casted.get_tensor().get_tensor()
+                        )
+                    else:
+                        scope_tensor._share_data_with(t_casted.get_tensor())
+
+        world_group = get_world_process_group()
+        if (
+            is_comm
+            and world_group.nranks > 1
+            and paddle.distributed.get_world_size() > 1
+        ):
+            paddle.disable_static()
+            barrier_tensor = paddle.full([1], 1, dtype="int32")
+            # barrier is not available in xpu for now
+            if not paddle.framework.core.is_compiled_with_xpu():
+                paddle._legacy_C_ops.barrier(
+                    barrier_tensor, barrier_tensor, 'ring_id', 0
+                )
+            paddle.enable_static()
+
+    def cache_whole_graph_dist_attr(self, all_params):
+        for param_value in all_params:
+            dist_attr = param_value.dist_attr()
+            if dist_attr:
+                process_mesh = dist_attr.process_mesh
+                dims_mapping = dist_attr.dims_mapping
+                self._all_params_dist_attr[param_value.name] = [
+                    process_mesh,
+                    dims_mapping,
+                ]
+
+    @property
+    def concrete_program(self):
+        return self.static_func().concrete_program
+
+    @property
+    def main_program(self):
+        return self.concrete_program.main_program
+
+    @property
+    def startup_program(self):
+        try:
+            return self.proxy_layer.startup_program
+        except Exception as err:
+            self._logger.warning(
+                "The startup_program is not built by `lazy init`."
+            )
+            if isinstance(err, AssertionError):
+                return self.concrete_program.startup_program
+            raise err
+
+    @property
+    def input_vars(self):
+        return to_list(self.proxy_layer.input_vars)
+
+    @property
+    def output_vars(self):
+        return to_list(self.proxy_layer.output_vars)
+
+    @property
+    def label_vars(self):
+        return to_list(self.proxy_layer.label_vars)
+
+    @property
+    def loss_vars(self):
+        return to_list(self.proxy_layer.loss_vars)
+
+    @property
+    def loss_names(self):
+        return to_list(self.proxy_layer.loss_names)
+
+    @property
+    def metric_vars(self):
+        return to_list(self.proxy_layer.metric_vars)
+
+    def named_parameters(self):
+        static_func = self.static_func()
+        partial_program = static_func.get_concrete_program(
+            self.inputs_spec, self.labels_spec
+        )[-1]
+        # TODO(xiongkun): support pir in the feature.
+        return {param.name: param for param in partial_program._params}

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

@@ -0,0 +1,342 @@
+#   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 functools
+import operator
+import os
+from collections import deque
+
+import paddle
+import paddle.distributed as dist
+
+from .cluster import DeviceType
+from .graph import Graph
+from .process_group import get_process_group
+
+
+def is_collective_comm_op(op):
+    comm_list = [
+        "all_gather",
+        "all_reduce",
+        "broadcast",
+    ]
+    reduce_type = [
+        dist.ReduceOp.SUM,
+        dist.ReduceOp.MIN,
+        dist.ReduceOp.MAX,
+        dist.ReduceOp.PROD,
+    ]
+    if op.type == "reduce" and op.attr("reduce_type") in reduce_type:
+        return True
+    if op.type in comm_list:
+        return True
+    else:
+        return False
+
+
+def is_p2p_comm_op(op):
+    comm_list = ["send_v2", "recv_v2"]
+    if op.type in comm_list:
+        return True
+    else:
+        return False
+
+
+def get_dtype_bytes(dtype):
+    num_bytes = 0
+    if dtype == paddle.float64:
+        num_bytes = 8
+    elif dtype == paddle.float32:
+        num_bytes = 4
+    elif dtype == paddle.float16:
+        num_bytes = 2
+    elif dtype == paddle.bfloat16:
+        num_bytes = 2
+    elif dtype == paddle.int64:
+        num_bytes = 8
+    elif dtype == paddle.int32:
+        num_bytes = 4
+    elif dtype == paddle.int16:
+        num_bytes = 2
+    elif dtype == paddle.int8:
+        num_bytes = 1
+    elif dtype == paddle.uint8:
+        num_bytes = 1
+    else:
+        raise ValueError(f"Unrecognized dtype {dtype}.")
+    return num_bytes
+
+
+def get_comm_volume(comm_op, src_rank, tgt_rank):
+    comm_volume = None
+    if src_rank == tgt_rank:
+        return comm_volume
+    comm_op_type = comm_op.type
+    if comm_op_type != "recv_v2":
+        tensor_name = comm_op.input_arg_names[0]
+    else:
+        tensor_name = comm_op.output_arg_names[0]
+    tensor = comm_op.block._find_var_recursive(tensor_name)
+    assert tensor is not None
+    tensor_shape = tensor.shape
+    # Skip the batch dim
+    new_tensor_shape = []
+    for val in tensor_shape:
+        if val == -1:
+            print("Warning: -1 in the tensor shape.")
+            new_tensor_shape.append(1)
+        else:
+            new_tensor_shape.append(val)
+    tensor_size = functools.reduce(operator.mul, new_tensor_shape, 1)
+    tensor_bytes = tensor_size * get_dtype_bytes(tensor.dtype)
+    if "c_allreduce" in comm_op_type or "all_reduce" in comm_op_type:
+        comm_volume = 2 * tensor_bytes
+    elif "all_gather" in comm_op_type:
+        comm_volume = tensor_bytes
+    elif "broadcast" in comm_op_type:
+        if comm_op.attr("root") == src_rank:
+            comm_volume = tensor_bytes
+        else:
+            comm_volume = None
+    elif "c_reduce" in comm_op_type:
+        if comm_op.attr("root_id") == src_rank:
+            comm_volume = None
+        else:
+            comm_volume = tensor_bytes
+    elif "reduce" == comm_op_type:
+        if comm_op.attr("root_id") == src_rank:
+            comm_volume = None
+        else:
+            comm_volume = tensor_bytes
+    elif "send_v2" in comm_op_type:
+        if comm_op.attr("peer") == tgt_rank:
+            comm_volume = tensor_bytes
+        else:
+            comm_volume = None
+    elif "recv_v2" in comm_op_type:
+        comm_volume = None
+    else:
+        raise ValueError("Unrecognized communication operator.")
+    return comm_volume
+
+
+def analyze_comm_requirements_from_op(op, rank, g_process_group_map):
+    comm_requirements_to_ranks = {}
+    if is_collective_comm_op(op):
+        process_group_id = op.attr("ring_id")
+        process_group = get_process_group(process_group_id, g_process_group_map)
+        if rank not in process_group.ranks:
+            return comm_requirements_to_ranks
+        for tgt_rank in process_group.ranks:
+            comm_volume = get_comm_volume(op, rank, tgt_rank)
+            if comm_volume is not None:
+                comm_requirements_to_ranks[tgt_rank] = {}
+                comm_requirements_to_ranks[tgt_rank]["comm_volume"] = (
+                    comm_volume
+                )
+    elif is_p2p_comm_op(op):
+        tgt_rank = op.attr("peer")
+        comm_volume = get_comm_volume(op, rank, tgt_rank)
+        if comm_volume is not None:
+            comm_requirements_to_ranks[tgt_rank] = {}
+            comm_requirements_to_ranks[tgt_rank]["comm_volume"] = comm_volume
+    else:
+        comm_requirements_to_ranks = {}
+    return comm_requirements_to_ranks
+
+
+def analyze_requirements_for_program(src_info, rank):
+    program = src_info[0]
+    g_process_group_map = src_info[1]
+    resource_requirements = {}
+    comm_requirements_to_ranks = {}
+    # only support device_type and only support GPU for now
+    resource_requirements["device_type"] = DeviceType.GPU
+    for block in program.blocks:
+        for op in block.ops:
+            cur_comm_requirements_to_ranks = analyze_comm_requirements_from_op(
+                op, rank, g_process_group_map
+            )
+            for tgt_rank, link_info in cur_comm_requirements_to_ranks.items():
+                if tgt_rank in comm_requirements_to_ranks:
+                    comm_requirements_to_ranks[tgt_rank]["comm_volume"] += (
+                        link_info["comm_volume"]
+                    )
+                else:
+                    comm_requirements_to_ranks[tgt_rank] = {}
+                    comm_requirements_to_ranks[tgt_rank]["comm_volume"] = (
+                        link_info["comm_volume"]
+                    )
+    return resource_requirements, comm_requirements_to_ranks
+
+
+def build_process_graph(distributed_program):
+    graph = Graph()
+    for src_rank, src_info in distributed_program.items():
+        (
+            resource_requirements,
+            comm_requirements_to_ranks,
+        ) = analyze_requirements_for_program(src_info, src_rank)
+        graph.add_node(src_rank, resource_requirements=resource_requirements)
+        for tgt_rank, comm_requirements in comm_requirements_to_ranks.items():
+            graph.add_edge(
+                src_rank, tgt_rank, comm_requirements=comm_requirements
+            )
+    return graph
+
+
+def build_cluster_graph(cluster):
+    graph = Graph()
+    cuda_visible_devices_env = os.getenv("CUDA_VISIBLE_DEVICES")
+    cuda_visible_devices = []
+    if cuda_visible_devices_env is not None and cuda_visible_devices_env != "":
+        cuda_visible_devices = [
+            int(d.strip()) for d in cuda_visible_devices_env.split(",")
+        ]
+    for machine in cluster.machines.values():
+        for device in machine.devices.values():
+            graph.add_node(device.global_id, device=device)
+            if (
+                cuda_visible_devices
+                and device.local_id not in cuda_visible_devices
+            ):
+                graph.nodes[device.global_id]["occupied"] = True
+            else:
+                graph.nodes[device.global_id]["occupied"] = False
+        for link in machine.links.values():
+            graph.add_edge(
+                link.source.global_id, link.target.global_id, link=link
+            )
+    return graph
+
+
+def mapping(distributed_program, cluster):
+    # A very simple mapping algorithm only for GPUs.
+    # Here we assume one process will be mapped to one GPU.
+    # In the future, more mapping configurations and algorithms will be supported.
+    process_graph = build_process_graph(distributed_program)
+
+    cluster_graph = build_cluster_graph(cluster)
+
+    for cur_rank_node in process_graph:
+        cur_rank_node["visited"] = False
+
+    def sort_by_comm_volume(rank_edge):
+        return rank_edge["comm_requirements"]["comm_volume"]
+
+    def sort_by_comm_bandwidth(device_edge):
+        return device_edge["link"].bandwidth
+
+    def select_unvisited_rank_node(rank_node_list):
+        selected_rank_node = None
+        for rank_node in rank_node_list:
+            if rank_node["visited"] is False:
+                selected_rank_node = rank_node
+        return selected_rank_node
+
+    queue = deque()
+    root_rank_node = select_unvisited_rank_node(
+        list(process_graph.nodes.values())
+    )
+    while root_rank_node is not None:
+        queue.append(root_rank_node)
+        while queue:
+            cur_rank_node = queue.popleft()
+            if cur_rank_node["visited"]:
+                continue
+            device_type = cur_rank_node["resource_requirements"]["device_type"]
+            cur_device_node = None
+            for device_node in cluster_graph.nodes.values():
+                if (device_node["device"].type == device_type) and (
+                    not device_node["occupied"]
+                ):
+                    device_node["occupied"] = True
+                    cur_rank_node["visited"] = True
+                    cur_rank_node["device"] = device_node["device"]
+                    cur_device_node = device_node
+                    break
+            assert cur_device_node, (
+                "Cannot find a device to satisfy the requirement."
+            )
+
+            nbr_rank_edges = []
+            for nbr_rank_node_id, nbr_rank_edge in process_graph.adjs[
+                cur_rank_node.id
+            ].items():
+                assert (
+                    nbr_rank_edge.src_id == cur_rank_node.id
+                    and nbr_rank_edge.tgt_id == nbr_rank_node_id
+                )
+                queue.append(process_graph.nodes[nbr_rank_node_id])
+                nbr_rank_edges.append(nbr_rank_edge)
+            nbr_rank_edges.sort(key=sort_by_comm_volume)
+
+            nbr_device_edges = []
+            for nbr_device_edge in cluster_graph.adjs[
+                cur_device_node.id
+            ].values():
+                nbr_device_edges.append(nbr_device_edge)
+            nbr_device_edges.sort(key=sort_by_comm_bandwidth)
+
+            for nbr_rank_edge in nbr_rank_edges:
+                src_rank_node = process_graph.nodes[nbr_rank_edge.src_id][
+                    "visited"
+                ]
+                if src_rank_node:
+                    continue
+                device_type = src_rank_node["resource_requirements"][
+                    "device_type"
+                ]
+                nbr_rank_node = process_graph.nodes[nbr_rank_edge.tgt_id]
+                for nbr_device_edge in nbr_device_edges:
+                    nbr_device_node = cluster_graph.nodes[
+                        nbr_device_edge.tgt_id
+                    ]
+                    if (nbr_device_node["device"].type == device_type) and (
+                        not nbr_device_node["occupied"]
+                    ):
+                        nbr_device_node["occupied"] = True
+                        nbr_rank_node["visited"] = True
+                        nbr_rank_node["device"] = nbr_device_node["device"]
+                        break
+        root_rank_node = select_unvisited_rank_node(
+            list(process_graph.nodes.values())
+        )
+
+    rank_mapping = {}
+    for rank, rank_node in process_graph.nodes.items():
+        device = rank_node["device"]
+        machine = device.machine
+        if machine.id in rank_mapping:
+            rank_mapping[machine.id]["hostname"] = machine.hostname
+            rank_mapping[machine.id]["addr"] = machine.addr
+            rank_mapping[machine.id]["port"] = machine.port
+            if rank not in rank_mapping[machine.id]["ranks"]:
+                rank_mapping[machine.id]["ranks"][rank] = []
+                rank_mapping[machine.id]["ranks"][rank].append(device.local_id)
+            else:
+                rank_mapping[machine.id]["ranks"][rank].append(device.local_id)
+        else:
+            rank_mapping[machine.id] = {}
+            rank_mapping[machine.id]["hostname"] = machine.hostname
+            rank_mapping[machine.id]["addr"] = machine.addr
+            rank_mapping[machine.id]["port"] = machine.port
+            rank_mapping[machine.id]["ranks"] = {}
+            rank_mapping[machine.id]["ranks"][rank] = []
+            rank_mapping[machine.id]["ranks"][rank].append(device.local_id)
+    for machine_mapping in rank_mapping.values():
+        for rank_devices in machine_mapping["ranks"].values():
+            rank_devices.sort()
+
+    return rank_mapping

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

@@ -0,0 +1,138 @@
+# 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 paddle
+
+from .reshard_funcs.base_reshard_func import is_replicated
+from .utils import _complete_op_dist_attr
+
+dist_skip_op_list = [
+    "builtin.combine",
+    "builtin.split",
+    "cf.yield",
+    "cf.tuple_push",
+    "cf.tuple_pop",
+    "cf.stack_create",
+    "pd_op.pylayer",
+]
+
+
+def verify_dist_block(block):
+    for op in block.ops:
+        if op.name() in dist_skip_op_list:
+            continue
+        if op.name() == "dist_op.shard_tensor":
+            raise RuntimeError("Block still contain shard_tensor_op.")
+        # Note (luchang): Temp fix, remove unused parameter 'op'.
+        # Will be removed in the future.
+        if op.name() == "builtin.parameter":
+            if op.result(0).use_empty():
+                op.erase()
+                continue
+
+
+def apply_mix2dist_pass(program, block=None):
+    if block is None:
+        block = program.global_block()
+    deleted_ops = []
+    for op in block.ops:
+        for inner_block in op.blocks():
+            apply_mix2dist_pass(program, block=inner_block)
+        if op.name() != "dist_op.shard_tensor":
+            continue
+        shard_operand_value = op.operand_source(0)
+        if not shard_operand_value.has_one_use():
+            raise RuntimeError(
+                f"shard_tensor is supposed to be called right after tensor is created, the use_count of tensor to be sharded is {shard_operand_value.use_count}, which is "
+                "not Supported in right now."
+            )
+        shard_result_value = op.result(0)
+        shard_result_value.replace_all_uses_with(shard_operand_value)
+        deleted_ops.append(op)
+        prev_op = shard_operand_value.get_defining_op()
+        if (
+            prev_op.name() == "builtin.parameter"
+            or prev_op.name() == "pd_op.data"
+        ):
+            prev_op.dist_attr = op.dist_attr
+            shard_operand_value.set_type(shard_result_value.type())
+            shard_operand_value.stop_gradient = shard_result_value.stop_gradient
+            shard_operand_value.persistable = shard_result_value.persistable
+        elif (
+            prev_op.name() == "pd_op.randint"
+            or prev_op.name() == "pd_op.gaussian"
+        ):
+            mesh = shard_result_value.dist_attr().process_mesh
+            # input
+            shape_value = prev_op.operand_source(0)
+            dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+                mesh, [-1 for _ in range(len(shape_value.shape))], {}
+            )
+            shape_value.update_dist_attr(dist_attr)
+            # op
+            prev_op.dist_attr = (
+                paddle.base.libpaddle.pir.create_op_dist_attribute(
+                    mesh, [dist_attr], [shard_result_value.dist_attr()]
+                )
+            )
+            # deal with full_int_array op
+            prev_prev_op = shape_value.get_defining_op()
+            prev_prev_op.dist_attr = (
+                paddle.base.libpaddle.pir.create_op_dist_attribute(
+                    mesh, [], [dist_attr]
+                )
+            )
+            # output
+            shard_operand_value.set_type(shard_result_value.type())
+            shard_operand_value.stop_gradient = shard_result_value.stop_gradient
+            shard_operand_value.persistable = shard_result_value.persistable
+        else:
+            dist_attr = shard_result_value.dist_attr()
+            if not is_replicated(dist_attr):
+                raise RuntimeError(
+                    f"{prev_op} is not support sharded by shard_tensor op in pir mode."
+                )
+            mesh = dist_attr.process_mesh
+            ops_list = [prev_op]
+            while len(ops_list) != 0:
+                cur_op = ops_list.pop()
+                if cur_op.dist_attr is not None:
+                    continue
+                operand_attrs = []
+                result_attrs = []
+                for input in cur_op.operands_source():
+                    dist_attr = (
+                        paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+                            mesh, [-1 for _ in range(len(input.shape))], {}
+                        )
+                    )
+                    operand_attrs.append(dist_attr)
+                    ops_list.append(input.get_defining_op())
+                for result in cur_op.results():
+                    dist_attr = (
+                        paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+                            mesh, [-1 for _ in range(len(result.shape))], {}
+                        )
+                    )
+                    result.update_dist_attr(dist_attr)
+                    result_attrs.append(dist_attr)
+                cur_op.dist_attr = (
+                    paddle.base.libpaddle.pir.create_op_dist_attribute(
+                        mesh, operand_attrs, result_attrs
+                    )
+                )
+    for op in deleted_ops:
+        op.erase()
+    _complete_op_dist_attr(program, block=block)
+    verify_dist_block(block)

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"),
+)

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

@@ -0,0 +1,539 @@
+#   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 json
+import logging
+import os
+import pathlib
+import pickle
+import shlex
+import subprocess
+import sys
+import time
+
+import paddle
+from paddle.distributed.passes import PassContext, new_pass
+from paddle.distributed.utils.log_utils import get_logger
+from paddle.framework import core
+from paddle.static import append_backward, program_guard
+
+from .cluster import Cluster
+from .completion import Completer
+from .dist_context import DistributedContext, set_default_distributed_context
+from .dist_op import DistributedOperator
+from .dist_tensor import DistributedTensor
+from .mapper import mapping
+from .partitioner import Partitioner
+from .planner import Planner
+from .process_group import (
+    ProcessGroup,
+    _g_process_group_map,
+    get_all_process_groups,
+    get_process_group,
+    get_world_process_group,
+)
+from .reshard import Resharder
+from .utils import SerialProgramInfo, make_data_unshard
+
+_logger = get_logger(logging.INFO)
+
+
+class AutoParallelizer:
+    """
+    AutoParallelizer is the main controller class to do the auto parallel process.
+    And the auto parallel process will be triggered in the wrapped parallelize function.
+    To facilitate the auto parallelization, it will contain information about program, cluster and the
+    related context. In this basic version, the program information will be retrieved from
+    Fleet object, and the cluster information can be retrieved in the new created Cluster object,
+    and the context information can be retrieved in the new created DistributedContext.
+    """
+
+    def __init__(self, fleet):
+        self._fleet = fleet
+        self._optimizer = self._fleet.user_defined_optimizer
+        self._dist_strategy = self._fleet._user_defined_strategy
+        self._dist_context = DistributedContext()
+        self._cluster = None
+        self._cluster_topo_path = os.getenv("PADDLE_CLUSTER_TOPO_PATH", None)
+        if self._cluster_topo_path is not None:
+            self._cluster = Cluster()
+            self._cluster.build_from_file(self._cluster_topo_path)
+        # Prepare information for auto mapping
+        self._rank_mapping_path = os.getenv("PADDLE_RANK_MAPPING_PATH", None)
+        enable_auto_mapping_env = os.getenv("PADDLE_ENABLE_AUTO_MAPPING", None)
+        if enable_auto_mapping_env is None:
+            self._enable_auto_mapping = False
+        else:
+            self._enable_auto_mapping = True
+        self._pass_context = PassContext()
+
+        self._need_rank_mapping = os.getenv("PADDLE_NEED_RANK_MAPPING")
+        self._need_rank_mapping = (
+            True
+            if self._need_rank_mapping
+            and self._need_rank_mapping.lower() == 'true'
+            else False
+        )
+        # self._pass_context = None
+
+    def _remove_distributed_attrs(self, main_program):
+        suffix = core.kAutoParallelSuffix()
+        # distributed attributes for variable have been removed
+        # in previous process.
+        for block in main_program.blocks:
+            for op in block.ops:
+                for attr_name in op.attr_names:
+                    if suffix in attr_name:
+                        op._remove_attr(attr_name)
+
+    def _apply_pre_optimization_passes(
+        self, main_program, startup_program, loss, params_grads, no_grad_set
+    ):
+        # apply amp pass
+        if self._dist_strategy.amp:
+            config = copy.deepcopy(self._dist_strategy.amp_configs)
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            config["loss"] = loss
+            if config["use_pure_fp16"]:
+                config["base_opt"] = self._optimizer
+                auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
+                auto_parallel_fp16_pass.apply(
+                    [main_program], [startup_program], self._pass_context
+                )
+                loss = auto_parallel_fp16_pass.get_loss()
+            else:
+                auto_parallel_amp_pass = new_pass("auto_parallel_amp", config)
+                auto_parallel_amp_pass.apply(
+                    [main_program], [startup_program], self._pass_context
+                )
+                loss = auto_parallel_amp_pass.get_loss()
+
+        # apply recompute pass
+        if self._dist_strategy.recompute:
+            config = copy.deepcopy(self._dist_strategy.recompute_configs)
+            config["dist_context"] = self._dist_context
+            config["no_grad_set"] = copy.deepcopy(no_grad_set)
+            config["loss"] = loss
+            auto_parallel_recompute_pass = new_pass(
+                "auto_parallel_recompute", config
+            )
+            auto_parallel_recompute_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+
+    def _generate_backward(
+        self,
+        main_program,
+        startup_program,
+        loss,
+        parameter_list,
+        no_grad_set,
+        callbacks,
+    ):
+        with program_guard(main_program, startup_program):
+            params_grads = append_backward(
+                loss,
+                parameter_list,
+                no_grad_set,
+                callbacks,
+                distop_context=self._dist_context.dist_op_context,
+            )
+        self._completer = Completer(self._dist_context)
+        self._completer.complete_backward_annotation(main_program)
+        self._dist_context.block_state.parse_backward_blocks(main_program)
+        return params_grads
+
+    def _apply_optimize(self, main_program, startup_program, params_grads):
+        optimizer = copy.deepcopy(self._optimizer)
+        with program_guard(main_program, startup_program):
+            optimize_ops = optimizer.apply_gradients(params_grads)
+
+        self._dist_context._serial_optimizer = optimizer
+        # update completion
+        self._completer = Completer(self._dist_context)
+        self._completer.complete_update_annotation(main_program)
+
+        return optimize_ops
+
+    def _apply_post_optimization_passes(
+        self, main_program, startup_program, rank, params_grads
+    ):
+        if self._dist_strategy.sharding:
+            config = copy.deepcopy(self._dist_strategy.sharding_configs)
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            config["global_rank"] = rank
+            auto_parallel_sharding_pass = new_pass(
+                "auto_parallel_sharding", config
+            )
+            auto_parallel_sharding_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+            params_grads = self._pass_context.get_attr("params_grads")
+
+        config = copy.deepcopy(self._dist_strategy.sharding_configs)
+        config["dist_context"] = self._dist_context
+        config["params_grads"] = params_grads
+        config["rank_id"] = rank
+        auto_parallel_clip_pass = new_pass("auto_parallel_grad_clip", config)
+        auto_parallel_clip_pass.apply(
+            [main_program], [startup_program], self._pass_context
+        )
+
+        if self._dist_strategy.gradient_merge:
+            config = copy.deepcopy(self._dist_strategy.gradient_merge_configs)
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            auto_parallel_gradient_merge_pass = new_pass(
+                "auto_parallel_gradient_merge_pass", config
+            )
+            auto_parallel_gradient_merge_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+
+    def _get_dist_program(self, rank, dist_context=None, relaunch_phase=False):
+        completed_main_program = None
+        serial_main_program = self._main_program.clone()
+        serial_startup_program = self._startup_program.clone()
+        serial_loss = serial_main_program.global_block().var(self._loss.name)
+
+        # generating serial
+        if dist_context is None:
+            # Annotation completion
+            self._dist_context = DistributedContext()
+            _logger.info("Start annotation dist attr.")
+            self._completer = Completer(self._dist_context)
+            completed_main_program = (
+                self._completer.complete_forward_annotation(serial_main_program)
+            )
+        else:
+            completed_main_program = serial_main_program
+            self._dist_context = copy.deepcopy(dist_context)
+
+        # parse forward sub block
+        self._dist_context.block_state.parse_forward_blocks(serial_main_program)
+
+        # serial backward pass
+        params_grads = self._generate_backward(
+            completed_main_program,
+            serial_startup_program,
+            serial_loss,
+            self._parameter_list,
+            self._no_grad_set,
+            self._callbacks,
+        )
+
+        # serial forward pass
+        self._apply_pre_optimization_passes(
+            completed_main_program,
+            serial_startup_program,
+            serial_loss,
+            params_grads,
+            self._no_grad_set,
+        )
+        # Logical partition
+        partitioner = Partitioner(self._dist_context, rank)
+        (
+            dist_main_prog,
+            dist_startup_prog,
+            dist_params_grads,
+        ) = partitioner.partition(
+            completed_main_program, serial_startup_program, params_grads
+        )
+
+        # TODO refactor the placement of optimizer
+        # generate optimize program
+        dist_optimize_ops = self._apply_optimize(
+            dist_main_prog, dist_startup_prog, dist_params_grads
+        )
+
+        make_data_unshard(dist_main_prog, dist_startup_prog, self._dist_context)
+
+        resharder = Resharder(
+            dist_main_prog,
+            dist_startup_prog,
+            rank,
+            self._dist_context,
+            dist_params_grads,
+        )
+        resharder.reshard()
+
+        self._apply_post_optimization_passes(
+            dist_main_prog, dist_startup_prog, rank, dist_params_grads
+        )
+        g_process_group_map = None
+        if not relaunch_phase:
+            g_process_group_map = copy.deepcopy(_g_process_group_map)
+            _g_process_group_map.clear()
+            _g_process_group_map[0] = ProcessGroup(0, [])
+            for process_mesh in self._dist_context._process_meshes:
+                _g_process_group_map[0].add_ranks(process_mesh.process_ids)
+        return (
+            dist_optimize_ops,
+            dist_params_grads,
+            dist_startup_prog,
+            dist_main_prog,
+            g_process_group_map,
+        )
+
+    def parallelize(
+        self,
+        loss,
+        startup_program,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        assert startup_program is not None
+        self._loss = loss
+        self._startup_program = startup_program
+        self._main_program = loss.block.program
+        self._parameter_list = parameter_list
+        self._no_grad_set = no_grad_set
+        self._callbacks = callbacks
+
+        if self._enable_auto_mapping and self._need_rank_mapping:
+            # Do the mapping pass before parallelization
+            assert self._cluster is not None, (
+                "The cluster must not be none when using auto mapping."
+            )
+            dist_programs = {}
+            world_process_group = get_world_process_group()
+            dist_context = None
+            # auto search
+            if self._dist_strategy.auto_search:
+                logging.info("Start searching dist attr.")
+                serial_program_info = SerialProgramInfo(
+                    self._main_program,
+                    self._startup_program,
+                    self._loss,
+                    self._optimizer,
+                    self._cluster,
+                )
+                planner = Planner(
+                    serial_program_info,
+                    self,
+                    algorithm_config={"name": "mcmc", "max_search_times": 5},
+                )
+                dist_context, _ = planner.search()
+                logging.info("End searching dist attr.")
+
+            # serialize the dist context by planner
+            if dist_context is not None:
+                logging.info("Start serialize searched dist attr")
+                cwd = pathlib.Path().cwd()
+                searched_dist_context_path = os.path.join(
+                    cwd, f"searched_dist_context_{time.time()}.pkl"
+                )
+                saved_dist_context = {}
+                ops_dist_attr = {}
+                tensors_dist_attr = {}
+                for key, dist_op in dist_context._dist_ops_for_program.items():
+                    ops_dist_attr[key] = dist_op.dist_attr
+                for (
+                    key,
+                    dist_tensor,
+                ) in dist_context._dist_tensors_for_program.items():
+                    tensors_dist_attr[key] = dist_tensor.dist_attr
+                saved_dist_context["ops_dist_attr"] = ops_dist_attr
+                saved_dist_context["tensors_dist_attr"] = tensors_dist_attr
+                saved_dist_context["process_meshes"] = (
+                    dist_context._process_meshes
+                )
+                with open(
+                    searched_dist_context_path, "wb"
+                ) as dist_context_file:
+                    pickle.dump(saved_dist_context, dist_context_file)
+                    os.environ['PADDLE_SEARCHED_DIST_CONTEXT_PATH'] = (
+                        searched_dist_context_path
+                    )
+                    logging.info(
+                        f"End serialize searched dist attr to {searched_dist_context_path}"
+                    )
+
+            for rank in world_process_group.ranks:
+                (
+                    dist_optimize_ops,
+                    dist_params_grads,
+                    dist_startup_prog,
+                    dist_main_prog,
+                    g_process_group_map,
+                ) = self._get_dist_program(rank, dist_context)
+                dist_programs[rank] = [dist_main_prog, g_process_group_map]
+
+            # Do the mapping between the distributed program graph and the cluster graph
+            rank_mapping_dict = mapping(dist_programs, self._cluster)
+            rank_mapping = list(rank_mapping_dict.values())
+
+            # Relaunch the training by using the rank mapping file
+            with open(self._rank_mapping_path, "w") as rank_mapping_file:
+                json.dump(rank_mapping, rank_mapping_file)
+
+            enable_elastic = os.getenv("PADDLE_ENABLE_ELASTIC")
+            enable_elastic = (
+                True
+                if enable_elastic and enable_elastic.lower() == 'true'
+                else False
+            )
+            if enable_elastic:
+                print("Auto mapping finished, now do elastic re-launch")
+                sys.exit(
+                    paddle.distributed.fleet.elastic.manager.ELASTIC_AUTO_PARALLEL_EXIT_CODE
+                )
+
+            original_cmd_args = os.getenv("PADDLE_ORIGINAL_CMD_ARGS")
+            rank_mapping_args = " ".join(
+                ["--rank_mapping_path", self._rank_mapping_path]
+            )
+            if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
+                coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
+            else:
+                coverage_args = []
+            new_cmd_args = (
+                "-m paddle.distributed.fleet.launch"
+                + " "
+                + rank_mapping_args
+                + " "
+                + original_cmd_args
+            )
+            new_cmd = [
+                sys.executable,
+                "-u",
+                *coverage_args,
+                *shlex.split(new_cmd_args),
+            ]
+            new_process = subprocess.Popen(new_cmd)
+            new_process.wait()
+            assert new_process.returncode == 0, (
+                "Launch failed with rank mapping"
+            )
+            print("Successfully do the second launch for auto mapping!")
+            sys.exit(0)
+        else:
+            # Parallelization after the mapping pass
+            rank = paddle.distributed.get_rank()
+            dist_context = None
+            searched_dist_context_path = os.getenv(
+                "PADDLE_SEARCHED_DIST_CONTEXT_PATH", None
+            )
+            if searched_dist_context_path is not None:
+                with open(
+                    searched_dist_context_path, "rb"
+                ) as dist_context_file:
+                    from paddle.framework.restricted_unpickler import (
+                        safe_load_pickle,
+                    )
+
+                    saved_dist_context = safe_load_pickle(dist_context_file)
+                    dist_context = DistributedContext()
+                    for op in self._main_program.global_block().ops:
+                        dist_attr = saved_dist_context["ops_dist_attr"][
+                            op.desc.id()
+                        ]
+                        dist_op = DistributedOperator(op, dist_attr)
+                        dist_context.add_dist_op_for_program(dist_op)
+
+                    vars = self._main_program.global_block().vars
+                    for var in vars.values():
+                        dist_attr = saved_dist_context["tensors_dist_attr"][
+                            var.desc.id()
+                        ]
+                        dist_tensor = DistributedTensor(var, dist_attr)
+                        dist_context.add_dist_tensor_for_program(dist_tensor)
+
+                    dist_context._process_meshes = saved_dist_context[
+                        "process_meshes"
+                    ]
+
+            else:
+                if self._dist_strategy.auto_search:
+                    serial_program_info = SerialProgramInfo(
+                        self._main_program,
+                        self._startup_program,
+                        self._loss,
+                        self._optimizer,
+                        cluster=self._cluster,
+                    )
+                    planner = Planner(
+                        serial_program_info,
+                        self,
+                        algorithm_config={
+                            "name": "mcmc",
+                            "max_search_times": 5,
+                        },
+                    )
+                    dist_context, _ = planner.search()
+
+            # rebuild g_process_group
+            if dist_context is not None:
+                pg0 = get_process_group(0)
+                for process_mesh in dist_context._process_meshes:
+                    pg0.add_ranks(process_mesh.process_ids)
+            (
+                dist_optimize_ops,
+                dist_params_grads,
+                dist_startup_prog,
+                dist_main_prog,
+                _,
+            ) = self._get_dist_program(rank, dist_context, relaunch_phase=True)
+
+            # NOTE: This is a trick to fix hang in pipeline mode when dist context is searched by planner
+            if self._dist_strategy.auto_search:
+                is_pipeline = False
+                for op in dist_main_prog.global_block().ops:
+                    if op.type == "send_v2" or op.type == "recv_v2":
+                        is_pipeline = True
+                        break
+                if is_pipeline:
+                    with paddle.static.program_guard(dist_main_prog):
+                        paddle.distributed.barrier()
+
+            # Traverse different rank programs and traverse each op of them,
+            # instantiate communication by process_mapping.
+            all_process_groups = get_all_process_groups()
+            for process_group in all_process_groups:
+                process_group.instantiate()
+
+            # Copy distributed info to the default context
+            set_default_distributed_context(self._dist_context)
+
+            # The last step: remove all distributed attributes to be compatible
+            # with inference.
+            self._remove_distributed_attrs(dist_main_prog)
+
+            return (
+                dist_optimize_ops,
+                dist_params_grads,
+                dist_startup_prog,
+                dist_main_prog,
+            )
+
+    def __deepcopy__(self, memo):
+        cls = self.__class__
+        result = cls.__new__(cls)
+        memo[id(self)] = result
+        for k, v in self.__dict__.items():
+            if (
+                k == "_main_program"
+                or k == "_startup_program"
+                or k == "_dist_context"
+                or k == "_fleet"
+                or k == "_loss"
+            ):
+                setattr(result, k, v)
+            else:
+                setattr(result, k, copy.deepcopy(v, memo))
+        return result

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

@@ -0,0 +1,553 @@
+#   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
+import logging
+import os
+import time
+
+from paddle.distributed.passes.pass_base import PassManager, new_pass
+from paddle.framework import get_flags
+from paddle.static import append_backward, program_guard
+
+from ...utils.log_utils import get_logger
+from ..random import init_auto_parallel_rng
+from .partitioner import Partitioner
+from .process_group import get_world_process_group
+from .reshard import Resharder
+from .utils import (
+    get_pp_stage,
+    is_sequential_run,
+)
+
+PIR_PASS = [
+    'fused_gemm_epilogue_pass',
+    'fused_linear_param_grad_add_pass',
+    'fuse_allreduce_split_to_reducescatter_pass',
+    'fused_dropout_add_pass',
+]
+
+PIR_PYTHON_PASS = [
+    'eliminate_transpose',
+]
+
+
+class Parallelizer:
+    def __init__(self, mode, completer, dist_context):
+        self._mode = mode
+        self._completer = completer
+        self._dist_context = dist_context
+        assert self._dist_context._is_initialized
+        self._pass_context = self._dist_context.pass_context
+        self._strategy = self._dist_context.strategy
+        self._logger = get_logger(logging.INFO)
+
+    @property
+    def is_train(self):
+        return self._mode == "train"
+
+    @property
+    def is_test(self):
+        return self._mode in ["eval", "predict"]
+
+    def parallel_all(self, parameter_list=None):
+        world_process_group = get_world_process_group()
+        all_ranks = world_process_group.ranks
+        for rank in all_ranks:
+            # self._dist_context._backup(serial=True, dist=True)
+            self.parallel(rank, parameter_list)
+            # self._dist_context._restore(serial=True, dist=True)
+
+    def parallel(self, rank, parameter_list=None):
+        serial_main_program = self._dist_context.serial_main_program
+        serial_startup_program = self._dist_context.serial_startup_program
+        serial_optimizer = self._dist_context.serial_optimizer
+        if self.is_train and serial_optimizer:
+            # Generate backward
+            serial_loss = self._dist_context.serial_loss
+            params_grads = self._generate_backward(
+                serial_main_program,
+                serial_startup_program,
+                serial_loss,
+                parameter_list,
+            )
+            # Apply pre optimization passes
+            time0 = time.time()
+            (
+                serial_main_program,
+                serial_startup_program,
+                params_grads,
+            ) = self._apply_pre_optimization(
+                serial_main_program,
+                serial_startup_program,
+                serial_loss,
+                serial_optimizer,
+                params_grads,
+            )
+            self._logger.debug(
+                f"within parallel apply_pre_optimization time: {time.time() - time0}, mode {self._mode}"
+            )
+            # Do logical partition
+            time0 = time.time()
+            partitioner = Partitioner(self._dist_context, rank)
+            (
+                dist_main_prog,
+                dist_startup_prog,
+                dist_params_grads,
+            ) = partitioner.partition(
+                serial_main_program, serial_startup_program, params_grads
+            )
+
+            init_auto_parallel_rng()
+
+            self._logger.debug(
+                f"within parallel partitioner time: {time.time() - time0}, mode {self._mode}"
+            )
+            # Generate optimizer
+            time0 = time.time()
+            self._generate_optimizer(
+                dist_main_prog,
+                dist_startup_prog,
+                serial_optimizer,
+                dist_params_grads,
+            )
+            self._logger.debug(
+                f"within parallel optimizer time: {time.time() - time0}, mode {self._mode}"
+            )
+
+            resharder = Resharder(
+                dist_main_prog,
+                dist_startup_prog,
+                rank,
+                self._dist_context,
+                dist_params_grads,
+            )
+            resharder.reshard()
+            self._logger.debug(
+                f"within parallel reshard time: {time.time() - time0}, mode {self._mode}"
+            )
+            # Apply post optimization passes
+            time0 = time.time()
+            self._apply_post_optimization(
+                dist_main_prog, dist_startup_prog, rank, dist_params_grads
+            )
+            self._logger.debug(
+                f"within parallel apply_post_optimization time: {time.time() - time0}, mode {self._mode}"
+            )
+        else:
+            # Apply pre optimization passes
+            time0 = time.time()
+            (
+                serial_main_program,
+                serial_startup_program,
+                params_grads,
+            ) = self._apply_pre_optimization(
+                serial_main_program, serial_startup_program, None, None, []
+            )
+            self._logger.debug(
+                f"within parallel apply_pre_optimization time: {time.time() - time0}, mode {self._mode}"
+            )
+            # Do logical partition
+            time0 = time.time()
+            partitioner = Partitioner(self._dist_context, rank)
+            (
+                dist_main_prog,
+                dist_startup_prog,
+                dist_params_grads,
+            ) = partitioner.partition(
+                serial_main_program, serial_startup_program, []
+            )
+            # Do reshard process
+            self._logger.debug(
+                f"within parallel partitioner time: {time.time() - time0}, mode {self._mode}"
+            )
+            time0 = time.time()
+            # Do reshard process
+            micro_bsz = (
+                1
+                if not self._strategy.pipeline.enable
+                else self._strategy.pipeline.micro_batch_size
+            )
+            resharder = Resharder(
+                dist_main_prog,
+                dist_startup_prog,
+                rank,
+                self._dist_context,
+                [],
+                micro_bsz,
+            )
+            resharder.reshard()
+            self._logger.debug(
+                f"within parallel reshard time: {time.time() - time0}, mode {self._mode}"
+            )
+            # Apply post optimization passes
+            time0 = time.time()
+            self._apply_post_optimization(
+                dist_main_prog, dist_startup_prog, rank, dist_params_grads
+            )
+            self._logger.debug(
+                f"within parallel apply_post_optimization time: {time.time() - time0}, mode {self._mode}"
+            )
+
+        # Clone program for test
+        if self.is_test:
+            pipeline_opt = dist_main_prog._pipeline_opt
+            dist_main_prog = dist_main_prog.clone(for_test=True)
+            dist_startup_prog = dist_startup_prog.clone(for_test=True)
+            dist_main_prog._pipeline_opt = pipeline_opt
+
+        # Store the distributed programs for further usages
+        self._dist_context.dist_main_programs[rank] = dist_main_prog
+        self._dist_context.dist_startup_programs[rank] = dist_startup_prog
+
+    def _generate_backward(
+        self, main_program, startup_program, loss, parameter_list=None
+    ):
+        # NOTE(zhaoyinglia):
+        # Guarantee the order of params_grads is same between dynamic mode and static mode
+        # by making parameter_list equal to model.parameters(),
+        # because the order affect the result of ClipGradByGLobalNorm.
+        # If parameter_list is not None, the order of params_grads is same with parameter_list.
+        # If parameter_list is None, params_grads will be as prog.global_block().all_parameters().
+        with program_guard(main_program, startup_program):
+            params_grads = append_backward(
+                loss,
+                parameter_list=parameter_list,
+                distop_context=self._dist_context.dist_op_context,
+            )
+        self._completer.complete_backward_annotation(main_program)
+        self._dist_context.block_state.parse_backward_blocks(main_program)
+        return params_grads
+
+    def _generate_optimizer(
+        self, main_program, startup_program, optimizer, params_grads
+    ):
+        # NOTE:
+        # 1. `apply_gradients` will add an Accumulator for a parameter only once,
+        #    but optimizer will be called repeatedly in re-launch, so optimizer need to be copied.
+        # 2. lr_scheduler cannot be deepcopy, cause 'deepcopy' will lead to difference of learning_rate between executor and engine.
+        learning_rate = optimizer._learning_rate
+        new_optimizer = copy.deepcopy(optimizer)
+        new_optimizer._learning_rate = learning_rate
+        new_optimizer._sorted = False
+        self._dist_context._serial_optimizer = optimizer
+        self._dist_context._serial_optimizer._learning_rate = learning_rate
+
+        with (
+            program_guard(main_program, startup_program),
+            main_program.switch_name_generator_guard("opt_"),
+        ):
+            optimizer_ops = new_optimizer.apply_gradients(params_grads)
+        self._completer.complete_update_annotation(main_program)
+        return optimizer_ops
+
+    def _apply_pre_optimization(
+        self, main_program, startup_program, loss, optimizer, params_grads
+    ):
+        if self._strategy is None:
+            return
+
+        # apply amp pass on train/eval/predict
+        if self._strategy.amp.enable:
+            config = copy.deepcopy(self._strategy.amp.to_dict())
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            config["loss"] = loss
+            config["input_data"] = (
+                self._dist_context.serial_feed_vars["inputs"]
+                + self._dist_context.serial_feed_vars["labels"]
+            )
+            self._logger.info(
+                "Applying AMP-{}-{} ...".format(
+                    config["dtype"], config['level']
+                ),
+            )
+            if config['level'] == "o1":
+                auto_parallel_amp_pass = new_pass("auto_parallel_amp", config)
+                auto_parallel_amp_pass.apply(
+                    [main_program], [startup_program], self._pass_context
+                )
+                loss = auto_parallel_amp_pass.get_loss()
+            elif config['level'] in ['o2', 'o3']:
+                config["base_opt"] = optimizer
+                auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
+                auto_parallel_fp16_pass.apply(
+                    [main_program], [startup_program], self._pass_context
+                )
+                loss = auto_parallel_fp16_pass.get_loss()
+            else:
+                raise ValueError("AMP level should be one of o1, o2, o3")
+
+        # apply quantization pass
+        # The pass can be applied when mode must be 'train'
+        if self.is_train and self._strategy.qat.enable:
+            config = copy.deepcopy(self._strategy.qat.to_dict())
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            config["mode"] = self._mode
+            config["loss"] = loss
+            auto_parallel_quantization_pass = new_pass(
+                "auto_parallel_quantization", config
+            )
+            auto_parallel_quantization_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+            main_program = self._pass_context.get_attr("main_program")
+            startup_program = self._pass_context.get_attr("startup_program")
+            params_grads = self._pass_context.get_attr("params_grads")
+            loss = self._pass_context.get_attr("loss")
+
+        # apply recompute pass
+        # recompute is then train-only optimization
+        if self.is_train and self._strategy.recompute.enable:
+            config = copy.deepcopy(self._strategy.recompute.to_dict())
+            config["dist_context"] = self._dist_context
+            config["no_grad_set"] = None
+            config["loss"] = loss
+            auto_parallel_recompute_pass = new_pass(
+                "auto_parallel_recompute", config
+            )
+            auto_parallel_recompute_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+
+        return main_program, startup_program, params_grads
+
+    def _check_dist_attr(self, program, num_model_chunks, dist_context):
+        for _, block in enumerate(program.blocks):
+            for _, op in enumerate(block.ops):
+                op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
+                if op_dist_attr is None:
+                    raise ValueError(
+                        f"There is not dist_attr for op[{op.type}]."
+                    )
+
+    def _apply_post_optimization(
+        self, main_program, startup_program, rank, params_grads
+    ):
+        if self._strategy is None:
+            return
+
+        # sequence parallel optimization
+        if self._strategy.sp_optimization.enable:
+            config = copy.deepcopy(self._strategy.sp_optimization.to_dict())
+            config["dist_context"] = self._dist_context
+            config["global_rank"] = rank
+            sp_pass = new_pass(
+                "auto_parallel_sequence_parallel_optimization", config
+            )
+            sp_pass.apply([main_program], [startup_program], self._pass_context)
+
+        # apply fused linear promotion pass
+        if (
+            self.is_train
+            and self._strategy.fused_linear_promotion.enable
+            and self._strategy.fused_passes.enable
+        ):
+            if (
+                len(self._strategy.fused_passes.fused_passes_list) > 0
+                and "fuse_gemm_epilogue"
+                in self._strategy.fused_passes.fused_passes_list
+            ):
+                amp_config = None
+                if self._strategy.amp.enable:
+                    amp_config = copy.deepcopy(self._strategy.amp.to_dict())
+                config = {}
+                config["dist_context"] = self._dist_context
+                config["global_rank"] = rank
+                config["enable_sp"] = self._strategy.sp_optimization.enable
+                config["params_grads"] = params_grads
+                config["amp_level"] = (
+                    amp_config['level'] if amp_config is not None else "o0"
+                )
+                fused_linear_promotion_pass = new_pass(
+                    "auto_parallel_fused_linear_promotion", config
+                )
+                fused_linear_promotion_pass.apply(
+                    [main_program], [startup_program], self._pass_context
+                )
+
+        # apply master grad pass
+        if self._strategy.amp.enable:
+            amp_config = copy.deepcopy(self._strategy.amp.to_dict())
+            config = {}
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            config["completer"] = self._completer
+            if amp_config['level'] == "o2" and amp_config["use_master_grad"]:
+                master_grad_pass = new_pass(
+                    "auto_parallel_master_grad_pass", config
+                )
+                master_grad_pass.apply(
+                    [main_program], [startup_program], self._pass_context
+                )
+
+        # data parallel optimization
+        if self._strategy.dp_optimization.enable:
+            config = copy.deepcopy(self._strategy.dp_optimization.to_dict())
+            config["dist_context"] = self._dist_context
+            config["global_rank"] = rank
+            config["use_sharding"] = self._strategy.sharding.enable
+            dp_pass = new_pass(
+                "auto_parallel_data_parallel_optimization", config
+            )
+            dp_pass.apply([main_program], [startup_program], self._pass_context)
+
+        gradient_sync_after_accumulate = (
+            self._strategy.dp_optimization.gradient_sync_after_accumulate
+        )
+        if gradient_sync_after_accumulate:
+            global_params_grads = params_grads
+
+        if self._strategy.sharding.enable:
+            config = copy.deepcopy(self._strategy.sharding.to_dict())
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            config["global_rank"] = rank
+            config["gradient_sync_after_accumulate"] = (
+                gradient_sync_after_accumulate
+            )
+            if self._strategy.amp.enable:
+                amp_config = copy.deepcopy(self._strategy.amp.to_dict())
+                config["amp_dtype"] = amp_config['dtype']
+            auto_parallel_sharding_pass = new_pass(
+                "auto_parallel_sharding", config
+            )
+            auto_parallel_sharding_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+            params_grads = self._pass_context.get_attr("params_grads")
+
+        if self._strategy.mp_optimization.allreduce_matmul_grad_overlapping:
+            if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
+                self._logger.warning(
+                    "You set mp_optimization.allreduce_matmul_grad_overlapping=True, but you did not set environment "
+                    "variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
+                    "loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
+                )
+
+            config = {
+                "dist_context": self._dist_context,
+            }
+            allreduce_matmul_grad_overlapping_pass = new_pass(
+                "allreduce_matmul_grad_overlapping", config
+            )
+            allreduce_matmul_grad_overlapping_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+
+        if self.is_train:
+            # GradClip is train-only optimization
+            config = copy.deepcopy(self._strategy.sharding.to_dict())
+            config["dist_context"] = self._dist_context
+            config["params_grads"] = params_grads
+            config["rank_id"] = rank
+            auto_parallel_clip_pass = new_pass(
+                "auto_parallel_grad_clip", config
+            )
+            auto_parallel_clip_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+
+        if not is_sequential_run():
+            # deps for newexe
+            config = {}
+            config["dist_context"] = self._dist_context
+            APSED_pass = new_pass(
+                "auto_parallel_supplement_explicit_dependencies", config
+            )
+            APSED_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+
+        if self.is_train and self._strategy.pipeline.enable:
+            self._strategy.gradient_merge.enable = True
+            self._strategy.gradient_merge.k_steps = (
+                self._strategy.pipeline.accumulate_steps
+            )
+            self._strategy.gradient_merge.avg = True
+
+        # gradient_merge is then train-only optimization
+        grad_to_global_grad = {}
+        if self.is_train and self._strategy.gradient_merge.enable:
+            config = copy.deepcopy(self._strategy.gradient_merge.to_dict())
+            config["dist_context"] = self._dist_context
+            config["grad_to_global_grad"] = grad_to_global_grad
+            config["pipeline_mode"] = self._strategy.pipeline.schedule_mode
+            if gradient_sync_after_accumulate:
+                config["params_grads"] = global_params_grads
+                config["gradient_sync_after_accumulate"] = (
+                    gradient_sync_after_accumulate
+                )
+            else:
+                config["params_grads"] = params_grads
+            auto_parallel_gradient_merge_pass = new_pass(
+                "auto_parallel_gradient_merge_pass", config
+            )
+            auto_parallel_gradient_merge_pass.apply(
+                [main_program], [startup_program], self._pass_context
+            )
+
+        self._check_dist_attr(
+            main_program,
+            self._strategy.pipeline.vpp_degree,
+            self._dist_context,
+        )
+
+        enable_ir = get_flags("FLAGS_enable_pir_in_executor")[
+            'FLAGS_enable_pir_in_executor'
+        ]
+        ir_pass_list = []
+        if self.is_train and self._strategy.fused_passes.enable:
+            if len(self._strategy.fused_passes.fused_passes_list) > 0:
+                program_pass_list = []
+                for p in self._strategy.fused_passes.fused_passes_list:
+                    if enable_ir and p in (PIR_PASS + PIR_PYTHON_PASS):
+                        ir_pass_list.append(p)
+                    else:
+                        program_pass_list.append(new_pass(p))
+                pass_manager = PassManager(program_pass_list)
+                pass_manager.apply([main_program], [startup_program])
+
+        main_program._pass_opt = {}
+        main_program._pass_opt['pass_list'] = ir_pass_list
+
+        if self.is_train and self._strategy.pipeline.enable:
+            enable_send_recv_overlap = (
+                self._strategy.pipeline.enable_send_recv_overlap
+            )
+            if (
+                enable_send_recv_overlap
+                and int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1
+            ):
+                self._logger.warning(
+                    "You set pipeline.enable_send_recv_overlap=True, but you did not set environment "
+                    "variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
+                    "loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
+                )
+
+            main_program._pipeline_opt = {}
+            main_program._pipeline_opt["standalone_opt"] = {
+                "enable_send_recv_overlap": enable_send_recv_overlap,
+                "schedule_mode": self._strategy.pipeline.schedule_mode,
+                "num_micro_batches": self._strategy.pipeline.accumulate_steps,
+                "pp_degree": len(self._dist_context.process_meshes),
+                "pp_stage": get_pp_stage(self._dist_context, rank),
+                "vpp_degree": self._strategy.pipeline.vpp_degree,
+                "dist_context": self._dist_context,
+                "program_runtimes": self._strategy.pipeline.program_runtimes,
+                "memory_limit_times": self._strategy.pipeline.memory_limit_times,
+                "split_backward": self._strategy.pipeline.split_backward,
+                "grad_to_global_grad": grad_to_global_grad,
+            }

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

@@ -0,0 +1,543 @@
+# 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
+from collections import defaultdict
+
+import paddle
+from paddle.distributed.auto_parallel.static.dist_context import (
+    DistributedContext,
+)
+from paddle.distributed.auto_parallel.static.operators.common import (
+    get_distributed_operator_impl_container,
+)
+from paddle.framework import Program, core
+from paddle.static import Parameter
+
+from .dist_attribute import OperatorDistAttr
+from .operators.common import BACKWARD_ONLY_DIST_OPS
+from .utils import (
+    __no_shape_var_type__,
+    is_backward_op,
+    is_forward_op,
+    is_loss_op,
+    is_optimize_op,
+)
+
+__varname_not_in_block__ = ["lod_tensor_blocking_queue"]
+
+
+class Partitioner:
+    """
+    warning:: Partitioner is experimental and subject to change.
+
+    Partitioner convert a program into another program.
+    Given a serial program which has been auto completed with shard annotation, the Partitioner
+    convert the serial program into a "distributed" program. The Partitioner will  modify the serial
+    program in following two ways, which is also the major difference between serial and distributed program:
+        1. partition op: replace a serial op into its corresponding dist op inferred from the shard annotation
+        2. partition var: if a var is sharded, modify the shape of var according to its shard annotation
+
+    Partitioner is supposed to be call by the auto parallel framework, and not supposed to be directly called by user.
+    """
+
+    def __init__(self, dist_context, rank_id=0):
+        """
+        Args:
+            dist_context (DistributedContext): used to access the distributed_attr of var & op, every Partitioner object could maintain its own DistributedContext member, and partition program base on that shard scenario.
+            rank_id (int): global rank id to which the partitioned distributed program belong.
+        """
+        if not isinstance(dist_context, DistributedContext):
+            raise TypeError(
+                f"dist_context be DistributedContext, got {type(dist_context)} here"
+            )
+
+        self._dist_context = dist_context
+        self._rank_id = rank_id
+        self._serial2dist_varname_mapping = defaultdict(
+            dict
+        )  # block_id -> serial_varname -> dist_varname
+        self._dist_varname_suffix = ""
+        self._forward_op_id2forward_op = {}
+
+    def partition(
+        self, serial_main_program, serial_startup_program, params_grads
+    ):
+        if not isinstance(serial_main_program, (Program)):
+            raise TypeError(
+                f"main_program be paddle.framework.Program, got {type(serial_main_program)} here"
+            )
+
+        # check if shard annotated serial program valid
+        if not self._is_valid_annotated_program(serial_main_program):
+            raise RuntimeError(
+                "Not all vars or ops are annotated in main program !"
+            )
+
+        # init distop helper
+        dist_op_context = self._dist_context.dist_op_context
+        dist_op_context.varname_mapping = self._serial2dist_varname_mapping
+        dist_op_context.rank_id = self._rank_id
+
+        # partition startup program
+        if serial_startup_program is None:
+            partitioned_startup_prog = None
+        else:
+            partitioned_startup_prog = self.partition_startup_program(
+                serial_main_program, serial_startup_program
+            )
+        dist_op_context.dst_startup_program = partitioned_startup_prog
+
+        # partition main program
+        (
+            partitioned_main_prog,
+            partitioned_params_grads,
+        ) = self.partition_main_program(serial_main_program, params_grads)
+
+        return (
+            partitioned_main_prog,
+            partitioned_startup_prog,
+            partitioned_params_grads,
+        )
+
+    def partition_startup_program(
+        self, serial_main_program, serial_startup_program
+    ):
+        if not isinstance(serial_startup_program, (Program)):
+            raise TypeError(
+                f"dist_context be paddle.framework.Program, got {type(serial_startup_program)} here"
+            )
+
+        partitioned_startup_prog = paddle.framework.Program()
+        partitioned_startup_prog._name_generator = (
+            serial_startup_program._name_generator.clone()
+        )
+        ref_block = serial_main_program.global_block()
+        target_block = partitioned_startup_prog.global_block()
+        var2shape = {}
+        temp_varname_map = {}
+
+        # tensors
+        for var in serial_startup_program.list_vars():
+            assert var.persistable
+            new_name = var.name + self._dist_varname_suffix
+            temp_varname_map[var.name] = new_name
+            target_shape = _partition_var(
+                self._dist_context, ref_block, target_block, var.name, new_name
+            )
+            var2shape[new_name] = target_shape
+
+        # ops
+        for op in serial_startup_program.global_block().ops:
+            # TODO if var not belong to this rank, should be filtered
+            output_vars = op.desc.output_arg_names()
+            assert len(output_vars) == 1, (
+                f"initializer should output only ONE variable, but got [{op.desc}]"
+            )
+            assert temp_varname_map[output_vars[0]] in var2shape, (
+                f"try to initialize [{output_vars[0]}] which is not a persistable var"
+            )
+            new_op_desc = target_block.desc.append_op()
+            new_op_desc.copy_from(op.desc)
+            new_op_desc._rename_output(
+                output_vars[0], temp_varname_map[output_vars[0]]
+            )
+            new_op_desc._set_attr(
+                "shape", var2shape[temp_varname_map[output_vars[0]]]
+            )
+            target_block._sync_with_cpp()
+
+            # set distribute attribute
+            new_op = target_block.ops[-1]
+            assert new_op.type == new_op_desc.type()
+            assert new_op.desc == new_op_desc
+            output_var = target_block.var(output_vars[0])
+            output_var_attr = (
+                self._dist_context.get_tensor_dist_attr_for_program(output_var)
+            )
+            op_attr = OperatorDistAttr()
+            op_attr.process_mesh = output_var_attr.process_mesh
+            op_attr.set_output_dims_mapping(
+                output_var.name, output_var_attr.dims_mapping
+            )
+            op_attr.set_input_dims_mapping(
+                output_var.name, output_var_attr.dims_mapping
+            )
+            self._dist_context.set_op_dist_attr_for_program(new_op, op_attr)
+
+        return partitioned_startup_prog
+
+    def partition_main_program(self, serial_main_program, params_and_grads):
+        """
+        1. partition variables
+        2. replace local op with corresponding dist op
+        """
+
+        partitioned_main_prog = paddle.framework.Program()
+        partitioned_main_prog._name_generator = (
+            serial_main_program._name_generator.clone()
+        )
+        dist_op_context = self._dist_context.dist_op_context
+        dist_op_context.dst_main_program = partitioned_main_prog
+
+        for idx in range(self._dist_context.block_state.nblock):
+            ref_block = serial_main_program.blocks[idx]
+
+            if idx == 0:
+                target_block = partitioned_main_prog.blocks[0]
+            else:
+                target_block = partitioned_main_prog._create_block(
+                    parent_idx=ref_block.parent_idx
+                )
+                assert ref_block.idx == target_block.idx
+                target_block._set_forward_block_idx(ref_block.forward_block_idx)
+            dist_op_context.work_block = target_block
+            self.partition_block(ref_block, target_block)
+
+        partitioned_main_prog.current_block_idx = 0
+
+        # should reconnect the block_attr ptr to the correct block
+        for block_id in range(self._dist_context.block_state.nblock):
+            block = partitioned_main_prog.block(block_id)
+            for op in block.ops:
+                for attr_name in op.all_attrs():
+                    if op.attr_type(attr_name) == core.AttrType.BLOCK:
+                        relative_id = op._block_attr_id(attr_name)
+                        op._set_attr(
+                            attr_name, partitioned_main_prog.block(relative_id)
+                        )
+
+        partitioned_params_and_grads = []
+        for p, g in params_and_grads:
+            assert p.name in self._serial2dist_varname_mapping[0]
+            dist_p = self._get_dist_var_by_serial_var(
+                p, partitioned_main_prog, 0
+            )
+            if g is None:
+                dist_g = None
+            else:
+                assert g.name in self._serial2dist_varname_mapping[0]
+                dist_g = self._get_dist_var_by_serial_var(
+                    g, partitioned_main_prog, 0
+                )
+            partitioned_params_and_grads.append((dist_p, dist_g))
+
+        return partitioned_main_prog, partitioned_params_and_grads
+
+    def partition_block(self, ref_block, target_block):
+        dist_op_context = self._dist_context.dist_op_context
+
+        last_fwd_op_idx = -1
+        for idx, op in enumerate(ref_block.ops):
+            if is_loss_op(op):
+                last_fwd_op_idx = idx
+                break
+
+        if last_fwd_op_idx == -1:
+            last_fwd_op_idx = len(ref_block.ops)
+
+        for idx in range(len(ref_block.ops)):
+            if idx <= last_fwd_op_idx:
+                self._forward_op_id2forward_op[
+                    ref_block.ops[idx].desc.original_id()
+                ] = ref_block.ops[idx]
+
+        # partition
+        appended_grad_times = 0
+        for idx, op in enumerate(ref_block.ops):
+            op_dist_attr = self._dist_context.get_op_dist_attr_for_program(op)
+            if is_backward_op(op) and (
+                is_forward_op(ref_block.ops[idx - 1])
+                or is_loss_op(ref_block.ops[idx - 1])
+            ):
+                if not op_dist_attr.is_recompute:
+                    appended_grad_times += 1
+
+            # partition input variables
+            for serial_input_varname in op.desc.input_arg_names():
+                if (
+                    serial_input_varname
+                    not in self._serial2dist_varname_mapping[
+                        ref_block.forward_block_idx
+                    ]
+                    or serial_input_varname
+                    not in self._serial2dist_varname_mapping[ref_block.idx]
+                ):
+                    new_varname = (
+                        serial_input_varname + self._dist_varname_suffix
+                    )
+                    if ref_block.has_var(serial_input_varname):
+                        _partition_var(
+                            self._dist_context,
+                            ref_block,
+                            target_block,
+                            serial_input_varname,
+                            new_varname,
+                        )
+                    self._serial2dist_varname_mapping[ref_block.idx][
+                        serial_input_varname
+                    ] = new_varname
+
+            # partition output vars
+            for serial_output_varname in op.desc.output_arg_names():
+                if (
+                    serial_output_varname
+                    not in self._serial2dist_varname_mapping[
+                        ref_block.forward_block_idx
+                    ]
+                    or serial_output_varname
+                    not in self._serial2dist_varname_mapping[ref_block.idx]
+                ):
+                    new_varname = (
+                        serial_output_varname + self._dist_varname_suffix
+                    )
+                    if ref_block.has_var(serial_output_varname):
+                        _partition_var(
+                            self._dist_context,
+                            ref_block,
+                            target_block,
+                            serial_output_varname,
+                            new_varname,
+                        )
+                    self._serial2dist_varname_mapping[ref_block.idx][
+                        serial_output_varname
+                    ] = new_varname
+
+            # partition op
+            if is_forward_op(op) or op_dist_attr.is_recompute:
+                kinputs, koutputs = dist_op_context.prepare_context(op)
+                dist_op_forward_impl = _get_dist_op_forward_implement(
+                    op, self._dist_context
+                )
+                dist_op_forward_impl.forward(
+                    self._dist_context, **kinputs, **koutputs
+                )
+
+            elif is_backward_op(op):
+                kinputs, koutputs = dist_op_context.prepare_context(op)
+                dist_op_backward_impl = _get_dist_op_backward_implement(
+                    op, self._dist_context, self._forward_op_id2forward_op
+                )
+                grad_var_to_var = (
+                    self._dist_context.dist_op_context.grad_var_to_var[
+                        appended_grad_times
+                    ]
+                )
+                dist_op_backward_impl.backward(
+                    self._dist_context,
+                    **kinputs,
+                    **koutputs,
+                    **{"grad_var_to_var": grad_var_to_var},
+                )
+            elif is_optimize_op(op):
+                # NOTE: BACKWARD_ONLY_DIST_OPS's op_role must be 2 because of 1F1B PASS
+                kinputs, koutputs = dist_op_context.prepare_context(op)
+                dist_op_opt_impl = _get_dist_op_backward_implement(
+                    op, self._dist_context, self._forward_op_id2forward_op
+                )
+                dist_op_opt_impl.backward(
+                    self._dist_context,
+                    **kinputs,
+                    **koutputs,
+                    **{"grad_var_to_var": {}},
+                )
+            else:
+                raise NotImplementedError(
+                    f"partitioner only support forward and backward, optimize ops, but got {op}"
+                )
+
+    def _is_valid_annotated_program(self, program):
+        # TODO (ZJ-LIANG) should check all block
+        ops = program.global_block().ops
+        vars_ = program.list_vars()
+        op_dist_attrs = [
+            self._dist_context.get_op_dist_attr_for_program(op) for op in ops
+        ]
+        var_dist_attrs = [
+            self._dist_context.get_tensor_dist_attr_for_program(var)
+            for var in vars_
+            if (var.type not in __no_shape_var_type__)
+        ]
+
+        all_ops_annotated = all(
+            dist_attr is not None for dist_attr in op_dist_attrs
+        )
+        all_vars_annotated = all(
+            dist_attr is not None for dist_attr in var_dist_attrs
+        )
+
+        return all_ops_annotated and all_vars_annotated
+
+    def _get_dist_var_by_serial_var(
+        self, serial_var, partitioned_main_prog, block_id
+    ):
+        block_idx = serial_var.block.idx
+        target_block = partitioned_main_prog.blocks[block_idx]
+        dist_var_name = self._serial2dist_varname_mapping[block_id][
+            serial_var.name
+        ]
+        assert target_block.has_var(dist_var_name)
+        return target_block.var(dist_var_name)
+
+
+def _get_dist_shape(var, dist_attr):
+    var_shape = var.shape
+    mapping = dist_attr.dims_mapping
+    mesh = dist_attr.process_mesh.shape
+    if mapping == []:
+        return var_shape
+
+    assert len(var_shape) == len(mapping), (
+        f"variable shape [{var_shape}] and dim_mapping [{mapping}] is NOT match !"
+    )
+    new_shape = []
+    for idx in range(len(var_shape)):
+        if var_shape[idx] == -1 or mapping[idx] == -1:
+            new_shape.append(var_shape[idx])
+        else:
+            assert var_shape[idx] % mesh[mapping[idx]] == 0, (
+                f"un-event partition: var_shape[idx]=[{var_shape[idx]}], mesh[{mesh[mapping[idx]]}], {var.name}, {var_shape}, {mesh}, {mapping}"
+            )
+            new_shape.append(var_shape[idx] // mesh[mapping[idx]])
+
+    return new_shape
+
+
+def _partition_parameter(
+    dist_context, src_var, dst_block, dst_varname, dst_shape
+):
+    # NOTE hack to copied Parameter
+    # not initialized parameter, need to initialize it
+    copied_kwargs = {}
+    copied_kwargs['trainable'] = src_var.trainable
+    copied_kwargs['optimize_attr'] = src_var.optimize_attr
+    copied_kwargs['regularizer'] = src_var.regularizer
+    copied_kwargs['do_model_average'] = src_var.do_model_average
+    copied_kwargs['need_clip'] = src_var.need_clip
+
+    param = Parameter(
+        block=dst_block,
+        type=src_var.type,
+        name=dst_varname,
+        shape=dst_shape,
+        dtype=src_var.dtype,
+        lod_level=src_var.lod_level,
+        error_clip=src_var.error_clip,
+        stop_gradient=src_var.stop_gradient,
+        is_data=src_var.is_data,
+        belong_to_optimizer=src_var.belong_to_optimizer,
+        **copied_kwargs,
+    )
+
+    return param
+
+
+def _partition_intermediate_var(
+    dist_context, src_var, dst_block, dst_varname, dst_shape
+):
+    var = dst_block.create_var(
+        type=src_var.type,
+        name=dst_varname,
+        shape=dst_shape,
+        dtype=src_var.dtype,
+        lod_level=src_var.lod_level,
+        persistable=src_var.persistable,
+        error_clip=src_var.error_clip,
+        stop_gradient=src_var.stop_gradient,
+        is_data=src_var.is_data,
+        belong_to_optimizer=src_var.belong_to_optimizer,
+    )
+
+    return var
+
+
+def _partition_var(
+    dist_context, src_block, dst_block, src_varname, dst_varname
+):
+    """
+    partition include: split + replicate
+    """
+    src_var = src_block.var(src_varname)
+
+    if src_var.type in __no_shape_var_type__:
+        persist = getattr(src_var, 'persistable', False)
+        new_var = dst_block.create_var(
+            type=src_var.type,
+            name=dst_varname,
+            persistable=persist,
+            stop_gradient=True,
+        )
+        target_shape = None
+    else:
+        dist_attr = dist_context.get_tensor_dist_attr_for_program(src_var)
+        target_shape = _get_dist_shape(src_var, dist_attr)
+
+        if isinstance(src_var, Parameter):
+            new_var = _partition_parameter(
+                dist_context, src_var, dst_block, dst_varname, target_shape
+            )
+        else:
+            new_var = _partition_intermediate_var(
+                dist_context, src_var, dst_block, dst_varname, target_shape
+            )
+
+    dist_attr = copy.deepcopy(
+        dist_context.get_tensor_dist_attr_for_program(src_var)
+    )
+    assert dist_attr is not None
+    dist_context.set_tensor_dist_attr_for_program(new_var, dist_attr)
+
+    return target_shape
+
+
+def _get_dist_op_backward_implement(
+    backward_op, dist_context, forward_op_id2forward_op
+):
+    dist_op_context = dist_context.dist_op_context
+    if backward_op.desc.original_id() in dist_op_context.grad_op_id_to_op_id:
+        forward_op_id = dist_op_context.grad_op_id_to_op_id[
+            backward_op.desc.original_id()
+        ]
+        forward_op = forward_op_id2forward_op[forward_op_id]
+        forward_op_dist_attr = dist_context.get_op_dist_attr_for_program(
+            forward_op
+        )
+        dist_op_impl_container = get_distributed_operator_impl_container(
+            forward_op_dist_attr.impl_type
+        )
+        dist_op_impl = dist_op_impl_container.get_impl(
+            forward_op_dist_attr.impl_idx
+        )
+        return dist_op_impl
+
+    # # NOTE trick for dist ops that only have backward implement
+    if backward_op.type in BACKWARD_ONLY_DIST_OPS:
+        op_dist_attr = dist_context.get_op_dist_attr_for_program(backward_op)
+        assert op_dist_attr.impl_idx >= 0
+        dist_op_impl = get_distributed_operator_impl_container(
+            op_dist_attr.impl_type
+        ).get_impl(op_dist_attr.impl_idx)
+        return dist_op_impl
+
+    dist_op = get_distributed_operator_impl_container("default")
+    return dist_op.get_impl(0)
+
+
+def _get_dist_op_forward_implement(forward_op, dist_context):
+    dist_attr = dist_context.get_op_dist_attr_for_program(forward_op)
+    dist_op_impl_container = get_distributed_operator_impl_container(
+        dist_attr.impl_type
+    )
+    dist_op_impl = dist_op_impl_container.get_impl(dist_attr.impl_idx)
+    return dist_op_impl