chore: import upstream snapshot with attribution

python/paddle/distributed/fleet/.gitignore

@@ -0,0 +1 @@
+proto

python/paddle/distributed/fleet/__init__.py

@@ -0,0 +1,113 @@
+#   Copyright (c) 2020 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.
+
+# TODO: define distributed api under this directory,
+from . import metrics  # noqa: F401
+from .base.distributed_strategy import DistributedStrategy
+from .base.role_maker import (
+    PaddleCloudRoleMaker,
+    Role,
+    UserDefinedRoleMaker,
+)
+from .base.topology import CommunicateTopology, HybridCommunicateGroup
+from .base.util_factory import UtilBase
+from .data_generator.data_generator import (
+    MultiSlotDataGenerator,
+    MultiSlotStringDataGenerator,
+)
+from .dataset import (  # noqa: F401
+    DatasetBase,
+    FileInstantDataset,
+    InMemoryDataset,
+    QueueDataset,
+)
+from .fleet import Fleet
+from .model import distributed_model
+from .optimizer import distributed_optimizer
+from .scaler import distributed_scaler
+from .utils import log_util
+
+__all__ = [
+    "CommunicateTopology",
+    "UtilBase",
+    "HybridCommunicateGroup",
+    "MultiSlotStringDataGenerator",
+    "UserDefinedRoleMaker",
+    "DistributedStrategy",
+    "Role",
+    "MultiSlotDataGenerator",
+    "PaddleCloudRoleMaker",
+    "Fleet",
+]
+
+fleet = Fleet()
+_final_strategy = fleet._final_strategy
+_get_applied_meta_list = fleet._get_applied_meta_list
+_get_applied_graph_list = fleet._get_applied_graph_list
+init = fleet.init
+is_first_worker = fleet.is_first_worker
+worker_index = fleet.worker_index
+worker_num = fleet.worker_num
+node_num = fleet.node_num
+rank = fleet.worker_index
+nranks = fleet.worker_num
+world_size = fleet.worker_num
+# device id in current trainer
+local_device_ids = fleet.local_device_ids
+# device ids in world
+world_device_ids = fleet.world_device_ids
+# rank in node
+local_rank = fleet.local_rank
+rank_in_node = local_rank
+is_worker = fleet.is_worker
+is_coordinator = fleet.is_coordinator
+init_coordinator = fleet.init_coordinator
+make_fl_strategy = fleet.make_fl_strategy
+get_fl_client = fleet.get_fl_client
+worker_endpoints = fleet.worker_endpoints
+server_num = fleet.server_num
+server_index = fleet.server_index
+server_endpoints = fleet.server_endpoints
+is_server = fleet.is_server
+util = UtilBase()
+barrier_worker = fleet.barrier_worker
+all_reduce = fleet.all_reduce
+init_worker = fleet.init_worker
+init_server = fleet.init_server
+run_server = fleet.run_server
+stop_worker = fleet.stop_worker
+distributed_optimizer = distributed_optimizer
+save_inference_model = fleet.save_inference_model
+save_persistables = fleet.save_persistables
+save_cache_model = fleet.save_cache_model
+check_save_pre_patch_done = fleet.check_save_pre_patch_done
+save_one_table = fleet.save_one_table
+save_dense_params = fleet.save_dense_params
+load_model = fleet.load_model
+load_inference_model = fleet.load_inference_model
+load_one_table = fleet.load_one_table
+set_date = fleet.set_date
+print_table_stat = fleet.print_table_stat
+minimize = fleet.minimize
+distributed_model = distributed_model
+shrink = fleet.shrink
+get_hybrid_communicate_group = fleet.get_hybrid_communicate_group
+distributed_scaler = distributed_scaler
+set_log_level = log_util.set_log_level
+get_log_level_code = log_util.get_log_level_code
+get_log_level_name = log_util.get_log_level_name
+check_memory_usage = log_util.check_memory_usage
+save_cache_table = fleet.save_cache_table
+collective_perf = fleet.collective_perf
+from .. import auto_parallel as auto  # noqa: F401

python/paddle/distributed/fleet/base/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2020 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/fleet/base/graphviz.py

@@ -0,0 +1,296 @@
+#   Copyright (c) 2018 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 logging
+import os
+import random
+import subprocess
+
+
+def crepr(v):
+    if isinstance(v, str):
+        return f'"{v}"'
+    return str(v)
+
+
+class Rank:
+    def __init__(self, kind, name, priority):
+        '''
+        kind: str
+        name: str
+        priority: int
+        '''
+        self.kind = kind
+        self.name = name
+        self.priority = priority
+        self.nodes = []
+
+    def __str__(self):
+        if not self.nodes:
+            return ''
+
+        return (
+            '{'
+            + f'rank={self.kind};'
+            + ','.join([node.name for node in self.nodes])
+            + '}'
+        )
+
+
+class Graph:
+    rank_counter = 0
+
+    def __init__(self, title, **attrs):
+        self.title = title
+        self.attrs = attrs
+        self.nodes = []
+        self.edges = []
+        self.rank_groups = {}
+
+    def code(self):
+        return self.__str__()
+
+    def rank_group(self, kind, priority):
+        name = f"rankgroup-{Graph.rank_counter}"
+        Graph.rank_counter += 1
+        rank = Rank(kind, name, priority)
+        self.rank_groups[name] = rank
+        return name
+
+    def node(self, label, prefix, description="", **attrs):
+        node = Node(label, prefix, description, **attrs)
+
+        if 'rank' in attrs:
+            rank = self.rank_groups[attrs['rank']]
+            del attrs['rank']
+            rank.nodes.append(node)
+        self.nodes.append(node)
+        return node
+
+    def edge(self, source, target, **attrs):
+        edge = Edge(source, target, **attrs)
+        self.edges.append(edge)
+        return edge
+
+    def compile(self, dot_path):
+        file = open(dot_path, 'w')
+        file.write(self.__str__())
+        image_path = os.path.join(
+            os.path.dirname(dot_path), dot_path[:-3] + "pdf"
+        )
+        cmd = ["dot", "-Tpdf", dot_path, "-o", image_path]
+        subprocess.Popen(
+            cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        logging.warning(f"write block debug graph to {image_path}")
+        return image_path
+
+    def show(self, dot_path):
+        image = self.compile(dot_path)
+        cmd = ["open", image]
+        subprocess.Popen(
+            cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+    def _rank_repr(self):
+        ranks = sorted(
+            self.rank_groups.items(),
+            key=functools.cmp_to_key(
+                lambda a, b: a[1].priority > b[1].priority
+            ),
+        )
+        repr = []
+        for x in ranks:
+            repr.append(str(x[1]))
+        return '\n'.join(repr) + '\n'
+
+    def __str__(self):
+        reprs = [
+            'digraph G {',
+            f'title = {crepr(self.title)}',
+        ]
+
+        for attr in self.attrs:
+            reprs.append(f"{attr}={crepr(self.attrs[attr])};")
+
+        reprs.append(self._rank_repr())
+
+        random.shuffle(self.nodes)
+        reprs += [str(node) for node in self.nodes]
+
+        for x in self.edges:
+            reprs.append(str(x))
+
+        reprs.append('}')
+        return '\n'.join(reprs)
+
+
+class Node:
+    counter = 1
+
+    def __init__(self, label, prefix, description="", **attrs):
+        self.label = label
+        self.name = f"{prefix}_{Node.counter}"
+        self.description = description
+        self.attrs = attrs
+        Node.counter += 1
+
+    def __str__(self):
+        reprs = '{name} [label={label} {extra} ];'.format(
+            name=self.name,
+            label=self.label,
+            extra=(
+                ','
+                + ','.join(
+                    f"{key}={crepr(value)}" for key, value in self.attrs.items()
+                )
+                if self.attrs
+                else ""
+            ),
+        )
+        return reprs
+
+
+class Edge:
+    def __init__(self, source, target, **attrs):
+        '''
+        Link source to target.
+        :param source: Node
+        :param target: Node
+        :param graph: Graph
+        :param attrs: dic
+        '''
+        self.source = source
+        self.target = target
+        self.attrs = attrs
+
+    def __str__(self):
+        repr = "{source} -> {target} {extra}".format(
+            source=self.source.name,
+            target=self.target.name,
+            extra=(
+                ""
+                if not self.attrs
+                else "["
+                + ','.join(
+                    f"{attr[0]}={crepr(attr[1])}" for attr in self.attrs.items()
+                )
+                + "]"
+            ),
+        )
+        return repr
+
+
+class GraphPreviewGenerator:
+    '''
+    Generate a graph image for ONNX proto.
+    '''
+
+    def __init__(self, title):
+        # init graphviz graph
+        self.graph = Graph(
+            title,
+            layout="dot",
+            concentrate="true",
+            rankdir="TB",
+        )
+
+        self.op_rank = self.graph.rank_group('same', 2)
+        self.param_rank = self.graph.rank_group('same', 1)
+        self.arg_rank = self.graph.rank_group('same', 0)
+
+    def __call__(self, path='temp.dot', show=False):
+        if not show:
+            self.graph.compile(path)
+        else:
+            self.graph.show(path)
+
+    def add_param(self, name, data_type, highlight=False):
+        label = '\n'.join(
+            [
+                '<<table cellpadding="5">',
+                '  <tr>',
+                '    <td bgcolor="#2b787e">',
+                '    <b>',
+                name,
+                '    </b>',
+                '    </td>',
+                '  </tr>',
+                '  <tr>',
+                '    <td>',
+                str(data_type),
+                '    </td>  </tr>',
+                '</table>>',
+            ]
+        )
+        return self.graph.node(
+            label,
+            prefix="param",
+            description=name,
+            shape="none",
+            style="rounded,filled,bold",
+            width="1.3",
+            color="#148b97" if not highlight else "orange",
+            fontcolor="#ffffff",
+            fontname="Arial",
+        )
+
+    def add_op(self, opType, **kwargs):
+        highlight = False
+        if 'highlight' in kwargs:
+            highlight = kwargs['highlight']
+            del kwargs['highlight']
+        return self.graph.node(
+            f"<<B>{opType}</B>>",
+            prefix="op",
+            description=opType,
+            shape="box",
+            style="rounded, filled, bold",
+            color="#303A3A" if not highlight else "orange",
+            fontname="Arial",
+            fontcolor="#ffffff",
+            width="1.3",
+            height="0.84",
+        )
+
+    def add_arg(self, name, highlight=False):
+        return self.graph.node(
+            crepr(name),
+            prefix="arg",
+            description=name,
+            shape="box",
+            style="rounded,filled,bold",
+            fontname="Arial",
+            fontcolor="#999999",
+            color="#dddddd" if not highlight else "orange",
+        )
+
+    def add_edge(self, source, target, **kwargs):
+        highlight = False
+        if 'highlight' in kwargs:
+            highlight = kwargs['highlight']
+            del kwargs['highlight']
+        return self.graph.edge(
+            source,
+            target,
+            color="#00000" if not highlight else "orange",
+            **kwargs,
+        )

python/paddle/distributed/fleet/base/meta_optimizer_factory.py

@@ -0,0 +1,41 @@
+#   Copyright (c) 2020 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 ..meta_optimizers import *  # noqa: F403
+
+__all__ = []
+
+meta_optimizer_names = list(
+    filter(lambda name: name.endswith("Optimizer"), dir())
+)
+
+# Because HybridParallelOptimizer is dygraph optimizer, it
+# should be removed
+meta_optimizer_names.remove("HybridParallelOptimizer")
+meta_optimizer_names.remove("HeterParallelOptimizer")
+meta_optimizer_names.remove("DGCMomentumOptimizer")
+meta_optimizer_names.remove("MuonShardingOptimizer")
+
+
+class MetaOptimizerFactory:
+    def __init__(self):
+        pass
+
+    def _get_valid_meta_optimizers(self, user_defined_optimizer, skip_names=[]):
+        opt_list = []
+        for opt_name in meta_optimizer_names:
+            if opt_name in skip_names:
+                continue
+            opt_list.append(globals()[opt_name](user_defined_optimizer))
+        return opt_list

python/paddle/distributed/fleet/base/orthogonal_strategy.py

@@ -0,0 +1,209 @@
+# 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 collections
+import functools
+import itertools
+
+import paddle.distributed as dist
+from paddle.distributed.fleet.base.strategy_group import StrategyGroupBase
+
+
+class OrthogonalStrategy:
+    """
+    A hybrid of multiple distributed strategies. Strategies need to be orthogonal, means the ranks are organized like
+    a square if there are two strategies, a cube if there are three strategies, etc.
+
+    Args:
+        list_of_strategy(list): Strategy in the list should be represented as tuple, format as (strategy_name, degree, strategy_class).
+        fused_strategy_dict(dict, optional): Exist strategies can be fused to new strategy. Use the name of new strategy as key, a list of
+            strategy names you want to fuse as value.
+
+    Returns:
+        The instance of strategy.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> from paddle.distributed.fleet.base.strategy_group import DPGroup, MPGroup, PPGroup
+            >>> from paddle.distributed.fleet.base.orthogonal_strategy import OrthogonalStrategy
+
+            >>> dist.init_parallel_env()
+            >>> strategy = OrthogonalStrategy(
+            ...     [
+            ...         ("dp", 2, DPGroup),
+            ...         ("mp", 2, MPGroup),
+            ...         ("pp", 2, PPGroup),
+            ...     ],
+            ...     fused_strategy_dict={"check": ["mp", "pp"]},
+            ... )
+
+    """
+
+    def __init__(
+        self, list_of_strategy, fused_strategy_dict={}, strategy_rank_list=None
+    ):
+        self._list_of_strategy = list_of_strategy
+        self._fused_strategy_dict = fused_strategy_dict
+        self._strategy_rank_list = (
+            strategy_rank_list
+            if strategy_rank_list is not None
+            else list(range(dist.get_world_size()))
+        )
+        self._name_to_group_dict = {}
+        self._name_to_degree_dict = {}
+        self._list_of_strategy_name = [
+            strategy[0] for strategy in list_of_strategy
+        ]
+        self._list_of_degree = [strategy[1] for strategy in list_of_strategy]
+        self._coordinate = collections.namedtuple(
+            'Coordinate', self._list_of_strategy_name
+        )
+        self._check_valid_strategy()
+
+        ranges = [range(degree) for degree in self._list_of_degree]
+        list_of_coord = [
+            self._coordinate(*coord) for coord in itertools.product(*ranges)
+        ]
+
+        self._coord_to_rank_dict = dict(
+            zip(list_of_coord, self._strategy_rank_list)
+        )
+
+        for idx, strategy in enumerate(list_of_strategy):
+            strategy_name = strategy[0]
+            self._name_to_degree_dict[strategy_name] = strategy[1]
+            rank_list = self._calc_rank_list(idx)
+            self._name_to_group_dict[strategy_name] = strategy[2](
+                rank_list,
+            )
+
+        self._name_to_fused_group_dict = {}
+        self._create_fused_group()
+
+    def strategy_group(self, name):
+        """
+        Get strategy group with specific name.
+
+        Args:
+            name: The name of strategy group
+
+        Returns:
+            An instance of specific strategy group.
+        """
+        assert name in self._list_of_strategy_name, (
+            f"Strategy group {name} is not created."
+        )
+        return self._name_to_group_dict[name]
+
+    def fused_strategy_group(self, name):
+        """
+        Get fused strategy group with specific name.
+
+        Args:
+            name: The name of fused strategy group
+
+        Returns:
+            (StrategyGroupBase): An instance of strategy group.
+        """
+        assert name in self._name_to_fused_group_dict, (
+            f"Fused strategy group {name} is not created."
+        )
+        return self._name_to_fused_group_dict[name]
+
+    def rank_in_strategy(self, name):
+        """
+        Get local rank in strategy group with specific name.
+
+        Args:
+            name: The name of strategy group
+
+        Returns:
+            (Integer): Local rank in specific strategy.
+        """
+        assert name in self._list_of_strategy_name, (
+            f"Strategy group {name} is not created."
+        )
+        return self._name_to_group_dict[name].group.rank
+
+    def _check_valid_strategy(self):
+        assert len(self._list_of_strategy_name) == len(
+            set(self._list_of_strategy_name)
+        ), f"Defined duplicated strategies: {self._list_of_strategy_name}"
+        num_of_ranks = functools.reduce(
+            lambda x, y: x * y, self._list_of_degree
+        )
+
+        assert num_of_ranks == len(self._strategy_rank_list), (
+            f"There are total {len(self._strategy_rank_list)} ranks, but need {num_of_ranks} ranks in this strategy."
+        )
+
+        for fused_strategy in self._fused_strategy_dict.values():
+            for strategy in fused_strategy:
+                assert strategy in self._list_of_strategy_name, (
+                    f"Can not fuse strategy {strategy} without defined previous."
+                )
+
+    def _create_fused_group(self):
+        for name in self._fused_strategy_dict:
+            fused_strategy = self._fused_strategy_dict[name]
+            non_fused_strategy = list(
+                set(self._list_of_strategy_name).difference(fused_strategy)
+            )
+            non_fused_ranges = []
+            for strategy in non_fused_strategy:
+                non_fused_ranges.append(
+                    range(self._name_to_degree_dict[strategy])
+                )
+            fused_ranges = []
+            for strategy in fused_strategy:
+                fused_ranges.append(range(self._name_to_degree_dict[strategy]))
+
+            rank_list = []
+            for non_fused_ranks in itertools.product(*non_fused_ranges):
+                coord_dict = {}
+                ranks = []
+                for i, non_fused_rank in enumerate(non_fused_ranks):
+                    coord_dict[non_fused_strategy[i]] = non_fused_rank
+                for fused_ranks in itertools.product(*fused_ranges):
+                    for i, fused_rank in enumerate(fused_ranks):
+                        coord_dict[fused_strategy[i]] = fused_rank
+                    ranks.append(
+                        self._coord_to_rank_dict[self._coordinate(**coord_dict)]
+                    )
+                rank_list.append(ranks)
+            self._name_to_fused_group_dict[name] = StrategyGroupBase(rank_list)
+
+    def _calc_rank_list(self, strategy_axis):
+        ranges = []
+        for idx, degree in enumerate(self._list_of_degree):
+            if idx == strategy_axis:
+                continue
+            ranges.append(range(degree))
+
+        rank_list = []
+        for coord in itertools.product(*ranges):
+            ranks = []
+            for val in range(self._list_of_degree[strategy_axis]):
+                coord_list = list(coord)
+                coord_list.insert(strategy_axis, val)
+                ranks.append(
+                    self._coord_to_rank_dict[self._coordinate(*coord_list)]
+                )
+            rank_list.append(ranks)
+
+        return rank_list

python/paddle/distributed/fleet/base/private_helper_function.py

@@ -0,0 +1,33 @@
+#   Copyright (c) 2020 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__ = []
+
+
+def wait_server_ready(endpoints):
+    """
+    Wait until parameter servers are ready, use connext_ex to detect
+    port readiness.
+
+    Args:
+    endpoints (list|tuple): endpoints string list, like:
+    ["127.0.0.1:8080", "127.0.0.1:8081"]
+
+    Examples:
+        .. code-block:: pycon
+
+             >>> wait_server_ready(["127.0.0.1:8080", "127.0.0.1:8081"])
+    """
+    return

python/paddle/distributed/fleet/base/runtime_factory.py

@@ -0,0 +1,39 @@
+#   Copyright (c) 2020 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 ...ps.the_one_ps import TheOnePSRuntime
+from ..runtime.collective_runtime import CollectiveRuntime
+
+__all__ = []
+
+
+class RuntimeFactory:
+    def __init__(self):
+        pass
+
+    def _create_runtime(self, context):
+        # add collective && pslib mode
+        if context.get("use_fleet_ps"):
+            ps_runtime = TheOnePSRuntime()
+            ps_runtime._set_basic_info(context)
+            return ps_runtime
+        if context["role_maker"]._is_collective:
+            collective_runtime = CollectiveRuntime()
+            collective_runtime._set_basic_info(context)
+            return collective_runtime
+
+        k_steps = context["valid_strategy"].a_sync_configs["k_steps"]
+        if not context["role_maker"]._is_collective and k_steps >= 0:
+            ps_runtime = TheOnePSRuntime()
+            ps_runtime._set_basic_info(context)
+            return ps_runtime

python/paddle/distributed/fleet/base/strategy_compiler.py

@@ -0,0 +1,227 @@
+#   Copyright (c) 2020 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__ = []
+
+
+def create_graph(optimizer_list):
+    nsize = len(optimizer_list)
+
+    edge = [[0] * nsize for _ in range(nsize)]  # adjacency matrix
+    indegree = [0] * nsize
+    for i, opt in enumerate(optimizer_list):
+        for j, opt_inner in enumerate(optimizer_list):
+            if opt._can_update(opt_inner):
+                edge[i][j] = 1  # weight
+                indegree[j] += 1
+
+    return edge, indegree
+
+
+def topo_sort(edge, indegree):
+    nsize = len(indegree)
+
+    topo = [-1] * nsize
+    for i in range(nsize):
+        j = 0
+        while j < nsize and indegree[j] != 0:
+            j += 1
+        assert j < nsize, 'The combination of meta optimizers contains ring'
+
+        topo[i] = j
+        indegree[j] = -1
+        for k in range(nsize):
+            if edge[j][k] != 0:
+                indegree[k] -= 1
+    return topo
+
+
+def floyd(edge):
+    nsize = len(edge)
+    max_len = -1
+    max_edge = [-1, -1]
+
+    max_path = [[[] for _ in range(nsize)] for _ in range(nsize)]
+    for i in range(nsize):
+        for j in range(nsize):
+            if edge[i][j] > 0:
+                max_path[i][j] = [j]
+
+                if edge[i][j] > max_len:
+                    max_len = edge[i][j]
+                    max_edge = [i, j]
+
+    # use floyd algorithm to find max_path
+    for k in range(nsize):
+        for i in range(nsize):
+            for j in range(nsize):
+                # if a-->b-->c, but a-/->c, can only apply a-->b or b-->c,
+                # however if a-->b-->c, and a-->c, can apply a->b->c
+                if edge[i][j] == 0:
+                    continue
+
+                if edge[i][k] == 0 or edge[k][j] == 0:
+                    continue
+
+                if edge[i][j] < edge[i][k] + edge[k][j]:
+                    edge[i][j] = edge[i][k] + edge[k][j]
+                    max_path[i][j] = max_path[i][k] + max_path[k][j]
+
+                    max_len = edge[i][j]
+                    max_edge = [i, j]
+
+    if max_len == -1:
+        return [0]
+
+    return [max_edge[0]] + max_path[max_edge[0]][max_edge[1]]
+
+
+def maximum_path_len_algo(optimizer_list):
+    if len(optimizer_list) == 0:
+        return None
+
+    edge, indegree = create_graph(optimizer_list)
+    topo_sort(edge, indegree)
+    max_path = floyd(edge)
+    candidate = []
+    for idx in max_path:
+        candidate.append(optimizer_list[idx])
+
+    for idx, opt in enumerate(candidate[:-1]):
+        opt._update_inner_optimizer(candidate[idx + 1])
+
+    return candidate
+
+
+class StrategyCompilerBase:
+    def __init__(self):
+        pass
+
+
+class StrategyCompiler(StrategyCompilerBase):
+    """
+    StrategyCompiler is responsible for meta optimizers combination
+    Generally, a user can define several distributed strategies that
+    can generate several meta optimizer. The combination of these
+    meta optimizers should have the right order to apply the optimizers'
+    minimize function.
+    This class is responsible for the executable distributed optimizer
+    generation.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._meta_optimizers = []
+        self._graph_optimizers = []
+        self._valid_optimizer_list = None
+        self._user_defined_strategy = None
+        self._meta_optimizer_candidates = []
+        self._graph_optimizer_candidates = []
+
+    def _get_applied_meta_optimizer(self):
+        return self._meta_optimizers
+
+    def _get_applied_meta_list(self):
+        return [type(opt).__name__ for opt in self._meta_optimizers]
+
+    def _get_applied_graph_list(self):
+        return [type(opt).__name__ for opt in self._graph_optimizers]
+
+    def _get_valid_strategy(self, dist_strategy, can_not_apply_optimizer_list):
+        import copy
+
+        valid_strategy = copy.deepcopy(dist_strategy)
+        invalid_optimizers = []
+        for candidate in self._meta_optimizer_candidates:
+            is_valid = False
+            for valid in self._meta_optimizers:
+                if candidate.__class__.__name__ == valid.__class__.__name__:
+                    is_valid = True
+                    break
+            if not is_valid:
+                invalid_optimizers.append(candidate)
+        for opt in invalid_optimizers:
+            opt._disable_strategy(valid_strategy)
+        for opt in can_not_apply_optimizer_list:
+            opt._disable_strategy(valid_strategy)
+        return valid_strategy
+
+    """
+    Meta Optimizer Type A: rewrite forward, backward. e.g. recompute, async, sync, pipeline.
+                           results will be split in async, sync, pipeline
+    Meta Optimizer Type B: rewrite forward,
+                           e.g. AMP and the corresponding backward is generated by rewritten forward
+    Meta Optimizer Type B: rewrite backward. e.g. gradient fusion
+    Meta Optimizer Type D: rewrite optimize. e.g. lars, lamb, localsgd, gradient merge, dgc
+    Meta Optimizer Type E: only transpile to Graph structure for runtime,
+                           currently, grad fusion and kernel fusion, sync batch-norm included.
+                           we will remove grad fusion and sync batch-norm
+    """
+
+    def generate_optimizer(
+        self,
+        loss,
+        role_maker,
+        optimizer,
+        user_defined_strategy,
+        meta_optimizer_list,
+        graph_optimizer_list,
+    ):
+        self._user_defined_strategy = user_defined_strategy
+        self._meta_optimizer_candidates = meta_optimizer_list
+        self._graph_optimizer_candidates = graph_optimizer_list
+
+        if len(meta_optimizer_list) == 0 and len(graph_optimizer_list) == 0:
+            return optimizer, None
+        else:
+            # currently, we use heuristic algorithm to select
+            # meta optimizers combinations
+            meta_optimizers = maximum_path_len_algo(meta_optimizer_list)
+            graph_optimizers = maximum_path_len_algo(graph_optimizer_list)
+            # should design a distributed strategy update interface
+            # when we have finally decided the combination of meta_optimizer
+            # and graph_optimizer, the corresponding distributed strategy
+            # should be updated.
+
+            self._meta_optimizers = (
+                [] if meta_optimizers is None else meta_optimizers
+            )
+            self._graph_optimizers = (
+                [] if graph_optimizers is None else graph_optimizers
+            )
+
+            return_meta = (
+                None if meta_optimizers is None else meta_optimizers[0]
+            )
+            return_graph = (
+                None if graph_optimizers is None else graph_optimizers[0]
+            )
+
+            if meta_optimizers is None or graph_optimizers is None:
+                return return_meta, return_graph
+
+            # do heuristic filter here, if any meta optimizer in graph optimizers is in
+            # any meta optimizers' black list, set return_graph to None
+            need_graph_opt = True
+            for graph_opt in graph_optimizers:
+                for program_opt in meta_optimizers:
+                    if (
+                        graph_opt.__class__.__name__
+                        in program_opt.meta_optimizers_black_list
+                    ):
+                        need_graph_opt = False
+            if not need_graph_opt:
+                return_graph = None
+
+            return return_meta, return_graph

python/paddle/distributed/fleet/base/strategy_group.py

@@ -0,0 +1,271 @@
+# 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 paddle.distributed as dist
+from paddle.distributed.fleet.layers.mpu import RNGStatesTracker
+
+
+class StrategyGroupBase:
+    """
+    The base class of communication group with distributed strategy.
+
+    Args:
+        list_of_ranks: A 2D-array, such as `[[0, 1, 2, 3], [4, 5, 6, 7]]`. Ranks in sublist represents
+    they are in the same communication group.
+
+    Returns:
+        The instance of strategy group.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle.distributed as dist
+            >>> from paddle.distributed.fleet.base.strategy_group import StrategyGroupBase
+
+            >>> dist.init_parallel_env()
+            >>> strategy_group = dist.fleet.base.strategy_group.StrategyGroupBase([[0, 1], [2, 3]])
+            >>> print(strategy_group.world_size)
+            2
+
+
+    """
+
+    def __init__(self, list_of_ranks):
+        """
+        Initialize the communication group.
+        """
+        assert dist.is_initialized(), (
+            "The global communication group need to be initialized."
+        )
+        assert len(list_of_ranks), "The list_of_ranks can not be empty."
+        self._rank = dist.get_rank()
+        self._list_of_ranks = list_of_ranks
+        self._group = self._create_group()
+        self.random_states_tracker = RNGStatesTracker()
+
+    def add_random_seed(self, name, seed):
+        """
+        Add random seed for current rank.
+        """
+        self.random_states_tracker.add(name, seed)
+
+    def get_random_states_tracker(self):
+        """
+        Get the random states tracker.
+        """
+        return self.random_states_tracker
+
+    @property
+    def world_size(self):
+        """
+        The world size of communication group.
+
+        Returns:
+            Integer if the world_size of each group are equal, or a list of world_size if they are not equal.
+        """
+        world_size_list = []
+        for ranks in self._list_of_ranks:
+            world_size_list.append(len(ranks))
+        is_value = all(
+            world_size == world_size_list[0] for world_size in world_size_list
+        )
+        return world_size_list[0] if is_value else world_size_list
+
+    @property
+    def group(self):
+        """
+        The communication group which current rank belongs to.
+
+        Returns:
+            Group if current rank only belong to single communication group, or a list of Group if it belongs many.
+        """
+        return self._group
+
+    def _create_group(self):
+        self.list_of_group = []
+        for ranks in self._list_of_ranks:
+            group = dist.new_group(ranks=ranks)
+            if self._rank in ranks:
+                self.list_of_group.append(group)
+
+        if not self.list_of_group:
+            return None
+        else:
+            return (
+                self.list_of_group[0]
+                if len(self.list_of_group) == 1
+                else self.list_of_group
+            )
+
+    def __repr__(self):
+        debug_str = f"seed: {self._seed}; "
+        if not self.list_of_group:
+            return debug_str + "No group."
+        for i in range(len(self.list_of_group)):
+            debug_str += f"Group[{i}]: {self.list_of_group[i]}; "
+        return debug_str
+
+
+class DPGroup(StrategyGroupBase):
+    """
+    The communication group strategy for data parallel.
+
+    Args:
+        list_of_ranks: A 2D-array, such as `[[0, 1, 2, 3], [4, 5, 6, 7]]`. Ranks in sublist represents
+    they are in the same communication group.
+
+    Returns:
+        The instance of data parallel strategy group.
+    """
+
+    def __init__(self, list_of_ranks):
+        super().__init__(list_of_ranks)
+        assert not isinstance(self.group, list), (
+            f"Rank {self._rank} belongs to multi dp groups"
+        )
+
+
+class MPGroup(StrategyGroupBase):
+    """
+    The communication group strategy for model parallel.
+
+    Args:
+        list_of_ranks: A 2D-array, such as `[[0, 1, 2, 3], [4, 5, 6, 7]]`. Ranks in sublist represents
+    they are in the same communication group.
+
+    Returns:
+        The instance of model parallel strategy group.
+    """
+
+    def __init__(self, list_of_ranks):
+        super().__init__(list_of_ranks)
+        assert not isinstance(self.group, list), (
+            f"Rank {self._rank} belongs to multi mp groups"
+        )
+
+
+class ShardingGroup(StrategyGroupBase):
+    """
+    The communication group strategy for sharding parallel.
+
+    Args:
+        list_of_ranks: A 2D-array, such as `[[0, 1, 2, 3], [4, 5, 6, 7]]`. Ranks in sublist represents
+    they are in the same communication group.
+
+    Returns:
+        The instance of sharding parallel strategy group.
+    """
+
+    def __init__(self, list_of_ranks):
+        super().__init__(list_of_ranks)
+        assert not isinstance(self.group, list), (
+            f"Rank {self._rank} belongs to multi sharding groups"
+        )
+
+
+class PPGroup(StrategyGroupBase):
+    """
+    The communication group strategy for pipeline parallel.
+
+    Args:
+        list_of_ranks: A 2D-array, such as `[[0, 1, 2, 3], [4, 5, 6, 7]]`. Ranks in sublist represents
+    they are in the same communication group.
+
+    Returns:
+        The instance of pipeline parallel strategy group.
+    """
+
+    def __init__(self, list_of_ranks):
+        super().__init__(list_of_ranks)
+        assert not isinstance(self.group, list), (
+            f"Rank {self._rank} belongs to multi pp groups"
+        )
+
+        self._send_next_group = None
+        self._send_prev_group = None
+        self._recv_next_group = None
+        self._recv_prev_group = None
+        self._rank_of_next_stage = None
+        self._rank_of_prev_stage = None
+
+        if self.world_size > 1:
+            self._create_p2p_group()
+
+    @property
+    def rank_of_prev_stage(self):
+        """
+        Rank of the previous pp stage.
+
+        Returns:
+            The global rank of previous pp stage. `None` if without previous.
+        """
+        return self._rank_of_prev_stage
+
+    @property
+    def rank_of_next_stage(self):
+        """
+        Rank of the next pp stage.
+
+        Returns:
+            The global rank of next pp stage. `None` if without next.
+        """
+        return self._rank_of_next_stage
+
+    @property
+    def p2p_groups(self):
+        """
+        Communication subgroup in order to switch data with previous and next stage.
+
+        Returns:
+            Four subgroups including send/recv to/from prev/next.
+        """
+        return (
+            self._send_next_group,
+            self._send_prev_group,
+            self._recv_next_group,
+            self._recv_prev_group,
+        )
+
+    def _create_p2p_group(self):
+        degree = self.world_size
+        for ranks in self._list_of_ranks:
+            for idx, rank in enumerate(ranks):
+                next_rank = ranks[(idx + 1) % degree]
+                prev_rank = ranks[(idx - 1) % degree]
+
+                if self._rank == rank:
+                    self._rank_of_next_stage = next_rank
+                    self._rank_of_prev_stage = prev_rank
+
+                next_group = dist.new_group(ranks=[rank, next_rank])
+
+                if self._rank == rank:
+                    self._send_next_group = next_group
+                elif self._rank == next_rank:
+                    self._recv_prev_group = next_group
+
+                prev_group = dist.new_group(ranks=[prev_rank, rank])
+                if self._rank == rank:
+                    self._send_prev_group = prev_group
+                elif self._rank == prev_rank:
+                    self._recv_next_group = prev_group
+
+        assert (
+            self._send_next_group
+            and self._send_prev_group
+            and self._recv_next_group
+            and self._recv_prev_group
+        ), f"Error occurs while creating p2p group for rank {self._rank}."

python/paddle/distributed/fleet/base/util_factory.py

@@ -0,0 +1,778 @@
+#   Copyright (c) 2020 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
+
+"""Fleet Utils."""
+"""distributed operations"""
+"""basic collective operations in python"""
+"""remote file system"""
+
+import os
+import re
+import subprocess
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Any, Literal
+
+import numpy as np
+from google.protobuf import text_format
+
+import paddle
+from paddle import framework
+from paddle.base import core
+from paddle.base.proto import framework_pb2
+from paddle.static import Program
+
+from ..utils.fs import FS
+from .graphviz import GraphPreviewGenerator
+
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+    from paddle import Tensor
+    from paddle._typing import NestedNumericSequence
+    from paddle.base.framework import Block
+    from paddle.distributed.fleet.base.distributed_strategy import (
+        DistributedStrategy,
+    )
+    from paddle.distributed.fleet.base.role_maker import PaddleCloudRoleMaker
+
+
+__all__ = []
+
+
+class UtilFactory:
+    def _create_util(self, context=None):
+        util = UtilBase()
+        if context is not None and "valid_strategy" in context:
+            util._set_strategy(context["valid_strategy"])
+        if context is not None and "role_maker" in context:
+            util._set_role_maker(context["role_maker"])
+        return util
+
+
+class UtilBase:
+    def __init__(self) -> None:
+        self.role_maker: PaddleCloudRoleMaker | None = None
+        self.dist_strategy: DistributedStrategy | None = None
+
+    def _set_strategy(self, dist_strategy: DistributedStrategy | None) -> None:
+        self.dist_strategy = dist_strategy
+
+    def _set_role_maker(self, role_maker: PaddleCloudRoleMaker | None) -> None:
+        self.role_maker = role_maker
+
+    def _set_file_system(self, fs_client: FS) -> None:
+        assert isinstance(fs_client, FS), (
+            "fs_client must be the instance of paddle.distributed.fleet.utils.FS"
+        )
+        self.fs_client = fs_client
+
+    def all_reduce(
+        self,
+        input: NestedNumericSequence | npt.NDArray[Any],
+        mode: Literal["sum", "min", "max"] = "sum",
+        comm_world: Literal["worker", "server", "all"] = "worker",
+    ) -> npt.NDArray[Any] | None:
+        """
+        All reduce `input` between specified collection. This is a distributed API.
+
+        Args:
+            input (list|tuple|numpy.array): The input variable to do all_reduce between specified collection.
+            mode (str): "sum" or "min" or "max".
+            comm_world (str, optional): Collection used to execute all_reduce operation. Supported collections include `worker` , `server` and `all` . The default is `worker` .
+
+        Returns:
+            output(Numpy.array|None): A numpy array with the same shape as the `input` .
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+                >>> # Save the following code in `train.py` , and then execute the command `fleetrun --server_num 2 --worker_num 2 train.py` .
+                >>> import paddle.distributed.fleet as fleet
+                >>> from paddle.distributed.fleet import PaddleCloudRoleMaker
+                >>> import sys
+                >>> import numpy as np
+                >>> import os
+
+                >>> os.environ["PADDLE_WITH_GLOO"] = "2"
+
+                >>> def train():
+                ...     role = PaddleCloudRoleMaker(
+                ...         is_collective=False,
+                ...         init_gloo=True,
+                ...         path="./tmp_gloo",
+                ...     )
+                ...     fleet.init(role)
+                ...
+                ...     if fleet.is_server():
+                ...         input = np.array([1, 2])
+                ...         output = fleet.util.all_reduce(input, "sum", "server")
+                ...         print(output)  # [2, 4]
+                ...     elif fleet.is_worker():
+                ...         input = np.array([3, 4])
+                ...         output = fleet.util.all_reduce(input, "sum", "worker")
+                ...         print(output)  # [6, 8]
+                ...     output = fleet.util.all_reduce(input, "sum", "all")
+                ...     print(output)  # [8, 12]
+
+                >>> if __name__ == "__main__":
+                ...     train()
+        """
+        if isinstance(input, tuple):
+            input = list(input)
+        return self.role_maker._all_reduce(input, mode, comm_world)
+
+    def barrier(
+        self, comm_world: Literal["worker", "server", "all"] = "worker"
+    ) -> None:
+        """
+        Barrier between specified collection.
+
+        Args:
+            comm_world (str, optional): Collection used to execute barrier operation. Supported collections include `worker` , `server` and `all` . The default is `worker` .
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+                >>> # Save the following code in `train.py` , and then execute the command `fleetrun --server_num 2 --worker_num 2 train.py` .
+                >>> import paddle.distributed.fleet as fleet
+                >>> from paddle.distributed.fleet import PaddleCloudRoleMaker
+                >>> import sys
+                >>> import os
+
+                >>> os.environ["PADDLE_WITH_GLOO"] = "2"
+
+                >>> def train():
+                ...     role = PaddleCloudRoleMaker(
+                ...         is_collective=False,
+                ...         init_gloo=True,
+                ...         path="./tmp_gloo",
+                ...     )
+                ...     fleet.init(role)
+                ...
+                ...     if fleet.is_server():
+                ...         fleet.util.barrier("server")
+                ...         print("all server arrive here")  # all server arrive here
+                ...     elif fleet.is_worker():
+                ...         fleet.util.barrier("worker")
+                ...         print("all server arrive here")  # all server arrive here
+                ...     fleet.util.barrier("all")
+                ...     print("all servers and workers arrive here")  # all servers and workers arrive here
+
+                >>> if __name__ == "__main__":
+                ...     train()
+        """
+        self.role_maker._barrier(comm_world)
+
+    def all_gather(
+        self,
+        input: float,
+        comm_world: Literal["worker", "server", "all"] = "worker",
+    ) -> list[float]:
+        """
+        All gather `input` between specified collection.
+
+        Args:
+            input (Int|Float): The input variable to do all_gather between specified collection.
+            comm_world (str, optional): Collection used to execute all_reduce operation. Supported collections include `worker` , `server` and `all` . The default is `worker` .
+
+        Returns:
+            output (List): A list of gathered values.
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+                >>> # Save the following code in `train.py` , and then execute the command `fleetrun --server_num 2 --worker_num 2 train.py` .
+                >>> import paddle.distributed.fleet as fleet
+                >>> from paddle.distributed.fleet import PaddleCloudRoleMaker
+                >>> import sys
+                >>> import os
+
+                >>> os.environ["PADDLE_WITH_GLOO"] = "2"
+
+                >>> def train():
+                ...     role = PaddleCloudRoleMaker(
+                ...         is_collective=False,
+                ...         init_gloo=True,
+                ...         path="./tmp_gloo",
+                ...     )
+                ...     fleet.init(role)
+                ...
+                ...     if fleet.is_server():
+                ...         input = fleet.server_index()
+                ...         output = fleet.util.all_gather(input, "server")
+                ...         print(output)  # [0, 1]
+                ...     elif fleet.is_worker():
+                ...         input = fleet.worker_index()
+                ...         output = fleet.util.all_gather(input, "worker")
+                ...         print(output)  # [0, 1]
+                ...     output = fleet.util.all_gather(input, "all")
+                ...     print(output)  # [0, 1, 0, 1]
+
+                >>> if __name__ == "__main__":
+                ...     train()
+        """
+
+        return self.role_maker._all_gather(input, comm_world)
+
+    def _broadcast(self) -> None:
+        pass
+
+    def _scatter(self) -> None:
+        pass
+
+    def get_heter_file_shard(self, files: list[str]) -> list[str]:
+        if not isinstance(files, list):
+            raise TypeError("files should be a list of file need to be read.")
+        trainers = self.role_maker._worker_num()
+        trainer_id = self.role_maker._worker_index() - trainers
+        remainder = len(files) % trainers
+        blocksize = int(len(files) / trainers)
+
+        blocks = [blocksize] * trainers
+        for i in range(remainder):
+            blocks[i] += 1
+
+        trainer_files = [[]] * trainers
+        begin = 0
+        for i in range(trainers):
+            trainer_files[i] = files[begin : begin + blocks[i]]
+            begin += blocks[i]
+
+        return trainer_files[trainer_id]
+
+    def get_file_shard(self, files: list[str]) -> list[str]:
+        """
+        Split files before distributed training, and return filelist assigned to the current trainer.
+
+        .. code-block:: text
+
+            example 1: files is [a, b, c ,d, e]  and trainer_num = 2, then trainer
+                    0 gets [a, b, c] and trainer 1 gets [d, e].
+            example 2: files is [a, b], and trainer_num = 3, then trainer 0 gets
+                    [a], trainer 1 gets [b],  trainer 2 gets []
+
+        Args:
+            files(list): File list need to be read.
+
+        Returns:
+            List: Files belong to this worker.
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+                >>> import paddle.distributed.fleet as fleet
+                >>> from paddle.distributed.fleet import UserDefinedRoleMaker
+
+                >>> role = UserDefinedRoleMaker(
+                ...     is_collective=False,
+                ...     init_gloo=False,
+                ...     current_id=0,
+                ...     role=fleet.Role.WORKER,
+                ...     worker_endpoints=["127.0.0.1:6003", "127.0.0.1:6004"],
+                ...     server_endpoints=["127.0.0.1:6001", "127.0.0.1:6002"],
+                ... )
+                >>> fleet.init(role)
+
+                >>> files = fleet.util.get_file_shard(["file1", "file2", "file3"])
+                >>> print(files)
+                ["file1", "file2"]
+        """
+        if not isinstance(files, list):
+            raise TypeError("files should be a list of file need to be read.")
+
+        trainer_id = self.role_maker._worker_index()
+        trainers = self.role_maker._worker_num()
+
+        remainder = len(files) % trainers
+        blocksize = int(len(files) / trainers)
+
+        blocks = [blocksize] * trainers
+        for i in range(remainder):
+            blocks[i] += 1
+
+        trainer_files = [[]] * trainers
+        begin = 0
+        for i in range(trainers):
+            trainer_files[i] = files[begin : begin + blocks[i]]
+            begin += blocks[i]
+
+        return trainer_files[trainer_id]
+
+    def print_on_rank(self, message: str, rank_id: int) -> None:
+        """
+        Worker of rank `rank_id` print some message.
+
+        Args:
+            message(str): Log to be printed.
+            rank_id(int): trainer id.
+
+        Examples:
+
+            .. code-block:: pycon
+
+                >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+                >>> import paddle.distributed.fleet as fleet
+                >>> from paddle.distributed.fleet import UserDefinedRoleMaker
+
+                >>> role = UserDefinedRoleMaker(
+                ...     is_collective=False,
+                ...     init_gloo=False,
+                ...     current_id=0,
+                ...     role=fleet.Role.WORKER,
+                ...     worker_endpoints=["127.0.0.1:6003", "127.0.0.1:6004"],
+                ...     server_endpoints=["127.0.0.1:6001", "127.0.0.1:6002"],
+                ... )
+                >>> fleet.init(role)
+
+                >>> fleet.util.print_on_rank("I'm worker 0", 0)
+                I'm worker 0
+        """
+        if self.role_maker._worker_index() != rank_id:
+            return
+        print(message)
+
+    def _save_program(
+        self,
+        program: Program,
+        model_filename: str = '__model__',
+        is_text: bool = False,
+    ) -> None:
+        if is_text:
+            with open(model_filename, "w") as f:
+                f.write(str(program))
+        else:
+            with open(model_filename, "wb") as f:
+                f.write(program.desc.serialize_to_string())
+
+    def _load_program(self, path: str, is_text: bool) -> Program:
+        def load_program_binary(path):
+            """load program from binary string file"""
+            with open(path, "rb") as f:
+                program_desc_str = f.read()
+            return Program.parse_from_string(program_desc_str)
+
+        def load_program_text(path):
+            """load program from human-readable text file"""
+            with open(path, "r") as f:
+                program_desc_text = f.read()
+
+            prog_desc = framework_pb2.ProgramDesc()
+            text_format.Merge(program_desc_text, prog_desc)
+            return Program.parse_from_string(prog_desc.SerializeToString())
+
+        if is_text:
+            return load_program_text(path)
+        else:
+            return load_program_binary(path)
+
+    def _program_type_trans(
+        self, prog_dir: str, prog_fn: str, is_text: bool
+    ) -> str:
+        prog = self._load_program(os.path.join(prog_dir, prog_fn), is_text)
+        prog_out_fn = prog_fn + ".bin" if is_text else prog_fn + ".pbtxt"
+        self._save_program(
+            prog, os.path.join(prog_dir, prog_out_fn), 1 - is_text
+        )
+        return prog_out_fn
+
+    def _visualize_graphviz(
+        self, program: Program, output_dir: str, output_filename: str
+    ) -> None:
+        block = program.global_block()
+        dot_path = os.path.join(output_dir, output_filename + '.dot')
+        pdf_path = os.path.join(output_dir, output_filename + '.pdf')
+        draw_block_graphviz(block, path=dot_path)
+        cmd = ["dot", "-Tpdf", dot_path, "-o", pdf_path]
+        p = subprocess.Popen(
+            cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        p.wait()
+
+    def _proto_check(self, config: Any) -> bool:
+        train_prog = self._load_program(
+            config.train_prog_path, config.is_text_train_program
+        )
+        pruned_prog = self._load_program(
+            config.pruned_prog_path, config.is_text_pruned_program
+        )
+
+        is_match = True
+
+        pruned_vars = [
+            (v.name, v)
+            for v in pruned_prog.list_vars()
+            if paddle.static.io.is_persistable(v)
+        ]
+        pruned_vars = OrderedDict(pruned_vars)
+        pruned_vars_name = list(pruned_vars)
+        print(f"persistable vars in pruned program: {pruned_vars_name}")
+
+        # feed and fetch op is added in pruned program when pruning, not need to be found in train program
+        feed_fetch_type_list = [
+            core.VarDesc.VarType.FEED_MINIBATCH,
+            core.VarDesc.VarType.FETCH_LIST,
+        ]
+
+        for var_name in pruned_vars:
+            var = pruned_vars[var_name]
+            # feed and fetch op is added in pruned program when pruning, not need to be found in train program
+            if var.type in feed_fetch_type_list:
+                break
+            try:
+                train_prog_var = train_prog.global_block().var(var_name)
+            except ValueError as e:
+                print(
+                    f"Not find variable '{var_name}' in train program. please check pruning."
+                )
+                is_match = False
+                continue
+            if (
+                var.shape != train_prog_var.shape
+                or var.dtype != train_prog_var.dtype
+            ):
+                print(
+                    f"variable: {var_name} not match. in pruned program shape: {var.shape} dtype:{var.dtype}, in train program shape: {train_prog_var.shape} dtype: {train_prog_var.dtype}"
+                )
+                is_match = False
+        return is_match
+
+    def _params_check(
+        self, config: Any
+    ) -> list[Tensor] | list[npt.NDArray[Any]] | Literal[False]:
+        def feed_gen(batch_size, feeded_vars_dims, feeded_vars_filelist):
+            def reader(batch_size, fn, dim):
+                data = []
+                if isinstance(dim, (list, tuple)):
+                    shape = list(dim)
+                    _temp = 1
+                    for x in dim:
+                        _temp = _temp * x
+                    dim = _temp
+                else:
+                    shape = [dim]
+
+                shape = [batch_size, *shape]
+                dim = dim * batch_size
+
+                for line in open(fn, 'r'):
+                    fields = line.strip().split(' ')
+                    fields = [float(d) for d in fields]
+                    while len(fields) >= dim:
+                        tmp = fields[:dim]
+                        fields = fields[dim:]
+                        data.append(np.array(tmp).reshape(shape))
+                return data
+
+            batch_feed = []
+            for i, fn in enumerate(feeded_vars_filelist):
+                batch_feed.append(reader(batch_size, fn, feeded_vars_dims[i]))
+            return batch_feed
+
+        prog = self._load_program(
+            os.path.join(config.dump_model_dir, config.dump_program_filename),
+            config.is_text_dump_program,
+        )
+        if config.is_text_dump_program:
+            model_filename = self._program_type_trans(
+                config.dump_model_dir,
+                config.dump_program_filename,
+                config.is_text_dump_program,
+            )
+
+        saved_params = [
+            v for v in prog.list_vars() if paddle.static.io.is_persistable(v)
+        ]
+        print(
+            f"persistable vars in dump program: {[v.name for v in saved_params]}"
+        )
+
+        def check_not_expected_ops(prog, not_expected_op_types):
+            op_types_set = set()
+            for op in prog.global_block().ops:
+                if (
+                    op.type in not_expected_op_types
+                    and op.type not in op_types_set
+                ):
+                    op_types_set.add(op.type)
+            return op_types_set
+
+        not_expected_op_types = check_not_expected_ops(prog, ["lookup_table"])
+        if len(not_expected_op_types) > 0:
+            print(
+                f"find op type '{list(not_expected_op_types)}' in program, please check if your program is pruned correctly !"
+            )
+            return False
+
+        place = framework.CPUPlace()
+        exe = paddle.static.Executor(place)
+        scope = paddle.static.Scope()
+        with paddle.static.scope_guard(scope):
+            (
+                inference_program,
+                feed_target_names,
+                fetch_targets,
+            ) = paddle.distributed.io.load_inference_model_distributed(
+                config.dump_model_dir,
+                exe,
+                model_filename=model_filename,
+                params_filename=config.save_params_filename,
+            )
+
+            # check program vars and saved vars shape
+            orig_para_shape = {
+                each_var.name: tuple(each_var.desc.shape())
+                for each_var in saved_params
+            }
+            for each_var in saved_params:
+                var_temp = paddle.static.global_scope().find_var(each_var.name)
+                assert var_temp is not None, (
+                    "can't not find var: " + each_var.name
+                )
+                new_shape = (np.array(var_temp.get_tensor())).shape
+                assert each_var.name in orig_para_shape, (
+                    each_var.name + "MUST in var list"
+                )
+                orig_shape = orig_para_shape.get(each_var.name)
+                if new_shape != orig_shape:
+                    raise RuntimeError(
+                        f"Shape not matching: the Program requires a parameter with a shape of ({orig_shape}), "
+                        f"while the loaded parameter (namely [ {each_var.name} ]) has a shape of  ({new_shape})."
+                    )
+
+            # check feed/fetch vars in program and config
+            feed_config = config.feed_config
+            fetch_config = config.fetch_config
+            fetch_targets_names = [v.name for v in fetch_targets]
+            if not feed_target_names:
+                print("warning! no feed targets in program.")
+            if not fetch_targets_names:
+                print("warning! no fetch targets in program.")
+            fetch_list = fetch_targets
+            feed_name_list = feed_target_names
+            if (
+                feed_config.feeded_vars_names is not None
+                and feed_target_names != feed_config.feeded_vars_names
+            ):
+                print(
+                    f"warning! feed vars in program and config are diff: feed in program: {feed_target_names}. feed in config {feed_config.feeded_vars_names}."
+                )
+                feed_name_list = feed_config.feeded_vars_names
+                # remove feed op in inference_program. new feed op will be added in exe.run
+                global_block = inference_program.global_block()
+                need_to_remove_op_index = []
+                for i, op in enumerate(global_block.ops):
+                    op.desc.set_is_target(False)
+                    if op.type == "feed":  # only remove feed op here
+                        need_to_remove_op_index.append(i)
+                for index in need_to_remove_op_index[::-1]:
+                    global_block._remove_op(index)
+            if (
+                fetch_config.fetch_vars_names is not None
+                and fetch_targets_names != fetch_config.fetch_vars_names
+            ):
+                print(
+                    f"warning! fetch vars in program and config are diff: fetch in program: {fetch_targets_names}. fetch in config {fetch_config.fetch_vars_names}."
+                )
+                fetch_list = [
+                    inference_program.global_block().var(i)
+                    for i in fetch_config.fetch_vars_names
+                ]
+                # remove fetch op in inference_program. new fetch op will be added in exe.run
+                global_block = inference_program.global_block()
+                need_to_remove_op_index = []
+                for i, op in enumerate(global_block.ops):
+                    op.desc.set_is_target(False)
+                    if op.type == "fetch":  # only remove fetch op here
+                        need_to_remove_op_index.append(i)
+                for index in need_to_remove_op_index[::-1]:
+                    global_block._remove_op(index)
+
+            # if fetch_list have lod tensor
+            return_numpy = all(v.lod_level == 0 for v in fetch_list)
+
+            # try dump fetch_targets
+            feed_tensors = []
+            assert (
+                len(feed_config.feeded_vars_names)
+                == len(feed_config.feeded_vars_dims)
+                == len(feed_config.feeded_vars_types)
+            )
+            # check program vars and feed tensor shape in config
+            for i in range(len(feed_config.feeded_vars_names)):
+                var = inference_program.global_block().var(
+                    feed_config.feeded_vars_names[i]
+                )
+                if not isinstance(
+                    feed_config.feeded_vars_dims[i], (list, tuple)
+                ):
+                    tensor_shape = (feed_config.feeded_vars_dims[i],)
+                else:
+                    tensor_shape = tuple(feed_config.feeded_vars_dims[i])
+                feed_config.feeded_vars_dims[i] = tensor_shape
+                var_shape = var.shape[1:]
+                if tensor_shape != var_shape:
+                    raise RuntimeError(
+                        f"feed variable '{feed_config.feeded_vars_names[i]}' shape not match. infer program  shape: {var_shape}. feed tensor shape: {tensor_shape}"
+                    )
+
+            if not feed_config.feeded_vars_filelist:
+                print("generate random feed vars.")
+                for i in range(len(feed_config.feeded_vars_names)):
+                    var = inference_program.global_block().var(
+                        feed_config.feeded_vars_names[i]
+                    )
+                    # create fake feed tensor. if lod_level > 1, should create_lod_tensor()
+                    if var.lod_level == 0:
+                        feed_tensors.append(
+                            np.array(
+                                np.random.random(
+                                    (
+                                        config.batch_size,
+                                        *feed_config.feeded_vars_dims[i],
+                                    )
+                                ),
+                                dtype=feed_config.feeded_vars_types[i],
+                            )
+                        )
+                    elif var.lod_level == 1:
+                        t = np.array(
+                            np.random.random(
+                                (
+                                    config.batch_size,
+                                    *feed_config.feeded_vars_dims[i],
+                                )
+                            ),
+                            dtype=feed_config.feeded_vars_types[i],
+                        )
+                        feed_tensors.append(
+                            paddle.base.create_lod_tensor(
+                                t, [[1] * config.batch_size], place
+                            )
+                        )
+                    else:
+                        raise RuntimeError(
+                            "vars with lod_level >= 2 is not supported now in this infer program check tool."
+                        )
+                results = exe.run(
+                    inference_program,
+                    feed={
+                        name: feed_tensors[i]
+                        for i, name in enumerate(feed_name_list)
+                    },
+                    fetch_list=fetch_list,
+                    return_numpy=return_numpy,
+                )
+            else:
+                print(
+                    f"load feed vars from files: {feed_config.feeded_vars_filelist}."
+                )
+                feed_vars = [
+                    inference_program.global_block().var(
+                        feed_config.feeded_vars_names[i]
+                    )
+                    for i in range(len(feed_config.feeded_vars_names))
+                ]
+                feeder = paddle.base.DataFeeder(
+                    feed_list=feed_vars, place=place
+                )
+                batch_feed = feed_gen(
+                    config.batch_size,
+                    feed_config.feeded_vars_dims,
+                    feed_config.feeded_vars_filelist,
+                )
+                slots = [batch_feed]
+                results = exe.run(
+                    inference_program,
+                    feed=feeder.feed(slots),
+                    fetch_list=fetch_list,
+                    return_numpy=return_numpy,
+                )
+            for i, v in enumerate(fetch_list):
+                print(f"fetch_targets name: {v.name}")
+                print(f"fetch_targets: {results[i]}")
+            return results
+
+
+def draw_block_graphviz(
+    block: Block, highlights: list[str] | None = None, path: str = "./temp.dot"
+) -> None:
+    '''
+    Generate a debug graph for block.
+    Args:
+        block(Block): a block.
+    '''
+    graph = GraphPreviewGenerator("some graph")
+    # collect parameters and args
+    protostr = block.desc.serialize_to_string()
+    desc = framework_pb2.BlockDesc.FromString(bytes(protostr))
+
+    def need_highlight(name: str) -> bool:
+        if highlights is None:
+            return False
+        for pattern in highlights:
+            assert type(pattern) is str
+            if re.match(pattern, name):
+                return True
+        return False
+
+    # draw parameters and args
+    vars = {}
+    for var in desc.vars:
+        # TODO(gongwb): format the var.type
+        # create var
+        if var.persistable:
+            var_name = graph.add_param(
+                var.name,
+                str(var.type).replace("\n", "<br />", 1),
+                highlight=need_highlight(var.name),
+            )
+        else:
+            var_name = graph.add_arg(
+                var.name, highlight=need_highlight(var.name)
+            )
+        vars[var.name] = var_name
+
+    def add_op_link_var(op, var, op2var=False):
+        for arg in var.arguments:
+            if arg not in vars:
+                # add missing variables as argument
+                vars[arg] = graph.add_arg(arg, highlight=need_highlight(arg))
+            var_name = vars[arg]
+            highlight = need_highlight(op.description) or need_highlight(
+                var_name.description
+            )
+            if op2var:
+                graph.add_edge(op, var_name, highlight=highlight)
+            else:
+                graph.add_edge(var_name, op, highlight=highlight)
+
+    for op in desc.ops:
+        opn = graph.add_op(op.type, highlight=need_highlight(op.type))
+        for var in op.inputs:
+            add_op_link_var(opn, var, False)
+        for var in op.outputs:
+            add_op_link_var(opn, var, True)
+
+    graph(path, show=False)

python/paddle/distributed/fleet/cloud_utils.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2020 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 paddle.distributed.fleet.launch_utils import get_cluster, logger
+
+__all__ = []
+
+
+def get_cloud_cluster(
+    args_node_ips, device_mode, devices_per_proc, args_port=6170
+):
+    """
+    args_node_ips:string, device_mode:DeviceMode(Int), device_per_proc:list, args_port: int
+    """
+    # you can automatically get ip info while using paddlecloud multi nodes mode.
+    node_ips = os.getenv("PADDLE_TRAINERS")
+    assert node_ips is not None, "PADDLE_TRAINERS should not be None"
+
+    node_ip = os.getenv("POD_IP")
+    assert node_ip is not None, "POD_IP should not be None"
+
+    node_rank = os.getenv("PADDLE_TRAINER_ID")
+    assert node_rank is not None, "PADDLE_TRAINER_ID should not be None"
+
+    paddle_ports_num = int(os.getenv("TRAINER_PORTS_NUM"))
+    assert paddle_ports_num is not None, "TRAINER_PORTS_NUM should not be None"
+
+    node_ips = node_ips.split(",")
+    num_nodes = len(node_ips)
+    node_rank = int(node_rank)
+
+    if args_node_ips != "127.0.0.1" and args_node_ips != ",".join(node_ips):
+        logger.warning(
+            f"Please NOTE: When using paddlecloud, cluster_node_ips is \
+automatically got from PADDLE_TRAINERS(multi nodes) or POD_IP(single node).\
+Your input cluster_node_ips: {args_node_ips} doesn't equals to IPs: {node_ips} from \
+paddlecloud environment."
+        )
+
+    # DISTRIBUTED_TRAINER_ENDPOINTS: new environment since paddlecloud 1.8.4
+    # e.g: DISTRIBUTED_TRAINER_ENDPOINTS="ip1:port1,ip1:port2,ip1:port3,ip1:port4,ip2:port5,ip2:port6,ip2:port7,ip2:port8"
+    trainer_endpoints = os.getenv("DISTRIBUTED_TRAINER_ENDPOINTS")
+    if trainer_endpoints is None:
+        started_port = args_port
+        if num_nodes > 1:
+            try:
+                paddle_port = int(os.getenv("PADDLE_PORT", ""))
+
+                if (
+                    paddle_ports_num >= len(devices_per_proc)
+                    and paddle_port != args_port
+                ):
+                    logger.warning(f"Use Cloud specified port:{paddle_port}.")
+                    started_port = paddle_port
+
+            except Exception as e:
+                print(e)
+
+        if started_port is None:
+            started_port = 6170
+        ports = list(range(started_port, started_port + len(devices_per_proc)))
+        trainer_endpoints = []
+        for ip in node_ips:
+            trainer_endpoints.append([f"{ip}:{port}" for port in ports])
+    else:
+        trainer_endpoints_ori = trainer_endpoints.split(",")
+        trainer_endpoints = []
+        assert num_nodes * paddle_ports_num == len(trainer_endpoints_ori)
+        for i in range(num_nodes):
+            trainer_endpoints.append(
+                trainer_endpoints_ori[
+                    i * paddle_ports_num : (i + 1) * paddle_ports_num
+                ]
+            )
+
+    logger.debug(
+        f"parsed from args: node_ips:{node_ips} \
+        node_ip:{node_ip} node_rank:{node_rank} trainer_endpoints:{trainer_endpoints}"
+    )
+
+    cluster, pod = get_cluster(
+        node_ips, node_ip, trainer_endpoints, device_mode, devices_per_proc
+    )
+    return cluster, cluster.pods[node_rank]
+
+
+def use_paddlecloud():
+    node_ips = os.getenv("PADDLE_TRAINERS")
+    node_ip = os.getenv("POD_IP")
+    node_rank = os.getenv("PADDLE_TRAINER_ID")
+    paddle_ports_num = os.getenv("TRAINER_PORTS_NUM")
+    if (
+        node_ips is None
+        or node_ip is None
+        or node_rank is None
+        or paddle_ports_num is None
+    ):
+        return False
+    else:
+        return True
+
+
+def get_trainers_num():
+    return int(os.getenv("PADDLE_TRAINERS_NUM", "1"))

python/paddle/distributed/fleet/data_generator/__init__.py

@@ -0,0 +1,16 @@
+#   Copyright (c) 2020 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
+
+from .data_generator import DataGenerator, MultiSlotDataGenerator  # noqa: F401
+
+__all__ = []

python/paddle/distributed/fleet/data_generator/data_generator.py

@@ -0,0 +1,391 @@
+# Copyright (c) 2020 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 sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+__all__ = []
+
+
+class DataGenerator:
+    """
+    DataGenerator is a general Base class for user to inherit
+    A user who wants to define his/her own python processing logic
+    with paddle.distributed.InMemoryDataset/QueueDataset should
+    inherit this class.
+    """
+
+    def __init__(self):
+        self._proto_info = None
+        self.batch_size_ = 32
+
+    def set_batch(self, batch_size):
+        '''
+        Set batch size of current DataGenerator
+        This is necessary only if a user wants to define generator_batch
+
+        Example:
+
+            .. code-block:: pycon
+
+                >>> import paddle.distributed.fleet.data_generator as dg
+                >>> class MyData(dg.DataGenerator):
+                ...     def generate_sample(self, line):
+                ...         def local_iter():
+                ...             int_words = [int(x) for x in line.split()]
+                ...             yield ("words", int_words)
+                ...
+                ...         return local_iter
+                ...
+                ...     def generate_batch(self, samples):
+                ...         def local_iter():
+                ...             for s in samples:
+                ...                 yield ("words", s[1].extend([s[1][0]]))
+                >>> mydata = MyData()
+                >>> mydata.set_batch(128)
+
+        '''
+        self.batch_size_ = batch_size
+
+    def run_from_memory(self):
+        '''
+        This function generator data from memory, it is usually used for
+        debug and benchmarking
+
+        Example:
+            .. code-block:: pycon
+
+                >>> # doctest: +SKIP('raise NotImplementedError')
+                >>> import paddle.distributed.fleet.data_generator as dg
+                >>> class MyData(dg.DataGenerator):
+                ...     def generate_sample(self, line):
+                ...         def local_iter():
+                ...             yield ("words", [1, 2, 3, 4])
+                ...
+                ...         return local_iter
+                >>> mydata = MyData()
+                >>> mydata.run_from_memory()
+        '''
+        batch_samples = []
+        line_iter = self.generate_sample(None)
+        for user_parsed_line in line_iter():
+            if user_parsed_line is None:
+                continue
+            batch_samples.append(user_parsed_line)
+            if len(batch_samples) == self.batch_size_:
+                batch_iter = self.generate_batch(batch_samples)
+                for sample in batch_iter():
+                    sys.stdout.write(self._gen_str(sample))
+                batch_samples = []
+        if len(batch_samples) > 0:
+            batch_iter = self.generate_batch(batch_samples)
+            for sample in batch_iter():
+                sys.stdout.write(self._gen_str(sample))
+
+    def run_from_stdin(self):
+        '''
+        This function reads the data row from stdin, parses it with the
+        process function, and further parses the return value of the
+        process function with the _gen_str function. The parsed data will
+        be wrote to stdout and the corresponding protofile will be
+        generated.
+
+        Example:
+
+            .. code-block:: pycon
+
+                >>> import paddle.distributed.fleet.data_generator as dg
+                >>> class MyData(dg.DataGenerator):
+                ...     def generate_sample(self, line):
+                ...         def local_iter():
+                ...             int_words = [int(x) for x in line.split()]
+                ...             yield ("words", [int_words])
+                ...
+                ...         return local_iter
+                >>> mydata = MyData()
+                >>> mydata.run_from_stdin()
+
+        '''
+        batch_samples = []
+        for line in sys.stdin:
+            line_iter = self.generate_sample(line)
+            for user_parsed_line in line_iter():
+                if user_parsed_line is None:
+                    continue
+                batch_samples.append(user_parsed_line)
+                if len(batch_samples) == self.batch_size_:
+                    batch_iter = self.generate_batch(batch_samples)
+                    for sample in batch_iter():
+                        sys.stdout.write(self._gen_str(sample))
+                    batch_samples = []
+        if len(batch_samples) > 0:
+            batch_iter = self.generate_batch(batch_samples)
+            for sample in batch_iter():
+                sys.stdout.write(self._gen_str(sample))
+
+    def _gen_str(self, line):
+        '''
+        Further processing the output of the process() function rewritten by
+        user, outputting data that can be directly read by the datafeed,and
+        updating proto_info information.
+
+        Args:
+            line(str): the output of the process() function rewritten by user.
+
+        Returns:
+            Return a string data that can be read directly by the datafeed.
+        '''
+        raise NotImplementedError(
+            "pls use MultiSlotDataGenerator or PairWiseDataGenerator"
+        )
+
+    def generate_sample(self, line):
+        '''
+        This function needs to be overridden by the user to process the
+        original data row into a list or tuple.
+
+        Args:
+            line(str): the original data row
+
+        Returns:
+            Returns the data processed by the user.
+            The data format is list or tuple:
+            [(name, [feasign, ...]), ...]
+            or ((name, [feasign, ...]), ...)
+
+            For example:
+            [("words", [1926, 08, 17]), ("label", [1])]
+            or (("words", [1926, 08, 17]), ("label", [1]))
+
+        Note:
+            The type of feasigns must be in int or float. Once the float
+            element appears in the feasign, the type of that slot will be
+            processed into a float.
+
+        Example:
+
+            .. code-block:: pycon
+
+                >>> import paddle.distributed.fleet.data_generator as dg
+                >>> class MyData(dg.DataGenerator):
+                ...     def generate_sample(self, line):
+                ...         def local_iter():
+                ...             int_words = [int(x) for x in line.split()]
+                ...             yield ("words", [int_words])
+                ...
+                ...         return local_iter
+        '''
+        raise NotImplementedError(
+            "Please rewrite this function to return a list or tuple: "
+            + "[(name, [feasign, ...]), ...] or ((name, [feasign, ...]), ...)"
+        )
+
+    def generate_batch(self, samples):
+        '''
+        This function needs to be overridden by the user to process the
+        generated samples from generate_sample(self, str) function
+        It is usually used as batch processing when a user wants to
+        do preprocessing on a batch of samples, e.g. padding according to
+        the max length of a sample in the batch
+
+        Args:
+            samples(list tuple): generated sample from generate_sample
+
+        Returns:
+            a python generator, the same format as return value of generate_sample
+
+        Example:
+
+            .. code-block:: pycon
+
+                >>> import paddle.distributed.fleet.data_generator as dg
+                >>> class MyData(dg.DataGenerator):
+                ...     def generate_sample(self, line):
+                ...         def local_iter():
+                ...             int_words = [int(x) for x in line.split()]
+                ...             yield ("words", int_words)
+                ...
+                ...         return local_iter
+                ...
+                ...     def generate_batch(self, samples):
+                ...         def local_iter():
+                ...             for s in samples:
+                ...                 yield ("words", s[1].extend([s[1][0]]))
+                >>> mydata = MyData()
+                >>> mydata.set_batch(128)
+        '''
+
+        def local_iter():
+            yield from samples
+
+        return local_iter
+
+
+# TODO: guru4elephant
+# add more generalized DataGenerator that can adapt user-defined slot
+# for example, [(name, float_list), (name, str_list), (name, int_list)]
+class MultiSlotStringDataGenerator(DataGenerator):
+    def _gen_str(
+        self,
+        line: Sequence[tuple[str, list[str]]],
+    ) -> str:
+        '''
+        Further processing the output of the process() function rewritten by
+        user, outputting data that can be directly read by the MultiSlotDataFeed,
+        and updating proto_info information.
+
+        The input line will be in this format:
+            >>> [(name, [str(feasign), ...]), ...]
+            >>> or ((name, [str(feasign), ...]), ...)
+        The output will be in this format:
+            >>> [ids_num id1 id2 ...] ...
+
+        For example, if the input is like this:
+            >>> [("words", ["1926", "08", "17"]), ("label", ["1"])]
+            >>> or (("words", ["1926", "08", "17"]), ("label", ["1"]))
+        the output will be:
+            >>> 3 1234 2345 3456 1 1
+
+        Args:
+            line(str): the output of the process() function rewritten by user.
+
+        Returns:
+            Return a string data that can be read directly by the MultiSlotDataFeed.
+        '''
+        if isinstance(line, zip):
+            line = list(line)
+
+        if not isinstance(line, list) and not isinstance(line, tuple):
+            raise ValueError(
+                "the output of process() must be in list or tuple type"
+                "Examples: [('words', ['1926', '08', '17']), ('label', ['1'])]"
+            )
+        output = ""
+        for index, item in enumerate(line):
+            name, elements = item
+            if output:
+                output += " "
+            out_str = []
+            out_str.append(str(len(elements)))
+            out_str.extend(elements)
+            output += " ".join(out_str)
+        return output + "\n"
+
+
+class MultiSlotDataGenerator(DataGenerator):
+    def _gen_str(
+        self,
+        line: Sequence[tuple[str, list[float]]],
+    ) -> str:
+        '''
+        Further processing the output of the process() function rewritten by
+        user, outputting data that can be directly read by the MultiSlotDataFeed,
+        and updating proto_info information.
+
+        The input line will be in this format:
+            >>> [(name, [feasign, ...]), ...]
+            >>> or ((name, [feasign, ...]), ...)
+        The output will be in this format:
+            >>> [ids_num id1 id2 ...] ...
+        The proto_info will be in this format:
+            >>> [(name, type), ...]
+
+        For example, if the input is like this:
+            >>> [("words", [1926, 08, 17]), ("label", [1])]
+            >>> or (("words", [1926, 08, 17]), ("label", [1]))
+        the output will be:
+            >>> 3 1234 2345 3456 1 1
+        the proto_info will be:
+            >>> [("words", "uint64"), ("label", "uint64")]
+
+        Args:
+            line(str): the output of the process() function rewritten by user.
+
+        Returns:
+            Return a string data that can be read directly by the MultiSlotDataFeed.
+        '''
+        if isinstance(line, zip):
+            line = list(line)
+
+        if not isinstance(line, list) and not isinstance(line, tuple):
+            raise ValueError(
+                "the output of process() must be in list or tuple type"
+                "Example: [('words', [1926, 08, 17]), ('label', [1])]"
+            )
+        output = ""
+
+        if self._proto_info is None:
+            self._proto_info = []
+            for item in line:
+                name, elements = item
+                if not isinstance(name, str):
+                    raise ValueError(f"name{type(name)} must be in str type")
+                if not isinstance(elements, list):
+                    raise ValueError(
+                        f"elements{type(elements)} must be in list type"
+                    )
+                if not elements:
+                    raise ValueError(
+                        "the elements of each field can not be empty, you need padding it in process()."
+                    )
+                self._proto_info.append((name, "uint64"))
+                if output:
+                    output += " "
+                output += str(len(elements))
+                for elem in elements:
+                    if isinstance(elem, float):
+                        self._proto_info[-1] = (name, "float")
+                    elif not isinstance(elem, int):
+                        raise ValueError(
+                            f"the type of element{type(elem)} must be in int or float"
+                        )
+                    output += " " + str(elem)
+        else:
+            if len(line) != len(self._proto_info):
+                raise ValueError(
+                    "the complete field set of two given line are inconsistent."
+                )
+            for index, item in enumerate(line):
+                name, elements = item
+                if not isinstance(name, str):
+                    raise ValueError(f"name{type(name)} must be in str type")
+                if not isinstance(elements, list):
+                    raise ValueError(
+                        f"elements{type(elements)} must be in list type"
+                    )
+                if not elements:
+                    raise ValueError(
+                        "the elements of each field can not be empty, you need padding it in process()."
+                    )
+                if name != self._proto_info[index][0]:
+                    raise ValueError(
+                        f"the field name of two given line are not match: require<{self._proto_info[index][0]}>, get<{name}>."
+                    )
+                if output:
+                    output += " "
+                output += str(len(elements))
+                for elem in elements:
+                    if self._proto_info[index][1] != "float":
+                        if isinstance(elem, float):
+                            self._proto_info[index] = (name, "float")
+                        elif not isinstance(elem, int):
+                            raise ValueError(
+                                f"the type of element{type(elem)} must be in int or float"
+                            )
+                    output += " " + str(elem)
+        return output + "\n"

python/paddle/distributed/fleet/dataset/__init__.py

@@ -0,0 +1,22 @@
+#   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
+
+from .dataset import (  # noqa: F401
+    DatasetBase,
+    FileInstantDataset,
+    InMemoryDataset,
+    QueueDataset,
+)
+from .index_dataset import TreeIndex  # noqa: F401
+
+__all__ = []

python/paddle/distributed/fleet/dataset/index_dataset.py

@@ -0,0 +1,104 @@
+# 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 typing import Any
+
+from paddle.base import core
+
+__all__ = []
+
+
+class Index:
+    def __init__(self, name: str) -> None:
+        self._name = name
+
+
+class TreeIndex(Index):
+    def __init__(self, name: str, path: str) -> None:
+        super().__init__(name)
+        self._wrapper = core.IndexWrapper()
+        self._wrapper.insert_tree_index(name, path)
+        self._tree = self._wrapper.get_tree_index(name)
+        self._height = self._tree.height()
+        self._branch = self._tree.branch()
+        self._total_node_nums = self._tree.total_node_nums()
+        self._emb_size = self._tree.emb_size()
+        self._layerwise_sampler = None
+
+    def height(self) -> int:
+        return self._height
+
+    def branch(self) -> int:
+        return self._branch
+
+    def total_node_nums(self) -> int:
+        return self._total_node_nums
+
+    def emb_size(self) -> int:
+        return self._emb_size
+
+    def get_all_leaves(self) -> list[Any]:
+        return self._tree.get_all_leaves()
+
+    def get_nodes(self, codes: list[int]) -> list[Any]:
+        return self._tree.get_nodes(codes)
+
+    def get_layer_codes(self, level: int) -> list[int]:
+        return self._tree.get_layer_codes(level)
+
+    def get_travel_codes(self, id: int, start_level: int = 0) -> list[int]:
+        return self._tree.get_travel_codes(id, start_level)
+
+    def get_ancestor_codes(self, ids: list[int], level: int) -> list[int]:
+        return self._tree.get_ancestor_codes(ids, level)
+
+    def get_children_codes(self, ancestor: int, level: int) -> list[int]:
+        return self._tree.get_children_codes(ancestor, level)
+
+    def get_travel_path(self, child: int, ancestor: int) -> list[int]:
+        res = []
+        while child > ancestor:
+            res.append(child)
+            child = int((child - 1) / self._branch)
+        return res
+
+    def get_pi_relation(self, ids: list[int], level: int) -> dict[int, int]:
+        codes = self.get_ancestor_codes(ids, level)
+        return dict(zip(ids, codes))
+
+    def init_layerwise_sampler(
+        self,
+        layer_sample_counts: list[int],
+        start_sample_layer: int = 1,
+        seed: int = 0,
+    ) -> None:
+        assert self._layerwise_sampler is None
+        self._layerwise_sampler = core.IndexSampler("by_layerwise", self._name)
+        self._layerwise_sampler.init_layerwise_conf(
+            layer_sample_counts, start_sample_layer, seed
+        )
+
+    def layerwise_sample(
+        self,
+        user_input: list[list[int]],
+        index_input: list[int],
+        with_hierarchy: bool = False,
+    ) -> list[list[int]]:
+        if self._layerwise_sampler is None:
+            raise ValueError("please init layerwise_sampler first.")
+        return self._layerwise_sampler.sample(
+            user_input, index_input, with_hierarchy
+        )

python/paddle/distributed/fleet/elastic/__init__.py

@@ -0,0 +1,89 @@
+# 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
+import signal
+import sys
+
+from paddle.distributed.fleet.launch_utils import DistributeMode  # noqa: F401
+
+from .collective import CollectiveLauncher
+from .manager import (
+    ELASTIC_EXIT_CODE,
+    ElasticLevel,  # noqa: F401
+    ElasticManager,
+    ElasticStatus,
+)
+
+
+def enable_elastic(args, distribute_mode):
+    # elastic_level = os.getenv('PADDLE_ELASTIC_FAULT_TOLERANC_LEVEL')
+    # if not elastic_level and (elastic_level != ElasticLevel.FAULT_TOLERANCE and
+    #                          elastic_level != ElasticLevel.ELASTIC):
+    #    return False
+
+    # if distribute_mode != DistributeMode.COLLECTIVE:
+    #    return False
+
+    if not args.elastic_server and not os.getenv('PADDLE_ELASTIC_SERVER'):
+        return False
+
+    if not args.job_id and not os.getenv('PADDLE_ELASTIC_JOB_ID'):
+        return False
+
+    if not args.np and not os.getenv('PADDLE_ELASTIC_NP'):
+        return False
+
+    return True
+
+
+def launch_elastic(args, distribute_mode):
+    server = args.elastic_server or os.getenv('PADDLE_ELASTIC_SERVER')
+    srv, port = server.split(':')
+    import etcd3
+
+    etcd_client = etcd3.client(host=srv, port=port)
+    elastic = ElasticManager(args, etcd_client)
+
+    signal.signal(signal.SIGTERM, elastic.signal_handler)
+    signal.signal(signal.SIGABRT, elastic.signal_handler)
+    signal.signal(signal.SIGINT, elastic.signal_handler)
+
+    while True:
+        # wait for all nodes ready to run
+        elastic.wait()
+
+        # execute pre hook action, eg: run shell
+        elastic.pre_hook()
+
+        # run self with specified launcher
+        elastic.run(CollectiveLauncher)
+
+        # keep wathing the health status of self and being notified for other's failure
+        ret = elastic.watch()
+        if ret == ElasticStatus.COMPLETED:
+            break
+        if ret == ElasticStatus.HOLD:
+            continue
+        if ret == ElasticStatus.EXIT:
+            break
+        if ret == ElasticStatus.ERROR:
+            sys.exit(3)
+        if ret == ElasticStatus.RESTART:
+            sys.exit(ELASTIC_EXIT_CODE)
+
+    if int(elastic.sigint) > 0:
+        sys.exit(128 + int(elastic.sigint))
+    else:
+        sys.exit(0)

python/paddle/distributed/fleet/elastic/collective.py

@@ -0,0 +1,67 @@
+# 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
+import shutil
+import tempfile
+
+import paddle
+from paddle.distributed.fleet.elastic.manager import LauncherInterface
+from paddle.distributed.fleet.launch_utils import (
+    logger,
+    pull_worker_log,
+    start_local_trainers,
+)
+
+
+class CollectiveLauncher(LauncherInterface):
+    def __init__(self, args):
+        self.args = args
+        self.procs = []
+
+    def launch(self):
+        logger.info("collective launcher launch ...")
+        args = self.args
+        self.tmp_dir = tempfile.mkdtemp()
+        cluster, pod = paddle.distributed.fleet.launch.get_cluster_info(args)
+        global_envs = paddle.distributed.fleet.launch.get_global_envs(
+            args, self.tmp_dir
+        )
+
+        self.procs = start_local_trainers(
+            cluster,
+            pod,
+            training_script=args.training_script,
+            training_script_args=args.training_script_args,
+            log_dir=args.log_dir,
+            envs=global_envs,
+        )
+
+        for idx, proc in enumerate(self.procs):
+            logger.info(f"launch proc_id:{proc.proc.pid} idx:{idx}")
+
+    def stop(self):
+        logger.info("collective launcher stop ...")
+        if not self._terminate_procs():
+            logger.error("kill process failed")
+        if os.path.exists(self.tmp_dir):
+            shutil.rmtree(self.tmp_dir)
+
+    def watch(self):
+        logger.debug("collective launcher watch ...")
+        for p in self.procs:
+            if p.log_fn and p.local_rank == 0:
+                pull_worker_log(p)
+        ret = self._check_procs()
+        return ret

python/paddle/distributed/fleet/elastic/manager.py

@@ -0,0 +1,635 @@
+# 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 os
+import random
+import signal
+import socket
+import subprocess
+import threading
+import time
+import traceback
+
+from paddle.distributed.fleet import cloud_utils, launch_utils
+from paddle.distributed.utils.log_utils import get_logger
+
+from ...backup_env import getenv_or_backup
+
+logger = get_logger("INFO", "ELASTIC")
+
+ELASTIC_EXIT_CODE = 101
+ELASTIC_AUTO_PARALLEL_EXIT_CODE = 102
+
+# wait for timeout, unit: seconds
+ELASTIC_TIMEOUT = 2 * 60
+
+# keepalived ttl, unit: seconds
+ELASTIC_TTL = 60
+
+
+# 1: Fault tolerance, 2: Elastic
+class ElasticLevel:
+    FAULT_TOLERANCE = 1
+    ELASTIC = 2
+
+
+class ElasticStatus:
+    COMPLETED = "completed"
+    ERROR = "error"
+    HOLD = "hold"
+    RESTART = "restart"
+    EXIT = "exit"
+
+
+class LauncherInterface:
+    def __init__(self, args):
+        self.args = args
+        self.procs = []
+
+    def _terminate_procs(self):
+        # try to terminate process by group, this happened in multiprocess scenario in user process
+        if os.name != 'nt':
+            for p in self.procs:
+                if p.proc.poll() is None:
+                    os.killpg(os.getpgid(p.proc.pid), signal.SIGTERM)
+                    if p.log_fn:
+                        p.log_fn.close()
+                    logger.info(f"terminate process group gid:{p.proc.pid}")
+
+            time.sleep(1)
+        for p in self.procs:
+            if p.proc.poll() is None:
+                p.proc.terminate()
+                if p.log_fn:
+                    p.log_fn.close()
+                logger.info(f"terminate process id:{p.proc.pid}")
+
+        for step in range(0, 50):
+            alive = False
+            for p in self.procs:
+                if p.proc.poll() is None:  # not terminate
+                    os.kill(p.proc.pid, signal.SIGKILL)
+                    alive = True
+
+            if not alive:
+                logger.info("terminated all the procs")
+                return True
+
+            time.sleep(1)
+        return False
+
+    def _check_procs(self):
+        alive = False
+        result = None
+        for p in self.procs:
+            ret = p.proc.poll()
+            if ret is None:
+                alive = True
+            elif ret != 0:
+                if ret == ELASTIC_AUTO_PARALLEL_EXIT_CODE:
+                    logger.info("return form elastic auto parallel re-launch")
+                    return ret
+                logger.error("ABORT!!! ABORT!!! ABORT!!!")
+                logger.error(
+                    f"ERROR rank {p.rank} error with exit code {ret}, check log for detail."
+                )
+                result = ret
+        if not alive and result is None:
+            return 0
+        else:
+            return result
+
+    def launch(self):
+        raise NotImplementedError
+
+    def stop(self):
+        raise NotImplementedError
+
+    def watch(self):
+        raise NotImplementedError
+
+
+class ElasticManager:
+    def __init__(self, args, etcd_client):
+        self.args = args
+        server = args.elastic_server or os.getenv('PADDLE_ELASTIC_SERVER')
+        name = args.job_id or os.getenv('PADDLE_ELASTIC_JOB_ID')
+        self.min_np, self.max_np = self._parse_np(args.np)
+        host = args.host or os.getenv('POD_IP')
+        scale = args.scale or int(os.getenv('PADDLE_ELASTIC_SCALE', 0))
+        force = args.force or os.getenv('PADDLE_ELASTIC_FORCE')
+
+        self.host = host if host else self._get_host()
+
+        (
+            self.device_mode,
+            self.devices_per_proc,
+        ) = launch_utils.get_device_proc_info(args)
+
+        self.elastic_timeout = int(
+            os.getenv('PADDLE_ELASTIC_TIMEOUT', ELASTIC_TIMEOUT)
+        )
+        elastic_ttl = int(os.getenv('PADDLE_ELASTIC_TTL', ELASTIC_TTL))
+
+        self.start_port = None
+        if cloud_utils.use_paddlecloud():
+            self.trainers = os.getenv('PADDLE_TRAINERS', '')
+            self.np = len(self.trainers.split(","))
+            self.start_port = int(os.getenv("PADDLE_PORT", "6170"))
+            self.dist_endpoints = os.getenv('DISTRIBUTED_TRAINER_ENDPOINTS', '')
+            trainer_endpoints = getenv_or_backup('PADDLE_TRAINER_ENDPOINTS', '')
+            self.trainer_endpoints_list = trainer_endpoints.split(",")
+        else:
+            self.trainers = args.ips or os.getenv('PADDLE_TRAINERS', '')
+            node_ips = self.trainers.split(",")
+            self.np = len(node_ips)
+            self.start_port = int(os.getenv("FLAGS_START_PORT", "6170"))
+            self.dist_endpoints = self._host_to_endpoints(
+                node_ips, self.devices_per_proc, self.start_port
+            )
+            self.trainer_endpoints_list = [
+                f"{ip}:{self.start_port}" for ip in node_ips
+            ]
+
+        self.curr_host = f"{self.host}:{self.start_port}"
+        logger.info(f'start job with np={self.np}')
+        logger.info(
+            f"trainers={self.trainers}, trainer_endpoints_list={self.trainer_endpoints_list}"
+        )
+
+        # auto correct the value of elastic_level
+        # 1: Fault tolerant, 2: Elastic
+        self.elastic_level = int(
+            os.getenv(
+                'PADDLE_ELASTIC_FAULT_TOLERANC_LEVEL',
+                ElasticLevel.FAULT_TOLERANCE,
+            )
+        )
+        if self.min_np == self.max_np or (self.min_np > 0 and self.max_np == 0):
+            self.elastic_level = ElasticLevel.FAULT_TOLERANCE
+            logger.info('start job with ElasticLevel.FAULT_TOLERANCE')
+        if self.min_np > 0 and self.max_np > self.min_np:
+            self.elastic_level = ElasticLevel.ELASTIC
+            logger.info('start job with ElasticLevel.ELASTIC')
+
+        # compatible with kubernetes service discovery
+        if (
+            not server
+            and os.getenv('PADDLE_ELASTIC_ETCD_SERVICE_HOST')
+            and os.getenv('PADDLE_ELASTIC_ETCD_SERVICE_PORT')
+        ):
+            server = '{}:{}'.format(
+                os.getenv('PADDLE_ELASTIC_ETCD_SERVICE_HOST'),
+                os.getenv('PADDLE_ELASTIC_ETCD_SERVICE_PORT'),
+            )
+
+        logger.debug(f'init with server {server} host {host}')
+
+        self.hosts = []
+        self.stopped = False
+
+        self.sigint = 0
+        self.need_sync = False
+
+        self.elastic_startup_time = None
+
+        if not server or ':' not in server or not name or not self.np:
+            logger.info(
+                f'Elastic is not enabled with server {server} name {name} and np {self.np}'
+            )
+            self.enable = False
+            return
+        else:
+            self.enable = True
+
+        self.etcd = etcd_client
+
+        # etcd data
+        self.prefix = "/paddle/" + name
+        self.node_prefix = self.prefix + '/nodes'
+        self.np_path = self.prefix + '/np'
+        self.endpoints_path = self.prefix + '/endpoints'
+
+        node_tag = ''.join(
+            random.choice('abcdefghijklmnopqrstuvwxyz') for _ in range(6)
+        )
+        self.host_path = f'{self.node_prefix}/{node_tag}{time.time()}'
+        '''
+        0 group mode, be aware of healthy status of other workers
+        1 decouple mode, check own status only
+        '''
+        self.etcd.put(self.prefix, b'0')
+
+        # register callback
+        def host_call_back(event):
+            self.hosts = [
+                i[0].decode() for i in self.etcd.get_prefix(self.node_prefix)
+            ]
+            self.hosts = list(set(self.hosts)) if self.hosts else self.hosts
+            logger.info(
+                f"host_call_back curr_host={self.curr_host}, hosts:{self.hosts}"
+            )
+            self.need_sync = True
+            self.elastic_startup_time = None
+
+        host_watch = self.etcd.add_watch_prefix_callback(
+            self.node_prefix, host_call_back
+        )
+        host_lease = self.etcd.lease(elastic_ttl)
+
+        # register etcd lease heartbeat
+        def lease_heartbeat():
+            while True:
+                try:
+                    host_lease.refresh()
+
+                    hosts = [
+                        i[0].decode()
+                        for i in self.etcd.get_prefix(self.node_prefix)
+                    ]
+                    hosts = list(set(hosts)) if hosts else hosts
+                    logger.info(
+                        f"[lease_heartbeat] curr_host={self.curr_host}, hosts={hosts}"
+                    )
+                    if self.curr_host not in hosts:
+                        logger.info(
+                            f"[lease_heartbeat] register host={self.curr_host}"
+                        )
+                        self.etcd.put(
+                            self.host_path,
+                            self.curr_host.encode('latin-1'),
+                            lease=host_lease,
+                        )
+                except Exception as e:
+                    logger.error(
+                        f"[lease_heartbeat] internal error:{e} {traceback.format_exc()}"
+                    )
+                    break
+                time.sleep(elastic_ttl / 3)
+
+        keepalived_thread = threading.Thread(
+            name='lease_heartbeat', target=lease_heartbeat, daemon=True
+        )
+        keepalived_thread.start()
+
+        self.etcd.put(
+            self.host_path, self.curr_host.encode('latin-1'), lease=host_lease
+        )
+
+        # endpoints handle DISTRIBUTED_TRAINER_ENDPOINTS and PADDLE_TRAINERS
+        self.etcd.put(
+            self.endpoints_path,
+            f'{self.dist_endpoints}|{self.trainers}'.encode('latin-1'),
+        )
+
+        def endpoints_call_back(event):
+            if not self.dist_endpoints:
+                return
+            value = self.etcd.get(self.endpoints_path)[0]
+            edps = value.decode() if value is not None else ''
+            self.dist_endpoints, self.trainers = edps.split('|')
+            logger.info(
+                f"set DISTRIBUTED_TRAINER_ENDPOINTS {self.dist_endpoints} "
+            )
+            logger.info(f"set PADDLE_TRAINERS {self.trainers} ")
+
+        endpoints_watch = self.etcd.add_watch_callback(
+            self.endpoints_path, endpoints_call_back
+        )
+
+        self.watches = [host_watch, endpoints_watch]
+        self.launcher = None
+
+    def _host_to_endpoints(
+        self, ip_port_list: list, devices_per_proc: list, start_port: int = 6170
+    ) -> str:
+        endpoint_list = []
+        for ip_port in ip_port_list:
+            endpoints = ip_port.split(":")
+            if len(endpoints) == 2:
+                ip = endpoints[0]
+                port = int(endpoints[1])
+            else:
+                ip = endpoints
+                port = start_port
+
+            ports = list(range(port, port + len(devices_per_proc)))
+            endpoint_list.extend([f"{ip}:{port}" for port in ports])
+
+        dist_endpoints = ','.join(endpoint_list)
+        return dist_endpoints
+
+    def exit(self, completed=False):
+        logger.info(f'manager exist completed {completed}')
+
+        if self.launcher:
+            self.launcher.stop()
+
+        if not self.enable:
+            return
+
+        if completed:
+            self.etcd.put(self.prefix, b'1')
+
+        for watch in self.watches:
+            self.etcd.cancel_watch(watch)
+        self.etcd.delete(self.host_path)
+
+        hosts = list(self.etcd.get_prefix(self.node_prefix))
+        if len(hosts) == 0:
+            self.etcd.delete_prefix(self.prefix)
+
+    def pre_hook(self):
+        if not self.args.elastic_pre_hook:
+            logger.info("skip pre_hook")
+            return
+        logger.info("execute pre_hook...")
+        current_env = copy.copy(os.environ.copy())
+        out, err = subprocess.Popen(
+            self.args.elastic_pre_hook,
+            env=current_env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            shell=True,
+        ).communicate()
+        if err:
+            logger.warning("pre_hook exec failed")
+        else:
+            logger.info(f"pre_hook exec result: {out.decode('utf-8').strip()}")
+
+    def _parse_np(self, np: str):
+        """
+        np format is "MIN" or "MIN:MAX"
+        """
+        np_str = np or os.getenv('PADDLE_ELASTIC_NP', "0")
+        np_dict = np_str.split(":")
+        min_np = max_np = 0
+        if len(np_dict) == 1:
+            # Fault tolerant
+            min_np = int(np_dict[0])
+            min_np = 1 if min_np <= 0 else min_np
+            max_np = 1
+        elif len(np_dict) == 2:
+            # Elastic
+            min_np = int(np_dict[0])
+            max_np = int(np_dict[1])
+            min_np = 1 if min_np <= 0 else min_np
+            max_np = max(max_np, min_np)
+        else:
+            raise ValueError(
+                f'the np={np} needs to be in "MIN" or "MIN:MAX" format'
+            )
+
+        return min_np, max_np
+
+    def _get_host(self):
+        try:
+            return socket.gethostbyname(socket.getfqdn(socket.gethostname()))
+        except:
+            return '127.0.0.1'
+
+    def _completed(self):
+        if not self.enable:
+            return True
+
+        return int(self.etcd.get(self.prefix)[0]) == 1
+
+    def _match(self, host_list: list | None = None):
+        if host_list:
+            self.hosts = host_list
+        else:
+            self.hosts = [
+                i[0].decode() for i in self.etcd.get_prefix(self.node_prefix)
+            ]
+        self.hosts = list(set(self.hosts)) if self.hosts else self.hosts
+
+        if self.elastic_level == ElasticLevel.FAULT_TOLERANCE:
+            if len(self.hosts) == self.np:
+                return True
+            else:
+                return False
+
+        if self.elastic_level == ElasticLevel.ELASTIC:
+            hosts_num = len(self.hosts)
+            if hosts_num == self.np:
+                return True
+
+            if not self.elastic_startup_time:
+                self.elastic_startup_time = time.time()
+            if hosts_num == self.max_np:
+                self.elastic_startup_time = None
+                return True
+            elif hosts_num >= self.min_np and hosts_num < self.max_np:
+                interval_time = time.time() - self.elastic_startup_time
+                if interval_time <= self.elastic_timeout:
+                    logger.info(
+                        f"wait for timeout, you can set value by PADDLE_ELASTIC_TIMEOUT, \
+                        hosts_num={hosts_num}, min_np={self.min_np}, \
+                        interval_time={interval_time}, elastic_timeout={self.elastic_timeout}"
+                    )
+                    return False
+                return True
+            else:
+                self.elastic_startup_time = None
+                return False
+
+        return False
+
+    def _update_endpoint(self, endpoints, hosts):
+        self.etcd.put(
+            self.endpoints_path,
+            f'{endpoints}|{hosts}'.encode('latin-1'),
+        )
+
+    def _update_fault_tolerance(self):
+        rank = int(os.getenv('PADDLE_TRAINER_ID', -1))
+        logger.debug(
+            f"self.curr_host={self.curr_host}, self.dist_endpoints={self.dist_endpoints}"
+        )
+        if self.curr_host in self.dist_endpoints:
+            os.environ['DISTRIBUTED_TRAINER_ENDPOINTS'] = self.dist_endpoints
+            os.environ['PADDLE_TRAINERS'] = self.trainers
+            logger.info(
+                f"update env DISTRIBUTED_TRAINER_ENDPOINTS {self.dist_endpoints} "
+            )
+            logger.info(f"update env PADDLE_TRAINERS {self.trainers} ")
+            return
+
+        # fault tolerance
+        idx = self.hosts.index(self.curr_host)
+
+        # swap if self.host not in the right position
+        if rank >= 0:
+            self.hosts[idx] = self.hosts[rank]
+            self.hosts[rank] = self.curr_host
+        else:
+            os.environ['PADDLE_TRAINER_ID'] = f'{idx}'
+        hosts = ','.join([host_port.split(":")[0] for host_port in self.hosts])
+        self.args.ips = hosts
+        os.environ['PADDLE_TRAINERS'] = hosts
+
+    def _update_elastic_scale_out(self):
+        host_endpoints = copy.deepcopy(self.trainer_endpoints_list)
+        logger.info(
+            f"elastic scale out, from {len(self.hosts)} to {self.np}, hosts={self.hosts}, host_endpoints={host_endpoints}"
+        )
+
+        for curr_host_port in self.hosts:
+            if curr_host_port not in host_endpoints:
+                host_endpoints.append(curr_host_port)
+
+        os.environ['PADDLE_TRAINER_ID'] = str(
+            host_endpoints.index(self.curr_host)
+        )
+        hosts = ','.join(
+            [host_port.split(":")[0] for host_port in host_endpoints]
+        )
+        self.args.ips = hosts
+        os.environ['PADDLE_TRAINERS'] = hosts
+        self.np = len(host_endpoints)
+        os.environ['PADDLE_TRAINER_ENDPOINTS'] = ','.join(host_endpoints)
+        os.environ['DISTRIBUTED_TRAINER_ENDPOINTS'] = self.dist_endpoints
+        self.trainer_endpoints_list = host_endpoints
+
+    def _update_elastic_scale_in(self):
+        host_endpoints = copy.deepcopy(self.trainer_endpoints_list)
+        logger.info(
+            f"elastic scale in, from {self.np} to {len(self.hosts)}, hosts={self.hosts}, host_endpoints={host_endpoints}"
+        )
+
+        # If scale in node from the first of the rank list, you need to minimize the movement of the rank
+        # eg:
+        #   the source trainers is:10.10.10.0,10.10.10.1,10.10.10.2,10.10.10.3
+        #   10.10.10.0 is removed
+        #   the new trainers is:10.10.10.3,10.10.10.1,10.10.10.2
+        #   In this case, the rank of 10.10.10.1 and 10.10.10.2 remains unchanged, while the rank of 10.10.10.3 is set to rank0
+        endpoints_dict = {}
+        unsorted_endpoints = []
+        for id, host_port in enumerate(self.hosts):
+            idx = host_endpoints.index(host_port)
+            if idx <= len(self.hosts) - 1 and not endpoints_dict.get(idx):
+                endpoints_dict[idx] = host_port
+            else:
+                unsorted_endpoints.append(host_port)
+
+        idle_index = 0
+        sorted_endpoints = []
+        for idx in range(len(self.hosts)):
+            if not endpoints_dict.get(idx) and len(unsorted_endpoints) > 0:
+                endpoints_dict[idx] = unsorted_endpoints[idle_index]
+                idle_index += 1
+
+            sorted_endpoints.append(endpoints_dict.get(idx))
+
+        logger.info(f"elastic scale in, sorted_endpoints={sorted_endpoints}")
+        self.trainer_endpoints_list = sorted_endpoints
+
+        ip_list = [ip_port.split(":")[0] for ip_port in sorted_endpoints]
+        hosts = ','.join(ip_list)
+        new_endpoints = self._host_to_endpoints(
+            sorted_endpoints, self.devices_per_proc
+        )
+
+        self.args.ips = hosts
+        os.environ['PADDLE_TRAINER_ID'] = str(
+            sorted_endpoints.index(self.curr_host)
+        )
+        os.environ['PADDLE_TRAINERS'] = hosts
+        self.np = len(sorted_endpoints)
+        os.environ['PADDLE_TRAINER_ENDPOINTS'] = ','.join(sorted_endpoints)
+        os.environ['DISTRIBUTED_TRAINER_ENDPOINTS'] = new_endpoints
+        self._update_endpoint(new_endpoints, hosts)
+
+    def _update_hosts(self):
+        assert len(self.hosts) != 0, 'hosts empty'
+        if self.elastic_level == ElasticLevel.FAULT_TOLERANCE:
+            self._update_fault_tolerance()
+        else:
+            # elastic
+            if len(self.hosts) == self.np:
+                logger.info(f"elastic startup, hosts={self.hosts}")
+                self._update_fault_tolerance()
+
+            elif len(self.hosts) > self.np:
+                # scale out
+                self._update_elastic_scale_out()
+            else:
+                # scale in
+                self._update_elastic_scale_in()
+
+    def wait(self):
+        if not self.enable:
+            return
+
+        idx = 1
+        while not self.stopped:
+            if self._match():
+                logger.info(f'ready with hosts {self.hosts}')
+                self._update_hosts()
+                return
+            logger.info(f'not ready for np {self.np} with hosts {self.hosts}')
+            idx += 1
+            time.sleep(2)
+        return
+
+    def run(self, launcher):
+        if self.stopped:
+            return
+
+        self.launcher = launcher(self.args)
+        self.launcher.launch()
+
+    def watch(self):
+        if self.need_sync:
+            self.need_sync = False
+
+        while not self.stopped:
+            ret = self.launcher.watch()
+            logger.debug(f"launcher.watch():{ret}")
+
+            if ret is not None:  # self terminated
+                logger.info(f'job exit with code {ret}')
+                if ret == ELASTIC_AUTO_PARALLEL_EXIT_CODE:
+                    logger.info('job re-launch for auto parallel')
+                    self.launcher.stop()
+                    return ElasticStatus.HOLD
+
+                # process is completed if ret >= 0 or error else
+                completed = True if ret == 0 else False
+                self.exit(completed=completed)
+                if completed:
+                    return ElasticStatus.COMPLETED
+                if self.elastic_level == ElasticLevel.FAULT_TOLERANCE:
+                    return ElasticStatus.RESTART
+                else:
+                    return ElasticStatus.ERROR
+
+            if not self._completed() and (not self._match() or self.need_sync):
+                self.launcher.stop()
+                return ElasticStatus.HOLD
+
+            time.sleep(2)
+
+        if self.launcher:
+            self.launcher.stop()
+
+        return ElasticStatus.EXIT
+
+    def signal_handler(self, sigint, frame):
+        if self.enable:
+            self.exit()
+        self.sigint = sigint
+        self.stopped = True

python/paddle/distributed/fleet/launch.py

@@ -0,0 +1,771 @@
+# 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.
+r"""
+fleetrun is a module that spawns multiple distributed
+process on each training node for gpu training and cpu training.
+Usage:
+    In both of single node training or multiple node training, this module
+launch a process on each of the given gpu card or cpu machine.
+    GPU training:
+    1. for single node training with all visible gpu cards:
+       fleetrun your_training_py (arg1 arg2 and all others)
+    2. for single node training with [0,4) cards
+       fleetrun --gpus="0,1,2,3" your_training_py (arg1 arg2 and all others)
+    3. for multiple node training such as two node:192.168.0.16, 192.168.0.17
+        on 192.168.0.16:
+            fleetrun --ips="192.168.0.16,192.168.0.17" \
+                your_training_py (arg1 arg2 and all others)
+        on 192.168.0.17:
+            fleetrun --ips="192.168.0.16,192.168.0.17" \
+                your_training_py (arg1 arg2 and all others)
+    CPU training:
+    1. for single node training with multi servers and workers:
+        fleetrun --server_num=2 --worker_num=2 your_training_py (arg1 arg2 and all others)
+    2. for multiple node training such as two node:192.168.0.16, 192.168.0.17 \
+        with 2 servers and 4 workers.
+        on 192.168.0.16:
+            fleetrun --servers="192.168.0.16:6170,192.168.0.17:6170" \
+                --workers="192.168.0.16,192.168.0.17,192.168.0.16,192.168.0.17" \
+                your_training_py (arg1 arg2 and all others)
+        on 192.168.0.17:
+            fleetrun --servers="192.168.0.16:6170,192.168.0.17:6171" \
+                --workers="192.168.0.16,192.168.0.17,192.168.0.16,192.168.0.17" \
+                your_training_py (arg1 arg2 and all others)
+    3. use gloo backend for multiple node training such as two node:192.168.0.16, 192.168.0.17 \
+        with 2 servers and 4 workers. (workers should set port)
+        on 192.168.0.16:
+            fleetrun --servers="192.168.0.16:6170,192.168.0.17:6170" \
+                --workers="192.168.0.16:6171,192.168.0.17:6171,192.168.0.16:6172,192.168.0.17:6172" \
+                your_training_py (arg1 arg2 and all others)
+        on 192.168.0.17:
+            fleetrun --servers="192.168.0.16:6170,192.168.0.17:6170" \
+                --workers="192.168.0.16:6171,192.168.0.17:6171,192.168.0.16:6172,192.168.0.17:6172" \
+                your_training_py (arg1 arg2 and all others)
+"""
+
+import copy
+import os
+import pathlib
+import shutil
+import sys
+import tempfile
+import time
+from argparse import REMAINDER, ArgumentParser
+
+from paddle import framework
+from paddle.distributed.fleet import cloud_utils, launch_utils
+from paddle.distributed.fleet.elastic import enable_elastic, launch_elastic
+from paddle.distributed.fleet.launch_utils import (
+    DeviceMode,
+    DistributeMode,
+    ParameterServerLauncher,
+    block_windows_and_macos,
+    check_backend,
+    direct_start,
+    find_free_ports,
+    get_cluster,
+    get_host_name_ip,
+    get_logger,
+    logger,
+    start_local_trainers,
+    terminate_local_procs,
+    watch_local_trainers,
+)
+
+__all__ = []
+
+
+def _print_arguments(args):
+    print("-----------  Configuration Arguments -----------")
+    for arg, value in sorted(vars(args).items()):
+        print(f"{arg}: {value}")
+    print("------------------------------------------------")
+
+
+def _parse_args():
+    """
+    Helper function parsing the command line options
+    @retval ArgumentParser
+    """
+    parser = ArgumentParser(
+        description='''start paddle training using multi-process mode.
+see: http://www.paddlepaddle.org/documentation/docs/zh/1.6/user_guides/howto/training/cluster_howto.html#permalink-8--nccl2-
+'''
+    )
+    base_group = parser.add_argument_group("Base Parameters")
+
+    base_group.add_argument(
+        "--log_dir",
+        type=str,
+        default="log",
+        help="The path for each process's log. Default --log_dir=log/",
+    )
+    base_group.add_argument(
+        "--backend",
+        type=str,
+        default=os.environ.get('PADDLE_DISTRI_BACKEND', 'auto'),
+        help="Specify the backend, can be gloo|nccl|bkcl|auto|heter. "
+        "Default value is auto which prefers nccl or bkcl.",
+    )
+    base_group.add_argument(
+        "--nproc_per_node",
+        type=int,
+        default=None,
+        help="The number of processes to launch on a node."
+        "In gpu training, it should be less or equal to the gpus number of you system(or you set by --gpus). And so each process can"
+        " bound to one or average number of gpus.",
+    )
+
+    base_group.add_argument(
+        "--run_mode",
+        type=str,
+        default=None,
+        help="run mode of job, can be:collective/ps/ps-heter",
+    )
+
+    if framework.core.is_compiled_with_cuda():
+        base_group.add_argument(
+            "--gpus",
+            type=str,
+            default=None,
+            help="It's for gpu training."
+            "For example:"
+            '--gpus="0,1,2,3" will launch four training processes each bound to one gpu.',
+        )
+        base_group.add_argument("--selected_gpus", dest="gpus")
+
+    if framework.core.is_compiled_with_xpu():
+        base_group.add_argument(
+            "--xpus",
+            type=str,
+            default=None,
+            help="It's for xpu training. For example: "
+            '--xpus="0,1,2,3" will launch four training processes each bound to one xpu.',
+        )
+        base_group.add_argument("--selected_xpus", dest="xpus")
+
+    base_group.add_argument(
+        "training_script",
+        type=str,
+        help="The full path to the single GPU training "
+        "program/script to be launched in parallel, "
+        "followed by all the arguments for the "
+        "training script",
+    )
+
+    base_group.add_argument('training_script_args', nargs=REMAINDER)
+
+    # Optional arguments for the launch helper
+    # for collective
+    collective_group = parser.add_argument_group("Collective Parameters")
+    collective_group.add_argument(
+        "--ips",
+        type=str,
+        default="127.0.0.1",
+        help="Paddle cluster nodes ips, such as 192.168.0.16,192.168.0.17..",
+    )
+    collective_group.add_argument(
+        "--cluster_topo_path",
+        type=str,
+        default=None,
+        help="A json format file will be stored in this path which is used"
+        "to represent the cluster topology information for auto parallel.",
+    )
+    collective_group.add_argument(
+        "--rank_mapping_path",
+        type=str,
+        default=None,
+        help="A json format file will be stored in this path which is used"
+        "to map processes to machines for auto parallel.",
+    )
+    collective_group.add_argument(
+        "--enable_auto_mapping",
+        type=bool,
+        default=False,
+        help="Set true to enable the lazy launch for auto-parallel scenario.",
+    )
+
+    ps_group = parser.add_argument_group("Parameter-Server Parameters")
+    # for parameter server
+    ps_group.add_argument(
+        "--servers", type=str, default="", help="User defined servers ip:port"
+    )
+    ps_group.add_argument(
+        "--workers", type=str, default="", help="User defined workers ip:port"
+    )
+    ps_group.add_argument(
+        "--coordinators",
+        type=str,
+        default="",
+        help="User defined coordinators ip:port",
+    )
+    ps_group.add_argument(
+        "--heter_workers",
+        type=str,
+        default="",
+        help="User defined heter workers in each stage ip1:port1;ip2:port2",
+    )
+    ps_group.add_argument(
+        "--heter_devices",
+        type=str,
+        default="",
+        help="User defined heter devices in each stage cpu;gpu;cpu",
+    )
+
+    ps_group.add_argument("--worker_num", type=int, help="number of workers")
+    ps_group.add_argument(
+        "--coordinator_num", type=int, help="number of coordinators"
+    )
+    ps_group.add_argument("--server_num", type=int, help="number of servers")
+    ps_group.add_argument(
+        "--heter_worker_num",
+        type=str,
+        help="number of heter_workers in each stage 1;2;3",
+    )
+    ps_group.add_argument("--http_port", type=int, help="Gloo http Port")
+
+    # parameter elastic mode
+    elastic_group = parser.add_argument_group("Elastic Parameters")
+    elastic_group.add_argument(
+        "--elastic_server", type=str, help="etcd server host:port"
+    )
+    elastic_group.add_argument(
+        "--elastic_pre_hook", type=str, help="elastic pre_hook shell cmd"
+    )
+
+    elastic_group.add_argument("--job_id", type=str, help="job unique id")
+    elastic_group.add_argument("--np", type=int, help="job pod/node number")
+    elastic_group.add_argument("--scale", type=int, default=0, help="scale np")
+    elastic_group.add_argument(
+        "--host", type=str, help="bind host, default to POD_IP env"
+    )
+    elastic_group.add_argument(
+        "--force", type=bool, default=False, help="update np force"
+    )
+
+    known_args, _ = parser.parse_known_args()
+    return known_args
+
+
+def get_cluster_from_args(args, device_mode, devices_per_proc):
+    node_ips = [x.strip() for x in args.ips.split(',')]
+    if len(node_ips) == 1:
+        node_ip = node_ips[0]
+    else:
+        if args.host:
+            node_ip = args.host
+        else:
+            _, node_ip = get_host_name_ip()
+
+    assert node_ip in node_ips, (
+        f"Can't find your local ip {{{node_ip}}} in node_ips: {{{node_ips}}}"
+    )
+    node_rank = node_ips.index(node_ip)
+
+    logger.debug(
+        f"parsed from args: node_ips:{node_ips} node_ip:{node_ip} node_rank:{node_rank}"
+    )
+
+    free_ports = None
+    if (
+        not cloud_utils.use_paddlecloud()
+        and len(node_ips) <= 1
+        and os.environ.get('FLAGS_START_PORT') is None
+    ):
+        free_ports = find_free_ports(len(devices_per_proc))
+        if free_ports is not None:
+            free_ports = list(free_ports)
+            logger.info(f"find free ports:{free_ports}")
+    else:
+        start_port = 6070
+        if os.environ.get('FLAGS_START_PORT') is not None:
+            start_port = int(os.environ.get('FLAGS_START_PORT'))
+
+        free_ports = list(range(start_port, start_port + len(devices_per_proc)))
+
+    trainer_endpoints = []
+    for ip in node_ips:
+        trainer_endpoints.append([f"{ip}:{port}" for port in free_ports])
+    return get_cluster(
+        node_ips, node_ip, trainer_endpoints, device_mode, devices_per_proc
+    )
+
+
+def cpuonly_check(args):
+    if args.ips and len(args.ips.split(',')) > 1:
+        raise RuntimeError(
+            f"CPUONLY launch only support single trainer, that is len(ips)=1, but got {args.ips}."
+        )
+    if args.run_mode:
+        assert args.run_mode == 'cpuonly', (
+            "CPUONLY launch only support run mode is CPUONLY"
+        )
+    if args.servers:
+        raise RuntimeError("CPUONLY launch can't have --servers as arguments.")
+    return True
+
+
+def get_cluster_info(args):
+    # parse arguments, used for cloud-single-machine and local
+    if args.backend == 'gloo':
+        cpuonly_check(args)
+    if args.enable_auto_mapping:
+        (device_mode, devices_per_proc) = (DeviceMode.GPU, [])
+    else:
+        (device_mode, devices_per_proc) = launch_utils.get_device_proc_info(
+            args
+        )
+    trainers_num = cloud_utils.get_trainers_num()
+    logger.debug(
+        f"parsed from args trainers_num:{trainers_num} mode:{device_mode} devices:{devices_per_proc}"
+    )
+
+    cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+
+    cluster = None
+    pod = None
+
+    start_port = 6170
+    if os.environ.get('FLAGS_START_PORT') is not None:
+        start_port = os.environ.get('FLAGS_START_PORT')
+    # auto mapping between processes and devices for auto-parallel
+    if args.enable_auto_mapping:
+        assert args.cluster_topo_path is not None, (
+            "The cluster topology must be provided when enabling auto mapping."
+        )
+        rank_mapping_path = args.rank_mapping_path or os.getenv(
+            "PADDLE_RANK_MAPPING_PATH"
+        )
+        if not rank_mapping_path:
+            os.environ["PADDLE_NEED_RANK_MAPPING"] = str(True)
+            os.environ["PADDLE_ENABLE_ELASTIC"] = str(
+                enable_elastic(args, device_mode)
+            )
+            cwd = pathlib.Path().cwd()
+            rank_mapping_path = os.path.join(
+                cwd, "auto_parallel_rank_mapping.json"
+            )
+            os.environ["PADDLE_RANK_MAPPING_PATH"] = str(rank_mapping_path)
+
+            original_args = sys.argv[1:]
+            os.environ["PADDLE_ORIGINAL_CMD_ARGS"] = " ".join(original_args)
+            os.environ["PADDLE_CLUSTER_TOPO_PATH"] = str(args.cluster_topo_path)
+            os.environ["PADDLE_ENABLE_AUTO_MAPPING"] = str(
+                args.enable_auto_mapping
+            )
+            (
+                cluster,
+                pod,
+            ) = launch_utils.get_mapped_cluster_from_args_without_rank_mapping(
+                args, device_mode
+            )
+        else:
+            os.environ["PADDLE_NEED_RANK_MAPPING"] = str(False)
+            os.environ["PADDLE_ENABLE_ELASTIC"] = str(
+                enable_elastic(args, device_mode)
+            )
+
+            os.environ["PADDLE_CLUSTER_TOPO_PATH"] = str(args.cluster_topo_path)
+            os.environ["PADDLE_RANK_MAPPING_PATH"] = str(rank_mapping_path)
+            os.environ["PADDLE_ENABLE_AUTO_MAPPING"] = str(
+                args.enable_auto_mapping
+            )
+            (
+                cluster,
+                pod,
+            ) = launch_utils.get_mapped_cluster_from_args_with_rank_mapping(
+                args, device_mode
+            )
+    elif cloud_utils.use_paddlecloud() and trainers_num != 1:
+        cluster, pod = cloud_utils.get_cloud_cluster(
+            args.ips, device_mode, devices_per_proc, start_port
+        )
+        logger.debug(f"get cluster from cloud:{cluster}")
+    else:
+        # trainers_num = 1 or not use paddlecloud ips="a,b"
+        cluster, pod = get_cluster_from_args(
+            args, device_mode, devices_per_proc
+        )
+        logger.debug(f"get cluster from args:{cluster}")
+    return cluster, pod
+
+
+def get_global_envs(args, tmp_dir):
+    global_envs = copy.copy(os.environ.copy())
+    # add gloo env
+    global_envs["PADDLE_WITH_GLOO"] = str(os.getenv("PADDLE_WITH_GLOO", "0"))
+    global_envs["PADDLE_GLOO_RENDEZVOUS"] = "3"
+    global_envs["PADDLE_GLOO_FS_PATH"] = tmp_dir
+    global_envs["PADDLE_DISTRI_BACKEND"] = args.backend
+    return global_envs
+
+
+def launch_collective(args):
+    tmp_dir = tempfile.mkdtemp()
+    cluster, pod = get_cluster_info(args)
+    global_envs = get_global_envs(args, tmp_dir)
+
+    procs = start_local_trainers(
+        cluster,
+        pod,
+        training_script=args.training_script,
+        training_script_args=args.training_script_args,
+        log_dir=args.log_dir,
+        envs=global_envs,
+    )
+
+    for idx, proc in enumerate(procs):
+        print(f"launch proc_id:{proc.proc.pid} idx:{idx}")
+
+    while True:
+        try:
+            alive = watch_local_trainers(procs, cluster.trainers_nranks())
+
+            if not alive:
+                logger.info("Local processes completed.")
+                logger.debug(f"POD info:{pod}")
+                break
+
+            time.sleep(3)
+
+        except:
+            logger.warning("Terminating... exit")
+            terminate_local_procs(procs)
+            sys.exit(1)
+
+    if os.path.exists(tmp_dir):
+        shutil.rmtree(tmp_dir)
+
+
+def launch_ps(args, distribute_mode):
+    cloud_flag = cloud_utils.use_paddlecloud()
+
+    # for ps-cpu on paddlecloud
+    if cloud_flag and distribute_mode == DistributeMode.PS:
+        direct_start(args)
+        return
+    # elif cloud_flag and distribute_mode == DistributeMode.PS_HETER:
+    #    cloud_ps_heter_env_set(args)
+    #    args.workers = os.getenv("PADDLE_TRAINER_ENDPOINTS")
+    #    args.servers = os.getenv("PADDLE_PSERVERS_IP_PORT_LIST")
+    #    args.heter_workers = os.getenv("PADDLE_HETER_TRAINER_IP_PORT_LIST")
+
+    ps_launcher = ParameterServerLauncher(args, distribute_mode)
+    ps_launcher.start_ps()
+    return
+
+
+def infer_backend(args):
+    if args.backend != "auto":
+        return
+    if framework.core.is_compiled_with_cuda():
+        args.backend = 'nccl'
+    elif framework.core.is_compiled_with_xpu():
+        args.backend = 'bkcl'
+    else:
+        args.backend = 'gloo'
+
+
+def which_distributed_mode(args):
+    infer_backend(args)  # modify the args.backend
+    if args.run_mode is not None:
+        assert args.run_mode in ["collective", "ps", "ps-heter"]
+
+    if args.run_mode == "collective":
+        return DistributeMode.COLLECTIVE
+    elif args.run_mode == "ps":
+        return DistributeMode.PS
+    elif args.run_mode == "ps-heter":
+        return DistributeMode.PS_HETER
+
+    ps_args = [
+        '--worker_num',
+        '--server_num',
+        '--heter_worker_num',
+        '--servers',
+        '--workers',
+        '--heter_workers',
+        '--heter_devices',
+        '--http_port',
+    ]
+    collective_args = ['--ips']
+
+    ps_heter_args = ["--heter_worker_num", "--heter_workers", "--heter_devices"]
+
+    coordinator_args = ["--coordinator_num", "--coordinators"]
+
+    has_ps_args = [
+        ps_arg for ps_arg in ps_args if ps_arg in " ".join(sys.argv[1:-1])
+    ]
+    has_collective_args = [
+        co_arg
+        for co_arg in collective_args
+        if co_arg in " ".join(sys.argv[1:-1])
+    ]
+
+    if len(has_ps_args) > 1 and len(has_collective_args) > 1:
+        raise ValueError(
+            "Only one mode(Collective or Parameter-Server) can be selected at the same time, but more than one configuration was received."
+        )
+
+    if framework.core.is_compiled_with_cuda():
+        accelerators = framework.core.get_cuda_device_count()
+    elif framework.core.is_compiled_with_xpu():
+        accelerators = framework.core.get_xpu_device_count()
+    else:
+        accelerators = 0
+
+    if len(has_ps_args) > 0:
+        logger.info(
+            f"Run parameter-sever mode. pserver arguments:{has_ps_args}, accelerators count:{accelerators}"
+        )
+        has_ps_heter_args = list(set(has_ps_args) & set(ps_heter_args))
+        has_coordinator_args = list(set(has_ps_args) & set(coordinator_args))
+        if len(has_ps_heter_args) > 0:
+            return DistributeMode.PS_HETER
+        else:
+            return DistributeMode.PS
+    elif len(has_collective_args) > 0:
+        logger.info(
+            f"Run collective mode. gpu arguments:{has_collective_args}, cuda count:{accelerators}"
+        )
+        return DistributeMode.COLLECTIVE
+    else:
+        if (
+            not framework.core.is_compiled_with_cuda()
+            and not framework.core.is_compiled_with_xpu()
+        ):
+            if args.servers:
+                logger.warning(
+                    "Not found distinct arguments and not compiled with cuda or xpu. "
+                    "But found args.servers not empty, default use ps mode"
+                )
+                return DistributeMode.PS
+            else:
+                return DistributeMode.COLLECTIVE
+        else:
+            logger.warning(
+                "Not found distinct arguments and compiled with cuda or xpu. "
+                "Default use collective mode"
+            )
+            return DistributeMode.COLLECTIVE
+
+
+def launch():
+    """
+    Paddle distribution training entry ``python -m paddle.distributed.launch``.
+
+    Usage:
+        .. code-block:: bash
+            :name: code-block-bash1
+
+            python -m paddle.distributed.launch [-h] [--log_dir LOG_DIR] [--nproc_per_node NPROC_PER_NODE] [--run_mode RUN_MODE] [--gpus GPUS]
+                             [--selected_gpus GPUS] [--ips IPS] [--servers SERVERS] [--workers WORKERS] [--heter_workers HETER_WORKERS]
+                             [--worker_num WORKER_NUM] [--server_num SERVER_NUM] [--heter_worker_num HETER_WORKER_NUM]
+                             [--http_port HTTP_PORT] [--elastic_server ELASTIC_SERVER] [--job_id JOB_ID] [--np NP] [--scale SCALE]
+                             [--host HOST] [--force FORCE]
+                             training_script ...
+
+
+    Base Parameters:
+        - ``--log_dir``: The path for each process's log. e.g., ``--log_dir=output_dir``. Default ``--log_dir=log``.
+
+        - ``--nproc_per_node``: The number of processes to launch on a node. In gpu training, it should be less or equal to the gpus number of you system(or you set by --gpus).  e.g., ``--nproc_per_node=8``
+
+        - ``--run_mode``: run mode of job, can be:collective/ps/ps-heter. e.g., ``--run_mode=ps``. Default ``--run_mode=collective``.
+
+        - ``--gpus``: It's for gpu training. e.g., ``--gpus=0,1,2,3`` will launch four training processes each bound to one gpu.
+
+        - ``--selected_gpus``: gpus aliases, recommend to use ``--gpus``.
+
+        - ``--xpus``: It's for xpu training if xpu is available. e.g., ``--xpus=0,1,2,3``.
+
+        - ``--selected_xpus``: xpus aliases, recommend to use ``--xpus``.
+
+        - ``training_script``: The full path to the single GPU training program/script to be launched in parallel, followed by all the arguments for the training script. e.g., ``training.py``
+
+        - ``training_script_args``: The args of training_script. e.g., ``--lr=0.1``
+
+    Collective Parameters:
+        - ``--ips``: Paddle cluster nodes ips, e.g., ``--ips=192.168.0.16,192.168.0.17``. Default ``--ips=127.0.0.1``.
+
+    Parameter-Server Parameters:
+        - ``--servers``: User defined servers ip:port, e.g., ``--servers="192.168.0.16:6170,192.168.0.17:6170"``
+
+        - ``--workers``: User defined workers ip:port, e.g., ``--workers="192.168.0.16:6171,192.168.0.16:6172,192.168.0.17:6171,192.168.0.17:6172"``
+
+        - ``--heter_workers``: User defined heter workers ip1:port1;ip2:port2, e.g., ``--heter_workers="192.168.0.16:6172;192.168.0.17:6172"``
+
+        - ``--worker_num``: Number of workers (It recommend to set when in the emulated distributed environment using single node)
+
+        - ``--server_num``: Number of servers (It recommend to set when in the emulated distributed environment using single node)
+
+        - ``--heter_worker_num``: Number of heter_workers in each stage (It recommend to set when in the emulated distributed environment using single node)
+
+        - ``--heter_devices``: Type of heter_device in each stage
+
+        - ``--http_port``: Gloo http Port
+
+    Elastic Parameters:
+        - ``--elastic_server``: etcd server host:port, e.g., ``--elastic_server=127.0.0.1:2379``
+
+        - ``--job_id``: job unique id, e.g., ``--job_id=job1``
+
+        - ``--np``: job pod/node number, e.g., ``--np=2``
+
+        - ``--host``: bind host, default to POD_IP env.
+
+
+    Returns:
+        ``None``
+
+    Examples 1 (collective, single node):
+        .. code-block:: bash
+            :name: code-block-example-bash1
+
+            # For training on single node using 4 gpus.
+
+            python -m paddle.distributed.launch --gpus=0,1,2,3 train.py --lr=0.01
+
+    Examples 2 (collective, multi node):
+        .. code-block:: bash
+            :name: code-block-example-bash2
+
+            # The parameters of --gpus and --ips must be consistent in each node.
+
+            # For training on multiple nodes, e.g., 192.168.0.16, 192.168.0.17
+
+            # On 192.168.0.16:
+
+            python -m paddle.distributed.launch --gpus=0,1,2,3 --ips=192.168.0.16,192.168.0.17 train.py --lr=0.01
+
+            # On 192.168.0.17:
+            python -m paddle.distributed.launch --gpus=0,1,2,3 --ips=192.168.0.16,192.168.0.17 train.py --lr=0.01
+
+    Examples 3 (ps, cpu, single node):
+        .. code-block:: bash
+            :name: code-block-example-bash3
+
+            # To simulate distributed environment using single node, e.g., 2 servers and 4 workers.
+
+            python -m paddle.distributed.launch --server_num=2 --worker_num=4 train.py --lr=0.01
+
+    Examples 4 (ps, cpu, multi node):
+        .. code-block:: bash
+            :name: code-block-example-bash4
+
+            # For training on multiple nodes, e.g., 192.168.0.16, 192.168.0.17 where each node with 1 server and 2 workers.
+
+            # On 192.168.0.16:
+
+            python -m paddle.distributed.launch --servers="192.168.0.16:6170,192.168.0.17:6170" --workers="192.168.0.16:6171,192.168.0.16:6172,192.168.0.17:6171,192.168.0.17:6172" train.py --lr=0.01
+
+            # On 192.168.0.17:
+
+            python -m paddle.distributed.launch --servers="192.168.0.16:6170,192.168.0.17:6170" --workers="192.168.0.16:6171,192.168.0.16:6172,192.168.0.17:6171,192.168.0.17:6172" train.py --lr=0.01
+
+    Examples 5 (ps, gpu, single node):
+        .. code-block:: bash
+            :name: code-block-example-bash5
+
+           # To simulate distributed environment using single node, e.g., 2 servers and 4 workers, each worker use single gpu.
+
+            export CUDA_VISIBLE_DEVICES=0,1,2,3
+            python -m paddle.distributed.launch --server_num=2 --worker_num=4 train.py --lr=0.01
+
+    Examples 6 (ps, gpu, multi node):
+        .. code-block:: bash
+            :name: code-block-example-bash6
+
+            # For training on multiple nodes, e.g., 192.168.0.16, 192.168.0.17 where each node with 1 server and 2 workers.
+
+            # On 192.168.0.16:
+
+            export CUDA_VISIBLE_DEVICES=0,1
+            python -m paddle.distributed.launch --servers="192.168.0.16:6170,192.168.0.17:6170" --workers="192.168.0.16:6171,192.168.0.16:6172,192.168.0.17:6171,192.168.0.17:6172" train.py --lr=0.01
+
+            # On 192.168.0.17:
+
+            export CUDA_VISIBLE_DEVICES=0,1
+            python -m paddle.distributed.launch --servers="192.168.0.16:6170,192.168.0.17:6170" --workers="192.168.0.16:6171,192.168.0.16:6172,192.168.0.17:6171,192.168.0.17:6172" train.py --lr=0.01
+
+    Examples 7 (ps-heter, cpu + gpu, single node):
+        .. code-block:: bash
+            :name: code-block-example-bash7
+
+            # To simulate distributed environment using single node, e.g., 2 servers and 4 workers, two workers use gpu, two workers use cpu.
+
+            export CUDA_VISIBLE_DEVICES=0,1
+            python -m paddle.distributed.launch --server_num=2 --worker_num=2 --heter_worker_num=2 train.py --lr=0.01
+
+    Examples 8 (ps-heter, cpu + gpu, multi node):
+        .. code-block:: bash
+            :name: code-block-example-bash8
+
+            # For training on multiple nodes, e.g., 192.168.0.16, 192.168.0.17 where each node with 1 server, 1 gpu worker, 1 cpu worker.
+
+            # On 192.168.0.16:
+
+            export CUDA_VISIBLE_DEVICES=0
+            python -m paddle.distributed.launch --servers="192.168.0.16:6170,192.168.0.17:6170" --workers="192.168.0.16:6171,192.168.0.17:6171" --heter_workers="192.168.0.16:6172,192.168.0.17:6172" train.py --lr=0.01
+
+            # On 192.168.0.17:
+
+            export CUDA_VISIBLE_DEVICES=0
+            python -m paddle.distributed.launch --servers="192.168.0.16:6170,192.168.0.17:6170" --workers="192.168.0.16:6171,192.168.0.17:6171" --heter_workers="192.168.0.16:6172,192.168.0.17:6172" train.py --lr=0.01
+
+    Examples 9 (elastic):
+        .. code-block:: bash
+            :name: code-block-example-bash9
+
+            python -m paddle.distributed.launch --elastic_server=127.0.0.1:2379 --np=2 --job_id=job1  --gpus=0,1,2,3 train.py
+
+    """
+
+    args = _parse_args()
+    logger = get_logger()
+    _print_arguments(args)
+
+    if args.backend == 'auto':
+        distribute_mode = which_distributed_mode(
+            args
+        )  # which_distributed_mode must modify args.backend
+    else:
+        assert args.run_mode == 'collective' or args.run_mode is None, (
+            "When backend is not 'auto', run mode must be collective"
+        )
+        check_backend(args.backend)
+        distribute_mode = DistributeMode.COLLECTIVE
+
+    # assert args.backend in ['gloo', 'nccl', 'bkcl', 'heter', 'unknown']
+
+    if args.backend == 'gloo':
+        logger.warning("launch start with CPUONLY mode")
+
+    block_windows_and_macos(
+        args.backend
+    )  # raise error when using gloo on windows or macos
+
+    if enable_elastic(args, distribute_mode):
+        launch_elastic(args, distribute_mode)
+        return
+
+    if distribute_mode == DistributeMode.COLLECTIVE:
+        launch_collective(args)
+    else:
+        launch_ps(args, distribute_mode)
+
+
+if __name__ == "__main__":
+    launch()

python/paddle/distributed/fleet/layers/mpu/__init__.py

@@ -0,0 +1,26 @@
+#   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 .mp_layers import (  # noqa: F401
+    ColumnParallelLinear,
+    ParallelCrossEntropy,
+    RowParallelLinear,
+    VocabParallelEmbedding,
+)
+from .random import (  # noqa: F401
+    RNGStatesTracker,
+    dropout,
+    get_rng_state_tracker,
+    model_parallel_random_seed,
+)

python/paddle/distributed/fleet/layers/mpu/mp_layers.py

@@ -0,0 +1,886 @@
+#   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
+
+import paddle
+from paddle.autograd import PyLayer
+from paddle.base import core
+from paddle.distributed import fleet
+from paddle.nn import functional as F
+
+from ....communication.reduce import ReduceOp, _get_reduce_op
+from ....flex_checkpoint.dcp.sharded_weight import build_sharded_state_dict
+from ...base import topology as tp
+from ...utils.log_util import logger
+from . import mp_ops
+from .random import get_rng_state_tracker
+
+__all__ = []
+
+# Follow this paper to achieve the file:
+# Shoeybi M, Patwary M, Puri R, et al. Megatron-lm: Training multi-billion parameter
+# language models using model parallelism[J]. arXiv preprint arXiv:1909.08053, 2019. (https://arxiv.org/abs/1909.08053)
+
+
+def is_fused_matmul_bias_supported():
+    return hasattr(core.eager.ops.legacy, 'fused_gemm_epilogue')
+
+
+def is_fused_linear_param_grad_add_supported():
+    if (
+        paddle.is_compiled_with_cuda() and not paddle.is_compiled_with_rocm()
+    ) or paddle.is_compiled_with_xpu():
+        return hasattr(paddle._C_ops, 'fused_linear_param_grad_add')
+    else:
+        return False
+
+
+class VocabParallelEmbedding(paddle.nn.Layer):
+    """Embedding mp parallelized in the vocabulary dimension.
+    this class is used for splitting embedding in mp group.
+
+    Args:
+        num_embeddings(int): One element which indicate the size of the dictionary of embeddings.
+        embedding_dim(int): One element which indicate the size of each embedding vector respectively.
+        weight_attr(ParamAttr|None): To specify the weight parameter property. Default: None, which means the
+            default weight parameter property is used. See usage for details in :ref:`api_paddle_ParamAttr` . In addition,
+            user-defined or pre-trained word vectors can be loaded with the :attr:`param_attr` parameter.
+            The local word vector needs to be transformed into numpy format, and the shape of local word
+            vector should be consistent with :attr:`num_embeddings` . Then :ref:`api_paddle_nn_initializer_Assign`
+            is used to load custom or pre-trained word vectors. See code example for details.
+        mp_group(Group): The tensor parallel group.
+        name(str, optional): For detailed information, please refer
+               to :ref:`api_guide_Name`. Usually name is no need to set and
+               None by default.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.distributed import fleet
+
+            >>> class SimpleMPNet(paddle.nn.Layer):
+            ...     def __init__(self, vocab_size, hidden_size, inner_size, output_size):
+            ...         super().__init__()
+            ...         self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
+            ...             hidden_size,
+            ...             inner_size,
+            ...             gather_output=False,
+            ...             has_bias=True,
+            ...         )
+            ...         self.linear2 = fleet.meta_parallel.RowParallelLinear(
+            ...             inner_size,
+            ...             hidden_size,
+            ...             input_is_parallel=True,
+            ...             has_bias=True,
+            ...         )
+            ...         self.linear3 = paddle.nn.Linear(hidden_size, output_size)
+            ...         self.embedding = fleet.meta_parallel.VocabParallelEmbedding(vocab_size, hidden_size)
+            ...
+            ...     def forward(self, x):
+            ...         x = self.embedding(x)
+            ...         x = self.linear1(x)
+            ...         x = self.linear2(x)
+            ...         x = self.linear3(x)
+            ...         return x
+
+    """
+
+    def __init__(
+        self,
+        num_embeddings,
+        embedding_dim,
+        weight_attr=None,
+        mp_group=None,
+        name=None,
+    ):
+        super().__init__()
+
+        self.model_parallel_group = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
+            if mp_group is None
+            else mp_group
+        )
+        self.world_size = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
+            if mp_group is None
+            else mp_group.nranks
+        )
+        self.rank = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
+            if mp_group is None
+            else mp_group.rank
+        )
+
+        self.origin_num_embeddings = num_embeddings
+        self.is_mp = self.world_size > 1
+
+        assert num_embeddings % self.world_size == 0, (
+            "The length of the vocabulary must be divisible by the parallelism degree of MP"
+        )
+
+        per_part_size = num_embeddings // self.world_size
+
+        self.vocab_start_index = self.rank * per_part_size
+        self._dtype = self._helper.get_default_dtype()
+        self._size = [per_part_size, embedding_dim]
+        self._weight_attr = weight_attr
+        self._name = name
+        self.num_embeddings = num_embeddings
+
+        if self.is_mp and paddle.in_dynamic_mode():
+            with get_rng_state_tracker().rng_state():
+                self.weight = self.create_parameter(
+                    attr=self._weight_attr,
+                    shape=self._size,
+                    dtype=self._dtype,
+                    is_bias=False,
+                )
+        else:
+            self.weight = self.create_parameter(
+                attr=self._weight_attr,
+                shape=self._size,
+                dtype=self._dtype,
+                is_bias=False,
+            )
+
+        self.weight.is_distributed = True if self.is_mp else False
+        if self.weight.is_distributed:
+            self.weight.split_axis = 0
+
+    def forward(self, x):
+        if self.is_mp:
+            output_parallel = mp_ops._c_lookup_table(
+                self.weight,
+                x,
+                start_index=self.vocab_start_index,
+                vocab_size=self.num_embeddings,
+                name=self._name,
+            )
+            output = mp_ops._mp_allreduce(
+                output_parallel,
+                group=self.model_parallel_group,
+                use_calc_stream=True,
+                use_model_parallel=True,
+            )
+        else:
+            output = F.embedding(
+                x,
+                weight=self.weight,
+                padding_idx=None,
+                sparse=False,
+                name=self._name,
+            )
+        return output
+
+    def sharded_state_dict(
+        self,
+        structured_name_prefix: str = "",
+    ):
+        state_dict = self.state_dict(structured_name_prefix="")
+        return build_sharded_state_dict(
+            state_dict, {"weight": 0}, structured_name_prefix
+        )
+
+
+_raise_cuda_env_unset_warning = True
+
+
+class InnerOverlapLinear(paddle.autograd.PyLayer):
+    @staticmethod
+    def forward(
+        ctx,
+        x,
+        weight,
+        bias,
+        fuse_matmul_bias,
+        mp_async_allreduce,
+        mp_skip_c_identity,
+        mp_fused_linear_param_grad_add,
+        model_parallel_group,
+    ):
+        ctx.save_for_backward(x, weight, bias)
+        ctx.model_parallel_group = model_parallel_group
+        ctx.mp_fused_linear_param_grad_add = mp_fused_linear_param_grad_add
+        if mp_skip_c_identity is False:
+            x = paddle._legacy_C_ops.c_identity(
+                x,
+                'use_calc_stream',
+                True,
+                'ring_id',
+                model_parallel_group.id,
+                'use_model_parallel',
+                True,
+            )
+        if not fuse_matmul_bias:
+            return paddle._C_ops.linear(x, weight, bias)
+        else:
+            return paddle._legacy_C_ops.fused_gemm_epilogue(x, weight, bias)
+
+    @staticmethod
+    def backward(ctx, dy):
+        x, weight, bias = ctx.saved_tensor()
+        if dy.dtype == weight.dtype:
+            dx = paddle.matmul(dy, weight, transpose_y=True)
+        else:
+            dx = paddle.matmul(
+                dy, paddle.cast(weight, dtype=dy.dtype), transpose_y=True
+            )
+        op_type = _get_reduce_op(ReduceOp.SUM)
+        task = ctx.model_parallel_group.process_group.all_reduce(
+            dx, op_type, sync_op=False
+        )
+        # Using small operation to preempt GPU SMs for all_reduce to achieve overlap.
+        if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
+            global _raise_cuda_env_unset_warning
+            if _raise_cuda_env_unset_warning:
+                logger.warning(
+                    "You set mp_async_allreduce=True, but you forget to 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."
+                )
+            _raise_cuda_env_unset_warning = False
+            tmp = paddle.ones([512])
+
+        if ctx.mp_fused_linear_param_grad_add:
+            if not is_fused_linear_param_grad_add_supported():
+                raise NotImplementedError(
+                    "You set mp_fused_linear_param_grad_add=True, "
+                    "however, the paddle you are using not support this operation. "
+                    "Please unset fused_linear_param_grad_add or use paddle compiled "
+                    "with cuda 11.6 or higher."
+                )
+
+            if bias is None:
+                if hasattr(weight, "main_grad"):
+                    (
+                        weight.main_grad,
+                        _,
+                    ) = paddle._C_ops.fused_linear_param_grad_add(
+                        x, dy, weight.main_grad, None, True, False
+                    )
+                    task.wait()
+                    return dx, None
+                else:
+                    if weight.grad is not None:
+                        (
+                            weight.grad,
+                            _,
+                        ) = paddle._C_ops.fused_linear_param_grad_add(
+                            x, dy, weight.grad, None, False, False
+                        )
+                        task.wait()
+                        return dx, None
+                    else:
+                        (
+                            dw,
+                            _,
+                        ) = paddle._C_ops.fused_linear_param_grad_add(
+                            x, dy, None, None, False, False
+                        )
+                        task.wait()
+                        return dx, dw
+
+            if hasattr(weight, "main_grad") and hasattr(bias, "main_grad"):
+                (
+                    weight.main_grad,
+                    bias.main_grad,
+                ) = paddle._C_ops.fused_linear_param_grad_add(
+                    x,
+                    dy,
+                    weight.main_grad,
+                    bias.main_grad,
+                    True,
+                    True,
+                )
+                task.wait()
+                return dx, None, None
+            else:
+                if weight.grad is not None:
+                    assert bias.grad is not None
+                    (
+                        weight.grad,
+                        bias.grad,
+                    ) = paddle._C_ops.fused_linear_param_grad_add(
+                        x, dy, weight.grad, bias.grad, False, True
+                    )
+                    task.wait()
+                    return dx, None, None
+                else:
+                    # When main_grad is not enabled and gradient_accumulation is used, the grad is not initialized for the first acc step.
+                    (
+                        dw,
+                        dbias,
+                    ) = paddle._C_ops.fused_linear_param_grad_add(
+                        x, dy, None, None, False, True
+                    )
+                    task.wait()
+                    return dx, dw, dbias
+        else:
+            dy = dy.reshape([-1, dy.shape[-1]])
+            dw = paddle.matmul(
+                x.reshape([-1, x.shape[-1]]),
+                dy,
+                transpose_x=True,
+            )
+            if bias is None:
+                task.wait()
+                return dx, dw
+            else:
+                dbias = paddle.sum(dy, axis=0)
+                task.wait()
+                return dx, dw, dbias
+
+
+class ColumnParallelLinear(paddle.nn.Layer):
+    """Linear layer with mp parallelized(column).
+    this class is used for splitting Linear Layer in mp group, column split the weight of the Linear layer.
+
+    Args:
+        in_features(int): The number of input units.
+        out_features(int): The number of output units.
+        weight_attr(ParamAttr|None): The attribute for the learnable weight of this layer. The default value is None
+            and the weight will be initialized to zero. For detailed information, please refer to paddle.ParamAttr.
+        has_bias(bool): whether to add bias.
+        gather_output(bool): whether to do allgather for the output of each rank.
+        fuse_matmul_bias(bool): whether to fuse matmul and bias.
+        mp_group(Group): The tensor parallel group.
+        name(str, optional): Normally there is no need for user to set this parameter.
+            For detailed information, please refer to :ref:`api_guide_Name` .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.distributed import fleet
+
+            >>> class SimpleMPNet(paddle.nn.Layer):
+            ...     def __init__(self, vocab_size, hidden_size, inner_size, output_size):
+            ...         super().__init__()
+            ...         self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
+            ...             hidden_size,
+            ...             inner_size,
+            ...             gather_output=False,
+            ...             has_bias=True,
+            ...         )
+            ...         self.linear2 = fleet.meta_parallel.RowParallelLinear(
+            ...             inner_size,
+            ...             hidden_size,
+            ...             input_is_parallel=True,
+            ...             has_bias=True,
+            ...         )
+            ...         self.linear3 = paddle.nn.Linear(hidden_size, output_size)
+            ...         self.embedding = fleet.meta_parallel.VocabParallelEmbedding(vocab_size, hidden_size)
+            ...
+            ...     def forward(self, x):
+            ...         x = self.embedding(x)
+            ...         x = self.linear1(x)
+            ...         x = self.linear2(x)
+            ...         x = self.linear3(x)
+            ...         return x
+    """
+
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        weight_attr=None,
+        has_bias=None,
+        gather_output=True,
+        fuse_matmul_bias=False,
+        mp_group=None,
+        name=None,
+    ):
+        super().__init__()
+
+        self.model_parallel_group = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
+            if mp_group is None
+            else mp_group
+        )
+        self.world_size = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
+            if mp_group is None
+            else mp_group.nranks
+        )
+        self._name = name
+        self.is_mp = self.world_size > 1
+
+        self.gather_output = gather_output
+        assert out_features % self.world_size == 0, (
+            f"Number of column of the weight for linear ({out_features}) must be"
+            f" divisible by model parallel size ({self.world_size})"
+        )
+        self.output_size_per_partition = out_features // self.world_size
+
+        self._weight_attr = weight_attr
+        self._dtype = self._helper.get_default_dtype()
+
+        if self.is_mp and paddle.in_dynamic_mode():
+            with get_rng_state_tracker().rng_state():
+                self.weight = self.create_parameter(
+                    shape=[in_features, self.output_size_per_partition],
+                    attr=self._weight_attr,
+                    dtype=self._dtype,
+                    is_bias=False,
+                )
+        else:
+            self.weight = self.create_parameter(
+                shape=[in_features, self.output_size_per_partition],
+                attr=self._weight_attr,
+                dtype=self._dtype,
+                is_bias=False,
+            )
+
+        self.weight.is_distributed = True if self.is_mp else False
+
+        if self.weight.is_distributed:
+            self.weight.split_axis = 1
+
+        if has_bias:
+            # initialize bias to zero like Megatron
+            self.bias = self.create_parameter(
+                shape=[self.output_size_per_partition],
+                attr=paddle.nn.initializer.Constant(value=0.0),
+                dtype=self._dtype,
+                is_bias=True,
+            )
+            self.bias.is_distributed = True if self.is_mp else False
+            if self.bias.is_distributed:
+                self.bias.split_axis = 0
+        else:
+            self.bias = None
+
+        self.linear = F.linear
+
+        self.fuse_matmul_bias = fuse_matmul_bias
+
+        mp_configs = fleet.fleet._user_defined_strategy.hybrid_configs[
+            "mp_configs"
+        ]
+        self.mp_async_allreduce = self.is_mp and mp_configs.mp_async_allreduce
+        self.mp_skip_c_identity = (
+            self.is_mp
+            and mp_configs.mp_async_allreduce
+            and mp_configs.mp_skip_c_identity
+        )
+        self.mp_fused_linear_param_grad_add = (
+            self.is_mp
+            and mp_configs.mp_async_allreduce
+            and mp_configs.mp_fused_linear_param_grad_add
+        )
+        if (
+            self.mp_async_allreduce
+            or self.mp_skip_c_identity
+            or self.mp_fused_linear_param_grad_add
+        ):
+            assert paddle.in_dynamic_mode(), (
+                "mp_async_allreduce, mp_skip_c_identity and mp_fused_linear_param_grad_add are only available under dygraph mode"
+            )
+        if self.fuse_matmul_bias:
+            if not is_fused_matmul_bias_supported():
+                raise NotImplementedError(
+                    "You set fuse_matmul_bias=True in ColumnParallelLinear, "
+                    "however, the paddle you are using not support this operation. "
+                    "Please set fuse_matmul_bias=False or use paddle compiled "
+                    "with cuda 11.6 or higher."
+                )
+            from paddle.incubate.nn.functional import fused_linear
+
+            self.linear = fused_linear
+
+    def forward(self, x):
+        # use inner api to process identity
+
+        def _overlap_linear():
+            return InnerOverlapLinear.apply(
+                x,
+                self.weight,
+                self.bias,
+                self.fuse_matmul_bias,
+                self.mp_async_allreduce,
+                self.mp_skip_c_identity,
+                self.mp_fused_linear_param_grad_add,
+                self.model_parallel_group,
+            )
+
+        if self.mp_async_allreduce:
+            output_parallel = _overlap_linear()
+        else:
+            if self.is_mp:
+                input_parallel = mp_ops._c_identity(
+                    x,
+                    group=self.model_parallel_group,
+                    skip_c_identity_dynamic=self.mp_skip_c_identity,
+                )
+            else:
+                input_parallel = x
+
+            output_parallel = self.linear(
+                input_parallel, self.weight, self.bias, name=self._name
+            )
+
+        if self.gather_output and self.is_mp:
+            output = mp_ops._c_concat(
+                output_parallel, group=self.model_parallel_group
+            )
+        else:
+            output = output_parallel
+        return output
+
+    def sharded_state_dict(
+        self,
+        structured_name_prefix: str = "",
+    ):
+        state_dict = self.state_dict(structured_name_prefix="")
+        return build_sharded_state_dict(
+            state_dict, {"weight": 1, "bias": 0}, structured_name_prefix
+        )
+
+
+class MPScale(PyLayer):
+    @staticmethod
+    def forward(ctx, x, mp_degree):
+        out = paddle.scale(x, 1.0 / mp_degree)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        return dout
+
+
+class RowParallelLinear(paddle.nn.Layer):
+    """Linear layer with mp parallelized(row).
+    this class is used for splitting Linear Layer in mp group, row split the weight of the Linear layer.
+
+    Args:
+        in_features(int): The number of input units.
+        out_features(int): The number of output units.
+        weight_attr(ParamAttr|None): The attribute for the learnable weight of this layer. The default value is None
+            and the weight will be initialized to zero. For detailed information, please refer to paddle.ParamAttr.
+        has_bias(bool): whether to add bias.
+        input_is_parallel(bool): whether the input has already been split across the mp group.
+        fuse_matmul_bias(bool): whether to fuse matmul and bias.
+        mp_group(Group): The tensor parallel group.
+        name(str, optional): Normally there is no need for user to set this parameter.
+            For detailed information, please refer to :ref:`api_guide_Name` .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.distributed import fleet
+
+            >>> class SimpleMPNet(paddle.nn.Layer):
+            ...     def __init__(self, vocab_size, hidden_size, inner_size, output_size):
+            ...         super().__init__()
+            ...         self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
+            ...             hidden_size,
+            ...             inner_size,
+            ...             gather_output=False,
+            ...             has_bias=True,
+            ...         )
+            ...         self.linear2 = fleet.meta_parallel.RowParallelLinear(
+            ...             inner_size,
+            ...             hidden_size,
+            ...             input_is_parallel=True,
+            ...             has_bias=True,
+            ...         )
+            ...         self.linear3 = paddle.nn.Linear(hidden_size, output_size)
+            ...         self.embedding = fleet.meta_parallel.VocabParallelEmbedding(vocab_size, hidden_size)
+            ...
+            ...     def forward(self, x):
+            ...         x = self.embedding(x)
+            ...         x = self.linear1(x)
+            ...         x = self.linear2(x)
+            ...         x = self.linear3(x)
+            ...         return x
+
+    """
+
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        weight_attr=None,
+        has_bias=True,
+        input_is_parallel=False,
+        fuse_matmul_bias=False,
+        mp_group=None,
+        name=None,
+    ):
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.input_is_parallel = input_is_parallel
+        self._weight_attr = weight_attr
+        self._dtype = self._helper.get_default_dtype()
+        self._name = name
+
+        self.model_parallel_group = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
+            if mp_group is None
+            else mp_group
+        )
+        self.world_size = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
+            if mp_group is None
+            else mp_group.nranks
+        )
+        self.rank = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
+            if mp_group is None
+            else mp_group.rank
+        )
+
+        self.is_mp = self.world_size > 1
+        mp_configs = fleet.fleet._user_defined_strategy.hybrid_configs[
+            "mp_configs"
+        ]
+        self.mp_async_allreduce = self.is_mp and mp_configs.mp_async_allreduce
+        self.mp_skip_c_identity = (
+            self.is_mp
+            and mp_configs.mp_async_allreduce
+            and mp_configs.mp_skip_c_identity
+        )
+        self.mp_fused_linear_param_grad_add = (
+            self.is_mp
+            and mp_configs.mp_async_allreduce
+            and mp_configs.mp_fused_linear_param_grad_add
+        )
+        if (
+            self.mp_async_allreduce
+            or self.mp_skip_c_identity
+            or self.mp_fused_linear_param_grad_add
+        ):
+            assert paddle.in_dynamic_mode(), (
+                "mp_async_allreduce, mp_skip_c_identity and mp_fused_linear_param_grad_add are only available under dygraph mode"
+            )
+        assert in_features % self.world_size == 0, (
+            f"Number of row of the weight for linear ({in_features}) must be"
+            f" divisible by model parallel size ({self.world_size})"
+        )
+
+        self.input_size_per_partition = in_features // self.world_size
+
+        if self.is_mp and paddle.in_dynamic_mode():
+            with get_rng_state_tracker().rng_state():
+                self.weight = self.create_parameter(
+                    shape=[self.input_size_per_partition, self.out_features],
+                    attr=self._weight_attr,
+                    dtype=self._dtype,
+                    is_bias=False,
+                )
+        else:
+            self.weight = self.create_parameter(
+                shape=[self.input_size_per_partition, self.out_features],
+                attr=self._weight_attr,
+                dtype=self._dtype,
+                is_bias=False,
+            )
+
+        self.weight.is_distributed = True if self.is_mp else False
+        if self.weight.is_distributed:
+            self.weight.split_axis = 0
+
+        if has_bias:
+            self.bias = self.create_parameter(
+                shape=[self.out_features],
+                attr=paddle.nn.initializer.Constant(value=0.0),
+                dtype=self._dtype,
+                is_bias=True,
+            )
+        else:
+            self.bias = None
+
+        self.linear = F.linear
+
+        if fuse_matmul_bias:
+            if not is_fused_matmul_bias_supported():
+                raise NotImplementedError(
+                    "You set fuse_matmul_bias=True in RowParallelLinear, "
+                    "however, the paddle you are using not support this operation. "
+                    "Please set fuse_matmul_bias=False or use paddle compiled "
+                    "with cuda 11.6 or higher."
+                )
+            from paddle.incubate.nn.functional import fused_linear
+
+            self.linear = fused_linear
+        self.fuse_matmul_bias = fuse_matmul_bias
+
+    def forward(self, x):
+        if self.input_is_parallel or (not self.is_mp):
+            input_parallel = x
+        else:
+            # split last dim
+            input_parallel = mp_ops._c_split(x, group=self.model_parallel_group)
+
+        if self.is_mp:
+            if self.fuse_matmul_bias:
+                bias = MPScale.apply(self.bias, self.world_size)
+                output_parallel = self.linear(
+                    input_parallel, self.weight, bias, name=self._name
+                )
+                output = mp_ops._mp_allreduce(
+                    output_parallel,
+                    group=self.model_parallel_group,
+                    use_calc_stream=True,
+                    use_model_parallel=True,
+                    skip_c_identity_dynamic=self.mp_skip_c_identity,
+                )
+            else:
+                output_parallel = self.linear(
+                    input_parallel, self.weight, name=self._name
+                )
+                output_ = mp_ops._mp_allreduce(
+                    output_parallel,
+                    group=self.model_parallel_group,
+                    use_calc_stream=True,
+                    use_model_parallel=True,
+                    skip_c_identity_dynamic=self.mp_skip_c_identity,
+                )
+                output = (
+                    output_ + self.bias if self.bias is not None else output_
+                )
+        else:
+            output = self.linear(
+                input_parallel, self.weight, self.bias, name=self._name
+            )
+
+        return output
+
+    def sharded_state_dict(
+        self,
+        structured_name_prefix: str = "",
+    ):
+        state_dict = self.state_dict(structured_name_prefix="")
+        return build_sharded_state_dict(
+            state_dict, {"weight": 0}, structured_name_prefix
+        )
+
+
+class ParallelCrossEntropy(paddle.nn.Layer):
+    """CrossEntropy with mp parallelized.
+    this class is used for splitting softmax cross entropy in mp group.
+
+    Args:
+        mp_group(Group): The tensor parallel group.
+        name(str, optional): Normally there is no need for user to set this parameter.
+            For detailed information, please refer to :ref:`api_guide_Name` .
+        ignore_index (long int, optional):  Specifies a target value that is ignored and
+            does not contribute to the loss. A negative value means that no label value
+            needs to be ignored. Default is -100 .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('No img to demonstrate')
+            >>> from paddle.distributed.fleet.layers.mpu import ParallelCrossEntropy
+            >>> loss_func = ParallelCrossEntropy
+            >>> loss = loss_func(img, label)
+
+    """
+
+    def __init__(self, mp_group=None, name=None, ignore_index=-100):
+        super().__init__()
+        self.name = name
+        self.model_parallel_group = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
+            if mp_group is None
+            else mp_group
+        )
+        self.world_size = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
+            if mp_group is None
+            else mp_group.nranks
+        )
+        self.rank = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
+            if mp_group is None
+            else mp_group.rank
+        )
+        self.ignore_index = ignore_index
+
+    def forward(self, input, label):
+        loss = mp_ops._c_softmax_with_cross_entropy(
+            input,
+            label,
+            group=self.model_parallel_group,
+            ignore_index=self.ignore_index,
+        )
+        return loss
+
+
+class ParallelMultiLabelCrossEntropy(paddle.nn.Layer):
+    """CrossEntropy with mp parallelized.
+    this class is used for splitting softmax cross entropy in mp group.
+
+    Args:
+        mp_group(Group): The tensor parallel group.
+        name(str, optional): Normally there is no need for user to set this parameter.
+            For detailed information, please refer to :ref:`api_guide_Name` .
+        ignore_index (long int, optional):  Specifies a target value that is ignored and
+            does not contribute to the loss. A negative value means that no label value
+            needs to be ignored. Default is -100 .
+        sum_multi_label_loss (bool, optional): Whether to sum the loss. Default is True .
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('No img to demonstrate')
+            >>> from paddle.distributed.fleet.layers.mpu import ParallelMultiLabelCrossEntropy
+            >>> loss_func = ParallelMultiLabelCrossEntropy()
+            >>> loss = loss_func(img, label, smooth_weight)
+
+    """
+
+    def __init__(
+        self,
+        mp_group=None,
+        name=None,
+        ignore_index=-100,
+        sum_multi_label_loss=True,
+    ):
+        super().__init__()
+        self.name = name
+        self.model_parallel_group = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
+            if mp_group is None
+            else mp_group
+        )
+        self.world_size = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
+            if mp_group is None
+            else mp_group.nranks
+        )
+        self.rank = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
+            if mp_group is None
+            else mp_group.rank
+        )
+        self.ignore_index = ignore_index
+        self.sum_multi_label_loss = sum_multi_label_loss
+
+    def forward(self, input, label, smooth_weight):
+        loss = mp_ops._c_softmax_with_multi_label_cross_entropy(
+            input,
+            label,
+            smooth_weight,
+            group=self.model_parallel_group,
+            ignore_index=self.ignore_index,
+            sum_multi_label_loss=self.sum_multi_label_loss,
+        )
+        return loss

python/paddle/distributed/fleet/layers/mpu/random.py

@@ -0,0 +1,266 @@
+#   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 contextlib
+
+import numpy as np
+
+import paddle
+from paddle import _legacy_C_ops
+from paddle.base import core
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.common_ops_import import Variable
+from paddle.framework import LayerHelper, in_dynamic_mode
+
+__all__ = []
+
+MODEL_PARALLEL_RNG = 'model_parallel_rng'
+
+# This file is inspired by Megatron to control random states for MP:
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/random.py
+
+
+class RNGStatesTracker:
+    """
+    Tracker the RNG states.
+    """
+
+    def __init__(self):
+        # Map from name to the rng state.
+        self.states_ = {}
+        self.seeds_ = set()
+
+    def reset(self):
+        self.states_ = {}
+        self.seeds_ = set()
+
+    def add(self, name, seed):
+        if seed in self.seeds_:
+            raise ValueError(f'seed {seed} already exists')
+        self.seeds_.add(seed)
+        if name in self.states_:
+            raise ValueError(f'state {name} already exists')
+        orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
+        # register a new state and set that state with the seed, store the indices into states_
+        self.states_[name] = paddle.incubate.register_rng_state_as_index()
+        paddle.seed(seed)
+        paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
+
+    def get_states_tracker(self):
+        states = {}
+        orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
+        for name in self.states_:
+            # switch index to name
+            paddle.incubate.set_rng_state(self.states_[name], use_index=True)
+            # export the saved state
+            states[name] = paddle.get_rng_state()
+        paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
+        return states
+
+    def set_states_tracker(self, states):
+        orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
+        for name in states:
+            if name not in self.states_:
+                raise ValueError(f'state {name} does not exists')
+            # switch index to name
+            paddle.incubate.set_rng_state(self.states_[name], use_index=True)
+            # set the state to saved state
+            paddle.set_rng_state(states[name])
+
+        paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
+
+    @contextlib.contextmanager
+    def rng_state(self, name=MODEL_PARALLEL_RNG):
+        if name not in self.states_:
+            raise ValueError(f'state {name} does not exist')
+        orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
+        paddle.incubate.set_rng_state(self.states_[name], use_index=True)
+        try:
+            yield
+        finally:
+            self.states_[name] = paddle.incubate.get_rng_state(use_index=True)
+            paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
+
+
+RNG_STATE_TRACKER = RNGStatesTracker()
+
+
+def get_rng_state_tracker():
+    return RNG_STATE_TRACKER
+
+
+def model_parallel_random_seed(seed=None):
+    from paddle.distributed import fleet
+
+    hcg = fleet.get_hybrid_communicate_group()
+
+    mp_rank = hcg.get_model_parallel_rank()
+    mp_size = hcg.get_model_parallel_world_size()
+
+    pp_rank = hcg.get_stage_id()
+    pp_size = hcg.get_pipe_parallel_world_size()
+
+    if seed:
+        global_seed = seed
+        # dp/sharding seed is same
+        local_seed = seed + 1 + mp_rank * pp_size + pp_rank
+    else:
+        global_seed = np.random.randint(0, 10000)
+        local_seed = global_seed + 1 + mp_rank * pp_size + pp_rank
+
+    RNG_STATE_TRACKER.reset()
+    RNG_STATE_TRACKER.add(MODEL_PARALLEL_RNG, local_seed)
+    paddle.seed(global_seed)
+
+
+def dropout(
+    x,
+    p=0.5,
+    axis=None,
+    rng_name=None,
+    training=True,
+    mode="upscale_in_train",
+    name=None,
+):
+    """
+    Dropout is a regularization technique for reducing overfitting by preventing
+    neuron co-adaption during training. The dropout operator randomly sets the
+    outputs of some units to zero, while upscale others according to the given
+    dropout probability.
+
+    Args:
+        x (Tensor): The input tensor. The data type is float32 or float64.
+        p (float|int): Probability of setting units to zero. Default 0.5.
+        axis (int|list|tuple): The axis along which the dropout is performed. Default None.
+        rng_name (str): The random seed generator name, which used to obtain deterministic results.
+        training (bool): A flag indicating whether it is in train phrase or not. Default True.
+        mode(str): ['upscale_in_train'(default) | 'downscale_in_infer'].
+
+                           1. upscale_in_train(default), upscale the output at training time
+
+                              - train: out = input * mask / ( 1.0 - dropout_prob )
+                              - inference: out = input
+
+                           2. downscale_in_infer, downscale the output at inference
+
+                              - train: out = input * mask
+                              - inference: out = input * (1.0 - dropout_prob)
+        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        A Tensor representing the dropout, has same shape and data type as `x` .
+
+
+    Examples:
+        We use ``p=0.5`` in the following description for simplicity.
+
+        1. When ``axis=None`` , this is commonly used dropout, which dropout each element of x randomly.
+
+        ..  code-block:: text
+
+            Let's see a simple case when x is a 2d tensor with shape 2*3:
+            [[1 2 3]
+             [4 5 6]]
+            we generate mask with the same shape as x, which is 2*3. The value of mask is
+            sampled from a Bernoulli distribution randomly. For example, we may get such mask:
+            [[0 1 0]
+             [1 0 1]]
+            So the output is obtained from elementwise multiply of x and mask:
+            [[0 2 0]
+             [4 0 6]]
+            Using default setting, i.e. ``mode='upscale_in_train'`` ,
+            if in training phase, the final upscale output is:
+            [[0 4 0 ]
+             [8 0 12]]
+            if in test phase, the output is the same as input:
+            [[1 2 3]
+             [4 5 6]]
+            we can also set ``mode='downscale_in_infer'`` , then
+            if in training phase, the final output is:
+            [[0 2 0]
+             [4 0 6]]
+            if in test phase, the scale output is:
+            [[0.5 1.  1.5]
+             [2.  2.5 3. ]]
+
+    """
+    if rng_name is None:
+        return paddle.nn.functional.dropout(x, p, axis, training, mode, name)
+
+    if not isinstance(p, (float, int, Variable)):
+        raise TypeError("p argument should be a number(int|float) or Variable")
+
+    # fast return for p == 0
+    if isinstance(p, (int, float)) and p == 0:
+        return x
+
+    assert 0 <= p <= 1, ValueError("p argument should between 0 and 1")
+    assert mode in ('downscale_in_infer', 'upscale_in_train'), ValueError(
+        "mode argument should be 'downscale_in_infer' or 'upscale_in_train'"
+    )
+
+    assert axis is None, TypeError(
+        "unsupported axis when using random seed generator"
+    )
+
+    mode = (
+        'downgrade_in_infer' if mode == 'downscale_in_infer' else mode
+    )  # semantic transfer
+
+    # dygraph using tracker, doesn't need determinate seed
+    if in_dynamic_mode():
+        out, mask = _legacy_C_ops.dropout(
+            x,
+            'dropout_prob',
+            p,
+            'is_test',
+            not training,
+            'fix_seed',
+            False,
+            'seed',
+            0,
+            'dropout_implementation',
+            mode,
+        )
+        return out
+    else:
+        if isinstance(p, Variable) and not p.shape != [1]:
+            raise TypeError(
+                f"Required p.shape == [1] if type(p) is Variable, but received p.shape = {p.shape}"
+            )
+
+        helper = LayerHelper('dropout', **locals())
+        check_variable_and_dtype(
+            x, 'x', ['float16', 'float32', 'float64'], 'dropout'
+        )
+
+        seed = helper.create_variable_for_type_inference(dtype=paddle.int32)
+        helper.append_op(type='seed', outputs={'Out': seed})
+
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+        mask = helper.create_variable_for_type_inference(
+            dtype=core.VarDesc.VarType.UINT8, stop_gradient=True
+        )
+
+        helper.append_op(
+            type='dropout',
+            inputs={'X': [x], 'Seed': seed},
+            outputs={'Out': [out], 'Mask': [mask]},
+            attrs={
+                'dropout_prob': p,
+                'is_test': not training,
+                'dropout_implementation': mode,
+            },
+        )
+        return out

python/paddle/distributed/fleet/meta_optimizers/__init__.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 NVIDIA Corporation. 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
+
+from .amp_optimizer import AMPOptimizer  # noqa: F401
+from .asp_optimizer import ASPOptimizer  # noqa: F401
+from .dgc_optimizer import (  # noqa: F401
+    DGCMomentumOptimizer,
+    DGCOptimizer,
+)
+from .dygraph_optimizer import (  # noqa: F401
+    HeterParallelOptimizer,
+    HybridParallelGradScaler,
+    HybridParallelOptimizer,
+)
+from .fp16_allreduce_optimizer import FP16AllReduceOptimizer  # noqa: F401
+from .gradient_merge_optimizer import GradientMergeOptimizer  # noqa: F401
+from .lamb_optimizer import LambOptimizer  # noqa: F401
+from .lars_optimizer import LarsOptimizer  # noqa: F401
+from .localsgd_optimizer import (  # noqa: F401
+    AdaptiveLocalSGDOptimizer,
+    LocalSGDOptimizer,
+)
+from .muon_sharding_optimizer import MuonShardingOptimizer  # noqa: F401
+from .pipeline_optimizer import PipelineOptimizer  # noqa: F401
+from .ps_optimizer import ParameterServerOptimizer  # noqa: F401
+from .qat_optimizer import QATOptimizer  # noqa: F401
+from .raw_program_optimizer import RawProgramOptimizer  # noqa: F401
+from .recompute_optimizer import RecomputeOptimizer  # noqa: F401
+from .sharding_optimizer import ShardingOptimizer  # noqa: F401
+from .tensor_parallel_optimizer import TensorParallelOptimizer  # noqa: F401

python/paddle/distributed/fleet/meta_optimizers/amp_optimizer.py

@@ -0,0 +1,139 @@
+#   Copyright (c) 2020 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
+
+import paddle.static.amp as mixed_precision
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class AMPOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.wrapped_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+        ]
+        self.meta_optimizers_black_list = ["DGCOptimizer"]
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _init_wrapped_opt(self):
+        if self.wrapped_opt is not None:
+            return
+
+        config = self.user_defined_strategy.amp_configs
+
+        custom_white_list = set(config['custom_white_list'])
+        custom_black_list = set(config['custom_black_list'])
+        custom_black_varnames = set(config['custom_black_varnames'])
+        amp_lists = mixed_precision.AutoMixedPrecisionLists(
+            custom_white_list, custom_black_list, custom_black_varnames
+        )
+
+        self.wrapped_opt = mixed_precision.decorate(
+            self.inner_opt,
+            amp_lists,
+            config['init_loss_scaling'],
+            config['incr_every_n_steps'],
+            config['decr_every_n_nan_or_inf'],
+            config['incr_ratio'],
+            config['decr_ratio'],
+            config['use_dynamic_loss_scaling'],
+            config['use_pure_fp16'],
+            config['use_fp16_guard'],
+        )
+
+        # if worker_num > 1, all cards will communication with each other,
+        # add is_distributed to optimize amp, overlap communication and
+        # computation by split the check_finite_and_unscale op.
+        is_distributed = self.role_maker._worker_num() > 1
+        if self.user_defined_strategy.sharding:
+            # FIXME(wangxi). sharding failed when split check_finite_and_unscale
+            # FIXME(JZ-LIANG). To support Sharding-Megatron-AMP, Megatron should follow Sharding's behavior that to disable is_distributed.
+            is_distributed = False
+        self.wrapped_opt._set_distributed(is_distributed)
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.amp:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.amp = False
+        dist_strategy.amp_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.amp = True
+        dist_strategy.amp_configs = {
+            "init_loss_scaling": 32768.0,
+            "incr_every_n_steps": 1000,
+            "decr_every_n_nan_or_inf": 2,
+            "incr_ratio": 2.0,
+            "decr_ratio": 0.8,
+            "use_dynamic_loss_scaling": True,
+        }
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        # maybe inner_opt of other meta optimizer
+        self._init_wrapped_opt()
+        return self.wrapped_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_gradients(self, params_grads):
+        return self.wrapped_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.wrapped_opt.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self._init_wrapped_opt()
+        optimize_ops, params_grads = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads
+
+    def amp_init(
+        self, place, scope=None, test_program=None, use_fp16_test=False
+    ):
+        return self.wrapped_opt.amp_init(
+            place, scope, test_program, use_fp16_test
+        )
+
+    def get_loss_scaling(self):
+        return self.wrapped_opt.get_loss_scaling()

python/paddle/distributed/fleet/meta_optimizers/asp_optimizer.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 NVIDIA Corporation. 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
+
+
+from paddle.incubate.asp import ASPHelper
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class ASPOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+            "GradientMergeOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.asp:
+            return True
+
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.asp = False
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.asp = True
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = ASPHelper._minimize(
+            self.inner_opt,
+            loss,
+            startup_program=startup_program,
+            parameter_list=parameter_list,
+            no_grad_set=no_grad_set,
+        )
+
+        return optimize_ops, params_grads

python/paddle/distributed/fleet/meta_optimizers/common.py

@@ -0,0 +1,236 @@
+# Copyright (c) 2020 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.framework import core
+from paddle.utils import unique_name
+
+from ..base.private_helper_function import wait_server_ready
+
+__all__ = []
+
+OpRole = core.op_proto_and_checker_maker.OpRole
+
+OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
+OP_ROLE_VAR_KEY = core.op_proto_and_checker_maker.kOpRoleVarAttrName()
+
+
+def is_update_op(op):
+    return (
+        'Param' in op.input_names
+        and 'Grad' in op.input_names
+        and "LearningRate" in op.input_names
+    )
+
+
+def is_loss_grad_op(op):
+    if OP_ROLE_KEY not in op.attr_names:
+        return False
+    op_role = int(op.all_attrs()[OP_ROLE_KEY])
+    return op_role & int(OpRole.Backward) and op_role & int(OpRole.Loss)
+
+
+def is_backward_op(op):
+    return OP_ROLE_KEY in op.attr_names and int(
+        op.all_attrs()[OP_ROLE_KEY]
+    ) & int(OpRole.Backward)
+
+
+def is_optimizer_op(op):
+    return OP_ROLE_KEY in op.attr_names and int(
+        op.all_attrs()[OP_ROLE_KEY]
+    ) & int(OpRole.Optimize)
+
+
+class CollectiveHelper:
+    def __init__(self, role_maker, nrings=1, wait_port=True):
+        self.nrings = nrings
+        self.wait_port = wait_port
+        self.role_maker = role_maker
+
+    def update_startup_program(self, startup_program=None):
+        self.startup_program = startup_program
+        if startup_program is None:
+            self.startup_program = paddle.static.default_startup_program()
+
+        endpoints = self.role_maker._get_trainer_endpoints()
+        current_endpoint = endpoints[self.role_maker._worker_index()]
+        for ring_id in range(self.nrings):
+            self._init_communicator(
+                self.startup_program,
+                current_endpoint,
+                endpoints,
+                self.role_maker._worker_index(),
+                ring_id,
+                self.wait_port,
+            )
+        self._broadcast_params()
+
+    def _init_communicator(
+        self,
+        program,
+        current_endpoint,
+        endpoints,
+        rank,
+        ring_id,
+        wait_port,
+        global_ring_id=None,
+        sync=True,
+    ):
+        # if current_endpoint is None, it means just for sync,
+        # no group is created.
+        endpoints_str = ",".join(endpoints)
+        if current_endpoint:
+            nranks = len(endpoints)
+            other_endpoints = endpoints[:]
+            other_endpoints.remove(current_endpoint)
+
+        def _add_sync_by_allreduce(block):
+            sync_var = block.create_var(
+                name=unique_name.generate('sync_var'),
+                dtype=core.VarDesc.VarType.INT32,
+                persistable=False,
+                stop_gradient=True,
+            )
+            block.append_op(
+                type='fill_constant',
+                inputs={},
+                outputs={'Out': [sync_var]},
+                attrs={
+                    'shape': [1],
+                    'dtype': sync_var.dtype,
+                    'value': 1,
+                    'force_cpu': False,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='all_reduce',
+                inputs={'x': [sync_var]},
+                outputs={'out': [sync_var]},
+                attrs={
+                    'ring_id': global_ring_id,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='c_sync_calc_stream',
+                inputs={'X': sync_var},
+                outputs={'Out': sync_var},
+                attrs={OP_ROLE_KEY: OpRole.Forward},
+            )
+
+        block = program.global_block()
+        if current_endpoint is None:
+            assert endpoints is None
+            assert sync
+            _add_sync_by_allreduce(block)
+            return
+
+        comm_id_var = block.create_var(
+            name=unique_name.generate('comm_id'),
+            persistable=True,
+            type=core.VarDesc.VarType.RAW,
+        )
+        if core.is_compiled_with_cuda():
+            block.append_op(
+                type='c_gen_nccl_id',
+                inputs={},
+                outputs={'Out': comm_id_var},
+                attrs={
+                    'rank': rank,
+                    'endpoint': current_endpoint,
+                    'other_endpoints': other_endpoints,
+                    'ring_id': ring_id,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='c_comm_init',
+                inputs={'X': comm_id_var},
+                outputs={},
+                attrs={
+                    'nranks': nranks,
+                    'rank': rank,
+                    'ring_id': ring_id,
+                    'endpoints': endpoints_str,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+        elif core.is_compiled_with_xpu():
+            block.append_op(
+                type='c_gen_bkcl_id',
+                inputs={},
+                outputs={'Out': comm_id_var},
+                attrs={
+                    'rank': rank,
+                    'endpoint': current_endpoint,
+                    'other_endpoints': other_endpoints,
+                    'ring_id': ring_id,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='c_comm_init',
+                inputs={'X': comm_id_var},
+                outputs={},
+                attrs={
+                    'nranks': nranks,
+                    'rank': rank,
+                    'ring_id': ring_id,
+                    'endpoints': endpoints_str,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+        else:
+            raise ValueError(
+                "comm_id must be generated in paddlepaddle-xpu or paddlepaddle-xpu."
+            )
+        if sync:
+            _add_sync_by_allreduce(block)
+
+    def _wait(self, current_endpoint, endpoints):
+        assert self.wait_port
+        other_endpoints = endpoints[:]
+        other_endpoints.remove(current_endpoint)
+        wait_server_ready(other_endpoints)
+
+    def _broadcast_params(self):
+        block = self.startup_program.global_block()
+        ring_id = -1
+        for param in block.iter_parameters():
+            if param.is_distributed:
+                continue
+
+            ring_id = (ring_id + 1) % self.nrings
+            block.append_op(
+                type='broadcast',
+                inputs={'x': param},
+                outputs={'out': param},
+                attrs={
+                    'ring_id': ring_id,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+        for ring_id in range(self.nrings):
+            block.append_op(
+                type='c_sync_comm_stream',
+                inputs={'X': param},
+                outputs={'Out': param},
+                attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Forward},
+            )

python/paddle/distributed/fleet/meta_optimizers/dgc_optimizer.py

@@ -0,0 +1,595 @@
+#   Copyright (c) 2020 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
+
+import logging
+from functools import reduce
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+import paddle
+from paddle.base import framework
+from paddle.base.dygraph import base as imperative_base
+from paddle.common_ops_import import LayerHelper
+from paddle.framework import core, in_dynamic_mode
+from paddle.nn.clip import ClipGradByNorm, append_gradient_clip_ops
+from paddle.optimizer import Momentum, Optimizer
+from paddle.regularizer import L1Decay, L2Decay
+from paddle.static import create_global_var
+
+
+class DGCMomentumOptimizer(Optimizer):
+    _u_velocity_acc_str = "_dgc_u_"
+    _v_velocity_acc_str = "_dgc_v_"
+
+    def __init__(
+        self,
+        learning_rate,
+        momentum,
+        rampup_begin_step,
+        rampup_step=1,
+        sparsity=[0.999],
+        parameter_list=None,
+        use_nesterov=False,
+        num_trainers=None,
+        regularization=None,
+        grad_clip=None,
+        name=None,
+    ):
+        if in_dynamic_mode():
+            raise Exception("In dygraph, don't support DGCMomentumOptimizer.")
+
+        assert core.is_compiled_with_cuda(), (
+            "Paddle is not compiled with CUDA. DGC is only support GPU for now."
+        )
+
+        assert learning_rate is not None
+        assert momentum is not None
+        super().__init__(
+            learning_rate=learning_rate,
+            parameters=parameter_list,
+            weight_decay=regularization,
+            grad_clip=grad_clip,
+            name=name,
+        )
+        self.type = "dgc_momentum"
+        self._momentum = momentum
+        self._use_nesterov = bool(use_nesterov)
+
+        assert rampup_begin_step >= 0, "rampup_begin_step must >= 0"
+        self._rampup_begin_step = rampup_begin_step
+        self._rampup_step = rampup_step
+        self._sparsity = sparsity
+
+        self._rampup_begin_step_var = None
+        self._global_step_var = None
+
+        self._dgc_clip_norm = None
+        self._num_trainers = num_trainers
+        if grad_clip is not None:
+            if not isinstance(grad_clip, ClipGradByNorm):
+                raise TypeError(
+                    "The type of grad_clip should be 'ClipGradByNorm', because DGCMomentumOptimizer only support ClipGradByNorm"
+                )
+            assert isinstance(num_trainers, int), (
+                f"The type of num_trainers should be 'int', but received {type(num_trainers)}"
+            )
+            assert num_trainers > 0, (
+                "The value of num_trainers should be greater than 0!"
+            )
+
+            self._dgc_clip_norm = grad_clip.clip_norm * (num_trainers**-0.5)
+
+        self.regular_type, self.regular_coeff = self._get_regularization_param(
+            self.regularization
+        )
+
+    def _get_regularization_param(self, regularization):
+        regular_type = 0
+        regular_coeff = 0.0
+
+        if regularization is not None:
+            regular_coeff = regularization._coeff
+
+            if isinstance(regularization, L1Decay):
+                regular_type = 1
+            elif isinstance(regularization, L2Decay):
+                regular_type = 2
+            else:
+                raise AssertionError(
+                    "regularization must be None|L1Decay|L2Deacy"
+                )
+        return regular_type, regular_coeff
+
+    def _is_use_dgc(self, param_var, grad_var):
+        var_numel = abs(reduce(lambda x, y: x * y, param_var.shape, 1))
+        if (
+            var_numel < 16384
+            or param_var.type == core.VarDesc.VarType.SELECTED_ROWS
+            or grad_var.type == core.VarDesc.VarType.SELECTED_ROWS
+            or param_var.dtype != core.VarDesc.VarType.FP32
+        ):
+            return False
+        return True
+
+    def _append_optimize_op(self, block, param_and_grad):
+        assert isinstance(block, paddle.framework.Block)
+        velocity_acc = self._get_accumulator(
+            self._u_velocity_acc_str, param_and_grad[0]
+        )
+        assert velocity_acc is not None
+
+        inputs = {
+            "Param": param_and_grad[0],
+            "Grad": param_and_grad[1],
+            "Velocity": velocity_acc,
+            "LearningRate": self._create_param_lr(param_and_grad),
+        }
+        outputs = {
+            "ParamOut": param_and_grad[0],
+            "VelocityOut": velocity_acc,
+        }
+        attrs = {"mu": self._momentum, "use_nesterov": self._use_nesterov}
+
+        if not self._is_use_dgc(param_and_grad[0], param_and_grad[1]):
+            type = "momentum"
+        else:
+            type = "dgc_momentum"
+            inputs.update(
+                {
+                    "current_step": self._global_step_var,
+                    "nranks": self._nranks_var,
+                }
+            )
+            outputs.update({'Grad_out': param_and_grad[1]})
+            attrs.update({"rampup_begin_step": float(self._rampup_begin_step)})
+
+        # create the dgc momentum optimize op
+        dgc_momentum_op = block.append_op(
+            type=type,
+            inputs=inputs,
+            outputs=outputs,
+            attrs=attrs,
+            stop_gradient=True,
+        )
+        return dgc_momentum_op
+
+    def _add_auto_increment_var(self, counter_name, begin, step=1):
+        helper = LayerHelper('global_step_counter')
+        counter, is_new_var = helper.create_or_get_global_variable(
+            name=counter_name, dtype='float32', shape=[1], persistable=True
+        )
+        if is_new_var:
+            helper.set_variable_initializer(
+                counter,
+                initializer=paddle.nn.initializer.ConstantInitializer(
+                    value=float(begin - 1), force_cpu=True
+                ),
+            )
+            helper.main_program.global_block()._prepend_op(
+                type='increment',
+                inputs={'X': [counter]},
+                outputs={'Out': [counter]},
+                attrs={'step': float(step)},
+                stop_gradient=True,
+            )
+            counter.stop_gradient = True
+
+        return counter
+
+    def _add_nranks_var(self, name, value=-1):
+        helper = LayerHelper('global_step_counter')
+        counter, is_new_var = helper.create_or_get_global_variable(
+            name=name, dtype='float32', shape=[1], persistable=True
+        )
+        if is_new_var:
+            helper.set_variable_initializer(
+                counter,
+                initializer=paddle.nn.initializer.ConstantInitializer(
+                    value=float(value), force_cpu=True
+                ),
+            )
+            counter.stop_gradient = True
+
+        return counter
+
+    def _append_dgc_ops(self, param_and_grads):
+        main_program = paddle.static.default_main_program()
+        main_program._enable_dgc = True
+
+        # step counter
+        self._global_step_var = self._add_auto_increment_var(
+            counter_name=core.dgc.kDGCCounterName(), begin=0
+        )
+
+        self._nranks_var = self._add_nranks_var(
+            name=core.dgc.kDGCNRanksName(), value=self._num_trainers
+        )
+
+        # rampup begin step var for all_reduce_op_handle
+        self._rampup_begin_step_var = create_global_var(
+            shape=[1],
+            dtype=core.VarDesc.VarType.FP32,
+            persistable=True,
+            name=core.dgc.kDGCRampUpBeginStepName(),
+            value=self._rampup_begin_step * 1.0,
+            force_cpu=True,
+        )
+
+        self.helper = LayerHelper(self.__class__.__name__)
+
+        for param_var, grad_var in param_and_grads:
+            # reuse velocity in dgc_op and dgc_momentum_op
+            u_var = self._add_accumulator(self._u_velocity_acc_str, param_var)
+
+            if not self._is_use_dgc(param_var, grad_var):
+                continue
+
+            v_var = self._add_accumulator(self._v_velocity_acc_str, param_var)
+
+            k_var = create_global_var(
+                shape=[1],
+                dtype=param_var.dtype,
+                persistable=True,
+                name=param_var.name + core.dgc.kDGCKName(),
+                value=0.0,
+                force_cpu=True,
+            )
+
+            encoded_var = create_global_var(
+                shape=[1],
+                dtype=param_var.dtype,
+                persistable=True,
+                name=param_var.name + core.dgc.kDGCEncodedName(),
+                value=0.0,
+                force_cpu=False,
+            )
+
+            gather_var = create_global_var(
+                shape=[1],
+                dtype=param_var.dtype,
+                persistable=True,
+                name=param_var.name + core.dgc.kDGCGatherName(),
+                value=0.0,
+                force_cpu=False,
+            )
+
+            # del back oprolevarname
+            op_maker = core.op_proto_and_checker_maker
+            backward = core.op_proto_and_checker_maker.OpRole.Backward
+            for op in main_program.global_block().ops:
+                if not self._is_the_backward_op(op):
+                    continue
+
+                var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
+                if param_var.name not in var_attr:
+                    continue
+
+                var_attr.remove(param_var.name)
+                var_attr.remove(grad_var.name)
+                if len(var_attr) > 1:
+                    op._set_attr(op_maker.kOpRoleVarAttrName(), var_attr)
+                else:
+                    op._remove_attr(op_maker.kOpRoleVarAttrName())
+
+            clip_var = grad_var
+            if self._dgc_clip_norm is not None:
+                clip_var = self._append_clip_norm(grad_var, self._dgc_clip_norm)
+            self._dgc_op(
+                param_var,
+                clip_var,
+                grad_var,
+                u_var,
+                v_var,
+                k_var,
+                encoded_var,
+                gather_var,
+            )
+
+    def _is_the_backward_op(self, op):
+        op_maker = core.op_proto_and_checker_maker
+        backward = core.op_proto_and_checker_maker.OpRole.Backward
+        if op_maker.kOpRoleVarAttrName() in op.attr_names and int(
+            op.all_attrs()[op_maker.kOpRoleAttrName()]
+        ) == int(backward):
+            return True
+        return False
+
+    def _clip_by_norm(self, x, max_norm, name=None):
+        args = {'x': x, 'max_norm': max_norm, 'name': name}
+
+        helper = LayerHelper("dgc_clip_by_norm_op", **args)
+
+        if name is None:
+            name = paddle.base.unique_name.generate_with_ignorable_key(
+                ".".join([helper.name, 'tmp'])
+            )
+
+        out = helper.create_variable(
+            type=x.type, name=name, dtype=x.dtype, persistable=False
+        )
+
+        helper.append_op(
+            type="dgc_clip_by_norm",
+            inputs={"X": x, "current_step": self._global_step_var},
+            attrs={
+                "max_norm": max_norm,
+                "rampup_begin_step": float(self._rampup_begin_step),
+            },
+            outputs={"Out": out},
+        )
+        return out
+
+    def _append_clip_norm(self, grad_var, clip_norm):
+        with grad_var.block.program._backward_role_guard():
+            return self._clip_by_norm(
+                x=grad_var, max_norm=clip_norm, name=grad_var.name
+            )
+
+    def _dgc_op(
+        self,
+        param_var,
+        clip_var,
+        grad_var,
+        u_var,
+        v_var,
+        k_var,
+        encoded_var,
+        gather_var,
+    ):
+        block = paddle.static.default_main_program().global_block()
+        op_maker = core.op_proto_and_checker_maker
+
+        regular_type = self.regular_type
+        regular_coeff = self.regular_coeff
+        # The regularizer of the Parameters have higher priority
+        if param_var.regularizer is not None:
+            regular_type, regular_coeff = self._get_regularization_param(
+                param_var.regularizer
+            )
+
+        dgc_op = block.append_op(
+            type="dgc",
+            inputs={
+                "U": u_var,
+                "V": v_var,
+                "Grad": clip_var,
+                "Param": param_var,
+                "current_step": self._global_step_var,
+                "nranks": self._nranks_var,
+            },
+            outputs={
+                "U_out": u_var,
+                "V_out": v_var,
+                "EncodeGrad": encoded_var,
+                "k": k_var,
+                "Grad_out": grad_var,
+                "GatherBuff": gather_var,
+            },
+            attrs={
+                "m": self._momentum,
+                "sparsity": self._sparsity,
+                "use_nesterov": self._use_nesterov,
+                "rampup_begin_step": float(self._rampup_begin_step),
+                "rampup_step": float(self._rampup_step),
+                "regular_coeff": float(regular_coeff),
+                "regular_type": int(regular_type),
+            },
+            stop_gradient=True,
+        )
+
+        backward = op_maker.OpRole.Backward
+        dgc_op._set_attr(op_maker.kOpRoleAttrName(), backward)
+        dgc_op._set_attr(
+            op_maker.kOpRoleVarAttrName(), [param_var.name, grad_var.name]
+        )
+
+    def _process_distribute_lookuptable(self, param_grads):
+        """
+        Because distribute lookup table only support SGD optimizer for now, not support
+        other optimizer and regularization, so we should find the table parameter out,
+        and avoid to add regularization and other op for it, and add sgd optimize op
+        for it independently.
+        :param param_grads(list((Var, Var))): list of (param, grad) pair.
+        :param loss: the loss variable.
+        :param startup_program: the startup program
+        """
+        from paddle.distributed.distribute_lookup_table import (
+            find_distributed_lookup_table,
+        )
+
+        program = framework.default_main_program()
+        global_block = framework.default_main_program().global_block()
+        table_name = find_distributed_lookup_table(program)
+        table_param = None
+        table_grad = None
+        new_param_grads = []
+        for p, g in param_grads:
+            if p.name == table_name:
+                if table_param is not None:
+                    raise RuntimeError(
+                        "multi dist table var found, only support one now!"
+                    )
+                table_param = p
+                table_grad = g
+            else:
+                new_param_grads.append((p, g))
+        sgd_op = None
+        if table_param is not None:
+            param_and_grad = [table_param, table_grad]
+            with (
+                table_param.block.program._optimized_guard(param_and_grad),
+                framework.name_scope("optimizer"),
+            ):
+                self._create_global_learning_rate()
+                # create the optimize op
+                sgd_op = global_block.append_op(
+                    type='sgd',
+                    inputs={
+                        "Param": table_param,
+                        "Grad": table_grad,
+                        "LearningRate": self._create_param_lr(param_and_grad),
+                    },
+                    outputs={"ParamOut": param_and_grad[0]},
+                )
+        return new_param_grads, (table_param, table_grad), sgd_op
+
+    @imperative_base.no_grad()
+    def apply_gradients(self, params_grads):
+        # Note: since we can't use all_reduce_op now,
+        # dgc_op should be the last op of one grad.
+        # Maybe need a grad allreduce pass.
+        self._append_dgc_ops(params_grads)
+
+        params_grads = sorted(params_grads, key=lambda x: x[0].name)
+        (
+            params_grads,
+            table_param_and_grad,
+            table_optimize_op,
+        ) = self._process_distribute_lookuptable(params_grads)
+
+        not_dgc_params_grads = []
+        dgc_params_grads = []
+        # DGC clip and regularization in optimizer.backward
+        for param, grad in params_grads:
+            if not self._is_use_dgc(param, grad):
+                not_dgc_params_grads.append((param, grad))
+            else:
+                dgc_params_grads.append((param, grad))
+
+        # 'optimizer(grad_clip)' or 'set_gradient_clip'
+        if self._grad_clip is not None:
+            not_dgc_params_grads = self._grad_clip(not_dgc_params_grads)
+        else:
+            not_dgc_params_grads = append_gradient_clip_ops(
+                not_dgc_params_grads
+            )
+
+        not_dgc_params_grads = self.append_regularization_ops(
+            not_dgc_params_grads, self.regularization
+        )
+
+        params_grads = not_dgc_params_grads + dgc_params_grads
+        params_grads = sorted(params_grads, key=lambda x: x[0].name)
+
+        optimize_ops = self._create_optimization_pass(params_grads)
+        if table_optimize_op is not None:
+            optimize_ops.append(table_optimize_op)
+            params_grads.append(table_param_and_grad)
+
+        return optimize_ops
+
+
+class DGCOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.dgc_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = []
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _init_dgc_opt(self):
+        if self.dgc_opt is not None:
+            return
+
+        opt = self.inner_opt
+
+        if not self.role_maker._is_collective:
+            return
+
+        if not isinstance(opt, Momentum):
+            return
+
+        configs = self.user_defined_strategy.dgc_configs
+        if len(configs['sparsity']) == 0:
+            # default is [0.999]
+            configs['sparsity'] = [0.999]
+
+        self.dgc_opt = DGCMomentumOptimizer(
+            learning_rate=opt._learning_rate,
+            momentum=opt._momentum,
+            rampup_begin_step=configs['rampup_begin_step'],
+            rampup_step=configs['rampup_step'],
+            sparsity=configs['sparsity'],
+            parameter_list=opt._parameter_list,
+            use_nesterov=opt._use_nesterov,
+            num_trainers=self.role_maker._worker_num(),
+            regularization=opt.regularization,
+            grad_clip=opt._grad_clip,
+            name=opt._name,
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.dgc:
+            if not isinstance(self.inner_opt, Momentum):
+                logging.warning("dgc only works on Momentum optimizer")
+                return False
+            if self.role_maker._worker_num() <= 1:
+                logging.warning("dgc only works on multi cards")
+                return False
+
+            return True
+
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.dgc = False
+        dist_strategy.dgc_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.dgc = True
+        dist_strategy.dgc_configs = {"rampup_begin_step": 0, "rampup_step": 1}
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        self._init_dgc_opt()
+        return self.dgc_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_gradients(self, params_grads):
+        self._init_dgc_opt()
+        return self.dgc_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        self._init_dgc_opt()
+        return self.dgc_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self._init_dgc_opt()
+        optimize_ops, params_grads = self.dgc_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/__init__.py

@@ -0,0 +1,18 @@
+#   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
+from .dygraph_sharding_optimizer import DygraphShardingOptimizer  # noqa: F401
+from .heter_parallel_optimizer import HeterParallelOptimizer  # noqa: F401
+from .hybrid_parallel_gradscaler import HybridParallelGradScaler  # noqa: F401
+from .hybrid_parallel_optimizer import HybridParallelOptimizer  # noqa: F401
+
+__all__ = []

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/heter_parallel_optimizer.py

@@ -0,0 +1,65 @@
+# 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.autograd as imperative_base
+from paddle import framework
+
+__all__ = []
+
+
+def _obtain_optimizer_parameters_list(optimizer):
+    if getattr(optimizer, '_param_groups', None) and isinstance(
+        optimizer._param_groups[0], dict
+    ):
+        parameters_list = []
+        for group in optimizer._param_groups:
+            for param in group['params']:
+                parameters_list.append(param)
+    else:
+        parameters_list = list(optimizer._parameter_list)
+
+    return parameters_list
+
+
+class HeterParallelOptimizer:
+    # adapter wrapper for optimizer
+    def __init__(self, optimizer, strategy):
+        self._inner_opt = optimizer
+        self._strategy = strategy
+
+        # NOTE(liubo48): In pure DataParallel mode,
+        # the gradient synchronization is achieved through reducer.
+
+    @imperative_base.no_grad()
+    @framework.dygraph_only
+    def step(self):
+        parameters_list = _obtain_optimizer_parameters_list(self._inner_opt)
+        self._inner_opt.step()
+
+    @imperative_base.no_grad()
+    def minimize(
+        self, loss, startup_program=None, parameters=None, no_grad_set=None
+    ):
+        # minimize does not support parameters in the form of param_group,
+        # so no need use _obtain_optimizer_parameters_list
+        parameter_list = (
+            parameters if parameters else self._inner_opt._parameter_list
+        )
+
+        return self._inner_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+
+    def __getattr__(self, item):
+        return getattr(self._inner_opt, item)

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_gradscaler.py

@@ -0,0 +1,85 @@
+# 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
+import paddle.autograd as imperative_base
+from paddle import _legacy_C_ops
+
+from ...base.topology import ParallelMode
+
+__all__ = []
+
+
+class HybridParallelGradScaler:
+    def __init__(self, scaler, hcg):
+        self._scaler = scaler
+        self._hcg = hcg
+        self._use_dp_mode = (
+            self._hcg.get_parallel_mode() == ParallelMode.DATA_PARALLEL
+        )
+
+    def scale(self, var):
+        return self._scaler.scale(var)
+
+    def minimize(self, optimizer, *args, **kwargs):
+        if not self._enable:
+            return optimizer.minimize(*args, **kwargs)
+
+        #  unscale the grad
+        self._unscale(optimizer)
+
+        optimize_ops, params_grads = (None, None)
+
+        if hasattr(optimizer, "_set_auxiliary_var"):
+            optimizer._set_auxiliary_var('found_inf', self._found_inf)
+            optimize_ops, params_grads = optimizer.minimize(*args, **kwargs)
+            # TODO: Fix to _cache_found_inf after PaddleNLP update
+            self._cache_found_inf = optimizer._get_auxiliary_var('found_inf')
+        else:
+            if self._found_inf:
+                self._cache_found_inf = True
+            else:
+                optimize_ops, params_grads = optimizer.minimize(*args, **kwargs)
+                self._cache_found_inf = False
+
+        if self._use_dynamic_loss_scaling:
+            self._update()
+
+        return optimize_ops, params_grads
+
+    @imperative_base.no_grad()
+    def _unscale(self, optimizer):
+        if not self._enable:
+            return
+        param_grads = [
+            param._grad_ivar()
+            for param in optimizer._parameter_list
+            if param._grad_ivar() is not None
+        ]
+        _legacy_C_ops.check_finite_and_unscale(
+            param_grads, self._scale, param_grads, self._found_inf
+        )
+        # allreduce_max found_inf in check_group
+        if not self._use_dp_mode:
+            self._found_inf = paddle.cast(self._found_inf, dtype="int32")
+            # TODO(shenliang03) Since the minimize call in the optimizer is
+            # after the grad scaler, check_finite needs to synchronize global
+            # information. In the future, we should use check_group
+            paddle.distributed.all_reduce(
+                self._found_inf, op=paddle.distributed.ReduceOp.MAX, group=None
+            )
+            self._found_inf = paddle.cast(self._found_inf, dtype="bool")
+
+    def __getattr__(self, item):
+        return getattr(self._scaler, item)

python/paddle/distributed/fleet/meta_optimizers/fp16_allreduce_optimizer.py

@@ -0,0 +1,158 @@
+#   Copyright (c) 2020 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
+
+import paddle
+from paddle.framework import core
+from paddle.utils import unique_name
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class FP16AllReduceOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+            "LocalSGDOptimizer",
+            "GradientMergeOptimizer",
+            "AdaptiveLocalSGDOptimizer",
+        ]
+        self.meta_optimizers_black_list = ["DGCOptimizer"]
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.fp16_allreduce:
+            return True
+
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.fp16_allreduce = False
+
+    def _enable_strategy(self, dist_strategy, context=None):
+        dist_strategy.fp16_allreduce = True
+
+    @staticmethod
+    def fp16_compression(param_and_grads):
+        """
+        Compress fp32 gradients to fp16 during allreduce.
+        """
+        op_maker = core.op_proto_and_checker_maker
+
+        new_param_and_grads = []  # param, grad, is_cast
+        # cast grad from fp32->fp16 before allreduce,
+        for param, grad in param_and_grads:
+            if grad is None or grad.dtype != core.VarDesc.VarType.FP32:
+                new_param_and_grads.append((param, grad, False))
+                continue
+
+            op = grad.op
+            block = grad.block
+            var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
+            if param.name not in var_attr:
+                new_param_and_grads.append((param, grad, False))
+                continue
+
+            # remove (param, grad) from op_role_var
+            var_attr.remove(param.name)
+            var_attr.remove(grad.name)
+            if len(var_attr) > 1:
+                op._set_attr(op_maker.kOpRoleVarAttrName(), var_attr)
+            else:
+                op._remove_attr(op_maker.kOpRoleVarAttrName())
+
+            new_grad = block.create_var(
+                name=unique_name.generate(grad.name + ".cast_fp16"),
+                dtype=core.VarDesc.VarType.FP16,
+                persistable=False,
+                stop_gradient=True,
+            )
+
+            with block.program._backward_role_guard():
+                cast_op = block.append_op(
+                    type="cast",
+                    inputs={"X": grad},
+                    outputs={"Out": new_grad},
+                    attrs={
+                        "in_dtype": core.VarDesc.VarType.FP32,
+                        "out_dtype": core.VarDesc.VarType.FP16,
+                    },
+                    stop_gradient=True,
+                )
+
+                backward = op_maker.OpRole.Backward
+                cast_op._set_attr(op_maker.kOpRoleAttrName(), backward)
+                cast_op._set_attr(
+                    op_maker.kOpRoleVarAttrName(), [param.name, new_grad.name]
+                )
+                new_grad.op = cast_op
+
+            new_param_and_grads.append((param, new_grad, True))
+
+        ret_param_and_grads = []
+        # cast grad from fp16->fp32 after allreduce.
+        # NOTE. Now we split fp16 compression into two for loops,
+        # if we do not separate them, fuse allreduce will wrong.
+        # This must be the problem of fuse allreduce pass, need
+        # fixed in future.
+        for param, grad, cast in new_param_and_grads:
+            if not cast:
+                ret_param_and_grads.append((param, grad))
+                continue
+
+            block = grad.block
+            new_grad = block.create_var(
+                name=unique_name.generate(grad.name + ".cast_fp32"),
+                dtype=core.VarDesc.VarType.FP32,
+                persistable=False,
+                stop_gradient=True,
+            )
+
+            with (
+                block.program._optimized_guard([param, grad]),
+                paddle.static.name_scope('fp16_allreduce'),
+            ):
+                cast_op = block.append_op(
+                    type="cast",
+                    inputs={"X": grad},
+                    outputs={"Out": new_grad},
+                    attrs={
+                        "in_dtype": core.VarDesc.VarType.FP16,
+                        "out_dtype": core.VarDesc.VarType.FP32,
+                    },
+                    stop_gradient=True,
+                )
+            ret_param_and_grads.append((param, new_grad))
+
+        return ret_param_and_grads
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        new_params_grads = self.fp16_compression(params_grads)
+        return self.inner_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=new_params_grads
+        )

python/paddle/distributed/fleet/meta_optimizers/gradient_merge_optimizer.py

@@ -0,0 +1,75 @@
+#   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
+
+from paddle.incubate.optimizer import GradientMergeOptimizer as GM
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class GradientMergeOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.wrapped_opt = None
+        self.meta_optimizers_white_list = [
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _init_wrapped_opt(self):
+        config = self.user_defined_strategy.gradient_merge_configs
+        self.wrapped_opt = GM(self.inner_opt)
+        self.wrapped_opt._set_k_steps(
+            self.user_defined_strategy.gradient_merge_configs["k_steps"]
+        )
+        self.wrapped_opt._set_avg(
+            self.user_defined_strategy.gradient_merge_configs["avg"]
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        can_apply = (
+            self.user_defined_strategy.gradient_merge
+        ) and self.user_defined_strategy.gradient_merge_configs["k_steps"] > 1
+        return can_apply
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.gradient_merge = False
+        dist_strategy.gradient_merge_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        # we currently do not support auto-enable GradientMerge
+        return
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self._init_wrapped_opt()
+        optimize_ops, params_grads = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

python/paddle/distributed/fleet/meta_optimizers/lamb_optimizer.py

@@ -0,0 +1,121 @@
+#   Copyright (c) 2020 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
+
+import logging
+
+import paddle
+from paddle.optimizer import Adam
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class LambOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.lamb_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = []
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+        opt = self.inner_opt
+        if not isinstance(opt, Adam):
+            return
+
+        configs = self.user_defined_strategy.lamb_configs
+        if len(configs['exclude_from_weight_decay']) == 0:
+            _exclude_from_weight_decay_fn = None
+        else:
+
+            def exclude_fn(param):
+                exclude_list = configs['exclude_from_weight_decay']
+                for name in exclude_list:
+                    if param.name.endswith(name):
+                        return True
+                return False
+
+            _exclude_from_weight_decay_fn = exclude_fn
+
+        self.lamb_opt = paddle.optimizer.Lamb(
+            learning_rate=opt._learning_rate,
+            lamb_weight_decay=configs['lamb_weight_decay'],
+            beta1=opt._beta1,
+            beta2=opt._beta2,
+            epsilon=opt._epsilon,
+            parameters=opt._parameter_list,
+            grad_clip=opt._grad_clip,
+            exclude_from_weight_decay_fn=_exclude_from_weight_decay_fn,
+            name=opt._name,
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.lamb:
+            if not isinstance(self.inner_opt, Adam):
+                logging.warning(
+                    f"lamb need the inner optimizer to be AdamOptimizer optimizer but got {self.inner_opt.type}."
+                )
+                return False
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.lamb = False
+        dist_strategy.lamb_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.lamb = True
+        dist_strategy.lamb_configs = {
+            "lamb_weight_decay": 0.01,
+            "exclude_from_weight_decay": [],
+        }
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self.lamb_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    # the following function will be used by AMP if both LARS and AMP are turn on together.
+    def apply_gradients(self, params_grads):
+        return self.lamb_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.lamb_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.lamb_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

python/paddle/distributed/fleet/meta_optimizers/lars_optimizer.py

@@ -0,0 +1,110 @@
+#   Copyright (c) 2020 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
+
+import logging
+
+from paddle.incubate.optimizer import LarsMomentumOptimizer
+from paddle.optimizer import Momentum
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class LarsOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.lars_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = []
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+        opt = self.inner_opt
+        if not isinstance(opt, Momentum):
+            return
+
+        configs = self.user_defined_strategy.lars_configs
+
+        self.lars_opt = LarsMomentumOptimizer(
+            learning_rate=opt._learning_rate,
+            momentum=opt._momentum,
+            lars_coeff=configs['lars_coeff'],
+            lars_weight_decay=configs['lars_weight_decay'],
+            parameter_list=opt._parameter_list,
+            regularization=opt.regularization,
+            grad_clip=opt._grad_clip,
+            name=opt._name,
+            exclude_from_weight_decay=configs['exclude_from_weight_decay'],
+            epsilon=configs['epsilon'],
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.lars:
+            if not isinstance(self.inner_opt, Momentum):
+                logging.warning(
+                    f"lars need the inner optimizer to be Momentum optimizer but got {self.inner_opt.type}."
+                )
+                return False
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.lars = False
+        dist_strategy.lars_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.lars = True
+        dist_strategy.lars_configs = {
+            "lars_coeff": 0.01,
+            "lars_weight_decay": 0.0005,
+        }
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self.lars_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    # the following function will be used by AMP if both LARS and AMP are turn on together.
+    def apply_gradients(self, params_grads):
+        return self.lars_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.lars_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.lars_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

python/paddle/distributed/fleet/meta_optimizers/localsgd_optimizer.py

@@ -0,0 +1,494 @@
+# Copyright (c) 2020 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.static import (
+    default_main_program,
+    default_startup_program,
+    program_guard,
+)
+
+from .common import OP_ROLE_KEY, CollectiveHelper, OpRole
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class LocalSGDOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = ['AMPOptimizer']
+        self.meta_optimizers_black_list = [
+            "AdaptiveLocalSGDOptimizer",
+        ]
+        self.snapshot_key = '@SNAPSHOT'
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if not self.user_defined_strategy.localsgd:
+            return False
+
+        if self.role_maker._worker_num() <= 1:
+            return False
+
+        return isinstance(
+            self.inner_opt,
+            (
+                paddle.optimizer.momentum.Momentum,
+                paddle.optimizer.sgd.SGD,
+            ),
+        )
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.localsgd = False
+        dist_strategy.localsgd_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.localsgd = True
+        dist_strategy.localsgd_configs = {"k_steps": 1, "begin_step": 1}
+
+    def snapshot_name(self, param_name):
+        return param_name + self.snapshot_key
+
+    def create_snapshot_vars(self, program):
+        block = program.global_block()
+
+        non_dist_params = []
+        for param in block.iter_parameters():
+            if not param.is_distributed:
+                non_dist_params.append(param)
+
+        p2s = []
+        for param in non_dist_params:
+            snapshot = block.create_var(
+                name=self.snapshot_name(param.name),
+                shape=param.shape,
+                persistable=True,
+                stop_gradient=True,
+                dtype=param.dtype,
+            )
+            p2s.append([param, snapshot])
+        return p2s
+
+    def init_snapshot_vars(self, startup_program, param2snapshot):
+        with program_guard(startup_program):
+            for param, snapshot in param2snapshot:
+                paddle.assign(param, snapshot)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        minimized = self.inner_opt.minimize(
+            loss, startup_program=startup_program
+        )
+
+        k_steps_value = self.user_defined_strategy.localsgd_configs['k_steps']
+        begin_step_value = self.user_defined_strategy.localsgd_configs[
+            'begin_step'
+        ]
+
+        if startup_program is None:
+            startup_program = default_startup_program()
+        main_block = loss.block
+
+        self.nrings = 2
+        collective_helper = CollectiveHelper(self.role_maker, self.nrings)
+        collective_helper.update_startup_program(startup_program)
+        p2s = self.create_snapshot_vars(startup_program)
+        self.init_snapshot_vars(startup_program, p2s)
+
+        p2s = self.create_snapshot_vars(main_block.program)
+        with program_guard(main_block.program, startup_program):
+            step = paddle.optimizer.lr.autoincreased_step_counter(begin=1)
+            k_steps = paddle.static.create_global_var(
+                name="k_steps",
+                shape=[1],
+                value=k_steps_value,
+                dtype='int64',
+                persistable=True,
+            )
+
+            begin_step = paddle.static.create_global_var(
+                name="begin_step",
+                shape=[1],
+                value=begin_step_value,
+                dtype='int64',
+                persistable=True,
+            )
+
+            last_step = paddle.static.create_global_var(
+                name="last_step",
+                shape=[1],
+                value=begin_step_value,
+                dtype='int64',
+                persistable=True,
+            )
+
+            def communicate():
+                sub_block = default_main_program().current_block()
+                ring_id = -1
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='c_sync_calc_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    ring_id = (ring_id + 1) % self.nrings
+                    sub_block.append_op(
+                        type='all_reduce',
+                        inputs={'x': [param]},
+                        outputs={'out': [param]},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for ring_id in range(self.nrings):
+                    sub_block.append_op(
+                        type='c_sync_comm_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={
+                            'ring_id': ring_id,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='scale',
+                        inputs={'X': [param]},
+                        outputs={'Out': [param]},
+                        attrs={
+                            'scale': 1.0 / self.role_maker._worker_num(),
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='assign',
+                        inputs={'X': [param]},
+                        outputs={'Out': [snapshot]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                paddle.assign(step, last_step)
+
+            def begin_localsgd():
+                paddle.static.nn.cond(step - last_step == k_steps, communicate)
+
+            paddle.static.nn.cond(
+                step > begin_step, begin_localsgd, communicate
+            )
+        return minimized
+
+
+class AdaptiveLocalSGDOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = ['AMPOptimizer']
+        self.meta_optimizers_black_list = [
+            "LocalSGDOptimizer",
+        ]
+        self.snapshot_key = '@SNAPSHOT'
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if not self.user_defined_strategy.adaptive_localsgd:
+            return False
+
+        if self.role_maker._worker_num() <= 1:
+            return False
+
+        return isinstance(
+            self.inner_opt,
+            (
+                paddle.optimizer.Momentum,
+                paddle.optimizer.sgd.SGD,
+            ),
+        )
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.adaptive_localsgd = False
+        dist_strategy.adaptive_localsgd_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.adaptive_localsgd = True
+        dist_strategy.adaptive_localsgd_configs = {
+            "init_k_steps": 1,
+            "begin_step": 1,
+        }
+
+    def snapshot_name(self, param_name):
+        return param_name + self.snapshot_key
+
+    def create_snapshot_vars(self, program):
+        block = program.global_block()
+
+        non_dist_params = []
+        for param in block.iter_parameters():
+            if not param.is_distributed:
+                non_dist_params.append(param)
+
+        p2s = []
+        for param in non_dist_params:
+            snapshot = block.create_var(
+                name=self.snapshot_name(param.name),
+                shape=param.shape,
+                persistable=True,
+                stop_gradient=True,
+                dtype=param.dtype,
+            )
+            p2s.append([param, snapshot])
+        return p2s
+
+    def init_snapshot_vars(self, startup_program, param2snapshot):
+        with program_guard(startup_program):
+            for param, snapshot in param2snapshot:
+                paddle.assign(param, snapshot)
+
+    def _generate_avg_loss(self, program_block, loss, avg_loss):
+        program_block.append_op(
+            type='all_reduce',
+            inputs={'x': [loss]},
+            outputs={'out': [avg_loss]},
+            attrs={
+                'ring_id': 0,
+                OP_ROLE_KEY: OpRole.Optimize,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
+            },
+        )
+        program_block.append_op(
+            type='c_sync_calc_stream',
+            inputs={'X': [avg_loss]},
+            outputs={'Out': [avg_loss]},
+            attrs={OP_ROLE_KEY: OpRole.Optimize},
+        )
+
+        program_block.append_op(
+            type='scale',
+            inputs={'X': [avg_loss]},
+            outputs={'Out': [avg_loss]},
+            attrs={
+                'scale': 1.0 / self.role_maker._worker_num(),
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        minimized = self.inner_opt.minimize(
+            loss, startup_program=startup_program
+        )
+
+        init_k_steps = self.user_defined_strategy.adaptive_localsgd_configs[
+            'init_k_steps'
+        ]
+        begin_step_value = self.user_defined_strategy.adaptive_localsgd_configs[
+            'begin_step'
+        ]
+
+        if startup_program is None:
+            startup_program = default_startup_program()
+        main_block = loss.block
+
+        self.nrings = 2
+        collective_helper = CollectiveHelper(self.role_maker, self.nrings)
+        collective_helper.update_startup_program(startup_program)
+        p2s = self.create_snapshot_vars(startup_program)
+        self.init_snapshot_vars(startup_program, p2s)
+
+        p2s = self.create_snapshot_vars(main_block.program)
+        with program_guard(main_block.program, startup_program):
+            step = paddle.optimizer.lr.autoincreased_step_counter(begin=1)
+
+            k_steps = paddle.static.create_global_var(
+                name="k_steps",
+                shape=[1],
+                value=int(init_k_steps),
+                dtype='int64',
+                persistable=True,
+            )
+
+            begin_step = paddle.static.create_global_var(
+                name="begin_step",
+                shape=[1],
+                value=int(begin_step_value),
+                dtype='int64',
+                persistable=True,
+            )
+
+            last_step = paddle.static.create_global_var(
+                name="last_step",
+                shape=[1],
+                value=0,
+                dtype='int64',
+                persistable=True,
+            )
+
+            avg_loss = paddle.static.create_global_var(
+                name="avg_loss",
+                shape=[1],
+                value=float(0),
+                dtype=loss.dtype,
+                persistable=True,
+            )
+
+            lr_0 = paddle.static.create_global_var(
+                name="lr_0",
+                shape=[1],
+                value=float(0),
+                dtype='float32',
+                persistable=True,
+            )
+
+            loss_0 = paddle.static.create_global_var(
+                name="loss_0",
+                shape=[1],
+                value=float(0),
+                dtype='float32',
+                persistable=True,
+            )
+
+            global_lr = self.inner_opt._global_learning_rate()
+
+            def initialize():
+                self._generate_avg_loss(main_block, loss, avg_loss)
+                paddle.assign(avg_loss, loss_0)
+                paddle.assign(global_lr, lr_0)
+
+            paddle.static.nn.cond(step == 1, initialize)
+
+            def communicate():
+                sub_block = default_main_program().current_block()
+                ring_id = -1
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='c_sync_calc_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    ring_id = (ring_id + 1) % self.nrings
+                    sub_block.append_op(
+                        type='all_reduce',
+                        inputs={'x': [param]},
+                        outputs={'out': [param]},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for ring_id in range(self.nrings):
+                    sub_block.append_op(
+                        type='c_sync_comm_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={
+                            'ring_id': ring_id,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='scale',
+                        inputs={'X': [param]},
+                        outputs={'Out': [param]},
+                        attrs={
+                            'scale': 1.0 / self.role_maker._worker_num(),
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='assign',
+                        inputs={'X': [param]},
+                        outputs={'Out': [snapshot]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                paddle.assign(step, last_step)
+
+            def communicate_avg_loss():
+                communicate()
+                self._generate_avg_loss(main_block, loss, avg_loss)
+
+                next_local_steps = paddle.cast(
+                    paddle.ceil(
+                        paddle.sqrt(
+                            lr_0
+                            * avg_loss
+                            / (global_lr * loss_0)
+                            * float(init_k_steps)
+                        )
+                    ),
+                    dtype='int64',
+                )
+                max_local_steps = paddle.full(
+                    shape=[1], dtype='int64', fill_value=16
+                )
+                min_local_steps = paddle.full(
+                    shape=[1], dtype='int64', fill_value=1
+                )
+                next_local_steps = paddle.minimum(
+                    next_local_steps, max_local_steps
+                )
+                next_local_steps = paddle.maximum(
+                    next_local_steps, min_local_steps
+                )
+                paddle.assign(next_local_steps, k_steps)
+
+            def begin_localsgd():
+                paddle.static.nn.cond(
+                    step - last_step == k_steps, communicate_avg_loss
+                )
+
+            paddle.static.nn.cond(
+                step > begin_step, begin_localsgd, communicate
+            )
+
+        return minimized

python/paddle/distributed/fleet/meta_optimizers/meta_optimizer_base.py

@@ -0,0 +1,106 @@
+#   Copyright (c) 2020 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.optimizer import Optimizer
+
+__all__ = []
+
+
+class MetaOptimizerBase(Optimizer):
+    def __init__(self, optimizer):
+        self.inner_opt = optimizer
+        self._learning_rate = self.inner_opt._learning_rate
+        self._learning_rate_map = self.inner_opt._learning_rate_map
+        self.meta_optimizers_white_list = []
+        self.meta_optimizers_black_list = []
+
+    def _set_auxiliary_var(self, key, val):
+        super()._set_auxiliary_var(key, val)
+        self.inner_opt._set_auxiliary_var(key, val)
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        self.loss = loss
+        self.role_maker = role_maker
+        self.user_defined_optimizer = user_defined_optimizer
+        self.user_defined_strategy = user_defined_strategy
+
+    def _update_inner_optimizer(self, optimizer):
+        self.inner_opt = optimizer
+
+    def _can_apply(self):
+        return False
+
+    def _is_graph_out(self):
+        return False
+
+    def _can_update(self, optimizer):
+        if str(optimizer.__class__.__name__) in self.meta_optimizers_white_list:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        raise NotImplementedError(
+            f"you should implement disable strategy in {type(self).__name__}"
+        )
+
+    def _enable_strategy(self, dist_strategy, context=None):
+        raise NotImplementedError(
+            f"you should implement enable strategy in {type(self).__name__}"
+        )
+
+    def apply_gradients(self, params_grads):
+        return self.inner_opt.apply_gradients(params_grads=params_grads)
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self.inner_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.inner_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        params_grads = self.backward(
+            loss,
+            startup_program=startup_program,
+            parameter_list=parameter_list,
+            no_grad_set=no_grad_set,
+        )
+
+        optimize_ops = self.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+        return optimize_ops, params_grads
+
+    def minimize(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.minimize_impl(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

python/paddle/distributed/fleet/meta_optimizers/parameter_server_optimizer.py

@@ -0,0 +1,450 @@
+#   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
+
+import os
+import platform
+import re
+import subprocess
+
+import paddle
+from paddle.framework import core
+
+from ..base.private_helper_function import wait_server_ready
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class ParameterServerOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+        # self.micro_batch_size = user_defined_strategy.pipeline_configs[
+        #    'micro_batch_size']
+        self.num_microbatches = user_defined_strategy.pipeline_configs[
+            'accumulate_steps'
+        ]
+
+    def _is_graph_out(self):
+        return False
+
+    def _can_apply(self):
+        if self.role_maker._is_collective:
+            return False
+
+        k_steps = self.user_defined_strategy.a_sync_configs["k_steps"]
+        return True if k_steps >= 0 else False
+
+    def get_dist_env(self):
+        trainer_id = int(os.getenv('PADDLE_TRAINER_ID', '0'))
+        trainer_endpoints = ''
+        current_endpoint = ''
+        num_trainers = 0
+        if os.getenv('PADDLE_TRAINER_ENDPOINTS'):
+            trainer_endpoints = os.getenv('PADDLE_TRAINER_ENDPOINTS')
+            current_endpoint = trainer_endpoints.split(',')[trainer_id]
+            num_trainers = len(trainer_endpoints.split(','))
+
+        return {
+            'trainer_id': trainer_id,
+            'num_trainers': num_trainers,
+            'current_endpoint': current_endpoint,
+            'trainer_endpoints': trainer_endpoints,
+        }
+
+    def _get_distributed_strategy(self):
+        from paddle.incubate.distributed.fleet.parameter_server.distribute_transpiler.distributed_strategy import (
+            StrategyFactory,
+        )
+
+        k_steps = self.user_defined_strategy.a_sync_configs["k_steps"]
+        strategy = None
+
+        if not self.user_defined_strategy.a_sync and k_steps == 0:
+            strategy = StrategyFactory.create_sync_strategy()
+
+        if self.user_defined_strategy.a_sync and k_steps == 0:
+            strategy = StrategyFactory.create_async_strategy()
+
+        if self.user_defined_strategy.a_sync and k_steps > 0:
+            strategy = StrategyFactory.create_geo_strategy(k_steps)
+
+        if not strategy:
+            raise ValueError("k_steps must be invalid value, please check")
+
+        return strategy
+
+    def _build_trainer_programs(self, compiled_config):
+        from paddle.incubate.distributed.fleet.parameter_server.ir import (
+            trainer_pass as worker,
+        )
+
+        _main = compiled_config.origin_main_program.clone()
+        _startup = compiled_config.origin_startup_program.clone()
+
+        use_ps_gpu = self.user_defined_strategy.a_sync_configs["use_ps_gpu"]
+
+        if not compiled_config.is_geo_mode():
+            from paddle.incubate.distributed.fleet.parameter_server.ir.public import (
+                _add_lr_decay_table_pass,
+            )
+
+            _add_lr_decay_table_pass(
+                _main,
+                compiled_config,
+                self.user_defined_strategy.a_sync_configs["lr_decay_steps"],
+            )
+
+            # for main program
+            _main = worker.distributed_ops_pass(
+                _main, compiled_config, use_ps_gpu
+            )
+            if not use_ps_gpu:
+                _main = worker.delete_optimizer_pass(_main, compiled_config)
+                _main = worker.append_send_ops_pass(_main, compiled_config)
+                _startup = worker.delete_extra_optimizes_pass(
+                    _startup, compiled_config
+                )
+
+                # for startup program
+            _startup = worker.fake_init_ops_pass(_startup, compiled_config)
+            if use_ps_gpu:
+                _main = worker.ps_gpu_pass(_main)
+                from paddle.distributed.transpiler.collective import (
+                    SingleProcessMultiThread,
+                )
+
+                t = SingleProcessMultiThread()
+                env = self.get_dist_env()
+                t.transpile(
+                    startup_program=_startup,
+                    main_program=_main,
+                    rank=env["trainer_id"],
+                    endpoints=env["trainer_endpoints"],
+                    current_endpoint=env['current_endpoint'],
+                    wait_port=False,
+                )
+
+            compiled_config.set_origin_ps_main_program(_main)
+            compiled_config.set_origin_ps_startup_program(_startup)
+            # for heter program
+            if self.role_maker._is_heter_parameter_server_mode:
+                from paddle.incubate.distributed.fleet.parameter_server.ir import (
+                    heter_trainer_pass as heter_worker,
+                )
+
+                if self.role_maker._is_heter_worker():
+                    # for heter worker
+                    stage_id = self.role_maker._get_stage_id()
+                    device = self.role_maker._heter_device_type().lower()
+                    _main = heter_worker.split_heter_worker_ops_pass(
+                        _main, compiled_config, stage_id, device
+                    )
+                else:
+                    # for default worker
+                    _main = heter_worker.split_trainer_ops_pass(
+                        _main, compiled_config
+                    )
+        else:
+            _main = worker.append_send_ops_pass(_main, compiled_config)
+            _startup = _startup
+            compiled_config.set_origin_ps_main_program(_main)
+            compiled_config.set_origin_ps_startup_program(_startup)
+
+        launch_barrier = self.user_defined_strategy.a_sync_configs[
+            "launch_barrier"
+        ]
+        launch_barrier_flag = int(os.getenv("FLAGS_LAUNCH_BARRIER", "1"))
+        if launch_barrier and launch_barrier_flag:
+            # for trainer wait server ready
+            wait_server_ready(self.role_maker._get_pserver_endpoints())
+
+            # for ps-heter mode, wait heter worker ready
+            # if self.role_maker._is_heter_parameter_server_mode and self.role_maker._is_worker(
+            # ):
+            #     wait_server_ready(self.role_maker._get_heter_worker_endpoints())
+
+        return _main, _startup
+
+    def _build_pserver_programs(self, compiled_config):
+        _main = paddle.static.Program()
+        _startup = paddle.static.Program()
+
+        from paddle.incubate.distributed.fleet.parameter_server.ir import (
+            pserver_pass as server,
+        )
+
+        if not compiled_config.is_geo_mode():
+            from paddle.incubate.distributed.fleet.parameter_server.ir.public import (
+                _get_optimize_ops,
+            )
+
+            is_sgd_adam = False
+
+            main_program = compiled_config.get_origin_main_program()
+            ops = _get_optimize_ops(main_program)
+
+            if len(ops) == 0:
+                return _main, _startup
+
+            from paddle.incubate.distributed.fleet.parameter_server.ir.public import (
+                _add_lr_decay_table_pass,
+            )
+
+            lr_decay_steps = self.user_defined_strategy.a_sync_configs[
+                "lr_decay_steps"
+            ]
+            _add_lr_decay_table_pass(
+                main_program, compiled_config, lr_decay_steps
+            )
+
+            for op in ops:
+                if op.type in ["sgd", "adam"]:
+                    is_sgd_adam = True
+                    break
+
+            if is_sgd_adam:
+                return _main, _startup
+
+            _main = server.add_listen_and_serv_pass(_main, compiled_config)
+            _main = server.add_rpc_global_flags_pass(_main, compiled_config)
+            _main = server.add_optimizer_pass(_main, compiled_config)
+            _main = server.large_scale_sparse_pass(
+                _main, _main, compiled_config, False
+            )
+            _startup = server.build_pserver_startup_program_pass(
+                _startup, _main, compiled_config
+            )
+            _startup = server.large_scale_sparse_pass(
+                _startup, _main, compiled_config, True
+            )
+
+            if not compiled_config.is_sync_mode():
+                _main = server.delete_unused_in_main_pass(
+                    _main, compiled_config
+                )
+
+            _startup = server.delete_unused_in_startup_pass(
+                _startup, _main, compiled_config
+            )
+        else:
+            _main = server.add_listen_and_serv_pass(_main, compiled_config)
+            _main = server.add_rpc_global_flags_pass(_main, compiled_config)
+            _main = server.add_geo_optimizer_pass(_main, compiled_config)
+            _startup = server.build_pserver_startup_program_pass(
+                _startup, _main, compiled_config
+            )
+            _startup = server.delete_unused_in_startup_pass(
+                _startup, _main, compiled_config
+            )
+
+        return _main, _startup
+
+    def _can_apply_geo(self, dist_strategy, program):
+        def get_sys_free_mem():
+            plat = platform.system()
+            if platform.system() == "Darwin":
+                vm = subprocess.Popen(
+                    ['vm_stat'], stdout=subprocess.PIPE
+                ).communicate()[0]
+                # Process vm_stat
+                vmLines = vm.split('\n')
+                sep = re.compile(r':[\s]+')
+                vmStats = {}
+                for row in range(1, len(vmLines) - 2):
+                    rowText = vmLines[row].strip()
+                    rowElements = sep.split(rowText)
+                    vmStats[(rowElements[0])] = (
+                        int(rowElements[1].strip(r'\.')) * 4096
+                    )
+                return vmStats["Pages free"]
+            elif platform.system() == "Linux":
+                mems = {}
+                with open('/proc/meminfo', 'rb') as f:
+                    for line in f:
+                        fields = line.split()
+                        mems[fields[0]] = int(fields[1]) * 1024
+                free = mems[b'MemFree:']
+                return free
+            else:
+                raise ValueError(
+                    f"{platform.system()} platform is unsupported is parameter server optimizer"
+                )
+
+        if not isinstance(self.inner_opt, paddle.optimizer.SGD):
+            return False
+
+        free = get_sys_free_mem()
+
+        from paddle.incubate.distributed.fleet.parameter_server.ir import (
+            vars_metatools,
+        )
+
+        processed_var_names = {"@EMPTY@"}
+        param_memory_size = 0
+        for varname in program.global_block().vars:
+            var = program.global_block().vars[varname]
+            if (
+                not var.persistable
+                or var.desc.type() != core.VarDesc.VarType.DENSE_TENSOR
+            ):
+                continue
+            param = vars_metatools.create_var_struct(var)
+            param_memory_size += param.m_size
+            processed_var_names.add(varname)
+
+        upper_mem_use = param_memory_size * 5.0
+
+        program_tmp_vars = {}
+        eval_batch_size = 1024
+        for op in program.global_block().ops:
+            for var_name in op.output_arg_names:
+                if var_name in processed_var_names:
+                    continue
+                processed_var_names.add(var_name)
+                var = program.global_block().vars[var_name]
+
+                if var.desc.type() != core.VarDesc.VarType.DENSE_TENSOR:
+                    continue
+
+                data_count = 1
+                neg_dim_count = 0
+                for x in var.shape:
+                    if x < 0:
+                        if neg_dim_count >= 1:
+                            raise ValueError(
+                                f"Var {var_name} has more than one negative dim."
+                            )
+                        neg_dim_count += 1
+                        data_count *= -x
+                    else:
+                        data_count *= x
+                program_tmp_vars[var_name] = (
+                    data_count,
+                    neg_dim_count,
+                    vars_metatools.dtype_to_size[var.dtype],
+                )
+
+        for varname in program_tmp_vars:
+            data_count, neg_dim_count, type_size = program_tmp_vars[varname]
+            if neg_dim_count == 1:
+                data_count *= eval_batch_size
+            var_memory = data_count * type_size
+            upper_mem_use += var_memory
+
+        if upper_mem_use < free:
+            return True
+        else:
+            return False
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self.inner_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        strategy = self._get_distributed_strategy()
+
+        _origin_main_program = loss.block.program
+        _origin_startup_program = startup_program
+        from paddle.incubate.distributed.fleet.parameter_server.ir import public
+
+        compiled_config = public.CompileTimeStrategy(
+            _origin_main_program,
+            _origin_startup_program,
+            strategy,
+            self.role_maker,
+        )
+        compiled_config.strategy = strategy
+
+        if self.role_maker._is_worker() or self.role_maker._is_heter_worker():
+            main_program, startup_program = self._build_trainer_programs(
+                compiled_config
+            )
+            if self.role_maker._is_heter_parameter_server_mode:
+                _origin_startup_program._heter_pipeline_opt = {
+                    "startup_program": startup_program,
+                    "pipeline_stage": int(self.role_maker._get_stage_id()) - 1,
+                    "heter_place": self.role_maker._heter_device(),
+                }
+
+                loss.block.program._heter_pipeline_opt = {
+                    "trainer": "HeterPipelineTrainer",
+                    "device_worker": "HeterSection",
+                    "trainers": self.role_maker._get_stage_trainers(),  # trainer num in each stage
+                    "trainer_id": int(self.role_maker._role_id()),
+                    "pipeline_stage": int(self.role_maker._get_stage_id()) - 1,
+                    "num_pipeline_stages": int(
+                        self.role_maker._get_num_stage()
+                    ),
+                    "section_program": main_program,
+                    "num_microbatches": self.num_microbatches,
+                    "heter_place": self.role_maker._heter_device(),
+                }
+            else:
+                loss.block.program = main_program
+                paddle.framework.switch_startup_program(startup_program)
+
+        elif self.role_maker._is_server():
+            main_program, startup_program = self._build_pserver_programs(
+                compiled_config
+            )
+            loss.block.program = main_program
+            paddle.framework.switch_startup_program(startup_program)
+        return None, None
+
+    def _disable_strategy(self, dist_strategy):
+        # if self.role_maker._is_heter_parameter_server_mode:
+        #    dist_strategy.pipeline = False
+        #    dist_strategy.pipeline_configs = {
+        #        "micro_batch_size": 1,
+        #        "accumulate_steps": 1,
+        #    }
+        dist_strategy.a_sync = False
+        a_sync_configs = dist_strategy.a_sync_configs
+        a_sync_configs["k_steps"] = -1
+        dist_strategy.a_sync_configs = a_sync_configs
+
+    def _enable_strategy(self, dist_strategy, context):
+        # if self.role_maker._is_heter_parameter_server_mode:
+        #    dist_strategy.pipeline = True
+        #    dist_strategy.pipeline_configs = {
+        #        "micro_batch_size": 1,
+        #        "accumulate_steps": 1,
+        #    }
+        a_sync_configs = dist_strategy.a_sync_configs
+        if a_sync_configs["k_steps"] >= 0:
+            return
+
+        dist_strategy.a_sync = True
+        a_sync_configs = dist_strategy.a_sync_configs
+
+        is_geo = self._can_apply_geo(
+            dist_strategy, context["origin_main_program"]
+        )
+
+        if is_geo:
+            a_sync_configs["k_steps"] = 800
+        else:
+            a_sync_configs["k_steps"] = 0
+        dist_strategy.a_sync_configs = a_sync_configs

python/paddle/distributed/fleet/meta_optimizers/pipeline_optimizer.py

@@ -0,0 +1,319 @@
+#   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
+
+import paddle
+from paddle.incubate.optimizer import PipelineOptimizer as PO
+
+from .common import (
+    OP_ROLE_KEY,
+    OP_ROLE_VAR_KEY,
+    CollectiveHelper,
+    OpRole,
+    is_backward_op,
+    is_loss_grad_op,
+)
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class PipelineOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = [
+            "RecomputeOptimizer",
+            "AMPOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+        self.global_ring_id = 1
+        self.dp_ring_id = 2
+        self.start_pipeline_ring_id = 20  # Just a magic number
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+        self.micro_batch_size = user_defined_strategy.pipeline_configs[
+            'micro_batch_size'
+        ]
+        self.num_microbatches = user_defined_strategy.pipeline_configs[
+            'accumulate_steps'
+        ]
+        self.schedule_mode = user_defined_strategy.pipeline_configs[
+            'schedule_mode'
+        ]
+        self.use_sharding = user_defined_strategy.sharding
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        # FIXME revise for hybrid parallelism
+        if self.use_sharding:
+            return False
+
+        if self.user_defined_strategy.pipeline:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.pipeline = False
+        dist_strategy.pipeline_configs = {
+            "micro_batch_size": 1,
+            "accumulate_steps": 1,
+            "schedule_mode": "1F1B",
+        }
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.pipeline = True
+        dist_strategy.pipeline_configs = {
+            "micro_batch_size": 1,
+            "accumulate_steps": 1,
+            "schedule_mode": "1F1B",
+        }
+
+    def _broadcast_params(self, ring_id):
+        block = self.startup_program.global_block()
+        param = None
+        for param in block.iter_parameters():
+            if param.is_distributed:
+                continue
+
+            block.append_op(
+                type='broadcast',
+                inputs={'x': param},
+                outputs={'out': param},
+                attrs={
+                    'ring_id': ring_id,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+        if not param:
+            return  # no parameter on this device
+        block.append_op(
+            type='c_sync_comm_stream',
+            inputs={'X': param},
+            outputs={'Out': param},
+            attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Forward},
+        )
+
+    def _get_process_group_info(self):
+        # global ring info
+        self.global_endpoints = self.endpoints
+        self.global_rank = self.rank
+        self.global_nranks = self.nranks
+
+        # data parallel ring info
+        if self.pipeline_num > 1:
+            self.dp_rank = self.rank // self.inner_parallelism
+            self.dp_nranks = self.nranks // self.inner_parallelism
+            start_index = self.rank % self.inner_parallelism
+            self.dp_endpoints = [
+                self.endpoints[start_index + i * self.inner_parallelism]
+                for i in range(self.pipeline_num)
+            ]
+
+    def _init_process_group(self, pipeline_pair, pipeline_ring_map):
+        self._get_process_group_info()
+        collective_helper = CollectiveHelper(self.role_maker, wait_port=False)
+        # Create global ring for all gpus (ring_id = 0)
+        collective_helper._init_communicator(
+            self.startup_program,
+            self.current_endpoint,
+            self.global_endpoints,
+            self.global_rank,
+            self.global_ring_id,
+            True,
+            self.global_ring_id,
+            True,
+        )
+        # Create pipeline rings
+        if self.inner_parallelism > 1:
+            pipeline_id = self.rank // self.inner_parallelism
+            start_index = pipeline_id * self.inner_parallelism
+            for pair in pipeline_pair:
+                pair_key = pair[0] * 1000 + pair[1]
+                ring_id = pipeline_ring_map[pair_key]
+                assert ring_id >= self.start_pipeline_ring_id
+                first_node = pair[0] + start_index
+                second_node = pair[1] + start_index
+                if self.rank != first_node and self.rank != second_node:
+                    collective_helper._init_communicator(
+                        self.startup_program,
+                        None,
+                        None,
+                        None,
+                        None,
+                        False,
+                        self.global_ring_id,
+                        True,
+                    )
+                    continue
+                pipeline_endpoints = [
+                    self.endpoints[first_node],
+                    self.endpoints[second_node],
+                ]
+                pipeline_rank = 0 if self.rank == first_node else 1
+                pipeline_nranks = 2
+                collective_helper._init_communicator(
+                    self.startup_program,
+                    self.current_endpoint,
+                    pipeline_endpoints,
+                    pipeline_rank,
+                    ring_id,
+                    False,
+                    self.global_ring_id,
+                    True,
+                )
+
+        # Create dp rings
+        if self.pipeline_num > 1:
+            collective_helper._init_communicator(
+                self.startup_program,
+                self.current_endpoint,
+                self.dp_endpoints,
+                self.dp_rank,
+                self.dp_ring_id,
+                True,
+                self.global_ring_id,
+                True,
+            )
+            self._broadcast_params(self.dp_ring_id)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self.endpoints = self.role_maker._get_trainer_endpoints()
+        self.current_endpoint = self.endpoints[self.role_maker._worker_index()]
+        self.rank = self.role_maker._worker_index()
+        self.nranks = self.role_maker._worker_num()
+
+        self.wrapped_opt = PO(
+            self.inner_opt, num_microbatches=self.num_microbatches
+        )
+        orig_startup_program = (
+            startup_program
+            if startup_program
+            else paddle.static.default_startup_program()
+        )
+        block = loss.block
+        program = block.program
+
+        program._pipeline_opt = {}
+        program._pipeline_opt['local_rank'] = self.rank
+        program._pipeline_opt['global_ring_id'] = self.global_ring_id
+        program._pipeline_opt['ring_id'] = self.start_pipeline_ring_id
+        program._pipeline_opt['micro_batch_size'] = self.micro_batch_size
+        program._pipeline_opt['schedule_mode'] = self.schedule_mode
+        program._pipeline_opt['use_sharding'] = False
+        program._pipeline_opt['mp_degree'] = 1
+        program._pipeline_opt['mp_rank'] = 0
+        (
+            optimize_ops,
+            params_grads,
+            prog_list,
+            pp_pair,
+            ring_map,
+        ) = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        self.startup_program = orig_startup_program._pipeline_opt[
+            'startup_program'
+        ]
+        self.inner_parallelism = program._pipeline_opt['inner_parallelism']
+        assert self.nranks % self.inner_parallelism == 0
+        assert prog_list
+        self.pipeline_num = len(self.endpoints) // self.inner_parallelism
+
+        self._init_process_group(pp_pair, ring_map)
+
+        self.main_program_list = prog_list
+        self.main_program = program
+        if self.pipeline_num > 1:
+            self._transpile_main_program(loss)
+        return optimize_ops, params_grads
+
+    def _transpile_main_program(self, loss):
+        self._insert_loss_grad_ops(loss, self.pipeline_num)
+        self._insert_allreduce_ops(self.dp_ring_id)
+
+    def _insert_loss_grad_ops(self, loss, pipeline_num):
+        """
+        In order to keep the learning rate consistent in different numbers of
+        training workers, we scale the loss grad by the number of workers
+        """
+        block = self.main_program_list[-1].global_block()
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_loss_grad_op(op):
+                loss_grad_var = block.vars[op.output_arg_names[0]]
+                block._insert_op(
+                    idx + 1,
+                    type='scale',
+                    inputs={'X': loss_grad_var},
+                    outputs={'Out': loss_grad_var},
+                    attrs={
+                        'scale': 1.0 / pipeline_num,
+                        OP_ROLE_KEY: OpRole.Backward,
+                    },
+                )
+
+    def _insert_allreduce_ops(self, ring_id):
+        block = self.main_program._pipeline_opt[
+            'section_program'
+        ].global_block()
+        origin_block = self.main_program.global_block()
+        grad = None
+        processed_param_name = set()
+        first_optimize_op_idx = None
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_backward_op(op) and not first_optimize_op_idx:
+                first_optimize_op_idx = idx + 1
+                # no optimize phase
+                if first_optimize_op_idx == len(block.ops):
+                    return
+            if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
+                op_role_var = op.all_attrs()[OP_ROLE_VAR_KEY]
+                if len(op_role_var) == 0:
+                    continue
+                assert len(op_role_var) % 2 == 0
+                offset = 0
+                for i in range(0, len(op_role_var), 2):
+                    param_name = op_role_var[i]
+                    param = block.vars[op_role_var[i]]
+                    if param_name in processed_param_name:
+                        continue
+                    processed_param_name.add(param_name)
+                    grad_name = op_role_var[i + 1]
+                    if 'MERGED' not in grad_name:
+                        grad_name += '@MERGED'
+                    grad = block.vars[grad_name]
+                    origin_param = origin_block.vars[op_role_var[i]]
+                    if origin_param.is_distributed:
+                        continue
+
+                    block._insert_op(
+                        first_optimize_op_idx + offset,
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )

python/paddle/distributed/fleet/meta_optimizers/ps_optimizer.py

@@ -0,0 +1,299 @@
+#   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
+
+import copy
+import os
+import platform
+import re
+import subprocess
+
+import paddle.distributed.passes
+from paddle.distributed.passes import PassContext
+from paddle.distributed.ps.utils.ps_factory import PsProgramBuilderFactory
+from paddle.distributed.ps.utils.public import (
+    TrainerRuntimeConfig,
+    build_var_distributed,
+    dtype_to_size,
+    get_dist_env,
+    get_var_mem_size,
+    logger,
+)
+from paddle.framework import core
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+
+class ParameterServerOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "RecomputeOptimizer",
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+            "ASPOptimizer",
+        ]
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _set_origin_programs(self, losses):
+        self.origin_main_programs = []
+        for loss in losses:
+            self.origin_main_programs.append(loss.block.program)
+
+    def _init_ps_pass_context(self, loss, startup_program):
+        self.pass_ctx = PassContext()
+        attrs = {}
+        # trainer
+        attrs["env"] = get_dist_env()
+
+        attrs['loss'] = loss
+        attrs['min_block_size'] = 81920
+        attrs['origin_main_program'] = loss.block.program
+        attrs['origin_startup_program'] = startup_program
+
+        attrs['origin_main_programs'] = self.origin_main_programs
+
+        attrs['cloned_main'] = attrs['origin_main_program'].clone()
+        attrs['cloned_startup'] = attrs['origin_startup_program'].clone()
+
+        attrs['user_defined_strategy'] = self.user_defined_strategy
+        attrs['valid_strategy'] = self.user_defined_strategy
+        attrs['trainer'] = TrainerRuntimeConfig(self.user_defined_strategy)
+        attrs['ps_mode'] = attrs['trainer'].mode
+        logger.info("ps_mode: {}".format(attrs['ps_mode']))
+        attrs['role_maker'] = self.role_maker
+        attrs['is_heter_ps_mode'] = (
+            self.role_maker._is_heter_parameter_server_mode
+        )
+        attrs['is_worker'] = self.role_maker._is_worker()
+        attrs['is_server'] = self.role_maker._is_server()
+        attrs['is_heter_worker'] = self.role_maker._is_heter_worker()
+        logger.info(
+            "this process is heter? {}".format(attrs['is_heter_worker'])
+        )
+        attrs['use_ps_gpu'] = self.user_defined_strategy.a_sync_configs[
+            "use_ps_gpu"
+        ]
+        attrs['use_gpu_graph'] = self.user_defined_strategy.a_sync_configs[
+            "use_gpu_graph"
+        ]
+        attrs['lr_decay_steps'] = self.user_defined_strategy.a_sync_configs[
+            "lr_decay_steps"
+        ]
+        # FL
+        attrs['local_sparse'] = attrs[
+            "user_defined_strategy"
+        ].trainer_desc_configs["local_sparse"]
+        attrs['remote_sparse'] = attrs[
+            "user_defined_strategy"
+        ].trainer_desc_configs["remote_sparse"]
+        attrs['is_fl_ps_mode'] = self.user_defined_strategy.is_fl_ps_mode
+        attrs['with_coordinator'] = (
+            self.user_defined_strategy.is_with_coordinator
+        )
+
+        attrs['k_steps'] = self.user_defined_strategy.a_sync_configs["k_steps"]
+        attrs['launch_barrier'] = self.user_defined_strategy.a_sync_configs[
+            "launch_barrier"
+        ]
+
+        attrs['launch_barrier_flag'] = int(
+            os.getenv("FLAGS_LAUNCH_BARRIER", "1")
+        )
+
+        build_var_distributed(attrs)
+
+        # server
+        attrs['_main_server'] = paddle.static.Program()
+        attrs['_startup_server'] = paddle.static.Program()
+        attrs['tensor_table'] = {}
+
+        self.pass_ctx._attrs = attrs
+
+    def _is_graph_out(self):
+        return False
+
+    def _can_apply(self):
+        if self.role_maker._is_collective:
+            return False
+
+        k_steps = self.user_defined_strategy.a_sync_configs["k_steps"]
+        return True if k_steps >= 0 else False
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.inner_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        if startup_program is None:
+            startup_program = paddle.static.default_startup_program()
+        #        print("program after inner optimizer minimize:",
+        #              str(loss.block.program))
+        self._set_origin_programs([loss])
+        self._init_ps_pass_context(loss, startup_program)
+        ps_builder = PsProgramBuilderFactory()._create_ps_program_builder(
+            self.pass_ctx
+        )
+        ps_builder._build_programs()
+        return optimize_ops, params_grads
+
+    def minimize_losses_impl(
+        self,
+        losses,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+    ):
+        self.inner_opts = [self.inner_opt]
+        for idx, loss in enumerate(losses):
+            if idx == 0:
+                continue
+            tmp_opt = copy.deepcopy(self.inner_opt)
+            self.inner_opts.append(tmp_opt)
+        if parameter_list is None:
+            parameter_list = [None] * len(losses)
+        for idx, loss in enumerate(losses):
+            startup_prog = startup_program[idx]
+            parameters = parameter_list[idx]
+            self.inner_opts[idx].minimize(
+                loss, startup_prog, parameters, no_grad_set
+            )
+        self._set_origin_programs(losses)
+        for idx, loss in enumerate(losses):
+            print("ps_optimizer idx loss:", idx, loss)
+            startup_prog = startup_program[idx]
+            self._init_ps_pass_context(loss, startup_prog)
+            ps_builder = PsProgramBuilderFactory()._create_ps_program_builder(
+                self.pass_ctx
+            )
+            ps_builder._build_programs()
+            startup_program[idx] = self.pass_ctx._attrs['cloned_startup']
+        return None, None
+
+    def _can_apply_geo(self, program):
+        def get_sys_free_mem():
+            plat = platform.system()
+            if platform.system() == "Darwin":
+                vm = subprocess.Popen(
+                    ['vm_stat'], stdout=subprocess.PIPE
+                ).communicate()[0]
+                # Process vm_stat
+                vmLines = vm.split('\n')
+                sep = re.compile(r':[\s]+')
+                vmStats = {}
+                for row in range(1, len(vmLines) - 2):
+                    rowText = vmLines[row].strip()
+                    rowElements = sep.split(rowText)
+                    vmStats[(rowElements[0])] = (
+                        int(rowElements[1].strip(r'\.')) * 4096
+                    )
+                return vmStats["Pages free"]
+            elif platform.system() == "Linux":
+                mems = {}
+                with open('/proc/meminfo', 'rb') as f:
+                    for line in f:
+                        fields = line.split()
+                        mems[fields[0]] = int(fields[1]) * 1024
+                free = mems[b'MemFree:']
+                return free
+            else:
+                raise ValueError(
+                    f"{platform.system()} platform is unsupported is parameter server optimizer"
+                )
+
+        if not isinstance(self.inner_opt, paddle.optimizer.SGD):
+            return False
+
+        free = get_sys_free_mem()
+        processed_var_names = {"@EMPTY@"}
+        param_memory_size = 0
+        for varname in program.global_block().vars:
+            var = program.global_block().vars[varname]
+            if (
+                not var.persistable
+                or var.desc.type() != core.VarDesc.VarType.DENSE_TENSOR
+            ):
+                continue
+            param_memory_size += get_var_mem_size(var)
+            processed_var_names.add(varname)
+
+        upper_mem_use = param_memory_size * 5.0
+
+        program_tmp_vars = {}
+        eval_batch_size = 1024
+        for op in program.global_block().ops:
+            for var_name in op.output_arg_names:
+                if var_name in processed_var_names:
+                    continue
+                processed_var_names.add(var_name)
+                var = program.global_block().vars[var_name]
+
+                if var.desc.type() != core.VarDesc.VarType.DENSE_TENSOR:
+                    continue
+
+                data_count = 1
+                neg_dim_count = 0
+                for x in var.shape:
+                    if x < 0:
+                        if neg_dim_count >= 1:
+                            raise ValueError(
+                                f"Var {var_name} has more than one negative dim."
+                            )
+                        neg_dim_count += 1
+                        data_count *= -x
+                    else:
+                        data_count *= x
+                program_tmp_vars[var_name] = (
+                    data_count,
+                    neg_dim_count,
+                    dtype_to_size[var.dtype],
+                )
+
+        for varname in program_tmp_vars:
+            data_count, neg_dim_count, type_size = program_tmp_vars[varname]
+            if neg_dim_count == 1:
+                data_count *= eval_batch_size
+            var_memory = data_count * type_size
+            upper_mem_use += var_memory
+
+        if upper_mem_use < free:
+            return True
+        else:
+            return False
+
+    def _enable_strategy(self, dist_strategy, context):
+        a_sync_configs = dist_strategy.a_sync_configs
+        if dist_strategy.a_sync_configs["k_steps"] >= 0:
+            return
+        dist_strategy.a_sync = True
+        a_sync_configs = dist_strategy.a_sync_configs
+
+        is_geo = self._can_apply_geo(context["origin_main_program"])
+
+        a_sync_configs["k_steps"] = 800 if is_geo else 0
+        dist_strategy.a_sync_configs = a_sync_configs
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.a_sync = False
+        a_sync_configs = dist_strategy.a_sync_configs
+        dist_strategy.a_sync_configs["k_steps"] = -1
+        dist_strategy.a_sync_configs = a_sync_configs

python/paddle/distributed/fleet/meta_optimizers/qat_optimizer.py

@@ -0,0 +1,123 @@
+# 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
+
+import copy
+
+import paddle
+from paddle.static.quantization.quanter import (
+    _quant_config_default,
+    quant_aware,
+)
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+
+class QATOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+            "GraphExecutionOptimizer",
+            "RecomputeOptimizer",
+            "GradientMergeOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.qat:
+            return True
+
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.qat = False
+        dist_strategy.qat_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.qat = True
+        dist_strategy.qat_configs = {
+            'channel_wise_abs_max': True,
+            'weight_bits': 8,
+            'activation_bits': 8,
+            'not_quant_pattern': [],
+            'algo': "",
+        }
+
+    def _gen_qat_config(self):
+        # Align the config to auto_parallel quantization pass
+        config = self.user_defined_strategy.qat_configs
+        qat_config = copy.deepcopy(_quant_config_default)
+        qat_config['quantize_op_types'] = [
+            'conv2d',
+            'depthwise_conv2d',
+            'mul',
+            'matmul',
+            'matmul_v2',
+        ]
+        qat_config['weight_quantize_type'] = (
+            'channel_wise_abs_max'
+            if config['channel_wise_abs_max']
+            else 'abs_max'
+        )
+        qat_config['weight_bits'] = config['weight_bits']
+        qat_config['activation_bits'] = config['activation_bits']
+        qat_config['not_quant_pattern'] = list(config['not_quant_pattern'])
+        return qat_config
+
+    def _replace_program(self, main_program, refer_program):
+        main_program._rebuild_from_desc(refer_program.desc)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.inner_opt.minimize(
+            loss,
+            startup_program,
+            parameter_list,
+            no_grad_set,
+        )
+        device = paddle.device.get_device()
+        place = paddle.set_device(device)
+        qat_config = self._gen_qat_config()
+        qat_program = quant_aware(
+            loss.block.program, place, config=qat_config, return_program=True
+        )
+        self._replace_program(loss.block.program, qat_program)
+        return optimize_ops, params_grads
+
+    def qat_init(self, place, scope=None, test_program=None):
+        if test_program is not None:
+            qat_config = self._gen_qat_config()
+            qat_program = quant_aware(
+                test_program,
+                place,
+                scope=scope,
+                config=qat_config,
+                for_test=True,
+                return_program=True,
+            )
+            self._replace_program(test_program, qat_program)

python/paddle/distributed/fleet/meta_optimizers/raw_program_optimizer.py

@@ -0,0 +1,566 @@
+#   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
+
+import os
+
+import paddle
+from paddle import static
+from paddle.base import core
+from paddle.framework.ir import apply_build_strategy
+from paddle.utils import unique_name
+
+from .common import (
+    OP_ROLE_KEY,
+    OP_ROLE_VAR_KEY,
+    CollectiveHelper,
+    OpRole,
+    is_backward_op,
+    is_loss_grad_op,
+    is_optimizer_op,
+)
+from .meta_optimizer_base import MetaOptimizerBase
+
+
+def evaluate_flag_apply_pass_to_program(val: str) -> bool:
+    val = val.lower()
+    if val in ('false', 'off', '0'):
+        return False
+    else:
+        return True
+
+
+class RawProgramOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = [
+            "RecomputeOptimizer",
+            "AMPOptimizer",
+            "GradientMergeOptimizer",
+            "LambOptimizer",
+            "LarsOptimizer",
+            "DGCOptimizer",
+            "LocalSGDOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+        self.global_ring_id = 0
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+        self.without_graph_optimization = (
+            user_defined_strategy.without_graph_optimization
+        )
+        self.fuse_all_reduce_ops = user_defined_strategy.fuse_all_reduce_ops
+        if self.fuse_all_reduce_ops:
+            self.fuse_grad_size_in_num = (
+                user_defined_strategy.fuse_grad_size_in_num
+            )
+            self.calc_comm_same_stream = (
+                user_defined_strategy._calc_comm_same_stream
+            )
+            self.sync_before_allreduce = os.environ.get(
+                'FLAGS_sync_before_allreduce', None
+            )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+        if self.user_defined_strategy.tensor_parallel:
+            return False
+        if self.user_defined_strategy.sharding:
+            return False
+
+        if self.without_graph_optimization:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.without_graph_optimization = False
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.without_graph_optimization = True
+
+    def _broadcast_params(self, ring_id):
+        block = self.startup_program.global_block()
+        param = None
+        for param in block.iter_parameters():
+            if param.is_distributed:
+                continue
+
+            block.append_op(
+                type='broadcast',
+                inputs={'x': param},
+                outputs={'out': param},
+                attrs={
+                    'ring_id': ring_id,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+        if not param:
+            return  # no parameter on this device
+        block.append_op(
+            type='c_sync_comm_stream',
+            inputs={'X': param},
+            outputs={'Out': param},
+            attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Forward},
+        )
+
+    def _get_process_group_info(self):
+        # global ring info
+        self.global_endpoints = self.endpoints
+        self.global_rank = self.rank
+        self.global_nranks = self.nranks
+
+    def _init_process_group(self):
+        self._get_process_group_info()
+        collective_helper = CollectiveHelper(self.role_maker, wait_port=False)
+        # Create global ring for all gpus (ring_id = 0)
+        collective_helper._init_communicator(
+            self.startup_program,
+            self.current_endpoint,
+            self.global_endpoints,
+            self.global_rank,
+            self.global_ring_id,
+            True,
+            self.global_ring_id,
+            True,
+        )
+        self._broadcast_params(self.global_ring_id)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self.endpoints = self.role_maker._get_trainer_endpoints()
+        self.current_endpoint = self.endpoints[self.role_maker._worker_index()]
+        self.rank = self.role_maker._worker_index()
+        self.nranks = self.role_maker._worker_num()
+        if startup_program is None:
+            startup_program = static.default_startup_program()
+        self.startup_program = startup_program
+
+        block = loss.block
+        program = block.program
+        self.main_program = program
+
+        optimize_ops, params_grads = self.inner_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        # Not apply pass only when FLAGS_apply_pass_to_program explicitly set to False
+        is_apply_pass_to_program = os.environ.get(
+            'FLAGS_apply_pass_to_program', '1'
+        )
+        if evaluate_flag_apply_pass_to_program(is_apply_pass_to_program):
+            pass_attrs = {"use_cuda": True}
+            build_strategy = self.user_defined_strategy.build_strategy._copy()
+            build_strategy.fuse_all_optimizer_ops = False
+            build_strategy.fuse_all_reduce_ops = False
+            apply_build_strategy(
+                self.main_program,
+                self.startup_program,
+                build_strategy,
+                pass_attrs,
+            )
+            self.main_program._pass_applied = True
+        if self.nranks == 1:
+            return optimize_ops, params_grads
+        self._init_process_group()
+
+        self.main_program = program
+        if self.nranks > 1:
+            self._transpile_main_program(loss)
+        return optimize_ops, params_grads
+
+    def _find_gradient_merge_block(self):
+        GRAD_MERGE_COND_NAME = "grad_merge_cond_name"
+        gm_cond_var_name = None
+        for op in self.main_program.global_block().ops:
+            if GRAD_MERGE_COND_NAME not in op.attr_names:
+                continue
+            if gm_cond_var_name is None:
+                gm_cond_var_name = op.attr(GRAD_MERGE_COND_NAME)
+            else:
+                assert gm_cond_var_name == op.attr(GRAD_MERGE_COND_NAME), (
+                    "multiple gradient merge condition found"
+                )
+        if gm_cond_var_name is None:
+            return None
+
+        cond_op = (
+            None  # false_fn of gm is None, so we should only find one block
+        )
+        for op in self.main_program.global_block().ops:
+            if op.type != 'conditional_block' or 'Cond' not in op.input_names:
+                continue
+            cond_vars = op.input('Cond')
+            if not cond_vars or cond_vars[0] != gm_cond_var_name:
+                continue
+            assert cond_op is None, "multiple gradient merge block found"
+            cond_op = op
+        assert cond_op is not None, "cannot find gradient merge block"
+        return cond_op._block_attr("sub_block")
+
+    def _insert_allreduce_ops_for_gm(self, gm_block):
+        block = self.main_program.global_block()
+
+        first_optimize_op_idx = None
+        for i, op in reversed(list(enumerate(gm_block.ops))):
+            if is_backward_op(op) and first_optimize_op_idx is None:
+                first_optimize_op_idx = i + 1
+                break
+        if first_optimize_op_idx is None:
+            first_optimize_op_idx = 0
+
+        param_vars = []
+        grad_vars = []
+        for op in block.ops:
+            if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
+                op_role_var = op.attr(OP_ROLE_VAR_KEY)
+                assert len(op_role_var) % 2 == 0
+                for i in range(0, len(op_role_var), 2):
+                    param = block.var(op_role_var[i])
+                    grad = block.var(op_role_var[i + 1])
+                    if param.is_distributed:
+                        continue
+                    param_vars.append(param)
+                    grad_vars.append(grad)
+
+        if not grad_vars:
+            return
+
+        gm_block._insert_op(
+            first_optimize_op_idx,
+            type="c_sync_calc_stream",
+            inputs={'X': grad_vars[0]},
+            outputs={'Out': grad_vars[0]},
+            attrs={OP_ROLE_KEY: OpRole.Backward},
+        )
+
+        insert_op_num = 1
+        ring_id = self.global_ring_id
+
+        # NOTE: can perform fuse allreduce inside the loop in the future
+        for i, (p, g) in enumerate(zip(param_vars, grad_vars)):
+            gm_block._insert_op(
+                first_optimize_op_idx + insert_op_num,
+                type="all_reduce",
+                inputs={'x': g},
+                outputs={'out': g},
+                attrs={
+                    'ring_id': ring_id,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
+                    OP_ROLE_KEY: OpRole.Backward,
+                },
+            )
+            insert_op_num += 1
+
+        gm_block._insert_op(
+            first_optimize_op_idx + insert_op_num,
+            type="c_sync_comm_stream",
+            inputs={'X': grad_vars},
+            outputs={'Out': grad_vars},
+            attrs={
+                'ring_id': ring_id,
+                OP_ROLE_KEY: OpRole.Backward,
+            },
+        )
+
+    def _transpile_main_program(self, loss):
+        self._insert_loss_grad_ops(loss)
+        gm_block = self._find_gradient_merge_block()
+        if gm_block is not None:
+            # TODO(zjl): support fuse allreduce
+            self._insert_allreduce_ops_for_gm(gm_block)
+            return
+
+        if self.fuse_all_reduce_ops and self.fuse_grad_size_in_num > 1:
+            self._allreduce_fusion_program()
+        else:
+            self._insert_allreduce_ops()
+
+    def _insert_loss_grad_ops(self, loss):
+        """
+        In order to keep the learning rate consistent in different numbers of
+        training workers, we scale the loss grad by the number of workers
+        """
+        block = self.main_program.global_block()
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_loss_grad_op(op):
+                loss_grad_var = block.vars[op.output_arg_names[0]]
+                block._insert_op(
+                    idx + 1,
+                    type='scale',
+                    inputs={'X': loss_grad_var},
+                    outputs={'Out': loss_grad_var},
+                    attrs={
+                        'scale': 1.0 / self.nranks,
+                        OP_ROLE_KEY: OpRole.Backward,
+                    },
+                )
+
+    def _insert_allreduce_ops(self):
+        block = self.main_program.global_block()
+        ring_id = self.global_ring_id
+        grad = None
+        grad_vars = []
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
+                op_role_var = op.attr(OP_ROLE_VAR_KEY)
+                if len(op_role_var) == 0:
+                    continue
+                assert len(op_role_var) % 2 == 0
+                offset = 1
+                for i in range(0, len(op_role_var), 2):
+                    param_name = op_role_var[i]
+                    param = block.var(param_name)
+                    grad_name = op_role_var[i + 1]
+                    grad = block.var(grad_name)
+                    if param.is_distributed:
+                        continue
+
+                    block._insert_op(
+                        idx + offset,
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Backward,
+                        },
+                    )
+
+        if grad is None:
+            return
+
+    # This function helps reduce the number of allreduce by integrating op, which can save communication time.
+    # to use allreduce fuse, follow these codes:
+    # strategy = paddle.distributed.fleet.DistributedStrategy()
+    # strategy.without_graph_optimization = True
+    # strategy.fuse_all_reduce_ops = True
+    # strategy.calc_comm_same_stream = False
+    # strategy.fuse_grad_size_in_num = 8
+    def _allreduce_fusion_program(self):
+        block = self.main_program.global_block()
+        ring_id = self.global_ring_id
+        param_grads = []
+        first_backward_idx = -1
+
+        # find all grad params
+        for idx, op in enumerate(block.ops):
+            if first_backward_idx == -1 and is_backward_op(op):
+                first_backward_idx = idx
+            if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
+                op_role_var = op.attr(OP_ROLE_VAR_KEY)
+                if len(op_role_var) == 0:
+                    continue
+                assert len(op_role_var) % 2 == 0, (
+                    "vars need to be one param var followed by one grad var, "
+                    "but got odd number of vars"
+                )
+                for i in range(0, len(op_role_var), 2):
+                    param_name = op_role_var[i]
+                    param = block.var(param_name)
+                    grad_name = op_role_var[i + 1]
+                    grad = block.var(grad_name)
+                    if param.is_distributed:
+                        continue
+                    param_grads.append((param, grad))
+
+        outputs_name_to_idx = self.__get_outputs_name_to_idx(
+            first_backward_idx, block
+        )
+
+        # structure of grad_param_segments is
+        # [([grad0, grad1], [param0, param1]), ([grad2, grad3], [param2, param3])]
+        # each entry of the list is a tuple stores the grads segment list and
+        # the corresponding params segment list
+
+        # its type is: dict[dtype, list[tuple[list[grad], list[param]]]]
+        grad_param_segments_by_dtype = {}
+        # split the grad based on dtype and fused size
+        for param, grad in param_grads:
+            if grad.dtype not in grad_param_segments_by_dtype:
+                grad_param_segments_by_dtype[grad.dtype] = [([], [])]
+            grad_segment, param_segment = grad_param_segments_by_dtype[
+                grad.dtype
+            ][-1]
+            if len(param_segment) == self.fuse_grad_size_in_num:
+                grad_param_segments_by_dtype[grad.dtype].append(([], []))
+                grad_segment, param_segment = grad_param_segments_by_dtype[
+                    grad.dtype
+                ][-1]
+            param_segment.append(param)
+            grad_segment.append(grad)
+
+        grad_param_segments = []
+        for _, group in grad_param_segments_by_dtype.items():
+            grad_param_segments.extend(group)
+
+        if len(grad_param_segments) == 0:
+            return
+
+        # because the regroup operation make the relative order invalid,
+        # we need to reorder these fuse group by after_idx
+        def get_after_idx_of_fuse_group(grad_param_segments):
+            grad_segment, param_segment = grad_param_segments
+            return max([outputs_name_to_idx[grad][1] for grad in grad_segment])
+
+        grad_param_segments.sort(key=get_after_idx_of_fuse_group)
+
+        fused_vars = [None] * len(grad_param_segments)
+        for i in range(len(grad_param_segments) - 1, -1, -1):
+            # travers the grad_param_segments in backward
+            # not to use reversed since needs the absolute index value
+            grad_segment, param_segment = grad_param_segments[i]
+            # insert coalesce tensor
+            fused_var = block.create_var(
+                name=unique_name.generate(
+                    f'FusedOutput_{grad_segment[0].name}'
+                ),
+                dtype=grad_segment[0].dtype,
+                persistable=False,
+                stop_gradient=True,
+            )
+            fused_vars[i] = fused_var
+            after_idx = max(
+                [outputs_name_to_idx[grad][1] for grad in grad_segment]
+            )
+            block._insert_op_without_sync(
+                after_idx + 1,
+                type='all_reduce',
+                inputs={'x': fused_var},
+                outputs={'out': fused_var},
+                attrs={
+                    'ring_id': ring_id,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
+                    OP_ROLE_KEY: OpRole.Backward,
+                },
+            )
+            if not self.calc_comm_same_stream and self.sync_before_allreduce:
+                block._insert_op_without_sync(
+                    after_idx + 1,
+                    type='c_sync_calc_stream',
+                    inputs={'X': fused_var},
+                    outputs={'Out': fused_var},
+                    attrs={OP_ROLE_KEY: OpRole.Backward},
+                )
+        idx = 0
+        if not self.calc_comm_same_stream and not self.sync_before_allreduce:
+            for i in range(len(grad_param_segments)):
+                while (
+                    block.ops[idx].type != 'c_allreduce_sum'
+                    and (
+                        not (
+                            block.ops[idx].type == 'all_reduce'
+                            and block.ops[idx].attr('reduce_type')
+                            == paddle.distributed.ReduceOp.SUM
+                        )
+                    )
+                ) or fused_vars[i].name not in block.ops[idx].input_arg_names:
+                    idx += 1
+                grad_segment, param_segment = grad_param_segments[i]
+                for grad in grad_segment:
+                    block._insert_op_without_sync(
+                        idx + 1,
+                        type='depend',
+                        inputs={'X': grad, 'Dep': fused_vars[i]},
+                        outputs={'Out': grad},
+                    )
+                    idx += 1
+
+        # update the outputs_name_to_idx after insertion of sync/allreduce ops
+        outputs_name_to_idx = self.__get_outputs_name_to_idx(
+            first_backward_idx, block
+        )
+        # the before_idx is not guaranteed sorted, therefore we have to find the
+        # topology to insert the coalesce ops
+        pos_for_coalesce = {}
+        for i in range(len(grad_param_segments) - 1, -1, -1):
+            # We separate the insertion of coalesce op and the insertion of sync/allreduce op,
+            # since that the coalesce op's insertion may invalidate the outputs_name_to_idx
+            grad_segment, param_segment = grad_param_segments[i]
+            before_idx = len(block.ops)
+            for grad in outputs_name_to_idx:
+                before_idx = min(before_idx, outputs_name_to_idx[grad][0])
+            pos_for_coalesce[i] = before_idx
+
+        # insert the coalesce op based on the sorted before_idx
+        pos_for_coalesce = sorted(
+            pos_for_coalesce.items(),
+            key=lambda kv: (kv[1], kv[0]),
+            reverse=True,
+        )
+        for i, before_idx in pos_for_coalesce:
+            grad_segment, param_segment = grad_param_segments[i]
+            fused_var = fused_vars[i]
+            block._insert_op_without_sync(
+                before_idx,
+                type="coalesce_tensor",
+                inputs={"Input": param_segment},
+                outputs={"Output": grad_segment, "FusedOutput": fused_var},
+                attrs={
+                    "copy_data": False,
+                    "use_align": True,
+                    "dtype": grad_segment[0].dtype,
+                    OP_ROLE_KEY: OpRole.Backward,
+                },
+            )
+
+        if self.calc_comm_same_stream or not self.sync_before_allreduce:
+            block._sync_with_cpp()
+            return
+
+        # insert the sync comm op
+        for idx, op in enumerate(block.ops):
+            if is_optimizer_op(op):
+                block._insert_op_without_sync(
+                    idx,
+                    type='c_sync_comm_stream',
+                    inputs={'X': fused_vars},
+                    outputs={'Out': fused_vars},
+                    attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Backward},
+                )
+                break
+        block._sync_with_cpp()
+
+    def __get_outputs_name_to_idx(self, first_backward_idx, block):
+        # Each item of outputs_name_to_idx is a pair of idx.
+        # The first entry of this pair is the idx of the first op generates the grad,
+        # which is used to indicate the position to insert coalesce op.
+        # The second entry of this pair is the idx of the last op generates the grad,
+        # which is used to indicate the position to insert sync and allreduce op.
+        outputs_name_to_idx = {}
+        for idx in range(first_backward_idx, len(block.ops)):
+            op = block.ops[idx]
+            if is_optimizer_op(op):
+                break
+            for name in op.output_arg_names:
+                if name == core.kEmptyVarName():
+                    continue
+                var = block.var(name)
+                if not outputs_name_to_idx.get(var):
+                    # if the grad only be generated by one op
+                    # the first idx and the last ids are identical
+                    outputs_name_to_idx[var] = (idx, idx)
+                else:
+                    outputs_name_to_idx[var] = (
+                        outputs_name_to_idx[var][0],
+                        idx,
+                    )
+        return outputs_name_to_idx

python/paddle/distributed/fleet/meta_optimizers/recompute_optimizer.py

@@ -0,0 +1,104 @@
+#   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
+
+from paddle.incubate.optimizer import RecomputeOptimizer as RO
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class RecomputeOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.wrapped_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "LarsOptimizer",
+            "LambOptimizer",
+            "DGCOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _init_wrapped_opt(self):
+        if self.wrapped_opt is not None:
+            return
+
+        configs = self.user_defined_strategy.recompute_configs
+        self.wrapped_opt = RO(self.inner_opt)
+        self.wrapped_opt._set_checkpoints(list(configs["checkpoints"]))
+        if configs["enable_offload"]:
+            self.wrapped_opt._enable_offload()
+            # TODO(JZ-LIANG) might found a way to infer the checkpoint shape automatically
+            checkpoint_shapes = list(configs["checkpoint_shape"])
+            self.wrapped_opt.checkpoint_shape = checkpoint_shapes
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.recompute:
+            if (
+                len(self.user_defined_strategy.recompute_configs["checkpoints"])
+                == 0
+            ):
+                return False
+            else:
+                return True
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.recompute = False
+        dist_strategy.recompute_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        # we do not support automatically recompute checkpoints currently
+        return
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        # maybe inner_opt of other meta optimizer
+        self._init_wrapped_opt()
+        return self.wrapped_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_gradients(self, params_grads):
+        return self.wrapped_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.wrapped_opt.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self._init_wrapped_opt()
+        optimize_ops, params_grads = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

python/paddle/distributed/fleet/meta_optimizers/sharding/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2020 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/fleet/meta_optimizers/sharding/fp16_helper.py

@@ -0,0 +1,273 @@
+# Copyright (c) 2020 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,
+    is_optimizer_op,
+)
+from paddle.framework import core
+
+__all__ = []
+
+
+class FP16Utils:
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def is_fp16_cast_op(block, op, params):
+        if op.type != "cast":
+            return False
+        if is_optimizer_op(op):
+            return False
+        assert len(op.desc.input_arg_names()) == 1
+        assert len(op.desc.output_arg_names()) == 1
+        input_name, output_name = (
+            op.desc.input_arg_names()[0],
+            op.desc.output_arg_names()[0],
+        )
+        if input_name not in params:
+            return False
+        input_var = block.var(input_name)
+        output_var = block.var(output_name)
+        if (
+            input_var.dtype != core.VarDesc.VarType.FP32
+            or output_var.dtype != core.VarDesc.VarType.FP16
+        ):
+            return False
+        return True
+
+    @staticmethod
+    def is_fp32_cast_op(block, op):
+        if op.type != "cast":
+            return False
+        if not is_optimizer_op(op):
+            return False
+        assert len(op.desc.input_arg_names()) == 1
+        assert len(op.desc.output_arg_names()) == 1
+        input_name, output_name = (
+            op.desc.input_arg_names()[0],
+            op.desc.output_arg_names()[0],
+        )
+        input_var = block.var(input_name)
+        output_var = block.var(output_name)
+        if (
+            input_var.dtype != core.VarDesc.VarType.FP16
+            or output_var.dtype != core.VarDesc.VarType.FP32
+        ):
+            return False
+        return True
+
+    @staticmethod
+    def remove_cast_op(block, params, segment, offset):
+        inserted_op_num = 0
+        for op_idx in reversed(
+            range(offset + segment._start_idx, offset + segment._end_idx)
+        ):
+            op = block.ops[op_idx]
+            if FP16Utils.is_fp16_cast_op(block, op, params):
+                block._remove_op(op_idx, sync=False)
+                inserted_op_num -= 1
+        block._sync_with_cpp()
+        return inserted_op_num
+
+    @staticmethod
+    def prune_fp16(block, shard, reduced_grads_to_param, ring_ids):
+        """
+        1. prune all cast_fp16_to_fp32 ops if the param not belongs to this shard
+        2. revise amp inifine grad checking for sharding
+        """
+        # remove cast
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not FP16Utils.is_fp32_cast_op(block, op):
+                continue
+            output_name = op.desc.output_arg_names()[0]
+            # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+            param_name = output_name.removesuffix("@MERGED").removesuffix(
+                "@GRAD"
+            )
+            if param_name not in shard.global_params:
+                raise ValueError(
+                    "Output 'X' of cast_op must be a grad of"
+                    f"model param, but {output_name} is not a grad"
+                )
+            if output_name in reduced_grads_to_param:
+                continue
+            if shard.has_param(param_name):
+                continue
+            block._remove_op(idx, sync=False)
+            block._remove_var(output_name, sync=False)
+
+        block._sync_with_cpp()
+        update_loss_scaling_op_idx = -1
+        inf_var_name = ''
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if op.type == "update_loss_scaling":
+                update_loss_scaling_op_idx = idx
+                inf_var_name = op.desc.input('FoundInfinite')[0]
+            if op.type in ["check_finite_and_unscale", "update_loss_scaling"]:
+                reversed_x = []
+                reversed_x_paramname = []
+                for input_name in op.desc.input('X'):
+                    # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+                    param_name = input_name.removesuffix(
+                        "@MERGED"
+                    ).removesuffix("@GRAD")
+                    if param_name not in shard.global_params:
+                        raise ValueError(
+                            "Input 'X' of check_finite_and_unscale must"
+                            f"be grads, but {input_name} is not a grad"
+                        )
+                    if shard.has_param(param_name):
+                        reversed_x.append(input_name)
+                        reversed_x_paramname.append(param_name)
+                op.desc.set_input('X', reversed_x)
+                op.desc.set_output('Out', reversed_x)
+
+                # the grad checking should take the all and only param in the current shard
+                to_check_param = set(reversed_x_paramname)
+                should_check_param = set(shard.global_params).intersection(
+                    {
+                        param
+                        for param, worker_idx in shard.global_param2device.items()
+                        if worker_idx == shard.worker_idx
+                    }
+                )
+                assert to_check_param == should_check_param, (
+                    f"amp \
+                    check_finite_and_unscale checking miss [{should_check_param - to_check_param}] and got unexpected [{to_check_param - should_check_param}]"
+                )
+
+        if update_loss_scaling_op_idx == -1:
+            return
+        inf_var = block.var(inf_var_name)
+        inf_var_int32 = block.create_var(
+            name=inf_var_name + "@cast_int32",
+            shape=inf_var.shape,
+            dtype=core.VarDesc.VarType.INT32,
+        )
+
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            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,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+
+        # allreduce(mp)->allreduce(sharding)->allreduce(pp)
+        for ring_id in ring_ids:
+            if ring_id == -1:
+                continue
+            # this allreduce communication should not overlap with calc
+            block._insert_op_without_sync(
+                update_loss_scaling_op_idx,
+                type='all_reduce',
+                inputs={'x': inf_var_int32},
+                outputs={'out': inf_var_int32},
+                attrs={
+                    'ring_id': ring_id,
+                    'op_type': paddle.distributed.ReduceOp.MAX,
+                    OP_ROLE_KEY: OpRole.Optimize,
+                },
+            )
+            update_loss_scaling_op_idx += 1
+
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            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,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+        block._sync_with_cpp()
+
+    # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+    @staticmethod
+    def sync_amp_check_nan_inf(block, ring_ids):
+        update_loss_scaling_op_idx = -1
+
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if op.type == "update_loss_scaling":
+                update_loss_scaling_op_idx = idx
+                inf_var_name = op.desc.input('FoundInfinite')[0]
+                break
+
+        # not use amp
+        if update_loss_scaling_op_idx == -1:
+            return
+        # 0. inf_var_int32 = cast(inf_var)
+        # 1. inf_var_int32 = allreduce_max(inf_var_int32)
+        # 3. inf_var = cast(inf_var_int32)
+        inf_var = block.var(inf_var_name)
+        inf_var_int32 = block.create_var(
+            name=inf_var_name + "@cast_int32",
+            shape=inf_var.shape,
+            dtype=core.VarDesc.VarType.INT32,
+        )
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            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,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+
+        # allreduce(mp)->allreduce(pp)
+        for ring_id in ring_ids:
+            if ring_id == -1:
+                continue
+            block._insert_op_without_sync(
+                update_loss_scaling_op_idx,
+                type='all_reduce',
+                inputs={'x': inf_var_int32},
+                outputs={'out': inf_var_int32},
+                attrs={
+                    'ring_id': ring_id,
+                    'op_type': paddle.distributed.ReduceOp.MAX,
+                    OP_ROLE_KEY: OpRole.Optimize,
+                },
+            )
+            update_loss_scaling_op_idx += 1
+
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            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,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+        block._sync_with_cpp()

python/paddle/distributed/fleet/meta_optimizers/sharding/gradient_clip_helper.py

@@ -0,0 +1,259 @@
+# Copyright (c) 2020 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
+
+__all__ = []
+
+
+class GradientClipHelper:
+    def __init__(self, mp_ring_id):
+        self.mp_ring_id = mp_ring_id
+
+    def _is_gradient_clip_op(self, op):
+        return op.desc.has_attr("op_namescope") and op.desc.attr(
+            "op_namescope"
+        ).startswith("/gradient_clip")
+
+    def prune_gradient_clip(self, block, shard, ring_ids):
+        """
+        prune gradient_clip related ops for params that not belong to cur shard
+        prune: square, reduce_sum, elementwise_mul
+        keep: sum, sqrt, elementwise_max, elementwise_div
+        """
+        deprecated_vars = set()
+        deprecate_op_idx = set()
+        reversed_x_paramname = []
+        global_norm_sum_op_idx = -1
+        for idx, op in enumerate(block.ops):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == "sum":
+                global_norm_sum_op_idx = idx
+                continue
+            deprecate_op = False
+            for input_name in op.desc.input_arg_names():
+                if input_name in deprecated_vars:
+                    deprecate_op = True
+                # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+                param_name = input_name.removesuffix("@MERGED").removesuffix(
+                    "@GRAD"
+                )
+                if shard.is_param(param_name) and not shard.has_param(
+                    param_name
+                ):
+                    deprecate_op = True
+                elif shard.is_param(param_name):
+                    reversed_x_paramname.append(param_name)
+
+            if deprecate_op:
+                deprecate_op_idx.add(idx)
+                for output_name in op.desc.output_arg_names():
+                    if output_name not in op.desc.input_arg_names():
+                        deprecated_vars.add(output_name)
+
+        # NOTE(wangxi): If only have 2 sharding, and 1 param.
+        # sharding 0 will not deprecated_vars, will return, only
+        # sharding 1 will insert allreduce, then hang.
+        if not deprecated_vars and global_norm_sum_op_idx == -1:
+            # got no gradient_clip op
+            return
+
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if idx in deprecate_op_idx:
+                block._remove_op(idx, sync=False)
+                continue
+            if op.type == "sum":
+                reversed_inputs = []
+                global_norm_sum_op_idx = idx
+                for input_name in op.desc.input_arg_names():
+                    if input_name not in deprecated_vars:
+                        reversed_inputs.append(input_name)
+
+                op.desc.set_input("X", reversed_inputs)
+                assert len(op.desc.output_arg_names()) == 1
+                sum_res = op.desc.output_arg_names()[0]
+
+                # NOTE(wangxi): If we have 2 param, but sharding is 4,
+                # then the sum op in some cards will not have input.
+                # So we use fill_constant_op to set `sum_var` to zero,
+                # which does not affect correctness.
+                if len(reversed_inputs) == 0:
+                    sum_var = block.var(sum_res)
+                    namescope = op.attr("op_namescope")
+
+                    block._remove_op(idx, sync=False)
+                    op = block._insert_op_without_sync(
+                        idx,
+                        type='fill_constant',
+                        inputs={},
+                        outputs={'Out': sum_res},
+                        attrs={
+                            'shape': sum_var.shape,
+                            'dtype': sum_var.dtype,
+                            'value': 0.0,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    op._set_attr('op_namescope', namescope)
+
+                # allreduce(mp)->allreduce(sharding)->allreduce(pp)
+                idx_offset = 1
+                for ring_id in ring_ids:
+                    if ring_id == -1:
+                        continue
+                    # this allreduce should not overlap with calc and should be scheduled in calc stream
+                    block._insert_op_without_sync(
+                        idx + idx_offset,
+                        type='all_reduce',
+                        inputs={'x': sum_res},
+                        outputs={'out': sum_res},
+                        attrs={
+                            'ring_id': ring_id,
+                            'op_namescope': "/gradient_clip_model_parallelism",
+                            'reduce_type': paddle.distributed.ReduceOp.Sum,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    idx_offset += 1
+
+        # the grad sum here should take the all and only param in the current shard
+        to_check_param = set(reversed_x_paramname)
+        should_check_param = set(shard.global_params).intersection(
+            {
+                param
+                for param, worker_idx in shard.global_param2device.items()
+                if worker_idx == shard.worker_idx
+            }
+        )
+        assert to_check_param == should_check_param, (
+            f"amp check_finite_and_unscale \
+        checking miss [{should_check_param - to_check_param}] and got unexpected [{to_check_param - should_check_param}]"
+        )
+
+        for var_name in deprecated_vars:
+            block._remove_var(var_name, sync=False)
+        block._sync_with_cpp()
+        return
+
+    # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+    def sync_global_norm(self, block, ring_ids, mp_rank):
+        """
+        prune gradient_clip related ops for params that not belong to cur shard
+        prune: square, reduce_sum, elementwise_mul
+        keep: sum, sqrt, elementwise_max, elementwise_div
+        """
+        is_clip_grad_by_global_norm = False
+        for idx, op in list(enumerate(block.ops)):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == 'sum':
+                is_clip_grad_by_global_norm = True
+                break
+        if not is_clip_grad_by_global_norm:
+            # TODO(Yuang Liu): need some extra handles when clip_grad_norm for mp
+            return
+
+        removed_op_idx = set()
+        removed_tmp_var = set()
+        for idx, op in list(enumerate(block.ops)):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == 'sum':
+                break
+            for input_name in op.input_arg_names:
+                input_var = block.var(input_name)
+                # NOTE: when mp_degree > 1, some vars will be split into each mp rank.
+                # However, there still some vars such as Scale, Bias are not split.
+                # Those not be split vars should only be counted once during grad clip
+                # by global norm. Those vars either doesn't have is_distributed attr
+                # or the is_distributed attr has been set as False.
+                # Therefore, we prune those duplicated vars for grad clip.
+                if mp_rank >= 1 and (
+                    not (
+                        hasattr(input_var, 'is_distributed')
+                        and input_var.is_distributed
+                    )
+                ):
+                    removed_op_idx.add(idx)
+                    for output_name in op.output_arg_names:
+                        removed_tmp_var.add(output_name)
+
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if idx in removed_op_idx:
+                block._remove_op(idx, sync=False)
+
+        for var_name in removed_tmp_var:
+            block._remove_var(var_name, sync=False)
+
+        for idx, op in list(enumerate(block.ops)):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == 'sum':
+                # If mp_rank == 0, no extra handles, just allreduce
+                # If mp_rank >= 1, some extra handles is needed
+                sum_rst_var = block.var(op.output_arg_names[0])
+                if mp_rank >= 1:
+                    reserved_vars = []
+                    for input_name in op.input_arg_names:
+                        if input_name not in removed_tmp_var:
+                            reserved_vars.append(input_name)
+
+                    if len(reserved_vars) > 0:
+                        op.desc.set_input("X", reserved_vars)
+                    else:
+                        # If all input of sum op should be removed, then remove the sum op.
+                        # And set the output's value of sum to 0.
+                        namescope = op.attr("op_namescope")
+                        block._remove_op(idx, sync=False)
+                        fill_constant_op = block._insert_op_without_sync(
+                            idx,
+                            type='fill_constant',
+                            inputs={},
+                            outputs={'Out': sum_rst_var},
+                            attrs={
+                                'shape': sum_rst_var.shape,
+                                'dtype': sum_rst_var.dtype,
+                                'value': 0.0,
+                                OP_ROLE_KEY: OpRole.Optimize,
+                            },
+                        )
+                        fill_constant_op._set_attr('op_namescope', namescope)
+                self._insert_allreduce(block, ring_ids, idx, sum_rst_var)
+                break
+
+    @staticmethod
+    def _insert_allreduce(block, ring_ids, idx, var):
+        for ring_id in ring_ids:
+            if ring_id == -1:
+                continue
+
+            idx = idx + 1
+            block._insert_op_without_sync(
+                idx,
+                type='all_reduce',
+                inputs={'x': var},
+                outputs={'out': var},
+                attrs={
+                    'ring_id': ring_id,
+                    'op_namescope': "/gradient_clip_model_parallelism",
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
+                    OP_ROLE_KEY: OpRole.Optimize,
+                },
+            )

python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py

@@ -0,0 +1,575 @@
+# Copyright (c) 2020 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.framework import core
+from paddle.utils import unique_name
+
+from ..common import OP_ROLE_KEY, OpRole, is_optimizer_op, is_update_op
+
+__all__ = []
+
+
+class PlaceType:
+    # sync with memcpy op, maybe not a good design
+    CPU = 0
+    CUDA = 1
+    CUDA_PINNED = 2
+    XPU = 3  # unsupported for now
+
+    @staticmethod
+    def default_device():
+        if core.is_compiled_with_cuda():
+            return PlaceType.CUDA
+        return PlaceType.CPU
+
+    @staticmethod
+    def default_pinned():
+        if core.is_compiled_with_cuda():
+            return PlaceType.CUDA_PINNED
+        return PlaceType.CPU
+
+
+class OffloadHelper:
+    cpu_place_type = 0
+    cuda_place_type = PlaceType.default_device()
+    cuda_pinned_place_type = PlaceType.default_pinned()
+
+    def __init__(self, mp_ring_id=None, dp_ring_id=None):
+        self.mp_ring_id = mp_ring_id
+        self.dp_ring_id = dp_ring_id
+
+    def _insert_cast_op(self, block, idx, src_name, dst_name):
+        src_var = block.var(src_name)
+        if not block.has_var(dst_name):
+            block.create_var(
+                name=dst_name,
+                shape=src_var.shape,
+                dtype=core.VarDesc.VarType.FP16,
+                persistable=True,
+            )
+        dst_var = block.var(dst_name)
+        assert dst_var.dtype == paddle.float16
+        block._insert_op_without_sync(
+            idx,
+            type='cast',
+            inputs={'X': src_var},
+            outputs={'Out': dst_var},
+            attrs={
+                'in_dtype': src_var.dtype,
+                'out_dtype': dst_var.dtype,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+
+    def _insert_broadcast_op(self, block, idx, param_name):
+        rings = []
+
+        if self.dp_ring_id is not None:
+            rings.append(self.dp_ring_id)
+
+        # need sync non distributed param in mp group
+        if self.mp_ring_id is not None:
+            param = block.var(param_name)
+            if not hasattr(param, 'is_distributed') or not param.is_distributed:
+                rings.append(self.mp_ring_id)
+
+        # the insert op order is: mp, dp
+        for ring in rings:
+            block._insert_op_without_sync(
+                idx,
+                type="broadcast",
+                inputs={'x': param_name},
+                outputs={'out': param_name},
+                attrs={
+                    'ring_id': ring,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+    def _insert_memcpy_op(self, block, idx, src_name, dst_name, dst_place_type):
+        src_var = block.var(src_name)
+        dst_var = block.var(dst_name)
+        block._insert_op_without_sync(
+            idx,
+            type='memcpy',
+            inputs={'X': src_var},
+            outputs={'Out': dst_var},
+            attrs={
+                'dst_place_type': dst_place_type,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+
+    def _insert_fetch_op(self, block, idx, src_name, dst_name):
+        self._insert_memcpy_op(
+            block, idx, src_name, dst_name, OffloadHelper.cuda_place_type
+        )
+
+    def _insert_offload_op(self, block, idx, src_name, dst_name):
+        self._insert_memcpy_op(
+            block, idx, src_name, dst_name, OffloadHelper.cuda_pinned_place_type
+        )
+
+    def _get_offload_var_name(self, name):
+        return unique_name.generate(name + '@offload')
+
+    def _create_offload_var(self, var_name, offload_var_name, blocks):
+        for block in blocks:
+            var = block.var(var_name)
+            var.persistable = False
+            offload_var = block.create_var(
+                name=offload_var_name,
+                shape=var.shape,
+                dtype=var.dtype,
+                persistable=True,
+            )
+
+    def offload_fp32param(self, block, startup_block, offload=True):
+        """
+        (p_fp16) = cast(p)
+        (p_fp16_recompute) = cast(p)
+        (pout,) = adam(p)
+        ===========================>
+        rename(p_fp16_recompute, p_fp16)
+
+        (p,) = prefetch(p@offload)
+        (pout,) = adam(p)
+        (p_fp16) = cast(p)
+        (p@offload) = memcpy(p)
+        """
+        param_to_idx = {}
+        param_to_fp16 = {}
+        # recompute_var which need rename to fp16_param
+        fp16_param_to_recompute = {}
+        recompute_to_fp16 = {}
+
+        def remove_param(input_name):
+            param_to_idx.pop(input_name)
+            if input_name in param_to_fp16:
+                fp16_param = param_to_fp16.pop(input_name)
+                if fp16_param in fp16_param_to_recompute:
+                    recompute = fp16_param_to_recompute.pop(fp16_param)
+                    recompute_to_fp16.pop(recompute)
+
+        # step1: record param
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_update_op(op):
+                param = op.desc.input("Param")[0]
+                param_to_idx[param] = idx
+
+        # step2: remove param which can't offload and
+        #        record param->fp16param, fp16param->recompute_var
+        for idx, op in enumerate(block.ops):
+            if is_optimizer_op(op):
+                break
+            # TODO (Yuang Liu): tmp solution for fuse_grad_merge + optimize_cast
+            if not offload and op.type == 'coalesce_tensor':
+                continue
+            for input_name in op.desc.input_arg_names():
+                if input_name not in param_to_idx:
+                    continue
+
+                # param which will be used by fp32 op
+                if op.type != 'cast':
+                    remove_param(input_name)
+                    continue
+
+                # param is only used by cast op,
+                # which to cast fp32_param to fp16_param
+                output_name = op.output_arg_names[0]
+                if 'cast_fp16' not in output_name:
+                    remove_param(input_name)
+                    continue
+
+                if 'subprog' not in output_name:
+                    assert output_name == input_name + '.cast_fp16'
+                    assert input_name not in param_to_fp16, (
+                        "There must be only one cast op from fp32 param to fp16 param."
+                    )
+                    param_to_fp16[input_name] = output_name
+                else:
+                    # fp16-->recompute_var
+                    assert input_name in param_to_fp16, (
+                        "param must first be cast to fp16"
+                    )
+                    fp16_param = param_to_fp16[input_name]
+                    fp16_param_to_recompute[fp16_param] = output_name
+                    recompute_to_fp16[output_name] = fp16_param
+
+        param_name_to_offload_name = {}
+        # step3: main_block add offload, cast op
+        # change recompute to fp16, remove cast(param) to fp16
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_update_op(op):
+                param = op.desc.input("Param")[0]
+                if param not in param_to_idx:
+                    continue
+                # step3.1: create offload_var
+                offload_var_name = self._get_offload_var_name(param)
+                param_name_to_offload_name[param] = offload_var_name
+                if offload:
+                    self._create_offload_var(
+                        param, offload_var_name, [block, startup_block]
+                    )
+
+                    # step3.2: insert cast op and offload op
+                    self._insert_offload_op(
+                        block, idx + 1, param, offload_var_name
+                    )
+
+                assert param in param_to_fp16
+                fp16_param_name = param_to_fp16[param]
+                fp16_param_var = block.var(fp16_param_name)
+                fp16_param_var.persistable = True
+                self._insert_cast_op(
+                    block, idx + 1, param, param_to_fp16[param]
+                )
+
+                if offload:
+                    # step3.3: insert fetch op
+                    self._insert_fetch_op(block, idx, offload_var_name, param)
+                continue
+
+            # step3.4: remove cast op
+            if op.type == 'cast':
+                input_name = op.desc.input_arg_names()[0]
+                if input_name in param_to_idx:
+                    block._remove_op(idx, sync=False)
+                    continue
+
+            # step3.5: change recompute_param to fp16_param
+            for input_name in op.desc.input_arg_names():
+                if input_name in recompute_to_fp16:
+                    op._rename_input(input_name, recompute_to_fp16[input_name])
+            for output_name in op.desc.output_arg_names():
+                if output_name in recompute_to_fp16:
+                    op._rename_output(
+                        output_name, recompute_to_fp16[output_name]
+                    )
+
+        # step4: remove recompute_param
+        for name in recompute_to_fp16.keys():
+            block._remove_var(name, sync=False)
+
+        # step5: startup_block add offload
+        visited_vars = set()
+        # FIXME(wangxi): should insert in idx, need move comm init to the head.
+        insert_idx = len(startup_block.ops)
+        for idx, op in reversed(list(enumerate(startup_block.ops))):
+            for out_name in op.output_arg_names:
+                if out_name in visited_vars:
+                    continue
+
+                if out_name in param_name_to_offload_name:
+                    var_name = out_name
+                    if offload:
+                        offload_var_name = param_name_to_offload_name[var_name]
+                        self._insert_offload_op(
+                            startup_block,
+                            insert_idx,
+                            var_name,
+                            offload_var_name,
+                        )
+                    self._insert_cast_op(
+                        startup_block,
+                        insert_idx,
+                        var_name,
+                        param_to_fp16[var_name],
+                    )
+                    # NOTE(wangxi): cast and offload should insert after broadcast param.
+                    # the insert op order is: {mp, dp}broadcast, cast, offload
+                    self._insert_broadcast_op(
+                        startup_block, insert_idx, var_name
+                    )
+
+                visited_vars.add(out_name)
+
+        block._sync_with_cpp()
+        startup_block._sync_with_cpp()
+
+    def cast_fp32param_in_optimize(self, block, startup_block):
+        """
+        (p_fp16) = cast(p)
+        (p_fp16_recompute) = cast(p)
+        (pout,) = adam(p)
+        ===========================>
+        rename(p_fp16_recompute, p_fp16)
+
+        (pout,) = adam(p)
+        (p_fp16) = cast(p)
+        """
+        self.offload_fp32param(block, startup_block, offload=False)
+
+    def offload(self, block, startup_block):
+        """
+        (m1, m2) = prefetch(m1@offload, m2@offload)
+        (m1out, m2out, pout) = adam(m1, m2, p)
+        (m1@offload, m2@offload) = memcpy(m1, m2)
+        """
+        vars_name_to_offload_name = {}
+
+        # main_block add offload
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not is_optimizer_op(op):
+                break
+
+            vars_name = []
+            if op.type == "adam" or op.type == "adamw":
+                # {Moment1Out = [''], Moment2Out = [''], ParamOut = ['']} =
+                # adam(inputs={Moment1 = [''], Moment2 = [''], Param = ['']})
+                vars_name.append(op.desc.input("Moment1")[0])
+                vars_name.append(op.desc.input("Moment2")[0])
+            elif op.type == 'momentum':
+                pass
+            elif op.type == 'lars':
+                pass
+            elif op.type == 'lamb':
+                pass
+
+            # step1: create and init offload_var
+            for var_name in vars_name:
+                assert var_name not in vars_name_to_offload_name
+
+                offload_var_name = self._get_offload_var_name(var_name)
+                vars_name_to_offload_name[var_name] = offload_var_name
+
+                self._create_offload_var(
+                    var_name, offload_var_name, [block, startup_block]
+                )
+
+            # step2: insert offload op
+            for var_name in vars_name:
+                offload_var_name = vars_name_to_offload_name[var_name]
+                self._insert_offload_op(
+                    block, idx + 1, var_name, offload_var_name
+                )
+
+            # step3: insert fetch op
+            for var_name in vars_name:
+                offload_var_name = vars_name_to_offload_name[var_name]
+                self._insert_fetch_op(block, idx, offload_var_name, var_name)
+
+        # startup_block add offload
+        visited_vars = set()
+        for idx, op in reversed(list(enumerate(startup_block.ops))):
+            for out_name in op.output_arg_names:
+                if out_name in visited_vars:
+                    continue
+
+                if out_name in vars_name_to_offload_name:
+                    var_name = out_name
+                    offload_var_name = vars_name_to_offload_name[var_name]
+                    # insert offload op after var is generated
+                    self._insert_offload_op(
+                        startup_block, idx + 1, var_name, offload_var_name
+                    )
+                visited_vars.add(out_name)
+
+        block._sync_with_cpp()
+        startup_block._sync_with_cpp()
+
+    def opt_sharding_cast_fp32param(
+        self, block, startup_block, params, offload=False
+    ):
+        """
+        (p_fp16) = cast(p)
+        (p_fp16_recompute) = cast(p)
+        (pout,) = adam(p)
+        ===========================>
+        rename(p_fp16_recompute, p_fp16)
+
+        (pout,) = adam(p)
+        (p_fp16) = cast(p)
+        broadcast(p_fp16)
+        """
+        global_params = set()
+        local_params = set()
+        param_to_fp16 = {}
+        # recompute_var which need rename to fp16_param
+        fp16_param_to_recompute = {}
+        recompute_to_fp16 = {}
+
+        def remove_param(input_name):
+            global_params.remove(input_name)
+            if input_name in local_params:  # noqa: FURB132
+                local_params.remove(input_name)
+            if input_name in param_to_fp16:
+                fp16_param = param_to_fp16.pop(input_name)
+                if fp16_param in fp16_param_to_recompute:
+                    recompute = fp16_param_to_recompute.pop(fp16_param)
+                    recompute_to_fp16.pop(recompute)
+
+        # step1: record param
+        global_params = set(params)
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_update_op(op):
+                param = op.desc.input("Param")[0]
+                local_params.add(param)
+
+        # step2: remove param which can't offload and
+        #        record param->fp16param, fp16param->recompute_var
+        for idx, op in enumerate(block.ops):
+            if is_optimizer_op(op):
+                break
+            # TODO (Yuang Liu): tmp solution for fuse_grad_merge + optimize_cast
+            if op.type == 'coalesce_tensor':
+                continue
+            for input_name in op.desc.input_arg_names():
+                if input_name not in global_params:
+                    continue
+
+                # param which will be used by fp32 op
+                if op.type != 'cast':
+                    remove_param(input_name)
+                    continue
+
+                # param is only used by cast op,
+                # which to cast fp32_param to fp16_param
+                output_name = op.output_arg_names[0]
+                if 'cast_fp16' not in output_name:
+                    remove_param(input_name)
+                    continue
+
+                if 'subprog' not in output_name:
+                    assert output_name == input_name + '.cast_fp16'
+                    assert input_name not in param_to_fp16, (
+                        "There must be only one cast op from fp32 param to fp16 param."
+                    )
+                    param_to_fp16[input_name] = output_name
+                else:
+                    # fp16-->recompute_var
+                    assert input_name in param_to_fp16, (
+                        "param must first be cast to fp16"
+                    )
+                    fp16_param = param_to_fp16[input_name]
+                    fp16_param_to_recompute[fp16_param] = output_name
+                    recompute_to_fp16[output_name] = fp16_param
+
+        param_name_to_offload_name = {}
+        # step3: main_block add offload, cast op
+        # change recompute to fp16, remove cast(param) to fp16
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_update_op(op):
+                param = op.desc.input("Param")[0]
+                if param not in global_params:
+                    continue
+                # step3.1: create offload_var
+                offload_var_name = self._get_offload_var_name(param)
+                param_name_to_offload_name[param] = offload_var_name
+                if offload:
+                    self._create_offload_var(
+                        param, offload_var_name, [block, startup_block]
+                    )
+
+                    # step3.2: insert cast op and offload op
+                    self._insert_offload_op(
+                        block, idx + 1, param, offload_var_name
+                    )
+
+                assert param in param_to_fp16
+                fp16_param_name = param_to_fp16[param]
+                fp16_param_var = block.var(fp16_param_name)
+                fp16_param_var.persistable = True
+                self._insert_cast_op(
+                    block, idx + 1, param, param_to_fp16[param]
+                )
+
+                if offload:
+                    # step3.3: insert fetch op
+                    self._insert_fetch_op(block, idx, offload_var_name, param)
+
+                continue
+
+            # step3.4: remove cast op
+            if op.type == 'cast':
+                input_name = op.desc.input_arg_names()[0]
+                if input_name in global_params:
+                    block._remove_op(idx, sync=False)
+                    continue
+
+            # step3.5: change recompute_param to fp16_param
+            for input_name in op.desc.input_arg_names():
+                if input_name in recompute_to_fp16:
+                    op._rename_input(input_name, recompute_to_fp16[input_name])
+            for output_name in op.desc.output_arg_names():
+                if output_name in recompute_to_fp16:
+                    op._rename_output(
+                        output_name, recompute_to_fp16[output_name]
+                    )
+
+        # step4: remove recompute_param
+        for name in recompute_to_fp16.keys():
+            block._remove_var(name, sync=False)
+
+        # step5: remove fp32 param which not need
+        for idx, op in enumerate(block.ops):
+            if op.type not in ['coalesce_tensor', 'c_broadcast', 'broadcast']:
+                continue
+            for input_name in op.desc.input_arg_names():
+                if input_name in param_to_fp16:
+                    op._rename_input(input_name, param_to_fp16[input_name])
+            for output_name in op.desc.output_arg_names():
+                if output_name in param_to_fp16:
+                    op._rename_output(output_name, param_to_fp16[output_name])
+
+        for param in global_params:
+            assert param in param_to_fp16
+            fp16_param_name = param_to_fp16[param]
+            fp16_param_var = block.var(fp16_param_name)
+            fp16_param_var.persistable = True
+
+            if param not in local_params:
+                block._remove_var(param, sync=False)
+
+        # step6: startup_block add offload
+        visited_vars = set()
+        insert_idx = len(startup_block.ops)
+        for idx, op in reversed(list(enumerate(startup_block.ops))):
+            for out_name in op.output_arg_names:
+                if out_name in visited_vars:
+                    continue
+
+                if out_name in param_to_fp16:
+                    var_name = out_name
+                    if offload:
+                        self._insert_offload_op(
+                            startup_block,
+                            idx + 1,
+                            var_name,
+                            param_name_to_offload_name[var_name],
+                        )
+
+                    self._insert_cast_op(
+                        startup_block,
+                        insert_idx,
+                        var_name,
+                        param_to_fp16[var_name],
+                    )
+
+                    # NOTE(wangxi): cast and offload should insert after broadcast param.
+                    # the insert op order is: {mp, dp}broadcast, cast, offload
+                    self._insert_broadcast_op(
+                        startup_block, insert_idx, var_name
+                    )
+
+                    if var_name not in local_params:
+                        param = startup_block.var(out_name)
+                        param.persistable = False
+
+                visited_vars.add(out_name)
+
+        block._sync_with_cpp()
+        startup_block._sync_with_cpp()

python/paddle/distributed/fleet/meta_optimizers/sharding/prune.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2020 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__ = []
+
+
+class ProgramDeps:
+    def __init__(self, block, start_vars, end_vars):
+        self._block = block
+        # vars where to start to build the deps
+        self._start_vars = start_vars
+        # vars where to stop to build the deps
+        self._end_vars = end_vars
+        # var name -> op idxs which depends on this var
+        self._var_to_use_op = {}
+        # sub block deps which is a subset of this topo
+        self._sub_block_deps = {}
+        # var name -> op idxs which generate var
+        self._var_to_generate_op = {}
+        self._should_removed_var = set()
+        self._father_block_deps = None
+        self._build_deps()
+
+    def get_sub_block_deps(self, idx):
+        if idx in self._sub_block_deps:
+            return self._sub_block_deps[idx]
+        else:
+            return None
+
+    def get_var_deps(self, var_name):
+        if var_name in self._var_to_use_op:
+            return self._var_to_use_op[var_name]
+        else:
+            return None
+
+    def _build_deps(
+        self,
+    ):
+        for var_name in self._start_vars:
+            self._var_to_use_op[var_name] = []
+            self._var_to_generate_op[var_name] = []
+
+        for idx, op in enumerate(self._block.ops):
+            if op.type in [
+                "c_sync_comm_stream",
+                "c_calc_comm_stream",
+                'all_reduce',
+            ]:
+                continue
+            input_vars = op.desc.input_arg_names()
+            output_vars = op.desc.output_arg_names()
+            deps_reduce = False
+            for input_name in input_vars:
+                if input_name in self._var_to_use_op:
+                    deps_reduce = True
+            if not deps_reduce:
+                continue
+            for input_name in input_vars:
+                if input_name in self._var_to_use_op:
+                    self._var_to_use_op[input_name].append(idx)
+            for output_name in output_vars:
+                if output_name not in self._var_to_use_op:
+                    self._var_to_use_op[output_name] = []
+                if output_name not in self._var_to_generate_op:
+                    self._var_to_generate_op[output_name] = [idx]
+                else:
+                    self._var_to_generate_op[output_name].append(idx)
+            if op.type == "conditional_block":
+                # subblock
+                assert op.desc.has_attr("sub_block")
+                subblock_idx = op.desc.attr("sub_block").id
+                subblock_deps = ProgramDeps(
+                    self._block.program.block(subblock_idx),
+                    op.desc.input_arg_names(),
+                    op.desc.output_arg_names(),
+                )
+                self._sub_block_deps[subblock_idx] = subblock_deps
+                subblock_deps._father_block_deps = self
+
+    def crop_input_var_from_op(self, op_idx, var_name):
+        if var_name in self._var_to_use_op:
+            # update var -> dep_var_op
+            if self._var_to_use_op[var_name] != []:
+                if op_idx not in self._var_to_use_op[var_name]:
+                    raise ValueError(
+                        f"op_idx: {op_idx} is not in self._var_to_use_op[{var_name}], "
+                        f"self._var_to_use_op[{var_name}] is {self._var_to_use_op[var_name]}"
+                    )
+                self._var_to_use_op[var_name].remove(op_idx)
+            # update _should_removed_var
+            if var_name in self._start_vars:
+                self._should_removed_var.discard(var_name)
+            elif (
+                self._var_to_use_op[var_name] == []
+            ):  # no more deps of this var
+                self._should_removed_var.add(var_name)
+            elif (
+                self._var_to_generate_op[var_name][-1]
+                >= self._var_to_use_op[var_name][-1]
+            ):
+                # there are circle in the graph
+                self._should_removed_var.add(var_name)
+            else:  # input_name should not be deleted
+                self._should_removed_var.discard(var_name)
+
+    def crop_output_var_from_op(self, op_idx, var_name):
+        if var_name in self._var_to_generate_op:
+            assert op_idx in self._var_to_generate_op[var_name]
+            self._var_to_generate_op[var_name].remove(op_idx)
+        if self._block.has_var(var_name):
+            if (
+                var_name not in self._var_to_generate_op
+                or self._var_to_generate_op[var_name] == []
+            ):
+                self._block._remove_var(var_name, sync=False)
+
+    def remove_op(self, op_idx, reserved_vars=None):
+        # update deps
+        op = self._block.ops[op_idx]
+        for input_name in op.desc.input_arg_names():
+            if reserved_vars is not None and input_name in reserved_vars:
+                continue
+            self.crop_input_var_from_op(op_idx, input_name)
+        for output_name in op.desc.output_arg_names():
+            if reserved_vars is not None and output_name in reserved_vars:
+                continue
+            self.crop_output_var_from_op(op_idx, output_name)
+        self._block._remove_op(op_idx, sync=False)
+
+    def should_remove_op(self, op_idx):
+        op = self._block.ops[op_idx]
+
+        # NOTE: At present, it is found that the OP without output is
+        # only send_v2 and partial_send op, which will be used in
+        # all device
+        if len(op.desc.output_arg_names()) == 0:
+            return False
+
+        for output_name in op.desc.output_arg_names():
+            if output_name not in self._should_removed_var:
+                return False
+        return True

python/paddle/distributed/fleet/meta_optimizers/sharding/shard.py

@@ -0,0 +1,175 @@
+# Copyright (c) 2020 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 re
+
+from paddle.distributed.fleet.meta_optimizers.common import is_optimizer_op
+from paddle.distributed.fleet.meta_optimizers.sharding.fp16_helper import (
+    FP16Utils,
+)
+from paddle.distributed.fleet.meta_optimizers.sharding.utils import get_var_size
+
+__all__ = []
+
+
+class Shard:
+    def __init__(
+        self,
+    ):
+        self.global_params = set()
+        self.worker_idx = -1
+        self.worker_num = -1
+        self.global_param2device = {}
+        self.device2global_params = {}
+
+    def setup(self, params_grads, worker_idx, worker_num):
+        # param names of all devices
+        self.global_params = {x[0].name for x in params_grads}
+        # _param(str) -> device_id(int)
+        self.worker_idx = worker_idx
+        self.worker_num = worker_num
+        # global_param2device contains fp32 params and fp16 params
+        # device2global_params only contains fp32 params
+        (
+            self.global_param2device,
+            self.device2global_params,
+        ) = self._split_params(params_grads, worker_idx, worker_num)
+
+    def has_param(self, var_name):
+        return (
+            var_name in self.global_param2device
+            and self._var_device_id(var_name) == self.worker_idx
+        )
+
+    def has_opt_var(self, var_name):
+        return self._var_device_id(var_name) == self.worker_idx
+
+    def has_var(self, var_name):
+        return (
+            self._var_device_id(var_name) == -1
+            or self._var_device_id(var_name) == self.worker_idx
+        )
+
+    def _split_params(self, params_grads, worker_idx, worker_num):
+        param2device = {}
+        total_param_mem = 0.0
+        param2mem = []
+        for param in [x[0] for x in params_grads]:
+            mem = get_var_size(param)
+            total_param_mem += mem
+            param2mem.append((param.name, mem))
+        device2params = {x: [] for x in range(worker_num)}
+        device_idx = 0
+        mem_accu = 0.0
+        for param_name, mem in param2mem:
+            if mem_accu > total_param_mem * 1.0 * (device_idx + 1) / worker_num:
+                device_idx += 1
+            device2params[device_idx].append(param_name)
+            param2device[param_name] = device_idx
+            mem_accu += mem
+        return param2device, device2params
+
+    def _var_device_id(self, var_name):
+        if var_name in self.global_param2device:
+            return self.global_param2device[var_name]
+        for suffix in [
+            "_moment1_0",
+            "_moment2_0",
+            "_beta1_pow_acc_0",
+            "_beta2_pow_acc_0",
+            "_velocity_0",
+        ]:
+            base_name = re.sub(suffix, '', var_name)
+            if base_name in self.global_param2device:
+                return self.global_param2device[base_name]
+        return -1
+
+    def find_broadcast_params(self, block):
+        broadcast_vars = set()
+        fp16_params = set()
+        fp16_to_fp32 = {}
+
+        param_usage = dict.fromkeys(self.global_params, 0)
+        for op in block.ops:
+            if is_optimizer_op(op):
+                continue
+            for input_name in op.desc.input_arg_names():
+                if input_name in self.global_params:
+                    param_usage[input_name] += 1
+
+        for op in block.ops:
+            if not FP16Utils.is_fp16_cast_op(block, op, self.global_params):
+                continue
+            input_name = op.input_arg_names[0]
+            output_name = op.output_arg_names[0]
+            broadcast_vars.add(output_name)
+            fp16_params.add(output_name)
+            fp16_to_fp32[output_name] = input_name
+            param_usage[input_name] -= 1
+            self.global_param2device[output_name] = self.global_param2device[
+                input_name
+            ]
+
+        for param, usage in param_usage.items():
+            if usage > 0:
+                broadcast_vars.add(param)
+        return broadcast_vars
+
+    def device(self, var_name):
+        return self._var_device_id(var_name)
+
+    def is_param(self, var_name):
+        return var_name in self.global_params
+
+    def is_opti_var(self, var_name):
+        if var_name in self.global_params:
+            return True
+        for suffix in [
+            "_moment1_0",
+            "_moment2_0",
+            "_beta1_pow_acc_0",
+            "_beta2_pow_acc_0",
+            "_velocity_0",
+        ]:
+            base_name = re.sub(suffix, '', var_name)
+            if base_name in self.global_params:
+                return True
+        return False
+
+    def filter_grads(self, grads):
+        grads_in_shard = []
+        for grad in grads:
+            param = grad.split("@")[0]
+            if self.has_param(param):
+                grads_in_shard.append(grad)
+        return grads_in_shard
+
+
+class ProgramSegment:
+    def __init__(self, block):
+        self._block = block
+        self._allreduce_vars = []
+        # sub program start idx
+        self._start_idx = -1
+        # sub program end idx
+        self._end_idx = -1
+        # param name to broadcast name
+        self._param2broadcast = {}
+        self._broadcast_vars = []
+        # cast op pairs, fp16 name (str) -> fp32 name (str)
+        self._cast_ops = {}
+        # fill constant vars
+        self._fill_constant_vars = []
+        # parameter mems
+        self._param_mem = 0.0

python/paddle/distributed/fleet/meta_optimizers/sharding/weight_decay_helper.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2020 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 OP_ROLE_VAR_KEY
+
+__all__ = []
+
+
+class WeightDecayHelper:
+    def __init__(self):
+        pass
+
+    def _is_weight_decay_op(self, op):
+        return op.desc.has_attr("op_namescope") and op.desc.attr(
+            "op_namescope"
+        ).startswith("/regularization")
+
+    def prune_weight_decay(self, block, shard):
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not self._is_weight_decay_op(op):
+                continue
+            if OP_ROLE_VAR_KEY not in op.attr_names:
+                raise ValueError(
+                    "The Weight Decay op should hold op_role_var attribute"
+                    f"but the {op.type} op does not hold op_role_var"
+                )
+            op_role_var = op.all_attrs()[OP_ROLE_VAR_KEY]
+            if not shard.has_param(op_role_var[0]):
+                block._remove_op(idx, sync=False)
+        block._sync_with_cpp()

python/paddle/distributed/fleet/meta_optimizers/tensor_parallel_optimizer.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
+
+import paddle
+from paddle import static
+
+from .common import (
+    OP_ROLE_KEY,
+    OP_ROLE_VAR_KEY,
+    CollectiveHelper,
+    OpRole,
+    is_backward_op,
+    is_loss_grad_op,
+    is_optimizer_op,
+)
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class TensorParallelOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = [
+            "RecomputeOptimizer",
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+        self.mp_ring_id = 0
+        self.global_ring_id = 1
+        self.dp_ring_id = 2
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+        self.mp_degree = user_defined_strategy.tensor_parallel_configs[
+            'tensor_parallel_degree'
+        ]
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.tensor_parallel:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.tensor_parallel = False
+        dist_strategy.tensor_parallel_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.tensor_parallel = True
+        dist_strategy.tensor_parallel_configs = {
+            "tensor_parallel_degree": 1,
+        }
+
+    def _broadcast_params(self, ring_id, mp_mode):
+        block = self.startup_program.global_block()
+        param = None
+        for param in block.iter_parameters():
+            if param.is_distributed and mp_mode:
+                continue
+
+            block.append_op(
+                type='broadcast',
+                inputs={'x': param},
+                outputs={'out': param},
+                attrs={
+                    'ring_id': ring_id,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+        if not param:
+            return  # no parameter on this device
+        block.append_op(
+            type='c_sync_comm_stream',
+            inputs={'X': param},
+            outputs={'Out': param},
+            attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Forward},
+        )
+
+    def _get_process_group_info(self):
+        # global ring info
+        self.global_endpoints = self.endpoints
+        self.global_rank = self.rank
+        self.global_nranks = self.nranks
+
+        # model parallel ring info
+        self.mp_rank = self.rank % self.mp_degree
+        self.mp_nranks = self.mp_degree
+        mp_group = self.rank // self.mp_degree
+        self.mp_endpoints = [
+            self.endpoints[i]
+            for i in range(self.global_nranks)
+            if i // self.mp_degree == mp_group
+        ]
+
+        # data parallel ring info
+        if self.nranks > self.mp_degree:
+            self.dp_rank = self.rank // self.mp_degree
+            self.dp_nranks = self.nranks // self.mp_degree
+            start_index = self.rank % self.mp_degree
+            self.dp_endpoints = [
+                self.endpoints[start_index + i * self.mp_degree]
+                for i in range(self.dp_nranks)
+            ]
+
+    def _init_process_group(self):
+        self._get_process_group_info()
+        collective_helper = CollectiveHelper(self.role_maker, wait_port=False)
+
+        # Create global ring for all gpus
+        collective_helper._init_communicator(
+            self.startup_program,
+            self.current_endpoint,
+            self.global_endpoints,
+            self.global_rank,
+            self.global_ring_id,
+            True,
+            self.global_ring_id,
+            True,
+        )
+
+        # Create model parallel ring for all gpus
+        collective_helper._init_communicator(
+            self.startup_program,
+            self.current_endpoint,
+            self.mp_endpoints,
+            self.mp_rank,
+            self.mp_ring_id,
+            True,
+            self.global_ring_id,
+            True,
+        )
+        self._broadcast_params(self.mp_ring_id, mp_mode=True)
+
+        # Create dp rings
+        if self.nranks > self.mp_degree:
+            collective_helper._init_communicator(
+                self.startup_program,
+                self.current_endpoint,
+                self.dp_endpoints,
+                self.dp_rank,
+                self.dp_ring_id,
+                True,
+                self.global_ring_id,
+                True,
+            )
+            self._broadcast_params(self.dp_ring_id, mp_mode=False)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self.endpoints = self.role_maker._get_trainer_endpoints()
+        self.current_endpoint = self.endpoints[self.role_maker._worker_index()]
+        self.startup_program = startup_program
+        if startup_program is None:
+            self.startup_program = static.default_startup_program()
+
+        optimize_ops, params_grads = self.inner_opt.minimize(
+            loss, self.startup_program, parameter_list, no_grad_set
+        )
+
+        self.main_program = loss.block.program
+        self.nranks = len(self.endpoints)
+        self.rank = self.role_maker._worker_index()
+
+        self._init_process_group()
+
+        assert self.nranks % self.mp_degree == 0
+
+        if self.nranks > self.mp_degree:
+            # data parallelism
+            dp_degree = self.nranks // self.mp_degree
+            self._transpile_main_program(loss, dp_degree)
+        return optimize_ops, params_grads
+
+    def _transpile_main_program(self, loss, dp_degree):
+        self._insert_loss_grad_ops(loss, dp_degree)
+        self._insert_allreduce_ops(loss, self.dp_ring_id)
+
+    def _insert_loss_grad_ops(self, loss, dp_degree):
+        """
+        In order to keep the learning rate consistent in different numbers of
+        training workers, we scale the loss grad by the number of workers
+        """
+        block = loss.block
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_loss_grad_op(op):
+                loss_grad_var = block.vars[op.output_arg_names[0]]
+                block._insert_op(
+                    idx + 1,
+                    type='scale',
+                    inputs={'X': loss_grad_var},
+                    outputs={'Out': loss_grad_var},
+                    attrs={
+                        'scale': 1.0 / dp_degree,
+                        OP_ROLE_KEY: OpRole.Backward,
+                    },
+                )
+                break
+
+    def _insert_allreduce_ops(self, loss, ring_id):
+        block = loss.block
+        grad = None
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
+                op_role_var = op.attr(OP_ROLE_VAR_KEY)
+                if len(op_role_var) == 0:
+                    continue
+                assert len(op_role_var) % 2 == 0
+                offset = idx
+                for i in range(0, len(op_role_var), 2):
+                    param = block.vars[op_role_var[i]]
+                    grad = block.vars[op_role_var[i + 1]]
+                    if offset == idx:
+                        offset += 1
+                        block._insert_op(
+                            offset,
+                            type='c_sync_calc_stream',
+                            inputs={'X': grad},
+                            outputs={'Out': grad},
+                            attrs={OP_ROLE_KEY: OpRole.Backward},
+                        )
+                        offset += 1
+
+                    block._insert_op(
+                        offset,
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Backward,
+                        },
+                    )
+
+        if grad is None:
+            return
+
+        for idx, op in list(enumerate(block.ops)):
+            if is_optimizer_op(op):
+                block._insert_op(
+                    idx,
+                    type='c_sync_comm_stream',
+                    inputs={'X': grad},
+                    outputs={'Out': grad},
+                    attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Backward},
+                )
+                break

python/paddle/distributed/fleet/meta_parallel/__init__.py

@@ -0,0 +1,57 @@
+# 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.
+
+# isort: skip_file
+from .meta_parallel_base import MetaParallelBase  # noqa: F401
+
+from .parallel_layers import (  # noqa: F401
+    ColumnParallelLinear,
+    LayerDesc,
+    LocalSharedLayerDesc,
+    ParallelCrossEntropy,
+    PipelineLayer,
+    RNGStatesTracker,
+    RowParallelLinear,
+    SharedLayerDesc,
+    VocabParallelEmbedding,
+    get_rng_state_tracker,
+    model_parallel_random_seed,
+    LayerSpec,
+    import_spec_layer,
+    get_spec_layer,
+    build_spec_layer,
+)
+from .pipeline_parallel import (  # noqa: F401
+    NoPipelineParallel,
+    PipelineParallel,
+    PipelineParallelMicroStepLocations,
+    PipelineParallelWithInterleave,
+    PipelineParallelWithInterleaveFthenB,
+    PipelineDatasetPreprocessor,
+    VPPFhenBInBalancedMemory,
+    register_global_pipeline_parallel_hook,
+)
+from .dualpipev import DualPipeVParallel  # noqa: F401
+from .segment_parallel import SegmentParallel  # noqa: F401
+from .sharding_parallel import ShardingParallel  # noqa: F401
+from .tensor_parallel import TensorParallel  # noqa: F401
+from .pp_utils.forward_backward_overlap_utils import (  # noqa: F401
+    ScheduleNode,
+    ScheduleChunk,
+)
+from .pp_utils.utils import (  # noqa: F401
+    dict_to_tuple_helper,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/dualpipev.py

@@ -0,0 +1,851 @@
+# 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.
+
+# The file has been adapted from DeepSeek DualPipe project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DualPipe/blob/main/LICENSE
+
+from __future__ import annotations
+
+import paddle
+from paddle import framework
+from paddle.distributed.communication.batch_isend_irecv import (
+    P2POp,
+    batch_isend_irecv,
+)
+
+try:
+    from paddle.distributed.communication import deep_ep
+except ImportError:
+    deep_ep = None
+
+from ..utils.log_util import logger
+from .pipeline_parallel import (
+    FakeMicroDataset,
+    HybridParallelOptimizer,
+    PipelineDatasetPreprocessor,
+    PipelineParallel,
+)
+from .pp_utils.batch_comm_helper import BatchCommHelper
+from .pp_utils.forward_backward_overlap_utils import ScheduleChunk
+from .zero_bubble_utils import EventStore, WeightGradStore
+
+__all__ = []
+
+
+def detach_and_requires_grad(x):
+    o = x.detach()
+    o.stop_gradient = False
+    return o
+
+
+class DualPipeVParallel(PipelineParallel):
+    """
+    An implementation of the DualPipeV, based on
+    https://github.com/deepseek-ai/DualPipe/blob/main/dualpipe/dualpipe.py.
+    """
+
+    def __init__(self, layers, hcg, strategy):
+        super().__init__(layers=layers, hcg=hcg, strategy=strategy)
+        self.overlapped_forward_backward = hasattr(
+            type(self._layers), "overlapped_forward_backward"
+        )
+        logger.info(
+            f"Using DualPipeVParallel with overlapping forward backward={self.overlapped_forward_backward}"
+        )
+
+        self.num_ranks = self.num_stages
+        self.group_rank = self.pp_group.rank
+        self.prev_rank = self.pp_group.ranks[
+            (self.group_rank - 1) % self.pp_group.world_size
+        ]
+        self.next_rank = self.pp_group.ranks[
+            (self.group_rank + 1) % self.pp_group.world_size
+        ]
+
+        # NOTE(zhangyuqin1998): The first rank has to broadcast the meta information
+        # of the P2P communication after the first forward.
+        self.need_broadcast_meta = self.is_pipeline_first_stage()
+        self.need_recv_meta = not self.is_pipeline_first_stage()
+        self._p2p_helper = BatchCommHelper(self._using_cache)
+
+    def is_pipeline_first_stage(self):
+        return self.group_rank == 0
+
+    def is_pipeline_last_stage(self):
+        return self.group_rank == self.num_ranks - 1
+
+    def _reset_states(self):
+        self.input_tensors = ([], [])
+        self.output_tensors = ([], [])
+        self.input_grad_tensors = ([], [])
+        self.output_grad_tensors = ([], [])
+        self.loss_tensors: list[paddle.Tensor] = []
+        self.schedule_chunks = ([], [])
+        self.loss_fn_chunks = []
+
+        # The first value in the list corresponds to phase 0, and the second value corresponds to phase 1.
+        self.current_f_acc_id = [0, 0]
+        self.current_b_acc_id = [0, 0]
+        self.current_send_f_acc_id = [0, 0]
+        self.current_send_b_acc_id = [0, 0]
+        self.current_recv_f_acc_id = [0, 0]
+        self.current_recv_b_acc_id = [0, 0]
+        self.comm_forward_ops: list[P2POp] = []
+        self.comm_backward_ops: list[P2POp] = []
+        self.to_free: list[paddle.Tensor] = []
+
+    def _get_forward_inputs(self, micro_datasets, phase, acc_id):
+        is_first_stage = self.is_pipeline_first_stage() and phase == 0
+        if is_first_stage:
+            assert micro_datasets is not None
+            self.input_tensors[phase].append(next(micro_datasets[phase])[0])
+        if self.forward_only:
+            self.input_tensors[phase][acc_id] = None
+        return self.input_tensors[phase][acc_id]
+
+    def _get_forward_labels(self, micro_datasets, phase, acc_id):
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if is_last_stage and self._compute_loss:
+            assert micro_datasets is not None
+            labels = next(micro_datasets[phase])[1]
+            self._check_micro_batch_data_valid(labels)
+            return labels
+        else:
+            return None
+
+    def _loss_compute(self, micro_datasets, phase, acc_id, logits):
+        labels = self._get_forward_labels(micro_datasets, phase, acc_id)
+        loss_fn_node = None
+        if not self.overlapped_forward_backward:
+            loss_tensor = self._layers._loss_fn[0](logits, labels)
+        else:
+            loss_fn_node = self._layers._loss_fn[0].build_schedule_node()
+            loss_fn_node.labels = labels
+            loss_tensor = loss_fn_node.forward(logits)
+        self._store_forward_loss(phase, loss_tensor, loss_fn_node)
+
+    def _store_forward_tensors(self, phase, outputs, schedule_chunk):
+        self.schedule_chunks[phase].append(schedule_chunk)
+        if self.is_pipeline_last_stage() and phase == 0:
+            self.input_tensors[1].append(
+                [detach_and_requires_grad(output) for output in outputs]
+            )
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if not is_last_stage:
+            self.output_tensors[phase].append(outputs)
+
+    def _forward_compute(self, phase: int, micro_datasets=None) -> None:
+        acc_id = self.current_f_acc_id[phase]
+        self.current_f_acc_id[phase] += 1
+
+        inputs = self._get_forward_inputs(micro_datasets, phase, acc_id)
+
+        if self.overlapped_forward_backward:
+            schedule_chunk = self._layers.get_schedule_chunk(chunk_id=phase)
+            outputs = schedule_chunk.forward(inputs)
+        else:
+            schedule_chunk = None
+            outputs = self._layers.forward(inputs, chunk_id=phase)
+        outputs = [outputs] if isinstance(outputs, paddle.Tensor) else outputs
+
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if is_last_stage and self._compute_loss:
+            self._loss_compute(micro_datasets, phase, acc_id, outputs)
+        self._store_forward_tensors(phase, outputs, schedule_chunk)
+
+    def _get_backward_inputs(self, phase, acc_id):
+        outputs = self.output_tensors[phase][acc_id]
+        self.output_tensors[phase][acc_id] = None
+        output_grads = self.output_grad_tensors[phase][acc_id]
+        self.output_grad_tensors[phase][acc_id] = None
+        non_empty = [
+            (t, g) for t, g in zip(outputs, output_grads) if g is not None
+        ]
+        outputs, output_grads = list(zip(*non_empty))
+        return outputs, output_grads
+
+    def _store_backward_tensors(self, phase, acc_id, input_grads=None):
+        if input_grads is None:
+            inputs = self.input_tensors[phase][acc_id]
+            input_grads = [
+                t.grad
+                for t in inputs
+                if (t is not None and not t.stop_gradient)
+            ]
+        self.input_tensors[phase][acc_id] = None
+
+        if isinstance(input_grads, paddle.Tensor):
+            input_grads = (input_grads,)
+        if self.is_pipeline_last_stage() and phase == 1:
+            self.output_grad_tensors[0].append(input_grads)
+        else:
+            self.input_grad_tensors[phase].append(input_grads)
+
+    def _store_forward_loss(self, phase, loss_tensor, loss_fn_node=None):
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if is_last_stage and self._compute_loss:
+            if isinstance(loss_tensor, (tuple, list)):
+                assert len(loss_tensor) == 1
+                loss_tensor = loss_tensor[0]
+            assert isinstance(loss_tensor, paddle.Tensor), (
+                "Currently, loss_fn should obtain Paddle.Tensor dtype"
+            )
+
+            self.loss_tensors.append(loss_tensor)
+            self.loss_fn_chunks.append(loss_fn_node)
+
+    def _backward_compute(self, phase: int, enable_zb: bool = False) -> None:
+        if self.forward_only:
+            return
+
+        acc_id = self.current_b_acc_id[phase]
+        self.current_b_acc_id[phase] += 1
+
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+
+        WeightGradStore.enabled = enable_zb
+        input_grads = None
+        with paddle.amp.auto_cast(enable=False):
+            if is_last_stage:
+                loss = self.loss_tensors[acc_id]
+                if self.overlapped_forward_backward:
+                    loss_fn_node = self.loss_fn_chunks[acc_id]
+                    backward_chunk = self.schedule_chunks[phase][acc_id]
+                    _, _, input_grads = (
+                        self._layers.overlapped_forward_backward(
+                            ScheduleChunk([]),  # forward_chunk
+                            None,  # forward_inputs
+                            None,  # forward_loss_fn_node
+                            backward_chunk,
+                            loss_fn_node,
+                            None,  # input_grads
+                            self.scaler,
+                            combine_bw_event_to_wait=None,
+                            pp_stream=None,
+                        )
+                    )
+                    self.loss_fn_chunks[acc_id] = None
+                    self.schedule_chunks[phase][acc_id] = None
+                else:
+                    if self.scaler:
+                        paddle.autograd.backward(self.scaler.scale(loss))
+                    else:
+                        paddle.autograd.backward(loss)
+            else:
+                outputs, output_grads = self._get_backward_inputs(phase, acc_id)
+                if self.overlapped_forward_backward:
+                    backward_chunk = self.schedule_chunks[phase][acc_id]
+                    _, _, input_grads = (
+                        self._layers.overlapped_forward_backward(
+                            ScheduleChunk([]),  # forward_chunk
+                            None,  # forward_inputs
+                            None,  # forward_loss_fn_node
+                            backward_chunk,
+                            None,  # backward_loss_fn_node
+                            output_grads,
+                            None,  # scaler
+                            combine_bw_event_to_wait=None,
+                            pp_stream=None,
+                        )
+                    )
+                    self.schedule_chunks[phase][acc_id] = None
+                else:
+                    if len(outputs) > 0:
+                        outputs = [t for t in outputs if not t.stop_gradient]
+                        paddle.autograd.backward(
+                            tensors=outputs,
+                            grad_tensors=output_grads,
+                        )
+        WeightGradStore.enabled = False
+        if enable_zb:
+            WeightGradStore.flush()
+
+        self._store_backward_tensors(phase, acc_id, input_grads=input_grads)
+
+    def _forward_backward_compute(
+        self,
+        forward_phase: int,
+        backward_phase: int,
+        micro_datasets=None,
+        combine_backward_event_to_wait=None,
+        pass_pp_stream=False,
+    ) -> None:
+        if self.forward_only:
+            self._forward_compute(forward_phase, micro_datasets)
+            return
+
+        if not self.overlapped_forward_backward:
+            self._forward_compute(forward_phase, micro_datasets)
+            self._backward_compute(backward_phase)
+            return
+
+        # pre-forward
+        forward_acc_id = self.current_f_acc_id[forward_phase]
+        self.current_f_acc_id[forward_phase] += 1
+
+        forward_inputs = self._get_forward_inputs(
+            micro_datasets, forward_phase, forward_acc_id
+        )
+        forward_labels = self._get_forward_labels(
+            micro_datasets, forward_phase, forward_acc_id
+        )
+        if forward_labels is not None:
+            forward_loss_fn_node = self._layers._loss_fn[
+                0
+            ].build_schedule_node()
+            forward_loss_fn_node.labels = forward_labels
+        else:
+            forward_loss_fn_node = None
+
+        # pre-backward
+        backward_acc_id = self.current_b_acc_id[backward_phase]
+        self.current_b_acc_id[backward_phase] += 1
+
+        is_last_stage1 = self.is_pipeline_first_stage() and backward_phase == 1
+        if is_last_stage1:
+            backward_loss_fn_node = self.loss_fn_chunks[backward_acc_id]
+            backward_grads = None
+        else:
+            backward_loss_fn_node = None
+            _, backward_grads = self._get_backward_inputs(
+                backward_phase, backward_acc_id
+            )
+
+        # event_to_wait = deep_ep.get_event_from_custom_stream(paddle.device.current_stream().stream_base)
+
+        # forward & backward
+        forward_chunk = self._layers.get_schedule_chunk(chunk_id=forward_phase)
+        backward_chunk = self.schedule_chunks[backward_phase][backward_acc_id]
+        forward_outputs, forward_loss, backward_input_grads = (
+            self._layers.overlapped_forward_backward(
+                forward_chunk,
+                forward_inputs,
+                forward_loss_fn_node,
+                backward_chunk,
+                backward_loss_fn_node,
+                backward_grads,
+                self.scaler,
+                combine_bw_event_to_wait=combine_backward_event_to_wait,
+                pp_stream=(
+                    self.pp_group.process_group.get_stream(
+                        paddle.framework._current_expected_place_()
+                    )
+                    if pass_pp_stream
+                    else None
+                ),
+            )
+        )
+        self.schedule_chunks[backward_phase][backward_acc_id] = None
+
+        # post-forward
+        self._store_forward_tensors(
+            forward_phase, forward_outputs, forward_chunk
+        )
+        self._store_forward_loss(
+            forward_phase, forward_loss, forward_loss_fn_node
+        )
+
+        # post-backward
+        self._store_backward_tensors(
+            backward_phase, backward_acc_id, input_grads=backward_input_grads
+        )
+
+    def _commit_and_wait_comm(
+        self, p2p_overlap=False, use_outer_event_wait=False
+    ) -> None:
+        common_forward_ops_num = (
+            len(self.comm_forward_ops)
+            if self.comm_forward_ops is not None
+            else 0
+        )
+        common_backward_ops_num = (
+            len(self.comm_backward_ops)
+            if self.comm_backward_ops is not None
+            else 0
+        )
+        if common_forward_ops_num == 0 and common_backward_ops_num == 0:
+            if EventStore.event is not None:
+                e_t = EventStore.event
+                EventStore.event = None
+                return e_t
+            return deep_ep.get_event_from_custom_stream(
+                paddle.device.current_stream().stream_base
+            )
+
+        use_stream_wait_event = (
+            p2p_overlap and self._overlap_p2p_comm and deep_ep is not None
+        )
+
+        pp_raw_stream = self.pp_group.process_group.get_stream(
+            paddle.framework._current_expected_place_()
+        )
+        if use_outer_event_wait:
+            self.pp_group.process_group.set_outer_wait(True)
+
+        if common_forward_ops_num > 0:
+            fwd_reqs = batch_isend_irecv(self.comm_forward_ops)
+
+            if not use_stream_wait_event:
+                for req in fwd_reqs:
+                    req.wait()
+
+        if use_outer_event_wait:
+            self.pp_group.process_group.set_outer_wait(False)
+
+        if use_stream_wait_event:
+            forward_event_to_wait = deep_ep.get_event_from_custom_stream(
+                pp_raw_stream
+            )
+
+        backward_outer_event_wait = False
+        if EventStore.event is not None:
+            with paddle.device.stream_guard(
+                paddle.device.Stream(stream_base=pp_raw_stream)
+            ):
+                EventStore.event.current_stream_wait()
+
+            EventStore.set(None)
+            self.pp_group.process_group.set_outer_wait(True)
+
+            backward_outer_event_wait = True
+
+        if common_backward_ops_num > 0:
+            bwd_reqs = batch_isend_irecv(self.comm_backward_ops)
+
+            if not use_stream_wait_event:
+                for req in bwd_reqs:
+                    req.wait()
+
+        if backward_outer_event_wait:
+            self.pp_group.process_group.set_outer_wait(False)
+
+        if use_stream_wait_event:
+            forward_event_to_wait.current_stream_wait()
+
+            combine_bw_event_to_wait = deep_ep.get_event_from_custom_stream(
+                pp_raw_stream
+            )
+        else:
+            combine_bw_event_to_wait = deep_ep.get_event_from_custom_stream(
+                paddle.device.current_stream().stream_base
+            )
+
+        self.comm_forward_ops = []
+        self.comm_backward_ops = []
+
+        self._free_tensors()
+
+        return combine_bw_event_to_wait
+
+    def _weight_pass(self) -> None:
+        if self.forward_only:
+            return
+
+        self._commit_and_wait_comm()
+
+        # Assume FIFO
+        WeightGradStore.pop()
+
+    def _free_tensors(self) -> None:
+        self._release_output(self.to_free)
+        self.to_free = []
+
+    def _recv_forward(self, phase: int) -> None:
+        if (self.is_pipeline_first_stage() and phase == 0) or (
+            self.is_pipeline_last_stage() and phase == 1
+        ):
+            return
+
+        self.current_recv_f_acc_id[phase] += 1
+
+        tensors = self._p2p_helper.append_irecv(
+            self.comm_forward_ops,
+            self.prev_rank if phase == 0 else self.next_rank,
+            self.pp_group,
+            alloc_on_comm_stream=self._overlap_p2p_comm,
+        )
+        self.input_tensors[phase].append(tensors)
+
+    def _send_forward(self, phase: int) -> None:
+        if (self.is_pipeline_first_stage() and phase == 1) or (
+            self.is_pipeline_last_stage() and phase == 0
+        ):
+            return
+
+        acc_id = self.current_send_f_acc_id[phase]
+        self.current_send_f_acc_id[phase] += 1
+        tensors = self.output_tensors[phase][acc_id]
+
+        self._p2p_helper.append_isend(
+            self.comm_forward_ops,
+            tensors,
+            self.next_rank if phase == 0 else self.prev_rank,
+            self.pp_group,
+            self.need_broadcast_meta,
+        )
+        self.need_broadcast_meta = False
+
+        self.to_free.extend(tensors)
+
+    def _recv_backward(self, phase: int) -> None:
+        if self.forward_only:
+            return
+
+        if (self.is_pipeline_first_stage() and phase == 1) or (
+            self.is_pipeline_last_stage() and phase == 0
+        ):
+            return
+
+        self.current_recv_b_acc_id[phase] += 1
+        tensors = self._p2p_helper.append_irecv(
+            self.comm_backward_ops,
+            self.next_rank if phase == 0 else self.prev_rank,
+            self.pp_group,
+            alloc_on_comm_stream=self._overlap_p2p_comm,
+        )
+        self.output_grad_tensors[phase].append(tensors)
+
+    def _send_backward(self, phase: int) -> None:
+        if self.forward_only:
+            return
+
+        if (self.is_pipeline_first_stage() and phase == 0) or (
+            self.is_pipeline_last_stage() and phase == 1
+        ):
+            return
+
+        acc_id = self.current_send_b_acc_id[phase]
+        self.current_send_b_acc_id[phase] += 1
+        tensors = self.input_grad_tensors[phase][acc_id]
+        self.input_grad_tensors[phase][acc_id] = None
+
+        self._p2p_helper.append_isend(
+            self.comm_backward_ops,
+            tensors,
+            self.prev_rank if phase == 0 else self.next_rank,
+            self.pp_group,
+        )
+
+    def _forward_pass(
+        self,
+        phase: int,
+        micro_datasets=None,
+        recv: bool = True,
+        send: bool = True,
+    ) -> None:
+        if recv:
+            self._recv_forward(phase)
+        self._commit_and_wait_comm()
+
+        self._forward_compute(phase, micro_datasets)
+
+        if send:
+            self._send_forward(phase)
+
+    def _backward_pass(
+        self,
+        phase: int,
+        enable_zb: bool = False,
+        recv: bool = True,
+        send: bool = True,
+    ) -> None:
+        if recv:
+            self._recv_backward(phase)
+        self._commit_and_wait_comm()
+
+        self._backward_compute(phase, enable_zb)
+
+        if send:
+            self._send_backward(phase)
+
+    def _forward_backward_pass(
+        self,
+        forward_phase: int,
+        backward_phase: int,
+        micro_datasets=None,
+        recv0: bool = True,
+        first_chunk=False,
+        last_chunk=False,
+        main_stage=False,
+        last_stage_and_first_chunk=False,
+    ) -> None:
+        if recv0:
+            self._recv_forward(forward_phase)
+        self._recv_backward(backward_phase)
+
+        need_send_forward = not (
+            self.is_pipeline_first_stage() and forward_phase == 1
+        ) or (self.is_pipeline_last_stage() and forward_phase == 0)
+        need_send_backward = not (
+            self.is_pipeline_first_stage() and backward_phase == 0
+        ) or (self.is_pipeline_last_stage() and backward_phase == 1)
+
+        use_outer_event_wait = (
+            main_stage
+            and not first_chunk
+            and self._overlap_p2p_comm
+            and deep_ep is not None
+            and (need_send_forward and need_send_backward)
+        )
+
+        pass_pp_stream = (
+            main_stage
+            and not last_chunk
+            and self._overlap_p2p_comm
+            and deep_ep is not None
+            and (need_send_forward and need_send_backward)
+            and (not last_stage_and_first_chunk)
+        )
+
+        combine_bw_wait_event = self._commit_and_wait_comm(
+            not last_chunk, use_outer_event_wait
+        )
+
+        self._forward_backward_compute(
+            forward_phase,
+            backward_phase,
+            micro_datasets,
+            combine_backward_event_to_wait=combine_bw_wait_event,
+            pass_pp_stream=pass_pp_stream,
+        )
+
+        self._send_forward(forward_phase)
+        self._send_backward(backward_phase)
+
+    def _wrap_data(self, data, phase):
+        """
+        for backward compatibility, wrap data to Fake FakeMicroDataset if it is of type list or tuple
+        """
+        if isinstance(data, PipelineDatasetPreprocessor):
+            data = data()
+
+        if (not isinstance(data, tuple)) and (not isinstance(data, list)):
+            return data
+
+        micro_dataset = FakeMicroDataset(
+            data,
+            self.is_pipeline_first_stage() and phase == 0,
+            self.is_pipeline_first_stage() and phase == 1,
+            self.accumulate_steps,
+            self.micro_batch_size,
+        )
+        return micro_dataset
+
+    def _prepare_training(self, data, optimizer, lr_scheduler):
+        assert isinstance(optimizer, HybridParallelOptimizer), (
+            'optimizer should be HybridParallelOptimizer subclass.'
+        )
+
+        assert framework._dygraph_tracer()._has_grad, (
+            'Please enable the generation of gradients.'
+        )
+
+        if self.is_pipeline_first_stage():
+            assert data is not None, (
+                "For the first and the last stage, the data must be set."
+            )
+        else:
+            data = None
+
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+        self._layers.train()
+        self.register_sharding_comm_overlap_hook(optimizer)
+
+        return data
+
+    def _broadcast_final_loss(self):
+        loss_sum_tensor = paddle.zeros([1], "float32")
+        if self.is_pipeline_first_stage():
+            assert len(self.loss_tensors) > 0, (
+                "train_batch() in last stage should obtain valid loss"
+            )
+            for loss in self.loss_tensors:
+                loss_sum_tensor += loss.detach().astype("float32")
+            loss_sum_tensor /= self.accumulate_steps
+
+        paddle.distributed.all_reduce(
+            loss_sum_tensor, group=self.pp_group, sync_op=True
+        )
+        return loss_sum_tensor
+
+    def forward_backward_pipeline(
+        self,
+        data,
+        scaler,
+        forward_only=False,
+        compute_loss=True,
+    ):
+        self.scaler = scaler
+
+        rank = self.group_rank
+        num_ranks = self.num_ranks
+        assert (
+            self.accumulate_steps > 0 and self.accumulate_steps >= num_ranks * 2
+        ), f"{self.accumulate_steps=}, {num_ranks=}"
+        self.forward_only = forward_only
+
+        self._reset_states()
+
+        # NOTE(zhangyuqin1998): Tensors to be sent or received must have a
+        # consistent shape and data type throughout the entire pipeline. We
+        # broadcast the meta info in the first forward of the first rank.
+        self._p2p_helper.recv_meta_from_head(self.pp_group, self.need_recv_meta)
+        self.need_recv_meta = False
+
+        micro_dataset_phase0 = self._wrap_data(data, 0)
+        micro_dataset_phase1 = self._wrap_data(data, 1)
+        micro_datasets = [micro_dataset_phase0, micro_dataset_phase1]
+
+        # Step 1: nF0
+        step_1 = (num_ranks - rank - 1) * 2
+        for i in range(step_1):
+            self._forward_pass(0, micro_datasets)
+
+        # Step 2: nF0F1
+        step_2 = rank + 1
+        self._recv_forward(0)
+        for i in range(step_2):
+            self._forward_pass(0, micro_datasets, recv=False, send=False)
+            self._recv_forward(0)
+            self._forward_pass(
+                1,
+                micro_datasets,
+                send=(not self.is_pipeline_last_stage()) or (i < step_2 - 1),
+            )
+            self._send_forward(0)
+
+        # Step 3: nB1W1F1 (Use zero bubble)
+        step_3 = num_ranks - rank - 1
+        for i in range(step_3):
+            self._backward_pass(1, enable_zb=True)
+            self._recv_forward(1)
+            self._weight_pass()
+            self._forward_pass(1, micro_datasets, recv=False)
+
+        # Step 4 (Main step): nF0B1F1B0
+        step_4 = self.accumulate_steps - num_ranks * 2 + rank + 1
+        have_step5 = num_ranks - rank - 1 > 0
+        # Update code to support send/recv overlap
+        # Only support send/recv overlap in MainStep
+        for i in range(step_4):
+            is_last_chunk = i + 1 == step_4
+            if i == 0:
+                if self.is_pipeline_last_stage():
+                    # NOTE: We don't overlap these two passes to further reduce bubble size.
+                    self._forward_pass(
+                        0, micro_datasets, recv=False, send=False
+                    )
+                    self._send_forward(1)
+                    self._backward_pass(1, send=False)
+                    self._send_forward(0)
+                    self._send_backward(1)
+
+                    self._forward_backward_pass(
+                        1,
+                        0,
+                        micro_datasets,
+                        first_chunk=True,
+                        last_chunk=is_last_chunk,
+                        main_stage=True,
+                    )
+                else:
+                    self._forward_backward_pass(
+                        0,
+                        1,
+                        micro_datasets,
+                        recv0=False,
+                        first_chunk=True,
+                        main_stage=True,
+                    )
+
+                    self._forward_backward_pass(
+                        1,
+                        0,
+                        micro_datasets,
+                        last_chunk=is_last_chunk,
+                        main_stage=True,
+                    )
+            else:
+                self._forward_backward_pass(
+                    0,
+                    1,
+                    micro_datasets,
+                    main_stage=True,
+                    last_stage_and_first_chunk=self.is_pipeline_last_stage(),
+                )
+                self._forward_backward_pass(
+                    1,
+                    0,
+                    micro_datasets,
+                    last_chunk=is_last_chunk,
+                    main_stage=True,
+                )
+
+        # Step 5: nB1F1B0
+        step_5 = num_ranks - rank - 1
+        for i in range(step_5):
+            self._backward_pass(1)
+            self._forward_backward_pass(1, 0, micro_datasets)
+
+        # Step 6: nB1B0 (The second half of the passes use zero bubble)
+        step_6 = rank + 1
+        enable_zb = False
+        for i in range(step_6):
+            if i == step_6 // 2 and rank % 2 == 1:
+                enable_zb = True
+            self._backward_pass(1, enable_zb=enable_zb)
+            if i == step_6 // 2 and rank % 2 == 0:
+                enable_zb = True
+            self._backward_pass(0, enable_zb=enable_zb)
+
+        # Step 7: nWB0 (Use zero bubble)
+        step_7 = num_ranks - rank - 1
+        for i in range(step_7):
+            self._weight_pass()
+            self._backward_pass(0, enable_zb=True)
+
+        # Step 8: nW
+        step_8 = rank + 1
+        for i in range(step_8):
+            self._weight_pass()
+        assert WeightGradStore.funcs_queue.empty()
+
+        self._commit_and_wait_comm()
+
+        self._layers.allreduce_shared_weight_gradients()
+
+        with paddle.amp.auto_cast(enable=False):
+            train_loss = self._broadcast_final_loss()
+
+        self._reset_states()
+        return train_loss
+
+    def train_batch(
+        self,
+        data,
+        optimizer,
+        lr_scheduler=None,
+        scaler=None,
+    ):
+        data = self._prepare_training(data, optimizer, lr_scheduler)
+
+        train_loss = self.forward_backward_pipeline(data, scaler)
+
+        # optimizer
+        with paddle.amp.auto_cast(enable=False):
+            self._optimizer_step()
+
+        return train_loss

python/paddle/distributed/fleet/meta_parallel/meta_parallel_base.py

@@ -0,0 +1,44 @@
+#   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 import nn
+
+__all__ = []
+
+
+class MetaParallelBase(nn.Layer):
+    def __init__(self, layers, hcg, strategy):
+        super().__init__(layers.full_name() + "_meta_parallel_base")
+        self._layers = layers
+        self._hcg = hcg
+        self._strategy = strategy
+        self._prepare_for_model()
+
+    def _prepare_for_model(self):
+        pass
+
+    def _pre_forward(self, *inputs, **kwargs):
+        pass
+
+    def forward(self, *inputs, **kwargs):
+        self._pre_forward(*inputs, **kwargs)
+
+        output = self._layers(*inputs, **kwargs)
+
+        self._post_forward(output)
+
+        return output
+
+    def _post_forward(self, output):
+        pass

python/paddle/distributed/fleet/meta_parallel/parallel_layers/__init__.py

@@ -0,0 +1,39 @@
+# 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 .mp_layers import (  # noqa: F401
+    ColumnParallelLinear,
+    ParallelCrossEntropy,
+    RowParallelLinear,
+    VocabParallelEmbedding,
+)
+from .pp_layers import (  # noqa: F401
+    LayerDesc,
+    LocalSharedLayerDesc,
+    PipelineLayer,
+    SharedLayerDesc,
+)
+from .random import (  # noqa: F401
+    RNGStatesTracker,
+    get_rng_state_tracker,
+    model_parallel_random_seed,
+)
+from .spec_utils import (
+    LayerSpec as LayerSpec,
+    build_spec_layer as build_spec_layer,
+    get_spec_layer as get_spec_layer,
+    import_spec_layer as import_spec_layer,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/parallel_layers/mp_layers.py

@@ -0,0 +1,22 @@
+#   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 ...layers.mpu.mp_layers import (  # noqa: F401
+    ColumnParallelLinear,
+    ParallelCrossEntropy,
+    RowParallelLinear,
+    VocabParallelEmbedding,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/parallel_layers/random.py

@@ -0,0 +1,22 @@
+#   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 ...layers.mpu.random import (  # noqa: F401
+    RNGStatesTracker,
+    dropout,
+    get_rng_state_tracker,
+    model_parallel_random_seed,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/parallel_layers/spec_utils.py

@@ -0,0 +1,141 @@
+# 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.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from __future__ import annotations
+
+import types
+import warnings
+from dataclasses import dataclass, field
+
+
+@dataclass
+class LayerSpec:
+    """This is a Layer Specification dataclass.
+
+    Specification defines the location of the layer (to import dynamically)
+    or the imported layer itself. It also defines the extra_kwargs that need to be
+    passed to initialize the layer.
+
+    Args:
+        layer (tuple | type): A tuple describing the location of the
+            layer class e.g. `(layer.location, LayerClass)` or the imported
+            layer class itself e.g. `LayerClass` (which is already imported
+            using `from layer.location import LayerClass`).
+        extra_kwargs (dict): A dictionary of extra_kwargs that need to be passed while init.
+
+    """
+
+    layer: tuple | type
+    extra_kwargs: dict = field(default_factory=lambda: {})
+    sublayers_spec: type = None
+
+    def __repr__(self):
+        rst = ""
+        if isinstance(self.layer, tuple):
+            for sub_layer in self.layer:
+                rst = rst + repr(sub_layer) + ","
+        else:
+            rst = repr(self.layer) + repr(self.extra_kwargs)
+        return rst
+
+
+def import_spec_layer(layer_path: tuple[str]):
+    """Import a named object from a layer in the context of this function."""
+    base_path, name = layer_path
+    try:
+        layer = __import__(base_path, globals(), locals(), [name])
+    except ImportError as e:
+        print(f"couldn't import layer due to {e}")
+        return None
+    return vars(layer)[name]
+
+
+def get_spec_layer(spec_or_layer: LayerSpec | type, **additional_kwargs):
+    # If a layer class is already provided return it as is
+    if isinstance(spec_or_layer, (type, types.FunctionType)):
+        return spec_or_layer
+
+    # If the layer is provided instead of layer path, then return it as is
+    if isinstance(spec_or_layer.layer, (type, types.FunctionType)):
+        return spec_or_layer.layer
+
+    # Otherwise, return the dynamically imported layer from the layer path
+    return import_spec_layer(spec_or_layer.layer)
+
+
+def build_spec_layer(spec_or_layer: LayerSpec | type, *args, **kwargs):
+    # If the passed `spec_or_layer` is
+    # a `Function`, then return it as it is
+    # NOTE: to support an already initialized layer add the following condition
+    # `or isinstance(spec_or_layer, paddle.nn.Layer)` to the following if check
+    if isinstance(spec_or_layer, types.FunctionType):
+        return spec_or_layer
+
+    # If the passed `spec_or_layer` is actually a spec (instance of
+    # `LayerSpec`) and it specifies a `Function` using its `layer`
+    # field, return the `Function` as it is
+    if isinstance(spec_or_layer, LayerSpec) and isinstance(
+        spec_or_layer.layer, types.FunctionType
+    ):
+        return spec_or_layer.layer
+
+    # Check if a layer class is provided as a spec or if the layer path
+    # itself is a class
+    if isinstance(spec_or_layer, type):
+        layer = spec_or_layer
+    elif hasattr(spec_or_layer, "layer") and isinstance(
+        spec_or_layer.layer, type
+    ):
+        layer = spec_or_layer.layer
+    else:
+        # Otherwise, dynamically import the layer from the layer path
+        layer = import_spec_layer(spec_or_layer.layer)
+
+    # If the imported layer is actually a `Function` return it as it is
+    if isinstance(layer, types.FunctionType):
+        return layer
+
+    # Finally return the initialized layer with extra_kwargs from the spec as well
+    # as those passed as **kwargs from the code
+
+    # Add the `sublayers_spec` argument to the layer init call if it exists in the
+    # spec.
+    if (
+        hasattr(spec_or_layer, "sublayers_spec")
+        and spec_or_layer.sublayers_spec is not None
+    ):
+        kwargs["sublayers_spec"] = spec_or_layer.sublayers_spec
+    if hasattr(spec_or_layer, "extra_kwargs"):
+        for key in spec_or_layer.extra_kwargs.keys():
+            if key in kwargs:
+                warnings.warn(
+                    f"Got same key {key} in extra_kwargs and kwargs during init {layer.__name__}. Will keep the value ing extra_kwargs."
+                )
+                kwargs.pop(key)
+    try:
+        return layer(
+            *args,
+            **spec_or_layer.extra_kwargs
+            if hasattr(spec_or_layer, "extra_kwargs")
+            else {},
+            **kwargs,
+        )
+    except Exception as e:
+        # improve the error message since we hide the layer name in the line above
+        import sys
+
+        raise type(e)(
+            f"{e!s} when instantiating {layer.__name__}"
+        ).with_traceback(sys.exc_info()[2])

python/paddle/distributed/fleet/meta_parallel/pipeline_hooks.py

@@ -0,0 +1,56 @@
+# 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
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class PipelineHook:
+    def __init__(self):
+        self.hooks: dict[int, list[Callable]] = defaultdict(list)
+        self._hooks_capacity = 0
+        self.reset_current_id()
+
+    def reset_current_id(self):
+        self._current_id = 0
+
+    def set_hooks_capacity(self, capacity: int):
+        self._hooks_capacity = capacity
+
+    def register_hook(self, hook_id: int, hook: Callable):
+        assert hook_id < self._hooks_capacity, (
+            f"hook_id {hook_id} is out of range, maximum capacity is {self._hooks_capacity}."
+        )
+        self.hooks[hook_id].append(hook)
+
+    def run_hook(self):
+        assert self._current_id < self._hooks_capacity, (
+            f"hook_id {self._current_id} is out of range, maximum capacity is {self._hooks_capacity}."
+        )
+        for hook in self.hooks[self._current_id]:
+            hook(self._current_id)
+        self._current_id += 1
+
+    @property
+    def current_id(self):
+        return self._current_id
+
+    @property
+    def hooks_capacity(self):
+        return self._hooks_capacity

python/paddle/distributed/fleet/meta_parallel/pp_utils/__init__.py

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

python/paddle/distributed/fleet/meta_parallel/pp_utils/batch_comm_helper.py

@@ -0,0 +1,133 @@
+# 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 paddle
+from paddle.distributed.communication.batch_isend_irecv import (
+    P2POp,
+)
+
+from .p2p_communication import SendRecvMeta
+from .utils import (
+    number_2_dtype,
+    paddle_2_number,
+)
+
+
+class BatchCommHelper:
+    # NOTE(zhangyuqin1998): Tensors to be sent or received must have a
+    # consistent shape and data type throughout the entire pipeline.
+    def __init__(self, use_cache=True):
+        self._send_recv_meta = SendRecvMeta()
+        self._use_cache = use_cache
+
+    def clear_meta_cache(self):
+        self._send_recv_meta.init_or_erase_meta()
+
+    def _send_meta(self, tensors, group, broadcast=False):
+        self._send_recv_meta.set_send_message(tensors)
+        self._send_recv_meta.send_meta(tensors, group, broadcast=broadcast)
+        self._send_recv_meta.recv_shape_message = (
+            self._send_recv_meta.send_shape_message
+        )
+        self._send_recv_meta.recv_dtype_message = (
+            self._send_recv_meta.send_dtype_message
+        )
+
+    def _recv_meta(self, group, broadcast=False):
+        self._send_recv_meta.recv_meta(group, broadcast=broadcast)
+
+    def _build_from_meta(self):
+        shape_message = self._send_recv_meta.recv_shape_message
+        dtype_message = self._send_recv_meta.recv_dtype_message
+        stop_gradient = self._send_recv_meta.recv_stop_gradient
+        assert (shape_message is not None) and (dtype_message is not None), (
+            "Failed to build from meta."
+        )
+
+        res = []
+        if isinstance(shape_message, tuple):
+            for idx, shape in enumerate(shape_message):
+                tmp = paddle.empty(
+                    shape=shape, dtype=number_2_dtype(dtype_message[idx])
+                )
+                tmp.stop_gradient = (
+                    stop_gradient[idx] if stop_gradient is not None else False
+                )
+                res.append(tmp)
+        else:
+            tmp = paddle.empty(
+                shape=shape_message, dtype=number_2_dtype(dtype_message)
+            )
+            tmp.stop_gradient = stop_gradient
+            res.append(tmp)
+        return res
+
+    def _check_valid(self, tensors):
+        shape_message = self._send_recv_meta.recv_shape_message
+        dtype_message = self._send_recv_meta.recv_dtype_message
+
+        assert (shape_message is not None) and (dtype_message is not None), (
+            "Failed to build from meta."
+        )
+
+        if isinstance(shape_message, tuple):
+            assert isinstance(tensors, (list, tuple))
+            assert len(tensors) == len(shape_message)
+            for idx, (shape, dtype, tensor) in enumerate(
+                zip(shape_message, dtype_message, tensors)
+            ):
+                assert tensor.shape == shape, "Invalid shape."
+                assert number_2_dtype(
+                    paddle_2_number(tensor.dtype)
+                ) == number_2_dtype(dtype), "Invalid dtype."
+        else:
+            if isinstance(tensors, (list, tuple)):
+                assert len(tensors) == 1
+                tensors = tensors[0]
+
+            assert tensors.shape == shape_message, "Invalid shape."
+            assert number_2_dtype(
+                paddle_2_number(tensors.dtype)
+            ) == number_2_dtype(dtype_message), "Invalid dtype."
+
+    def recv_meta_from_head(self, group, need_recv_meta):
+        if not need_recv_meta:
+            return
+        self._recv_meta(group, broadcast=True)
+
+    def append_irecv(self, ops, src, group, alloc_on_comm_stream=False):
+        if alloc_on_comm_stream:
+            send_recv_stream = paddle.device.Stream(
+                stream_base=group.process_group.get_stream(
+                    paddle.framework._current_expected_place_()
+                )
+            )
+            with paddle.device.stream_guard(send_recv_stream):
+                tensors = self._build_from_meta()
+        else:
+            tensors = self._build_from_meta()
+        for tensor in tensors:
+            if tensor is not None:
+                ops.append(P2POp(paddle.distributed.irecv, tensor, src, group))
+        return tensors
+
+    def append_isend(self, ops, tensors, dst, group, need_broadcast_meta=False):
+        if need_broadcast_meta:
+            self._send_meta(tensors, group, broadcast=True)
+        self._check_valid(tensors)
+        for tensor in tensors:
+            if tensor is not None:
+                ops.append(P2POp(paddle.distributed.isend, tensor, dst, group))

python/paddle/distributed/fleet/meta_parallel/pp_utils/forward_backward_overlap_utils.py

@@ -0,0 +1,201 @@
+# 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
+
+from .utils import dict_to_tuple_helper
+
+
+class ScheduleChunk:
+    # NOTE(zhangyuqin): ScheduleChunk is the atomic unit of pipeline scheduling.
+    # A ScheduleChunk can contain several ScheduleNodes
+    def __init__(self, nodes):
+        self.nodes = nodes
+        self._check_nodes_valid()
+
+    def forward(self, inputs):
+        for n in self.nodes:
+            inputs = n.forward(inputs)
+        return inputs
+
+    def backward(self, output_grad):
+        for n in reversed(self.nodes):
+            output_grad = n.backward(output_grad)
+        return output_grad
+
+    def _check_nodes_valid(self):
+        for n in self.nodes:
+            assert isinstance(n, (ScheduleNode, ScheduleChunk))
+
+
+def detach_and_requires_grad(inputs):
+    if isinstance(inputs, (tuple, list)):
+        is_tuple = isinstance(inputs, tuple)
+        ret = []
+        for input in inputs:
+            if isinstance(input, (tuple, list)):
+                ret.append(detach_and_requires_grad(input))
+            elif isinstance(input, paddle.Tensor):
+                tmp = input.detach() if input is not None else None
+                if tmp is not None:
+                    tmp.stop_gradient = input.stop_gradient
+                ret.append(tmp)
+            else:
+                ret.append(input)
+        if is_tuple:
+            ret = tuple(ret)
+        return ret
+    elif isinstance(inputs, dict):
+        ret = {}
+        for key in inputs.keys():
+            input = inputs[key]
+            tmp = input.detach() if input is not None else None
+            if tmp is not None:
+                tmp.stop_gradient = input.stop_gradient
+            ret[key] = tmp
+        return ret
+    else:
+        tmp = inputs.detach()
+        tmp.stop_gradient = inputs.stop_gradient
+        return tmp
+
+
+def clone_and_clear_dataptr(outputs, clear_dataptr=False):
+    if isinstance(outputs, (tuple, list)):
+        is_tuple = isinstance(outputs, tuple)
+        ret = [
+            FakeClone.apply(o)
+            for o in outputs
+            if o is not None and isinstance(o, paddle.Tensor)
+        ]
+
+        if clear_dataptr:
+            for o in ret:
+                o._clear_dataptr()
+        if is_tuple:
+            ret = tuple(ret)
+        return ret
+    elif isinstance(outputs, dict):
+        ret = {}
+        for key in outputs.keys():
+            o = outputs[key]
+            if o is not None and isinstance(o, paddle.Tensor):
+                ret[key] = FakeClone.apply(o)
+        if clear_dataptr:
+            for key in ret:
+                ret[key]._clear_dataptr()
+        return ret
+    else:
+        ret = FakeClone.apply(outputs)
+        if clear_dataptr:
+            ret._clear_dataptr()
+        return ret
+
+
+class FakeClone(paddle.autograd.PyLayer):
+    # NOTE(zhangyuqin): Some input tensors may not be used in the forward function, but their gradients
+    # need to be retained. Therefore, we need a clone here. To avoid the DtoD copy, we need a FakeClone
+    @staticmethod
+    def forward(ctx, input):
+        return paddle.empty_like(input)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class ScheduleNode:
+    # NOTE(zhangyuqin): ScheduleNode is a subgraph of the pipeline, capable of independently calling
+    # forward and backward. Users should not use paddle.autograd.backward on the results of ScheduleNode.forward.
+    # Instead, they should use ScheduleNode.backward. Otherwise, resource leakage may occur.
+    def __init__(self, fwd_func, name=""):
+        self.name = name
+        self.fwd_func = fwd_func
+        self.inputs = None
+        self.outputs = None
+
+        self.labels = None
+        self.scale_loss_factor = None
+
+    def forward(self, inputs=(), **kwargs):
+        detached_inputs = detach_and_requires_grad(inputs)
+        self.inputs = detached_inputs
+        if self.labels is not None:
+            outputs = self.fwd_func(self.inputs, self.labels, **kwargs)
+        else:
+            outputs = self.fwd_func(self.inputs, **kwargs)
+        if self.scale_loss_factor is not None:
+            outputs /= self.scale_loss_factor
+
+        # Do not release the loss tensor.
+        clear_dataptr = self.labels is None
+        self.outputs = clone_and_clear_dataptr(outputs, clear_dataptr)
+        return outputs
+
+    def backward(self, output_grad=None, scaler=None):
+        if output_grad is None:
+            if isinstance(self.outputs, (tuple, list)):
+                assert len(self.outputs) == 1
+                outputs = self.outputs[0]
+            else:
+                outputs = self.outputs
+            assert isinstance(outputs, paddle.Tensor)
+            if scaler is not None:
+                paddle.autograd.backward(scaler.scale(outputs))
+            else:
+                paddle.autograd.backward(outputs)
+        else:
+            # Record the original type (tuple or list) to preserve it after filtering
+            is_output_grad_tuple = isinstance(output_grad, tuple)
+            if not isinstance(output_grad, (tuple, list)):
+                is_output_grad_tuple = True  # Single value becomes tuple
+                output_grad = (output_grad,)
+
+            outputs = dict_to_tuple_helper(self.outputs)
+            if not isinstance(outputs, (tuple, list)):
+                outputs = (outputs,)
+            outputs = [t for t in outputs if not t.stop_gradient]
+
+            # Filter None values from output_grad
+            output_grad = [grad for grad in output_grad if grad is not None]
+            # Preserve original type (tuple or list)
+            output_grad = (
+                tuple(output_grad)
+                if is_output_grad_tuple
+                else list(output_grad)
+            )
+
+            assert len(outputs) == len(output_grad), (
+                f"{len(outputs)} of {type(outputs[0])} vs {len(output_grad)} of {type(output_grad[0])}"
+            )
+
+            paddle.autograd.backward(outputs, output_grad)
+
+        inputs = dict_to_tuple_helper(self.inputs)
+        if not isinstance(inputs, (tuple, list)):
+            inputs = (inputs,)
+        grad = tuple([e.grad if e is not None else None for e in inputs])
+        # grad = tuple([e.grad if e is not None and not e.stop_gradient else None for e in inputs])
+        self._reset_states()
+
+        # if len(grad) == 1:
+        #     grad = grad[0]
+        return grad
+
+    def _reset_states(self):
+        self.inputs = None
+        self.outputs = None
+        self.labels = None
+        self.scale_loss_factor = None

python/paddle/distributed/fleet/meta_parallel/pp_utils/four_directions_p2p_communication.py

@@ -0,0 +1,870 @@
+# 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
+
+import numpy as np
+
+import paddle
+from paddle import framework
+
+from ...utils import timer_helper as timer
+from ...utils.log_util import logger
+from .utils import number_2_dtype, paddle_2_number
+
+_hcg = None
+_enable_partial_send_recv = True
+_timers = None
+
+_xpu_comm_group_started = False
+
+_sync_send = os.environ.get("PADDLE_P2P_SYNC_SEND", "0")
+_sync_send = _sync_send.lower() in ['1', 'true']
+
+
+def _xpu_comm_group_start():
+    if not paddle.is_compiled_with_xpu():
+        return
+    global _xpu_comm_group_started
+    assert not _xpu_comm_group_started
+    framework.core.ProcessGroupBKCL.group_start()
+    _xpu_comm_group_started = True
+
+
+def _xpu_comm_group_end():
+    if not paddle.is_compiled_with_xpu():
+        return
+    global _xpu_comm_group_started
+    if _xpu_comm_group_started:
+        framework.core.ProcessGroupBKCL.group_end()
+        _xpu_comm_group_started = False
+
+
+def initialize_p2p_groups(
+    hcg, enable_partial_send_recv=True, enable_timer=False
+):
+    global _hcg, _enable_partial_send_recv, _timers
+    _hcg = hcg
+    _enable_partial_send_recv = enable_partial_send_recv
+    if enable_timer:
+        _timers = timer.get_timers()
+    (
+        send_next_group,
+        send_prev_group,
+        recv_next_group,
+        recv_prev_group,
+    ) = _hcg.get_p2p_groups()
+
+    debug_str = (
+        f"P2pInfo: send_next_group: {send_next_group!r}, send_prev_group: {send_prev_group!r}, "
+        f"recv_next_group: {recv_next_group!r}, recv_prev_group: {recv_prev_group!r}"
+    )
+    logger.info(debug_str)
+
+
+class SendRecvMeta:
+    """Mainly used to help p2p communication context information"""
+
+    def __init__(self):
+        self.send_shape_message = None
+        self.send_dtype_message = None
+
+        self.recv_shape_message = None
+        self.recv_dtype_message = None
+        self.recv_stop_gradient = None
+
+        self.has_send_meta = False
+        self.has_recv_meta = False
+
+    def _recv_shape_dtype(self, group):
+        # recv len(shape)
+        dims = paddle.to_tensor([0])
+        src_rank = _hcg._get_p2p_prev_rank()
+
+        paddle.distributed.recv(dims, src=src_rank, group=group)
+        dims = dims.item()
+
+        # recv shape
+        shape = paddle.to_tensor([0] * dims)
+        paddle.distributed.recv(shape, src=src_rank, group=group)
+
+        # recv dtype
+        dtype = paddle.to_tensor([0])
+        paddle.distributed.recv(dtype, src=src_rank, group=group)
+
+        # recv stop_gradient
+        stop_grad = paddle.to_tensor([0])
+        paddle.distributed.recv(stop_grad, src=src_rank, group=group)
+        return shape.tolist(), dtype.item(), stop_grad.item()
+
+    def recv_meta(self, group):
+        tensor_type = paddle.to_tensor([0])
+        src_rank = _hcg._get_p2p_prev_rank()
+
+        paddle.distributed.recv(tensor_type, src=src_rank, group=group)
+        tensor_type = tensor_type.item()
+
+        if tensor_type == 0:
+            shape, dtype, stop_grad = self._recv_shape_dtype(group)
+            self.recv_shape_message = shape
+            self.recv_dtype_message = dtype
+            self.recv_stop_gradient = bool(stop_grad)
+
+        elif tensor_type == 1:
+            num = paddle.to_tensor([0])
+            paddle.distributed.recv(num, src=src_rank, group=group)
+            num = num.item()
+            shapes = []
+            dtypes = []
+            stop_grads = []
+            for i in range(num):
+                shape, dtype, stop_grad = self._recv_shape_dtype(group)
+                shapes.append(shape)
+                dtypes.append(dtype)
+                stop_grads.append(bool(stop_grad))
+
+            self.recv_shape_message = tuple(shapes)
+            self.recv_dtype_message = tuple(dtypes)
+            self.recv_stop_gradient = tuple(stop_grads)
+
+    def _send_dims_shape_dtype(self, tensor, group):
+        # send len(shape)
+        dims = paddle.to_tensor([len(tensor.shape)])
+        dst_rank = _hcg._get_p2p_next_rank()
+
+        paddle.distributed.send(dims, dst=dst_rank, group=group)
+
+        # send shape
+        shape = paddle.to_tensor(tensor.shape)
+        paddle.distributed.send(shape, dst=dst_rank, group=group)
+
+        # send dtype
+        dtype = paddle.to_tensor([paddle_2_number(tensor.dtype)])
+        paddle.distributed.send(dtype, dst=dst_rank, group=group)
+
+        # send trainable
+        stop_grad = paddle.to_tensor([int(tensor.stop_gradient)])
+        paddle.distributed.send(stop_grad, dst=dst_rank, group=group)
+
+    def send_meta(self, tensor, group):
+        dst_rank = _hcg._get_p2p_next_rank()
+
+        if isinstance(tensor, paddle.Tensor):
+            tensor_type = paddle.to_tensor([0])
+            # send tensor type
+            paddle.distributed.send(tensor_type, dst=dst_rank, group=group)
+
+            self._send_dims_shape_dtype(tensor, group)
+        elif isinstance(tensor, tuple):
+            tensor_type = paddle.to_tensor([1])
+            # send tensor type
+            paddle.distributed.send(tensor_type, dst=dst_rank, group=group)
+
+            nums = paddle.to_tensor([len(tensor)])
+            paddle.distributed.send(nums, dst=dst_rank, group=group)
+
+            for d in tensor:
+                assert isinstance(d, paddle.Tensor)
+                self._send_dims_shape_dtype(d, group=group)
+
+    def set_send_message(self, tensor):
+        if isinstance(tensor, paddle.Tensor):
+            self.send_shape_message = tensor.shape
+            self.send_dtype_message = paddle_2_number(tensor.dtype)
+        elif isinstance(tensor, tuple):
+            self.send_shape_message = tuple(
+                [d.shape for d in tensor if not d.stop_gradient]
+            )
+            self.send_dtype_message = tuple(
+                [
+                    paddle_2_number(d.dtype)
+                    for d in tensor
+                    if not d.stop_gradient
+                ]
+            )
+
+
+def _is_valid_send_recv_partial(tensor, mp_degree):
+    if not _enable_partial_send_recv:
+        return False
+    tensor_numel = np.prod(tensor.shape)
+    assert tensor_numel != 0, "can't send/recv zero element"
+    return mp_degree > 1 and tensor_numel % mp_degree == 0
+
+
+def _partial_send_op(
+    tensor, group, use_calc_stream, ring_id, dst, nranks, rank_id
+):
+    dst_rank_in_group = dst if group is None else group.get_group_rank(dst)
+    if framework.in_dynamic_mode():
+        group = (
+            paddle.distributed.collective._get_default_group()
+            if group is None
+            else group
+        )
+        comm_op = (
+            group.process_group.send_partial_on_calc_stream
+            if use_calc_stream
+            else group.process_group.send_partial
+        )
+        return comm_op(tensor, dst_rank_in_group, nranks, rank_id)
+
+
+def send_partial(
+    tensor, dst=0, nranks=1, rank_id=0, group=None, use_calc_stream=True
+):
+    # dst: local rank in group
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+
+    dst_rank = (
+        _hcg._get_p2p_next_rank() if dst == 1 else _hcg._get_p2p_prev_rank()
+    )
+
+    if _is_valid_send_recv_partial(tensor, nranks):
+        return _partial_send_op(
+            tensor, group, use_calc_stream, ring_id, dst_rank, nranks, rank_id
+        )
+    else:
+        send_op = paddle.distributed.isend
+        return send_op(tensor.detach(), dst=dst_rank, group=group)
+
+
+def _partial_recv_op(
+    tensor, group, use_calc_stream, ring_id, src, nranks, rank_id
+):
+    src_rank_in_group = src if group is None else group.get_group_rank(src)
+    group = (
+        paddle.distributed.collective._get_default_group()
+        if group is None
+        else group
+    )
+    comm_op = (
+        group.process_group.recv_partial_on_calc_stream
+        if use_calc_stream
+        else group.process_group.recv_partial
+    )
+    return comm_op(tensor, src_rank_in_group, nranks, rank_id)
+
+
+def recv_partial(
+    tensor, src=0, nranks=1, rank_id=0, group=None, use_calc_stream=True
+):
+    # src: local rank in group
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+
+    src_rank = (
+        _hcg._get_p2p_prev_rank() if src == 0 else _hcg._get_p2p_next_rank()
+    )
+
+    if _is_valid_send_recv_partial(tensor, nranks):
+        return _partial_recv_op(
+            tensor, group, use_calc_stream, ring_id, src_rank, nranks, rank_id
+        )
+    else:
+        if use_calc_stream:
+            recv_op = paddle.distributed.recv
+        elif framework.in_dynamic_mode():
+            recv_op = paddle.distributed.irecv
+        return recv_op(tensor.detach(), src=src_rank, group=group)
+
+
+def _partial_allgather_op(
+    tensor, group, use_calc_stream, ring_id, nranks, rank_id
+):
+    group = (
+        paddle.distributed.collective._get_default_group()
+        if group is None
+        else group
+    )
+    comm_op = (
+        group.process_group.all_gather_partial_on_calc_stream
+        if use_calc_stream
+        else group.process_group.all_gather_partial
+    )
+    return comm_op(tensor, tensor, nranks, rank_id)
+
+
+def allgather_partial(
+    tensor, nranks=1, rank_id=0, group=None, use_calc_stream=True
+):
+    if not _is_valid_send_recv_partial(tensor, nranks):
+        return tensor
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+
+    return _partial_allgather_op(
+        tensor, group, use_calc_stream, ring_id, nranks, rank_id
+    )
+
+
+def _p2p_helper(
+    tensor_send_next,
+    tensor_send_prev,
+    recv_prev,
+    recv_next,
+    sync_recv=True,
+    send_recv_meta=None,
+):
+    global _hcg
+
+    tensor_recv_prev = None
+    tensor_recv_next = None
+
+    # send / recv message
+    assert send_recv_meta is not None, "send_recv_meta should not be None"
+    recv_shape_msg = send_recv_meta.recv_shape_message
+    recv_dtype_msg = send_recv_meta.recv_dtype_message
+    recv_stop_gradient = send_recv_meta.recv_stop_gradient
+
+    send_shape_msg = send_recv_meta.send_shape_message
+    send_dtype_msg = send_recv_meta.send_dtype_message
+
+    # model parallel message
+    mp_group = _hcg.get_model_parallel_group()
+    mp_degree = _hcg.get_model_parallel_world_size()
+    mp_rank = _hcg.get_model_parallel_rank()
+
+    if recv_prev:
+        if isinstance(recv_shape_msg, tuple):
+            tensor_recv_prev = []
+            for idx, shape in enumerate(recv_shape_msg):
+                tmp = paddle.empty(
+                    shape=shape, dtype=number_2_dtype(recv_dtype_msg[idx])
+                )
+                tmp.stop_gradient = recv_stop_gradient[idx]
+                tensor_recv_prev.append(tmp)
+            tensor_recv_prev = tuple(tensor_recv_prev)
+        else:
+            tensor_recv_prev = paddle.empty(
+                shape=recv_shape_msg, dtype=number_2_dtype(recv_dtype_msg)
+            )
+            tensor_recv_prev.stop_gradient = recv_stop_gradient
+
+    if recv_next:
+        if isinstance(send_shape_msg, tuple):
+            tensor_recv_next = []
+            for idx, shape in enumerate(send_shape_msg):
+                tensor_recv_next.append(
+                    paddle.empty(
+                        shape=shape, dtype=number_2_dtype(send_dtype_msg[idx])
+                    )
+                )
+            tensor_recv_next = tuple(tensor_recv_next)
+        else:
+            tensor_recv_next = paddle.empty(
+                shape=send_shape_msg, dtype=number_2_dtype(send_dtype_msg)
+            )
+
+    # TODO(Yuang Liu): use batch_isend_irecv replace all these comm ops
+    tasks = []
+    # start to p2p communicate
+
+    if _sync_send:
+        # Some devices(NPU for example) do not support asynchronized send op, So the order is
+        # recv_prev -> send_next -> recv_next -> send_prev
+        # When using this order, the environment variable
+        # 'PADDLE_P2P_SYNC_SEND' should be set True
+        if tensor_recv_prev is not None:
+            if isinstance(tensor_recv_prev, tuple):
+                for d in tensor_recv_prev:
+                    task = recv_partial(
+                        d,
+                        src=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_prev_group,
+                        use_calc_stream=sync_recv,
+                    )
+                    if sync_recv:
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+            else:
+                task = recv_partial(
+                    tensor_recv_prev,
+                    src=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_prev_group,
+                    use_calc_stream=sync_recv,
+                )
+
+                if sync_recv:
+                    allgather_partial(
+                        tensor_recv_prev,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+
+        if tensor_send_next is not None:
+            if isinstance(tensor_send_next, tuple):
+                for d in tensor_send_next:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_next_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_next, use_calc_stream=True)
+                send_partial(
+                    tensor_send_next,
+                    dst=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_next_group,
+                    use_calc_stream=False,
+                )
+
+        if tensor_recv_next is not None:
+            if isinstance(tensor_recv_next, tuple):
+                for d in tensor_recv_next:
+                    task = recv_partial(
+                        d,
+                        src=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_next_group,
+                        use_calc_stream=sync_recv,
+                    )
+
+                    if sync_recv:
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+
+            else:
+                task = recv_partial(
+                    tensor_recv_next,
+                    src=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_next_group,
+                    use_calc_stream=sync_recv,
+                )
+                if sync_recv:
+                    allgather_partial(
+                        tensor_recv_next,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+
+        if tensor_send_prev is not None:
+            if isinstance(tensor_send_prev, tuple):
+                for d in tensor_send_prev:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_prev_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_prev, use_calc_stream=True)
+                send_partial(
+                    tensor_send_prev,
+                    dst=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_prev_group,
+                    use_calc_stream=False,
+                )
+    else:
+        _xpu_comm_group_start()
+        if tensor_send_prev is not None:
+            if isinstance(tensor_send_prev, tuple):
+                for d in tensor_send_prev:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_prev_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_prev, use_calc_stream=True)
+                send_partial(
+                    tensor_send_prev,
+                    dst=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_prev_group,
+                    use_calc_stream=False,
+                )
+
+        if tensor_recv_prev is not None:
+            if isinstance(tensor_recv_prev, tuple):
+                for d in tensor_recv_prev:
+                    task = recv_partial(
+                        d,
+                        src=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_prev_group,
+                        use_calc_stream=sync_recv,
+                    )
+                    if sync_recv:
+                        _xpu_comm_group_end()
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+            else:
+                task = recv_partial(
+                    tensor_recv_prev,
+                    src=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_prev_group,
+                    use_calc_stream=sync_recv,
+                )
+
+                if sync_recv:
+                    _xpu_comm_group_end()
+                    allgather_partial(
+                        tensor_recv_prev,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+
+        if tensor_send_next is not None:
+            if isinstance(tensor_send_next, tuple):
+                for d in tensor_send_next:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_next_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_next, use_calc_stream=True)
+                send_partial(
+                    tensor_send_next,
+                    dst=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_next_group,
+                    use_calc_stream=False,
+                )
+
+        if tensor_recv_next is not None:
+            if isinstance(tensor_recv_next, tuple):
+                for d in tensor_recv_next:
+                    task = recv_partial(
+                        d,
+                        src=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_next_group,
+                        use_calc_stream=sync_recv,
+                    )
+
+                    if sync_recv:
+                        _xpu_comm_group_end()
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+
+            else:
+                task = recv_partial(
+                    tensor_recv_next,
+                    src=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_next_group,
+                    use_calc_stream=sync_recv,
+                )
+                if sync_recv:
+                    _xpu_comm_group_end()
+                    allgather_partial(
+                        tensor_recv_next,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+        _xpu_comm_group_end()
+    if not sync_recv:
+        if framework.in_dynamic_mode():
+            # wait irecv tasks in eager dygraph mode with new comm library
+            for task in tasks:
+                assert task is not None
+                task.wait()
+
+        tensors_for_all_gather = []
+        if tensor_recv_prev is not None:
+            if isinstance(tensor_recv_prev, tuple):
+                for d in tensor_recv_prev:
+                    tensors_for_all_gather.append(d)
+            else:
+                tensors_for_all_gather.append(tensor_recv_prev)
+        if tensor_recv_next is not None:
+            if isinstance(tensor_recv_next, tuple):
+                for d in tensor_recv_next:
+                    tensors_for_all_gather.append(d)
+            else:
+                tensors_for_all_gather.append(tensor_recv_next)
+
+        for tensor in tensors_for_all_gather:
+            allgather_partial(
+                tensor,
+                nranks=mp_degree,
+                rank_id=mp_rank,
+                group=mp_group,
+                use_calc_stream=True,
+            )
+
+    return tensor_recv_prev, tensor_recv_next
+
+
+class P2pHelper:
+    def __init__(self, use_cache=True):
+        self._send_recv_meta = SendRecvMeta()
+        self._use_cache = use_cache
+
+    def _send_meta(self, output_tensor, skip_check_meta=False):
+        if not self._send_recv_meta.has_send_meta:
+            self._send_recv_meta.set_send_message(output_tensor)
+            self._send_recv_meta.send_meta(
+                output_tensor, _hcg.get_pipe_parallel_group()
+            )
+            self._send_recv_meta.has_send_meta = self._use_cache
+
+    def _recv_meta(self):
+        if not self._send_recv_meta.has_recv_meta:
+            self._send_recv_meta.recv_meta(_hcg.get_pipe_parallel_group())
+            self._send_recv_meta.has_recv_meta = self._use_cache
+
+    def recv_forward(self, pp_first_stage, sync_recv=True):
+        global _timers
+        if _timers is not None:
+            _timers("recv_forward").start()
+        if pp_first_stage:
+            input_tensor = None
+        else:
+            self._recv_meta()
+
+            input_tensor, _ = _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=None,
+                recv_prev=True,
+                recv_next=False,
+                sync_recv=sync_recv,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("recv_forward").stop()
+        return input_tensor
+
+    def recv_backward(self, pp_last_stage, sync_recv=True):
+        global _timers
+        if _timers is not None:
+            _timers("recv_backward").start()
+        if pp_last_stage:
+            output_tensor_grad = None
+        else:
+            _, output_tensor_grad = _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=None,
+                recv_prev=False,
+                recv_next=True,
+                sync_recv=sync_recv,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("recv_backward").stop()
+        return output_tensor_grad
+
+    def send_forward(self, output_tensor, pp_last_stage, skip_check_meta=False):
+        global _timers
+        if _timers is not None:
+            _timers("send_forward").start()
+        if not pp_last_stage:
+            self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
+
+            _p2p_helper(
+                tensor_send_next=output_tensor,
+                tensor_send_prev=None,
+                recv_prev=False,
+                recv_next=False,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_forward").stop()
+
+    def send_backward(self, input_tensor_grad, pp_first_stage):
+        global _timers
+        if _timers is not None:
+            _timers("send_backward").start()
+        if not pp_first_stage:
+            _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=input_tensor_grad,
+                recv_prev=False,
+                recv_next=False,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_backward").stop()
+
+    def send_forward_recv_backward(self, output_tensor, pp_last_stage):
+        global _timers
+        if _timers is not None:
+            _timers("send_forward_recv_backward").start()
+        if pp_last_stage:
+            output_tensor_grad = None
+        else:
+            _, output_tensor_grad = _p2p_helper(
+                tensor_send_next=output_tensor,
+                tensor_send_prev=None,
+                recv_prev=False,
+                recv_next=True,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_forward_recv_backward").stop()
+        return output_tensor_grad
+
+    def send_backward_recv_forward(self, input_tensor_grad, pp_first_stage):
+        global _timers
+        if _timers is not None:
+            _timers("send_backward_recv_forward").start()
+        if pp_first_stage:
+            input_tensor = None
+        else:
+            input_tensor, _ = _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=input_tensor_grad,
+                recv_prev=True,
+                recv_next=False,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_backward_recv_forward").stop()
+        return input_tensor
+
+    def send_forward_backward_recv_forward_backward(
+        self, output_tensor, input_tensor_grad, recv_prev, recv_next
+    ):
+        # always have to send dtype info to downstream
+        global _timers
+        if _timers is not None:
+            _timers("send_forward_backward_recv_forward_backward").start()
+
+        self._send_meta(output_tensor)
+        if recv_prev:
+            self._recv_meta()
+        input_tensor, output_tensor_grad = _p2p_helper(
+            tensor_send_next=output_tensor,
+            tensor_send_prev=input_tensor_grad,
+            recv_prev=recv_prev,
+            recv_next=recv_next,
+            sync_recv=False,
+            send_recv_meta=self._send_recv_meta,
+        )
+        if _timers is not None:
+            _timers("send_forward_backward_recv_forward_backward").stop()
+        return input_tensor, output_tensor_grad
+
+    def send_forward_recv_forward(self, output_tensor, recv_prev):
+        # always have to send dtype info to downstream
+        global _timers
+        if _timers is not None:
+            _timers("send_forward_recv_forward").start()
+
+        if output_tensor is not None:
+            self._send_meta(output_tensor)
+
+        if recv_prev:
+            self._recv_meta()
+
+        input_tensor, _ = _p2p_helper(
+            tensor_send_next=output_tensor,
+            tensor_send_prev=None,
+            recv_prev=recv_prev,
+            recv_next=False,
+            sync_recv=False,
+            send_recv_meta=self._send_recv_meta,
+        )
+        if _timers is not None:
+            _timers("send_forward_recv_forward").stop()
+        return input_tensor
+
+    def send_backward_recv_backward(self, input_tensor_grad, recv_next):
+        global _timers
+        if _timers is not None:
+            _timers("send_backward_recv_backward").start()
+        _, output_tensor_grad = _p2p_helper(
+            tensor_send_next=None,
+            tensor_send_prev=input_tensor_grad,
+            recv_prev=False,
+            recv_next=recv_next,
+            sync_recv=False,
+            send_recv_meta=self._send_recv_meta,
+        )
+        if _timers is not None:
+            _timers("send_backward_recv_backward").stop()
+        return output_tensor_grad

python/paddle/distributed/fleet/meta_parallel/pp_utils/profiler_helper.py

@@ -0,0 +1,41 @@
+# 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 os
+
+
+def main():
+    all_record = []
+    all_files = os.listdir('./')
+    all_files = sorted(
+        filter(
+            lambda file: file.startswith("profile_record_tmp_file_for_rank_"),
+            all_files,
+        )
+    )
+
+    for files in all_files:
+        with open(files, 'r') as f:
+            for line in f:
+                all_record.append(line.strip())
+
+    with open('pipeline_profile.json', 'w') as f:
+        f.write('[ ')
+        f.writelines(all_record[i] + ',\n' for i in range(len(all_record) - 1))
+        f.write(all_record[-1])
+        f.write(' ]\n')
+
+
+if __name__ == "__main__":
+    main()

python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py

@@ -0,0 +1,160 @@
+# 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 paddle
+from paddle import _C_ops
+
+__all__ = []
+
+
+FLOAT_TYPE_DICT = {
+    paddle.float16: "float16",
+    paddle.float32: "float32",
+    paddle.float64: "float64",
+    paddle.bfloat16: "bfloat16",
+    paddle.bool: "bool",
+}
+
+PADDLE_TO_NUMBER = {
+    paddle.float16: 0,
+    paddle.float32: 1,
+    paddle.float64: 2,
+    paddle.int32: 3,
+    paddle.int64: 4,
+    paddle.bfloat16: 5,
+    paddle.bool: 6,
+}
+
+NUMBER_TO_DTYPE = {
+    0: "float16",
+    1: "float32",
+    2: "float64",
+    3: "int32",
+    4: "int64",
+    5: "bfloat16",
+    6: "bool",
+}
+
+
+def is_float_tensor(tensor):
+    """Is a float tensor"""
+    return tensor.dtype in FLOAT_TYPE_DICT.keys()
+
+
+def get_tensor_dtype(dtype):
+    assert dtype in FLOAT_TYPE_DICT.keys()
+    return FLOAT_TYPE_DICT[dtype]
+
+
+def paddle_2_number(dtype):
+    assert dtype in PADDLE_TO_NUMBER.keys()
+    return PADDLE_TO_NUMBER[dtype]
+
+
+def number_2_dtype(number):
+    assert number in NUMBER_TO_DTYPE.keys()
+    return NUMBER_TO_DTYPE[number]
+
+
+def get_tensor_bytes(tensor):
+    """Get the bytes a tensor occupied."""
+    elem_size = None
+    if tensor.dtype == paddle.float32:
+        elem_size = 4
+    elif tensor.dtype == paddle.float64:
+        elem_size = 8
+    elif tensor.dtype == paddle.int64:
+        elem_size = 8
+    elif tensor.dtype == paddle.int32:
+        elem_size = 4
+    elif tensor.dtype == paddle.float16:
+        elem_size = 2
+    elif tensor.dtype == paddle.int8:
+        elem_size = 1
+    else:
+        raise ValueError(f"unknown data type: {tensor.dtype}")
+    return tensor.numel() * elem_size
+
+
+def _all_gather(tensor, group=None, use_calc_stream=True):
+    """
+    The main difference with paddle.distributed.all_gather:
+    no need to pass in tensor_list, the returned tensor is spliced
+    """
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+    nranks = (
+        paddle.distributed.collective._get_global_group().nranks
+        if group is None
+        else group.nranks
+    )
+    return _C_ops.all_gather(
+        tensor,
+        ring_id,
+        nranks,
+    )
+
+
+def tuple_to_dict_helper(input_tensor):
+    # recv tuple -> fwd input dict
+    use_dict = False
+    if isinstance(input_tensor, tuple):
+        use_dict = hasattr(input_tensor[0], "key")
+    else:  # single tensor
+        use_dict = hasattr(input_tensor, "key")
+    if use_dict:
+        input_tensor = convert_tensor_tuple_to_dict(input_tensor)
+    return input_tensor, use_dict
+
+
+def dict_to_tuple_helper(output_tensor):
+    if isinstance(output_tensor, dict):
+        output_tensor_tuple = convert_tensor_dict_to_tuple(
+            output_tensor_dict=output_tensor
+        )
+    else:  # single tensor or tensor tuple
+        output_tensor_tuple = output_tensor
+    return output_tensor_tuple
+
+
+def convert_tensor_dict_to_tuple(output_tensor_dict):
+    output_tensor = []
+    for key, tensor in output_tensor_dict.items():
+        if isinstance(tensor, (list, tuple)):
+            for idx, t in enumerate(tensor):
+                t.key = key + " " + str(idx)
+                output_tensor.append(t)
+        else:  # single tensor
+            tensor.key = key
+            output_tensor.append(tensor)
+
+    return tuple(output_tensor)
+
+
+def convert_tensor_tuple_to_dict(input_tensor_tuple):
+    input_tensor_dict = {}
+    for tensor in input_tensor_tuple:
+        key = tensor.key
+        if " " in key:
+            real_key, _ = key.split(" ")
+            if real_key in input_tensor_dict.keys():
+                input_tensor_dict[real_key].append(tensor)
+            else:
+                input_tensor_dict[real_key] = [tensor]
+        else:
+            input_tensor_dict[key] = tensor
+        delattr(tensor, "key")
+    return input_tensor_dict

python/paddle/distributed/fleet/meta_parallel/segment_parallel.py

@@ -0,0 +1,40 @@
+#   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 ..utils.hybrid_parallel_util import (
+    broadcast_dp_parameters,
+    broadcast_sep_parameters,
+    broadcast_sharding_parameters,
+)
+from ..utils.log_util import logger
+from .meta_parallel_base import MetaParallelBase
+
+__all__ = []
+
+
+class SegmentParallel(MetaParallelBase):
+    def __init__(self, layers, hcg, **kwargs):
+        super().__init__(layers, hcg, **kwargs)
+
+    def _prepare_for_model(self):
+        logger.info("start broadcast sep parameters")
+        broadcast_sep_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_sharding_parallel_world_size() > 1:
+            logger.info("start broadcast sharding parameters")
+            broadcast_sharding_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_data_parallel_world_size() > 1:
+            logger.info("start broadcast dp parameters")
+            broadcast_dp_parameters(self._layers, self._hcg)

python/paddle/distributed/fleet/meta_parallel/sharding/__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/fleet/meta_parallel/sharding/group_sharded_fully_shard.py

@@ -0,0 +1,730 @@
+# 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 logging
+from collections import OrderedDict
+from types import MethodType
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import collective, fleet
+from paddle.framework import core
+from paddle.nn import ClipGradByGlobalNorm
+
+from .group_sharded_stage3 import (
+    ForwardPostHooks,
+    ForwardPreHooks,
+    OrderedSet,
+    TaskFlow,
+    _current_layer_params,
+    _PartitionParam,
+    _TensorWrapper,
+    _UnsliceParam,
+    align,
+    alignment,
+)
+from .group_sharded_storage import GradStorage
+from .group_sharded_utils import GroupShardedClipGrad, Type, device_guard
+
+
+def _OptimizerWrapper(optimizer, offload, group, update_params_slice):
+    if not hasattr(optimizer, "_optim"):
+        optimizer._optim = optimizer
+        optimizer.offload = offload
+        optimizer._group = group
+        optimizer.update_scaler = None
+        optimizer.update_slice = update_params_slice
+    return optimizer
+
+
+class FullyShardOptimizer:
+    def __init__(
+        self,
+        optimizer,
+        group=None,
+        sync_buffers=False,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        segment_size=2**20,
+        pretrain_sync_models=True,
+        offload=False,
+        sync_comm=False,
+        dp_group=None,
+        exclude_layer=None,
+    ):
+        self._default_device = device
+        self.__sync_buffers = sync_buffers
+        self._offload = offload
+        self._sync_comm = sync_comm
+
+        # stage3 support some layer set by users to be unslice
+        # _exclude_layer=[layer_name or id(layer)]
+        self._exclude_layer = [] if exclude_layer is None else exclude_layer
+        assert isinstance(self._exclude_layer, (list, tuple)), (
+            "the exclude_layers must be a list with layers' name or layers' id"
+        )
+
+        # segmentation size
+        assert segment_size >= 0, "segment_size must be GE than 0."
+        self._segment_size = segment_size
+
+        global param2dtype
+        param2dtype = {}
+
+        hcg = fleet.get_hybrid_communicate_group()
+        group = hcg.get_sharding_parallel_group()
+        # Communication group establishment
+        self._group = (
+            collective.new_group(collective._get_global_group().ranks)
+            if group is None
+            else group
+        )
+        self._dp_group = dp_group
+        self._world_size_scaling = 1.0 / self._group.nranks
+        assert self._group.nranks > 1, (
+            "Training must be distributed, ranks must be greater than 1."
+        )
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[
+            0
+        ]  # picking ranks index 0 as the reference
+
+        # Parameter segmentation for global ranks
+        self._unslice_params = OrderedSet()  # param's numel <= segment_size
+        self._unslice_params2align = {}  # {param.name: param's align}
+        self._grad_storages = {}  # {param.dtype: GradStorage}
+
+        assert not isinstance(optimizer, list), (
+            "Multiple optimizers are not supported now."
+        )
+        self._optim = _OptimizerWrapper(
+            optimizer,
+            self._offload,
+            self._group,
+            self._update_params_slice,
+        )
+        self._ori_parameter_list = self._optim._parameter_list
+        self._ori_param_groups = self._optim._param_groups
+
+        for p in self._ori_parameter_list:
+            del p._need_shard
+            if p._numel() > self._segment_size:
+                pass
+            elif p.trainable:
+                self._unslice_params.add(_UnsliceParam(p))
+
+        # check main_grad
+        self._check_main_grad()
+
+        # Replace optimizer's _grad_clip
+        if isinstance(self._optim._grad_clip, ClipGradByGlobalNorm):
+            logging.warning(
+                "While using ClipGradByGlobalNorm in GroupShardedStage3, the grad clip of original optimizer will be changed."
+            )
+            if self.use_main_grad:
+                self._optim._inner_opt._grad_clip = GroupShardedClipGrad(
+                    self._optim._inner_opt._grad_clip,
+                    paddle.get_device(),
+                    self._group,
+                )
+            else:
+                self._optim._grad_clip = GroupShardedClipGrad(
+                    self._optim._grad_clip, paddle.get_device(), self._group
+                )
+            if self._optim._parameter_list and isinstance(
+                self._optim._parameter_list[0], dict
+            ):
+                for item in self._optim._param_groups:
+                    if "grad_clip" in item.keys():
+                        item["grad_clip"] = self._optim._grad_clip
+
+        # Add unslice params to master_weight in fp16
+        self._setup_master_weights_for_unslice()
+
+        # Redefine optimizer step and clear function
+        self._redefine_opt_step()
+        self._redefine_opt_clear()
+
+    def _check_main_grad(self):
+        self.use_main_grad = None
+        for param in self._ori_parameter_list:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+
+    def _setup_master_weights_for_unslice(self):
+        for param in self._unslice_params:
+            # Update optimizer master weights
+            if (
+                param.dtype == Type.fp16.value or param.dtype == Type.bf16.value
+            ) and not self._offload:
+                master_tensor = paddle.cast(param, Type.fp32.value)
+                master_tensor.name = param.name
+                self._optim._master_weights[param.name] = master_tensor
+
+    def _clear_gradients(self):
+        current_layer_params = self._ori_parameter_list
+        # 1.Handle param's slice
+        trainable_params = list(
+            filter(
+                lambda p: p.trainable and p not in self._unslice_params,
+                current_layer_params,
+            )
+        )
+        for param in trainable_params:
+            if not hasattr(param, "fw_storage"):
+                continue
+            assert hasattr(param, "fw_storage"), (
+                f"Find {param.name} don't have fw_storage attribute."
+            )
+            if self.use_main_grad:
+                param.fw_storage.main_grad._clear()
+                param.fw_storage.main_grad = None
+            else:
+                param.fw_storage.clear_gradient(False)
+            param.bw_storage._clear()
+            param.bw_storage = None
+
+    # Update param memory slice
+    def _update_params_slice(self):
+        update_list = self._update_params()
+
+        if not isinstance(self._optim._param_groups[0], dict):
+            slice_params = [param.fw_storage for param in update_list]
+            self._optim._parameter_list = slice_params + list(
+                self._unslice_params
+            )
+            self._optim._param_groups = slice_params + list(
+                self._unslice_params
+            )
+        else:
+            for param_group in self._optim._param_groups:
+                p_group = []
+                for p in param_group['params']:
+                    if hasattr(p, "fw_storage"):
+                        p_group.append(p.fw_storage)
+                    else:
+                        p_group.append(p)
+
+                param_group['params'] = p_group
+
+    def _update_params(self):
+        """
+        Update parameters to optimizer memory slice.
+        """
+        update_list = []
+        current_layer_params = self._ori_parameter_list
+        trainable_params = list(
+            filter(
+                lambda p: p.trainable and p not in self._unslice_params,
+                current_layer_params,
+            )
+        )
+        # 1.Handle param's slice
+        for param in trainable_params:
+            assert hasattr(param, "fw_storage"), (
+                f"Find {param.name} don't have fw_storage attribute"
+            )
+
+            param.fw_storage = _TensorWrapper(param)
+            if self.use_main_grad:
+                param.fw_storage.main_grad = param.bw_storage
+            else:
+                assert param.fw_storage.grad is None
+                param.fw_storage._copy_gradient_from(param.bw_storage)
+            update_list.append(param)
+
+        return update_list
+
+    def _redefine_opt_step(self):
+        params_slice_func = self._update_params_slice
+        opt_step = self._optim.step
+
+        def _opt_step(self):
+            if not self.update_scaler:
+                params_slice_func()
+            opt_step()
+
+        self._optim.step = MethodType(_opt_step, self._optim)
+
+    def _redefine_opt_clear(self):
+        clear_func = self._clear_gradients
+
+        def _opt_clear(self):
+            clear_func()
+
+        self._optim.clear_grad = MethodType(_opt_clear, self._optim)
+
+
+class FullyShard(nn.Layer):
+    """
+    A wrapper for Sharding Stage3 Layer in Dygraph.
+
+    .. warning: GroupShardedStage3 encapsulates the layer strategy and integrates it into the nn.Layer.
+
+    .. ZeRO: https://arxiv.org/pdf/1910.02054.pdf.
+    """
+
+    def __init__(
+        self,
+        layer,
+        group=None,
+        sync_buffers=False,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        segment_size=2**20,
+        pretrain_sync_models=True,
+        offload=False,
+        sync_comm=False,
+        dp_group=None,
+        exclude_layer=None,
+    ):
+        super().__init__()
+
+        # Default configs
+        assert (
+            core.is_compiled_with_cuda()
+            or core.is_compiled_with_xpu()
+            or (device in core.get_all_custom_device_type())
+        ), "Only support CUDA / XPU / CustomDevice."
+
+        self._layer = layer
+        self._default_device = device
+        self.__sync_buffers = sync_buffers
+        self._offload = offload
+        self._sync_comm = sync_comm
+
+        # stage3 support some layer set by users to be unslice
+        self._exclude_layer = [] if exclude_layer is None else exclude_layer
+        assert isinstance(self._exclude_layer, (list, tuple)), (
+            "the exclude_layers must be a list with layers' name or layers' id"
+        )
+
+        # segmentation size
+        assert segment_size >= 0, "segment_size must be GE than 0."
+        self._segment_size = segment_size
+
+        global DEV
+        DEV = (
+            "cpu"
+            if paddle.get_device() == "cpu"
+            else paddle.get_device().split(":")[0]
+        )
+        global DEV_ID
+        DEV_ID = (
+            0
+            if paddle.get_device() == "cpu"
+            else int(paddle.get_device().split(":")[1])
+        )
+        global param2dtype
+        param2dtype = {}
+
+        hcg = fleet.get_hybrid_communicate_group()
+        group = hcg.get_sharding_parallel_group()
+
+        self._group = (
+            collective.new_group(collective._get_global_group().ranks)
+            if group is None
+            else group
+        )
+        self._dp_group = dp_group
+        self._world_size_scaling = 1.0 / self._group.nranks
+        assert self._group.nranks > 1, (
+            "Training must be distributed, ranks must be greater than 1."
+        )
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[
+            0
+        ]  # picking ranks index 0 as the reference
+
+        # Parameter segmentation for global ranks
+        # After flatten -> self._param2buffer_size, self._param2buffer, self._trainable_params
+        self._param2buffer_size = {}  # {param.name: size}
+        self._param2buffer = {}  # {param.name: [(start0, end0),(start1, end1), ...]}
+        self._trainable_params = {}  # {id(layer): [trainable_params]}
+        self._unslice_params = OrderedSet()  # param's numel <= segment_size
+        self._unslice_params2align = {}  # {param.name: param's align}
+        self._grad_storages = {}  # {param.dtype: GradStorage}
+
+        self._ori_parameter_list = self._layer.parameters()
+        for param in self._ori_parameter_list:
+            param._need_shard = True
+        # check main_grad
+        self._check_main_grad()
+
+        # Synchronous all ranks models
+        if pretrain_sync_models:
+            self._sync_params_and_buffers()
+
+        self._segment_rank_params(self._layer)
+
+        # In the first step, record the execution order of the layer
+        self._order_tracer = OrderedDict()
+        self._order_tracer["order"] = 0
+        self._order_tracer["layer"] = []
+
+        # Add unslice params GradStorage
+        self._handle_unslice_params()
+
+        # Register task flow
+        self._task_flow = TaskFlow()
+
+        # Register forward hooks
+        self._register_forward_hooks(self._layer)
+
+        # Register backward parameter hooks
+        self._register_backward_hooks()
+
+    def _handle_unslice_params(self):
+        buffer_size = {}
+        buffer_size[Type.bf16.value] = 0
+        buffer_size[Type.fp32.value] = 0
+        buffer_size[Type.fp16.value] = 0
+        for param in self._unslice_params:
+            param2dtype[param.name] = param.dtype
+            p_align = self._param2align(param)
+            self._unslice_params2align[param.name] = p_align
+            buffer_size[param.dtype] += param._numel() + p_align
+
+        # Create unslice_params'grad
+        for param in sorted(self._unslice_params, key=lambda p: p.name):
+            if param.dtype not in self._grad_storages.keys():
+                self._grad_storages[param.dtype] = GradStorage(
+                    buffer_size[param.dtype],
+                    dtype=(
+                        param.dtype
+                        if not self.use_main_grad
+                        else paddle.float32
+                    ),
+                    device=self._default_device,
+                    destination=self._rank,
+                    param2align=self._unslice_params2align,
+                )
+            self._grad_storages[param.dtype].add_grad(
+                param, self._unslice_params2align[param.name]
+            )
+
+    def _check_main_grad(self):
+        self.use_main_grad = None
+        for param in self._layer.parameters():
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+
+    @paddle.autograd.no_grad()
+    def _sync_params_and_buffers(self):
+        """
+        Sync all model states for all ranks
+        """
+
+        for p in self._layer.parameters():
+            dist.broadcast(
+                p, src=self._global_root_rank, group=self._group, sync_op=True
+            )
+            if self._dp_group is not None and self._dp_group.nranks > 1:
+                dist.broadcast(
+                    p,
+                    src=self._dp_group.ranks[0],
+                    group=self._dp_group,
+                    sync_op=True,
+                )
+
+    def _sync_grad_storages_hook(self):
+        for grad_storage in self._grad_storages.values():
+            grad_storage.buffer.scale_(scale=self._world_size_scaling)
+            dist.all_reduce(tensor=grad_storage.buffer, group=self._group)
+            if self._dp_group is not None and self._dp_group.nranks > 1:
+                grad_storage.buffer.scale_(scale=(1.0 / self._dp_group.nranks))
+                dist.all_reduce(
+                    tensor=grad_storage.buffer, group=self._dp_group
+                )
+
+    def forward(self, *inputs, **kwargs):
+        """
+        A wrapper for Sharding Stage3 layer.
+        """
+        # add hook to sync grad storage
+        for grad_storage in self._grad_storages.values():
+            grad_storage.buffer.zero_()
+            grad_storage.manual_release()
+            grad_storage.rebuild()
+        core.eager._add_backward_final_hook(self._sync_grad_storages_hook)
+
+        # 1.Sync layer's buffers state
+        if self.__sync_buffers:
+            self._sync_buffers()
+
+        # 2.Normal FW on the base model
+        fw = self._layer(*inputs, **kwargs)
+
+        return fw
+
+    def _segment_rank_params(self, layer, name="last_layer"):
+        """
+        Flatten parameters according to layer.
+        """
+        current_layer_params = _current_layer_params(layer)
+        if current_layer_params:
+            self._flatten_layer_params(layer, current_layer_params)
+
+        for name, sub_layer in layer.named_children():
+            self._segment_rank_params(sub_layer, name)
+
+    def _flatten_layer_params(self, layer, current_layer_params):
+        """
+        Parameter segmentation and memory integration.
+        """
+
+        if id(layer) in self._trainable_params.keys():
+            return
+
+        def _add_manage_info(trainable_param):
+            return _PartitionParam(trainable_param)
+
+        current_params = []
+        for p in current_layer_params:
+            if p._numel() > self._segment_size:
+                current_params.append(_add_manage_info(p))
+            elif p.trainable:
+                self._unslice_params.add(_UnsliceParam(p))
+
+        self._trainable_params[id(layer)] = current_params
+
+        for param in self._trainable_params[id(layer)]:
+            if param.name in self._param2buffer.keys():
+                continue
+            self._param2buffer[param.name] = []
+            # 1.Params alignment
+            align_ = self._param2align(param)
+
+            offset = align_ + param._numel()
+            buffer_size = (
+                offset
+                if offset % self._group.nranks == 0
+                else offset + self._group.nranks - (offset % self._group.nranks)
+            )
+            self._param2buffer_size[param.name] = buffer_size
+
+            # 2.Combination param buffer
+            assert buffer_size % self._group.nranks == 0
+            pre_buffer = buffer_size // self._group.nranks
+
+            for rank_ in range(self._group.nranks):
+                self._param2buffer[param.name].append(
+                    (rank_ * pre_buffer, (rank_ + 1) * pre_buffer)
+                )
+
+            # Record param's dtype
+            param2dtype[param.name] = param.dtype
+            # 3.Flatten layer params and release other rank buffer
+            self._param_storage(param, buffer_size)
+
+    def _param_storage(self, param, buffer_size):
+        """
+        This is a function to simplify the handling of parameter InternalStorages.
+        """
+        assert isinstance(buffer_size, int)
+        value = (
+            np.zeros(buffer_size, dtype=np.float16)
+            if (
+                Type.fp16.value == param.dtype or Type.bf16.value == param.dtype
+            )
+            else np.zeros(buffer_size, dtype=np.float32)
+        )
+        buffer = core.eager.Tensor(value=value, place=core.CPUPlace())
+        if Type.bf16.value == param.dtype:
+            buffer = buffer.cast(Type.bf16.value)
+
+        param_shape = param.shape
+        origin_state = param.stop_gradient
+        param.stop_gradient = True
+        param.flatten_()
+        param.stop_gradient = origin_state
+        start, end = self._param2buffer[param.name][self._rank]
+
+        # Copy the current param value
+        with device_guard():
+            tmp_var = buffer._slice(0, param._numel())
+        param_cpu = param.cpu()
+        tmp_var.get_tensor().set(param_cpu.get_tensor(), core.CPUPlace())
+        del tmp_var
+        param.get_tensor()._set_dims(param_shape)
+
+        # Current rank param_storage
+        param.fw_storage = core.eager.Tensor(
+            value=buffer._slice(start, end), name="slice@" + param.name
+        )
+        param.status = "part"
+        param._clear_data()
+
+    def _register_forward_hooks(self, layer):
+        """
+        Register PyLayer to manage memory slices.
+        There are four stages:
+        FW
+        1. Before the forward layers, synchronize the full parameters.
+        2. After the forward layers, release the full parameter and keep the parameter slice.
+        BW
+        3. Before the backward layers, synchronize the full parameters and create param's grad.
+        4. After the gradient accumulation, release the full parameter and keep the parameter slice.
+        """
+        current_layer_params = _current_layer_params(layer)
+        if current_layer_params:
+            # the layer in self._exclude_layer will be added hooks.
+            if not (
+                id(layer) in self._exclude_layer
+                or layer.__class__.__name__ in self._exclude_layer
+            ):
+                self._register_forward_all_hooks(layer, self._task_flow)
+
+        for _, sub_layer in layer.named_children():
+            self._register_forward_hooks(sub_layer)
+
+    def _register_forward_all_hooks(self, sub_layer, task_flow):
+        def _forward_pre_hook(layer, inputs):
+            return ForwardPreHooks(
+                layer,
+                self._order_tracer,
+                self._trainable_params,
+                self._param2buffer_size,
+                self._group,
+                self._sync_comm,
+                self._offload,
+                task_flow,
+            )
+
+        def _forward_post_hook(layer, inputs, outputs):
+            if isinstance(outputs, paddle.Tensor):
+                outputs = (outputs,)
+            return ForwardPostHooks.apply(
+                *outputs,
+                layer=layer,
+                order_tracer=self._order_tracer,
+                trainable_params=self._trainable_params,
+                param2buffer=self._param2buffer,
+                param2buffer_size=self._param2buffer_size,
+                rank=self._rank,
+                group=self._group,
+                sync_comm=self._sync_comm,
+                offload=self._offload,
+                task_flow=task_flow,
+            )
+
+        # register previous forward hooks
+        sub_layer.register_forward_pre_hook(_forward_pre_hook)
+
+        # register post forward hooks
+        sub_layer.register_forward_post_hook(_forward_post_hook)
+
+    @paddle.autograd.no_grad()
+    def _sync_buffers(self):
+        """
+        Sync all the param buffers from all ranks (exp: batch norm statistics).
+        """
+
+        for buffer in self._layer.buffers(include_sublayers=True):
+            dist.broadcast(
+                buffer, self._global_root_rank, self._group, sync_op=True
+            )
+            if self._dp_group is not None and self._dp_group.nranks > 1:
+                dist.broadcast(
+                    buffer,
+                    self._dp_group.ranks[0],
+                    self._dp_group,
+                    sync_op=True,
+                )
+
+    def __getattr__(self, name):
+        """Forward missing attributes to wrapped layer."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._layer, name)
+
+    def _register_backward_hooks(self):
+        current_layer_params = self._layer.parameters(include_sublayers=True)
+        trainable_params = list(
+            filter(
+                lambda p: p.trainable and p not in self._unslice_params,
+                current_layer_params,
+            )
+        )
+
+        for param in trainable_params:
+            allreduce_function = self._get_allreduce_fn(param)
+            param._register_backward_hook(allreduce_function)
+
+    def _get_allreduce_fn(self, param):
+        @paddle.autograd.no_grad()
+        def allreduce_(*_):
+            assert param.trainable, (
+                "the param must be trainable for grad allreduced"
+            )
+            if param.name in self._task_flow.full_grad.keys():
+                full_grad = self._task_flow.full_grad[param.name]
+                # Only support sync allreduce current rank's layer now
+                full_grad.scale_(scale=self._world_size_scaling)
+                dist.all_reduce(tensor=full_grad, group=self._group)
+                if self._dp_group is not None and self._dp_group.nranks > 1:
+                    full_grad.scale_(scale=1.0 / self._dp_group.nranks)
+                    dist.all_reduce(tensor=full_grad, group=self._dp_group)
+
+                start, end = self._param2buffer[param.name][self._rank]
+                if param.bw_storage is None:
+                    param.bw_storage = (
+                        full_grad._slice(start, end).detach().clone()
+                    )
+                else:
+                    param.bw_storage = paddle.add(
+                        param.bw_storage,
+                        full_grad._slice(start, end).detach().clone(),
+                    )
+
+                if self.use_main_grad:
+                    param.main_grad = None
+                else:
+                    param.clear_gradient(False)
+                del self._task_flow.full_grad[param.name]
+
+            if param.name in self._task_flow.full_param.keys():
+                if param.status == "all":
+                    param.use_count = 0
+                    param._clear_data()
+                    start, end = self._param2buffer[param.name][self._rank]
+                    param.fw_storage = (
+                        self._task_flow.full_param[param.name][0]
+                        ._slice(start, end)
+                        .detach()
+                        .clone()
+                    )
+                    param.status = "part"
+                    del self._task_flow.full_param[param.name]
+
+        return allreduce_
+
+    def _param2align(self, param):
+        # CUDA alignment 256 bytes
+        size = param._numel() * align[param.dtype]
+        device_alignment = alignment[self._default_device]
+        remaining = size % device_alignment
+        ali = 0 if remaining == 0 else device_alignment - remaining
+        align_ = ali // align[param.dtype]
+        return align_

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_optimizer_stage2.py

@@ -0,0 +1,792 @@
+#   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.
+
+# The file has been adapted from fairscale file:
+# https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/oss.py
+# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
+# We retain the following license from the original files:
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import warnings
+from collections import OrderedDict
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import ParallelMode, fleet
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    ShardedStateDict,
+    ShardedWeight,
+    create_sharded_weight_with_new_local,
+)
+from paddle.framework import core
+from paddle.nn import ClipGradByGlobalNorm
+from paddle.optimizer import Optimizer
+
+HybridParallelClipGrad = fleet.meta_optimizers.dygraph_optimizer.hybrid_parallel_optimizer.HybridParallelClipGrad
+from paddle.distributed.collective import _get_global_group, new_group
+
+from .group_sharded_storage import GradStorage, ParamStorage
+from .group_sharded_utils import GroupShardedClipGrad, Type, device_guard
+
+# CUDA alignment 256 bytes, cpu alignment 4096 bytes
+alignment = {"gpu": 256, "cpu": 4096, "xpu": 256}
+align = {
+    Type.fp16.value: 2,
+    Type.bf16.value: 2,
+    Type.fp32.value: 4,
+}
+
+
+class GroupShardedOptimizerStage2(Optimizer):
+    """
+    A wrapper for Sharding Stage2 Optimizer in Dygraph.
+
+    .. warning: ShardingOptimizer encapsulates the optimization strategy and integrates it into the optimizer.
+
+    .. ZeRO: 1.https://arxiv.org/pdf/1910.02054.pdf 2.https://arxiv.org/pdf/1910.02054.pdf.
+
+    """
+
+    # TODO (Baibaifan)
+    # Feature Notes:
+    # 1. Unified memory for parameters and parameters.grad to InternalStorage.
+    # 2. Support the segmentation of optimizer parameters and partial updating of parameters.
+    # 3. Dynamically adjust training parameters and models.
+    # 4. Support offload function.
+    # 5. Support the establishment of independent communication groups.
+    # 6. Broadcast_fp16 is not supported now.
+    def __init__(
+        self,
+        params,
+        optim,
+        group=None,
+        offload=False,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        pretrain_sync_models=True,
+        dp_group=None,
+        **kw,
+    ):
+        super().__init__(learning_rate=optim._learning_rate, parameters=params)
+        assert (
+            core.is_compiled_with_cuda()
+            or core.is_compiled_with_xpu()
+            or (device in core.get_all_custom_device_type())
+        ), "Only GPU and XPU and CustomDevice is supported now"
+
+        # Segmentation information
+        self._dtype_rank_params = (
+            OrderedDict()
+        )  # {dtype:[param1,param2]} device, rank, params
+        self._param2rank = {}
+        self.__segment_params = []
+        self._rank_buffer_size = {}  # {dtype: {rank: numel+alignment}}
+        self._param2align = {}  # {param.name: align}
+
+        # Default information
+        self._optim = optim
+
+        # sharing stage 2 comm overlap flag
+        self._reduce_overlap = False
+        # record the last task used for comm overlap for sharding stage 2
+        self._comm_task = None
+
+        assert hasattr(self._optim, "_master_weights"), (
+            "Must use optimizer with _master_weights attribute"
+        )
+
+        # Support parameter group and parameter list
+        self._local_params = []
+        if isinstance(params[0], dict):
+            for param_group in params:
+                self._local_params.extend(list(param_group["params"]))
+        else:
+            self._local_params.extend(list(params))
+
+        self.use_main_grad = None
+        for param in self._local_params:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+        if self.use_main_grad:
+            assert not offload, "offload not support main_grad for now"
+
+        self._default_device = device
+        self._pfp16 = (
+            len(
+                list(
+                    filter(
+                        lambda x: x.trainable and x.dtype == Type.fp16.value,
+                        self._local_params,
+                    )
+                )
+            )
+            > 0
+        )
+        self._pbf16 = (
+            len(
+                list(
+                    filter(
+                        lambda x: x.trainable and x.dtype == Type.bf16.value,
+                        self._local_params,
+                    )
+                )
+            )
+            > 0
+        )
+
+        self._broadcast_overlap = False
+        self._forward_pre_hook_remove_helper = []
+        try:
+            # The fp32 params such as layer_norm_0.w_0 will be at the end of param_list.
+            # Have to sort the params to make sure all params are in the forward using order.
+            self._broadcast_order_params = sorted(
+                self.local_params,
+                key=lambda x: int(x.name.split('.')[0].split('_')[-1]),
+            )
+        except ValueError:
+            self._broadcast_order_params = None
+
+        self._group = (
+            new_group(_get_global_group().ranks) if group is None else group
+        )
+
+        # only support to combine stage2 and dp hybrid parallel now.
+        self._dp_group = dp_group
+        self.world_size = self._group.nranks
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[0]
+
+        if self._dp_group is not None and self._dp_group.nranks > 1:
+            assert not offload, (
+                "Not support! when using offload with sharding stage2, please use pure sharding stage2, exclude data parallel."
+            )
+
+        # Synchronous all ranks models
+        if pretrain_sync_models:
+            self._sync_params_and_buffers()
+
+        self.param_storages = {}  # {dtype: {rank: InternalStorage}}
+
+        if isinstance(self._optim._grad_clip, ClipGradByGlobalNorm):
+            logging.warning(
+                "While using ClipGradByGlobalNorm in GroupShardedOptimizerStage2, the grad clip of original optimizer will be changed."
+            )
+
+            hcg = fleet.fleet._hcg if hasattr(fleet.fleet, "_hcg") else None
+            if (
+                hcg
+                and hcg.get_parallel_mode() is not ParallelMode.DATA_PARALLEL
+                and not offload
+            ):
+                if self.use_main_grad:
+                    self._optim._inner_opt._grad_clip = HybridParallelClipGrad(
+                        self._optim._inner_opt._grad_clip, hcg
+                    )
+                else:
+                    self._optim._grad_clip = HybridParallelClipGrad(
+                        self._optim._grad_clip, hcg
+                    )
+            else:
+                if self.use_main_grad:
+                    self._optim._inner_opt._grad_clip = GroupShardedClipGrad(
+                        self._optim._inner_opt._grad_clip,
+                        paddle.get_device(),
+                        self._group,
+                    )
+                else:
+                    self._optim._grad_clip = GroupShardedClipGrad(
+                        self._optim._grad_clip, paddle.get_device(), self._group
+                    )
+
+            if self._optim._parameter_list and isinstance(
+                self._optim._parameter_list[0], dict
+            ):
+                for item in self._optim._param_groups:
+                    if "grad_clip" in item.keys():
+                        item["grad_clip"] = self._optim._grad_clip
+
+        if offload:
+            assert self._pfp16, (
+                "Only support offload strategy while using 'Adam', 'AdamW' and 'Momentum' optimizer with AMP/Pure FP16"
+            )
+
+        self.offload = offload  # Using for offload
+        self.offload_device = "cpu"
+        self.offload_buffer_size = 0
+        self.offload_param2align = {}
+        self.offload_params = None
+        self.offload_grads = None
+        self.dev_id = int(paddle.get_device().split(":")[1])
+
+        self._master_params = {}
+
+        # Update optimizer parameters and adjust parameter storage and use according to rank.
+        self._update_opt_status()
+
+    def _set_auxiliary_var(self, key, val):
+        super()._set_auxiliary_var(key, val)
+        self._optim._set_auxiliary_var(key, val)
+
+    @paddle.autograd.no_grad()
+    def _sync_params_and_buffers(self):
+        """
+        Sync all model states for all ranks
+        """
+
+        for p in self._local_params:
+            dist.broadcast(
+                p, src=self._global_root_rank, group=self._group, sync_op=True
+            )
+
+            if self._dp_group:
+                dist.broadcast(
+                    p,
+                    src=self._dp_group.ranks[0],
+                    group=self._dp_group,
+                    sync_op=True,
+                )
+
+    def _update_task(self, task):
+        if self._reduce_overlap:
+            assert task is not None
+        # Only track of the last reduce task.
+        # Since all tasks are on the same stream, only need to wait the last one.
+        # After waiting for the last reduce task, all reduce tasks before have already finished.
+        self._comm_task = task
+
+    def _set_reduce_overlap(self, reduce_overlap):
+        # Enable gradients' reduces overlap with backward calculation.
+        self._reduce_overlap = reduce_overlap
+
+    def _set_broadcast_overlap(
+        self, broadcast_overlap, layers=None, num_groups=None
+    ):
+        # Enable post optimizer broadcasts overlap with the forward calculation of next batch.
+        self._broadcast_overlap = broadcast_overlap
+        if self._broadcast_overlap:
+            assert layers is not None, (
+                "To enable broadcast overlap forward, please pass the module to the function."
+            )
+            self._layers = layers
+            warnings.warn(
+                "Setting overlap broadcast means the `paddle.device.cuda.synchronize()` "
+                "must be called manually before calling `paddle.save()` and before and inference."
+            )
+            if self._broadcast_order_params is None:
+                # Params' names should be like column_linear_32.w_0 pattern to get the best performance.
+                warnings.warn(
+                    r"The param name passed to the optimizer doesn't follow .+_[0-9]+\..+ pattern, "
+                    "overlap broadcast may harm the performance."
+                )
+                self._broadcast_order_params = self._local_params
+
+        if num_groups is None or num_groups > len(self._broadcast_order_params):
+            warnings.warn(
+                "The num_groups for broadcast is larger than the number of params to be broadcast. "
+                "It will set to default value: 1 (use the default sharding group)."
+            )
+            num_groups = 1
+
+        assert isinstance(num_groups, int) and num_groups > 0, (
+            "num_groups should be a positive integer"
+        )
+
+        self._number_of_broadcast_groups = num_groups
+        self._broadcast_groups = [
+            None for _ in range(self._number_of_broadcast_groups)
+        ]
+        self._broadcast_groups[0] = self._group
+
+        ranks = self._group.ranks
+        for i in range(1, self._number_of_broadcast_groups):
+            self._broadcast_groups[i] = new_group(ranks)
+
+    def _generate_master_params(self, trainable_params):
+        if self.offload:
+            for param in trainable_params:
+                if param.name not in self._master_params.keys():
+                    self._master_params[param.name] = core.eager.Tensor(
+                        name=param.name,
+                        value=param.cast(dtype=Type.fp32.value).numpy(),
+                        place=core.CPUPlace(),
+                        stop_gradient=param.stop_gradient,
+                    )
+        else:
+            for param in trainable_params:
+                if (
+                    param.dtype == Type.fp16.value
+                    or param.dtype == Type.bf16.value
+                ):
+                    master_tensor = paddle.cast(param, Type.fp32.value)
+                    master_tensor.name = param.name
+                    self._optim._master_weights[param.name] = master_tensor
+
+    def _update_opt_status(self):
+        """Update optimizer status and parameter storage information, and special functions to be developed."""
+        # func 1
+        self._integration_params()
+
+    # Segment helpers
+
+    def _segment_params(self):
+        """
+        Divide all optimizer parameters equally into rank.
+        """
+        if len(self.__segment_params) == 0:
+            self.__segment_params, param_lists = (
+                [[] for _ in range(self.world_size)],
+                [[] for _ in range(self.world_size)],
+            )
+            sizes = [0] * self.world_size
+            for param in self._local_params:
+                # Add this param to rank with smallest size.
+                rank = sizes.index(min(sizes))
+                param_lists[rank].append(param)
+
+                # Statistical real numels
+                sizes[rank] += param._numel() if param.trainable else 0
+
+            for rank, params in enumerate(param_lists):
+                self.__segment_params[rank].extend(params)
+        return self.__segment_params
+
+    @property
+    def local_params(self):
+        return self._local_params
+
+    @property
+    def param2rank(self):
+        """Map the params to the rank which owns them"""
+        if len(self._param2rank) == 0:
+            for rank, params in enumerate(self._segment_params()):
+                for param in params:
+                    self._param2rank[param.name] = rank
+        return self._param2rank
+
+    @property
+    def dtype_rank_params(self):
+        """
+        Divide the parameters into groups according to rank and dtype.
+        """
+        if len(self._dtype_rank_params) == 0:
+            # Assign the parameters of each rank according to the type
+            trainable_params = list(
+                filter(lambda x: x.trainable, self._local_params)
+            )
+            for param in trainable_params:
+                if param.dtype not in self._dtype_rank_params.keys():
+                    self._dtype_rank_params[param.dtype] = [
+                        [] for _ in range(self.world_size)
+                    ]
+                self._dtype_rank_params[param.dtype][
+                    self.param2rank[param.name]
+                ].append(param)
+
+            # Sort per rank params by size
+            for dtype in self._dtype_rank_params.keys():
+                for rank_params in self._dtype_rank_params[dtype]:
+                    rank_params.sort(key=lambda x: x._numel())
+
+        return self._dtype_rank_params
+
+    @property
+    def rank_buffer_size(self):
+        """
+        Count the memory size of the parameters corresponding to rank under the corresponding dtype.
+        """
+        # CUDA alignment 256 bytes
+        if self._default_device in core.get_all_custom_device_type():
+            device_alignment = core.libpaddle._get_device_min_chunk_size(
+                self._default_device
+            )
+        else:
+            device_alignment = alignment[self._default_device]
+
+        if len(self._rank_buffer_size) == 0:
+            for dtype in self.dtype_rank_params.keys():
+                if dtype not in self._rank_buffer_size.keys():
+                    self._rank_buffer_size[dtype] = {}
+                for dst_rank, per_rank_params in enumerate(
+                    self.dtype_rank_params[dtype]
+                ):
+                    if dst_rank not in self._rank_buffer_size[dtype].keys():
+                        self._rank_buffer_size[dtype][dst_rank] = 0
+                    for param in per_rank_params:
+                        if not param.trainable:
+                            continue
+                        size = param._numel() * align[dtype]
+                        remaining = size % device_alignment
+                        ali = (
+                            0
+                            if remaining == 0
+                            else device_alignment - remaining
+                        )
+                        align_ = ali // align[dtype]
+                        self._rank_buffer_size[dtype][dst_rank] += (
+                            param._numel() + align_
+                        )
+                        self._param2align[param.name] = align_
+
+        return self._rank_buffer_size
+
+    def _integration_params(self):
+        """
+        Integrate the parameters into a continuous memory according to rank, and support the update of training parameters.
+        """
+
+        for dtype, per_rank_params in self.dtype_rank_params.items():
+            if dtype not in self.param_storages.keys():
+                self.param_storages[dtype] = {}
+
+            for dst_rank, params in enumerate(per_rank_params):
+                if len(params) > 0:
+                    # Merge all the trainable params in a single InternalStorage
+                    trainable_params = list(
+                        filter(lambda x: x.trainable, params)
+                    )
+                    if (self._pfp16 or self._pbf16) and dst_rank == self._rank:
+                        self._generate_master_params(trainable_params)
+                    if trainable_params:
+                        param_storage = ParamStorage(
+                            size=self.rank_buffer_size[dtype][dst_rank],
+                            dtype=dtype,
+                            device=self._default_device,
+                        )
+
+                        param_storage.add_rank_params(
+                            trainable_params, self._param2align
+                        )
+                        self.param_storages[dtype][dst_rank] = param_storage
+
+        # Clear the InternalStorage keys which are not in use anymore
+        dtype_in_use = list(self.dtype_rank_params.keys())
+        dtype_to_pop = list(
+            filter(lambda x: x not in dtype_in_use, self.param_storages.keys())
+        )
+        for d in dtype_to_pop:
+            self.param_storages.pop(d)
+
+        if self.offload:
+            self._optim._master_weights = self._master_params
+            cpu_master_params = list(self._master_params.values())
+            if self._default_device in core.get_all_custom_device_type():
+                device_alignment = core.libpaddle._get_device_min_chunk_size(
+                    self._default_device
+                )
+            else:
+                device_alignment = alignment[self._default_device]
+
+            for param in cpu_master_params:
+                size = param._numel() * align[Type.fp32.value]
+                remaining = size % device_alignment
+                ali = 0 if remaining == 0 else device_alignment - remaining
+                align_ = ali // align[Type.fp32.value]
+                self.offload_buffer_size += param._numel() + align_
+                self.offload_param2align[param.name] = align_
+
+            if cpu_master_params:
+                with device_guard(self._rank, self.offload_device):
+                    self.offload_params = ParamStorage(
+                        size=self.offload_buffer_size,
+                        dtype=Type.fp32.value,
+                        device=self.offload_device,
+                    )
+                    self.offload_params.buffer.name = "offload_buffer"
+                    self.offload_params.add_rank_params(
+                        cpu_master_params, self.offload_param2align, False
+                    )
+                    self.offload_params.buffer.stop_gradient = False
+
+                    self.offload_grads = GradStorage(
+                        size=self.offload_buffer_size,
+                        dtype=Type.fp32.value,
+                        device=self.offload_device,
+                        destination=self._rank,
+                        param2align=self.offload_param2align,
+                        convert_cpu=True,
+                    )
+                    for p in cpu_master_params:
+                        self.offload_grads.add_grad(
+                            p, self.offload_param2align[p.name]
+                        )
+
+                    self._optim._master_weights[
+                        self.offload_params.buffer.name
+                    ] = self.offload_params.buffer
+
+    def _offload_acc_grad(self, param_name, grad_fp32_cpu):
+        """accumulate grads with offload strategy"""
+        with device_guard(self._rank, self.offload_device):
+            if param_name in self._master_params.keys():
+                if self._master_params[param_name].grad is None:
+                    self._master_params[param_name]._copy_gradient_from(
+                        grad_fp32_cpu
+                    )
+                else:
+                    self._master_params[param_name].grad.add_(grad_fp32_cpu)
+
+        self.offload_params.buffer._copy_gradient_from(
+            self.offload_grads.buffer
+        )
+
+    def _offload_scale_grad(self, scale_size):
+        """scale grads with offload strategy"""
+        with device_guard(self._rank, self.offload_device):
+            self.offload_grads.buffer.scale_(scale=scale_size)
+
+    def _offload_clear_grad(self):
+        """clear grads with offload strategy"""
+        with device_guard(self._rank, self.offload_device):
+            self.offload_grads.buffer.zero_()
+
+    def _step(self):
+        if self._broadcast_overlap:
+            # Clear the pre forward hook in the optimizer step.
+            for hook_remove in self._forward_pre_hook_remove_helper:
+                hook_remove.remove()
+            self._forward_pre_hook_remove_helper = []
+
+        if self.offload:
+            params_list = [self.offload_params.buffer]
+
+            # TODO(Baibaifan): Offload will support param_groups later
+            if not isinstance(self._optim._param_groups[0], dict):
+                self._optim._parameter_list = params_list
+                self._optim._param_groups = params_list
+
+        # Run the optimizer of the current rank step
+        if self.offload:
+            with device_guard(device=self.offload_device):
+                self._optim.step()
+
+            for param in self._local_params:
+                if param.name in self._master_params.keys():
+                    if (
+                        self._default_device
+                        in core.get_all_custom_device_type()
+                    ):
+                        param.set_value(
+                            self._master_params[param.name]
+                            ._copy_to(
+                                paddle.CustomPlace(
+                                    self._default_device, self.dev_id
+                                ),
+                                True,
+                            )
+                            .cast(dtype=param.dtype)
+                        )
+                    elif self._default_device == "xpu":
+                        param.set_value(
+                            self._master_params[param.name]
+                            .to("xpu:" + str(self.dev_id))
+                            .cast(dtype=param.dtype)
+                        )
+                    else:
+                        param.set_value(
+                            self._master_params[param.name]
+                            .cuda(self.dev_id)
+                            .cast(dtype=param.dtype)
+                        )
+        else:
+            self._optim.step()
+
+        # Synchronize all the updated shards in between the ranks
+        self._broadcast_params()
+
+    def step(self):
+        """
+        A wrapper for Optimizer's step function to finish the update operation of the optimizer.
+        """
+        # This method won't be called directly by opt.step()!
+        # The _redefine_opt_step() in class GroupShardedStage2 will wrap this function.
+        self._step()
+
+    def minimize(self):
+        raise RuntimeError(
+            "optimizer.minimize() not support now, please use optimizer.step()"
+        )
+
+    def set_state_dict(self, state_dict):
+        self._optim.set_state_dict(state_dict)
+
+    def state_dict(self):
+        return self._optim.state_dict()
+
+    def _clear_cache(self):
+        self.__segment_params.clear()
+        self._dtype_rank_params.clear()
+        self._param2rank.clear()
+
+    @paddle.autograd.no_grad()
+    def _broadcast_params(self):
+        """Broadcast the parameters of the current rank to each rank"""
+
+        # Exchange all the shards with the other ranks
+        if self._broadcast_overlap:
+            self._broadcast_params_overlap_forward()
+        else:
+            for dtype_per_rank in self.param_storages.values():
+                for dst_rank, internal_storage in dtype_per_rank.items():
+                    dist.broadcast(
+                        tensor=internal_storage.buffer,
+                        src=self._group.ranks[dst_rank],
+                        group=self._group,
+                        sync_op=True,
+                    )
+
+    def _forward_pre_hook_function(self, tasks):
+        # Since the layers will call pre hook by `forward_pre_hook(self, inputs)`,
+        # the helper functions needs the x and y to take those params.
+        def __impl__(x, y):
+            for task in tasks:
+                # Wait for broadcast task before using the result of the broadcast.
+                task.wait()
+
+        return __impl__
+
+    def set_lr(self, lr):
+        super().set_lr(lr)
+        self._optim.set_lr(lr)
+
+    def get_lr(self):
+        return self._optim.get_lr()
+
+    @paddle.autograd.no_grad()
+    def _broadcast_params_overlap_forward(self):
+        # Exchange all the shards with the other ranks,
+        # but overlap the broadcast with next batch's calculation.
+        group_idx = 0
+
+        param2task = {}
+        for x in self._broadcast_order_params:
+            if x.trainable:
+                group = self._broadcast_groups[group_idx]
+                group_idx = (group_idx + 1) % self._number_of_broadcast_groups
+                task = dist.broadcast(
+                    tensor=x,
+                    src=group.ranks[self._param2rank[x.name]],
+                    group=group,
+                    sync_op=False,
+                )
+                assert x.name not in param2task
+                param2task[x.name] = task
+
+        for layer in self._layers.sublayers():
+            if len(layer.sublayers()) == 0:
+                # Register forward pre hood for leaf layers. This will get the best performance.
+                tasks = []
+                for param in layer.parameters():
+                    if param.trainable:
+                        if param.name in param2task:
+                            tasks.append(param2task[param.name])
+                self._forward_pre_hook_remove_helper.append(
+                    layer.register_forward_pre_hook(
+                        self._forward_pre_hook_function(tasks)
+                    )
+                )
+
+    def sharded_state_dict(
+        self,
+        model_sharded_state_dict: ShardedStateDict,
+    ) -> ShardedStateDict:
+        """
+        Convert optimizer state dict to a sharded state dict based on model sharding information.
+
+        Args:
+            model_sharded_state_dict (dict): Sharded state dict of the model, containing tensor metadata.
+
+        Returns:
+            dict: A new optimizer state dict where weights are wrapped as ShardedWeight.
+        """
+
+        _FP32_MASTER = "fp32_master_0"
+        _MOMENT_NAME = "moment"
+        _optimizer_scalar_name = [
+            "beta1_pow_acc_0",
+            "beta2_pow_acc_0",
+        ]
+        _optimizer_non_scaler_name = [
+            "moment1_0",
+            "moment2_0",
+            "velocity_0",
+        ]
+
+        def _generate_base_static_name(vname):
+            if _FP32_MASTER in vname:
+                return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
+            for name in _optimizer_scalar_name + _optimizer_non_scaler_name:
+                if vname.endswith(name):
+                    return vname[: -(len(name) + 1)], name
+            raise ValueError(f"Cannot split variable name: {vname}.")
+
+        optimizer_sharded_state_dict = {}
+        optimizer_state_dict = self.state_dict()
+        # Build name mapping and remove non-tensor entries from optimizer state
+        static_to_struct_mapping = {}
+        model_sharded_state_dict = dict(
+            sorted(model_sharded_state_dict.items())
+        )
+        for k, v in model_sharded_state_dict.items():
+            # When shared weights exist, the v.local_tensor.name of shared parameters are identical, but only the first parameter has optimizer states. Therefore, only the key-value pairs of the first occurrence in the shared parameter group need to be retained.
+            if v.local_tensor.name not in static_to_struct_mapping:
+                static_to_struct_mapping[v.local_tensor.name] = k
+
+        master_weights = optimizer_state_dict.pop("master_weights", None)
+        optimizer_state_dict.pop("LR_Scheduler", None)
+
+        # Process main optimizer states
+        for key, tensor in optimizer_state_dict.items():
+            static_name, optim_state_type = _generate_base_static_name(key)
+            struct_name = static_to_struct_mapping[static_name]
+            sharded_weight = model_sharded_state_dict[struct_name]
+
+            unified_name = f"{struct_name}.{optim_state_type}"
+
+            # Determine tensor partitioning scheme
+            if _MOMENT_NAME in optim_state_type:
+                optimizer_sharded_state_dict[unified_name] = (
+                    create_sharded_weight_with_new_local(
+                        unified_name, tensor, sharded_weight
+                    )
+                )
+            else:  # Non-momentum parameters
+                optimizer_sharded_state_dict[unified_name] = ShardedWeight(
+                    key=unified_name,
+                    local_tensor=tensor,
+                    local_shape=(1,),
+                    global_shape=(1,),
+                    global_offset=(0,),
+                )
+
+        # Process master weights if using mixed precision
+        if master_weights is not None:
+            for key, tensor in master_weights.items():
+                struct_name = static_to_struct_mapping[key]
+                sharded_weight = model_sharded_state_dict[struct_name]
+                unified_name = f"{struct_name}.w_0"
+                optimizer_sharded_state_dict[unified_name] = (
+                    create_sharded_weight_with_new_local(
+                        unified_name, tensor, sharded_weight
+                    )
+                )
+
+        return optimizer_sharded_state_dict

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage2.py

@@ -0,0 +1,720 @@
+#   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.
+
+# The file has been adapted from fairscale file:
+# https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/data_parallel/sharded_ddp.py
+# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
+# We retain the following license from the original files:
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from functools import reduce
+from types import MethodType
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import collective
+from paddle.distributed.utils.log_utils import get_logger
+from paddle.framework import core
+
+from .group_sharded_optimizer_stage2 import GroupShardedOptimizerStage2
+from .group_sharded_storage import GradStorage
+from .group_sharded_utils import Type, device_guard
+
+logger_ = get_logger(logging.WARNING)
+
+
+def _trainable(param):
+    return param.trainable
+
+
+class GroupShardedStage2(nn.Layer):
+    """
+    A wrapper for Sharding Stage2 Layer in Dygraph.
+    .. warning: GroupShardedStage2 encapsulates the layer strategy and integrates it into the nn.Layer.
+    .. ZeRO: https://arxiv.org/pdf/1910.02054.pdf.
+    """
+
+    # TODO (Baibaifan)
+    # Feature Notes::
+    # 1. Unified memory for param and param.grad to InternalStorage.
+    # 2. Divide param.grad according to rank to centrally apply for and release GPU memory.
+    # 3. Dynamically adjust training parameters and models.
+    # 4. Support offload function.
+    # 5. Support the establishment of independent communication groups.
+
+    def __init__(
+        self,
+        layer,
+        sharding_optimizer,
+        group=None,
+        sync_buffers=False,
+        buffer_max_size=2**23,  # 8MB
+        auto_refresh_trainable=True,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        dp_group=None,
+    ):
+        super().__init__()
+
+        # training options
+        self._layer = layer
+        self._sharding_optimizers = (
+            [sharding_optimizer]
+            if not isinstance(sharding_optimizer, list)
+            else sharding_optimizer
+        )
+        assert all(
+            isinstance(opt, GroupShardedOptimizerStage2)
+            for opt in self._sharding_optimizers
+        ), "Please use GroupShardedOptimizerStage2 optimizer"
+        self._sync_buffers = sync_buffers
+        self._auto_refresh_trainable = auto_refresh_trainable
+
+        # Communication related attributes
+        self._group = (
+            collective.new_group(collective._get_global_group().ranks)
+            if group is None
+            else group
+        )
+        self._world_size_scaling = 1.0 / self._group.nranks
+        assert self._group.nranks > 1, (
+            "Training must be distributed, ranks must be greater than 1"
+        )
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[
+            0
+        ]  # picking ranks index 0 as the reference
+        self._default_device = device
+
+        self._dp_group = dp_group
+
+        # Global statistical parameters
+        self._all_params = []
+        for optim in self._sharding_optimizers:
+            self._all_params.extend(list(optim.local_params))
+        self.use_main_grad = None
+        for param in self._all_params:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+
+        # sharing stage 2 comm overlap flag
+        self._reduce_overlap = False
+
+        self._grad_reduced = []
+        self._trainable_param2rank = {}
+        self._trainable_param2align = {}
+        self._trainable_params = list(
+            filter(lambda x: x.trainable, self._all_params)
+        )
+        self._trainable_mask = list(map(_trainable, self._trainable_params))
+        self._param_grads = []
+
+        # Set grad storage size & Display param sizes and model sizes
+        model_size = sum([p._numel() for p in self._layer.parameters()])
+        assert buffer_max_size >= 0, "buffer_max_size must be GE than 0."
+        self._buffer_max_size = self._rank_buffer_size(
+            buffer_max_size, model_size
+        )
+        self._use_grad_storage = buffer_max_size > 0
+        self._grad_storages = {}  # {dtype: {rank: GradStorage}}
+        self._has_grad_storage = []
+        self._grad_storage_list = []
+
+        # Offload
+        # TODO(haohongxiang): Now it's not be supported for multi-optimizers using Offload strategy
+        self._offload_optims = list(
+            filter(lambda optim: optim.offload, self._sharding_optimizers)
+        )
+        if len(self._offload_optims) > 0:
+            assert len(self._sharding_optimizers) == 1, (
+                "Only support offload strategy for single optimizer"
+            )
+
+        self._offload = len(self._offload_optims) > 0
+        self._offload_device = "cpu"
+
+        # Set backward pass hooks
+        self._bw_hooks = []
+
+        self.scale_in_opt = False
+
+        # TODO (Baibaifan) Set tasks flow support asynchronous communicate
+        # self._tasks_flow = deque()
+
+        # Define optimizer step and clear_grad
+        self._redefine_opt_step()
+        self._redefine_opt_clear()
+
+    def forward(self, *inputs, **kwargs):
+        """
+        A wrapper for Sharding Stage2 layer.
+        - Fresh trainable params or rebuild grad storage
+        - Sync layer's buffer params
+        - Clear all flags states
+        - Forward for origin layers
+        """
+
+        # Whether to need to reset trainable parameters
+        needs_fresh = len(self._bw_hooks) == 0 and self.training
+
+        if self._auto_refresh_trainable:
+            needs_fresh |= self._detect_train_change()
+
+        # Front hook
+        self._init_internal_storage(needs_fresh)
+
+        # Sync layer's buffers state
+        if self._sync_buffers:
+            self.__sync_buffers()
+
+        # Normal FW on the base model
+        fw = self._layer(*inputs, **kwargs)
+
+        return fw
+
+    def set_state_dict(self, state_dict, use_structured_name=True):
+        self._layer.set_state_dict(
+            state_dict, use_structured_name=use_structured_name
+        )
+
+    def state_dict(
+        self,
+        destination=None,
+        include_sublayers=True,
+        structured_name_prefix="",
+    ):
+        return self._layer.state_dict(
+            destination=destination,
+            include_sublayers=include_sublayers,
+            structured_name_prefix=structured_name_prefix,
+        )
+
+    def _clear_gradients(self):
+        """
+        Set zero to the gradient of the optimizer's current rank trainable parameters.
+        """
+        # Release grad storages
+        for dtype in self._grad_storages.keys():
+            if (
+                not self._offload
+                and self._rank in self._grad_storages[dtype].keys()
+            ):
+                self._grad_storages[dtype][self._rank].buffer.zero_()
+
+        # Release grads of params
+        for param in self._trainable_params:
+            if param.name in self._param_grads:
+                if self.use_main_grad and param.main_grad is not None:
+                    param.main_grad.zero_()
+                elif param.grad is not None:
+                    param._zero_grads()
+
+        # Release grads of master params with offload strategy
+        if self._offload:
+            self._sharding_optimizers[0]._offload_clear_grad()
+
+    def _grad_scale(self):
+        """
+        this function will do 2 things:
+        1.  Before the optimization, scale main_grad to support gradient merge if param has main_grad, or to support fused_linear_param_grad_add gradient merge.
+        2.  Before the optimization, scale the gradients before allreduce of dp_group.
+        """
+
+        need_dp_scale = self._dp_group is not None and self._dp_group.nranks > 1
+        if self.scale_in_opt:
+            scale_factor = self._world_size_scaling
+        else:
+            scale_factor = 1.0
+
+        if need_dp_scale:
+            dp_scale_factor = 1.0 / (self._dp_group.nranks)
+            scale_factor = scale_factor * dp_scale_factor
+
+        # Scale grad storages
+        for dtype in self._grad_storages.keys():
+            if (
+                not self._offload
+                and self._rank in self._grad_storages[dtype].keys()
+            ):
+                self._grad_storages[dtype][self._rank].buffer.scale_(
+                    scale=scale_factor
+                )
+        # Scale grads of params
+        with paddle.no_grad():
+            for param in self._trainable_params:
+                if param.name in self._param_grads:
+                    if self.use_main_grad and param.main_grad is not None:
+                        param.main_grad.scale_(scale=scale_factor)
+                    elif param.grad is not None:
+                        param.grad.scale_(scale=scale_factor)
+
+        # Scale grads of master params with offload strategy
+        if self._offload:
+            if need_dp_scale is False:
+                return
+            self._sharding_optimizers[0]._offload_scale_grad(
+                scale=1.0 / (self._dp_group.nranks)
+            )
+
+    def _init_internal_storage(self, needs_fresh):
+        """
+        Judge Fresh trainable params or rebuild grad storage.
+        """
+        if needs_fresh:
+            self._fresh_trainable()
+        else:
+            self._build_grad_storages()
+
+        # Clear all flags state
+        self._clear_counters()
+
+    def to(self, device=None, dtype=None, blocking=True):
+        """
+        Synchronously or asynchronously convert the data type of the layer, the device is not supported now.
+        """
+        assert isinstance(device, str), "Device must be type str"
+        assert device == self._default_device, (
+            "New devices are not supported, because of the optimizer state is not sync"
+        )
+
+        self._layer.to(device=device, dtype=dtype, blocking=blocking)
+
+        # Re-build the buckets, hooks, etc..
+        self._fresh_trainable()
+
+    def _fresh_trainable(self):
+        """Whether to update training parameters."""
+
+        # Make sure that this is not done while gradients are waiting to be reduced (if no_sync context for instance)
+        if reduce(lambda x, y: x or y, self._grad_reduced, False):
+            logging.warning("Grads waiting to be reduced.")
+
+        self._trainable_params = list(
+            filter(lambda x: x.trainable, self._all_params)
+        )
+        self._trainable_params.sort(key=lambda x: x._numel())
+
+        self._trainable_param2rank = {}
+        for optim in self._sharding_optimizers:
+            # Need to be wrapped for Sharding Stage2 Optimizer
+            if len(optim.param_storages.keys()) == 0:
+                optim._update_opt_status()
+
+            # Get the parameters split by the optimizer according to rank
+            for per_rank_params in (
+                optim.dtype_rank_params.values()
+            ):  # all the params from all ranks
+                for params in per_rank_params:
+                    for param in filter(lambda x: x.trainable, params):
+                        self._trainable_param2rank[param.name] = (
+                            optim.param2rank[param.name]
+                        )
+                        self._trainable_param2align[param.name] = (
+                            optim._param2align[param.name]
+                        )
+
+        # Create grad_storage
+        self._setup_use_grad_storage()
+        # setup backward hooks
+        self._setup_backward_hooks()
+
+    @paddle.autograd.no_grad()
+    def __sync_buffers(self):
+        """
+        Sync all the param buffers from all ranks (exp: batch norm statistics).
+        """
+
+        for buffer in self._layer.buffers(include_sublayers=True):
+            dist.broadcast(
+                buffer, self._global_root_rank, self._group, sync_op=True
+            )
+
+            if self._dp_group and self._dp_group.nranks > 1:
+                dist.broadcast(
+                    buffer,
+                    self._dp_group.ranks[0],
+                    self._dp_group,
+                    sync_op=True,
+                )
+
+    def __getattr__(self, name):
+        """Forward missing attributes to wrapped layer."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._layer, name)
+
+    @paddle.autograd.no_grad()
+    def _clear_counters(self):
+        """Reset all the grad reduce and call counters."""
+        if self.training:
+            self._grad_reduced = [True for _ in self._trainable_params]
+
+        if self._use_grad_storage:
+            for grad_storage in self._grad_storage_list:
+                grad_storage.reset_checked_in()
+
+    def _set_reduce_overlap(self, reduce_overlap):
+        # Hacky way to not add an extra parameter to the `group_sharded_parallel` funct.
+        # User should use this like:
+        # model, optimizer, scaler = group_sharded_parallel(...)
+        # model._set_reduce_overlap(True)
+        self._reduce_overlap = reduce_overlap
+        if self._reduce_overlap:
+            assert len(self._sharding_optimizers) == 1, (
+                "Only support comm overlap strategy for single optimizer"
+            )
+        self._sharding_optimizers[0]._set_reduce_overlap(reduce_overlap)
+
+    def _get_scaled_grad_fn(self, param):
+        @paddle.autograd.no_grad()
+        def scale(grad):
+            # do gradient scale separately
+            # For grad scale, we need to do it in the backward hook due to fp16 may overflow if we first add grad and then scale
+            # For main_grad scale and fused_linear_param_grad_add, we do scale in the optimizer.
+            if not self.scale_in_opt:
+                if (
+                    not hasattr(param, "main_grad")
+                    and grad is not None
+                    and grad.dtype == Type.fp16.value
+                ):
+                    assert grad._is_initialized(), (
+                        "grad should be initialized in stage2"
+                    )
+                    grad.scale_(self._world_size_scaling)
+                else:
+                    self.scale_in_opt = True
+
+        return scale
+
+    def _get_reduce_fn(self, index, param, dst_rank):
+        """
+        There are two ways to reduce gradient.
+        - 1. Do not use self._use_grad_storage or exceeded buffer_max_size will be reduced separately.
+        - 2. Use grad_storage Reduce the storage to get the full gradient from different ranks.
+        """
+
+        if not self._use_grad_storage or not self._has_grad_storage[index]:
+            # Direct reduction
+            @paddle.autograd.no_grad()
+            def reduce(*_):
+                # Skip gradient reduction, do not change status information
+                if self._grad_reduced[index]:
+                    assert (
+                        param.grad is not None or param.main_grad is not None
+                    ), "Parameter should have grad or main grad"
+
+                    # Change reduce information
+                    self._grad_reduced[index] = False
+
+                    # Clear the gradient that does not belong to the current rank through the callback function
+                    def cleanup():
+                        if dst_rank != self._rank:
+                            if self.use_main_grad:
+                                param.main_grad._clear_data()
+                                param.main_grad = None
+                            else:
+                                param.clear_gradient(False)
+                        elif self._offload:
+                            tmp_grad = param.grad.cast(
+                                dtype=Type.fp32.value
+                            ).cpu()
+
+                            self._sharding_optimizers[0]._offload_acc_grad(
+                                param.name, tmp_grad
+                            )
+                            del tmp_grad
+                            param.clear_gradient(False)
+
+                    # Synchronize the reduce parameter gradient asynchronize
+                    self._sharding_optimizers[0]._update_task(
+                        dist.reduce(
+                            tensor=(
+                                param.grad
+                                if not self.use_main_grad
+                                else param.main_grad
+                            ),
+                            dst=self._group.ranks[dst_rank],
+                            group=self._group,
+                            sync_op=not self._reduce_overlap,
+                        )
+                    )
+
+                    # Clear the task flow and trigger callback to clear the redundant gradient
+                    # self._clear_task_flow()
+
+                    cleanup()
+
+        else:
+            # Buffer reduction
+            @paddle.autograd.no_grad()
+            def reduce(*_):
+                # Skip gradient reduction, do not change status information
+                if self._grad_reduced[index]:
+                    assert (
+                        param.grad is not None or param.main_grad is not None
+                    ), "Parameter should have grad or main grad"
+
+                    # Change reduce information
+                    self._grad_reduced[index] = False
+                    grad_storage = self._grad_storages[param.dtype][dst_rank]
+                    grad_storage.params_checked_in += 1
+
+                    if grad_storage.all_checked_in:
+                        assert grad_storage.buffer is not None
+
+                        # Clearing up the grad_storage buffer
+                        def cleanup():
+                            if dst_rank != self._rank:
+                                for p in grad_storage._params:
+                                    if self.use_main_grad:
+                                        p.main_grad._clear_data()
+                                        p.main_grad = None
+                                    else:
+                                        p.clear_gradient(False)
+
+                                grad_storage.buffer._clear_data()
+                            elif self._offload:
+                                grad_storage.to(device=self._offload_device)
+                                for p in grad_storage._params:
+                                    with device_guard():
+                                        tmp_grad = p.grad.cast(
+                                            dtype=Type.fp32.value
+                                        )
+                                    self._sharding_optimizers[
+                                        0
+                                    ]._offload_acc_grad(p.name, tmp_grad)
+                                    p.clear_gradient(False)
+                                grad_storage._device = self._default_device
+                                grad_storage.buffer._clear_data()
+
+                        # Reduce the bucket
+                        grad_storage.sent = True
+                        # Synchronize the reduce parameter gradient asynchronize
+                        self._sharding_optimizers[0]._update_task(
+                            dist.reduce(
+                                tensor=grad_storage.buffer,
+                                dst=self._group.ranks[grad_storage.destination],
+                                group=self._group,
+                                sync_op=not self._reduce_overlap,
+                            )
+                        )
+
+                        cleanup()
+
+                    # Clear the task flow and trigger callback to clear the redundant gradient
+                    # self._clear_task_flow()
+
+        return reduce
+
+    def _setup_backward_hooks(self):
+        """
+        Set the backward hook to synchronize the gradients of all rank by reduce group ranks.
+        """
+
+        # Remove previous backward hooks
+        while len(self._bw_hooks) > 0:
+            self._bw_hooks.pop().remove()
+
+        # Go through the parameters, attach the hook
+        if not self.training:
+            return
+
+        for index, param in enumerate(self._trainable_params):
+            param._register_grad_hook(self._get_scaled_grad_fn(param))
+
+            dst_rank = self._trainable_param2rank[param.name]
+
+            reduce_function = self._get_reduce_fn(index, param, dst_rank)
+
+            self._bw_hooks.append(
+                param._register_backward_hook(reduce_function)
+            )
+
+    def _setup_use_grad_storage(self):
+        """
+        Integrate the parameters gradient into a continuous memory according to rank, and support the update of training parameters.
+        """
+
+        # According to parameters's numel sort, allocate memory of parameter gradient to continuous memory according to rank
+        self._grad_storages = {}
+        self._has_grad_storage = [False for _ in self._trainable_params]
+
+        for index, param in enumerate(self._trainable_params):
+            dst_rank = self._trainable_param2rank[param.name]
+
+            if param.dtype not in self._grad_storages.keys():
+                self._grad_storages[param.dtype] = {}
+
+            if dst_rank not in self._grad_storages[param.dtype].keys():
+                self._grad_storages[param.dtype][dst_rank] = GradStorage(
+                    self._buffer_max_size[param.dtype],
+                    dtype=(
+                        param.dtype
+                        if not self.use_main_grad
+                        else paddle.float32
+                    ),
+                    device=self._default_device,
+                    destination=dst_rank,
+                    param2align=self._trainable_param2align,
+                )
+
+            # Criteria to decide whether this parameter is to be put in GradStorage
+            if self._grad_storages[param.dtype][dst_rank].can_add_grad_view(
+                param, self._trainable_param2align[param.name]
+            ):
+                self._grad_storages[param.dtype][dst_rank].add_grad(
+                    param, self._trainable_param2align[param.name]
+                )
+                self._has_grad_storage[index] = True
+            else:
+                self._param_grads.append(param.name)
+
+        for dtype in self._grad_storages.keys():
+            self._grad_storage_list.extend(
+                list(self._grad_storages[dtype].values())
+            )
+
+    # def _clear_task_flow(self):
+    #     """Try to consume the previous tasks."""
+    #     while len(self._tasks_flow) > 0:
+    #         task = self._tasks_flow.popleft()
+    #         task.wait()
+    #         if task.callback is not None:
+    #             task.callback()
+
+    def _detect_train_change(self):
+        # Current trainable parameters
+        trainable_mask = list(map(_trainable, self._trainable_params))
+
+        # Whether parameters trainability changed
+        trainability_changed = trainable_mask != self._trainable_mask
+
+        if trainability_changed:
+            logging.warning(
+                "Trainable params changed, because of eval/train mode or parameter freezing/unfreeze."
+            )
+            self._trainable_mask = trainable_mask
+
+        return trainability_changed
+
+    def _build_grad_storages(self):
+        """
+        Rebuild grad storages.
+        """
+        # Rebuild fp16/fp32 grad storages
+        for dtype in self._grad_storages.keys():
+            for dst_rank, grad_storage in self._grad_storages[dtype].items():
+                if self._offload or dst_rank != self._rank:
+                    grad_storage.manual_release()
+                    grad_storage.rebuild()
+
+    def _rank_buffer_size(self, buffer_max_size, model_size):
+        """
+        Generate the minimum buffer size for each rank & Display param sizes and model sizes.
+        """
+
+        # Initialize buffer size
+        rank_buffer_size = {}
+        for shard_opt in self._sharding_optimizers:
+            if shard_opt.rank_buffer_size:
+                for dtype in shard_opt.rank_buffer_size.keys():
+                    sizes = max(shard_opt.rank_buffer_size[dtype].values())
+                    rank_buffer_size[dtype] = min(sizes, buffer_max_size)
+
+        if Type.fp16.value in rank_buffer_size.keys():
+            # FP16 GradStorage and model size
+            logger_.info(
+                f"====== FP16 GradStorage size: {rank_buffer_size[Type.fp16.value] / 2**19:.2f}M parameters, Model size {model_size / 2**19:.2f}M parameters ======"
+            )
+        if Type.bf16.value in rank_buffer_size.keys():
+            # FP16 GradStorage and model size
+            logger_.info(
+                f"====== BF16 GradStorage size: {rank_buffer_size[Type.bf16.value] / 2**19:.2f}M parameters, Model size {model_size / 2**19:.2f}M parameters ======"
+            )
+        if Type.fp32.value in rank_buffer_size.keys():
+            # FP32 GradStorage and model size
+            logger_.info(
+                f"====== FP32 GradStorage size: {rank_buffer_size[Type.fp32.value] / 2**18:.2f}M parameters, Model size {model_size / 2**18:.2f}M parameters ======"
+            )
+        return rank_buffer_size
+
+    def _dp_allreduce(self):
+        # do dp allreduce here for gradient merge.
+        if self._dp_group and self._dp_group.nranks > 1:
+            for dtype in self._grad_storages.keys():
+                for rank, g in sorted(
+                    self._grad_storages[dtype].items(), key=lambda x: x[0]
+                ):
+                    if g.destination == self._rank:
+                        assert g.buffer._is_initialized()
+                        dist.all_reduce(
+                            tensor=g.buffer,
+                            group=self._dp_group,
+                            sync_op=True,
+                        )
+            for param in self._trainable_params:
+                if param.name in self._param_grads:
+                    if self.use_main_grad and param.main_grad is None:
+                        continue
+                    elif param.grad is None:
+                        continue
+                    dst_rank = self._trainable_param2rank[param.name]
+                    if dst_rank == self._rank:
+                        dist.all_reduce(
+                            tensor=(
+                                param.grad
+                                if not self.use_main_grad
+                                else param.main_grad
+                            ),
+                            group=self._dp_group,
+                            sync_op=True,
+                        )
+
+    def _redefine_opt_step(self):
+        grad_func = self._grad_scale
+        dp_allreduce_func = self._dp_allreduce
+
+        for opt in self._sharding_optimizers:
+            opt_step = opt.step
+
+            def _opt_step(self):
+                if self._reduce_overlap:
+                    # Wait for the last reduce task. This wait must before grad scale function.
+                    assert self._comm_task is not None
+                    self._comm_task.wait()
+
+                grad_func()
+                dp_allreduce_func()
+                opt_step()
+
+            opt.step = MethodType(_opt_step, opt)
+
+    def _redefine_opt_clear(self):
+        clear_func = self._clear_gradients
+
+        def _opt_clear(self):
+            clear_func()
+
+        for opt in self._sharding_optimizers:
+            opt.clear_grad = MethodType(_opt_clear, opt)

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_storage.py

@@ -0,0 +1,363 @@
+#   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.
+
+# The file has been adapted from fairscale file:
+# https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/misc/param_bucket.py
+# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
+# We retain the following license from the original files:
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+
+from .group_sharded_utils import Type, cvt_to_device, device_guard
+
+
+class BufferWarper(core.eager.Tensor):
+    def __init__(self):
+        super().__init__()
+        self.need_clip = True
+        self.is_distributed = False
+        self.trainable = True
+
+
+class InternalStorage:
+    """
+    This is a basic class, which is responsible for consolidating the basic storage tensor.
+
+    """
+
+    # Support integration parameter tensor
+    def __init__(self, size, dtype, device, convert_cpu=False):
+        self._params = []
+        self._param_ids = []
+        self._fill = 0
+        self._device = device
+        self._dtype = dtype
+
+        # The flatten tensor
+        size = [size] if isinstance(size, int) else size
+        if convert_cpu:
+            value = (
+                np.zeros(size, dtype=np.float16)
+                if Type.fp16.value == dtype
+                else np.zeros(size, dtype=np.float32)
+            )
+            self.buffer = core.eager.Tensor(value=value, place=core.CPUPlace())
+            if dtype == Type.bf16.value:
+                self.buffer = paddle.cast(self.buffer, dtype=paddle.bfloat16)
+        else:
+            self.buffer = paddle.zeros(size, dtype=dtype)
+
+        self.dev_id = (
+            0
+            if paddle.get_device() == "cpu"
+            else int(paddle.get_device().split(":")[1])
+        )
+
+    def to(self, device, dtype=None, keep_alignment=True):
+        """
+        Move the underlying buffer
+        """
+        assert self.buffer is not None, (
+            "Cannot move a collapsed bucket, please rebuild it"
+        )
+        assert dtype == Type.fp32.value or Type.fp16.value, (
+            "Conversion type is not supported now"
+        )
+
+        if self._device != device:
+            if device in paddle.device.get_all_custom_device_type():
+                tmp_buffer = self.buffer._copy_to(
+                    paddle.CustomPlace(device, self.dev_id), True
+                )
+            else:
+                tmp_buffer = (
+                    cvt_to_device(self.buffer, self.dev_id)
+                    if device in ["gpu", "xpu"]
+                    else self.buffer.cpu()
+                )
+            for param in self._params:
+                param.clear_gradient(False)
+
+            del self.buffer
+            self.buffer = tmp_buffer
+            self._device = device
+
+        if dtype is not None:
+            self.buffer = self.buffer.cast(dtype=dtype)
+            self._dtype = dtype
+
+    def warp_buffer(self):
+        tmp_buffer = BufferWarper()
+        self._buffer = self.buffer
+        tmp_buffer.get_tensor()._share_data_with(self.buffer.get_tensor())
+        self.buffer = tmp_buffer
+
+
+class ParamStorage(InternalStorage):
+    """
+    This is a basic class to simplify the handling of parameter InternalStorages.
+    """
+
+    def __init__(self, size, dtype, device):
+        super().__init__(size, dtype, device, convert_cpu=True)
+        self.param2align = None
+
+    def to(self, device, dtype=None, keep_alignment=True):
+        """
+        Move the underlying buffer
+        """
+
+        super().to(device, dtype)
+
+        if keep_alignment:
+            self._array_params()
+
+    @paddle.autograd.no_grad()
+    def add_rank_params(self, trainable_params, param2align, convert_gpu=True):
+        """
+        Add new parameters to the InternalStorage. Params becomes a view of this InternalStorage buffer.
+        """
+
+        assert all(
+            id(param) not in self._param_ids for param in trainable_params
+        ), "The same param cannot be checked in twice"
+        assert self.buffer is not None
+
+        self.param2align = param2align
+
+        cpu_param_shape = []
+        for param in trainable_params:
+            p_shape = self._add_param_as_view(
+                param, param2align[param.name], convert_gpu
+            )
+            cpu_param_shape.append(p_shape)
+
+        if convert_gpu:
+            if self._device in paddle.device.get_all_custom_device_type():
+                self.buffer = self.buffer._copy_to(
+                    paddle.CustomPlace(self._device, self.dev_id), True
+                )
+            else:
+                # buffer convert from cpu to cuda
+                self.buffer = cvt_to_device(self.buffer, self.dev_id)
+
+        self._fill = 0
+
+        for idx, param in enumerate(trainable_params):
+            self._convert_buffer(
+                param, cpu_param_shape[idx], param2align[param.name]
+            )
+            self._params.append(param)
+            self._param_ids.append(id(param))
+
+    @paddle.autograd.no_grad()
+    def _add_param_as_view(self, param, align, convert_gpu=True):
+        assert param.dtype == self.buffer.dtype, (
+            f"Different types for the InternalStorage and the param, cannot proceed: {param.dtype} - {self.buffer.dtype}"
+        )
+
+        var_end = self._fill + param._numel()
+        offset = var_end + align
+        assert offset <= self.buffer._numel()
+
+        p_shape = param.shape
+
+        origin_state = param.stop_gradient
+        param.stop_gradient = True
+        param.flatten_()
+        param.stop_gradient = origin_state
+
+        # Copy the current param value
+
+        with device_guard(self.dev_id, "cpu"):
+            tmp_var = self.buffer._slice(self._fill, var_end)
+            if convert_gpu:
+                param_cpu = param.cpu()
+                param._clear_data()
+                tmp_var.set_value(param_cpu)
+            else:
+                tmp_var.set_value(param)
+            del tmp_var
+
+        self._fill = offset
+        return p_shape
+
+    @paddle.autograd.no_grad()
+    def _convert_buffer(self, param, p_shape, align):
+        var_end = self._fill + np.prod(p_shape).tolist()
+        offset = var_end + align
+        assert offset <= self.buffer._numel()
+
+        # Convert the param value
+        with device_guard(self.dev_id, self._device):
+            tmp_tensor = self.buffer._slice(self._fill, var_end)
+            tmp_tensor._share_buffer_to(param)
+            param.get_tensor()._set_dims(p_shape)
+
+        self._fill = offset
+
+    @paddle.autograd.no_grad()
+    def _array_params(self):
+        """
+        Given the parameters which have been registered previously, rebuild the whole InternalStorage.
+        """
+        assert len(self._params) > 0
+        assert self.param2align is not None
+
+        self._fill = 0
+        for p in self._params:
+            self._convert_buffer(p, p.shape, self.param2align[p.name])  # modify
+
+
+class GradStorage(InternalStorage):
+    """
+    This is a basic class to simplify the handling of gradient InternalStorages
+    """
+
+    def __init__(
+        self, size, dtype, device, destination, param2align, convert_cpu=False
+    ):
+        if isinstance(size, np.int64):
+            size = size.tolist()
+        super().__init__(size, dtype, device, convert_cpu)
+
+        self._max_size = size
+        self._release = False
+
+        self.params_checked_in = 0
+        self.destination = destination
+        self._param2align = param2align
+        self.sent = False
+
+    def reset_checked_in(self):
+        """Reset the counter of the parameter grads which have been checked in"""
+        self.params_checked_in = 0
+        self.sent = False
+
+    @property
+    def all_checked_in(self):
+        """Judge all the expected gradient check-in happened"""
+        return len(self._params) == self.params_checked_in
+
+    def can_add_grad_view(self, param, align):
+        """Is there enough InternalStorage to add this parameter gradient, and whether this param have already checked in."""
+        return (
+            self._fill + param._numel() + align <= self._max_size
+            and id(param) not in self._param_ids
+        )
+
+    def to(self, device, dtype=None, keep_alignment=True):
+        """
+        Move the underlying buffer
+        """
+        if self._release:
+            self.rebuild()
+
+        super().to(device, dtype)
+
+        if keep_alignment:
+            self._array_grads()
+
+    @paddle.autograd.no_grad()
+    def add_grad(self, param, align):
+        """
+        Add a new parameter gradient to the InternalStorage. Param.grad becomes a view of this InternalStorage buffer.
+        """
+
+        assert id(param) not in self._param_ids, (
+            "The same gradients cannot be checked in twice"
+        )
+
+        self._add_grad_as_view(param, align)
+        self._params.append(param)
+        self._param_ids.append(id(param))
+
+    @paddle.autograd.no_grad()
+    def manual_release(self):
+        """
+        Release the buffer from InternalStorage. The InternalStorage will need to be rebuilt before use.
+        """
+        if not self._release:
+            for p in self._params:
+                use_main_grad = hasattr(p, "main_grad")
+                if use_main_grad and p.main_grad is not None:
+                    p.main_grad._clear_data()
+                    p.main_grad = None
+                elif p.grad is not None:
+                    p.clear_gradient(False)
+
+            self.buffer = None
+            self._fill = 0
+            self.params_checked_in = 0
+            self._release = True
+
+    @paddle.autograd.no_grad()
+    def rebuild(self):
+        """
+        Given the parameter gradients which have been registered previously, rebuild the whole InternalStorage.
+        """
+
+        if self._release:
+            self.buffer = paddle.zeros([self._max_size], dtype=self._dtype)
+
+            for p in self._params:
+                self._add_grad_as_view(p, self._param2align[p.name])
+
+            self._release = False
+
+    @paddle.autograd.no_grad()
+    def _array_grads(self):
+        """
+        Given the parameters gradients which have been registered previously, rebuild the whole InternalStorage.
+        """
+        if len(self._params) > 0:
+            self._fill = 0
+            for p in self._params:
+                self._add_grad_as_view(p, self._param2align[p.name])
+
+    @paddle.autograd.no_grad()
+    def _add_grad_as_view(self, param, align):
+        assert param._numel() > 0, (
+            "Cannot add a gradient to a released InternalStorage, please rebuild"
+        )
+
+        use_main_grad = hasattr(param, "main_grad")
+        if use_main_grad:
+            assert self.buffer.dtype == paddle.float32
+        else:
+            assert param.dtype == self.buffer.dtype
+
+        grad_end = self._fill + param._numel()
+        offset = grad_end + align
+        assert offset <= self.buffer._numel()
+
+        # Copy the current grad value to InternalStorage
+        with device_guard(self.dev_id, self._device):
+            tmp_var = self.buffer._slice(self._fill, grad_end)
+            tmp_var.get_tensor()._set_dims(param.shape)
+            if not use_main_grad:
+                param._copy_gradient_from(tmp_var)
+            else:
+                param.main_grad = tmp_var
+            del tmp_var
+
+        self._fill = offset

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_utils.py

@@ -0,0 +1,352 @@
+#   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 contextlib
+from enum import Enum
+from types import MethodType
+
+import numpy as np
+
+import paddle
+from paddle import _C_ops, _legacy_C_ops
+from paddle.base import core
+from paddle.common_ops_import import dygraph_only
+from paddle.nn import clip
+
+
+class Taskflow:
+    """
+    Task flows, one way linked list for task acquisition.
+    """
+
+    def __init__(self, task, callback):
+        self.task = task
+        self.callback = callback
+
+
+class Type(Enum):
+    """
+    Type of trainable parameters
+    """
+
+    fp16 = paddle.float16
+    bf16 = paddle.bfloat16
+    fp32 = paddle.float32
+
+
+class GroupShardedClipGrad:
+    def __init__(self, clip, device, group):
+        self._clip = clip
+        self._device = device
+        self._group = group
+
+    @paddle.autograd.no_grad()
+    def _dygraph_clip(self, params_grads):
+        sum_square_fp32, sum_square_fp16, sum_square_bfp16 = [], [], []
+        unslice_params_fp32, unslice_params_fp16, unslice_params_bfp16 = (
+            [],
+            [],
+            [],
+        )
+
+        for p, g in params_grads:
+            p_slice = True  # using for slice parameter in sharding stage3
+            if g is None or getattr(p, 'need_clip', True) is False:
+                continue
+            if hasattr(p, "unslice"):
+                p_slice = False
+
+            merge_grad = g
+            if g.type == core.VarDesc.VarType.SELECTED_ROWS:
+                merge_grad = clip.get_tensor_from_selected_rows(
+                    clip.merge_selected_rows(g)
+                )
+            square = paddle.square(merge_grad)
+            sum_square = paddle.sum(square)
+
+            if p.dtype == paddle.float16:
+                if p_slice:
+                    sum_square_fp16.append(sum_square)
+                else:
+                    unslice_params_fp16.append(sum_square)
+            elif p.dtype == paddle.float32:
+                if p_slice:
+                    sum_square_fp32.append(sum_square)
+                else:
+                    unslice_params_fp32.append(sum_square)
+            elif p.dtype == paddle.bfloat16:
+                if p_slice:
+                    sum_square_bfp16.append(sum_square)
+                else:
+                    unslice_params_bfp16.append(sum_square)
+
+        # global norm of non-distributed FP16 params_and_grads
+        if len(sum_square_fp16) == 0:
+            global_norm_fp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_norm_fp16 = paddle.add_n(sum_square_fp16)
+            global_norm_fp16 = paddle.cast(
+                global_norm_fp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed BFP16 params_and_grads
+        if len(sum_square_bfp16) == 0:
+            global_norm_bfp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_norm_bfp16 = paddle.add_n(sum_square_bfp16)
+            global_norm_bfp16 = paddle.cast(
+                global_norm_bfp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed FP16 params_and_grads for unslice parameters
+        if len(unslice_params_fp16) == 0:
+            global_unslice_fp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_unslice_fp16 = paddle.add_n(unslice_params_fp16)
+            global_unslice_fp16 = paddle.cast(
+                global_unslice_fp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed BFP16 params_and_grads for unslice parameters
+        if len(unslice_params_bfp16) == 0:
+            global_unslice_bfp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_unslice_bfp16 = paddle.add_n(unslice_params_bfp16)
+            global_unslice_bfp16 = paddle.cast(
+                global_unslice_bfp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed FP32 params_and_grads
+        if len(sum_square_fp32) == 0:
+            global_norm_fp32 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_norm_fp32 = paddle.add_n(sum_square_fp32)
+
+        # global norm of non-distributed FP32 params_and_grads for unslice parameters
+        if len(unslice_params_fp32) == 0:
+            global_unslice_fp32 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_unslice_fp32 = paddle.add_n(unslice_params_fp32)
+
+        global_unslice_var = (
+            global_unslice_fp16 + global_unslice_fp32 + global_unslice_bfp16
+        )
+
+        global_norm_var = (
+            global_norm_fp16 + global_norm_fp32 + global_norm_bfp16
+        )
+
+        # add all reduce to get global norm of distributed params_and_grads
+        dev_id = int(self._device.split(":")[1])
+        dev_type = self._device.split(':')[0]
+        if paddle.device.get_device() == "cpu":
+            if dev_type in paddle.device.get_all_custom_device_type():
+                global_norm_var = global_norm_var._copy_to(
+                    paddle.CustomPlace(dev_type, dev_id), True
+                )
+            elif dev_type == "xpu":
+                global_norm_var = global_norm_var.to(self._device)
+            else:
+                global_norm_var = global_norm_var.cuda(dev_id)
+
+        with device_guard(dev_id, self._device.split(":")[0]):
+            paddle.distributed.all_reduce(global_norm_var, group=self._group)
+
+        global_norm_var = paddle.sqrt(global_norm_var + global_unslice_var)
+        max_global_norm = paddle.full(
+            shape=[], dtype=global_norm_var.dtype, fill_value=self.clip_norm
+        )
+
+        clip_var = paddle.divide(
+            x=max_global_norm,
+            y=paddle.maximum(x=global_norm_var, y=max_global_norm),
+        )
+        clip_var_fp16 = paddle.cast(clip_var, paddle.float16)
+
+        for p, g in params_grads:
+            if getattr(p, 'need_clip', True) is False or g is None:
+                continue
+            origin_state = g.stop_gradient
+            g.stop_gradient = True
+            if p.dtype == paddle.float16:
+                g.scale_(clip_var_fp16)
+            else:
+                g.scale_(clip_var)
+            g.stop_gradient = origin_state
+            # p._reset_grad_inplace_version(True)
+
+        return params_grads
+
+    def __getattr__(self, item):
+        return getattr(self._clip, item)
+
+    def __call__(self, params_grads):
+        return self._dygraph_clip(params_grads)
+
+
+@contextlib.contextmanager
+def device_guard(dev_id=0, device="cpu"):
+    origin_device = paddle.device.get_device()
+    if device == "cpu":
+        paddle.set_device(device)
+    elif device in ["gpu", "xpu"]:
+        paddle.set_device(f"{device}:{dev_id}")
+    elif device in paddle.device.get_all_custom_device_type():
+        paddle.set_device(f"{device}:{dev_id}")
+
+    try:
+        yield
+    finally:
+        paddle.set_device(origin_device)
+
+
+@dygraph_only
+def GroupShardedScaler(scaler):
+    def unscale_method(self, optimizer):
+        if not self._enable:
+            return
+        param_grads = []
+        param_grads_bfp16 = []
+        param_grads_fp16 = []
+        param_grads_fp32 = []
+        if hasattr(optimizer, "update_slice"):
+            optimizer.update_slice()
+            optimizer.update_scaler = True
+
+        if getattr(optimizer._optim, '_param_groups', None) and isinstance(
+            optimizer._optim._param_groups[0], dict
+        ):
+            for group in optimizer._optim._param_groups:
+                for param in group['params']:
+                    tgt_grad = None
+                    if (
+                        hasattr(param, "main_grad")
+                        and param.main_grad is not None
+                    ):
+                        tgt_grad = param.main_grad
+                    elif param.grad is not None:
+                        tgt_grad = param.grad
+                    if tgt_grad is not None:
+                        param_grads.append(tgt_grad)
+                        if tgt_grad.dtype in [
+                            core.VarDesc.VarType.FP16,
+                            paddle.float16,
+                        ]:
+                            param_grads_fp16.append(tgt_grad)
+                        elif tgt_grad.dtype in [paddle.bfloat16]:
+                            param_grads_bfp16.append(tgt_grad)
+                        else:
+                            param_grads_fp32.append(tgt_grad)
+        else:
+            for param in optimizer._optim._parameter_list:
+                tgt_grad = None
+                if hasattr(param, "main_grad") and param.main_grad is not None:
+                    tgt_grad = param.main_grad
+                elif param.grad is not None:
+                    tgt_grad = param.grad
+                if tgt_grad is not None:
+                    param_grads.append(tgt_grad)
+                    if tgt_grad.dtype in [
+                        core.VarDesc.VarType.FP16,
+                        paddle.float16,
+                    ]:
+                        param_grads_fp16.append(tgt_grad)
+                    elif tgt_grad.dtype in [paddle.bfloat16]:
+                        param_grads_bfp16.append(tgt_grad)
+                    else:
+                        param_grads_fp32.append(tgt_grad)
+
+        temp_found_inf_fp16 = paddle.to_tensor(np.array([0]).astype(np.bool_))
+        temp_found_inf_bfp16 = paddle.to_tensor(np.array([0]).astype(np.bool_))
+        temp_found_inf_fp32 = paddle.to_tensor(np.array([0]).astype(np.bool_))
+
+        device = paddle.get_device().split(":")[0]
+        device = "cpu" if optimizer.offload else device
+        dev_id = (
+            0 if device == "cpu" else int(paddle.get_device().split(":")[1])
+        )
+
+        self._found_inf = self._temp_found_inf_value_false
+        with device_guard(dev_id, device):
+            if len(param_grads_bfp16):
+                _legacy_C_ops.check_finite_and_unscale(
+                    param_grads_bfp16,
+                    self._scale,
+                    param_grads_bfp16,
+                    temp_found_inf_bfp16,
+                )
+                self._found_inf = _C_ops.bitwise_or(
+                    self._found_inf, temp_found_inf_bfp16
+                )
+            if len(param_grads_fp16):
+                _legacy_C_ops.check_finite_and_unscale(
+                    param_grads_fp16,
+                    self._scale,
+                    param_grads_fp16,
+                    temp_found_inf_fp16,
+                )
+                self._found_inf = _C_ops.bitwise_or(
+                    self._found_inf, temp_found_inf_fp16
+                )
+            if len(param_grads_fp32):
+                _legacy_C_ops.check_finite_and_unscale(
+                    param_grads_fp32,
+                    self._scale,
+                    param_grads_fp32,
+                    temp_found_inf_fp32,
+                )
+                self._found_inf = _C_ops.bitwise_or(
+                    self._found_inf, temp_found_inf_fp32
+                )
+
+        self._found_inf = self._found_inf.cast("int32")
+
+        paddle.distributed.all_reduce(
+            self._found_inf, op=paddle.distributed.ReduceOp.MAX, group=None
+        )
+
+        self._found_inf = self._found_inf.cast("bool")
+
+    scaler._unscale = MethodType(unscale_method, scaler)
+    return scaler
+
+
+def cvt_to_device(x, dev_id, blocking=True):
+    """
+    Copy data in x from cpu memory to supported device
+    """
+    if paddle.is_compiled_with_cuda():
+        place = paddle.CUDAPlace(dev_id)
+    elif paddle.is_compiled_with_xpu():
+        place = paddle.XPUPlace(dev_id)
+    else:
+        supported_custom_devices = ["npu"]
+        place = paddle.framework._current_expected_place()
+        if place.get_device_type() not in supported_custom_devices:
+            raise OSError(
+                "Only supported compiled paddle with gpu/rocm and xpu, but current version is compiled with cpu."
+            )
+    return x._copy_to(place, blocking)

python/paddle/distributed/fleet/meta_parallel/sharding_parallel.py

@@ -0,0 +1,37 @@
+#   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.hybrid_parallel_util import (
+    broadcast_dp_parameters,
+    broadcast_sharding_parameters,
+)
+from ..utils.log_util import logger
+from .meta_parallel_base import MetaParallelBase
+
+__all__ = []
+
+
+class ShardingParallel(MetaParallelBase):
+    def __init__(self, layers, hcg, **kwargs):
+        super().__init__(layers, hcg, **kwargs)
+
+    def _prepare_for_model(self):
+        logger.info("start broadcast sharding parameters")
+        broadcast_sharding_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_data_parallel_world_size() > 1:
+            logger.info("start broadcast dp parameters")
+            broadcast_dp_parameters(self._layers, self._hcg)
+
+        logger.info("sharding's parameters is ready")

python/paddle/distributed/fleet/meta_parallel/tensor_parallel.py

@@ -0,0 +1,66 @@
+#   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.hybrid_parallel_util import (
+    broadcast_dp_parameters,
+    broadcast_input_data,
+    broadcast_moe_sharding_parameters,
+    broadcast_mp_parameters,
+    broadcast_sep_parameters,
+    broadcast_sharding_parameters,
+)
+from ..utils.log_util import logger
+from .meta_parallel_base import MetaParallelBase
+
+__all__ = []
+
+
+class TensorParallel(MetaParallelBase):
+    def __init__(self, layers, hcg, **kwargs):
+        super().__init__(layers, hcg, **kwargs)
+
+    def _prepare_for_model(self):
+        logger.info("start broadcast mp parameters")
+        broadcast_mp_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_sep_parallel_world_size() > 1:
+            logger.info("start broadcast sep parameters")
+            broadcast_sep_parameters(self._layers, self._hcg, fuse_params=False)
+
+        if self._hcg.get_sharding_parallel_world_size() > 1:
+            logger.info("start broadcast sharding parameters")
+            broadcast_sharding_parameters(
+                self._layers, self._hcg, fuse_params=False
+            )
+
+        if self._hcg.get_data_parallel_world_size() > 1:
+            logger.info("start broadcast dp parameters")
+            broadcast_dp_parameters(self._layers, self._hcg, fuse_params=False)
+
+        if self._hcg.get_moe_sharding_parallel_world_size() > 1:
+            logger.info("start broadcast moe sharding parameters")
+            broadcast_moe_sharding_parameters(
+                self._layers, self._hcg, fuse_params=False
+            )
+
+        logger.info("mp's parameters is ready")
+
+    def _pre_forward(self, *inputs, **kwargs):
+        need_broadcast_data = True
+        if self._strategy is not None:
+            mp_configs = self._strategy.hybrid_configs["mp_configs"]
+            need_broadcast_data = mp_configs.need_broadcast_data
+        if need_broadcast_data:
+            logger.debug("mp start broadcast input data")
+            return broadcast_input_data(self._hcg, *inputs, **kwargs)

python/paddle/distributed/fleet/meta_parallel/zero_bubble_utils.py

@@ -0,0 +1,133 @@
+# 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.
+
+# The file has been adapted from DeepSeek DualPipe project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DualPipe/blob/main/LICENSE
+
+
+import queue
+from functools import partial
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.autograd import PyLayer
+
+
+class WeightGradStore:
+    enabled = False
+    cache = []
+    funcs_queue = queue.Queue()
+
+    @classmethod
+    def put(cls, func) -> None:
+        cls.cache.append(func)
+
+    @classmethod
+    def flush(cls) -> None:
+        cls.funcs_queue.put(cls.cache)
+        cls.cache = []
+
+    @classmethod
+    def pop(cls) -> None:
+        assert not cls.funcs_queue.empty(), "Pop empty queue."
+        funcs = cls.funcs_queue.get()
+        for func in funcs:
+            func()
+
+    @classmethod
+    def clear(cls) -> None:
+        cls.cache = []
+        cls.funcs_queue = queue.Queue()
+
+
+class EventStore:
+    event = None
+
+    @classmethod
+    def set(cls, event) -> None:
+        cls.event = event
+
+
+def fold_init_dims(tensor):
+    # NOTE(zhangyuqin1998): Reshape a rank-3 tensor from P x M x N to (P * M) x N,
+    # to keep weight_grad in a correct rank. See phi::FoldInitDims.
+    if tensor.ndim == 3:
+        tensor = paddle.reshape(tensor, [-1, tensor.shape[-1]])
+    return tensor
+
+
+def grad_weight_fn(input, weight, out_grad, inplace_update_grad=True):
+    if weight.stop_gradient:
+        return
+    with paddle.no_grad():
+        weight_grad = paddle.matmul(
+            x=fold_init_dims(input),
+            y=fold_init_dims(out_grad),
+            transpose_x=True,
+            transpose_y=False,
+        )
+
+        if hasattr(weight, "main_grad"):
+            if weight.main_grad is None:
+                weight.main_grad = paddle.base.framework.core.eager.Tensor(
+                    value=weight_grad.cast(paddle.float32).value(),
+                    place=weight_grad.place,
+                    name="main_grad@" + weight.name,
+                )
+            else:
+                weight.main_grad.add_(weight_grad)
+            weight_grad._clear_data()
+        else:
+            if weight.grad is None:
+                weight.grad = paddle.zeros_like(weight, dtype=weight.dtype)
+            weight.grad = paddle.add(weight.grad, weight_grad)
+
+
+class SplitBWMatmul(PyLayer):
+    @staticmethod
+    def forward(ctx, input, weight, bias):
+        ctx.save_for_backward(input, weight, bias)
+        out = F.linear(x=input, weight=weight, bias=bias)
+        return out
+
+    @staticmethod
+    def backward(ctx, out_grad):
+        input, weight, bias = ctx.saved_tensor()
+
+        if WeightGradStore.enabled:
+            WeightGradStore.put(
+                partial(grad_weight_fn, input, weight, out_grad)
+            )
+        else:
+            grad_weight_fn(input, weight, out_grad)
+
+        input_grad = None
+        if not input.stop_gradient:
+            input_grad = paddle.matmul(
+                x=out_grad, y=weight, transpose_x=False, transpose_y=True
+            )
+        if bias is not None:
+            bias_grad = None
+            if not bias.stop_gradient:
+                bias_grad = paddle.sum(fold_init_dims(out_grad), axis=0)
+            return input_grad, None, bias_grad
+        else:
+            return input_grad, None
+
+
+class SplitBWLinear(nn.Linear):
+    def forward(self, input):
+        return SplitBWMatmul.apply(input, self.weight, bias=self.bias)

python/paddle/distributed/fleet/metrics/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2020 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 .metric import acc, auc, mae, max, min, mse, rmse, sum  # noqa: F401
+
+__all__ = []

python/paddle/distributed/fleet/metrics/metric.py

@@ -0,0 +1,446 @@
+#   Copyright (c) 2020 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.
+"""Fleet Metrics"""
+
+import math
+
+import numpy as np
+
+import paddle
+from paddle.common_ops_import import Variable
+
+__all__ = []
+
+
+def sum(input, scope=None, util=None):
+    """
+    distributed sum in fleet
+
+    Args:
+        input(numpy.array|Variable|string): output of a layer
+        scope(Scope): specific scope
+
+    Returns:
+        global_metric(numpy.array): sum array
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> input = paddle.cast(some_input, dtype='float32')
+            >>> cnt = paddle.sum(input)
+            >>> global_cnt = paddle.static.create_global_var(persistable=True, dtype='float32', shape=[], value=0)
+            >>> tmp = paddle.add(cnt, global_cnt)
+            >>> paddle.assign(tmp, global_cnt)
+
+            >>> # in train.py, after train or infer
+            >>> res = np.array(scope.find_var(global_cnt.name).get_tensor())
+            >>> print("sum array: ", paddle.distributed.fleet.sum(res))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+    if isinstance(input, Variable):
+        input = np.array(scope.find_var(input.name).get_tensor())
+    elif isinstance(input, str):
+        input = np.array(scope.find_var(input).get_tensor())
+    old_shape = np.array(input.shape)
+    output = np.copy(input) * 0
+    output = util.all_reduce(input, "sum")
+    output = output.reshape(old_shape)
+    return output
+
+
+def max(input, scope=None, util=None):
+    """
+    distributed max in fleet
+
+    Args:
+        input(numpy.array|Variable|string): output of a layer
+        scope(Scope): specific scope
+
+    Returns:
+        global_metric(numpy.array): max array
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> input = paddle.cast(some_input, dtype='float32')
+            >>> cnt = paddle.sum(input)
+            >>> global_cnt = paddle.static.create_global_var(persistable=True, dtype='float32', shape=[], value=0)
+            >>> tmp = paddle.maximum(cnt, global_cnt)
+            >>> paddle.assign(tmp, global_cnt)
+
+            >>> # in train.py, after train or infer
+            >>> res = np.array(scope.find_var(global_cnt.name).get_tensor())
+            >>> print("max array: ", paddle.distributed.fleet.max(res))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+    if isinstance(input, Variable):
+        input = np.array(scope.find_var(input.name).get_tensor())
+    elif isinstance(input, str):
+        input = np.array(scope.find_var(input).get_tensor())
+    old_shape = np.array(input.shape)
+    output = np.copy(input) * 0
+    output = util.all_reduce(input, "max")
+    output = output.reshape(old_shape)
+    return output
+
+
+def min(input, scope=None, util=None):
+    """
+    distributed min in fleet
+
+    Args:
+        input(numpy.array|Variable|string): output of a layer
+        scope(Scope): specific scope
+
+    Returns:
+        global_metric(numpy.array): min array
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> input = paddle.cast(some_input, dtype='float32')
+            >>> cnt = paddle.sum(input)
+            >>> global_cnt = paddle.static.create_global_var(persistable=True, dtype='float32', shape=[], value=0)
+            >>> tmp = paddle.minimum(cnt, global_cnt)
+            >>> paddle.assign(tmp, global_cnt)
+
+            >>> # in train.py, after train or infer
+            >>> res = np.array(scope.find_var(global_cnt.name).get_tensor())
+            >>> print("min array: ", paddle.distributed.fleet.min(res))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+    if isinstance(input, Variable):
+        input = np.array(scope.find_var(input.name).get_tensor())
+    elif isinstance(input, str):
+        input = np.array(scope.find_var(input).get_tensor())
+    old_shape = np.array(input.shape)
+    output = np.copy(input) * 0
+    output = util.all_reduce(input, "min")
+    output = output.reshape(old_shape)
+    return output
+
+
+def auc(stat_pos, stat_neg, scope=None, util=None):
+    """
+    distributed auc in fleet
+
+    Args:
+        stat_pos(numpy.array|Variable|string): stat_pos in output of paddle.static.auc
+        stat_neg(numpy.array|Variable|string): stat_neg in output of paddle.static.auc
+        scope(Scope): specific scope
+
+    Returns:
+        auc_value(float): auc value
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> similarity_norm = paddle.nn.functional.sigmoid(paddle.clip(output, min=-15.0, max=15.0))
+            >>> binary_predict = paddle.concat(
+            ...     input=[paddle.subtract(paddle.ceil(similarity_norm), similarity_norm), similarity_norm], axis=1
+            ... )
+            >>> self.auc, batch_auc, [batch_stat_pos, batch_stat_neg, stat_pos, stat_neg] =
+            ...     paddle.static.auc(input=binary_predict, label=label, curve='ROC', num_thresholds=4096)
+
+            >>> # in train.py, after train or infer
+            >>> pos = np.array(scope.find_var(stat_pos.name).get_tensor())
+            >>> neg = np.array(scope.find_var(stat_neg.name).get_tensor())
+            >>> print("auc: ", paddle.distributed.fleet.auc(pos, neg))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+
+    if isinstance(stat_pos, Variable):
+        stat_pos = np.array(scope.find_var(stat_pos.name).get_tensor())
+    elif isinstance(stat_pos, str):
+        stat_pos = np.array(scope.find_var(stat_pos).get_tensor())
+    if isinstance(stat_neg, Variable):
+        stat_neg = np.array(scope.find_var(stat_neg.name).get_tensor())
+    elif isinstance(stat_neg, str):
+        stat_neg = np.array(scope.find_var(stat_neg).get_tensor())
+    # auc pos bucket shape
+    old_pos_shape = np.array(stat_pos.shape)
+    # reshape to one dim
+    stat_pos = stat_pos.reshape(-1)
+    global_pos = np.copy(stat_pos) * 0
+    # mpi allreduce
+    global_pos = util.all_reduce(stat_pos, "sum")
+    global_pos = global_pos.reshape(old_pos_shape)
+
+    # auc neg bucket
+    old_neg_shape = np.array(stat_neg.shape)
+    stat_neg = stat_neg.reshape(-1)
+    global_neg = np.copy(stat_neg) * 0
+    global_neg = util.all_reduce(stat_neg, "sum")
+    global_neg = global_neg.reshape(old_neg_shape)
+
+    # calculate auc
+    num_bucket = len(global_pos[0])
+    area = 0.0
+    pos = 0.0
+    neg = 0.0
+    new_pos = 0.0
+    new_neg = 0.0
+    total_ins_num = 0
+    for i in range(num_bucket):
+        index = num_bucket - 1 - i
+        new_pos = pos + global_pos[0][index]
+        total_ins_num += global_pos[0][index]
+        new_neg = neg + global_neg[0][index]
+        total_ins_num += global_neg[0][index]
+        area += (new_neg - neg) * (pos + new_pos) / 2
+        pos = new_pos
+        neg = new_neg
+
+    auc_value = None
+    if pos * neg == 0 or total_ins_num == 0:
+        auc_value = 0.5
+    else:
+        auc_value = area / (pos * neg)
+
+    return auc_value
+
+
+def mae(abserr, total_ins_num, scope=None, util=None):
+    """
+    distributed mae in fleet
+
+    Args:
+        abserr(numpy.array|Variable|string): abserr in output of paddle.static.ctr_metric_bundle
+        total_ins_num(numpy.array|Variable|string): total variable
+        scope(Scope): specific scope
+
+    Returns:
+        mae(float): mae value
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> sqrerr, abserr, prob, q, pos, total = paddle.static.ctr_metric_bundle(
+            ...     similarity_norm, paddle.cast(x=label, dtype='float32')
+            ... )
+
+            >>> # in train.py, after train or infer
+            >>> res = np.array(scope.find_var(abserr.name).get_tensor())
+            >>> print("mae: ", paddle.distributed.fleet.mae(res, total_ins_num))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+
+    if isinstance(abserr, Variable):
+        abserr = np.array(scope.find_var(abserr.name).get_tensor())
+    elif isinstance(abserr, str):
+        abserr = np.array(scope.find_var(abserr).get_tensor())
+    if isinstance(total_ins_num, Variable):
+        total_ins_num = np.array(
+            scope.find_var(total_ins_num.name).get_tensor()
+        )
+    elif isinstance(total_ins_num, str):
+        total_ins_num = np.array(scope.find_var(total_ins_num).get_tensor())
+
+    old_metric_shape = np.array(abserr.shape)
+    abserr = abserr.reshape(-1)
+    global_metric = np.copy(abserr) * 0
+
+    global_metric = util.all_reduce(abserr, "sum")
+    global_metric = global_metric.reshape(old_metric_shape)
+    global_total_num = util.all_reduce(total_ins_num, "sum")
+
+    mae_value = float(global_metric[0]) / float(global_total_num[0])
+    return mae_value
+
+
+def rmse(sqrerr, total_ins_num, scope=None, util=None):
+    """
+    distributed rmse in fleet
+
+    Args:
+        sqrerr(numpy.array|Variable|string): sqrerr in output of paddle.static.ctr_metric_bundle
+        total_ins_num(numpy.array|Variable|string): total variable
+        scope(Scope): specific scope
+
+    Returns:
+        rmse(float): rmse value
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> sqrerr, abserr, prob, q, pos, total = paddle.static.ctr_metric_bundle(
+            ...     similarity_norm, paddle.cast(x=label, dtype='float32')
+            ... )
+
+            >>> # in train.py, after train or infer
+            >>> res = np.array(scope.find_var(sqrerr.name).get_tensor())
+            >>> print("rmse: ", paddle.distributed.fleet.rmse(res, total_ins_num))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+
+    if isinstance(sqrerr, Variable):
+        sqrerr = np.array(scope.find_var(sqrerr.name).get_tensor())
+    elif isinstance(sqrerr, str):
+        sqrerr = np.array(scope.find_var(sqrerr).get_tensor())
+    if isinstance(total_ins_num, Variable):
+        total_ins_num = np.array(
+            scope.find_var(total_ins_num.name).get_tensor()
+        )
+    elif isinstance(total_ins_num, str):
+        total_ins_num = np.array(scope.find_var(total_ins_num).get_tensor())
+    old_metric_shape = np.array(sqrerr.shape)
+    sqrerr = sqrerr.reshape(-1)
+    global_metric = np.copy(sqrerr) * 0
+
+    global_metric = util.all_reduce(sqrerr, "sum")
+    global_metric = global_metric.reshape(old_metric_shape)
+    global_total_num = util.all_reduce(total_ins_num, "sum")
+
+    rmse_value = math.sqrt(float(global_metric[0]) / float(global_total_num[0]))
+
+    return rmse_value
+
+
+def mse(sqrerr, total_ins_num, scope=None, util=None):
+    """
+    distributed mse in fleet
+
+    Args:
+        sqrerr(numpy.array|Variable|string): sqrerr in output of paddle.static.ctr_metric_bundle
+        total_ins_num(numpy.array|Variable|string): total variable
+        scope(Scope): specific scope
+
+    Returns:
+        mse(float): mse value
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> sqrerr, abserr, prob, q, pos, total = paddle.static.ctr_metric_bundle(
+            ...     similarity_norm, paddle.cast(x=label, dtype='float32')
+            ... )
+
+            >>> # in train.py, after train or infer
+            >>> metric = np.array(scope.find_var(sqrerr.name).get_tensor())
+            >>> print("mse: ", paddle.distributed.fleet.mse(metric, total_ins_num))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+
+    if isinstance(sqrerr, Variable):
+        sqrerr = np.array(scope.find_var(sqrerr.name).get_tensor())
+    elif isinstance(sqrerr, str):
+        sqrerr = np.array(scope.find_var(sqrerr).get_tensor())
+    if isinstance(total_ins_num, Variable):
+        total_ins_num = np.array(
+            scope.find_var(total_ins_num.name).get_tensor()
+        )
+    elif isinstance(total_ins_num, str):
+        total_ins_num = np.array(scope.find_var(total_ins_num).get_tensor())
+    old_metric_shape = np.array(sqrerr.shape)
+    sqrerr = sqrerr.reshape(-1)
+    global_metric = np.copy(sqrerr) * 0
+
+    global_metric = util.all_reduce(sqrerr, "sum")
+    global_metric = global_metric.reshape(old_metric_shape)
+    global_total_num = util.all_reduce(total_ins_num, "sum")
+
+    mse_value = float(global_metric[0]) / float(global_total_num[0])
+    return mse_value
+
+
+def acc(correct, total, scope=None, util=None):
+    """
+    distributed accuracy in fleet
+
+    Args:
+        correct(numpy.array|Variable|string): correct Variable
+        total(numpy.array|Variable): total Variable
+        scope(Scope): specific scope
+
+    Returns:
+        acc(float): accuracy value
+
+    Example:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # in model.py
+            >>> correct = paddle.static.create_global_var(dtype='float32', shape=[1], value=0)
+            >>> total = paddle.static.create_global_var(dtype='float32', shape=[1], value=0)
+            >>> acc = paddle.metric.accuracy(predict, label, k=1, correct=correct, total=total)
+
+            >>> global_correct = paddle.static.create_global_var(persistable=True, dtype='float32', shape=[1], value=0)
+            >>> tmp1 = paddle.minimum(correct, global_correct)
+            >>> paddle.assign(tmp1, global_correct)
+
+            >>> global_total = paddle.static.create_global_var(persistable=True, dtype='float32', shape=[1], value=0)
+            >>> tmp2 = paddle.minimum(total, global_total)
+            >>> paddle.assign(tmp2, global_total)
+
+            >>> # in train.py, after train or infer
+            >>> correct_num = np.array(scope.find_var(correct.name).get_tensor())
+            >>> total_num = np.array(scope.find_var(total.name).get_tensor())
+            >>> print("accuracy: ", paddle.distributed.fleet.acc(correct_num, total_num))
+    """
+    if scope is None:
+        scope = paddle.static.global_scope()
+    if util is None:
+        util = paddle.distributed.fleet.util
+
+    if isinstance(correct, Variable):
+        correct = np.array(scope.find_var(correct.name).get_tensor())
+    elif isinstance(correct, str):
+        correct = np.array(scope.find_var(correct).get_tensor())
+    if isinstance(total, Variable):
+        total = np.array(scope.find_var(total.name).get_tensor())
+    elif isinstance(total, str):
+        total = np.array(scope.find_var(total).get_tensor())
+
+    global_correct_num = np.copy(correct) * 0
+    global_total_num = np.copy(total) * 0
+
+    global_correct_num = util.all_reduce(correct, "sum")
+    global_total_num = util.all_reduce(total, "sum")
+
+    return float(global_correct_num[0]) / float(global_total_num[0])

python/paddle/distributed/fleet/model.py

@@ -0,0 +1,217 @@
+#   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 paddle
+from paddle.distributed import fleet
+
+from .base.topology import ParallelMode
+from .meta_parallel import (
+    DualPipeVParallel,
+    NoPipelineParallel,
+    PipelineLayer,
+    PipelineParallel,
+    PipelineParallelWithInterleave,
+    PipelineParallelWithInterleaveFthenB,
+    SegmentParallel,
+    ShardingParallel,
+    TensorParallel,
+    VPPFhenBInBalancedMemory,
+)
+
+_grad_scalar = None
+
+
+def distributed_model(model):
+    """
+    Return distributed data parallel model (Only work in dygraph mode)
+
+    Args:
+        model (Layer): the user-defined model which inherits Layer.
+
+    Returns:
+        distributed data parallel model which inherits Layer.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.nn as nn
+            >>> from paddle.distributed import fleet
+
+            >>> class LinearNet(nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self._linear1 = nn.Linear(10, 10)
+            ...         self._linear2 = nn.Linear(10, 1)
+            ...
+            ...     def forward(self, x):
+            ...         return self._linear2(self._linear1(x))
+
+            >>> # 1. initialize fleet environment
+            >>> fleet.init(is_collective=True)
+
+            >>> # 2. create layer & optimizer
+            >>> layer = LinearNet()
+            >>> loss_fn = nn.MSELoss()
+            >>> adam = paddle.optimizer.Adam(
+            ...     learning_rate=0.001,
+            ...     parameters=layer.parameters(),
+            ... )
+
+            >>> # 3. get data_parallel model using fleet
+            >>> adam = fleet.distributed_optimizer(adam)
+            >>> dp_layer = fleet.distributed_model(layer)
+
+            >>> # 4. run layer
+            >>> inputs = paddle.randn([10, 10], 'float32')
+            >>> outputs = dp_layer(inputs)
+            >>> labels = paddle.randn([10, 1], 'float32')
+            >>> loss = loss_fn(outputs, labels)
+            >>> print("loss:", loss.numpy())
+            >>> loss.backward()
+            >>> adam.step()
+            >>> adam.clear_grad()
+
+
+    """
+    fleet_env = fleet.fleet
+    strategy = fleet_env._user_defined_strategy
+    assert model is not None, "model should not be None"
+    if paddle.distributed.get_world_size() <= 1:
+        model = NoPipelineParallel(model, strategy=strategy)
+        return model
+
+    if strategy.amp:
+        level = (
+            "O2"
+            if strategy.amp_configs['use_pure_fp16']
+            or strategy.amp_configs['use_pure_bf16']
+            else "O1"
+        )
+
+        if level == "O2":
+            model = paddle.amp.decorate(
+                models=model,
+                optimizers=None,
+                level="O2",
+                master_weight=None,
+                save_dtype=None,
+                dtype=(
+                    "float16"
+                    if strategy.amp_configs['use_pure_fp16']
+                    else "bfloat16"
+                ),
+            )
+
+        init_loss_scaling = strategy.amp_configs['init_loss_scaling']
+        incr_ratio = strategy.amp_configs['incr_ratio']
+        decr_ratio = strategy.amp_configs['decr_ratio']
+        incr_every_n_steps = strategy.amp_configs['incr_every_n_steps']
+        decr_every_n_nan_or_inf = strategy.amp_configs[
+            'decr_every_n_nan_or_inf'
+        ]
+        use_dynamic_loss_scaling = strategy.amp_configs[
+            'use_dynamic_loss_scaling'
+        ]
+
+        global _grad_scalar
+        _grad_scalar = paddle.amp.GradScaler(
+            init_loss_scaling=init_loss_scaling,
+            incr_ratio=incr_ratio,
+            decr_ratio=decr_ratio,
+            incr_every_n_steps=incr_every_n_steps,
+            decr_every_n_nan_or_inf=decr_every_n_nan_or_inf,
+            use_dynamic_loss_scaling=use_dynamic_loss_scaling,
+        )
+
+    if fleet_env._hcg.get_parallel_mode() == ParallelMode.PIPELINE_PARALLEL:
+        assert isinstance(model, PipelineLayer), (
+            "For pipeline parallel, the model should an instance of PipelineLayer"
+        )
+        if strategy.hybrid_configs["pp_configs"].use_dualpipev:
+            model = DualPipeVParallel(model, fleet_env._hcg, strategy=strategy)
+        elif model.get_num_virtual_stages() == 1:
+            # 1f1b pipeline
+            model = PipelineParallel(model, fleet_env._hcg, strategy=strategy)
+        else:
+            accumulate_steps = strategy.pipeline_configs['accumulate_steps']
+            pp_degree = fleet_env._hcg.get_pipe_parallel_world_size()
+            if accumulate_steps >= 2 * pp_degree:
+                # interleave pipeline
+                model = PipelineParallelWithInterleave(
+                    model, fleet_env._hcg, strategy=strategy
+                )
+            elif pp_degree <= accumulate_steps < 2 * pp_degree:
+                if strategy.hybrid_configs[
+                    "pp_configs"
+                ].best_unbalanced_scheduler:
+                    model = VPPFhenBInBalancedMemory(
+                        model, fleet_env._hcg, strategy=strategy
+                    )
+                else:
+                    model = PipelineParallelWithInterleaveFthenB(
+                        model, fleet_env._hcg, strategy=strategy
+                    )
+            else:
+                raise ValueError(
+                    f"The accumulate_steps({accumulate_steps}) should be greater than or equal to pp_degree({pp_degree})"
+                )
+    else:
+        if isinstance(model, PipelineLayer):
+            # PaddleFleet Model
+            model = NoPipelineParallel(
+                model, strategy=strategy, hcg=fleet_env._hcg
+            )
+        else:
+            if strategy.heter_ccl_mode:
+                distributed_model = paddle.DataParallel(
+                    model,
+                    comm_buffer_size=strategy.fuse_grad_size_in_MB,
+                    last_comm_buffer_size=strategy.last_comm_group_size_MB,
+                    find_unused_parameters=strategy.find_unused_parameters,
+                )
+                return distributed_model
+
+            if (
+                fleet_env._hcg.get_parallel_mode()
+                == ParallelMode.SHARDING_PARALLEL
+            ):
+                model = ShardingParallel(
+                    model, fleet_env._hcg, strategy=strategy
+                )
+            elif (
+                fleet_env._hcg.get_parallel_mode() == ParallelMode.DATA_PARALLEL
+            ):
+                model = paddle.DataParallel(
+                    model,
+                    comm_buffer_size=strategy.fuse_grad_size_in_MB,
+                    last_comm_buffer_size=strategy.last_comm_group_size_MB,
+                    find_unused_parameters=strategy.find_unused_parameters,
+                    group=fleet_env._hcg.get_data_parallel_group(),
+                )
+            elif (
+                fleet_env._hcg.get_parallel_mode()
+                == ParallelMode.SEGMENT_PARALLEL
+            ):
+                model = SegmentParallel(
+                    model, fleet_env._hcg, strategy=strategy
+                )
+            elif (
+                fleet_env._hcg.get_parallel_mode()
+                == ParallelMode.TENSOR_PARALLEL
+            ):
+                model = TensorParallel(model, fleet_env._hcg, strategy=strategy)
+
+    return model

python/paddle/distributed/fleet/optimizer.py

@@ -0,0 +1,100 @@
+#   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.distributed import fleet
+from paddle.framework import in_dynamic_mode
+
+from .meta_optimizers import HeterParallelOptimizer, HybridParallelOptimizer
+from .utils.log_util import logger
+
+
+def _dygraph_distributed_optimizer(optimizer, strategy=None):
+    """
+    Optimizer for distributed training.
+    For the distributed training, this method would rebuild a new instance of DistributedOptimizer.
+    Which has basic Optimizer function and special features for distributed training.
+    Args:
+        optimizer(Optimizer): The executor to run for init server.
+        strategy(DistributedStrategy): Extra properties for distributed optimizer.
+            It is recommended to use DistributedStrategy in fleet.init(). The strategy
+            here is for compatibility. If the strategy in fleet.distributed_optimizer()
+            is not None, then it will overwrite the DistributedStrategy in fleet.init(),
+            which will take effect in distributed training.
+    Returns:
+        Fleet: instance of fleet.
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed.fleet as fleet
+            >>> fleet.init(is_collective=True)
+            >>> strategy = fleet.DistributedStrategy()
+            >>> linear = paddle.nn.Linear(10, 10)
+            >>> optimizer = paddle.optimizer.SGD(learning_rate=0.001, parameters=linear.parameters())
+            >>> optimizer = fleet.distributed_optimizer(optimizer, strategy=strategy)
+
+    """
+    fleet_env = fleet.fleet
+    fleet_env.user_defined_optimizer = optimizer
+
+    if strategy is not None:
+        if fleet_env._is_collective:
+            logger.warning(
+                "It is recommended to use DistributedStrategy "
+                "in fleet_env.init(). The strategy here is only for compatibility. "
+                "If the strategy in fleet_env.distributed_optimizer() is "
+                "not None, then it will overwrite the DistributedStrategy in fleet_env.init(), "
+                "which will take effect in distributed training."
+            )
+        fleet_env._user_defined_strategy = copy.deepcopy(strategy)
+
+    fleet_env._context = {}
+
+    if fleet_env.worker_num() > 1:
+        if not fleet_env._user_defined_strategy.heter_ccl_mode:
+            hp_optim = HybridParallelOptimizer(
+                optimizer, fleet_env._hcg, fleet_env._user_defined_strategy
+            )
+
+            if fleet_env._user_defined_strategy.hybrid_configs[
+                "pp_configs"
+            ].dp_comm_overlap:
+                # grad all-reduce of dp and sep with be fused
+                hp_optim._dp_enable = False
+                hp_optim._sep_enable = False
+
+            if fleet_env._user_defined_strategy.hybrid_configs[
+                "pp_configs"
+            ].sharding_comm_overlap:
+                hp_optim._sharding_enable = False
+                assert not hp_optim._sep_enable, (
+                    "sep parallel can not coexist with sharding_comm_overlap"
+                )
+
+            return hp_optim
+        else:
+            return HeterParallelOptimizer(
+                optimizer, fleet_env._user_defined_strategy
+            )
+    else:
+        return optimizer
+
+
+def distributed_optimizer(*args, **kwargs):
+    if in_dynamic_mode():
+        return _dygraph_distributed_optimizer(*args, **kwargs)
+    else:
+        return fleet.fleet.distributed_optimizer(*args, **kwargs)

python/paddle/distributed/fleet/recompute/__init__.py

@@ -0,0 +1,23 @@
+# 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 .recompute import (  # noqa: F401
+    custom_state_manager,
+    is_in_recompute,
+    recompute,
+    recompute_sequential,
+)
+from .recompute_hybrid import recompute_hybrid  # noqa: F401
+
+__all__ = []

python/paddle/distributed/fleet/recompute/recompute.py

@@ -0,0 +1,989 @@
+# 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 contextlib
+import copy
+import functools
+import inspect
+import random
+import threading
+import weakref
+from typing import TYPE_CHECKING, Any, TypedDict
+
+import numpy as np
+
+import paddle
+from paddle import framework
+from paddle.autograd import PyLayer
+from paddle.base.framework import EagerParamBase
+from paddle.base.wrapped_decorator import copy_signature
+from paddle.distributed.fleet.meta_parallel.parallel_layers.random import (
+    get_rng_state_tracker,
+)
+from paddle.framework import core, in_dynamic_mode
+from paddle.jit.dy2static.program_translator import StaticFunction
+
+from ..utils.log_util import logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from typing_extensions import NotRequired
+
+    from paddle.nn import Sequential
+
+    class _Ctx(TypedDict):
+        segments: int = 1
+        preserve_rng_state: NotRequired[bool]
+
+
+__all__ = []
+_SIGNATURE_CACHE = weakref.WeakKeyDictionary()
+
+
+class RecomputeContext:
+    """
+    A thread-safe context manager and decorator for tracking whether the current
+    execution is inside a recompute phase.
+
+    RecomputeContext uses a thread-local flag to mark when code is running within a
+    recompute region. It can be used as a context manager (``with`` statement) or as
+    a decorator to automatically set and clear the recompute-active state. This allows
+    downstream code to query ``is_in_recompute()`` and adapt its behavior accordingly
+    (e.g., skipping certain logging or side effects during recomputation).
+
+    Parameters:
+        None.
+
+    Returns:
+        RecomputeContext: A recompute context instance that can be used as a context
+            manager or decorator.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.distributed.fleet.utils import is_in_recompute
+
+            >>> # Usage as a context manager
+            >>> ctx = RecomputeContext()
+            >>> print(ctx.active)
+            False
+            >>> with ctx:
+            ...     print(ctx.active)
+            True
+            >>> print(ctx.active)
+            False
+
+            >>> # Usage as a decorator
+            >>> ctx = RecomputeContext()
+            >>> @ctx
+            ... def my_forward(x):
+            ...     return is_in_recompute()
+            >>> print(my_forward(None))
+            True
+    """
+
+    def __init__(self):
+        self._local = threading.local()
+
+    @property
+    def active(self) -> bool:
+        return getattr(self._local, 'active', False)
+
+    def __enter__(self):
+        self._local.active = True
+        return self
+
+    def __exit__(self, *_exc):
+        self._local.active = False
+        return False
+
+    def __call__(self, fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            with self:
+                return fn(*args, **kwargs)
+
+        copy_signature(fn, wrapper)
+
+        return wrapper
+
+
+_recompute_context = RecomputeContext()
+
+
+def is_in_recompute() -> bool:
+    """
+    Check whether the current thread is executing inside a recompute context.
+
+    This function inspects the global ``_recompute_context`` to determine if the
+    current thread is within an active recompute phase. It is typically used inside
+    forward computations to detect whether the execution is a normal forward pass
+    or a recompute (re-forward) pass triggered during backpropagation, so that
+    certain operations (e.g., logging, random state management) can be skipped or
+    adjusted accordingly.
+
+    Parameters:
+        None.
+
+    Returns:
+        bool: ``True`` if the current thread is inside a recompute context,
+            ``False`` otherwise.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.distributed.fleet.utils import is_in_recompute
+            >>> # Outside any recompute context
+            >>> print(is_in_recompute())
+            False
+
+            >>> from paddle.distributed.fleet.utils.__init__ import RecomputeContext
+            >>> ctx = RecomputeContext()
+            >>> with ctx:
+            ...     print(is_in_recompute())
+            True
+    """
+    return _recompute_context.active
+
+
+def _varbase_help(param):
+    state = copy.deepcopy(param.__dict__)
+    new_param = EagerParamBase(
+        shape=param.shape,
+        dtype=param.dtype,
+        trainable=param.trainable,
+        name=param.name,
+        **state,
+    )
+    param._share_buffer_to(new_param)
+    return new_param
+
+
+def detach_variable(inputs):
+    out = []
+    for inp in inputs:
+        if not isinstance(inp, core.eager.Tensor) and (
+            type(inp) is not tuple or not isinstance(inp[0], core.eager.Tensor)
+        ):
+            # the inp is not a tensor or not a tuple of tensors
+            out.append(inp)
+            continue
+
+        if isinstance(inp, EagerParamBase):
+            out.append(_varbase_help(inp))
+            continue
+
+        if type(inp) is tuple:
+            detach_inp = []
+            for i in inp:
+                # detach all tensors in the tuple
+                assert isinstance(i, core.eager.Tensor)
+
+                if isinstance(i, EagerParamBase):
+                    detach_inp.append(_varbase_help(i))
+                else:
+                    tmp_i = i.detach()
+                    tmp_i.stop_gradient = i.stop_gradient
+                    detach_inp.append(tmp_i)
+
+            out.append(tuple(detach_inp))
+            continue
+
+        x = inp.detach()
+        x.stop_gradient = inp.stop_gradient
+        out.append(x)
+    return tuple(out)
+
+
+def check_recompute_necessary(inputs):
+    necessary_for_each_input = []
+    for input_ in inputs:
+        if isinstance(input_, paddle.Tensor):
+            necessary_for_each_input.append(input_.stop_gradient)
+        elif type(input_) is tuple:
+            for i in input_:
+                # traverse all tensors in the tuple
+                if isinstance(i, paddle.Tensor):
+                    necessary_for_each_input.append(i.stop_gradient)
+    if all(necessary_for_each_input):
+        logger.warning(
+            "[Recompute]: None of the inputs to current recompute block need grad, "
+            "therefore there is NO need to recompute this block in backward !"
+        )
+
+
+def _closure_cell_values(run_function):
+    """Return cell contents of ``run_function``'s ``__closure__`` as a tuple.
+
+    Supports plain functions/lambdas and ``paddle.nn.Layer`` (uses ``forward``).
+    Deep Tensor extraction is done by the C++ side of ``_hold_tensors``.
+    """
+    fn = (
+        run_function.forward
+        if isinstance(run_function, paddle.nn.Layer)
+        else run_function
+    )
+    closure = getattr(fn, '__closure__', None) or ()
+    values = []
+    for cell in closure:
+        try:
+            values.append(cell.cell_contents)
+        except ValueError:  # empty cell
+            pass
+    return tuple(values)
+
+
+class CustomStatesManager:
+    """CustomStatesManager"""
+
+    def __init__(self):
+        """__init__"""
+        self.custom_get_state_func = None
+        self.custom_set_state_func = None
+
+    def set_custom_get_state_func(self, custom_get_state_func):
+        assert_msg = (
+            "The custom_state_manager does not support duplicate settings."
+        )
+        assert self.custom_get_state_func is None, assert_msg
+        self.custom_get_state_func = custom_get_state_func
+
+    def set_custom_set_state_func(self, custom_set_state_func):
+        assert_msg = (
+            "The custom_state_manager does not support duplicate settings."
+        )
+        assert self.custom_set_state_func is None, assert_msg
+        self.custom_set_state_func = custom_set_state_func
+
+
+custom_state_manager = CustomStatesManager()
+
+
+@contextlib.contextmanager
+def switch_rng_state_tracker(
+    rng_state,
+    tracker,
+    numpy_state,
+    random_state,
+    custom_state=None,
+    custom_get_state_func=None,
+    custom_set_state_func=None,
+):
+    orig_rng_state = paddle.get_rng_state()
+    orig_rng_tracker = get_rng_state_tracker().get_states_tracker()
+    paddle.set_rng_state(rng_state)
+    get_rng_state_tracker().set_states_tracker(tracker)
+
+    orig_numpy_state = None
+    orig_random_state = None
+
+    if numpy_state is not None:
+        orig_numpy_state = np.random.get_state()
+        np.random.set_state(numpy_state)
+    if random_state is not None:
+        orig_random_state = random.getstate()
+        random.setstate(random_state)
+
+    if custom_state is not None:
+        assert custom_get_state_func is not None
+        assert custom_set_state_func is not None
+        orig_custom_state = custom_get_state_func()
+        custom_set_state_func(custom_state)
+    try:
+        yield
+    finally:
+        paddle.set_rng_state(orig_rng_state)
+        get_rng_state_tracker().set_states_tracker(orig_rng_tracker)
+        if orig_numpy_state is not None:
+            np.random.set_state(orig_numpy_state)
+        if orig_random_state is not None:
+            random.setstate(orig_random_state)
+
+        if custom_state is not None:
+            custom_set_state_func(orig_custom_state)
+
+
+class RecomputeFunction(PyLayer):
+    @staticmethod
+    def forward(
+        ctx,
+        run_function,
+        preserve_rng_state,
+        preserve_external_rng_state,
+        offload_indices,
+        custom_get_state_func,
+        custom_set_state_func,
+        *args,
+        **kwargs,
+    ):
+        # store for recomputing
+        ctx.run_function = run_function
+        ctx.preserve_rng_state = preserve_rng_state
+        ctx.preserve_external_rng_state = preserve_external_rng_state
+        ctx.offload_indices = offload_indices
+        ctx.kwargs = kwargs
+
+        # NOTE the number of outputs of backward() should be equal to the number of tensors in forward()'s input
+        # the order of tensors in backward()'s output should be the same as tensors in forward()'s input
+        # None tensor inputs will be filtered in backward inputs.
+
+        # NOTE recompute with restore RNG only support one scenario where one process for one cuda gpu.
+        # one process with multiple gpu and mix-gpu-cpu scenarios are not support
+        if ctx.preserve_rng_state:
+            ctx.fw_rng_state = paddle.get_rng_state()
+            ctx.fwd_rng_state_tracker = (
+                get_rng_state_tracker().get_states_tracker()
+            )
+            if ctx.preserve_external_rng_state:
+                ctx.fwd_numpy_state = np.random.get_state()
+                ctx.fwd_random_state = random.getstate()
+            else:
+                ctx.fwd_numpy_state = None
+                ctx.fwd_random_state = None
+            ctx.fwd_custom_state = custom_get_state_func()
+            ctx.custom_get_state_func = custom_get_state_func
+            ctx.custom_set_state_func = custom_set_state_func
+
+        # TODO support AMP
+        tracer = framework._dygraph_tracer()
+        ctx.is_fw_autocast = (
+            False if tracer._amp_level == core.AmpLevel.O0 else True
+        )
+        if tracer._amp_level == core.AmpLevel.O2:
+            ctx.amp_level = 'O2'
+        elif tracer._amp_level in (core.AmpLevel.O1, core.AmpLevel.O0):
+            ctx.amp_level = 'O1'
+        else:
+            raise ValueError(f"unsupported amp level: {tracer._amp_level}")
+
+        if tracer._amp_dtype == 'float16':
+            ctx.amp_dtype = 'float16'
+        elif tracer._amp_dtype in ('bfloat16', 'float32'):
+            ctx.amp_dtype = 'bfloat16'
+        else:
+            raise ValueError(f"unsupported amp dtype: {tracer._amp_dtype}")
+
+        ctx.amp_white_list, ctx.amp_black_list = tracer._get_amp_op_list()
+
+        with paddle.no_grad(), _recompute_context:
+            outputs = run_function(*args, **kwargs)
+
+        # save input for backward
+        ctx.inputs = []
+        ctx.tensor_indices = []
+        ctx.duplicate_tensor = [False for _ in range(len(args))]
+        tensor_inputs = []
+        for i, arg in enumerate(args):
+            if paddle.is_tensor(arg):
+                if i in ctx.offload_indices:
+                    cpu_arg = (
+                        arg.pin_memory()
+                        if core.is_compiled_with_cuda()
+                        else arg.cpu()
+                    )
+                    cpu_arg._share_buffer_to(arg)
+                tensor_inputs.append(arg)
+                ctx.tensor_indices.append(i)
+                ctx.inputs.append(None)
+            elif type(arg) is tuple:
+                assert i not in ctx.offload_indices, (
+                    f"offload_indices should not contain tensor tuple in position{i}"
+                )
+                is_tensors = [paddle.is_tensor(a) for a in arg]
+                if all(is_tensors):
+                    # the tuple is a tuple of tensors
+                    tensors_stop_gradient = [a.stop_gradient for a in arg]
+                    if not all(tensors_stop_gradient) and any(
+                        tensors_stop_gradient
+                    ):
+                        # tensors in the tuple have different stop_gradient value, which pylayer doesn't support
+                        raise ValueError(
+                            "Recompute receive a tuple containing tensor holds different stop gradient."
+                        )
+                    tensor_inputs.append(arg)
+                    ctx.tensor_indices.append(i)
+                    # Mark the tuple is a tuple of tensors
+                    ctx.duplicate_tensor[i] = True
+                    ctx.inputs.append(None)
+                elif any(is_tensors):
+                    # the tuple contains tensors and non-tensor values
+                    raise ValueError(
+                        "Recompute receive a tuple containing tensor and non-tensor at same time."
+                    )
+                else:
+                    ctx.inputs.append(arg)
+            else:
+                ctx.inputs.append(arg)
+
+        ctx.save_for_backward(*tensor_inputs)
+
+        # Protect tensors captured in run_function's Python __closure__ against
+        # pipeline-parallel _clear_dataptr(); explicit tensor args are already
+        # covered by save_for_backward's tensor_hold_helper.
+        closure_values = _closure_cell_values(run_function)
+        ctx._has_held_tensors = bool(closure_values)
+        if closure_values:
+            ctx._hold_tensors(closure_values)
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *args):
+        with paddle.base.dygraph.guard():
+            # TODO need to check the recompute calling is valid or not
+
+            # Restore closure-captured tensors potentially emptied by
+            # pipeline-parallel _clear_dataptr() before re-running forward.
+            if getattr(ctx, '_has_held_tensors', False):
+                ctx._restore_held_tensors()
+
+            # Restore inputs
+            inputs = list(ctx.inputs)
+            tensor_indices = ctx.tensor_indices
+            duplicate_tensor = ctx.duplicate_tensor
+            tensors = ctx.saved_tensor()
+            for i, idx in enumerate(tensor_indices):
+                inputs[idx] = (
+                    tensors[i].to(
+                        paddle.base.framework._current_expected_place()
+                    )
+                    if i in ctx.offload_indices
+                    else tensors[i]
+                )
+                if i in ctx.offload_indices:
+                    # NOTE(zhiqiu): tensor.to(device) will set stop_gradient=True, which may break the gragh
+                    inputs[idx].stop_gradient = tensors[i].stop_gradient
+            # paddle.enable_grad()
+            tracer = framework._dygraph_tracer()
+            tracer._has_grad = True
+
+            # NOTE support AMP
+            # need restore auto_cast state as well as w/b list
+            if ctx.preserve_rng_state:
+                with (
+                    switch_rng_state_tracker(
+                        ctx.fw_rng_state,
+                        ctx.fwd_rng_state_tracker,
+                        ctx.fwd_numpy_state,
+                        ctx.fwd_random_state,
+                        ctx.fwd_custom_state,
+                        ctx.custom_get_state_func,
+                        ctx.custom_set_state_func,
+                    ),
+                    paddle.amp.auto_cast(
+                        enable=ctx.is_fw_autocast,
+                        custom_white_list=ctx.amp_white_list,
+                        custom_black_list=ctx.amp_black_list,
+                        level=ctx.amp_level,
+                        dtype=ctx.amp_dtype,
+                    ),
+                    _recompute_context,
+                ):
+                    detached_inputs = detach_variable(tuple(inputs))
+                    outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
+            else:
+                with (
+                    paddle.amp.auto_cast(
+                        enable=ctx.is_fw_autocast,
+                        custom_white_list=ctx.amp_white_list,
+                        custom_black_list=ctx.amp_black_list,
+                        level=ctx.amp_level,
+                        dtype=ctx.amp_dtype,
+                    ),
+                    _recompute_context,
+                ):
+                    detached_inputs = detach_variable(tuple(inputs))
+                    outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
+
+            if isinstance(outputs, core.eager.Tensor):
+                outputs = (outputs,)
+            assert len(outputs) == len(args)
+
+            # run backward() with only tensor that requires grad
+            forward_outputs_with_grad = []
+            # NOTE In Transformer-like network, if user put the attention mask into the recompute segment output,
+            # pylayer will force the stop_gradient of attention mask to be False, which will make the number of
+            # tensor that need grad does not match.
+            # the following backward_inputs_with_grad is used to avoid this case.
+            backward_inputs_with_grad = []
+            for i in range(len(outputs)):
+                if (
+                    isinstance(outputs[i], core.eager.Tensor)
+                    and not outputs[i].stop_gradient
+                ):
+                    forward_outputs_with_grad.append(outputs[i])
+                    backward_inputs_with_grad.append(args[i])
+
+            if len(forward_outputs_with_grad) == 0:
+                raise RuntimeError(
+                    "none of output has requires_grad=True, this recompute() is not necessary"
+                )
+
+            # actually backward
+            with paddle.amp.auto_cast(enable=False):
+                paddle.autograd.backward(
+                    forward_outputs_with_grad, backward_inputs_with_grad
+                )
+
+            grads = []
+            for idx, inp in enumerate(detached_inputs):
+                if isinstance(inp, core.eager.Tensor):
+                    grads.append(inp._grad_ivar())
+                elif type(inp) is tuple and duplicate_tensor[idx]:
+                    # input is a tuple and is a tuple of tensors
+                    if all(i.stop_gradient for i in inp):
+                        # all tensors in the tuple doesn't need grad, only return a None for the whole tuple
+                        grads.append(None)
+                    else:
+                        # all tensors in the tuple need grad, should return a tuple of grads
+                        grads.append(tuple(i._grad_ivar() for i in inp))
+
+            if in_dynamic_mode():
+                grads = tuple(grads)
+            else:
+                grads = list(grads)
+            return grads
+
+
+def _recompute_without_reentrant(
+    function,
+    custom_get_state_func,
+    custom_set_state_func,
+    preserve_rng_state=True,
+    preserve_external_rng_state=True,
+    *args,
+    **kwargs,
+):
+    """
+    recompute without reentrant, that means use hook to implement the recompute function rather than re-entrant autograd.
+    """
+
+    if preserve_rng_state:
+        cur_device = paddle.get_device()
+        if cur_device.startswith('gpu:'):
+            fw_cuda_rng_state = paddle.get_cuda_rng_state()
+        elif 'cpu' in cur_device:
+            fw_cuda_rng_state = paddle.get_rng_state()
+        elif 'xpu:' in cur_device:
+            fw_cuda_rng_state = paddle.get_rng_state()
+        elif (
+            cur_device.split(':')[0]
+            in paddle.device.get_all_custom_device_type()
+        ):
+            fw_cuda_rng_state = paddle.get_rng_state(cur_device)
+        else:
+            raise RuntimeError(
+                f"Recompute with RNG preserve is not support current device: {cur_device}."
+            )
+        fwd_cuda_rng_state_tracker = (
+            get_rng_state_tracker().get_states_tracker()
+        )
+        if preserve_external_rng_state:
+            fwd_numpy_state = np.random.get_state()
+            fwd_random_state = random.getstate()
+        else:
+            fwd_numpy_state = None
+            fwd_random_state = None
+        fwd_custom_state = custom_get_state_func()
+
+    tracer = framework._dygraph_tracer()
+    is_fw_autocast = False if tracer._amp_level == core.AmpLevel.O0 else True
+    if tracer._amp_level == core.AmpLevel.O2:
+        amp_level = 'O2'
+    elif tracer._amp_level in (core.AmpLevel.O1, core.AmpLevel.O0):
+        amp_level = 'O1'
+
+    if tracer._amp_dtype == 'float16':
+        amp_dtype = 'float16'
+    elif tracer._amp_dtype in ('bfloat16', 'float32'):
+        amp_dtype = 'bfloat16'
+
+    amp_white_list, amp_black_list = tracer._get_amp_op_list()
+
+    class Intermediate_Holder:
+        pass
+
+    storage = weakref.WeakKeyDictionary()
+    holder_list = []
+
+    def pack(x):
+        res = Intermediate_Holder()
+        holder_list.append(weakref.ref(res))
+        return res
+
+    def unpack(x):
+        unpack_counter = 0
+        if len(storage) == 0:
+
+            def inner_pack(inner_x):
+                nonlocal unpack_counter
+                unpack_counter += 1
+
+                if holder_list[unpack_counter - 1]() is None:
+                    return
+                if inner_x is None:
+                    storage[holder_list[unpack_counter - 1]()] = None
+                    return
+                if hasattr(inner_x, "main_grad") or inner_x.grad is not None:
+                    storage[holder_list[unpack_counter - 1]()] = inner_x
+                else:
+                    if inner_x.is_dist():
+                        tmp_tensor = core.eager.Tensor(inner_x)
+                    else:
+                        tmp_tensor = core.eager.Tensor(
+                            inner_x.dtype,
+                            inner_x.shape,
+                            inner_x.name + "cpy",
+                            core.VarDesc.VarType.DENSE_TENSOR,
+                            inner_x.persistable,
+                        )
+                    inner_x._unsafe_share_buffer_to(tmp_tensor)
+                    storage[holder_list[unpack_counter - 1]()] = tmp_tensor
+                return
+
+            def inner_unpack(inner_x):
+                raise Exception("An unexpected backward called on a tensor!")
+
+            if preserve_rng_state:
+                with (
+                    switch_rng_state_tracker(
+                        fw_cuda_rng_state,
+                        fwd_cuda_rng_state_tracker,
+                        fwd_numpy_state,
+                        fwd_random_state,
+                        fwd_custom_state,
+                        custom_get_state_func,
+                        custom_set_state_func,
+                    ),
+                    paddle.set_grad_enabled(True),
+                    paddle.amp.auto_cast(
+                        enable=is_fw_autocast,
+                        custom_white_list=amp_white_list,
+                        custom_black_list=amp_black_list,
+                        level=amp_level,
+                        dtype=amp_dtype,
+                    ),
+                    paddle.autograd.saved_tensors_hooks(
+                        inner_pack, inner_unpack
+                    ),
+                ):
+                    function(*args, **kwargs)
+            else:
+                with (
+                    paddle.set_grad_enabled(True),
+                    paddle.amp.auto_cast(
+                        enable=is_fw_autocast,
+                        custom_white_list=amp_white_list,
+                        custom_black_list=amp_black_list,
+                        level=amp_level,
+                        dtype=amp_dtype,
+                    ),
+                    paddle.autograd.saved_tensors_hooks(
+                        inner_pack, inner_unpack
+                    ),
+                ):
+                    function(*args, **kwargs)
+
+        if x not in storage:
+            raise Exception(
+                "Not supported to retrieve a tensor saved by autograd multiple times that is no need to recompute."
+            )
+
+        return storage.pop(x)
+
+    with paddle.autograd.saved_tensors_hooks(pack, unpack):
+        outputs = function(*args, **kwargs)
+
+    return outputs
+
+
+def recompute(function, *args, **kwargs):
+    """
+    recompute intermediate activations to save then memory.
+
+    Parameters:
+        function(paddle.nn.Layer): layer of sequence of layers that describes part of forward pass of the model
+            whose intermediate activations will be released to save memory in forward stage and will be recomputed
+            in backward stage for gradient calculation.
+        *args(Tensor): inputs to the function.
+        **kwargs(Dict): Kwargs should only contain two kinds of key-value params, the one is part of function's key-value params,
+            and the other contains 'preserve_rng_state', 'preserve_external_rng_state' and 'use_reentrant'.
+            The key-value pair of preserve_rng_state is used to indicate whether to save the forward rng. If it is True,
+            then the last forward rng value will be restored when the forward recalculation of backpropagation is performed,
+            its default value is True.
+            The key-value pair of preserve_external_rng_state is used to indicate whether to save and restore the external
+            random number generator states (numpy.random and python random). If your forward function does not use numpy.random
+            or python random, you can set this to False to improve performance. Its default value is True.
+            The key-value pair of use_reentrant is used to indicate which implementation of recompute you will be used.
+            'use_reentrant=True' means to use the PyLayer implementation of recompute, 'use_reentrant=False' means to
+            use the Hook implementation of recompute, its default value is True.
+    Returns:
+        Output of function on args.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED, env:GPU)
+            >>> import paddle
+            >>> from paddle.distributed.fleet.utils import recompute
+            >>> import random
+            >>> paddle.seed(2023)
+            >>> def get_fc_block(block_idx, input_size, is_last=False):
+            ...     block_name = "block_" + str(block_idx)
+            ...     block = paddle.nn.Sequential(
+            ...         (block_name + "_fc_0", paddle.nn.Linear(input_size, input_size, bias_attr=False)),
+            ...         (block_name + "_dropout", paddle.nn.Dropout(p=0.5)),
+            ...         (block_name + "_relu_1", paddle.nn.ReLU()),
+            ...         (block_name + "_fc_1", paddle.nn.Linear(input_size, input_size, bias_attr=False)),
+            ...         (block_name + "_relu_2", paddle.nn.ReLU()),
+            ...     )
+            ...     if is_last:
+            ...         block.add_sublayer(
+            ...             block_name + "_fc_2",
+            ...             paddle.nn.Linear(input_size, 1, bias_attr=False),
+            ...         )
+            ...     else:
+            ...         block.add_sublayer(
+            ...             block_name + "_fc_2",
+            ...             paddle.nn.Linear(input_size, input_size, bias_attr=False),
+            ...         )
+            ...     return block
+
+            >>> class Naive_fc_net(paddle.nn.Layer):
+            ...     def __init__(
+            ...         self,
+            ...         input_size=10,
+            ...         recompute_blocks=[1, 3],
+            ...         recompute_kwargs={},
+            ...     ):
+            ...         super().__init__()
+            ...         self.recompute_blocks = recompute_blocks
+            ...         self.recompute_kwargs = recompute_kwargs
+            ...         self.runfunc0 = get_fc_block(0, input_size, is_last=False)
+            ...         self.runfunc1 = get_fc_block(1, input_size, is_last=False)
+            ...         self.runfunc2 = get_fc_block(2, input_size, is_last=False)
+            ...         self.runfunc3 = get_fc_block(3, input_size, is_last=False)
+            ...         self.runfunc4 = get_fc_block(4, input_size, is_last=True)
+            ...         self.total_func = [self.runfunc0, self.runfunc1, self.runfunc2, self.runfunc3, self.runfunc4]
+            ...
+            ...     def forward(self, inputs):
+            ...         nums = len(self.total_func)
+            ...         for i in range(nums):
+            ...             if i in self.recompute_blocks:
+            ...                 inputs = recompute(self.total_func[i], inputs, **{"preserve_rng_state": True})
+            ...             else:
+            ...                 inputs = self.total_func[i](inputs)
+            ...         return inputs
+
+            >>> def run_model(cuda_state, recompute_block=[], recompute_kwargs={}):
+            ...     gen = paddle.seed(10)
+            ...     gen.manual_seed(10)
+            ...     random.seed(10)
+            ...     if cuda_state:
+            ...         paddle.set_cuda_rng_state(cuda_state)
+            ...     batch_size, input_size = 1, 10
+            ...     model = Naive_fc_net(
+            ...         input_size,
+            ...         recompute_blocks=recompute_block,
+            ...         recompute_kwargs=recompute_kwargs,
+            ...     )
+            ...     optimizer = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
+            ...     loss_ = []
+            ...     param_ = []
+            ...     grad_ = []
+            ...     for _ in range(5):
+            ...         x = paddle.rand(shape=[batch_size, input_size], dtype="float32")
+            ...         y_pred = model(x)
+            ...         loss = y_pred.mean()
+            ...         loss_.append(loss.item())
+            ...         loss.backward()
+            ...         optimizer.step()
+            ...         param_.append(model.parameters()[9])
+            ...         grad_.append(model.parameters()[3]._grad_ivar())
+            ...         optimizer.clear_grad()
+            ...     return loss_, param_, grad_
+
+            >>> cuda_state = paddle.get_cuda_rng_state()
+            >>> # without recompute
+            >>> loss_ref, param_ref, grad_ref = run_model(cuda_state, recompute_block=[])
+
+            >>> loss, param, grad = run_model(cuda_state, recompute_block=[1, 2])
+            >>> print("normal_loss: {}, recompute_loss: {}".format(loss_ref, loss))
+            >>> # The result of the recompute_loss should be the same as the normal_loss.
+            normal_loss: [0.0018744759727269411, 0.0, 0.035971127450466156, 0.0, 0.0], recompute_loss: [0.0018744759727269411, 0.0, 0.035971127450466156, 0.0, 0.0]
+
+    """
+    # Hack to mix *args with **kwargs in a python 2.7-compliant way
+    preserve = kwargs.pop('preserve_rng_state', True)
+    preserve_external_rng_state = kwargs.pop(
+        'preserve_external_rng_state', True
+    )
+
+    # whether to use reentrant method to implement recompute
+    use_reentrant = kwargs.pop('use_reentrant', True)
+
+    if custom_state_manager.custom_get_state_func is None:
+        assert custom_state_manager.custom_set_state_func is None
+        custom_get_state_func = lambda x=None: None
+        custom_set_state_func = lambda x=None: None
+    else:
+        custom_get_state_func = custom_state_manager.custom_get_state_func
+        custom_set_state_func = custom_state_manager.custom_set_state_func
+
+    if not in_dynamic_mode():
+        from paddle.distributed.auto_parallel.interface import (
+            recompute as static_auto_recompute,
+        )
+
+        return static_auto_recompute(function)(*args, **kwargs)
+
+    if framework._dygraph_tracer()._has_grad:
+        check_args = list(args)
+        check_args.extend(list(kwargs.values()))
+        check_recompute_necessary(check_args)
+
+    if use_reentrant:
+        offload_indices = kwargs.pop('offload_indices', [])
+        if not kwargs:  # fast path
+            return RecomputeFunction.apply(
+                function,
+                preserve,
+                preserve_external_rng_state,
+                offload_indices,
+                custom_get_state_func,
+                custom_set_state_func,
+                *args,
+            )
+
+        # rearrange `position-args + keyword-args` into `position-args`
+        target = (
+            function.forward
+            if isinstance(function, paddle.nn.Layer)
+            else function
+        )
+        if isinstance(target, StaticFunction):
+            target = target.dygraph_function
+
+        # Use getattr to get the cached signature. If it doesn't exist, parse and mount it to the target.
+        # This avoids the heavy overhead of inspect.signature during repeated executions.
+        cache_key = getattr(target, "__func__", target)
+        dyfunc_sig = _SIGNATURE_CACHE.get(cache_key)
+        if dyfunc_sig is None:
+            dyfunc_sig = inspect.signature(target)
+            _SIGNATURE_CACHE[cache_key] = dyfunc_sig
+
+        bound_args = dyfunc_sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        input_args = []
+        for arg, param in zip(
+            bound_args.arguments.values(), dyfunc_sig.parameters.values()
+        ):
+            if param.kind == param.VAR_POSITIONAL:
+                input_args.extend(arg)
+            elif param.kind in (
+                param.POSITIONAL_ONLY,
+                param.POSITIONAL_OR_KEYWORD,
+            ):
+                input_args.append(arg)
+            elif param.kind == param.VAR_KEYWORD:
+                input_args.extend(arg.values())
+            elif param.kind == param.KEYWORD_ONLY:
+                raise ValueError(
+                    "Currently, keyword-only arguments are not supported when you want to send kwargs(dict parameter) to function with use_reentrant=True."
+                )
+            else:
+                raise ValueError("Unknown parameter kind.")
+        return RecomputeFunction.apply(
+            function,
+            preserve,
+            preserve_external_rng_state,
+            offload_indices,
+            custom_get_state_func,
+            custom_set_state_func,
+            *input_args,
+        )
+    else:
+        return _recompute_without_reentrant(
+            function,
+            custom_get_state_func,
+            custom_set_state_func,
+            preserve,
+            preserve_external_rng_state,
+            *args,
+            **kwargs,
+        )
+
+
+def recompute_sequential(
+    ctx: _Ctx,
+    functions: Sequential | Sequence[Callable[..., Any]],
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    """
+    recompute intermediate activations to save the memory for 'Sequential' models. use 'ctx' to transmit some context params, it is similar to 'recompute_hybrid' API.
+
+    Parameters:
+        ctx(dict): include 'segments' and  'preserve_rng_state' keys, the key 'segments' (int, default 1), represents the number of chunks to create in the model,
+                   the key 'preserve_rng_state' (bool, optional, default=True) indicate whether to save the forward rng. If it is True, then the last forward rng value will be
+                   restored when the forward recalculation of backpropagation is performed.
+        functions(paddle.nn.Sequential): layer of sequence of layers that describes part of forward pass of the model
+              whose intermediate activations will be released to save memory in forward stage and will be recomputed
+              in backward stage for gradient calculation.
+        *args(Tensor): inputs(tuple) to the function.
+        **kwargs(Dict): inputs(dict) to the function.
+
+    Returns:
+        Output of function on args and kwargs.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle
+            >>> from paddle.incubate.distributed.fleet import recompute_sequential
+            >>> input = paddle.ones(shape=[8, 10])
+            >>> model = paddle.nn.Sequential(paddle.nn.Linear(10, 10), paddle.nn.Linear(10, 2))
+            >>> output = recompute_sequential({'segments': 1}, model, input)
+
+    """
+    segments = ctx.get('segments', 1)
+    preserve_rng_state = ctx.get('preserve_rng_state', True)
+
+    def _run_func(begin, end, funcs):
+        def do_run(input):
+            for i in range(begin, end + 1):
+                input = funcs[i](input)
+            return input
+
+        return do_run
+
+    if isinstance(functions, paddle.nn.Sequential):
+        functions = list(functions.children())
+
+    segment_size = len(functions) // segments
+
+    end = -1
+    for begin in range(0, segment_size * (segments - 1), segment_size):
+        end = begin + segment_size - 1
+        args = recompute(
+            _run_func(begin, end, functions),
+            *args,
+            preserve_rng_state=preserve_rng_state,
+            **kwargs,
+        )
+    return _run_func(end + 1, len(functions) - 1, functions)(*args)

python/paddle/distributed/fleet/recompute/recompute_hybrid.py

@@ -0,0 +1,347 @@
+# 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 random
+from typing import TYPE_CHECKING, Any, TypedDict
+
+import numpy as np
+
+import paddle
+from paddle import framework
+from paddle.autograd import PyLayer
+from paddle.framework import core
+
+from ..meta_parallel.parallel_layers.random import get_rng_state_tracker
+from ..meta_parallel.pp_utils import utils
+from .recompute import (
+    check_recompute_necessary,
+    custom_state_manager,
+    detach_variable,
+    switch_rng_state_tracker,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from typing_extensions import NotRequired
+
+    from paddle.distributed.communication.group import Group
+    from paddle.nn import Layer
+
+    class _Ctx(TypedDict):
+        mp_group: Group
+        offload: NotRequired[bool]
+        partition: NotRequired[bool]
+
+
+__all__ = []
+
+
+def _split_activation(tensor, mp_group):
+    mp_degree = mp_group.nranks
+    mp_rank = mp_group.rank
+    if mp_degree < 2:
+        return tensor
+
+    tensor_numel = paddle.numel(tensor)
+    assert tensor_numel != 0, "can't recompute zero element"
+    assert tensor_numel % mp_degree == 0, (
+        f"The capacity of the activation ({tensor_numel}) cannot be divisible by mp_degree({mp_degree})"
+    )
+
+    # use inplace operation to save memory
+    data = tensor.flatten_()
+
+    part_size = tensor_numel // mp_degree
+    start = part_size * mp_rank
+    end = start + part_size
+    return data[start:end]
+
+
+def _merge_activation(tensor, mp_group):
+    mp_degree = mp_group.nranks
+    mp_rank = mp_group.rank
+    if mp_degree < 2:
+        return tensor
+
+    # adapt to new dygraph
+    tensor_shape = list(tensor.shape)
+    tensor_shape[0] *= mp_group.nranks
+    out = paddle.empty(tensor_shape, tensor.dtype)
+    task = mp_group.process_group.all_gather(tensor.cuda(), out)
+    task.wait()
+    return out
+
+
+class _HPRecomputeFunction(PyLayer):
+    """
+    Compared with paddle.distributed.fleet.utils.recompute, there are the following differences:
+    1. In order to support PipeLineParallel, the input of recompute is modified to ensure that the input can be tuple type.
+    2. Offload support for activation
+    3. Support MP segmentation of activation to further reduce cuda memory
+    4. Adapt to the random state of MP
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        run_function,
+        all_outputs,
+        mp_group,
+        offload,
+        partition,
+        custom_get_state_func,
+        custom_set_state_func,
+        *args,
+        **kwargs,
+    ):
+        # store for recomputing
+        ctx.run_function = run_function
+
+        ctx.kwargs = kwargs
+
+        # store the rng states
+        ctx.fwd_rng_state = paddle.get_rng_state()
+        ctx.fwd_rng_state_tracker = get_rng_state_tracker().get_states_tracker()
+        ctx.fwd_numpy_state = np.random.get_state()
+        ctx.fwd_random_state = random.getstate()
+        ctx.fwd_custom_state = custom_get_state_func()
+        ctx.custom_get_state_func = custom_get_state_func
+        ctx.custom_set_state_func = custom_set_state_func
+
+        # save config info
+        ctx.mp_group = mp_group
+        ctx.offload = offload
+        ctx.partition = partition
+
+        # save input for backward
+        ctx.inputs = []
+        ctx.tensor_indices = []
+        ctx.tensor_shapes = []
+        tensor_inputs = []
+
+        cur_device = paddle.get_device()
+        assert (
+            'gpu:' in paddle.get_device()
+            or 'xpu:' in paddle.get_device()
+            or cur_device.split(':')[0]
+            in paddle.device.get_all_custom_device_type()
+        ), f"Recompute with RNG is not support current device: {cur_device}."
+
+        # TODO support AMP
+        tracer = framework._dygraph_tracer()
+        ctx.is_fw_autocast = (
+            False if tracer._amp_level == core.AmpLevel.O0 else True
+        )
+        if tracer._amp_level == core.AmpLevel.O2:
+            ctx.amp_level = 'O2'
+        elif tracer._amp_level in (core.AmpLevel.O1, core.AmpLevel.O0):
+            ctx.amp_level = 'O1'
+        else:
+            raise ValueError(f"unsupported amp level: {tracer._amp_level}")
+        ctx.amp_dtype = tracer._amp_dtype
+        ctx.amp_white_list, ctx.amp_black_list = tracer._get_amp_op_list()
+
+        with paddle.no_grad():
+            outputs = run_function(*args, **kwargs)
+
+        for i, arg in enumerate(args):
+            if paddle.is_tensor(arg):
+                state = arg.stop_gradient
+                if partition:
+                    ctx.tensor_shapes.append(arg.shape)
+                    partition = _split_activation(
+                        arg.detach(), mp_group
+                    ).clone()
+                    # TODO(shenliang03) not use calculate stream to D2H to speed
+                    arg = partition.cpu() if offload else partition
+                else:
+                    arg = arg.cpu() if offload else arg
+                arg.stop_gradient = state
+                tensor_inputs.append(arg)
+                ctx.tensor_indices.append(i)
+                ctx.inputs.append(None)
+
+                # In new dygraph mode, in some cases a subset of outputs is identity to the subset of inputs,
+                #  which is inplace operating. When the inputs' stop_gradient is True, an
+                #  error will occurs because the stop_gradient=True and inplace-op are not
+                #  supported in the same time. The solution is to mark the inputs non_differentiable
+                #  if its stop_gradient is True.
+                # Note:
+                #  If not marked non_differentiable, all output tensors' attr `stop gradient`
+                #  will be reset to `False` in c++ backend.
+                #  See https://github.com/PaddlePaddle/Paddle/blob/9d62efb0e6e5373823039d9eda96cd5905426c0a/paddle/fluid/pybind/eager_py_layer.cc#L388
+                if framework.in_dynamic_mode() and state:
+                    ctx.mark_non_differentiable(arg)
+            else:
+                ctx.inputs.append(arg)
+
+        ctx.save_for_backward(*tensor_inputs)
+
+        if paddle.is_tensor(outputs):
+            all_outputs += [outputs]
+            return outputs
+        else:
+            all_outputs += outputs
+            return tuple(outputs)
+
+    @staticmethod
+    def backward(ctx, *args):
+        with paddle.base.dygraph.guard():
+            # Restore inputs
+            inputs = list(ctx.inputs)
+            tensor_indices = ctx.tensor_indices
+            tensor_shapes = ctx.tensor_shapes
+            tensors = list(ctx.saved_tensor())
+
+            device_id = paddle.distributed.ParallelEnv().device_id
+            for i, idx in enumerate(tensor_indices):
+                if ctx.partition:
+                    state = tensors[i].stop_gradient
+                    tensors[i] = (
+                        _merge_activation(tensors[i], ctx.mp_group)
+                        .detach()
+                        .reshape_(tensor_shapes[i])
+                    )
+                    tensors[i].stop_gradient = state
+                inputs[idx] = (
+                    tensors[i].cuda(device_id) if ctx.offload else tensors[i]
+                )
+
+            tracer = framework._dygraph_tracer()
+            tracer._has_grad = True
+
+            # need restore auto_cast state as well as w/b list
+            with switch_rng_state_tracker(
+                ctx.fwd_rng_state,
+                ctx.fwd_rng_state_tracker,
+                ctx.fwd_numpy_state,
+                ctx.fwd_random_state,
+                ctx.fwd_custom_state,
+                ctx.custom_get_state_func,
+                ctx.custom_set_state_func,
+            ):
+                if ctx.is_fw_autocast:
+                    with paddle.amp.auto_cast(
+                        enable=ctx.is_fw_autocast,
+                        custom_white_list=ctx.amp_white_list,
+                        custom_black_list=ctx.amp_black_list,
+                        level=ctx.amp_level,
+                        dtype=ctx.amp_dtype,
+                    ):
+                        detached_inputs = detach_variable(tuple(inputs))
+                        outputs = ctx.run_function(
+                            *detached_inputs, **ctx.kwargs
+                        )
+                else:
+                    detached_inputs = detach_variable(tuple(inputs))
+                    outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
+
+            if isinstance(outputs, core.eager.Tensor):
+                outputs = (outputs,)
+            assert len(outputs) == len(args)
+
+            forward_outputs_with_grad = []
+            backward_inputs = []
+
+            for i in range(len(outputs)):
+                if (
+                    isinstance(outputs[i], core.eager.Tensor)
+                    and not outputs[i].stop_gradient
+                ):
+                    forward_outputs_with_grad.append(outputs[i])
+                    backward_inputs.append(args[i])
+
+            if len(forward_outputs_with_grad) == 0:
+                raise RuntimeError(
+                    "none of output has stop_gradient=False, this recompute() is not necessary"
+                )
+
+            # actually backward
+            paddle.autograd.backward(forward_outputs_with_grad, backward_inputs)
+            grads = tuple(
+                inp._grad_ivar()
+                for inp in detached_inputs
+                if isinstance(inp, core.eager.Tensor)
+            )
+            return grads
+
+
+def recompute_hybrid(
+    ctx: _Ctx, function: Layer | Callable[..., Any], *args: Any, **kwargs: Any
+) -> Any:
+    """
+    recompute intermediate activations to save the memory in hybrid parallel scene.
+    # NOTE(shenliang03)The current hybrid parallel recompute has limitations.
+    # It cannot handle the following situations:
+    # 1. The calculation output of recompute, there are tensors that do not require gradients.
+    # 2. The forward output tensor has no gradient. This problem can be solved temporarily by detach().
+    # 3. Here, we only use float dtype to distinguish whether a gradient is needed in output tensor
+
+    Parameters:
+        ctx(dict): include 'mp_group', 'offload', and 'partition' keys. the key 'mp_group' (Group), represents the activations are splitted
+                   in which group. the key 'offload' (bool, optional, default=False), represents whether to offload to cpu. the key 'partition' (bool, optional, default=False),
+                   represents whether to split activations in the mp_group.
+        function(paddle.nn.Layer): layer of sequence of layers that describes part of forward pass of the model
+              whose intermediate activations will be released to save memory in forward stage and will be recomputed
+              in backward stage for gradient calculation.
+        *args(Tensor): inputs(tuple) to the function.
+
+        **kwargs(Dict): inputs(dict) to the function.
+
+    Returns:
+        Output of function on args and kwargs.
+
+    """
+    mp_group = ctx.get('mp_group', None)
+    assert mp_group is not None, (
+        "ctx must contains mp_group and mp_group can not be None."
+    )
+
+    offload = ctx.get('offload', False)
+    partition = ctx.get('partition', False)
+
+    if framework._dygraph_tracer()._has_grad:
+        check_recompute_necessary(args)
+
+    if custom_state_manager.custom_get_state_func is None:
+        assert custom_state_manager.custom_set_state_func is None
+        custom_get_state_func = lambda x=None: None
+        custom_set_state_func = lambda x=None: None
+    else:
+        custom_get_state_func = custom_state_manager.custom_get_state_func
+        custom_set_state_func = custom_state_manager.custom_set_state_func
+
+    all_outputs = []
+    _HPRecomputeFunction.apply(
+        function,
+        all_outputs,
+        mp_group,
+        offload,
+        partition,
+        custom_get_state_func,
+        custom_set_state_func,
+        *args,
+        **kwargs,
+    )
+
+    if len(all_outputs) == 1:
+        return all_outputs[0]
+    else:
+        for output in all_outputs:
+            if paddle.is_tensor(output) and not utils.is_float_tensor(output):
+                output.stop_gradient = True
+
+        return tuple(all_outputs)

python/paddle/distributed/fleet/runtime/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2020 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 .collective_runtime import CollectiveRuntime  # noqa: F401
+from .parameter_server_runtime import ParameterServerRuntime  # noqa: F401
+from .the_one_ps import TheOnePSRuntime  # noqa: F401
+
+__all__ = []

python/paddle/distributed/fleet/runtime/collective_runtime.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2020 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 .runtime_base import RuntimeBase
+
+__all__ = []
+
+
+class CollectiveRuntime(RuntimeBase):
+    def __init__(self):
+        super().__init__()
+
+    def _init_worker(self):
+        logging.warning(
+            "You should not call 'init_worker' method for collective mode."
+        )
+
+    def _run_worker(self):
+        logging.warning(
+            "You should not call 'run_worker' method for collective mode."
+        )
+
+    def _init_server(self, *args, **kwargs):
+        logging.warning(
+            "You should not call 'init_server' method for collective mode."
+        )
+
+    def _run_server(self):
+        logging.warning(
+            "You should not call 'run_server' method for collective mode."
+        )
+
+    def _stop_worker(self):
+        logging.warning(
+            "You should not call 'stop_worker' method for collective mode."
+        )
+
+    # save inference model should be added here