chore: import upstream snapshot with attribution

python/paddle/incubate/distributed/__init__.py

@@ -0,0 +1,15 @@
+# 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 . import utils as utils

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

@@ -0,0 +1,20 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.distributed.fleet.recompute import (
+    recompute_hybrid,
+    recompute_sequential,
+)
+
+__all__ = ["recompute_sequential", "recompute_hybrid"]

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

@@ -0,0 +1,383 @@
+#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+
+from paddle import base
+from paddle.base.executor import Executor
+from paddle.distributed.fleet.base.role_maker import RoleMakerBase
+from paddle.optimizer import SGD
+from paddle.static.amp.decorator import OptimizerWithMixedPrecision
+
+__all__ = []
+
+
+class Mode:
+    """
+    There are various mode for fleet, each of them is designed for different model.
+    """
+
+    TRANSPILER = 1
+    PSLIB = 2
+    COLLECTIVE = 3
+
+
+class Fleet(metaclass=abc.ABCMeta):
+    """
+    Fleet is the base class, transpiler and pslib are implementation of Fleet.
+
+    Args:
+        mode(Mode): the implementation of Fleet's mode.
+
+    Returns:
+        None
+    """
+
+    def __init__(self, mode):
+        self._is_initialized = False
+        self._mode = mode
+        self._optimizer = None
+        self._role_maker = None
+        self._executor = None
+
+    def is_first_worker(self):
+        """
+        Check whether the node is the first instance of worker.
+
+        Returns:
+            bool: True if this is the first node of worker,
+                  False if not.
+        """
+        return self._role_maker.is_first_worker()
+
+    def worker_index(self):
+        """
+        Get current worker index.
+
+        Returns:
+            int: node id
+        """
+        return self._role_maker.worker_index()
+
+    def worker_num(self):
+        """
+        Get current total worker number.
+
+        Returns:
+            int: worker numbers
+        """
+        return self._role_maker.worker_num()
+
+    def is_worker(self):
+        """
+        Check whether the node is an instance of worker.
+
+        Returns:
+            bool: True if this is a node of worker,
+                  False if not.
+        """
+        return self._role_maker.is_worker()
+
+    def worker_endpoints(self, to_string=False):
+        """
+        Get current server endpoints, such as ["127.0.0.1:1001", "127.0.0.1:1002"].
+
+        Returns:
+            list/string: server endpoints
+        """
+
+        if to_string:
+            return ",".join(self._role_maker.get_trainer_endpoints())
+        else:
+            return self._role_maker.get_trainer_endpoints()
+
+    def server_num(self):
+        """
+        Get current total worker number.
+
+        Returns:
+            int: server number
+        """
+        return len(self._role_maker.get_pserver_endpoints())
+
+    def server_index(self):
+        """
+        Get current server index.
+
+        Returns:
+            int: node id
+        """
+        return self._role_maker.server_index()
+
+    def server_endpoints(self, to_string=False):
+        """
+        Get current server endpoints, such as ["127.0.0.1:1001", "127.0.0.1:1002"].
+
+        Returns:
+            list/string: server endpoints
+        """
+
+        if to_string:
+            return ",".join(self._role_maker.get_pserver_endpoints())
+        else:
+            return self._role_maker.get_pserver_endpoints()
+
+    def is_server(self):
+        """
+        Check whether the node is an instance of server.
+
+        Returns:
+            bool: True if this is a node of server,
+                  False if not
+        """
+        return self._role_maker.is_server()
+
+    def is_xpu(self):
+        """
+        Check whether the node is an instance of server.
+
+        Returns:
+            bool: True if this is a node of server,
+                  False if not.
+        """
+        return self._role_maker.is_xpu()
+
+    def split_files(self, files):
+        """
+        split files before distributed training,
+        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 belongs to this worker.
+        """
+        if not isinstance(files, list):
+            raise TypeError("files should be a list of file need to be read.")
+
+        trainer_id = self.worker_index()
+        trainers = self.worker_num()
+
+        remainder = len(files) % trainers
+        blocksize = 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 init(self, role_maker=None):
+        """
+        should be called only once in user's python scripts,
+        init() will initialize RoleMaker which is used for identifying
+            current node's role, e.g. worker, server, etc.
+
+        Args:
+            role_maker(RoleMakerBase): subclass of RoleMakerBase.
+
+        Returns:
+            None
+        """
+        self._executor = Executor(base.CPUPlace())
+
+        if role_maker and not isinstance(role_maker, RoleMakerBase):
+            from paddle.incubate.distributed.fleet.role_maker import (
+                RoleMakerBase as RoleMakerBaseIncubate,
+            )
+
+            if role_maker and not isinstance(role_maker, RoleMakerBaseIncubate):
+                raise TypeError(
+                    "role_maker must be an instance of RoleMakerBase"
+                )
+
+        self._role_maker = role_maker
+        self._role_maker.generate_role()
+        self._is_initialized = True
+
+    def all_reduce_worker(self, input, output):
+        """
+        all reduce between workers, only support array of one dim.
+
+        Args:
+            input(list|numpy.array): array of one dim
+            output(list|numpy.array): array of one dim
+        """
+        self._role_maker.all_reduce_worker(input, output)
+
+    def barrier_worker(self):
+        """
+        barrier between workers
+        """
+        self._role_maker.barrier_worker()
+
+    @abc.abstractmethod
+    def init_worker(self):
+        pass
+
+    @abc.abstractmethod
+    def init_server(self, model_dir=None, **kwargs):
+        pass
+
+    @abc.abstractmethod
+    def run_server(self):
+        pass
+
+    @abc.abstractmethod
+    def stop_worker(self):
+        pass
+
+    @abc.abstractmethod
+    def distributed_optimizer(self, optimizer, strategy=None):
+        pass
+
+    @abc.abstractmethod
+    def save_inference_model(
+        self,
+        executor,
+        dirname,
+        feeded_var_names,
+        target_vars,
+        main_program=None,
+        export_for_deployment=True,
+        legacy_format=False,
+    ):
+        pass
+
+    @abc.abstractmethod
+    def save_persistables(self, executor, dirname, main_program=None):
+        pass
+
+
+class DistributedOptimizer(metaclass=abc.ABCMeta):
+    """
+    DistributedOptimizer is a wrapper for paddle.base.optimizer
+    A user should pass a paddle.base.optimizer to DistributedOptimizer
+    minimize() function is implemented.
+    DistributedOptimizer is the starting point for a user who wants to
+    run distributed training. The optimized information will be stored in
+    Fleet() instance who holds the global information about current distributed
+    training.
+
+    Args:
+        optimizer(Optimizer): subclass of Optimizer.
+        strategy(any): the user define config for Optimizer.
+
+    Returns:
+        None
+
+    """
+
+    def __init__(self, optimizer, strategy=None):
+        if not isinstance(optimizer, SGD.__bases__) and not isinstance(
+            optimizer, OptimizerWithMixedPrecision
+        ):
+            raise TypeError("optimizer must be an instance of Optimizer")
+
+        self._optimizer = optimizer
+        self._strategy = strategy
+
+    @abc.abstractmethod
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        """
+        First part of `minimize`, do auto-diff to append backward ops for
+        the current program.
+
+        Args:
+            loss (Variable): loss variable to run optimizations.
+            startup_program (Program): startup_program for initializing parameters
+                in `parameter_list`.
+            parameter_list (list): list of Variables to update.
+            no_grad_set (set|None): set of Variables should be ignored.
+            callbacks (list|None): list of callables to run when appending backward
+                operator for one parameter.
+
+        Return:
+            list: list of (param, grad) pair, grad is the output of backward.
+
+        Examples:
+            See examples in `apply_gradients`.
+        """
+        pass
+
+    @abc.abstractmethod
+    def apply_gradients(self, params_grads):
+        """
+        Second part of `minimize`, appending optimization operators for
+        given `params_grads` pairs.
+
+        Args:
+            params_grads (list): list of (param, grad) pair to do optimization.
+
+        Returns:
+            list: A list of operators appended to the current program.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> # doctest: +SKIP('The network is not defined.')
+                >>> loss = network()
+                >>> optimizer = base.optimizer.SGD(learning_rate=0.1)
+                >>> params_grads = optimizer.backward(loss)
+                >>> # you may append operations for params_grads here
+                >>> # ...
+                >>> optimizer.apply_gradients(params_grads)
+        """
+        pass
+
+    @abc.abstractmethod
+    def minimize(
+        self,
+        losses,
+        scopes=None,
+        startup_programs=None,
+        parameter_list=None,
+        no_grad_set=None,
+    ):
+        """
+        Add operations to minimize `loss` by updating `parameter_list`.
+
+        This method combines interface `backward()` and
+        `apply_gradients()` into one.
+
+        Args:
+            losses (Variable|Variable List): loss variable to run optimizations.
+            scopes (Scope| Scope List): scope instance.
+            startup_programs (Program|Program List): startup_program for initializing parameters
+                in `parameter_list`.
+            parameter_list (list): list of Variables to update.
+            no_grad_set (set|None): set of Variables should be ignored.
+
+        Returns:
+            tuple: (optimize_ops, params_grads) which are, list of operators appended;
+            and list of (param, grad) Variables pair for optimization.
+        """
+        pass

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

