chore: import upstream snapshot with attribution

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)