@@ -0,0 +1,554 @@
+#   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 logging
+import os
+
+import paddle
+import paddle.distributed.transpiler.distribute_transpiler as dist_transpiler
+from paddle import base
+from paddle.base.compiler import CompiledProgram
+from paddle.base.executor import Executor
+from paddle.base.framework import Program
+from paddle.base.incubate.checkpoint.checkpoint_saver import (
+    CheckpointSaver,
+    PaddleModel,
+)
+from paddle.distributed.fleet.meta_optimizers import RawProgramOptimizer
+from paddle.incubate.distributed.fleet.base import (
+    DistributedOptimizer,
+    Fleet,
+    Mode,
+)
+from paddle.static import io
+
+
+class Collective(Fleet):
+    def __init__(self):
+        super().__init__(Mode.COLLECTIVE)
+        self._local_ip = 0
+
+        self.startup_program = None
+        self._origin_program = None
+        self._transpiled_program = None
+        self.main_program = None
+        self._checkpoint_prefix = "__paddle_fleet_checkpoint__"
+        self._param_file_name = "_paddle_fleet_param__"
+
+    def init_worker(self):
+        logging.warning(
+            "You should not call 'init_worker' method for collective mode."
+        )
+
+    def run_worker(self, main_programs=None, scopes=None):
+        logging.warning(
+            "You should not call 'run_worker' method for collective mode."
+        )
+
+    def init_server(self, model_dir=None):
+        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."
+        )
+
+    def distributed_optimizer(self, optimizer, strategy=None):
+        self._optimizer = CollectiveOptimizer(optimizer, strategy)
+        return self._optimizer
+
+    def save_inference_model(
+        self,
+        executor,
+        path_prefix,
+        feeded_vas=None,
+        fetch_vars=None,
+        program=None,
+        legacy_format=False,
+    ):
+        """
+        Prune the given `main_program` to build a new program especially for
+        inference, and then save it and all related parameters to given
+        `dirname` by the `executor`.
+        """
+        assert isinstance(executor, Executor), (
+            "In fleet.save_inference_model() function, executor must be as"
+            " Executor type."
+        )
+
+        if program is None:
+            program = self._origin_program
+        assert isinstance(program, Program), (
+            "In fleet.save_inference_model() function, main_program "
+            "must be as Program type."
+        )
+
+        io.save_inference_model(
+            path_prefix,
+            feeded_vas,
+            fetch_vars,
+            executor,
+            program=program,
+            legacy_format=legacy_format,
+        )
+
+    def save_persistables(
+        self, executor, dirname, main_program=None, filename=None
+    ):
+        """
+        This function filters out all variables with `persistable==True` from
+        the give `main_program` and then saves these variables to the folder
+        `dirname` or file `filename`.
+
+        The `dirname` is used to specify the folder where persistable variables
+        are going to be saved. If you would like to save variables in separate
+        files, set `filename` None; if you would like to save all variables in a
+        single file, use `filename` to specify the file name.
+        """
+        assert isinstance(executor, Executor), (
+            "In fleet.save_inference_model() function, executor must be as"
+            " Executor type."
+        )
+
+        if main_program is None:
+            main_program = self._origin_program
+
+        assert isinstance(main_program, Program), (
+            "In fleet.save_inference_model() function, main_program "
+            "must be as Program type."
+        )
+
+        paddle.distributed.io.save_persistables(
+            executor, dirname, main_program, filename=filename
+        )
+
+    def save_checkpoint(
+        self,
+        executor,
+        path,
+        trainer_id,
+        train_status,
+        fs,
+        main_program=None,
+        local_cache_path=".cache",
+        remain_all_checkpoint=True,
+    ):
+        """
+        This function save persistables and current epoch num to path.
+        """
+        if main_program is None:
+            main_program = self._transpiled_program
+
+        m = PaddleModel(executor, main_program)
+        t = train_status
+        c = CheckpointSaver(fs)
+        real_path, checkpoint_no = c.save_checkpoint(
+            path=path,
+            slists=[m, t],
+            trainer_id=trainer_id,
+            local_cache_path=local_cache_path,
+        )
+
+        if not remain_all_checkpoint:
+            c.clean_redundant_checkpoints(path)
+
+        return real_path, checkpoint_no
+
+    def load_checkpoint(
+        self,
+        executor,
+        path,
+        trainer_id,
+        train_status,
+        fs,
+        main_program=None,
+        local_cache_path=".cache",
+        ignore_empty=True,
+    ):
+        """
+        This function load persistables and current epoch num from path.
+        """
+
+        if main_program is None:
+            main_program = self._transpiled_program
+
+        m = PaddleModel(executor, main_program)
+        c = CheckpointSaver(fs)
+        return c.load_checkpoint(
+            path,
+            [m, train_status],
+            trainer_id=trainer_id,
+            ignore_empty=ignore_empty,
+            local_cache_path=local_cache_path,
+        )
+
+
+fleet = Collective()
+
+
+class DistributedStrategy(base.BuildStrategy):
+    """
+    Init function of DistributedStrategy
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.use_local_sgd = False
+        self.use_dist_fc = False
+
+        self.dist_fc_config = None  # DistFCConfig
+        self.mode = "nccl2"  # or collective
+        self.collective_mode = None  # local_sgd or grad_allreduce
+        self.nccl_comm_num = 1
+        self.forward_recompute = False  # use RecomputeOptimizer
+        self.recompute_checkpoints = []
+        self.use_amp = False  # use mixed precision optimizer
+        self.amp_loss_scaling = 2**15
+
+        # configurations below are used for unit test
+        self._ut4grad_allreduce = False
+
+
+class CollectiveOpBasedOptimizer(DistributedOptimizer):
+    """
+    Collective Operator Base Class For Distributed Optimizer
+    The class is invisible to a user
+    """
+
+    def __init__(self, optimizer, strategy=None):
+        assert isinstance(strategy, DistributedStrategy), (
+            "strategy must be DistributedStrategy"
+        )
+        super().__init__(optimizer, strategy)
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self._optimizer.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_gradients(self, params_grads):
+        return self._optimizer.apply_gradients(params_grads)
+
+
+class CollectiveOptimizer(DistributedOptimizer):
+    """
+    DistributedOptimizer is a wrapper for paddle.base.optimizer
+    A user should pass a paddle.base.optimizer to DistributedOptimizer
+    minimize() function is implemented.
+    DistributedOptimizer is the starting point for a user who wants to
+    run distributed training. The optimized information will be stored in
+    Fleet() instance who holds the global information about current distributed
+    training.
+    """
+
+    def __init__(self, optimizer, strategy=DistributedStrategy()):
+        if strategy is None:
+            strategy = DistributedStrategy()
+        super().__init__(optimizer, strategy)
+        self._forward_recompute = strategy.forward_recompute
+        if not isinstance(strategy.recompute_checkpoints, list):
+            raise ValueError(
+                "DistStrategy.recompute_checkpoints should be a List"
+            )
+        self._recompute_checkpoints = strategy.recompute_checkpoints
+        self._use_amp = strategy.use_amp
+        self._amp_loss_scaling = strategy.amp_loss_scaling
+        self.print_config = False
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self._optimizer.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_gradients(self, params_grads):
+        return self._optimizer.apply_gradients(params_grads)
+
+    def _check_condition(self, name, **kwargs):
+        for k, v in kwargs.items():
+            if v is True:
+                raise AssertionError(f"you can't use {name} and {k} together")
+
+    def _check_collective_mode(self, main_program, optimizer, strategy):
+        """
+        Check the conflict conditions.
+        """
+        if strategy.use_local_sgd:
+            strategy.mode = "collective"
+            strategy.collective_mode = "local_sgd"
+            self._check_condition(
+                "use_local_sgd",
+                use_dgc=main_program._enable_dgc,
+                use_dist_fc=strategy.use_dist_fc,
+                use_lamb=main_program._use_lamb,
+            )
+
+        if strategy.use_dist_fc:
+            self._check_condition(
+                "use_dist_fc",
+                use_dgc=main_program._enable_dgc,
+                use_local_sgd=strategy.use_local_sgd,
+                use_lamb=main_program._use_lamb,
+            )
+            assert strategy.dist_fc_config is not None, (
+                "DistributedStrategy.dist_fc_config should be set"
+            )
+
+        if strategy._ut4grad_allreduce:
+            strategy.mode = "collective"
+            strategy.collective_mode = "grad_allreduce"
+            self._check_condition(
+                "_ut4grad_allreduce",
+                use_dgc=main_program._enable_dgc,
+                use_lamb=main_program._use_lamb,
+            )
+
+        if (
+            self._strategy.collective_mode == "local_sgd"
+            or self._strategy.collective_mode == "grad_allreduce"
+        ):
+            assert self._strategy.mode == "collective", (
+                "local_sgd and grad_allreduce can be used under collective mode"
+            )
+
+    def _transpile(self, startup_program, main_program):
+        """
+        Transpile the programs to distributed programs. And add the variables.
+        """
+        worker_endpoints = fleet.worker_endpoints()
+        trainer_id = fleet.worker_index()
+        current_endpoint = fleet.worker_endpoints()[trainer_id]
+        worker_endpoints_env = ','.join(worker_endpoints)
+        trainers_num = fleet.worker_num()
+
+        if self.print_config:
+            print(
+                f"worker_endpoints:{worker_endpoints} trainers_num:{trainers_num} current_endpoint:{current_endpoint} \
+                  trainer_id:{trainer_id}"
+            )
+
+        # call transpiler
+        config = dist_transpiler.DistributeTranspilerConfig()
+        config.mode = self._strategy.mode
+        config.collective_mode = self._strategy.collective_mode
+
+        config.nccl_comm_num = self._strategy.nccl_comm_num
+        config.use_hierarchical_allreduce = (
+            self._strategy.use_hierarchical_allreduce
+        )
+        config.hierarchical_allreduce_inter_nranks = (
+            self._strategy.hierarchical_allreduce_inter_nranks
+        )
+
+        t = dist_transpiler.DistributeTranspiler(config=config)
+        t.transpile(
+            trainer_id=trainer_id,
+            trainers=worker_endpoints_env,
+            startup_program=startup_program,
+            program=main_program,
+            current_endpoint=current_endpoint,
+        )
+
+    def _get_node_ips_from_endpoints(self, endpoints):
+        ss = set()
+        ips = []
+        for ep in endpoints:
+            ip = ep.split(":")[0].strip()
+            if ip not in ss:
+                ss.add(ip)
+                ips.append(ip)
+            else:
+                continue
+
+        return ips
+
+    def _node_num(self):
+        worker_endpoints = fleet.worker_endpoints()
+        current_endpoint = fleet.worker_endpoints()[fleet.worker_index()]
+        worker_endpoints_env = ','.join(worker_endpoints)
+
+        node_ips = self._get_node_ips_from_endpoints(worker_endpoints)
+        node_ip = current_endpoint.split(":")[0].strip()
+
+        node_num = len(node_ips)
+
+        return node_num
+
+    def _try_to_compile(self, startup_program, main_program):
+        node_num = self._node_num()
+        assert node_num >= 1, f"nccl2 node_num must >= 1, now:{node_num}"
+
+        if node_num <= 1:
+            if self._strategy.nccl_comm_num > 1:
+                logging.warning(
+                    "set nccl_comm_num=1 since you only have 1 node."
+                )
+            self._strategy.nccl_comm_num = 1
+
+            if self._strategy.use_hierarchical_allreduce:
+                logging.warning(
+                    "set use_hierarchical_allreduce=False since you only have 1 node."
+                )
+            self._strategy.use_hierarchical_allreduce = False
+
+        sync_allreduce = os.getenv("FLAGS_sync_nccl_allreduce")
+
+        # NOTE. open sync_batch_norm will hang when use multi num_threads
+        sync_batch_norm = self._strategy.sync_batch_norm
+        if sync_batch_norm is not None and sync_batch_norm is True:
+            self._strategy.nccl_comm_num = 1
+            self._strategy.use_hierarchical_allreduce = False
+            logging.warning(
+                "use sync_batch_norm will hang when set num_threads > 1, so "
+                "set num_threads=1, nccl_comm_num=1, use_hierarchical_allreduce=False."
+            )
+
+        if self.print_config:
+            print(
+                "node_num:",
+                node_num,
+                "use_hierarchical_allreduce:",
+                self._strategy.use_hierarchical_allreduce,
+                "nccl_comm_num:",
+                self._strategy.nccl_comm_num,
+                "FLAGS_sync_nccl_allreduce:",
+                sync_allreduce,
+            )
+
+        self._transpile(startup_program, main_program)
+
+        if self._strategy.mode == "collective":
+            return main_program
+
+        self._strategy.num_trainers = fleet.worker_num()
+        self._strategy.trainer_id = fleet.worker_index()
+        self._strategy.trainers_endpoints = fleet.worker_endpoints()
+        self._strategy.enable_backward_optimizer_op_deps = True
+
+        comm_opt = RawProgramOptimizer(self._optimizer)
+        comm_opt.fuse_all_reduce_ops = True
+        comm_opt.fuse_grad_size_in_num = True
+        comm_opt.endpoints = self._strategy.trainers_endpoints
+        comm_opt.current_endpoint = comm_opt.endpoints[fleet.worker_index()]
+        comm_opt.rank = fleet.worker_index()
+        comm_opt.nranks = fleet.worker_num()
+        comm_opt.main_program = main_program
+        if comm_opt.nranks > 1:
+            comm_opt._transpile_main_program(self._loss)
+
+        self._compiled_program = CompiledProgram(
+            comm_opt.main_program, build_strategy=self._strategy
+        )
+
+        return self._compiled_program
+
+    def raiseOptimizeError(self, strategy_name, optimize_name):
+        raise ValueError(
+            f"can not use {optimize_name} when you set DistStrategy.{strategy_name} "
+            "as True"
+        )
+
+    def minimize(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        """
+        minimize a program through loss
+        Args:
+            loss (Variable|Variable List): loss variable or loss variable list to run optimization.
+            startup_program (Program): startup_program for initializing parameters
+                in `parameter_list`.
+            parameter_list (list): list of Variables to update.
+            no_grad_set (set|None): set of Variables should be ignored.
+        Returns:
+            tuple: (optimize_ops, params_grads) which are, list of operators appended;
+            and list of (param, grad) Variables pair for optimization.
+        Note that in parameter server mode, a worker will not get anything about optimize_os
+        Because optimizer algorithms run on pserver side. We will make this usable in pserver
+        process, but currently the optimization part is written into Fleet(). A user does not
+        need to care about how to startup a pserver node.
+        """
+
+        # check optimizer conflicts
+        if self._forward_recompute:
+            if self._recompute_checkpoints == []:
+                raise ValueError(
+                    "please set strategy.recompute_checkpoints"
+                    "when set strategy.forward_recompute as True"
+                )
+            if self._optimizer.__class__.__name__ in [
+                "RecomputeOptimizer",
+                "OptimizerWithMixedPrecision",
+            ]:
+                self.raiseOptimizeError(
+                    "forward_recompute", self._optimizer.__class__.__name__
+                )
+
+            self._optimizer = paddle.incubate.optimizer.RecomputeOptimizer(
+                self._optimizer
+            )
+            self._optimizer._set_checkpoints(self._recompute_checkpoints)
+
+        if self._use_amp:
+            if self._optimizer.__class__.__name__ in [
+                "OptimizerWithMixedPrecision",
+                "DGCMomentumOptimizer",
+            ]:
+                self.raiseOptimizeError(
+                    "mixed_precision", self._optimizer.__class__.__name__
+                )
+            self._optimizer = paddle.static.amp.decorate(
+                self._optimizer,
+                init_loss_scaling=self._amp_loss_scaling,
+                use_dynamic_loss_scaling=True,
+            )
+
+        main_program = loss.block.program
+        if startup_program is None:
+            startup_program = base.default_startup_program()
+        fleet.startup_program = startup_program
+
+        self._loss = loss
+
+        self._check_collective_mode(
+            main_program, self._optimizer, self._strategy
+        )
+
+        optimize_ops, param_grads = self._optimizer.minimize(
+            loss, startup_program, parameter_list, no_grad_set=no_grad_set
+        )
+
+        fleet._origin_program = main_program.clone(for_test=False)
+        fleet._transpiled_program = main_program
+        fleet.main_program = self._try_to_compile(startup_program, main_program)
+
+        return optimize_ops, param_grads

python/paddle/incubate/distributed/fleet/parameter_server/__init__.py

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

python/paddle/incubate/distributed/fleet/parameter_server/distribute_transpiler/__init__.py

@@ -0,0 +1,956 @@
+# 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.
+"""
+Convert the static program to distributed data-parallelism programs.
+"""
+
+import os
+import sys
+
+import paddle
+from paddle.base.compiler import CompiledProgram
+from paddle.distributed.fleet.base.private_helper_function import (
+    wait_server_ready,
+)
+from paddle.distributed.transpiler.distribute_transpiler import (
+    DistributeTranspilerConfig,
+)
+from paddle.framework import core
+from paddle.incubate.distributed.fleet.base import (
+    DistributedOptimizer,
+    Fleet,
+    Mode,
+)
+from paddle.incubate.distributed.fleet.parameter_server import version
+from paddle.incubate.distributed.fleet.parameter_server.distribute_transpiler.distributed_strategy import (
+    AsyncStrategy,
+    DistributedStrategy,
+    GeoStrategy,
+    HalfAsyncStrategy,
+    StrategyFactory,
+    SyncStrategy,
+    TrainerRuntimeConfig,  # noqa: F401
+)
+from paddle.incubate.distributed.fleet.parameter_server.ir import (
+    pserver_pass as server,
+    public,
+    trainer_pass as worker,
+)
+from paddle.incubate.distributed.fleet.parameter_server.ir.public import (
+    _get_lr_ops,
+    _has_global_step,
+    get_sparse_tablenames,
+)
+from paddle.incubate.distributed.fleet.parameter_server.mode import PSMode
+from paddle.incubate.distributed.fleet.parameter_server.pslib.optimizer_factory import (
+    DistributedAdam,  # noqa: F401
+)
+from paddle.incubate.distributed.fleet.role_maker import MPISymmetricRoleMaker
+from paddle.static import (
+    Executor,
+    Program,
+    default_main_program,
+    default_startup_program,
+)
+
+
+class FleetTranspiler(Fleet):
+    """
+    A subclass for compatibility with distributed.transpiler.DistributeTranspiler.
+    """
+
+    def __init__(self):
+        super().__init__(Mode.TRANSPILER)
+
+        self._inner_mode = None
+
+        if version.is_transpiler():
+            self._inner_mode = PSMode.TRANSPILER
+        else:
+            self._inner_mode = PSMode.PSLIB
+
+        self._strategy = None
+        self._transpiler = None
+        self._origin_main_program = None
+        self._origin_startup_program = None
+        self._communicator = None
+        self.startup_program = None
+        self.main_program = None
+
+        self._opt_info = None
+        self._local_ip = 0
+        self._fleet_ptr = None
+        self._main_programs = []
+        self._scopes = []
+        self._client2client_request_timeout_ms = 500000
+        self._client2client_connect_timeout_ms = 10000
+        self._client2client_max_retry = 3
+
+    def init(self, role_maker=None):
+        if role_maker is None:
+            role_maker = MPISymmetricRoleMaker()
+        super().init(role_maker)
+        if self._fleet_ptr is None:
+            self._fleet_ptr = core.Fleet()
+
+    def _init_transpiler_worker(self):
+        """
+        `init_worker` has many many functions to do before training,
+        first, wait for all parameter servers launch completely.
+        second, run executor to initialize startup program
+        third, wait for all worker initialize completely.
+
+        Returns:
+            None
+        """
+
+        def sync_strategy_envs():
+            kwargs = {}
+            kwargs["pserver_endpoints"] = (
+                self._role_maker.get_pserver_endpoints()
+            )
+            kwargs["trainer_id"] = self._role_maker.worker_index()
+            return kwargs
+
+        def geo_strategy_envs():
+            def get_sparse_attrs():
+                opt_init_map = {}
+                opt_init_map["gaussian_random"] = ["seed", "mean", "std"]
+                opt_init_map["fill_constant"] = ["value"]
+                opt_init_map["uniform_random"] = ["seed", "min", "max"]
+                opt_init_map["truncated_gaussian_random"] = [
+                    "seed",
+                    "mean",
+                    "std",
+                ]
+
+                dist_varnames = get_sparse_tablenames(
+                    self._origin_main_program, True
+                )
+                sparse_varnames = get_sparse_tablenames(
+                    self._origin_main_program, False
+                )
+
+                if len(dist_varnames) != 0:
+                    raise ValueError(
+                        "GeoStrategy can not support large scale embedding now, please use paddle.static.nn.embedding"
+                    )
+
+                init_attrs = []
+                for value_name in sparse_varnames:
+                    value_var = self._origin_main_program.global_block().vars[
+                        value_name
+                    ]
+                    value_attr = [
+                        value_name,
+                        ",".join([str(dim) for dim in value_var.shape]),
+                    ]
+                    for op in self._origin_startup_program.global_block().ops:
+                        if (
+                            op.type in opt_init_map.keys()
+                            and value_name == op.output("Out")[0]
+                        ):
+                            init_attr = [op.type]
+                            for attr in opt_init_map[op.type]:
+                                init_attr.append(str(op.attr(attr)))
+                            value_attr.append("&".join(init_attr))
+                            init_attrs.append(":".join(value_attr))
+                            break
+                return "#".join(init_attrs)
+
+            kwargs = {}
+            kwargs["trainers"] = self.worker_num()
+            kwargs["sparse_attrs"] = get_sparse_attrs()
+            return kwargs
+
+        # if MPISymmetricRoleMaker is defined
+        # we suppose a user wants to submit job on mpi cluster
+
+        if isinstance(self._role_maker, MPISymmetricRoleMaker):
+            # check whether server has been initialized
+            wait_server_ready(self.server_endpoints(to_string=False))
+
+        trainer_config = self._strategy.get_trainer_runtime_config()
+
+        print(trainer_config)
+
+        lrs = _has_global_step(_get_lr_ops(self._origin_main_program))
+
+        if lrs > 0:
+            kwargs = {"need_global_step": "1"}
+        else:
+            kwargs = {"need_global_step": "0"}
+
+        if isinstance(self._strategy, GeoStrategy):
+            geo_kwargs = geo_strategy_envs()
+            kwargs.update(geo_kwargs)
+        if isinstance(self._strategy, SyncStrategy):
+            sync_kwargs = sync_strategy_envs()
+            kwargs.update(sync_kwargs)
+
+        kwargs = kwargs if kwargs else None
+
+        send_ctx = fleet.compiled_config.get_communicator_send_context()
+
+        if self.compiled_config.is_geo_mode():
+            recv_ctx = fleet.compiled_config.get_communicator_recv_context(
+                recv_type=4
+            )
+        else:
+            recv_ctx = fleet.compiled_config.get_communicator_recv_context(
+                recv_type=1
+            )
+
+        from paddle.distributed.communicator import Communicator
+
+        self._communicator = Communicator(
+            trainer_config.mode, kwargs, trainer_config.get_communicator_flags()
+        )
+
+        self._communicator.init_with_ctx(send_ctx, recv_ctx)
+
+        if not self._communicator.is_running():
+            self._communicator.start()
+        else:
+            raise ValueError(
+                "Communicator can only be inited once, please check"
+            )
+
+    def init_worker(self):
+        """
+        `init_worker` has many many functions to do before training,
+        first, wait for all parameter servers launch completely.
+        second, run executor to initialize startup program
+        third, wait for all worker initialize completely.
+
+        Returns:
+            None
+        """
+        if self._inner_mode == PSMode.TRANSPILER:
+            self._init_transpiler_worker()
+        else:
+            raise NotImplementedError("add implement later")
+
+    def _init_transpiler_server(self, model_dir=None):
+        if not self.startup_program:
+            raise ValueError(
+                "startup_program is None, need invoke DistributedOptimizer.minimize first"
+            )
+
+        self._executor.run(self.startup_program)
+
+        if model_dir:
+            if not os.path.isdir(model_dir):
+                raise ValueError(f"There is no directory named '{model_dir}'")
+
+            sparse_varnames = self.compiled_config.get_sparse_varname_on_ps(
+                True
+            )
+            distributed_varnames = (
+                self.compiled_config.get_sparse_varname_on_ps(False)
+            )
+
+            remaining_vars = list(
+                filter(
+                    FleetTranspiler.__exclude_vars(
+                        sparse_varnames + distributed_varnames
+                    ),
+                    self.main_program.list_vars(),
+                )
+            )
+
+            paddle.static.load_vars(
+                self._executor,
+                main_program=self.main_program,
+                dirname=model_dir,
+                vars=remaining_vars,
+            )
+
+            self._load_sparse_params(
+                dirname=model_dir, varnames=sparse_varnames
+            )
+
+            # todo(tangwei12) load distributed vars
+            # self._load_sparse_params(dirname=model_dir, varnames=distributed_varnames)
+
+    def init_server(self, model_dir=None, **kwargs):
+        """
+        `init_server` has many many functions to do before start pserver,
+        first, run executor to initialize startup program,
+        second, if the `model_dir` is not empty, it will load parameters from it for increment training.
+
+        Args:
+            model_dir(str): The directory path.
+
+        Returns:
+            None
+        """
+
+        if self._inner_mode == PSMode.TRANSPILER:
+            self._init_transpiler_server(model_dir)
+        else:
+            raise NotImplementedError("add implement later")
+
+    def run_server(self):
+        """
+        `run_server` execute executor to start pserver main program.
+
+        Returns:
+            None
+        """
+
+        if self._inner_mode == PSMode.TRANSPILER:
+            if not self.main_program:
+                raise ValueError(
+                    "main_program is None, need invoke DistributedOptimizer.minimize first"
+                )
+
+            self._executor.run(self.main_program)
+        else:
+            raise NotImplementedError("add implement later")
+
+    def stop_worker(self):
+        """
+        Close this executor.
+
+        For the distributed training, this method would free the resource on PServers related to
+        the current Trainer.
+
+        Returns:
+            None
+        """
+
+        if self._inner_mode == PSMode.TRANSPILER:
+            self._communicator.stop()
+            if isinstance(self._role_maker, MPISymmetricRoleMaker):
+                self._role_maker._finalize()
+            self._executor.close()
+        else:
+            raise NotImplementedError("add implement later")
+
+    def distributed_optimizer(self, 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(DistributeTranspilerConfig): Extra properties for distributed optimizer.
+
+        Returns:
+            TranspilerOptimizer: subclass of DistributedOptimizer.
+        """
+
+        if not isinstance(optimizer, paddle.optimizer.Optimizer):
+            raise ValueError("optimizer must be an instance of Optimizer")
+        if not self._is_initialized:
+            raise ValueError(
+                "fleet.init(role) to initialize before optimizer.minimize(loss)"
+            )
+
+        if not strategy:
+            _strategy = StrategyFactory.create_async_strategy()
+
+        if isinstance(strategy, DistributedStrategy):
+            _strategy = strategy
+        elif isinstance(strategy, DistributeTranspilerConfig):
+            if strategy.sync_mode:
+                _strategy = SyncStrategy()
+            else:
+                if strategy.runtime_split_send_recv:
+                    if strategy.geo_sgd_mode:
+                        _strategy = GeoStrategy(strategy.geo_sgd_need_push_nums)
+                    elif strategy.half_async:
+                        _strategy = HalfAsyncStrategy()
+                    else:
+                        _strategy = AsyncStrategy()
+                else:
+                    _strategy = HalfAsyncStrategy()
+                    # for half_async compatibility
+                    strategy.half_async = True
+                    strategy.runtime_split_send_recv = True
+            _strategy.set_program_config(strategy)
+        elif isinstance(strategy, dict):
+            if self._inner_mode != PSMode.PSLIB:
+                raise TypeError("Dict strategy can only be used at PSLIB Mode")
+
+            _strategy = StrategyFactory.create_async_strategy()
+            _strategy.set_pslib_runtime_config(strategy)
+        else:
+            raise TypeError(
+                "strategy must be an instance of DistributeTranspilerConfig, DistributedStrategy"
+            )
+
+        self._strategy = _strategy
+        self._optimizer = ParameterServerOptimizer(optimizer, _strategy)
+        return self._optimizer
+
+    def save_inference_model(
+        self,
+        executor,
+        dirname,
+        feeded_var_names,
+        target_vars,
+        main_program=None,
+        export_for_deployment=True,
+        legacy_format=False,
+    ):
+        """
+        Prune the given `main_program` to build a new program especially for inference,
+        and then save it and all related parameters to given `dirname` by the `executor`.
+        """
+
+        if self._inner_mode == PSMode.PSLIB:
+            raise NotImplementedError("add implement later")
+
+        if not isinstance(executor, Executor):
+            raise TypeError(
+                "in fleet.save_inference_model() function, executor must be as Executor type"
+            )
+
+        # Todo(MrChengmo): support recv&save GPU-Kernel for ps-gpu model save
+        if not isinstance(executor.place, paddle.CPUPlace):
+            save_executor = Executor(paddle.CPUPlace())
+        else:
+            save_executor = executor
+
+        if main_program is not None:
+            if isinstance(main_program, CompiledProgram):
+                raise TypeError(
+                    "in fleet.save_inference_model() function, main_program must be as Program type, CompiledProgram is not allowed"
+                )
+            paddle.static.save_inference_model(
+                dirname,
+                feeded_var_names,
+                target_vars,
+                executor,
+                main_program,
+                None,
+                None,
+                export_for_deployment,
+                legacy_format=legacy_format,
+            )
+        else:
+            paddle.static.save_inference_model(
+                dirname,
+                feeded_var_names,
+                target_vars,
+                executor,
+                self._origin_main_program,
+                None,
+                None,
+                export_for_deployment,
+                True,
+                legacy_format=legacy_format,
+            )
+
+            model_basename = "__model__"
+            model_filename = os.path.join(dirname, model_basename)
+
+            with open(model_filename, "rb") as f:
+                program_desc_str = f.read()
+
+            program = Program.parse_from_string(program_desc_str)
+            program._copy_dist_param_info_from(self.main_program)
+            self.save_persistables(executor, dirname, program)
+
+    def _load_sparse_params(self, dirname, varnames):
+        from paddle.distributed.communicator import LargeScaleKV
+
+        scale_kv = LargeScaleKV()
+        for varname in varnames:
+            origin_varname, _, _ = public._get_varname_parts(varname)
+            sparse_dir = os.path.join(dirname, origin_varname, varname)
+            scale_kv.load(varname, sparse_dir)
+
+    def _get_optimizer_status(self, op, param_name):
+        supported_opts = [
+            "sgd",
+            "adam",
+            "adagrad",
+            "adamax",
+            "momentum",
+            "lars_momentum",
+            "rmsprop",
+            "decayed_adagrad",
+            "ftrl",
+        ]
+
+        reshaped_val_map = {}
+        reshaped_val_map["sgd"] = []
+        reshaped_val_map["adam"] = ["moment1_0", "moment2_0"]
+        reshaped_val_map["adagrad"] = ["moment_0"]
+        reshaped_val_map["adamax"] = ["moment_0", "inf_norm_0"]
+        reshaped_val_map["momentum"] = ["velocity_0"]
+        reshaped_val_map["lars_momentum"] = ["velocity_0"]
+        reshaped_val_map["rmsprop"] = [
+            "momentum_0",
+            "mean_square_0",
+            "mean_grad_0",
+        ]
+        reshaped_val_map["decayed_adagrad"] = ["moment_0"]
+        reshaped_val_map["ftrl"] = ["squared_0", "linear_0"]
+
+        orishaped_val_map = {}
+        orishaped_val_map["adam"] = ["beta1_pow_acc_0", "beta2_pow_acc_0"]
+        orishaped_val_map["adamax"] = ["beta1_pow_acc_0"]
+
+        if op not in supported_opts:
+            raise ValueError(
+                f"fleet can not support optimizer: {op}, only this can be supported: {supported_opts}"
+            )
+
+        reshaped_names = [
+            param_name + "_" + val for val in reshaped_val_map[op]
+        ]
+
+        if op not in orishaped_val_map:
+            origin_names = []
+        else:
+            origin_names = [
+                param_name + "_" + val for val in orishaped_val_map[op]
+            ]
+        return reshaped_names, origin_names
+
+    def _get_optimizer_op(self, param_name):
+        opts = public._get_optimize_ops(self._origin_main_program)
+        for op in opts:
+            if (
+                "Param" in op.input_names
+                and "LearningRate" in op.input_names
+                and op.input("Param")[0] == param_name
+            ):
+                return op
+
+    def _save_dense_params(self, executor, dirname, context, main_program):
+        self._communicator.recv()
+
+        prog = Program()
+        block = prog.global_block()
+        local_vars = []
+
+        for name, var_ctx in context.items():
+            if len(var_ctx.origin_varnames()) != 1:
+                raise ValueError("Dense can not support split now.")
+
+            varname = var_ctx.origin_varnames()[0]
+            local_vars.append(varname)
+
+            optimizer = self._get_optimizer_op(varname)
+            reshaped_varnames, origin_varnames = self._get_optimizer_status(
+                optimizer.type, varname
+            )
+
+            for var_name in [varname, *reshaped_varnames, *origin_varnames]:
+                var = self._origin_main_program.global_block().vars[var_name]
+                block.append_op(
+                    type='recv_save',
+                    attrs={
+                        "trainer_id": self._role_maker.worker_index(),
+                        "shape": var.shape,
+                        "slice_shapes": [",".join([str(i) for i in var.shape])],
+                        "slice_varnames": [var.name],
+                        "remote_varnames": [var.name],
+                        "is_sparse": False,
+                        "endpoints": var_ctx.split_endpoints(),
+                        "file_path": os.path.join(dirname, var.name),
+                    },
+                )
+
+        executor.run(prog)
+        return local_vars
+
+    def _save_sparse_params(self, executor, dirname, context, main_program):
+        prog = Program()
+        block = prog.global_block()
+        local_vars = []
+
+        for name, var_ctx in context.items():
+            if len(var_ctx.origin_varnames()) != 1:
+                raise ValueError("Dense can not support split now.")
+
+            varname = var_ctx.origin_varnames()[0]
+            local_vars.append(varname)
+
+            optimizer = self._get_optimizer_op(varname)
+            reshaped_varnames, origin_varnames = self._get_optimizer_status(
+                optimizer.type, varname
+            )
+
+            var = self._origin_main_program.global_block().vars[varname]
+            slice_shapes = []
+            dims1 = ",".join([str(i) for i in var.shape[1:]])
+
+            for section in var_ctx.sections():
+                slice_shapes.append(str(section) + dims1)
+
+            block.append_op(
+                type='recv_save',
+                attrs={
+                    "trainer_id": self._role_maker.worker_index(),
+                    "shape": var.shape,
+                    "slice_shapes": slice_shapes,
+                    "slice_varnames": var_ctx.split_varnames(),
+                    "remote_varnames": var_ctx.split_varnames(),
+                    "is_sparse": True,
+                    "endpoints": var_ctx.split_endpoints(),
+                    "pserver_num": len(
+                        self._role_maker.get_pserver_endpoints()
+                    ),
+                    "file_path": os.path.join(dirname, var.name),
+                },
+            )
+
+            for reshaped_varname in reshaped_varnames:
+                var = self._origin_main_program.global_block().vars[
+                    reshaped_varname
+                ]
+
+                slice_varnames = []
+                remote_varnames = []
+                for i in range(len(var_ctx.split_varnames())):
+                    slice_varnames.append(f"{reshaped_varname}.block{i}")
+                    remote_varnames.append(reshaped_varname)
+
+                block.append_op(
+                    type='recv_save',
+                    attrs={
+                        "trainer_id": self._role_maker.worker_index(),
+                        "shape": var.shape,
+                        "slice_shapes": slice_shapes,
+                        "slice_varnames": slice_varnames,
+                        "remote_varnames": remote_varnames,
+                        "is_sparse": True,
+                        "endpoints": var_ctx.split_endpoints(),
+                        "pserver_num": len(
+                            self._role_maker.get_pserver_endpoints()
+                        ),
+                        "file_path": os.path.join(dirname, var.name),
+                    },
+                )
+
+            for origin_varname in origin_varnames:
+                var = self._origin_main_program.global_block().vars[
+                    origin_varname
+                ]
+
+                block.append_op(
+                    type='recv_save',
+                    attrs={
+                        "trainer_id": self._role_maker.worker_index(),
+                        "shape": var.shape,
+                        "slice_shapes": [",".join([str(i) for i in var.shape])],
+                        "slice_varnames": [origin_varname],
+                        "remote_varnames": [origin_varname],
+                        "is_sparse": False,
+                        "endpoints": var_ctx.split_endpoints()[:1],
+                        "file_path": os.path.join(dirname, var.name),
+                    },
+                )
+        executor.run(prog)
+        return context.keys()
+
+    def _save_distributed_params(
+        self, executor, dirname, context, main_program
+    ):
+        prog = Program()
+        block = prog.global_block()
+
+        for name, var_ctx in context.items():
+            block.append_op(
+                type='checkpoint_notify',
+                attrs={
+                    "varname": name,
+                    "is_slice": True,
+                    "slice_varnames": var_ctx.split_varnames(),
+                    "remote_varnames": var_ctx.split_varnames(),
+                    "endpoints": var_ctx.split_endpoints(),
+                    "dirname": dirname,
+                },
+            )
+
+        executor.run(prog)
+        return context.keys()
+
+    def _save_distributed_persistables(self, executor, dirname, main_program):
+        dense_ctx = fleet.compiled_config.get_communicator_recv_context(
+            recv_type=1
+        )
+
+        sparse_ctx = fleet.compiled_config.get_communicator_recv_context(
+            recv_type=2
+        )
+
+        distributed_ctx = fleet.compiled_config.get_communicator_recv_context(
+            recv_type=3
+        )
+
+        recv_dense_varnames = self._save_dense_params(
+            executor, dirname, dense_ctx, main_program
+        )
+
+        recv_sparse_varnames = self._save_sparse_params(
+            executor, dirname, sparse_ctx, main_program
+        )
+
+        recv_distributed_varnames = self._save_distributed_params(
+            executor, dirname, distributed_ctx, main_program
+        )
+
+        saved_varnames = (
+            recv_dense_varnames
+            + list(recv_sparse_varnames)
+            + list(recv_distributed_varnames)
+        )
+
+        remaining_vars = list(
+            filter(
+                FleetTranspiler.__exclude_vars(saved_varnames),
+                main_program.list_vars(),
+            )
+        )
+
+        paddle.static.save_vars(
+            executor,
+            main_program=main_program,
+            dirname=dirname,
+            vars=remaining_vars,
+        )
+
+    def save_persistables(self, executor, dirname, main_program=None, **kwargs):
+        """
+                This function filters out all variables with `persistable==True` from the
+                give `main_program` and then saves these variables to the folder `dirname`
+                or file `filename`.
+
+                The `dirname` is used to specify the folder where persistable variables
+                are going to be saved. If you would like to save variables in separate
+                files, set `filename` None;
+        if you would like to save all variables in a
+                single file, use `filename` to specify the file name.
+        """
+
+        if self._inner_mode == PSMode.PSLIB:
+            raise NotImplementedError("add implement later")
+
+        if not isinstance(executor, Executor):
+            raise TypeError(
+                "in fleet.save_persistables() function, executor must be as Executor type"
+            )
+        # Todo(MrChengmo): support recv&save GPU-Kernel for ps-gpu model save
+        if not isinstance(executor.place, paddle.CPUPlace):
+            save_executor = Executor(paddle.CPUPlace())
+        else:
+            save_executor = executor
+
+        if main_program is None:
+            main_program = self.main_program
+
+        if isinstance(main_program, CompiledProgram):
+            raise TypeError(
+                "in fleet.save_persistables() function, main_program must be as Program type, CompiledProgram is not allowed"
+            )
+
+        self._save_distributed_persistables(
+            save_executor, dirname, main_program
+        )
+
+    @staticmethod
+    def __exclude_vars(exclude_var_names=[]):
+        def is_valid(var):
+            if var.name in exclude_var_names:
+                return False
+
+            origin_varname, _, _ = public._get_varname_parts(var.name)
+            if origin_varname.endswith("@GRAD"):
+                return False
+
+            if origin_varname == "learning_rate_0":
+                return False
+
+            if (
+                var.desc.type() == core.VarDesc.VarType.FEED_MINIBATCH
+                or var.desc.type() == core.VarDesc.VarType.FETCH_LIST
+                or var.desc.type() == core.VarDesc.VarType.READER
+            ):
+                return False
+            return var.persistable
+
+        return is_valid
+
+
+# fleet is a global instance for parameter server.
+fleet = FleetTranspiler()
+
+
+class ParameterServerOptimizer(DistributedOptimizer):
+    """
+    DistributedOptimizer is a wrapper for paddle.base.optimizer
+    A user should pass a paddle.base.optimizer to DistributedOptimizer
+    minimize() function is implemented.
+    DistributedOptimizer is the starting point for a user who wants to
+    run distributed training. The optimized information will be stored in
+    Fleet() instance who holds the global information about current distributed
+    training.
+
+    Args:
+        optimizer(Optimizer): subclass of Optimizer.
+        strategy(DistributeTranspilerConfig): instance of DistributeTranspilerConfig.
+
+    Returns:
+        None
+    """
+
+    def __init__(self, optimizer, strategy, mode=PSMode.TRANSPILER):
+        super().__init__(optimizer, strategy)
+        self._mode = mode
+        if self._mode == PSMode.PSLIB:
+            self._optimizer_name = f"Distributed{optimizer.type.capitalize()}"
+            if optimizer.type != "adam":
+                print(
+                    "Currently, distributed optimizer only support Adam"
+                    "Will config built-in adam for you."
+                    "We will support more functions in DistributedOptimizer",
+                    sys.stderr,
+                )
+                self._optimizer_name = "DistributedAdam"
+
+            self._optimizer = globals()[self._optimizer_name](optimizer)
+        else:
+            self._optimizer = optimizer
+
+        self._window = 1
+        self.type = "downpour"
+        self.data_norm_name = [
+            ".batch_size",
+            ".batch_square_sum",
+            ".batch_sum",
+            ".batch_size@GRAD",
+            ".batch_square_sum@GRAD",
+            ".batch_sum@GRAD",
+        ]
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        raise NotImplementedError
+
+    def apply_gradients(self, params_grads):
+        raise NotImplementedError
+
+    def _build_trainer_programs(self, compiled_config):
+        _main = fleet._origin_main_program.clone()
+        _startup = fleet._origin_startup_program.clone()
+
+        if not compiled_config.is_geo_mode():
+            # for main program
+            _main = worker.delete_optimizer_pass(_main, compiled_config)
+            _main = worker.distributed_ops_pass(_main, compiled_config)
+            _main = worker.append_send_ops_pass(_main, compiled_config)
+
+            # for startup program
+            _startup = worker.fake_init_ops_pass(_startup, compiled_config)
+            _startup = worker.init_from_server_pass(_startup, compiled_config)
+            _startup = worker.delete_extra_optimizes_pass(
+                _startup, compiled_config
+            )
+        else:
+            _main = worker.append_send_ops_pass(_main, compiled_config)
+            _startup = _startup
+
+        return _main, _startup
+
+    def _build_pserver_programs(self, compiled_config):
+        _main = paddle.static.Program()
+        _startup = paddle.static.Program()
+
+        if not compiled_config.is_geo_mode():
+            _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)
+            _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
+            )
+            _startup = server.delete_unused_in_startup_pass(
+                _startup, _main, compiled_config
+            )
+
+        return _main, _startup
+
+    def minimize(
+        self,
+        losses,
+        scopes=None,
+        startup_programs=None,
+        parameter_list=None,
+        no_grad_set=None,
+    ):
+        if isinstance(losses, list):
+            raise ValueError("need implement later")
+
+        self._optimizer.minimize(
+            losses, startup_programs, parameter_list, no_grad_set
+        )
+
+        fleet._origin_main_program = default_main_program().clone()
+        fleet._origin_startup_program = default_startup_program().clone()
+
+        compiled_config = public.CompileTimeStrategy(
+            fleet._origin_main_program,
+            fleet._origin_startup_program,
+            self._strategy,
+            fleet._role_maker,
+        )
+
+        fleet.compiled_config = compiled_config
+        fleet.main_program, fleet.startup_program = (
+            self._build_trainer_programs(compiled_config)
+            if fleet.is_worker()
+            else self._build_pserver_programs(compiled_config)
+        )

python/paddle/incubate/distributed/fleet/parameter_server/distribute_transpiler/distributed_strategy.py

@@ -0,0 +1,421 @@
+#   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.
+
+__all__ = []
+
+import os
+
+from paddle import base
+from paddle.distributed.transpiler.distribute_transpiler import (
+    DistributeTranspilerConfig,
+    ServerRuntimeConfig,
+)
+from paddle.incubate.distributed.fleet.parameter_server.mode import (
+    DistributedMode,
+)
+
+
+class TrainerRuntimeConfig:
+    def __init__(self):
+        self.mode = None
+        num_threads = os.getenv("CPU_NUM", "1")
+
+        self.runtime_configs = {}
+        self.runtime_configs['communicator_max_merge_var_num'] = os.getenv(
+            "FLAGS_communicator_max_merge_var_num", num_threads
+        )
+        self.runtime_configs['communicator_send_queue_size'] = os.getenv(
+            "FLAGS_communicator_send_queue_size", num_threads
+        )
+        self.runtime_configs['communicator_independent_recv_thread'] = (
+            os.getenv("FLAGS_communicator_independent_recv_thread", "1")
+        )
+        self.runtime_configs['communicator_min_send_grad_num_before_recv'] = (
+            os.getenv(
+                "FLAGS_communicator_min_send_grad_num_before_recv", num_threads
+            )
+        )
+        self.runtime_configs['communicator_thread_pool_size'] = os.getenv(
+            "FLAGS_communicator_thread_pool_size", "5"
+        )
+        self.runtime_configs['communicator_send_wait_times'] = os.getenv(
+            "FLAGS_communicator_send_wait_times", "5"
+        )
+        self.runtime_configs['communicator_is_sgd_optimizer'] = os.getenv(
+            "FLAGS_communicator_is_sgd_optimizer", "1"
+        )
+
+        # not used
+        self.runtime_configs['rpc_deadline'] = os.getenv(
+            "FLAGS_rpc_deadline", "180000"
+        )
+        self.runtime_configs['rpc_retry_times'] = os.getenv(
+            "FLAGS_rpc_retry_times", "3"
+        )
+
+    def get_communicator_flags(self):
+        need_keys = []
+        num_threads = os.getenv("CPU_NUM", "1")
+        mode_str = ""
+        if self.mode is None or self.mode == DistributedMode.ASYNC:
+            need_keys = self.runtime_configs.keys()
+            mode_str = "async"
+        elif (
+            self.mode == DistributedMode.SYNC
+            or self.mode == DistributedMode.HALF_ASYNC
+        ):
+            mode_str = "sync or half_async"
+            need_keys = [
+                'communicator_max_merge_var_num',
+                'communicator_send_wait_times',
+                'communicator_thread_pool_size',
+                'communicator_send_queue_size',
+            ]
+        elif self.mode == DistributedMode.GEO:
+            mode_str = "GEO"
+            need_keys = [
+                'communicator_thread_pool_size',
+                'communicator_send_wait_times',
+                'communicator_max_merge_var_num',
+                'communicator_send_queue_size',
+            ]
+        else:
+            raise ValueError("Unsupported Mode")
+
+        if (
+            self.mode == DistributedMode.SYNC
+            or self.mode == DistributedMode.HALF_ASYNC
+        ):
+            max_merge_var_num = self.runtime_configs[
+                'communicator_max_merge_var_num'
+            ]
+            send_queue_size = self.runtime_configs[
+                'communicator_send_queue_size'
+            ]
+            if max_merge_var_num != num_threads:
+                print(
+                    f'WARNING: In {mode_str} mode, communicator_max_merge_var_num '
+                    'must be equal to CPU_NUM. But received, '
+                    f'communicator_max_merge_var_num = {max_merge_var_num}, CPU_NUM = '
+                    f'{num_threads}. communicator_max_merge_var_num will be forced to {num_threads}.'
+                )
+                self.runtime_configs['communicator_max_merge_var_num'] = (
+                    num_threads
+                )
+            if send_queue_size != num_threads:
+                print(
+                    f'WARNING: In {mode_str} mode, communicator_send_queue_size '
+                    'must be equal to CPU_NUM. But received, '
+                    f'communicator_send_queue_size = {send_queue_size}, CPU_NUM = '
+                    f'{num_threads}. communicator_send_queue_size will be forced to {num_threads}.'
+                )
+                self.runtime_configs['communicator_send_queue_size'] = (
+                    num_threads
+                )
+
+        return {key: str(self.runtime_configs[key]) for key in need_keys}
+
+    def display(self, configs):
+        raw0, raw1, length = 45, 5, 50
+        h_format = "{:^45s}{:<5s}\n"
+        l_format = "{:<45s}{:<5s}\n"
+
+        border = "".join(["="] * length)
+        line = "".join(["-"] * length)
+
+        draws = ""
+        draws += border + "\n"
+        draws += h_format.format("TrainerRuntimeConfig Overview", "Value")
+        draws += line + "\n"
+
+        for k, v in configs.items():
+            draws += l_format.format(k, v)
+
+        draws += border
+
+        _str = f"\n{draws}\n"
+        return _str
+
+    def __repr__(self):
+        return self.display(self.get_communicator_flags())
+
+
+class PSLibRuntimeConfig:
+    def __init__(self):
+        self.runtime_configs = {}
+
+    def get_runtime_configs(self):
+        return self.runtime_configs
+
+
+class DistributedStrategy:
+    def __init__(self):
+        self._program_config = DistributeTranspilerConfig()
+        self._trainer_runtime_config = TrainerRuntimeConfig()
+        self._pslib_runtime_config = PSLibRuntimeConfig()
+        self._server_runtime_config = ServerRuntimeConfig()
+        num_threads = int(os.getenv("CPU_NUM", "1"))
+
+        self._build_strategy = base.BuildStrategy()
+
+        if num_threads > 1:
+            self._build_strategy.reduce_strategy = (
+                base.BuildStrategy.ReduceStrategy.Reduce
+            )
+        self.debug_opt = None
+        self.use_ps_gpu = False
+
+    def set_debug_opt(self, opt_info):
+        self.debug_opt = opt_info
+
+    def get_debug_opt(self):
+        opt_info = {}
+        if self.debug_opt is not None and isinstance(self.debug_opt, dict):
+            opt_info["dump_slot"] = bool(self.debug_opt.get("dump_slot", 0))
+            opt_info["dump_converter"] = str(
+                self.debug_opt.get("dump_converter", "")
+            )
+            opt_info["dump_fields"] = self.debug_opt.get("dump_fields", [])
+            opt_info["dump_file_num"] = self.debug_opt.get("dump_file_num", 16)
+            opt_info["dump_fields_path"] = self.debug_opt.get(
+                "dump_fields_path", ""
+            )
+            opt_info["dump_param"] = self.debug_opt.get("dump_param", [])
+        return opt_info
+
+    def get_program_config(self):
+        return self._program_config
+
+    def set_program_config(self, config):
+        if isinstance(config, DistributeTranspilerConfig):
+            self._program_config = config
+        elif isinstance(config, dict):
+            for key in config:
+                if hasattr(self._program_config, key):
+                    setattr(self._program_config, key, config[key])
+                else:
+                    raise ValueError(
+                        f"DistributeTranspilerConfig doesn't have key: {key}"
+                    )
+        else:
+            raise TypeError(
+                "program_config only accept input type: dict or DistributeTranspilerConfig"
+            )
+        self.check_program_config()
+
+    def check_program_config(self):
+        raise NotImplementedError(
+            "check_program_config must be implemented by derived class. You should use StrategyFactory to create DistributedStrategy."
+        )
+
+    def get_trainer_runtime_config(self):
+        return self._trainer_runtime_config
+
+    def set_trainer_runtime_config(self, config):
+        if isinstance(config, TrainerRuntimeConfig):
+            self._trainer_runtime_config = config
+        elif isinstance(config, dict):
+            for key, Value in config.items():
+                if key in self._trainer_runtime_config.runtime_configs:
+                    self._trainer_runtime_config.runtime_configs[key] = Value
+                else:
+                    raise ValueError(
+                        f"TrainerRuntimeConfig doesn't have key: {key}"
+                    )
+        else:
+            raise TypeError(
+                "trainer_runtime_config only accept input type: dict or TrainerRuntimeConfig"
+            )
+        self.check_trainer_runtime_config()
+
+    def check_trainer_runtime_config(self):
+        raise NotImplementedError(
+            "check_trainer_runtime_config must be implemented by derived class. You should use StrategyFactory to create DistributedStrategy."
+        )
+
+    def get_pslib_runtime_config(self):
+        return self._pslib_runtime_config
+
+    def set_pslib_runtime_config(self, config):
+        self._pslib_runtime_config.runtime_configs = config
+
+    def get_server_runtime_config(self):
+        return self._server_runtime_config
+
+    def set_server_runtime_config(self, config):
+        if isinstance(config, ServerRuntimeConfig):
+            self._server_runtime_config = config
+        elif isinstance(config, dict):
+            for key in config:
+                if hasattr(self._server_runtime_config, key):
+                    setattr(self._server_runtime_config, key, config[key])
+                else:
+                    raise ValueError(
+                        f"ServerRuntimeConfig doesn't have key: {key}"
+                    )
+        else:
+            raise TypeError(
+                "server_runtime_config only accept input type: dict or ServerRuntimeConfig"
+            )
+        self.check_server_runtime_config()
+
+    def check_server_runtime_config(self):
+        raise NotImplementedError(
+            "check_server_runtime_config must be implemented by derived class. You should use StrategyFactory to create DistributedStrategy."
+        )
+
+    def get_build_strategy(self):
+        return self._build_strategy
+
+    def set_build_strategy(self, config):
+        if isinstance(config, base.BuildStrategy):
+            self._build_strategy = config
+        elif isinstance(config, dict):
+            for key in config:
+                if hasattr(self._build_strategy, key):
+                    setattr(self._build_strategy, key, config[key])
+                else:
+                    raise ValueError(f"BuildStrategy doesn't have key: {key}")
+        else:
+            raise TypeError(
+                "build_strategy only accept input type: dict or BuildStrategy"
+            )
+        self.check_build_strategy()
+
+    def check_build_strategy(self):
+        raise NotImplementedError(
+            "check_build_strategy must be implemented by derived class. You should use StrategyFactory to create DistributedStrategy."
+        )
+
+
+class SyncStrategy(DistributedStrategy):
+    def __init__(self):
+        super().__init__()
+        self.check_program_config()
+        self.check_trainer_runtime_config()
+        self.check_server_runtime_config()
+        self.check_build_strategy()
+
+    def check_trainer_runtime_config(self):
+        self._trainer_runtime_config.mode = DistributedMode.SYNC
+
+    def check_program_config(self):
+        self._program_config.sync_mode = False
+        self._program_config.runtime_split_send_recv = True
+        self._program_config.half_async = True
+        self._program_config.completely_not_async = True
+
+    def check_server_runtime_config(self):
+        pass
+
+    def check_build_strategy(self):
+        self._build_strategy.async_mode = True
+
+
+class AsyncStrategy(DistributedStrategy):
+    def __init__(self):
+        super().__init__()
+        self.check_program_config()
+        self.check_trainer_runtime_config()
+        self.check_server_runtime_config()
+        self.check_build_strategy()
+
+    def check_trainer_runtime_config(self):
+        self._trainer_runtime_config.mode = DistributedMode.ASYNC
+
+    def check_program_config(self):
+        self._program_config.sync_mode = False
+        self._program_config.runtime_split_send_recv = True
+
+    def check_server_runtime_config(self):
+        pass
+
+    def check_build_strategy(self):
+        self._build_strategy.async_mode = True
+
+
+class HalfAsyncStrategy(DistributedStrategy):
+    def __init__(self):
+        super().__init__()
+        self.check_program_config()
+        self.check_trainer_runtime_config()
+        self.check_server_runtime_config()
+        self.check_build_strategy()
+
+    def check_trainer_runtime_config(self):
+        self._trainer_runtime_config.mode = DistributedMode.HALF_ASYNC
+
+    def check_program_config(self):
+        self._program_config.sync_mode = False
+        self._program_config.runtime_split_send_recv = True
+        self._program_config.half_async = True
+
+    def check_server_runtime_config(self):
+        pass
+
+    def check_build_strategy(self):
+        self._build_strategy.async_mode = True
+
+
+class GeoStrategy(DistributedStrategy):
+    def __init__(self, update_frequency=100):
+        super().__init__()
+        self._program_config.geo_sgd_need_push_nums = update_frequency
+        self.check_program_config()
+        self.check_trainer_runtime_config()
+        self.check_server_runtime_config()
+        self.check_build_strategy()
+
+    def check_program_config(self):
+        self._program_config.sync_mode = False
+        self._program_config.runtime_split_send_recv = True
+        self._program_config.geo_sgd_mode = True
+
+    def check_trainer_runtime_config(self):
+        self._trainer_runtime_config.mode = DistributedMode.GEO
+
+        self._trainer_runtime_config.runtime_configs[
+            'communicator_send_queue_size'
+        ] = self._program_config.geo_sgd_need_push_nums
+
+        self._trainer_runtime_config.runtime_configs[
+            'communicator_max_merge_var_num'
+        ] = self._program_config.geo_sgd_need_push_nums
+
+    def check_server_runtime_config(self):
+        pass
+
+    def check_build_strategy(self):
+        self._build_strategy.async_mode = True
+
+
+class StrategyFactory:
+    def __init_(self):
+        pass
+
+    @staticmethod
+    def create_sync_strategy():
+        return SyncStrategy()
+
+    @staticmethod
+    def create_half_async_strategy():
+        return HalfAsyncStrategy()
+
+    @staticmethod
+    def create_async_strategy():
+        return AsyncStrategy()
+
+    @staticmethod
+    def create_geo_strategy(update_frequency=100):
+        return GeoStrategy(update_frequency)

python/paddle/incubate/distributed/fleet/parameter_server/ir/__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/incubate/distributed/fleet/parameter_server/ir/heter_trainer_pass.py

@@ -0,0 +1,79 @@
+# 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 warnings
+
+import paddle
+from paddle.incubate.distributed.fleet.parameter_server.ir.trainer_pass import (
+    create_heter_program,
+    create_trainer_program,
+    find_block_joints,
+    find_heter_ops,
+    union_forward_gradient_op,
+)
+
+
+def split_heter_worker_ops_pass(program, config, stage_id, device):
+    """
+    split heter worker program from origin-program
+    1. find heter op (located on different device)
+    2. find input&output of every heter-block
+    3. create heter worker program, add listen&serv op
+    """
+    default_device = "cpu"
+    program, heter_ops, _, program_block_ops = find_heter_ops(
+        program, default_device
+    )
+    if len(heter_ops) == 0:
+        warnings.warn(
+            "Currently running in Heter Parameter Server mode, but no OP running on heterogeneous devices, Please check your code."
+        )
+        return program
+
+    program_block_ops = union_forward_gradient_op(program_block_ops)
+    block_vars_detail = find_block_joints(program, program_block_ops, heter_ops)
+    heter_program = paddle.static.Program()
+    create_heter_program(
+        program,
+        config,
+        heter_program,
+        program_block_ops,
+        heter_ops,
+        block_vars_detail,
+        device,
+        stage_id,
+    )
+    return heter_program
+
+
+def split_trainer_ops_pass(program, config, default_device="cpu"):
+    """
+    split cpu-trainer program from origin-program
+    1. find heter op (located on different device)
+    2. find input&output of every heter-block
+    3. create cpu-trainer program, add send&recv op
+    """
+    # Todo: support user define default_device (MrChengmo)
+    default_device_ = default_device
+    program, heter_ops, default_ops, program_block_ops = find_heter_ops(
+        program, default_device_
+    )
+    program_block_ops = union_forward_gradient_op(program_block_ops)
+
+    block_vars_detail = find_block_joints(program, program_block_ops, heter_ops)
+    trainer_program = program.clone()
+    create_trainer_program(
+        trainer_program, program, config, program_block_ops, block_vars_detail
+    )
+    return trainer_program

python/paddle/incubate/distributed/fleet/parameter_server/ir/ps_dispatcher.py

@@ -0,0 +1,127 @@
+#   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.
+
+
+class PSDispatcher:
+    """
+    PSDispatcher is the base class for dispatching vars
+    into different pserver instance.
+    You need to implement the `dispatch` interface.
+    """
+
+    def __init__(self, pserver_endpoints):
+        self._eps = pserver_endpoints
+        self._step = 0
+
+    @property
+    def eps(self):
+        return self._eps
+
+    def reset(self):
+        """
+        reset the step counter, set it zero.
+        """
+        self._step = 0
+
+    def dispatch(self, varlist):
+        """
+        Args:
+            varlist(list): a list of Variables
+        Returns:
+            a map of pserver endpoint -> varname
+        """
+        raise NotImplementedError("Interface has not been implemented.")
+
+
+class HashName(PSDispatcher):
+    """
+    Hash variable names to several endpoints using python
+    "hash()" function.
+
+    Args:
+        pserver_endpoints (list): list of endpoint(ip:port).
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.incubate.distributed.fleet.parameter_server.ir.ps_dispatcher import RoundRobin
+
+            >>> pserver_endpoints = ["127.0.0.1:6007", "127.0.0.1:6008"]
+            >>> vars = ["var1", "var2", "var3", "var4", "var5"]
+
+            >>> rr = HashName(pserver_endpoints)
+            >>> rr.dispatch(vars)
+
+    """
+
+    def __init__(self, pserver_endpoints):
+        super().__init__(pserver_endpoints)
+
+    def _hash_block(self, block_str, total):
+        return hash(block_str) % total
+
+    def dispatch(self, varlist):
+        """
+        use `HashName` method to dispatch variables with each parameter server.
+        Args:
+            varlist (list): a list of Variables
+
+        """
+        eplist = []
+        for var in varlist:
+            server_id = self._hash_block(var.name(), len(self._eps))
+            server_for_param = self._eps[server_id]
+            eplist.append(server_for_param)
+        return eplist
+
+
+class RoundRobin(PSDispatcher):
+    """
+    Distribute variables to several endpoints using
+    RondRobin<https://en.wikipedia.org/wiki/Round-robin_scheduling> method.
+
+    Args:
+        pserver_endpoints (list): list of endpoint(ip:port).
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.incubate.distributed.fleet.parameter_server.ir.ps_dispatcher import RoundRobin
+
+            >>> pserver_endpoints = ["127.0.0.1:6007", "127.0.0.1:6008"]
+            >>> vars = ["var1", "var2", "var3", "var4", "var5"]
+
+            >>> rr = RoundRobin(pserver_endpoints)
+            >>> rr.dispatch(vars)
+
+    """
+
+    def __init__(self, pserver_endpoints):
+        super().__init__(pserver_endpoints)
+
+    def dispatch(self, varlist):
+        """
+        use `RoundRobin` method to dispatch variables with each parameter server.
+        Args:
+            varlist (list): a list of Variables
+
+        """
+        eplist = []
+        for var in varlist:
+            server_for_param = self._eps[self._step]
+            eplist.append(server_for_param)
+            self._step += 1
+            if self._step >= len(self._eps):
+                self._step = 0
+        return eplist

python/paddle/incubate/distributed/fleet/parameter_server/ir/ufind.py

@@ -0,0 +1,64 @@
+# 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.
+
+
+class UnionFind:
+    """Union-find data structure.
+
+    Union-find is a data structure that keeps track of a set of elements partitioned
+    into a number of disjoint (non-overlapping) subsets.
+
+    Reference:
+    https://en.wikipedia.org/wiki/Disjoint-set_data_structure
+
+    Args:
+      elements(list): The initialize element list.
+    """
+
+    def __init__(self, elements=None):
+        self._parents = []  # index -> parent index
+        self._index = {}  # element -> index
+        self._curr_idx = 0
+        if not elements:
+            elements = []
+        for ele in elements:
+            self._parents.append(self._curr_idx)
+            self._index.update({ele: self._curr_idx})
+            self._curr_idx += 1
+
+    def find(self, x):
+        # Find the root index of given element x,
+        # execute the path compress while finding the root index
+        if x not in self._index:
+            return -1
+        idx = self._index[x]
+        while idx != self._parents[idx]:
+            t = self._parents[idx]
+            self._parents[idx] = self._parents[t]
+            idx = t
+        return idx
+
+    def union(self, x, y):
+        # Union two given element
+        x_root = self.find(x)
+        y_root = self.find(y)
+
+        if x_root == y_root:
+            return
+        self._parents[x_root] = y_root
+
+    def is_connected(self, x, y):
+        # If two given elements have the same root index,
+        # then they are connected.
+        return self.find(x) == self.find(y)

python/paddle/incubate/distributed/fleet/parameter_server/ir/vars_metatools.py

@@ -0,0 +1,206 @@
+# 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.
+from functools import reduce
+
+from paddle.framework import core
+from paddle.framework.io import Variable
+
+dtype_to_size = {
+    core.VarDesc.VarType.FP16: 2,
+    core.VarDesc.VarType.FP32: 4,
+    core.VarDesc.VarType.FP64: 8,
+    core.VarDesc.VarType.INT16: 2,
+    core.VarDesc.VarType.INT32: 4,
+    core.VarDesc.VarType.INT64: 8,
+    core.VarDesc.VarType.BOOL: 1,
+    core.VarDesc.VarType.UINT8: 1,
+}
+
+
+class VarBlock:
+    def __init__(self, varname, offset, size):
+        self.varname = varname
+        # NOTE: real offset is offset * size
+        self.offset = offset
+        self.size = size
+
+    def __str__(self):
+        return f"{self.varname}:{int(self.offset)}:{int(self.size)}"
+
+
+def create_var_struct(var):
+    if var.type == core.VarDesc.VarType.SELECTED_ROWS:
+        lod_level = None
+    elif var.type == core.VarDesc.VarType.DENSE_TENSOR:
+        lod_level = var.lod_level
+    else:
+        raise ValueError("can only support SELECTED_ROWS/DENSE_TENSOR now")
+
+    return VarStruct(
+        var.name, var.shape, var.dtype, var.type, lod_level, var.persistable
+    )
+
+
+class VarStruct:
+    """
+    record part properties of a Variable in python.
+    """
+
+    def __init__(self, name, shape, dtype, type, lod_level, persistable):
+        self.name = name
+        self.shape = shape
+        self.dtype = dtype
+        self.type = type
+        self.lod_level = lod_level
+        self.persistable = persistable
+        self.m_size = 1
+        self.m_size = reduce(lambda x, y: x * y, shape, 1)
+        self.m_size *= dtype_to_size[dtype]
+
+    def __str__(self):
+        return f"N: {self.name}, S: {self.shape}, D: {self.dtype}, T: {self.type}, LL: {self.lod_level}, P: {self.persistable}, M: {self.m_size}"
+
+
+class VarDistributed:
+    """
+    a class to record the var distributed on parameter servers.
+    the class will record the relationship between origin var and slice var.
+    the slice var's properties, such as type/shape/offset/endpoint.
+    """
+
+    def __init__(
+        self,
+        origin_var,
+        slice_var,
+        is_slice=None,
+        block_id=None,
+        offset=None,
+        vtype=None,
+        endpoint=None,
+    ):
+        """
+        Args:
+            origin_var(Variable|VarStruct): origin var properties
+            slice_var(Variable|VarStruct): slice var properties
+            is_slice(bool|None): slice or not, slice_var=True/False and its block size > 8192 are the judgement standard.
+            block_id(int|None): the number about the slice var.
+            offset(int|None): if the slice var is sliced, offset is the numel before the var.
+            vtype(str|None): a tag, such as Optimizer/Param/RemotePrefetch.
+            endpoint(str|None): which parameter the slice var on, such as "127.0.0.1:1001"
+        """
+
+        if isinstance(origin_var, Variable):
+            self.origin = create_var_struct(origin_var)
+        else:
+            self.origin = origin_var
+
+        if isinstance(slice_var, Variable):
+            self.slice = create_var_struct(slice_var)
+        else:
+            self.slice = slice_var
+
+        if self.equal(self.origin, self.slice):
+            self.is_slice = False
+            self.block_id = 0
+            self.offset = 0
+        else:
+            self.is_slice = True
+            self.block_id = 0
+            self.offset = 0
+
+        if is_slice is not None:
+            self.is_slice = is_slice
+        if block_id is not None:
+            self.block_id = block_id
+        if offset is not None:
+            self.offset = offset
+
+        self.vtype = vtype
+        self.endpoint = endpoint
+
+    @staticmethod
+    def equal(var1, var2):
+        """
+        the two var is equal or not.
+        Returns:
+            bool: equal will return True else False
+        """
+        assert isinstance(var1, VarStruct) and isinstance(var2, VarStruct)
+
+        return (
+            var1.name == var2.name
+            and var1.type == var2.type
+            and var1.shape == var2.shape
+            and var1.dtype == var2.dtype
+            and var1.lod_level == var2.lod_level
+            and var1.persistable == var2.persistable
+        )
+
+    def __str__(self):
+        origin_var_str = f"{self.origin.name} : base.{self.origin.type}.shape{self.origin.shape}.astype({self.origin.dtype})"
+
+        slice_var_str = (
+            f"{self.slice.name} : base.{self.slice.type}.shape{self.slice.shape}.astype({self.slice.dtype})"
+            f".slice({self.is_slice}).block({self.block_id}).offset({self.offset})"
+        )
+
+        return f"var owned: {self.vtype}, origin var: ( {origin_var_str} ), slice var: ( {slice_var_str} ), endpoint: {self.endpoint} "
+
+
+class VarsDistributed:
+    """
+    a gather about VarDistributed with many methods to find distributed vars.
+    through the class, we can get overview about the distributed parameters on parameter servers.
+    this class may centralized and convenient for developer to manage and get variable's distribute.
+    other module can also use this to find variables such io.py.
+    """
+
+    def __init__(self):
+        self.distributed_vars = []
+
+    def add_distributed_var(
+        self,
+        origin_var,
+        slice_var,
+        is_slice=None,
+        block_id=None,
+        offset=None,
+        vtype=None,
+        endpoint=None,
+    ):
+        """
+        add distributed var in this.
+
+        Args:
+            origin_var(Variable|VarStruct): origin var properties
+            slice_var(Variable|VarStruct): slice var properties
+            is_slice(bool|None): slice or not, slice_var=True/False and its block size > 8192 are the judgement standard.
+            block_id(int|None): the number about the slice var.
+            offset(int|None): if the slice var is sliced, offset is the numel before the var.
+            vtype(str|None): a tag, such as Optimizer/Param/RemotePrefetch.
+            endpoint(str|None): which parameter the slice var on, such as "127.0.0.1:1001"
+        Returns:
+            None
+        """
+        self.distributed_vars.append(
+            VarDistributed(
+                origin_var,
+                slice_var,
+                is_slice,
+                block_id,
+                offset,
+                vtype,
+                endpoint,
+            )
+        )

python/paddle/incubate/distributed/fleet/parameter_server/mode.py

@@ -0,0 +1,29 @@
+# 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.
+
+
+class PSMode:
+    """
+    There are various mode for fleet, each of them is designed for different model.
+    """
+
+    TRANSPILER = 1
+    PSLIB = 2
+
+
+class DistributedMode:
+    SYNC = 0
+    ASYNC = 1
+    HALF_ASYNC = 2
+    GEO = 3

python/paddle/incubate/distributed/fleet/parameter_server/pslib/.gitignore

@@ -0,0 +1 @@
+ps_pb2.py

python/paddle/incubate/distributed/fleet/parameter_server/pslib/node.py

@@ -0,0 +1,798 @@
+#   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
+"""Definition of Server and Worker."""
+
+# NOTE: reduce removed in functools in python3
+from functools import reduce
+
+from . import ps_pb2 as pslib
+
+
+class Server:
+    """
+    A Server basic class
+    it's a base class, does not have implementation
+    """
+
+    def __init__(self):
+        pass
+
+
+class Worker:
+    """
+    A Worker basic class.
+    it's a base class, does not have implementation
+    """
+
+    def __init__(self):
+        pass
+
+
+class DownpourServer(Server):
+    """
+    DownpourServer class is used to generate server program_desc
+    Args:
+        server: it is pslib.ServerParameter()
+    Examples:
+        server = DownpourServer()
+    """
+
+    def __init__(self):
+        self._server = pslib.ServerParameter()
+        self._server.downpour_server_param.service_param.server_class = (
+            "DownpourBrpcPsServer"
+        )
+        self._server.downpour_server_param.service_param.client_class = (
+            "DownpourBrpcPsClient"
+        )
+        self._server.downpour_server_param.service_param.service_class = (
+            "DownpourPsService"
+        )
+        self._server.downpour_server_param.service_param.start_server_port = 0
+        self._server.downpour_server_param.service_param.server_thread_num = 12
+
+    def add_sparse_table(self, table_id, strategy):
+        """
+        Args:
+            table_id(int): id of sparse params table
+            strategy(dict): the config dict.
+        Returns:
+            return None
+        """
+
+        for table in self._server.downpour_server_param.downpour_table_param:
+            if table.table_id == table_id:
+                if table.type == pslib.PS_SPARSE_TABLE:
+                    return
+                else:
+                    raise ValueError(
+                        f"expect table {table_id} type={pslib.PS_SPARSE_TABLE}, but actual type={table.type}"
+                    )
+        if strategy is None:
+            strategy = {}
+        table = self._server.downpour_server_param.downpour_table_param.add()
+        table.table_id = table_id
+        table.type = pslib.PS_SPARSE_TABLE
+
+        support_sparse_key_list = [
+            'sparse_table_class',
+            'sparse_compress_in_save',
+            'sparse_shard_num',
+            'sparse_accessor_class',
+            'sparse_learning_rate',
+            'sparse_initial_g2sum',
+            'sparse_initial_range',
+            'sparse_weight_bounds',
+            'sparse_embedx_dim',
+            'sparse_embedx_threshold',
+            'sparse_nonclk_coeff',
+            'sparse_click_coeff',
+            'sparse_base_threshold',
+            'sparse_delta_threshold',
+            'sparse_delta_keep_days',
+            'sparse_delete_after_unseen_days',
+            'sparse_show_click_decay_rate',
+            'sparse_delete_threshold',
+            'sparse_converter',
+            'sparse_deconverter',
+            'sparse_enable_cache',
+            'sparse_cache_rate',
+            'sparse_cache_file_num',
+            'sparse_beta1_decay_rate',
+            'sparse_beta2_decay_rate',
+            'sparse_ada_epsilon',
+            'sparse_optimizer',
+            'sparse_ssd_unseenday_threshold',
+            'embed_sparse_optimizer',
+            'embed_sparse_learning_rate',
+            'embed_sparse_weight_bounds',
+            'embed_sparse_initial_range',
+            'embed_sparse_initial_g2sum',
+            'embed_sparse_beta1_decay_rate',
+            'embed_sparse_beta2_decay_rate',
+            'embedx_sparse_optimizer',
+            'embedx_sparse_learning_rate',
+            'embedx_sparse_weight_bounds',
+            'embedx_sparse_initial_range',
+            'embedx_sparse_initial_g2sum',
+            'embedx_sparse_beta1_decay_rate',
+            'embedx_sparse_beta2_decay_rate',
+        ]
+
+        for key in strategy:
+            if key not in support_sparse_key_list:
+                raise ValueError(f"strategy key '{key}' not support")
+
+        support_table_class = ['DownpourSparseTable', 'DownpourSparseSSDTable']
+        if strategy.get('sparse_table_class') is not None:
+            table_class = strategy.get('sparse_table_class')
+            if table_class not in support_table_class:
+                raise ValueError(
+                    f"support sparse_table_class: [ 'DownpourSparseTable', 'DownpourSparseSSDTable'], \
+                        but actual {table_class}"
+                )
+        else:
+            table_class = 'DownpourSparseTable'
+
+        table.table_class = table_class
+
+        if (
+            table_class == 'DownpourSparseTable'
+            or table_class == 'DownpourSparseSSDTable'
+        ):
+            table.enable_sparse_table_cache = strategy.get(
+                'sparse_enable_cache', True
+            )
+            table.sparse_table_cache_rate = strategy.get(
+                'sparse_cache_rate', 0.00055
+            )
+            table.sparse_table_cache_file_num = strategy.get(
+                'sparse_cache_file_num', 16
+            )
+            table.compress_in_save = strategy.get(
+                'sparse_compress_in_save', True
+            )
+            table.shard_num = strategy.get('sparse_shard_num', 1000)
+            # DownpourFeatureValueAccessor: for ctr task, has cvm, embedding and sgd info
+            # DownpourCtrAccessor         : for ctr task, has cvm, slot, embedding and sgd info
+            # DownpourSparseValueAccessor : for general task, has embedding and sgd info
+            # DownpourCtrDoubleAccessor   : for ctr task, which show clk are in double
+            # DownpourUnitAccessor        : for ctr task, has cvm, slot, embedding and sgd info
+
+            support_accessor_class = [
+                'DownpourFeatureValueAccessor',
+                'DownpourCtrAccessor',
+                'DownpourCtrDymfAccessor',
+                'DownpourSparseValueAccessor',
+                'DownpourCtrDoubleAccessor',
+                'DownpourUnitAccessor',
+                'DownpourDoubleUnitAccessor',
+            ]
+            if strategy.get('sparse_accessor_class') is not None:
+                accessor_class = strategy.get('sparse_accessor_class')
+                if accessor_class not in support_accessor_class:
+                    raise ValueError(
+                        f"support sparse_accessor_class: ['DownpourFeatureValueAccessor', 'DownpourCtrAccessor', 'DownpourCtrDymfAccessor', \
+                        'DownpourSparseValueAccessor', 'DownpourCtrDoubleAccessor'], \
+                            but actual {accessor_class}"
+                    )
+            else:
+                accessor_class = 'DownpourCtrAccessor'
+
+            table.accessor.accessor_class = accessor_class
+
+            if (
+                accessor_class == 'DownpourFeatureValueAccessor'
+                or accessor_class == 'DownpourCtrAccessor'
+                or accessor_class == 'DownpourCtrDoubleAccessor'
+            ):
+                table.accessor.sparse_sgd_param.learning_rate = strategy.get(
+                    'sparse_learning_rate', 0.05
+                )
+                table.accessor.sparse_sgd_param.initial_g2sum = strategy.get(
+                    'sparse_initial_g2sum', 3
+                )
+                table.accessor.sparse_sgd_param.initial_range = strategy.get(
+                    'sparse_initial_range', 1e-4
+                )
+                if strategy.get('sparse_weight_bounds') is None:
+                    table.accessor.sparse_sgd_param.weight_bounds.extend(
+                        [-10, 10]
+                    )
+                else:
+                    table.accessor.sparse_sgd_param.weight_bounds.extend(
+                        strategy.get('sparse_weight_bounds')
+                    )
+                table.accessor.embedx_dim = strategy.get('sparse_embedx_dim', 8)
+                table.accessor.embedx_threshold = strategy.get(
+                    'sparse_embedx_threshold', 10
+                )
+                table.accessor.fea_dim = int(table.accessor.embedx_dim) + 3
+                table.accessor.downpour_accessor_param.nonclk_coeff = (
+                    strategy.get('sparse_nonclk_coeff', 0.1)
+                )
+                table.accessor.downpour_accessor_param.click_coeff = (
+                    strategy.get('sparse_click_coeff', 1)
+                )
+                table.accessor.downpour_accessor_param.base_threshold = (
+                    strategy.get('sparse_base_threshold', 1.5)
+                )
+                table.accessor.downpour_accessor_param.delta_threshold = (
+                    strategy.get('sparse_delta_threshold', 0.25)
+                )
+                table.accessor.downpour_accessor_param.delta_keep_days = (
+                    strategy.get('sparse_delta_keep_days', 16)
+                )
+                table.accessor.downpour_accessor_param.delete_after_unseen_days = strategy.get(
+                    'sparse_delete_after_unseen_days', 30
+                )
+                table.accessor.downpour_accessor_param.ssd_unseenday_threshold = strategy.get(
+                    'sparse_ssd_unseenday_threshold', 1
+                )
+                table.accessor.downpour_accessor_param.show_click_decay_rate = (
+                    strategy.get('sparse_show_click_decay_rate', 0.98)
+                )
+                table.accessor.downpour_accessor_param.delete_threshold = (
+                    strategy.get('sparse_delete_threshold', 0.8)
+                )
+                converter = strategy.get(
+                    'sparse_converter',
+                    "(scripts/xbox_compressor_mf.py | bin/xbox_pb_converter)",
+                )
+                deconverter = strategy.get(
+                    'sparse_deconverter',
+                    "(bin/xbox_pb_deconverter | scripts/xbox_decompressor_mf.awk)",
+                )
+
+                table1 = table.accessor.table_accessor_save_param.add()
+                table1.param = 1
+                table1.converter = converter
+                table1.deconverter = deconverter
+
+                table2 = table.accessor.table_accessor_save_param.add()
+                table2.param = 2
+                table2.converter = converter
+                table2.deconverter = deconverter
+            elif accessor_class == 'DownpourSparseValueAccessor':
+                optimizer_name = strategy.get("sparse_optimizer", "adam")
+                table.accessor.sparse_commonsgd_param.name = optimizer_name
+                table.accessor.embedx_dim = strategy.get('sparse_embedx_dim', 8)
+                table.accessor.fea_dim = int(table.accessor.embedx_dim)
+                if optimizer_name == "naive":
+                    table.accessor.sparse_commonsgd_param.naive.learning_rate = strategy.get(
+                        'sparse_learning_rate', 0.05
+                    )
+                    table.accessor.sparse_commonsgd_param.naive.initial_range = strategy.get(
+                        'sparse_initial_range', 1e-4
+                    )
+                    if strategy.get('sparse_weight_bounds') is None:
+                        table.accessor.sparse_commonsgd_param.naive.weight_bounds.extend(
+                            [-10, 10]
+                        )
+                    else:
+                        table.accessor.sparse_commonsgd_param.naive.weight_bounds.extend(
+                            strategy.get('sparse_weight_bounds')
+                        )
+                elif optimizer_name == "adagrad":
+                    table.accessor.sparse_commonsgd_param.adagrad.learning_rate = strategy.get(
+                        'sparse_learning_rate', 0.05
+                    )
+                    table.accessor.sparse_commonsgd_param.adagrad.initial_range = strategy.get(
+                        'sparse_initial_range', 1e-4
+                    )
+                    table.accessor.sparse_commonsgd_param.adagrad.initial_g2sum = strategy.get(
+                        'sparse_initial_g2sum', 3
+                    )
+                    if strategy.get('sparse_weight_bounds') is None:
+                        table.accessor.sparse_commonsgd_param.adagrad.weight_bounds.extend(
+                            [-10, 10]
+                        )
+                    else:
+                        table.accessor.sparse_commonsgd_param.adagrad.weight_bounds.extend(
+                            strategy.get('sparse_weight_bounds')
+                        )
+                elif optimizer_name == "adam":
+                    table.accessor.sparse_commonsgd_param.adam.learning_rate = (
+                        strategy.get('sparse_learning_rate', 0.001)
+                    )
+                    table.accessor.sparse_commonsgd_param.adam.initial_range = (
+                        strategy.get('sparse_initial_range', 1e-4)
+                    )
+                    table.accessor.sparse_commonsgd_param.adam.beta1_decay_rate = strategy.get(
+                        'sparse_beta1_decay_rate', 0.9
+                    )
+                    table.accessor.sparse_commonsgd_param.adam.beta2_decay_rate = strategy.get(
+                        'sparse_beta2_decay_rate', 0.999
+                    )
+                    table.accessor.sparse_commonsgd_param.adam.ada_epsilon = (
+                        strategy.get('sparse_ada_epsilon', 1e-8)
+                    )
+                    if strategy.get('sparse_weight_bounds') is None:
+                        table.accessor.sparse_commonsgd_param.adam.weight_bounds.extend(
+                            [-10, 10]
+                        )
+                    else:
+                        table.accessor.sparse_commonsgd_param.adam.weight_bounds.extend(
+                            strategy.get('sparse_weight_bounds')
+                        )
+                converter = strategy.get(
+                    'sparse_converter',
+                    "(scripts/xbox_compressor_mf.py | bin/xbox_pb_converter)",
+                )
+                deconverter = strategy.get(
+                    'sparse_deconverter',
+                    "(bin/xbox_pb_deconverter | scripts/xbox_decompressor_mf.awk)",
+                )
+
+                table1 = table.accessor.table_accessor_save_param.add()
+                table1.param = 1
+                table1.converter = converter
+                table1.deconverter = deconverter
+
+                table2 = table.accessor.table_accessor_save_param.add()
+                table2.param = 2
+                table2.converter = converter
+                table2.deconverter = deconverter
+            elif (
+                accessor_class == 'DownpourUnitAccessor'
+                or accessor_class == 'DownpourDoubleUnitAccessor'
+                or accessor_class == 'DownpourCtrDymfAccessor'
+            ):
+                self.add_sparse_table_common_config(table, strategy)
+                self.add_sparse_optimizer(
+                    table.accessor.embed_sgd_param, strategy, "embed_"
+                )
+                self.add_sparse_optimizer(
+                    table.accessor.embedx_sgd_param, strategy, "embedx_"
+                )
+
+    def add_dense_table(
+        self, table_id, param_var, grad_var, strategy, sparse_table_names
+    ):
+        """
+        Args:
+            table_id(int): id of sparse params table
+            param_var(list): param vars
+            grad_var(list): param grad vars
+            strategy(dict): the dense config dict
+            sparse_table_names(list): sparse table names
+        Returns:
+            return None
+        """
+        fea_dim = 0
+        dense_param_vars = []
+        for p in param_var:
+            if p.name not in sparse_table_names:
+                dense_param_vars.append(p)
+
+        for param in dense_param_vars:
+            fea_dim += reduce(lambda x, y: x * y, param.shape, 1)
+
+        for table in self._server.downpour_server_param.downpour_table_param:
+            if table.table_id == table_id:
+                if table.type == pslib.PS_DENSE_TABLE:
+                    table.accessor.fea_dim = fea_dim
+                    return
+                else:
+                    raise ValueError(
+                        f"expect table {table_id} type={pslib.PS_DENSE_TABLE}, but actual type={table.type}"
+                    )
+
+        if strategy is None:
+            strategy = {}
+        table = self._server.downpour_server_param.downpour_table_param.add()
+        table.table_id = table_id
+        support_dense_key_list = [
+            'dense_table_class',
+            'dense_compress_in_save',
+            'dense_accessor_class',
+            'dense_optimizer',
+            'dense_learning_rate',
+            'dense_avg_decay',
+            'dense_ada_decay',
+            'dense_ada_epsilon',
+            'dense_mom_decay',
+            'dense_naive_lr',
+        ]
+
+        for key in strategy:
+            if key not in support_dense_key_list:
+                raise ValueError(f"strategy key '{key}' not support")
+
+        table.table_class = strategy.get(
+            'dense_table_class', "DownpourDenseTable"
+        )
+        table.type = pslib.PS_DENSE_TABLE
+        table.compress_in_save = strategy.get('dense_compress_in_save', True)
+        table.accessor.accessor_class = strategy.get(
+            'dense_accessor_class', "DownpourDenseValueAccessor"
+        )
+        table.accessor.dense_sgd_param.name = strategy.get(
+            'dense_optimizer', "adam"
+        )
+        table.accessor.dense_sgd_param.adam.learning_rate = strategy.get(
+            'dense_learning_rate', 5e-06
+        )
+        table.accessor.dense_sgd_param.adam.avg_decay_rate = strategy.get(
+            'dense_avg_decay', 0.999993
+        )
+        table.accessor.dense_sgd_param.adam.ada_decay_rate = strategy.get(
+            'dense_ada_decay', 0.9999
+        )
+        table.accessor.dense_sgd_param.adam.ada_epsilon = strategy.get(
+            'dense_ada_epsilon', 1e-8
+        )
+        table.accessor.dense_sgd_param.adam.mom_decay_rate = strategy.get(
+            'dense_mom_decay', 0.99
+        )
+        table.accessor.dense_sgd_param.naive.learning_rate = strategy.get(
+            'dense_naive_lr', 0.0002
+        )
+        table.accessor.fea_dim = fea_dim
+
+    def add_data_norm_table(
+        self,
+        table_id,
+        learning_rate,
+        param_var,
+        grad_var,
+        strategy,
+        sparse_table_names,
+    ):
+        """
+        Args:
+            table_id(int): id of datanorm table
+            learning_rate(float): the learning rate used to update parameters
+            param_var(list): param vars
+            grad_var(list): param grad vars
+            strategy(dict): the datanorm config dict
+            sparse_table_names(list): sparse table names
+        Returns:
+            return None
+        """
+        fea_dim = 0
+        dense_param_vars = []
+        for p in param_var:
+            if p.name not in sparse_table_names:
+                dense_param_vars.append(p)
+
+        for param in dense_param_vars:
+            fea_dim += reduce(lambda x, y: x * y, param.shape, 1)
+
+        for table in self._server.downpour_server_param.downpour_table_param:
+            if table.table_id == table_id:
+                if table.type == pslib.PS_DENSE_TABLE:
+                    table.accessor.fea_dim = fea_dim
+                    return
+                else:
+                    raise ValueError(
+                        f"expect table {table_id} type={pslib.PS_DENSE_TABLE}, but actual type={table.type}"
+                    )
+        if strategy is None:
+            strategy = {}
+
+        support_datanorm_key_list = [
+            'datanorm_table_class',
+            'datanorm_compress_in_save',
+            'datanorm_accessor_class',
+            'datanorm_operation',
+            'datanorm_decay_rate',
+        ]
+
+        for key in strategy:
+            if key not in support_datanorm_key_list:
+                raise ValueError(f"strategy key '{key}' not support")
+
+        table = self._server.downpour_server_param.downpour_table_param.add()
+        table.table_id = table_id
+        table.table_class = strategy.get(
+            'datanorm_table_class', 'DownpourDenseTable'
+        )
+        table.type = pslib.PS_DENSE_TABLE
+        table.compress_in_save = strategy.get('datanorm_compress_in_save', True)
+        table.accessor.accessor_class = strategy.get(
+            'datanorm_accessor_class', 'DownpourDenseValueAccessor'
+        )
+        table.accessor.dense_sgd_param.name = strategy.get(
+            'datanorm_operation', 'summary'
+        )
+        table.accessor.dense_sgd_param.summary.summary_decay_rate = (
+            strategy.get('datanorm_decay_rate', 0.999999)
+        )
+        table.accessor.fea_dim = fea_dim
+
+    def add_sparse_optimizer(self, sgd, strategy, prefix):
+        optimizer_name = strategy.get(prefix + "sparse_optimizer", "adagrad")
+        sgd.name = optimizer_name
+        if optimizer_name == "naive":
+            sgd.naive.learning_rate = strategy.get(
+                prefix + 'sparse_learning_rate', 0.05
+            )
+            sgd.naive.initial_range = strategy.get(
+                prefix + 'sparse_initial_range', 1e-4
+            )
+            bounds = strategy.get(prefix + 'sparse_weight_bounds', [-10, 10])
+            sgd.naive.weight_bounds.extend(bounds)
+        elif optimizer_name == "adagrad":
+            sgd.adagrad.learning_rate = strategy.get(
+                prefix + 'sparse_learning_rate', 0.05
+            )
+            sgd.adagrad.initial_range = strategy.get(
+                prefix + 'sparse_initial_range', 1e-4
+            )
+            if prefix == "embed_":
+                sgd.adagrad.initial_range = 0
+            sgd.adagrad.initial_g2sum = strategy.get(
+                prefix + 'sparse_initial_g2sum', 3
+            )
+            bounds = strategy.get(prefix + 'sparse_weight_bounds', [-10, 10])
+            sgd.adagrad.weight_bounds.extend(bounds)
+        elif optimizer_name == "std_adagrad":
+            sgd.adagrad.learning_rate = strategy.get(
+                prefix + 'sparse_learning_rate', 0.05
+            )
+            sgd.adagrad.initial_range = strategy.get(
+                prefix + 'sparse_initial_range', 1e-4
+            )
+            if prefix == "embed_":
+                sgd.adagrad.initial_range = 0
+            sgd.adagrad.initial_g2sum = strategy.get(
+                prefix + 'sparse_initial_g2sum', 3
+            )
+            bounds = strategy.get(prefix + 'sparse_weight_bounds', [-10, 10])
+            sgd.adagrad.weight_bounds.extend(bounds)
+        elif optimizer_name == "adam":
+            sgd.adam.learning_rate = strategy.get(
+                prefix + 'sparse_learning_rate', 0.001
+            )
+            sgd.adam.initial_range = strategy.get(
+                prefix + 'sparse_initial_range', 1e-4
+            )
+            sgd.adam.beta1_decay_rate = strategy.get(
+                prefix + 'sparse_beta1_decay_rate', 0.9
+            )
+            sgd.adam.beta2_decay_rate = strategy.get(
+                prefix + 'sparse_beta2_decay_rate', 0.999
+            )
+            sgd.adam.ada_epsilon = strategy.get(
+                prefix + 'sparse_ada_epsilon', 1e-8
+            )
+            bounds = strategy.get(prefix + 'sparse_weight_bounds', [-10, 10])
+            sgd.adam.weight_bounds.extend(bounds)
+
+    def add_sparse_table_common_config(self, table, strategy):
+        table.accessor.embedx_dim = strategy.get('sparse_embedx_dim', 8)
+        table.accessor.embedx_threshold = strategy.get(
+            'sparse_embedx_threshold', 10
+        )
+        table.accessor.fea_dim = int(table.accessor.embedx_dim) + 3
+        table.accessor.downpour_accessor_param.nonclk_coeff = strategy.get(
+            'sparse_nonclk_coeff', 0.1
+        )
+        table.accessor.downpour_accessor_param.click_coeff = strategy.get(
+            'sparse_click_coeff', 1
+        )
+        table.accessor.downpour_accessor_param.base_threshold = strategy.get(
+            'sparse_base_threshold', 1.5
+        )
+        table.accessor.downpour_accessor_param.delta_threshold = strategy.get(
+            'sparse_delta_threshold', 0.25
+        )
+        table.accessor.downpour_accessor_param.delta_keep_days = strategy.get(
+            'sparse_delta_keep_days', 16
+        )
+        table.accessor.downpour_accessor_param.delete_after_unseen_days = (
+            strategy.get('sparse_delete_after_unseen_days', 30)
+        )
+        table.accessor.downpour_accessor_param.show_click_decay_rate = (
+            strategy.get('sparse_show_click_decay_rate', 0.98)
+        )
+        table.accessor.downpour_accessor_param.delete_threshold = strategy.get(
+            'sparse_delete_threshold', 0.8
+        )
+        converter = strategy.get(
+            'sparse_converter',
+            "(scripts/xbox_compressor_mf.py | bin/xbox_pb_converter)",
+        )
+        deconverter = strategy.get(
+            'sparse_deconverter',
+            "(bin/xbox_pb_deconverter | scripts/xbox_decompressor_mf.awk)",
+        )
+
+        table1 = table.accessor.table_accessor_save_param.add()
+        table1.param = 1
+        table1.converter = converter
+        table1.deconverter = deconverter
+
+        table2 = table.accessor.table_accessor_save_param.add()
+        table2.param = 2
+        table2.converter = converter
+        table2.deconverter = deconverter
+
+    def get_desc(self):
+        """
+        Return downpour server program_desc
+        """
+        return self._server
+
+
+class DownpourWorker(Worker):
+    """
+    DownpourWorker class is used to generate worker program_desc
+    Args:
+        window (int): push params frequency
+        worker: it is pslib.DownpourTrainerParameter
+    Examples:
+        worker = DownpourWorker(1)
+    """
+
+    def __init__(self, window):
+        self.window = window
+        self._worker = pslib.DownpourTrainerParameter()
+
+    def add_sparse_table(
+        self, table_id, slot_key_vars, slot_value_vars, slot_value_grads=None
+    ):
+        """
+        Args:
+            table_id(int): id of sparse params table
+            slot_key_vars(list): slot key id
+            slot_value_vars(list): slot key value after embedding
+            slot_value_grads(list): grad of all params, default is None
+        Returns:
+            return None
+        """
+        if slot_value_grads is None:
+            slot_value_grad_names = [
+                var.name + "@GRAD" for var in slot_value_vars
+            ]
+        else:
+            value_to_key = {}
+            for i in range(len(slot_key_vars)):
+                value_to_key[slot_value_vars[i].name] = slot_key_vars[i]
+            slot_value_grad_names = []
+            all_grad_names = [var.name for var in slot_value_grads]
+            for var in slot_value_vars:
+                if var.name + "@GRAD" in all_grad_names:
+                    slot_value_grad_names.append(var.name + "@GRAD")
+            sorted_slot_value_vars = [
+                i
+                for i in slot_value_vars
+                if i.name + "@GRAD" in slot_value_grad_names
+            ]
+            sorted_slot_value_vars += [
+                i
+                for i in slot_value_vars
+                if i.name + "@GRAD" not in slot_value_grad_names
+            ]
+            sorted_slot_key_vars = [
+                value_to_key[v.name] for v in sorted_slot_value_vars
+            ]
+
+        target_table = None
+        for table in self._worker.sparse_table:
+            if table.table_id == table_id:
+                keys = table.slot_key
+                key_names = [var.name for var in sorted_slot_key_vars]
+                for key_name in key_names:
+                    if key_name not in keys:
+                        raise ValueError(
+                            f"sparse table {table_id} slot_key error"
+                        )
+                target_table = table
+                break
+
+        table = target_table
+        if table is not None:
+            self._worker.sparse_table.remove(table)
+        table = self._worker.sparse_table.add()
+        table.table_id = table_id
+        table.slot_key.extend([var.name for var in sorted_slot_key_vars])
+        table.slot_value.extend([var.name for var in sorted_slot_value_vars])
+        table.slot_gradient.extend(slot_value_grad_names)
+
+    def add_dense_table(
+        self,
+        table_id,
+        learning_rate,
+        param_vars,
+        grad_vars,
+        dense_start_table_id,
+        sparse_table_names,
+    ):
+        r"""
+        Args:
+            table_id(int): id of sparse params table
+            learning_rate(float): the learning rate used to update parameters. \
+                Can be a float value
+            param_vars(list): all dense param. it is a list.
+            grad_vars(list): all dense grad param it is a list.
+            dense_start_table_id(int): dense table start index
+            sparse_table_names(list): sparse table names
+        Returns:
+            return None
+        """
+        sparse_table_name_grad = []
+        for name in sparse_table_names:
+            sparse_table_name_grad.append(name + "@GRAD")
+
+        dense_param_name = []
+        for p in param_vars:
+            if p.name not in sparse_table_names:
+                dense_param_name.append(p.name)
+
+        dense_grad_name = []
+        for g in grad_vars:
+            if g.name not in sparse_table_name_grad:
+                dense_grad_name.append(g.name)
+
+        dense_param_name.sort()
+        dense_grad_name.sort()
+
+        for table in self._worker.dense_table:
+            if table.table_id == table_id:
+                desc_dense_param_name = list(table.dense_variable_name)
+                desc_dense_param_name.sort()
+
+                if dense_param_name == desc_dense_param_name:
+                    desc_dense_grad_name = list(
+                        table.dense_gradient_variable_name
+                    )
+                    desc_dense_grad_name.sort()
+                    if dense_grad_name == desc_dense_grad_name:
+                        return
+                    else:
+                        raise ValueError(
+                            f"dense table {table_id} dense_gradient_variable_name "
+                            "error"
+                        )
+                else:
+                    raise ValueError(
+                        f"dense table {table_id} dense_variable_name error"
+                    )
+
+        table = self._worker.dense_table.add()
+        table.table_id = table_id
+
+        # def cmp_fc(x, y):
+        #    if x.startswith("fc_") and y.startswith("fc_"):
+        #        index_x = x.find('.')
+        #        index_y = y.find('.')
+        #        if index_x > 0 and index_y > 0:
+        #            num_x = x[3:index_x]
+        #            num_y = y[3:index_y]
+        #            if num_x.isdigit() and num_y.isdigit():
+        #                if int(num_x) < int(num_y):
+        #                    return -1
+        #                if int(num_x) > int(num_y):
+        #                    return 1
+        #                if x[index_x + 1] == 'w' and y[index_y + 1] == 'b':
+        #                    return -1
+        #                if x[index_x + 1] == 'b' and y[index_y + 1] == 'w':
+        #                    return 1
+        #    if x < y:
+        #        return -1
+        #    else:
+        #        return 1
+
+        # table.dense_variable_name.extend(sorted(dense_param_name, cmp_fc))
+        # table.dense_gradient_variable_name.extend(
+        #    sorted(dense_grad_name, cmp_fc))
+        table.dense_variable_name.extend(dense_param_name)
+        table.dense_gradient_variable_name.extend(dense_grad_name)
+
+    def get_desc(self):
+        """
+        Return downpour worker program_desc
+        """
+        return self._worker

python/paddle/incubate/distributed/fleet/utils.py

@@ -0,0 +1,511 @@
+# 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
+import os
+import subprocess
+from collections import OrderedDict
+
+import numpy as np
+from google.protobuf import text_format
+
+import paddle
+from paddle import base
+from paddle.base import core
+from paddle.base.framework import Program
+from paddle.base.proto import framework_pb2
+from paddle.distributed.fleet.base.util_factory import draw_block_graphviz
+from paddle.framework import io_utils
+
+__all__ = [
+    "load_program",
+    "save_program",
+    "program_type_trans",
+    "check_saved_vars_try_dump",
+    "parse_program",
+    "check_pruned_program_vars",
+    "graphviz",
+]
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+formatter = logging.Formatter(fmt='%(asctime)s - %(levelname)s - %(message)s')
+ch = logging.StreamHandler()
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+persistable_vars_out_fn = "vars_persistable.log"
+all_vars_out_fn = "vars_all.log"
+ops_out_fn = "ops.log"
+
+feed_fetch_type_list = [
+    core.VarDesc.VarType.FEED_MINIBATCH,
+    core.VarDesc.VarType.FETCH_LIST,
+]
+not_expected_op_types = ["lookup_table"]
+
+
+def load_program(model_filename, is_text=False):
+    if is_text:
+        return load_program_text(model_filename)
+    return load_program_binary(model_filename)
+
+
+def load_program_binary(model_filename):
+    """load program from binary string file"""
+    with open(model_filename, "rb") as f:
+        program_desc_str = f.read()
+    return Program.parse_from_string(program_desc_str)
+
+
+def load_program_text(model_filename):
+    """load program from human-readable text file"""
+    with open(model_filename, "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())
+
+
+def save_program(program, model_filename='__model__', is_text=False):
+    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 check_pruned_program_vars(train_prog, pruned_prog):
+    is_match = True
+
+    pruned_vars = [
+        (v.name, v)
+        for v in pruned_prog.list_vars()
+        if io_utils.is_persistable(v)
+    ]
+    pruned_vars = OrderedDict(pruned_vars)
+    pruned_vars_name = list(pruned_vars)
+    logger.info(f"persistable vars in pruned program: {pruned_vars_name}")
+
+    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:
+            logger.error(
+                f"not find variable '{var_name}' in train program. please check pruning."
+            )
+            logger.error(e)
+            continue
+        if (
+            var.shape != train_prog_var.shape
+            or var.dtype != train_prog_var.dtype
+        ):
+            logger.error(
+                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 graphviz(block, output_dir="", filename='debug'):
+    dot_path = os.path.join(output_dir, filename + '.dot')
+    pdf_path = os.path.join(output_dir, 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 program_type_trans(prog_dir, prog_fn, is_text):
+    prog = load_program(os.path.join(prog_dir, prog_fn), is_text)
+    prog_out_fn = prog_fn + ".bin" if is_text else prog_fn + ".pbtxt"
+    save_program(prog, os.path.join(prog_dir, prog_out_fn), 1 - is_text)
+    return prog_out_fn
+
+
+def append_save_op(block, var, path):
+    block.append_op(
+        type='save', inputs={'X': [var]}, outputs={}, attrs={'file_path': path}
+    )
+
+
+def append_load_op(block, var, path):
+    block.append_op(
+        type='load',
+        inputs={},
+        outputs={'Out': [var]},
+        attrs={'file_path': path},
+    )
+
+
+def save_var(np_array, var_name, shape_list, dtype, save_path):
+    program = base.Program()
+    place = base.CPUPlace()
+    exe = base.Executor(place)
+    shape = list(shape_list)
+    with base.program_guard(program):
+        d0_data = paddle.static.data(var_name, shape=shape, dtype=dtype)
+        append_save_op(program.global_block(), d0_data, save_path)
+        exe.run(feed={var_name: np_array}, fetch_list=[])
+
+
+def load_var(var_name, shape_list, dtype, save_path):
+    program = base.Program()
+    place = base.CPUPlace()
+    exe = base.Executor(place)
+    with base.program_guard(program):
+        d0_data = paddle.static.data(var_name, shape=shape_list, dtype=dtype)
+        append_load_op(program.global_block(), d0_data, save_path)
+        outs = exe.run(feed={}, fetch_list=[d0_data])
+        return outs
+
+
+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
+
+
+def feed_gen(batch_size, feeded_vars_dims, feeded_vars_filelist):
+    batch_feed = []
+    for i, fn in enumerate(feeded_vars_filelist):
+        batch_feed.append(reader(batch_size, fn, feeded_vars_dims[i]))
+    return batch_feed
+
+
+def try_load_model_vars(
+    dump_dir,
+    dump_prog_fn,
+    is_text_dump_program,
+    batch_size,
+    feed_config,
+    fetch_config,
+    save_filename,
+    saved_params,
+):
+    place = base.CPUPlace()
+    exe = base.Executor(place)
+    scope = base.core.Scope()
+    with base.scope_guard(scope):
+        if is_text_dump_program:
+            dump_prog_fn = program_type_trans(
+                dump_dir, dump_prog_fn, is_text_dump_program
+            )
+
+        [
+            inference_program,
+            feed_target_names,
+            fetch_targets,
+        ] = paddle.static.io.load_inference_model(
+            dump_dir,
+            exe,
+            model_filename=dump_prog_fn,
+            params_filename=save_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 = base.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
+        fetch_targets_names = [v.name for v in fetch_targets]
+        if not feed_target_names:
+            logger.warning("no feed targets in program.")
+        if not fetch_targets_names:
+            logger.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
+        ):
+            logger.warning(
+                f"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
+        ):
+            logger.warning(
+                f"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:
+            logger.info("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(
+                                (
+                                    batch_size,
+                                    *list(feed_config.feeded_vars_dims[i]),
+                                )
+                            ),
+                            dtype=feed_config.feeded_vars_types[i],
+                        )
+                    )
+                elif var.lod_level == 1:
+                    t = np.array(
+                        np.random.random(
+                            (
+                                batch_size,
+                                *list(feed_config.feeded_vars_dims[i]),
+                            )
+                        ),
+                        dtype=feed_config.feeded_vars_types[i],
+                    )
+                    feed_tensors.append(
+                        base.create_lod_tensor(t, [[1] * 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:
+            logger.info(
+                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 = base.DataFeeder(feed_list=feed_vars, place=place)
+            batch_feed = feed_gen(
+                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):
+            logger.info(f"fetch_targets name: {v.name}")
+            logger.info(f"fetch_targets: {results[i]}")
+        return results
+
+
+def check_not_expected_ops(prog):
+    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:
+            logger.warning(
+                f"find op type '{op.type}' in program, please check if your program is pruned correctly !"
+            )
+            op_types_set.add(op.type)
+
+
+def check_saved_vars_try_dump(
+    dump_dir,
+    dump_prog_fn,
+    is_text_dump_program,
+    feed_config,
+    fetch_config,
+    batch_size=1,
+    save_filename=None,
+):
+    dump_prog = load_program(
+        os.path.join(dump_dir, dump_prog_fn), is_text_dump_program
+    )
+    saved_params = [
+        v for v in dump_prog.list_vars() if io_utils.is_persistable(v)
+    ]
+    logger.info(
+        f"persistable vars in dump program: {[v.name for v in saved_params]}"
+    )
+
+    check_not_expected_ops(dump_prog)
+
+    return try_load_model_vars(
+        dump_dir,
+        dump_prog_fn,
+        is_text_dump_program,
+        batch_size,
+        feed_config,
+        fetch_config,
+        save_filename,
+        saved_params,
+    )
+
+
+def parse_program(program, output_dir):
+    # persistable vars
+    output = {}
+    persistable_vars = [
+        v for v in program.list_vars() if io_utils.is_persistable(v)
+    ]
+    output["persistable_vars"] = [
+        {
+            'name': str(v.name),
+            'shape': str(v.shape),
+            'lod_level': int(v.lod_level),
+            'dtype': str(v.dtype),
+            'type': str(v.type),
+        }
+        for v in persistable_vars
+    ]
+    with open(os.path.join(output_dir, persistable_vars_out_fn), 'w') as f:
+        f.write("persistable vars:\n")
+        for var in output["persistable_vars"]:
+            f.write(str(var))
+            f.write("\n")
+
+    # all vars
+    all_vars = list(program.list_vars())
+    output["all_vars"] = [
+        (
+            {
+                'name': str(v.name),
+                'shape': str(v.shape),
+                'lod_level': int(v.lod_level),
+                'dtype': str(v.dtype),
+            }
+            if v.type not in feed_fetch_type_list
+            else {'name': str(v.name), 'type': str(v.type)}
+        )
+        for v in all_vars
+    ]
+    with open(os.path.join(output_dir, all_vars_out_fn), 'w') as f:
+        f.write("all vars:\n")
+        for var in output["all_vars"]:
+            f.write(str(var))
+            f.write("\n")
+
+    # ops
+    ops = program.global_block().ops
+    output["ops"] = [
+        {
+            'type': op.type,
+            'input_arg_names': str(op.input_arg_names),
+            'output_arg_names': str(op.output_arg_names),
+        }
+        for op in ops
+    ]
+    with open(os.path.join(output_dir, ops_out_fn), 'w') as f:
+        f.write("ops:\n")
+        for op in output["ops"]:
+            f.write(str(op))
+            f.write("\n")

python/paddle/incubate/distributed/models/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

python/paddle/incubate/distributed/models/moe/__init__.py

@@ -0,0 +1,19 @@
+# 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 .gate import BaseGate, GShardGate, NaiveGate, SwitchGate  # noqa: F401
+from .grad_clip import ClipGradForMOEByGlobalNorm
+from .moe_layer import MoELayer  # noqa: F401
+
+ClipGradByGlobalNorm = ClipGradForMOEByGlobalNorm

python/paddle/incubate/distributed/models/moe/gate/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .base_gate import BaseGate  # noqa: F401
+from .gshard_gate import GShardGate  # noqa: F401
+from .naive_gate import NaiveGate  # noqa: F401
+from .switch_gate import SwitchGate  # noqa: F401

python/paddle/incubate/distributed/models/moe/gate/base_gate.py

@@ -0,0 +1,43 @@
+# 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 the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/base_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+from paddle import nn
+
+
+class BaseGate(nn.Layer):
+    def __init__(self, num_expert, world_size):
+        super().__init__()
+        self.world_size = world_size
+        self.num_expert = num_expert
+        self.tot_expert = world_size * num_expert
+        self.loss = None
+
+    def forward(self, x):
+        raise NotImplementedError("Please implement the forward function.")
+
+    def set_loss(self, loss):
+        self.loss = loss
+
+    def get_loss(self, clear=True):
+        loss = self.loss
+        if clear:
+            self.loss = None
+        return loss

python/paddle/incubate/distributed/models/moe/gate/gshard_gate.py

@@ -0,0 +1,84 @@
+# 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 the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/gshard_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import math
+
+import paddle
+import paddle.nn.functional as F
+
+from ..utils import limit_by_capacity
+from .naive_gate import NaiveGate
+
+
+class GShardGate(NaiveGate):
+    def __init__(
+        self,
+        d_model,
+        num_expert,
+        world_size,
+        topk=2,
+        capacity=(1.2, 2.4),
+        random_routing=True,
+        group=None,
+    ):
+        assert topk == 2, "topk should be 2 in gshard"
+        super().__init__(d_model, num_expert, world_size)
+        self.capacity = capacity
+        self.random_routing = random_routing
+        self.group = group
+
+    def forward(self, x):
+        topk_val, topk_idx, gate_score = super().forward(
+            x, return_all_scores=True
+        )
+        s = gate_score.shape[0]
+        top1_idx = topk_idx.flatten()
+        c_e = (
+            paddle.scatter(
+                paddle.zeros(shape=[self.tot_expert]),
+                top1_idx,
+                paddle.ones_like(top1_idx, dtype="float32"),
+                overwrite=False,
+            )
+            / s
+        )
+        m_e = paddle.mean(F.softmax(gate_score, axis=1), axis=0)
+        loss = paddle.mean(c_e * m_e) * (self.num_expert**2)
+        self.set_loss(loss)
+
+        cap_rate = self.capacity[0 if self.training else 1]
+        capacity = math.ceil(cap_rate * x.shape[0])
+        _new_lec, _new_gec, topk_idx = limit_by_capacity(
+            topk_idx,
+            self.num_expert,
+            self.world_size,
+            capacity,
+            group=self.group,
+        )
+
+        if self.random_routing:
+            rand_routing_prob = paddle.rand(
+                shape=[gate_score.shape[0]], dtype="float32"
+            )
+            topk_idx = paddle.distributed.models.moe.utils._random_routing(
+                topk_idx, topk_val, rand_routing_prob
+            )
+        return topk_val, topk_idx

python/paddle/incubate/distributed/models/moe/gate/naive_gate.py

@@ -0,0 +1,44 @@
+# 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 the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/naive_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import paddle
+from paddle import nn
+
+from .base_gate import BaseGate
+
+
+class NaiveGate(BaseGate):
+    def __init__(self, d_model, num_expert, world_size, topk=2):
+        super().__init__(num_expert, world_size)
+        self.gate = nn.Linear(d_model, self.tot_expert)
+        self.gate.weight.name = "gate_" + self.gate.weight.name
+        self.gate.bias.name = "gate_" + self.gate.bias.name
+        self.top_k = topk
+
+    def forward(self, inp, return_all_scores=False):
+        gate = self.gate(inp)
+        gate_top_k_val, gate_top_k_idx = paddle.topk(
+            gate, k=self.top_k, axis=-1, largest=True, sorted=False
+        )
+
+        if return_all_scores:
+            return gate_top_k_val, gate_top_k_idx, gate
+        return gate_top_k_val, gate_top_k_idx

python/paddle/incubate/distributed/models/moe/gate/switch_gate.py

@@ -0,0 +1,84 @@
+# 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 the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/gates/switch_gate.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import math
+
+import paddle
+import paddle.nn.functional as F
+
+from ..utils import limit_by_capacity
+from .naive_gate import NaiveGate
+
+
+class SwitchGate(NaiveGate):
+    def __init__(
+        self,
+        d_model,
+        num_expert,
+        world_size,
+        topk=1,
+        switch_eps=0.1,
+        capacity=(1.2, 2.4),
+        group=None,
+    ):
+        assert topk == 1, "topk should be 1 in switch"
+        super().__init__(d_model, num_expert, world_size, topk=1)
+        self.switch_eps = switch_eps
+        self.capacity = capacity
+        self.group = group
+
+    def forward(self, inp):
+        score = self.gate(inp)
+
+        if self.training:
+            noise = paddle.rand(shape=score.shape)
+            noise = noise * 2 * self.switch_eps + 1.0 - self.switch_eps
+            score += noise
+
+        score = F.softmax(score, axis=-1)
+        top1_score, top1_idx = paddle.topk(score, k=1, axis=-1, largest=True)
+
+        cap_rate = self.capacity[0 if self.training else 1]
+        capacity = math.ceil(cap_rate * inp.shape[0])
+        _new_lec, _new_gec, top1_idx = limit_by_capacity(
+            top1_idx,
+            self.num_expert,
+            self.world_size,
+            capacity,
+            group=self.group,
+        )
+        valid_idx = top1_idx[top1_idx > -1]
+        valid_idx_tmp = paddle.reshape(valid_idx, shape=[len(valid_idx), 1])
+        fraction_expert = (
+            paddle.scatter_nd_add(
+                x=paddle.zeros(shape=[self.tot_expert]),
+                index=valid_idx_tmp,
+                updates=paddle.ones_like(
+                    valid_idx, dtype=paddle.float32
+                ).reshape(shape=[len(valid_idx)]),
+            )
+            / valid_idx.numel()
+        )
+        prob_expert = score.sum(axis=0) / valid_idx.numel()
+        loss = (fraction_expert * prob_expert).sum() * self.tot_expert
+        self.set_loss(loss)
+
+        return top1_score, top1_idx

python/paddle/incubate/distributed/models/moe/grad_clip.py

@@ -0,0 +1,238 @@
+#   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 paddle
+import paddle.distributed as dist
+from paddle.autograd import no_grad
+from paddle.framework import core
+from paddle.nn import clip
+from paddle.nn.clip import ClipGradBase, _squared_l2_norm
+
+
+class ClipGradForMOEByGlobalNorm(ClipGradBase):
+    r"""
+    The Algorithm is the same as paddle.nn.ClipGradByGlobalNorm
+    Given a list of Tensor :math:`t\_list` , calculate the global norm for the elements of all tensors in
+    :math:`t\_list` , and limit it to ``clip_norm`` .
+
+    - If the global norm is greater than ``clip_norm`` , all elements of :math:`t\_list` will be compressed by a ratio.
+
+    - If the global norm is less than or equal to ``clip_norm`` , nothing will be done.
+
+    The list of Tensor :math:`t\_list` is not passed from this class, but the gradients of all parameters set in ``optimizer``.
+    If ``need_clip`` of specific param is ``False`` in its ``ParamAttr``, then the gradients of this param will not be clipped.
+
+    Gradient clip will takes effect after being set in ``optimizer`` , see the document ``optimizer``
+    (for example: :ref:`api_paddle_optimizer_SGD`).
+
+    The clipping formula is:
+
+    .. math::
+
+        t\_list[i] = t\_list[i] * \frac{clip\_norm}{\max(global\_norm, clip\_norm)}
+
+    where:
+
+    .. math::
+
+        global\_norm = \sqrt{\sum_{i=0}^{N-1}(l2norm(t\_list[i]))^2}
+
+    Note:
+        ``need_clip`` of ``ClipGradyGlobalNorm`` HAS BEEN DEPRECATED since 2.0.
+        Please use ``need_clip`` in ``ParamAttr`` to specify the clip scope.
+
+    Reference:
+        https://github.com/laekov/fastmoe/blob/master/examples/megatron/clip-grad-v2.2.patch
+        Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+
+
+    Args:
+        clip_norm (float): The maximum norm value.
+        is_expert_param_func (function): a function to decide whether a param should be put into moe_params_grads
+        moe_group (Group): group for moe experts communication.
+        group_name (str, optional): The group name for this clip. Default value is ``default_moe_group``.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x = paddle.uniform([10, 10], min=-1.0, max=1.0, dtype='float32')
+            >>> linear = paddle.nn.Linear(
+            ...     in_features=10,
+            ...     out_features=10,
+            ...     weight_attr=paddle.ParamAttr(need_clip=True),
+            ...     bias_attr=paddle.ParamAttr(need_clip=False),
+            ... )
+            >>> out = linear(x)
+            >>> loss = paddle.mean(out)
+            >>> loss.backward()
+
+            >>> clip = paddle.nn.ClipGradByGlobalNorm(
+            ...     clip_norm=1.0
+            ... )  # Cause paddle.nn hasn't this interface, so we use ClipGradByGlobalNorm here.
+            >>> sdg = paddle.optimizer.SGD(learning_rate=0.1, parameters=linear.parameters(), grad_clip=clip)
+            >>> sdg.step()
+    """
+
+    def __init__(
+        self,
+        clip_norm,
+        is_expert_param_func=None,
+        moe_group=None,
+        group_name="default_moe_group",
+    ):
+        super().__init__()
+        self.clip_norm = float(clip_norm)
+        self.group_name = group_name
+        self.moe_group = moe_group
+        if moe_group is not None and moe_group.nranks > 1:
+            assert is_expert_param_func is not None, (
+                "When moe group size > 1, a function for selecting expert params must be specified."
+            )
+        self.is_expert_param_func = is_expert_param_func
+
+    def __str__(self):
+        return f"Gradient Clip By GlobalNorm, global_norm={self.clip_norm:f}"
+
+    @staticmethod
+    def get_l2_norm_pow(params_grads, sum_dtype=None):
+        sum_square_list = []
+        sum_square_list_fp16 = []
+        sum_square_list_fp32 = []
+        for p, g in params_grads:
+            if g is None:
+                continue
+            if getattr(p, 'need_clip', True) is False:
+                continue
+            merge_grad = g
+            if g.type == core.VarDesc.VarType.SELECTED_ROWS:
+                merge_grad = clip.merge_selected_rows(g)
+                merge_grad = clip.get_tensor_from_selected_rows(merge_grad)
+            sum_square = _squared_l2_norm(merge_grad)
+            if sum_square.dtype == paddle.float16:
+                sum_square_list_fp16.append(sum_square)
+            elif sum_square.dtype == paddle.float32:
+                sum_square_list_fp32.append(sum_square)
+            else:
+                sum_square_list.append(sum_square)
+
+        # all parameters have been filtered out
+        if (
+            len(sum_square_list)
+            + len(sum_square_list_fp16)
+            + len(sum_square_list_fp32)
+            == 0
+        ):
+            return None, None
+        assert sum_dtype in [
+            "float64",
+            "float32",
+            None,
+        ], "sum's type must be float64/ float32 / None"
+        if sum_dtype != "float64":
+            sum_dtype = 'float64' if len(sum_square_list) > 0 else "float32"
+
+        global_norm_var = []
+        if len(sum_square_list_fp16) > 0:
+            global_norm_var_fp16 = paddle.add_n(sum_square_list_fp16)
+            global_norm_var.append(global_norm_var_fp16.astype(sum_dtype))
+        if len(sum_square_list_fp32) > 0:
+            global_norm_var_fp32 = paddle.add_n(sum_square_list_fp32)
+            if sum_dtype == 'float32':
+                global_norm_var.append(global_norm_var_fp32)
+            else:
+                global_norm_var.append(global_norm_var_fp32.astype(sum_dtype))
+        if len(sum_square_list) > 0:
+            global_norm_var_fp64 = paddle.add_n(sum_square_list)
+            global_norm_var.append(global_norm_var_fp64)
+        global_norm_var = paddle.add_n(global_norm_var)
+        return global_norm_var, sum_dtype
+
+    @no_grad()
+    def _dygraph_clip(self, params_grads):
+        normal_params_grads = []
+        moe_params_grads = []
+
+        # separate moe params from normal params
+        if self.moe_group is not None and self.moe_group.nranks > 1:
+            for p, g in params_grads:
+                if self.is_expert_param_func(p):
+                    moe_params_grads.append((p, g))
+                else:
+                    normal_params_grads.append((p, g))
+        else:
+            normal_params_grads = params_grads
+
+        # why to return sum_dtype?
+        # we will call `get_l2_norm_pow` twice and the precisions may be different.
+        # For convenience and simplification, we use sum_dtype directly instead of global_norm_var_normal.dtype
+        global_norm_var_normal, sum_dtype = self.get_l2_norm_pow(
+            normal_params_grads
+        )
+        global_norm_var_moe = None
+        if len(moe_params_grads) > 0:
+            global_norm_var_moe, _ = self.get_l2_norm_pow(
+                moe_params_grads, sum_dtype
+            )
+            if global_norm_var_moe is not None:
+                dist.all_reduce(
+                    global_norm_var_moe,
+                    op=dist.ReduceOp.SUM,
+                    group=self.moe_group,
+                )
+
+        if global_norm_var_normal is None and global_norm_var_moe is None:
+            return params_grads
+        elif global_norm_var_normal is None:
+            global_norm_var = global_norm_var_moe
+        elif global_norm_var_moe is None:
+            global_norm_var = global_norm_var_normal
+        else:
+            if global_norm_var_normal.dtype != global_norm_var_moe.dtype:
+                # compared with normal norm, moe norm is the later one,
+                # so its precision is no lower than normal norm
+                global_norm_var_normal = global_norm_var_normal.astype(
+                    global_norm_var_moe.dtype
+                )
+            global_norm_var = global_norm_var_normal + global_norm_var_moe
+
+        params_and_grads = []
+        global_norm_var = paddle.sqrt(global_norm_var)
+        max_global_norm = paddle.full(
+            shape=[1], 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),
+        )
+        for p, g in params_grads:
+            if g is None:
+                continue
+            if getattr(p, 'need_clip', True) is False:
+                params_and_grads.append((p, g))
+                continue
+            # TODO(wangxi): use inplace elementwise_mul
+            clip_input = (
+                clip_var.astype('float16')
+                if g.dtype == paddle.float16
+                else clip_var
+            )
+            new_grad = paddle.multiply(x=g, y=clip_input)
+            params_and_grads.append((p, new_grad))
+        return params_and_grads
+
+
+ClipGradByGlobalNorm = ClipGradForMOEByGlobalNorm

python/paddle/incubate/distributed/models/moe/moe_layer.py

@@ -0,0 +1,503 @@
+# 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.
+#
+# The file has been adapted from the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/layers.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import os
+
+import numpy as np
+
+import paddle
+from paddle import nn
+from paddle.autograd import PyLayer
+from paddle.distributed.utils.moe_utils import global_gather, global_scatter
+from paddle.distributed.utils.nccl_utils import check_nccl_version_for_p2p
+from paddle.framework import in_dynamic_mode
+from paddle.incubate.distributed.fleet import recompute_hybrid
+
+from .gate import BaseGate, GShardGate, NaiveGate, SwitchGate
+from .utils import count_by_gate
+
+
+def _local_scatter(inp, pos):
+    if pos.shape != [0]:
+        inp_buf = paddle.index_select(inp, pos, 0)
+    else:
+        inp_buf = paddle.empty([0, inp.shape[1]], dtype=inp.dtype)
+    return inp_buf
+
+
+def _local_gather(inp, pos, out_batch_size, maybe_overlap=True):
+    if pos.shape != [0]:
+        origin_dtype = inp.dtype
+        inp = paddle.cast(inp, dtype="float32")
+        inp_buf = paddle.scatter(
+            paddle.zeros(
+                shape=[out_batch_size, inp.shape[-1]], dtype="float32"
+            ),
+            pos,
+            inp,
+            overwrite=True,
+        )
+        inp_buf = paddle.cast(inp_buf, dtype=origin_dtype)
+    else:
+        inp_buf = paddle.zeros([out_batch_size, inp.shape[-1]], dtype=inp.dtype)
+    return inp_buf
+
+
+def _all_gather(tensor, group=None, use_calc_stream=True):
+    if group is not None and not group.is_member():
+        return
+
+    if in_dynamic_mode():
+        group = (
+            paddle.distributed.collective._get_default_group()
+            if group is None
+            else group
+        )
+        tensor_shape = list(tensor.shape)
+        tensor_shape[0] *= group.nranks
+        out = paddle.empty(tensor_shape, tensor.dtype)
+
+        task = group.process_group.all_gather(tensor, out)
+        task.wait()
+        return out
+    else:
+        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 paddle._C_ops.all_gather(
+            tensor,
+            ring_id,
+            nranks,
+        )
+
+
+class MoEScatter(PyLayer):
+    r"""
+    Scatter input samples from [batch x sequences] to contiguous alone experts.
+    If `world_size` is greater than 1, the samples will first be locally
+    scattered, and then exchanged across workers.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        inp,
+        pos,
+        local_expert_count,
+        global_expert_count,
+        fwd_batch_size,
+        world_size,
+        group=None,
+    ):
+        local_input_buf = _local_scatter(inp, pos)
+        if world_size > 1:
+            global_input_buf = global_scatter(
+                local_input_buf,
+                local_expert_count,
+                global_expert_count,
+                group=group,
+            )
+        else:
+            global_input_buf = local_input_buf
+
+        ctx.moe_args = inp.shape[0], world_size, group
+
+        variables = (pos, local_expert_count, global_expert_count)
+        ctx.save_for_backward(*variables)
+        return global_input_buf
+
+    @staticmethod
+    def backward(ctx, grad):
+        (pos, local_expert_count, global_expert_count) = ctx.saved_tensor()
+        (inp_batch_size, world_size, group) = ctx.moe_args
+
+        if world_size > 1:
+            local_grad_in = global_gather(
+                grad, local_expert_count, global_expert_count, group=group
+            )
+        else:
+            local_grad_in = grad
+        grad_in = _local_gather(local_grad_in, pos, inp_batch_size)
+        return grad_in, None, None, None
+
+
+class MoEGather(PyLayer):
+    r"""
+    Gather output samples from contiguous alone experts back to [batch x
+    sequences]. Works symmetrically with MoEScatter.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        global_output_buf,
+        pos,
+        local_expert_count,
+        global_expert_count,
+        local_batch_size,
+        world_size,
+        group=None,
+    ):
+        if world_size > 1:
+            local_output_buf = global_gather(
+                global_output_buf,
+                local_expert_count,
+                global_expert_count,
+                group=group,
+            )
+        else:
+            local_output_buf = global_output_buf
+        output = _local_gather(
+            local_output_buf, pos, local_batch_size, maybe_overlap=False
+        )
+
+        ctx.moe_args = (global_output_buf.shape[0], world_size, group)
+        variables = (pos, local_expert_count, global_expert_count)
+        ctx.save_for_backward(*variables)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        pos, local_expert_count, global_expert_count = ctx.saved_tensor()
+        fwd_batch_size, world_size, group = ctx.moe_args
+        grad_out_buf = _local_scatter(grad_out, pos)
+        if world_size > 1:
+            global_grad_out_buf = global_scatter(
+                grad_out_buf,
+                local_expert_count,
+                global_expert_count,
+                group=group,
+            )
+        else:
+            global_grad_out_buf = grad_out_buf
+        return global_grad_out_buf, None, None, None
+
+
+class AllGather(PyLayer):
+    r"""
+    A wrapper for the All-Gather function to support auto-differentiation.
+    """
+
+    @staticmethod
+    def forward(ctx, inp, rank, world_size, group):
+        tensor_list = []
+        paddle.distributed.all_gather(tensor_list, inp, group=group)
+        output = paddle.concat(tensor_list, axis=0)
+        ctx.args = rank, inp.shape[0]
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        rank, dim0 = ctx.args
+        return paddle.slice(
+            grad_out, axes=[0], starts=[rank * dim0], ends=[(rank + 1) * dim0]
+        )
+
+
+class Slice(PyLayer):
+    r"""
+    A wrapper for the Slice function to support auto-differentiation.
+    """
+
+    @staticmethod
+    def forward(ctx, inp, rank, world_size, group):
+        B = inp.shape[0]
+        local_batch_size = B // world_size
+        batch_start = local_batch_size * rank
+        batch_end = min(batch_start + local_batch_size, B)
+        inp = paddle.slice(
+            inp, axes=[0], starts=[batch_start], ends=[batch_end]
+        )
+        ctx.args = world_size, group
+        return inp
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        world_size, group = ctx.args
+        return _all_gather(grad_out, group=group)
+
+
+def prepare_forward(gate, num_expert, world_size, moe_group):
+    pos, local_expert_count, global_expert_count = count_by_gate(
+        gate, num_expert, world_size, group=moe_group
+    )
+    with paddle.no_grad():
+        fwd_expert_count = global_expert_count.reshape_(
+            [world_size, num_expert]
+        ).sum(axis=0)
+        fwd_batch_size = int(fwd_expert_count.sum().item())
+    return (
+        pos,
+        local_expert_count,
+        global_expert_count,
+        fwd_expert_count,
+        fwd_batch_size,
+    )
+
+
+class MoELayer(nn.Layer):
+    """MoE Layer
+    Args:
+        d_model (int): Model dimension.
+        experts (nn.LayerList): Expert networks list.
+        gate (dict|NaiveGate|SwitchGate|NaiveGate):
+
+            - If gate is a dict:
+              gate is a gate network config, containing 2 keys:
+              `type` (str) value can be: "naive", "gshard", "switch" or None, default is "gshard".
+              `top_k` (int) Default value is 2.
+            else gate is an instance of NaiveGate|SwitchGate|NaiveGate:
+
+        moe_group: moe group for experts communication.
+        mp_group: mp group for mp communication.
+        recompute_interval (int, optional): Whether to use recompute, default 0, means to disable recompute.
+        recompute_ctx (dict, optional): The context for recompute, if recompute_interval > 1, recompute_ctx must be given.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('Until Distributed move successfully, just skip it')
+            >>> from paddle.nn import layer, LayerList
+            >>> from paddle.distributed.moe import MoElayer
+            >>> from paddle.distributed.collective import Group
+            >>> from paddle.distributed import fleet
+
+            >>> moe_group = Group(
+            ...     fleet.worker_index(),
+            ...     0,
+            ...     list(range(fleet.worker_num())),
+            ... )
+            >>> mp_group = None
+
+            >>> num_experts = 8
+            >>> dim_feedforward = 512
+            >>> d_model = 8
+            >>> top_k = 2
+
+            >>> class ExpertLayer(Layer):
+            ...     def __init__(self, d_model, d_hidden, name=None, rank=0, windex=0, num_expert=1):
+            ...         super().__init__()
+            ...         self.htoh4 = nn.Linear(d_model, d_hidden)
+            ...         self.h4toh = nn.Linear(d_hidden, d_model)
+
+            ...     def forward(self, x):
+            ...         x = self.htoh4(x)
+            ...         x = self.h4toh(x)
+            ...         return x
+
+            >>> gate_config = {
+            ...     "type": "gshard",
+            ...     "top_k": top_k,
+            ... }
+
+            >>> experts_list = LayerList()
+            >>> for expi in range(num_experts):
+            ...     exp_layer = ExpertLayer(d_model, dim_feedforward // top_k, windex=expi, num_expert=num_experts)
+            ...     experts_list.append(exp_layer)
+
+            >>> moeLayer = MoELayer(
+            ...     d_model=d_model,
+            ...     experts=experts_list,
+            ...     gate=gate_config,
+            ...     moe_group=moe_group,
+            ...     mp_group=mp_group,
+            ...     recompute_interval=0,
+            ... )
+
+    """
+
+    def __init__(
+        self,
+        d_model,
+        experts,
+        gate=None,
+        moe_group=None,
+        mp_group=None,
+        recompute_interval=0,
+        recompute_ctx=None,
+    ):
+        super().__init__()
+
+        self.recompute_ctx = recompute_ctx
+
+        if gate is None:
+            gate = {}
+
+        assert isinstance(gate, (dict, BaseGate)), (
+            "gate config' type must be dict or an instance of BaseGate"
+        )
+        # only support mp/dp
+        self.group = moe_group
+
+        self.world_size = 1
+        if self.group is not None:
+            self.world_size = self.group.nranks
+        self.num_expert = len(experts)
+        self.recompute_interval = recompute_interval
+        assert experts is not None
+        self.experts = experts
+
+        if (
+            self.world_size > 1
+            and os.getenv("PADDLE_DISTRI_BACKEND", None) != "xccl"
+        ):
+            check_nccl_version_for_p2p()
+
+        self.mp_group = mp_group
+        self.d_model = d_model
+        if isinstance(gate, dict):
+            self.top_k = gate.get("top_k", 2)
+            gate = gate.get("type", "gshard")
+            if gate == "naive" or gate is None:
+                gate = NaiveGate(
+                    self.d_model,
+                    num_expert=len(experts),
+                    world_size=self.world_size,
+                    topk=self.top_k,
+                )
+            elif gate == "gshard":
+                gate = GShardGate(
+                    self.d_model,
+                    num_expert=len(experts),
+                    world_size=self.world_size,
+                    topk=self.top_k,
+                    group=self.group,
+                )
+            elif gate == "switch":
+                gate = SwitchGate(
+                    self.d_model,
+                    num_expert=len(experts),
+                    world_size=self.world_size,
+                    topk=self.top_k,
+                    group=self.group,
+                )
+            else:
+                raise AssertionError(
+                    f"We only support naive gate,                                 gshard gate and switch gate,                                 but you choose {gate} gate."
+                )
+        elif isinstance(gate, NaiveGate):
+            self.top_k = gate.top_k
+        elif isinstance(gate, BaseGate):
+            raise TypeError(f"Unimplemented gate type: {type(gate)}")
+        else:
+            raise TypeError("gate's type must be either dict or moe.BaseGate")
+        self.gate = gate
+
+    def forward(self, inp):
+        # inp shape: b * s * m
+        assert len(inp.shape) == 3
+        origin_shape = inp.shape
+        inp = inp.reshape_([-1, origin_shape[2]])
+
+        mp_rank = 0
+        mp_size = 1
+        if self.mp_group is not None:
+            mp_rank = self.mp_group.rank
+            mp_size = self.mp_group.nranks
+        if mp_size > 1:
+            inp = Slice.apply(inp, mp_rank, mp_size, self.mp_group)
+        value, gate = self.gate(inp)
+
+        (
+            pos,
+            local_expert_count,
+            global_expert_count,
+            fwd_expert_count,
+            fwd_batch_size,
+        ) = prepare_forward(gate, self.num_expert, self.world_size, self.group)
+
+        topk = 1
+        if len(gate.shape) == 2:
+            topk = gate.shape[1]
+
+        if pos.shape != [0]:
+            temp_pos = pos // topk
+        else:
+            temp_pos = pos
+        assert topk == self.top_k
+
+        x = MoEScatter.apply(
+            inp,
+            temp_pos,
+            local_expert_count,
+            global_expert_count,
+            fwd_batch_size,
+            self.world_size,
+            self.group,
+        )
+
+        d_model = self.d_model
+
+        def experts_fwd(x, fwd_expert_count, experts):
+            if x.shape[0] == 0:
+                return x
+            y = []
+            last_index = 0
+            assert isinstance(fwd_expert_count, np.ndarray)
+            assert len(experts) == len(fwd_expert_count)
+            for idx, expert_count in enumerate(fwd_expert_count):
+                if expert_count <= 0:
+                    continue
+                y.append(
+                    experts[idx](x[last_index : expert_count + last_index])
+                )
+                last_index = expert_count + last_index
+            return paddle.concat(y, axis=0)
+
+        if self.recompute_interval <= 0 or x.shape[0] == 0:
+            x = experts_fwd(x, fwd_expert_count.numpy(), self.experts)
+        else:
+            x = recompute_hybrid(
+                self.recompute_ctx,
+                experts_fwd,
+                x,
+                fwd_expert_count.numpy(),
+                self.experts,
+            )
+
+        out_batch_size = inp.shape[0]
+        if len(gate.shape) == 2:
+            out_batch_size *= gate.shape[1]
+
+        x = MoEGather.apply(
+            x,
+            pos,
+            local_expert_count,
+            global_expert_count,
+            out_batch_size,
+            self.world_size,
+            self.group,
+        )
+
+        x = x.reshape([-1, self.top_k, d_model])
+        value = value.reshape([x.shape[0], 1, self.top_k])
+        x = paddle.bmm(value, x).reshape([-1, d_model])
+
+        if mp_size > 1:
+            x = AllGather.apply(x, mp_rank, mp_size, self.mp_group)
+
+        x = paddle.reshape_(x, origin_shape)
+
+        return x

python/paddle/incubate/distributed/models/moe/utils.py

@@ -0,0 +1,87 @@
+# 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 the file:
+#     https://github.com/laekov/fastmoe/blob/master/fmoe/functions.py
+#     Git commit hash: 295a615aacce7e54a37e7935274ba15e901c78e4
+# We retain the following license from the original files:
+#     Copyright 2021, Jiaao He. All rights reserved.
+#   Licensed under the Apache License, Version 2.0 (the "License").
+
+import paddle
+from paddle.distributed.models.moe.utils import (
+    _assign_pos,
+    _limit_by_capacity,
+    _number_count,
+    _prune_gate_by_capacity,
+)
+from paddle.framework import in_dynamic_mode
+
+
+def _alltoall(in_tensor_list, group=None, use_calc_stream=True):
+    if group is not None and not group.is_member():
+        return
+
+    if in_dynamic_mode():
+        group = (
+            paddle.distributed.collective._get_default_group()
+            if group is None
+            else group
+        )
+        out = paddle.empty(in_tensor_list.shape, in_tensor_list.dtype)
+        task = group.process_group.alltoall(out, in_tensor_list)
+        task.wait()
+        return out
+    else:
+        ring_id = 0 if group is None else group.id
+        return paddle._C_ops.all_to_all(in_tensor_list, ring_id)
+
+
+def count_by_gate(gate, num_expert, world_size, require_pos=True, group=None):
+    total_expert_count = num_expert * world_size
+    with paddle.no_grad():
+        local_expert_count = _number_count(gate, total_expert_count)
+
+        if world_size > 1:
+            global_expert_count = _alltoall(local_expert_count, group=group)
+        else:
+            global_expert_count = local_expert_count
+        if not require_pos:
+            pos = None
+        else:
+            lec_cum = paddle.cumsum(local_expert_count, axis=0)
+            pos = _assign_pos(gate, lec_cum)
+    return pos, local_expert_count, global_expert_count
+
+
+def limit_by_capacity(topk_idx, num_expert, world_size, capacity, group=None):
+    with paddle.no_grad():
+        capacity = (
+            paddle.ones(shape=[num_expert], dtype=paddle.int64) * capacity
+        )
+        pos, lec, gec = count_by_gate(
+            topk_idx, num_expert, world_size, require_pos=False, group=group
+        )
+        new_gec = _limit_by_capacity(gec, capacity, world_size)
+        if world_size > 1:
+            assert group.nranks == world_size
+            new_lec = _alltoall(new_gec, group=group)
+        else:
+            new_lec = new_gec
+
+        topk_idx = _prune_gate_by_capacity(
+            topk_idx, new_lec, num_expert, world_size
+        )
+
+    return new_lec, new_gec, topk_idx

python/paddle/incubate/distributed/utils/__init__.py

@@ -0,0 +1,15 @@
+# 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 . import io as io

python/paddle/incubate/distributed/utils/io/__init__.py

@@ -0,0 +1,16 @@
+# 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 .dist_load import load  # noqa: F401
+from .dist_save import save, save_for_auto_inference  # noqa: F401

python/paddle/incubate/distributed/utils/io/dist_load.py

@@ -0,0 +1,121 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import re
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.framework import dygraph_only
+from paddle.distributed import fleet
+
+
+@dygraph_only
+def load(path, **configs):
+    """
+    Load an object can be used in paddle from specified path.
+    The file is saved by distributed.save
+
+    Note:
+        The file to load must be saved bu the API paddle.incubate.distributed.utils.io.save
+
+    Args:
+        path(str|BytesIO) : The path/buffer to load the target object. Generally, the path is the target
+            file path. When loading state_dict from the saved result of the API used to save
+            the inference model, the path may be a file prefix or directory.
+        **configs (dict, optional): other load configuration options for compatibility. We do not
+            recommend using these configurations, they may be removed in the future. If not necessary,
+            DO NOT use them. Default None.
+            The following options are currently supported:
+                (1) place: where to place the loaded state dict.
+                     If the state dict is too large, the place should be set 'cpu'.
+            Note:
+                Other config value may cause some error.Please don't use any more config options.
+    Returns:
+        Object(Object): a target object can be used in paddle
+
+    Examples:
+        import paddle
+        paddle.distributed.init_process_group(backend='nccl')
+        paddle.distributed.fleet.init(is_collective=True)
+
+        model = build_model()
+        optimizer = build_optimizer(model)
+
+        dist_model = paddle.distributed_optimizer(model)
+        dist_optimizer = paddle.distributed_optimizer(optimizer)
+
+
+        # load model state dict
+        model_state_dict = paddle.incubate.distributed.utils.io.load(path="path/to/load.pdparams")
+        dist_model.set_state_dict(model_state_dict)
+
+        # load optimizer state dict
+        optimizer_state_dict = paddle.incubate.distributed.utils.io.load(path="path/to/load.pdopt")
+        dist_optimizer.set_state_dict(optimizer_state_dict)
+
+    """
+    if dist.get_world_size() == 1:
+        return paddle.load(path, **configs)
+
+    hcg = fleet.get_hybrid_communicate_group()
+    assert (
+        hcg.get_model_parallel_world_size() == 1
+        and hcg.get_pipe_parallel_world_size() == 1
+    ), "Sharding and DP are supported only now"
+
+    # assert (
+    #     "place" in configs
+    # ), "the arg place ('cpu' or 'gpu:0', 'gpus:1' ...)must be passed"
+    if "place" not in configs:
+        configs["place"] = "cpu"
+    place = configs["place"]
+    assert isinstance(place, str), (
+        f"configs[place] must be a str, but this is a {type(place)}"
+    )
+
+    assert re.search("^(cpu|gpu:[0-9]*)$", place), (
+        "configs[place] must be cpu, gpu:0, gpu:1 ..."
+    )
+
+    return load_with_place(path, **configs)
+
+
+def load_with_place(path, **configs):
+    place = configs["place"]
+    if place is None:
+        return paddle.load(path)
+
+    origin_place = paddle.get_device()
+    paddle.set_device(place)
+
+    configs = _remove_not_supported_items(configs)
+    state_dict = paddle.load(path, **configs)
+
+    paddle.set_device(origin_place)
+
+    return state_dict
+
+
+def _remove_not_supported_items(configs):
+    __supported_by_load__ = [
+        "model_filename",
+        "params_filename",
+        "return_numpy",
+    ]
+    _configs = copy.copy(configs)
+    for k in configs.keys():
+        if k not in __supported_by_load__:
+            _configs.pop(k, None)
+    return _configs

python/paddle/incubate/distributed/utils/io/dist_save.py

@@ -0,0 +1,432 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import copy
+import re
+import sys
+from typing import TYPE_CHECKING, Any, Literal, TypedDict
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.framework import dygraph_only
+from paddle.distributed import fleet
+from paddle.distributed.fleet.utils.log_util import logger
+
+from .save_for_auto import save_for_auto_inference
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from io import BytesIO
+
+    from typing_extensions import Unpack
+
+    from paddle import Tensor
+    from paddle._typing import NestedStructure
+    from paddle.nn.layer.layers import _StateDict
+    from paddle.static import Program
+
+    class _SaveConfig(TypedDict, total=False):
+        use_binary_format: bool
+        gather_to: int | Sequence[int] | None
+        state_type: Literal['params', 'opt']
+        max_grouped_size: str | int
+
+
+__all__ = ["save", "save_for_auto_inference"]
+
+
+@dygraph_only
+def save(
+    state_dict: dict[str, Any] | _StateDict | NestedStructure[Tensor] | Program,
+    path: str | BytesIO,
+    **configs: Unpack[_SaveConfig],
+) -> None:
+    '''
+    Save a state dict to the specified path in both distributed and single-card environment.
+
+    Note:
+        Now supports saving ``state_dict`` of Layer/Optimizer, Tensor and nested structure containing Tensor, Program.
+
+    Note:
+        Different from ``paddle.jit.save``, since the save result of ``paddle.save`` is a single file,
+        there is no need to distinguish multiple saved files by adding a suffix. The argument ``path``
+        of ``paddle.save`` will be directly used as the saved file name instead of a prefix.
+        In order to unify the saved file name format, we recommend using the paddle standard suffix:
+        1. for ``Layer.state_dict`` , recommend to use ``.pdparams`` ;
+        2. for ``Optimizer.state_dict`` , recommend to use ``.pdopt`` .
+        For specific examples, please refer to API code examples.
+
+    Args:
+        obj(Object) : The object to be saved.
+        path(str|BytesIO) : The path/buffer of the object to be saved.
+            If saved in the current directory, the input path string will be used as the file name.
+        protocol(int, optional): The protocol version of pickle module must be greater than 1 and less than 5.
+            Default: 4.
+        **configs(dict, optional): optional keyword arguments. The following options are currently supported:
+
+            1. use_binary_format(bool):
+                To be used in paddle.save. When the saved object is static graph variable, you can specify ``use_binary_for_var``.
+                If True, save the file in the c++ binary format when saving a single static graph variable; otherwise, save it in pickle format.
+                Default: False.
+            2. gather_to(int|list|tuple|None):
+                To specify which global rank to save in.Default is None.
+                None value means distributed saving with no gathering to a single card.
+            3. state_type(str):
+                Value can be 'params' or 'opt', specifying to save parameters or optimizer state.
+            4. max_grouped_size(str|int):
+                To limit the max size(how many bits) a object group to be transferred a time.
+                If str, the format must be as num+'G/M/K', for example, 3G, 2K, 10M, etc. Default is 3G.
+
+    Returns:
+        None
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('TODO: the error will be fixed in the future')
+            >>> # type: ignore
+            >>> import paddle
+            >>> paddle.distributed.init_process_group(backend='nccl')
+            >>> paddle.distributed.fleet.init(is_collective=True)
+
+            >>> model = build_model()
+            >>> optimizer = build_optimizer(model)
+
+            >>> dist_optimizer = paddle.distributed_optimizer(optimizer)
+            >>> dist_model = paddle.distributed_optimizer(model)
+
+            >>> # gather params to rank 0 and then save
+            >>> paddle.incubate.distributed.utils.io.save(
+            ...     model.state_dict(), path="path/to/save.pdparams", gather_to=[0], state_type="params"
+            ... )
+
+            >>> # save whole params on all ranks
+            >>> paddle.incubate.distributed.utils.io.save(
+            ...     model.state_dict(), path="path/to/save.pdparams", gather_to=[0, 1], state_type="params"
+            ... )
+
+            >>> # save optimizer state dict on rank 0
+            >>> paddle.incubate.distributed.utils.io.save(optimizer.state_dict(), path="path/to/save.pdopt", gather=0, state_type="opt")
+
+    '''
+
+    gather_to = configs.get("gather_to", None)
+    if dist.get_world_size() == 1 or gather_to is None:
+        configs = _remove_not_supported_conf(configs)
+        return paddle.save(state_dict, path, **configs)
+
+    # gather_to is not None and world size > 1
+    state_type = configs.get("state_type", None)
+    assert isinstance(state_type, str), (
+        "must pass an arg state_type='params' or state_type='opt' to specify whether to save model state_dict or optimizer state_dict"
+    )
+    assert state_type in [
+        "params",
+        "opt",
+    ], "must pass an arg state_type='params' or state_type='opt'"
+
+    if re.search(f"{state_type}$", path) is None:
+        logger.warning(
+            f"You are saving {state_type}, while the path({path} does not end with {state_type})"
+        )
+
+    hcg = fleet.get_hybrid_communicate_group()
+    assert (
+        hcg.get_model_parallel_world_size() == 1
+        and hcg.get_pipe_parallel_world_size() == 1
+    ), (
+        f"Only DP and Sharding is supported now. However, current MP={hcg.get_model_parallel_world_size()} , PP={hcg.get_pipe_parallel_world_size()}"
+    )
+
+    sharding_group = hcg.get_sharding_parallel_group()
+    dp_group = hcg.get_data_parallel_group()
+
+    if state_type == "params":
+        if dp_group.nranks > 1:
+            assert _same_keys(state_dict, dp_group), (
+                "only sharding stage 1/2 and DP are supported now"
+            )
+        if sharding_group.nranks > 1:
+            assert _same_keys(state_dict, sharding_group), (
+                "only sharding stage 1/2 and DP are supported now"
+            )
+        configs = _remove_not_supported_conf(configs)
+        return paddle.save(state_dict, path, **configs)
+
+    # state_type == "opt"
+    if sharding_group.nranks == 1:
+        configs = _remove_not_supported_conf(configs)
+        return paddle.save(state_dict, path, **configs)
+    if _same_keys(state_dict, sharding_group):
+        return paddle.save(state_dict, path, **configs)
+    assert isinstance(gather_to, (list, tuple, int))
+    if isinstance(gather_to, int):
+        gather_to = [gather_to]
+    max_size = configs.get("max_grouped_size", "3G")
+    try:
+        logger.info("state_dict_keys:" + str(state_dict.keys()))
+        gathered_state_dict = _gather_state_dict(
+            state_dict, gather_to, sharding_group, max_size=max_size
+        )
+        logger.info("gathered_state_dict_keys:" + str(state_dict.keys()))
+        if dist.get_rank() in gather_to:
+            configs = _remove_not_supported_conf(configs)
+            paddle.save(gathered_state_dict, path, **configs)
+    except:
+        raise RuntimeError(
+            f'''Saving failed. Following are some suggestions:
+    1) pass the param max_grouped_size to turn the grouped size smaller (current value of max_grouped_size is {max_size})
+    2) if sharding stage is 1, use paddle.save rather than paddle.distributed.save
+    3) Concat the developers
+'''
+        )
+
+
+def _state_dict_groups(state_dict, max_size):
+    """
+    Description:
+        Generator of state dict groups to transfer.the size of each group is less than max_size.
+    """
+
+    # find the max size of a whole tensor
+    # now we only support to transfer  at least one whole tensor
+    max_tensor_size = 0
+    for k, v in state_dict.items():
+        if max_tensor_size < sys.getsizeof(v) + sys.getsizeof(k):
+            max_tensor_size = sys.getsizeof(v) + sys.getsizeof(k)
+
+    max_size = max(max_size, max_tensor_size)
+    logger.debug(f"max tensor size: {max_size}")
+
+    state_group = {}
+    k_list = list(state_dict.keys())
+    index = 0
+    bits = 0
+
+    # generate groups utils the end
+    while index < len(k_list):
+        bsize = sys.getsizeof(state_dict[k_list[index]]) + sys.getsizeof(
+            k_list[index]
+        )
+        if bits + bsize >= max_size:
+            yield state_group
+            state_group = {}
+            bits = 0
+
+        state_group[k_list[index]] = state_dict[k_list[index]]
+        index += 1
+        bits += bsize
+
+        if index == len(k_list) and bits > 0:
+            yield state_group
+
+
+def all_empty(dict_list):
+    """
+    Check if all items are empty
+    """
+    for v in dict_list:
+        if len(v) > 0:
+            return False
+    return True
+
+
+def _parse_mem_size_to_bits(max_size):
+    """
+    Parse an integer or a mem size str to an integer
+    convert xxxG to xxx * 1024^3
+    convert xxxM to xxx * 1024^2
+    convert xxxK to xxx * 1024^1
+    """
+    assert isinstance(max_size, (int, str))
+    if isinstance(max_size, str):
+        assert re.search("^[0-9]*[GMK]$", max_size), (
+            f"Wrong max_size 's format, the format ust be like 10K, 9M, 200G , etc, or an integer. However this is {max_size}"
+        )
+        num = int(max_size[:-1])
+        if max_size[-1] == "G":
+            max_size = num * 1024**3
+        elif max_size[-1] == "M":
+            max_size = num * 1024**2
+        else:
+            max_size = num * 1024
+    return max_size
+
+
+def _gather_state_dict(state_dict, dst, group, max_size="3G"):
+    """
+    Description:
+        Gather state dicts across all group ranks to dst, Depiring the same elements. including LR_Scheduler.
+    Args:
+        state_dict(dict):
+            local state dict
+        dst(int|list|tuple):
+            ranks the state dicts are gathered to
+        group(ProcessGroup):
+            group across which the state dicts are gathered
+        max_size(int|str):
+            The max limitation of the gathered tensor group size transformed a time. Default is 3G bits.
+            Each rank 's max tensor group before gathering is max_size // group.size
+    Returns:
+        Gathered state dict
+    """
+    assert isinstance(dst, (list, tuple, int)), (
+        "dst' type must be one of int, list and tuple"
+    )
+    if isinstance(dst, int):
+        dst = [dst]
+
+    max_size = _parse_mem_size_to_bits(max_size)
+    max_size //= dist.get_world_size(group)
+
+    logger.debug("len state_dict: len(state_dict)")
+
+    state_dict_ = copy.copy(state_dict)
+    mw = None
+    has_mw = False
+    has_lr = False
+
+    # Remove master_weights and LR_Scheduler to ensure that all the elements of the state dict are str->Tensor
+    if "master_weights" in state_dict_:
+        mw = state_dict_.pop("master_weights", None)
+        has_mw = True
+    if "LR_Scheduler" in state_dict_:
+        lr = state_dict_.pop("LR_Scheduler", None)
+        has_lr = True
+
+    # Gather optimizer state_dict
+    output = _grouped_gather_data_dict(state_dict_, dst, group, max_size)
+
+    # Gather master_weights if it exists
+    if isinstance(mw, dict):
+        masters = _grouped_gather_data_dict(mw, dst, group, max_size)
+    else:
+        assert mw is None, f"Wrong type of master weights . type: {type(mw)}"
+
+    # assign master_weights and LR_Scheduler
+    # Because LR_Schedulers are same across group, it just needs to be reset
+    if has_mw:
+        output["master_weights"] = masters
+    if has_lr:
+        output["LR_Scheduler"] = lr
+    return output
+
+
+def _grouped_gather_data_dict(state_data_dict, dst, group, max_size):
+    """
+    Description:
+        Gather state data dict by groups.
+    Args:
+        state__data_dict(dict):
+            local dict to transfer.The state_data_dict only contains the mapping: str->paddle.Tensor
+        dst(int|list|tuple):
+            ranks the state dicts are gathered to
+        group(ProcessGroup):
+            group across which the state dicts are gathered
+        max_size(int|str):
+            The max limitation of the gathered tensor group size transformed a time. Default is 3G bits.
+            Each rank 's max tensor group before gathering is max_size // group.size
+    Returns:
+        Gathered state_data_dict
+
+    """
+    numpy_dict = {}
+    logger.debug(f"len state_tict_ : {len(state_data_dict)}")
+
+    for k, v in state_data_dict.items():
+        try:
+            numpy_dict[k] = v.numpy()
+        except:
+            raise TypeError(
+                f"the object (type of {type(v)}) of '{k}' is neither tensor nor parameter"
+            )
+
+    total = 0
+    output_state = {}
+
+    logger.info("start all gather ...")
+    # gather all state_dict by groups
+    for state in _state_dict_groups(numpy_dict, max_size):
+        s_list = []
+        total += len(state)
+        logger.info(f"gen to gather: {total} / {len(numpy_dict)}")
+        dist.all_gather_object(s_list, state, group)
+        if dist.get_rank() in dst:
+            for s in s_list:
+                for k, v in s.items():
+                    logger.debug(f"gathered: {k}, {v.shape}")
+                output_state.update(s)
+
+        logger.debug(
+            f"s list size: {sum(len(s) for s in s_list)} output: {len(output_state)}"
+        )
+
+    # Because each size of groups may be different, here we should wait all objects gathered.
+    # The while block breaks until all objects from every rank are empty, which means all of the objects transforming is done.
+    while True:
+        s_list = []
+        state = {}
+        logger.debug("while True")
+        dist.all_gather_object(s_list, state, group)
+        if all_empty(s_list):
+            break
+        if dist.get_rank() in dst:
+            for s in s_list:
+                for k, v in s.items():
+                    logger.debug(f"gathered: {k}, {v.shape}")
+                output_state.update(s)
+        logger.debug(
+            f"s list size: {sum(len(s) for s in s_list)} output: {len(output_state)}"
+        )
+
+    logger.debug("all gathered ...")
+
+    if dist.get_rank() in dst:
+        # convert numpy.ndarray to Tensor in cpu place
+        place = paddle.CPUPlace()
+        for k in output_state.keys():
+            output_state[k] = paddle.to_tensor(output_state[k], place=place)
+            output_state[k].name = k
+        return output_state
+    return {}
+
+
+def _same_keys(state_dict, group):
+    """
+    Check whether all keys in each dict in the group are the same.
+    Used in sharding strategy to determine whether a dict needs to be gathered.
+    """
+    keys = list(state_dict.keys())
+    key_list = []
+    logger.info(keys)
+    dist.all_gather_object(key_list, keys, group=group)
+    for k in key_list:
+        if not k == keys:
+            return False
+    return True
+
+
+def _remove_not_supported_conf(configs):
+    """
+    Remove the config values not supported by paddle.save
+    """
+    __supported_by_save__ = ["use_binary_format"]
+    configs_ = copy.copy(configs)
+    for k in configs.keys():
+        if k not in __supported_by_save__:
+            configs_.pop(k, None)
+    return configs_

python/paddle/incubate/distributed/utils/io/save_for_auto.py

@@ -0,0 +1,368 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import copy
+import os
+import pickle
+import re
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.framework import dygraph_only
+from paddle.distributed import fleet
+from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_stage3 import (
+    GroupShardedStage3,
+)
+from paddle.distributed.fleet.utils.log_util import logger
+
+if TYPE_CHECKING:
+    from paddle.nn import Layer
+
+__all__ = ["save_for_auto_inference"]
+
+
+@dygraph_only
+def save_for_auto_inference(
+    path_prefix: str, dist_model: Layer, cvt2cpu: bool = False
+) -> None:
+    """
+    Description：
+        Save model parameters for auto parallel inference.
+        Supporting dp + mp + pp + sharding(stage1), dp + sharding stage2-3.
+        MoE not supported till MoE is supported in auto parallel mode.
+
+    Args:
+        path_prefix: path prefix to save. If `path_prefix` ends with path separator,
+            the path is processed as a directory and parameters will be saved in it,
+            automatically named saved_parameters. Otherwise, the parameters will be saved with name
+            path_prefix_dist{global_rank}.pdparams and path_prefix_dist{global_rank}.pdattrs.
+        dist_model: model in distributed model.
+        cvt2cpu: whether to move parameters to CPU when using sharding stage 3.
+            The var is invalid if not using sharding stage 3.
+
+    Returns:
+        None
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('model not exist')
+            >>> from paddle.incubate.distributed.utils.io import save_for_auto_inference
+            >>> dist_model = build_distributed_model()  # type: ignore[name-defined]
+            >>> path_prefix = "path/to/save_infer"
+            >>> save_for_auto_inference(path_prefix, dist_model=dist_model, cvt2cpu=False)
+
+    Outputs:
+        path/to/save_infer_dist0.pdparams path/to/save_infer_dist1.pdparams path/to/save_infer_dist2.pdparams ...
+        path/to/save_infer_dist0.pdattr  path/to/save_infer_dist1.pdattr   path/to/save_infer_dist2.pdattr ...
+
+    """
+
+    save_dir, basename_prefix = _get_abs_saved_prefix(path_prefix)
+
+    if isinstance(dist_model, GroupShardedStage3):
+        dist_model.get_all_parameters(cvt2cpu)
+
+    wrapped_dict = _get_wrapped_dist_state_dict(dist_model.state_dict())
+    global_rank = paddle.distributed.get_rank()
+
+    # save parameters
+    paddle.save(
+        wrapped_dict,
+        os.path.join(save_dir, f"{basename_prefix}_dist{global_rank}.pdparams"),
+    )
+
+    # save attributes
+    _save_param_attr(
+        wrapped_dict,
+        os.path.join(save_dir, f"{basename_prefix}_dist{global_rank}.pdattr"),
+    )
+
+    # unset dims mapping after saving attrs
+    for _, dist_param in wrapped_dict.items():
+        _unset_dims_mapping(dist_param)
+
+
+def _is_first_used(param):
+    return not hasattr(param, "is_firstly_shared") or param.is_firstly_shared
+
+
+def _get_all_ranks_of_pp(pp_rank, dp_degree, mp_degree, pp_degree):
+    """
+    Description:
+        get all global ranks involving given pp_rank
+    """
+
+    process_group = []
+
+    world_size = dp_degree * mp_degree * pp_degree
+
+    for i in range(dp_degree):
+        for k in range(mp_degree):
+            process_group.append(
+                i * world_size // dp_degree
+                + pp_rank * world_size // dp_degree // pp_degree
+                + k
+            )
+    return process_group
+
+
+def _save_param_attr(state_dict_, path, dims_mapping_dict=None):
+    """
+    Description:
+        save params' attr dict
+    Args:
+        state_dict_:
+            state for which to save attrs, when the state is optimizer state, the master and LRScheduler will be removed.
+        path:
+            path to save
+        dims_mapping_dict:
+            Dims mapping dict, mapping from parameter name in state_dict_ to dims_mapping.
+            If parameter in state_dict_ has attribute 'dims_mapping', the dims_mapping is ignored.
+            If parameter has no attribute 'dims_mapping', the dims mapping must contains the parameter's name.
+    """
+    state_dict = copy.copy(state_dict_)
+
+    # remove master_weights and LRScheduler, which needs no parameter attributes to save
+    state_dict.pop("master_weights", None)
+    state_dict.pop("LR_Scheduler", None)
+
+    if dims_mapping_dict is not None:
+        assert isinstance(dims_mapping_dict, dict), (
+            "dims_mapping_dict must be an instance of dict"
+        )
+        for k in state_dict.keys():
+            assert k in dims_mapping_dict, (
+                f"param {k} cannot find dims mapping in dims_mapping_dict"
+            )
+    if dist.get_world_size() > 1:
+        hcg = fleet.get_hybrid_communicate_group()
+        dp_degree = hcg.get_data_parallel_world_size()
+        mp_degree = hcg.get_model_parallel_world_size()
+        pp_degree = hcg.get_pipe_parallel_world_size()
+        sharding_degree = hcg.get_sharding_parallel_world_size()
+        dp_degree = dp_degree * sharding_degree
+
+        pp_group = hcg.get_pipe_parallel_group()
+    else:
+        pp_degree = 1
+        dp_degree = 1
+        mp_degree = 1
+        pp_group = None
+        hcg = None
+
+    logger.debug(f"dp degree * sharding degree : {dp_degree}")
+    logger.debug(f"mp degree: {mp_degree}")
+    logger.debug(f"pp degree: {pp_degree}")
+
+    pp_rank = dist.get_rank(pp_group)
+
+    # Why condition 'pp_rank < 0' exists?
+    # Because if pp_degree = 1, pp_rank is set -1
+    pp_rank = max(0, pp_rank)
+
+    if dist.get_world_size() > 1:
+        process_group = _get_all_ranks_of_pp(
+            pp_rank, dp_degree, mp_degree, pp_degree
+        )
+    else:
+        process_group = [0]
+
+    attr_dict = {}
+    for k, v in state_dict.items():
+        dims = len(v.shape)
+        logger.debug(f"shape: , {k}, {dims}")
+        attr_d = {
+            "process_shape": [dp_degree, mp_degree] if hcg else [1],
+            "process_group": process_group,
+            "dims_mapping": (
+                v.dims_mapping
+                if hasattr(v, "dims_mapping")
+                else [-1 for _ in v.shape]
+            ),
+        }
+        attr_dict[k] = attr_d
+
+    with open(path, "wb") as f:
+        pickle.dump(attr_dict, f)
+
+
+def _unset_dims_mapping(param):
+    if hasattr(param, "dims_mapping"):
+        delattr(param, "dims_mapping")
+
+
+def _get_dims_mapping(dist_parameter, mp_group):
+    """
+    Description:
+        return the splitting mapping:
+            {tensor_name: spiting_strategy}
+    Args:
+        dist_parameters(list): distributed model parameters
+        mp_group(ProcessGroup): Model Parallel communication group
+    Return:
+        The splitting mapping
+    Examples:
+        splitting_strategy's format (-1, -1, -1, 0), meaning the dims
+        of the tensor is 4 and it is splited along the first strategy axis in mesh
+
+    Mesh Examples: (2, 4) means dp=2, mp=4
+
+    """
+
+    import numpy as np
+
+    dist_shape = np.array(dist_parameter.shape)
+    if hasattr(dist_parameter, "split_axis"):
+        axis = dist_parameter.split_axis
+        mapping = [-1 for _ in dist_shape]
+        mapping[axis] = 1
+        logger.debug(
+            f"{dist_parameter.name} has attr split_axis: mapping: {mapping}"
+        )
+    else:
+        mapping = [-1 for _ in dist_shape]
+        logger.debug(f"normal parameter: {dist_parameter.name}")
+    return mapping
+
+
+def _get_abs_saved_prefix(path_prefix):
+    """
+    Description:
+        Get absolute dir path and basename prefix of path_prefix, with making path_prefix's directories.
+        If path_prefix is a directory name, basename is set 'saved_parameters'.
+        If path_prefix is a file name, basename is extracted from path_prefix.
+    Args:
+        path_prefix: str
+    Return:
+        (dirpath: str, basename: str)
+    """
+    abs_prefix = os.path.abspath(path_prefix)
+    if abs_prefix[-1] == os.path.sep:
+        save_dir = abs_prefix
+        basename_prefix = "saved_parameters"
+    else:
+        save_dir = os.path.dirname(abs_prefix)
+        basename_prefix = os.path.basename(abs_prefix)
+    os.makedirs(save_dir, exist_ok=True)
+    return save_dir, basename_prefix
+
+
+def _name_mapping_dist2single(state_dict, pp_group):
+    key_list = []
+    param_keys = [
+        v.name
+        for _, v in state_dict.items()
+        if isinstance(v, paddle.Tensor) and _is_first_used(v)
+    ]
+
+    if pp_group.nranks == 1:
+        return {k: k for k in param_keys}
+
+    dist.all_gather_object(key_list, param_keys, pp_group)
+
+    # find how many a op in a each pp:
+    # {"linear:"[0, 2,0,1,1,...]}
+    param_types = {}
+
+    matcher = re.compile(r"^\w+_\d+(?=\.)")
+
+    for pp, keys in enumerate(key_list):
+        param_type_idx = {}
+        for k in keys:
+            matched = matcher.search(k)
+            logger.debug(f"matched: {k}: {matched}")
+            assert matched is not None, (
+                f"the name of param, '{k}', is not satisfied the format 'name_idx.xxx'"
+            )
+            name_idx = k[matched.start() : matched.end()]
+            logger.debug(f"get param_type_idx: {name_idx}")
+
+            if name_idx in param_type_idx:
+                continue
+
+            name = "_".join(name_idx.split("_")[:-1])
+            idx = int(name_idx.split("_")[-1])
+            param_type_idx.update({name_idx: (name, idx)})
+            if name not in param_types:
+                param_types[name] = [0] * pp_group.nranks
+            param_types[name][pp] += 1
+
+        # check if continuous
+        types_idx = {}
+        for _, v in param_type_idx.items():
+            if v[0] not in types_idx:
+                types_idx.update({v[0]: [v[1]]})
+            else:
+                types_idx[v[0]].append(v[1])
+        for k, v in types_idx.items():
+            assert v == list(range(v[0], v[-1] + 1)), (
+                f"{k} is not continuous: {v}"
+            )
+
+    logger.debug(f"param type: {param_types}")
+
+    # analyse starting index
+    for k in param_types.keys():
+        param_types[k] = np.cumsum([0, *param_types[k][:-1]])
+
+    logger.debug(f"params type: {param_types}")
+
+    name_mapping = {}
+    pp_rank = dist.get_rank(pp_group)
+    for k in key_list[pp_rank]:
+        matched = matcher.search(k)
+        name_idx = k[matched.start() : matched.end()]
+        name = "_".join(name_idx.split("_")[:-1])
+        idx = int(name_idx.split("_")[-1])
+        logger.debug(f"idx: {idx}")
+
+        new_idx = param_types[name][pp_rank] + idx
+        logger.debug(f"new idx: {new_idx}")
+        new_name_idx = name + "_" + str(new_idx)
+        name_mapping[k] = new_name_idx + k[matched.end() :]
+
+    return name_mapping
+
+
+def _get_wrapped_dist_state_dict(dist_state_dict):
+    wrapped_state_dict = {}
+    if dist.get_world_size() <= 1:
+        for _, v in dist_state_dict.items():
+            wrapped_state_dict[v.name] = v
+        return wrapped_state_dict
+
+    hcg = fleet.get_hybrid_communicate_group()
+
+    pp_group = hcg.get_pipe_parallel_group()
+    mp_group = hcg.get_model_parallel_group()
+    logger.debug("execute _name_mapping_dist2single")
+
+    name_mapping = _name_mapping_dist2single(dist_state_dict, pp_group)
+    for _, v in dist_state_dict.items():
+        if not _is_first_used(v):
+            logger.debug(f"not first used : {v.name}")
+            continue
+        wrapped_state_dict[name_mapping[v.name]] = v
+        v.dims_mapping = _get_dims_mapping(v, mp_group)
+        logger.debug(
+            f"saving param: {v.name} -> {name_mapping[v.name]} shape: {v.shape}"
+        )
+    return wrapped_state_dict