chore: import upstream snapshot with attribution

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

@@ -0,0 +1,41 @@
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 NVIDIA Corporation. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+from .amp_optimizer import AMPOptimizer  # noqa: F401
+from .asp_optimizer import ASPOptimizer  # noqa: F401
+from .dgc_optimizer import (  # noqa: F401
+    DGCMomentumOptimizer,
+    DGCOptimizer,
+)
+from .dygraph_optimizer import (  # noqa: F401
+    HeterParallelOptimizer,
+    HybridParallelGradScaler,
+    HybridParallelOptimizer,
+)
+from .fp16_allreduce_optimizer import FP16AllReduceOptimizer  # noqa: F401
+from .gradient_merge_optimizer import GradientMergeOptimizer  # noqa: F401
+from .lamb_optimizer import LambOptimizer  # noqa: F401
+from .lars_optimizer import LarsOptimizer  # noqa: F401
+from .localsgd_optimizer import (  # noqa: F401
+    AdaptiveLocalSGDOptimizer,
+    LocalSGDOptimizer,
+)
+from .muon_sharding_optimizer import MuonShardingOptimizer  # noqa: F401
+from .pipeline_optimizer import PipelineOptimizer  # noqa: F401
+from .ps_optimizer import ParameterServerOptimizer  # noqa: F401
+from .qat_optimizer import QATOptimizer  # noqa: F401
+from .raw_program_optimizer import RawProgramOptimizer  # noqa: F401
+from .recompute_optimizer import RecomputeOptimizer  # noqa: F401
+from .sharding_optimizer import ShardingOptimizer  # noqa: F401
+from .tensor_parallel_optimizer import TensorParallelOptimizer  # noqa: F401

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

@@ -0,0 +1,139 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+import paddle.static.amp as mixed_precision
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class AMPOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.wrapped_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+        ]
+        self.meta_optimizers_black_list = ["DGCOptimizer"]
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _init_wrapped_opt(self):
+        if self.wrapped_opt is not None:
+            return
+
+        config = self.user_defined_strategy.amp_configs
+
+        custom_white_list = set(config['custom_white_list'])
+        custom_black_list = set(config['custom_black_list'])
+        custom_black_varnames = set(config['custom_black_varnames'])
+        amp_lists = mixed_precision.AutoMixedPrecisionLists(
+            custom_white_list, custom_black_list, custom_black_varnames
+        )
+
+        self.wrapped_opt = mixed_precision.decorate(
+            self.inner_opt,
+            amp_lists,
+            config['init_loss_scaling'],
+            config['incr_every_n_steps'],
+            config['decr_every_n_nan_or_inf'],
+            config['incr_ratio'],
+            config['decr_ratio'],
+            config['use_dynamic_loss_scaling'],
+            config['use_pure_fp16'],
+            config['use_fp16_guard'],
+        )
+
+        # if worker_num > 1, all cards will communication with each other,
+        # add is_distributed to optimize amp, overlap communication and
+        # computation by split the check_finite_and_unscale op.
+        is_distributed = self.role_maker._worker_num() > 1
+        if self.user_defined_strategy.sharding:
+            # FIXME(wangxi). sharding failed when split check_finite_and_unscale
+            # FIXME(JZ-LIANG). To support Sharding-Megatron-AMP, Megatron should follow Sharding's behavior that to disable is_distributed.
+            is_distributed = False
+        self.wrapped_opt._set_distributed(is_distributed)
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.amp:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.amp = False
+        dist_strategy.amp_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.amp = True
+        dist_strategy.amp_configs = {
+            "init_loss_scaling": 32768.0,
+            "incr_every_n_steps": 1000,
+            "decr_every_n_nan_or_inf": 2,
+            "incr_ratio": 2.0,
+            "decr_ratio": 0.8,
+            "use_dynamic_loss_scaling": True,
+        }
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        # maybe inner_opt of other meta optimizer
+        self._init_wrapped_opt()
+        return self.wrapped_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_gradients(self, params_grads):
+        return self.wrapped_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.wrapped_opt.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self._init_wrapped_opt()
+        optimize_ops, params_grads = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads
+
+    def amp_init(
+        self, place, scope=None, test_program=None, use_fp16_test=False
+    ):
+        return self.wrapped_opt.amp_init(
+            place, scope, test_program, use_fp16_test
+        )
+
+    def get_loss_scaling(self):
+        return self.wrapped_opt.get_loss_scaling()

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

@@ -0,0 +1,70 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 NVIDIA Corporation. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+
+from paddle.incubate.asp import ASPHelper
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class ASPOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+            "GradientMergeOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.asp:
+            return True
+
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.asp = False
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.asp = True
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = ASPHelper._minimize(
+            self.inner_opt,
+            loss,
+            startup_program=startup_program,
+            parameter_list=parameter_list,
+            no_grad_set=no_grad_set,
+        )
+
+        return optimize_ops, params_grads

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

@@ -0,0 +1,236 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import paddle
+from paddle.framework import core
+from paddle.utils import unique_name
+
+from ..base.private_helper_function import wait_server_ready
+
+__all__ = []
+
+OpRole = core.op_proto_and_checker_maker.OpRole
+
+OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
+OP_ROLE_VAR_KEY = core.op_proto_and_checker_maker.kOpRoleVarAttrName()
+
+
+def is_update_op(op):
+    return (
+        'Param' in op.input_names
+        and 'Grad' in op.input_names
+        and "LearningRate" in op.input_names
+    )
+
+
+def is_loss_grad_op(op):
+    if OP_ROLE_KEY not in op.attr_names:
+        return False
+    op_role = int(op.all_attrs()[OP_ROLE_KEY])
+    return op_role & int(OpRole.Backward) and op_role & int(OpRole.Loss)
+
+
+def is_backward_op(op):
+    return OP_ROLE_KEY in op.attr_names and int(
+        op.all_attrs()[OP_ROLE_KEY]
+    ) & int(OpRole.Backward)
+
+
+def is_optimizer_op(op):
+    return OP_ROLE_KEY in op.attr_names and int(
+        op.all_attrs()[OP_ROLE_KEY]
+    ) & int(OpRole.Optimize)
+
+
+class CollectiveHelper:
+    def __init__(self, role_maker, nrings=1, wait_port=True):
+        self.nrings = nrings
+        self.wait_port = wait_port
+        self.role_maker = role_maker
+
+    def update_startup_program(self, startup_program=None):
+        self.startup_program = startup_program
+        if startup_program is None:
+            self.startup_program = paddle.static.default_startup_program()
+
+        endpoints = self.role_maker._get_trainer_endpoints()
+        current_endpoint = endpoints[self.role_maker._worker_index()]
+        for ring_id in range(self.nrings):
+            self._init_communicator(
+                self.startup_program,
+                current_endpoint,
+                endpoints,
+                self.role_maker._worker_index(),
+                ring_id,
+                self.wait_port,
+            )
+        self._broadcast_params()
+
+    def _init_communicator(
+        self,
+        program,
+        current_endpoint,
+        endpoints,
+        rank,
+        ring_id,
+        wait_port,
+        global_ring_id=None,
+        sync=True,
+    ):
+        # if current_endpoint is None, it means just for sync,
+        # no group is created.
+        endpoints_str = ",".join(endpoints)
+        if current_endpoint:
+            nranks = len(endpoints)
+            other_endpoints = endpoints[:]
+            other_endpoints.remove(current_endpoint)
+
+        def _add_sync_by_allreduce(block):
+            sync_var = block.create_var(
+                name=unique_name.generate('sync_var'),
+                dtype=core.VarDesc.VarType.INT32,
+                persistable=False,
+                stop_gradient=True,
+            )
+            block.append_op(
+                type='fill_constant',
+                inputs={},
+                outputs={'Out': [sync_var]},
+                attrs={
+                    'shape': [1],
+                    'dtype': sync_var.dtype,
+                    'value': 1,
+                    'force_cpu': False,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='all_reduce',
+                inputs={'x': [sync_var]},
+                outputs={'out': [sync_var]},
+                attrs={
+                    'ring_id': global_ring_id,
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='c_sync_calc_stream',
+                inputs={'X': sync_var},
+                outputs={'Out': sync_var},
+                attrs={OP_ROLE_KEY: OpRole.Forward},
+            )
+
+        block = program.global_block()
+        if current_endpoint is None:
+            assert endpoints is None
+            assert sync
+            _add_sync_by_allreduce(block)
+            return
+
+        comm_id_var = block.create_var(
+            name=unique_name.generate('comm_id'),
+            persistable=True,
+            type=core.VarDesc.VarType.RAW,
+        )
+        if core.is_compiled_with_cuda():
+            block.append_op(
+                type='c_gen_nccl_id',
+                inputs={},
+                outputs={'Out': comm_id_var},
+                attrs={
+                    'rank': rank,
+                    'endpoint': current_endpoint,
+                    'other_endpoints': other_endpoints,
+                    'ring_id': ring_id,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='c_comm_init',
+                inputs={'X': comm_id_var},
+                outputs={},
+                attrs={
+                    'nranks': nranks,
+                    'rank': rank,
+                    'ring_id': ring_id,
+                    'endpoints': endpoints_str,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+        elif core.is_compiled_with_xpu():
+            block.append_op(
+                type='c_gen_bkcl_id',
+                inputs={},
+                outputs={'Out': comm_id_var},
+                attrs={
+                    'rank': rank,
+                    'endpoint': current_endpoint,
+                    'other_endpoints': other_endpoints,
+                    'ring_id': ring_id,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+            block.append_op(
+                type='c_comm_init',
+                inputs={'X': comm_id_var},
+                outputs={},
+                attrs={
+                    'nranks': nranks,
+                    'rank': rank,
+                    'ring_id': ring_id,
+                    'endpoints': endpoints_str,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+        else:
+            raise ValueError(
+                "comm_id must be generated in paddlepaddle-xpu or paddlepaddle-xpu."
+            )
+        if sync:
+            _add_sync_by_allreduce(block)
+
+    def _wait(self, current_endpoint, endpoints):
+        assert self.wait_port
+        other_endpoints = endpoints[:]
+        other_endpoints.remove(current_endpoint)
+        wait_server_ready(other_endpoints)
+
+    def _broadcast_params(self):
+        block = self.startup_program.global_block()
+        ring_id = -1
+        for param in block.iter_parameters():
+            if param.is_distributed:
+                continue
+
+            ring_id = (ring_id + 1) % self.nrings
+            block.append_op(
+                type='broadcast',
+                inputs={'x': param},
+                outputs={'out': param},
+                attrs={
+                    'ring_id': ring_id,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+        for ring_id in range(self.nrings):
+            block.append_op(
+                type='c_sync_comm_stream',
+                inputs={'X': param},
+                outputs={'Out': param},
+                attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Forward},
+            )

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

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

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

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

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

@@ -0,0 +1,65 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle.autograd as imperative_base
+from paddle import framework
+
+__all__ = []
+
+
+def _obtain_optimizer_parameters_list(optimizer):
+    if getattr(optimizer, '_param_groups', None) and isinstance(
+        optimizer._param_groups[0], dict
+    ):
+        parameters_list = []
+        for group in optimizer._param_groups:
+            for param in group['params']:
+                parameters_list.append(param)
+    else:
+        parameters_list = list(optimizer._parameter_list)
+
+    return parameters_list
+
+
+class HeterParallelOptimizer:
+    # adapter wrapper for optimizer
+    def __init__(self, optimizer, strategy):
+        self._inner_opt = optimizer
+        self._strategy = strategy
+
+        # NOTE(liubo48): In pure DataParallel mode,
+        # the gradient synchronization is achieved through reducer.
+
+    @imperative_base.no_grad()
+    @framework.dygraph_only
+    def step(self):
+        parameters_list = _obtain_optimizer_parameters_list(self._inner_opt)
+        self._inner_opt.step()
+
+    @imperative_base.no_grad()
+    def minimize(
+        self, loss, startup_program=None, parameters=None, no_grad_set=None
+    ):
+        # minimize does not support parameters in the form of param_group,
+        # so no need use _obtain_optimizer_parameters_list
+        parameter_list = (
+            parameters if parameters else self._inner_opt._parameter_list
+        )
+
+        return self._inner_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+
+    def __getattr__(self, item):
+        return getattr(self._inner_opt, item)

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

@@ -0,0 +1,85 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+import paddle.autograd as imperative_base
+from paddle import _legacy_C_ops
+
+from ...base.topology import ParallelMode
+
+__all__ = []
+
+
+class HybridParallelGradScaler:
+    def __init__(self, scaler, hcg):
+        self._scaler = scaler
+        self._hcg = hcg
+        self._use_dp_mode = (
+            self._hcg.get_parallel_mode() == ParallelMode.DATA_PARALLEL
+        )
+
+    def scale(self, var):
+        return self._scaler.scale(var)
+
+    def minimize(self, optimizer, *args, **kwargs):
+        if not self._enable:
+            return optimizer.minimize(*args, **kwargs)
+
+        #  unscale the grad
+        self._unscale(optimizer)
+
+        optimize_ops, params_grads = (None, None)
+
+        if hasattr(optimizer, "_set_auxiliary_var"):
+            optimizer._set_auxiliary_var('found_inf', self._found_inf)
+            optimize_ops, params_grads = optimizer.minimize(*args, **kwargs)
+            # TODO: Fix to _cache_found_inf after PaddleNLP update
+            self._cache_found_inf = optimizer._get_auxiliary_var('found_inf')
+        else:
+            if self._found_inf:
+                self._cache_found_inf = True
+            else:
+                optimize_ops, params_grads = optimizer.minimize(*args, **kwargs)
+                self._cache_found_inf = False
+
+        if self._use_dynamic_loss_scaling:
+            self._update()
+
+        return optimize_ops, params_grads
+
+    @imperative_base.no_grad()
+    def _unscale(self, optimizer):
+        if not self._enable:
+            return
+        param_grads = [
+            param._grad_ivar()
+            for param in optimizer._parameter_list
+            if param._grad_ivar() is not None
+        ]
+        _legacy_C_ops.check_finite_and_unscale(
+            param_grads, self._scale, param_grads, self._found_inf
+        )
+        # allreduce_max found_inf in check_group
+        if not self._use_dp_mode:
+            self._found_inf = paddle.cast(self._found_inf, dtype="int32")
+            # TODO(shenliang03) Since the minimize call in the optimizer is
+            # after the grad scaler, check_finite needs to synchronize global
+            # information. In the future, we should use check_group
+            paddle.distributed.all_reduce(
+                self._found_inf, op=paddle.distributed.ReduceOp.MAX, group=None
+            )
+            self._found_inf = paddle.cast(self._found_inf, dtype="bool")
+
+    def __getattr__(self, item):
+        return getattr(self._scaler, item)

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

@@ -0,0 +1,158 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+import paddle
+from paddle.framework import core
+from paddle.utils import unique_name
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class FP16AllReduceOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+            "LocalSGDOptimizer",
+            "GradientMergeOptimizer",
+            "AdaptiveLocalSGDOptimizer",
+        ]
+        self.meta_optimizers_black_list = ["DGCOptimizer"]
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.fp16_allreduce:
+            return True
+
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.fp16_allreduce = False
+
+    def _enable_strategy(self, dist_strategy, context=None):
+        dist_strategy.fp16_allreduce = True
+
+    @staticmethod
+    def fp16_compression(param_and_grads):
+        """
+        Compress fp32 gradients to fp16 during allreduce.
+        """
+        op_maker = core.op_proto_and_checker_maker
+
+        new_param_and_grads = []  # param, grad, is_cast
+        # cast grad from fp32->fp16 before allreduce,
+        for param, grad in param_and_grads:
+            if grad is None or grad.dtype != core.VarDesc.VarType.FP32:
+                new_param_and_grads.append((param, grad, False))
+                continue
+
+            op = grad.op
+            block = grad.block
+            var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()]
+            if param.name not in var_attr:
+                new_param_and_grads.append((param, grad, False))
+                continue
+
+            # remove (param, grad) from op_role_var
+            var_attr.remove(param.name)
+            var_attr.remove(grad.name)
+            if len(var_attr) > 1:
+                op._set_attr(op_maker.kOpRoleVarAttrName(), var_attr)
+            else:
+                op._remove_attr(op_maker.kOpRoleVarAttrName())
+
+            new_grad = block.create_var(
+                name=unique_name.generate(grad.name + ".cast_fp16"),
+                dtype=core.VarDesc.VarType.FP16,
+                persistable=False,
+                stop_gradient=True,
+            )
+
+            with block.program._backward_role_guard():
+                cast_op = block.append_op(
+                    type="cast",
+                    inputs={"X": grad},
+                    outputs={"Out": new_grad},
+                    attrs={
+                        "in_dtype": core.VarDesc.VarType.FP32,
+                        "out_dtype": core.VarDesc.VarType.FP16,
+                    },
+                    stop_gradient=True,
+                )
+
+                backward = op_maker.OpRole.Backward
+                cast_op._set_attr(op_maker.kOpRoleAttrName(), backward)
+                cast_op._set_attr(
+                    op_maker.kOpRoleVarAttrName(), [param.name, new_grad.name]
+                )
+                new_grad.op = cast_op
+
+            new_param_and_grads.append((param, new_grad, True))
+
+        ret_param_and_grads = []
+        # cast grad from fp16->fp32 after allreduce.
+        # NOTE. Now we split fp16 compression into two for loops,
+        # if we do not separate them, fuse allreduce will wrong.
+        # This must be the problem of fuse allreduce pass, need
+        # fixed in future.
+        for param, grad, cast in new_param_and_grads:
+            if not cast:
+                ret_param_and_grads.append((param, grad))
+                continue
+
+            block = grad.block
+            new_grad = block.create_var(
+                name=unique_name.generate(grad.name + ".cast_fp32"),
+                dtype=core.VarDesc.VarType.FP32,
+                persistable=False,
+                stop_gradient=True,
+            )
+
+            with (
+                block.program._optimized_guard([param, grad]),
+                paddle.static.name_scope('fp16_allreduce'),
+            ):
+                cast_op = block.append_op(
+                    type="cast",
+                    inputs={"X": grad},
+                    outputs={"Out": new_grad},
+                    attrs={
+                        "in_dtype": core.VarDesc.VarType.FP16,
+                        "out_dtype": core.VarDesc.VarType.FP32,
+                    },
+                    stop_gradient=True,
+                )
+            ret_param_and_grads.append((param, new_grad))
+
+        return ret_param_and_grads
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        new_params_grads = self.fp16_compression(params_grads)
+        return self.inner_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=new_params_grads
+        )

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

@@ -0,0 +1,75 @@
+#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+from paddle.incubate.optimizer import GradientMergeOptimizer as GM
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class GradientMergeOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.wrapped_opt = None
+        self.meta_optimizers_white_list = [
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+            "RecomputeOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _init_wrapped_opt(self):
+        config = self.user_defined_strategy.gradient_merge_configs
+        self.wrapped_opt = GM(self.inner_opt)
+        self.wrapped_opt._set_k_steps(
+            self.user_defined_strategy.gradient_merge_configs["k_steps"]
+        )
+        self.wrapped_opt._set_avg(
+            self.user_defined_strategy.gradient_merge_configs["avg"]
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        can_apply = (
+            self.user_defined_strategy.gradient_merge
+        ) and self.user_defined_strategy.gradient_merge_configs["k_steps"] > 1
+        return can_apply
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.gradient_merge = False
+        dist_strategy.gradient_merge_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        # we currently do not support auto-enable GradientMerge
+        return
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self._init_wrapped_opt()
+        optimize_ops, params_grads = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

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

@@ -0,0 +1,121 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+import logging
+
+import paddle
+from paddle.optimizer import Adam
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class LambOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.lamb_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = []
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+        opt = self.inner_opt
+        if not isinstance(opt, Adam):
+            return
+
+        configs = self.user_defined_strategy.lamb_configs
+        if len(configs['exclude_from_weight_decay']) == 0:
+            _exclude_from_weight_decay_fn = None
+        else:
+
+            def exclude_fn(param):
+                exclude_list = configs['exclude_from_weight_decay']
+                for name in exclude_list:
+                    if param.name.endswith(name):
+                        return True
+                return False
+
+            _exclude_from_weight_decay_fn = exclude_fn
+
+        self.lamb_opt = paddle.optimizer.Lamb(
+            learning_rate=opt._learning_rate,
+            lamb_weight_decay=configs['lamb_weight_decay'],
+            beta1=opt._beta1,
+            beta2=opt._beta2,
+            epsilon=opt._epsilon,
+            parameters=opt._parameter_list,
+            grad_clip=opt._grad_clip,
+            exclude_from_weight_decay_fn=_exclude_from_weight_decay_fn,
+            name=opt._name,
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.lamb:
+            if not isinstance(self.inner_opt, Adam):
+                logging.warning(
+                    f"lamb need the inner optimizer to be AdamOptimizer optimizer but got {self.inner_opt.type}."
+                )
+                return False
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.lamb = False
+        dist_strategy.lamb_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.lamb = True
+        dist_strategy.lamb_configs = {
+            "lamb_weight_decay": 0.01,
+            "exclude_from_weight_decay": [],
+        }
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self.lamb_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    # the following function will be used by AMP if both LARS and AMP are turn on together.
+    def apply_gradients(self, params_grads):
+        return self.lamb_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.lamb_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.lamb_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

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

@@ -0,0 +1,110 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+import logging
+
+from paddle.incubate.optimizer import LarsMomentumOptimizer
+from paddle.optimizer import Momentum
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class LarsOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.lars_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = []
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+        opt = self.inner_opt
+        if not isinstance(opt, Momentum):
+            return
+
+        configs = self.user_defined_strategy.lars_configs
+
+        self.lars_opt = LarsMomentumOptimizer(
+            learning_rate=opt._learning_rate,
+            momentum=opt._momentum,
+            lars_coeff=configs['lars_coeff'],
+            lars_weight_decay=configs['lars_weight_decay'],
+            parameter_list=opt._parameter_list,
+            regularization=opt.regularization,
+            grad_clip=opt._grad_clip,
+            name=opt._name,
+            exclude_from_weight_decay=configs['exclude_from_weight_decay'],
+            epsilon=configs['epsilon'],
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.lars:
+            if not isinstance(self.inner_opt, Momentum):
+                logging.warning(
+                    f"lars need the inner optimizer to be Momentum optimizer but got {self.inner_opt.type}."
+                )
+                return False
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.lars = False
+        dist_strategy.lars_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.lars = True
+        dist_strategy.lars_configs = {
+            "lars_coeff": 0.01,
+            "lars_weight_decay": 0.0005,
+        }
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self.lars_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    # the following function will be used by AMP if both LARS and AMP are turn on together.
+    def apply_gradients(self, params_grads):
+        return self.lars_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.lars_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.lars_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

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

@@ -0,0 +1,494 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+from paddle.static import (
+    default_main_program,
+    default_startup_program,
+    program_guard,
+)
+
+from .common import OP_ROLE_KEY, CollectiveHelper, OpRole
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class LocalSGDOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = ['AMPOptimizer']
+        self.meta_optimizers_black_list = [
+            "AdaptiveLocalSGDOptimizer",
+        ]
+        self.snapshot_key = '@SNAPSHOT'
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if not self.user_defined_strategy.localsgd:
+            return False
+
+        if self.role_maker._worker_num() <= 1:
+            return False
+
+        return isinstance(
+            self.inner_opt,
+            (
+                paddle.optimizer.momentum.Momentum,
+                paddle.optimizer.sgd.SGD,
+            ),
+        )
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.localsgd = False
+        dist_strategy.localsgd_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.localsgd = True
+        dist_strategy.localsgd_configs = {"k_steps": 1, "begin_step": 1}
+
+    def snapshot_name(self, param_name):
+        return param_name + self.snapshot_key
+
+    def create_snapshot_vars(self, program):
+        block = program.global_block()
+
+        non_dist_params = []
+        for param in block.iter_parameters():
+            if not param.is_distributed:
+                non_dist_params.append(param)
+
+        p2s = []
+        for param in non_dist_params:
+            snapshot = block.create_var(
+                name=self.snapshot_name(param.name),
+                shape=param.shape,
+                persistable=True,
+                stop_gradient=True,
+                dtype=param.dtype,
+            )
+            p2s.append([param, snapshot])
+        return p2s
+
+    def init_snapshot_vars(self, startup_program, param2snapshot):
+        with program_guard(startup_program):
+            for param, snapshot in param2snapshot:
+                paddle.assign(param, snapshot)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        minimized = self.inner_opt.minimize(
+            loss, startup_program=startup_program
+        )
+
+        k_steps_value = self.user_defined_strategy.localsgd_configs['k_steps']
+        begin_step_value = self.user_defined_strategy.localsgd_configs[
+            'begin_step'
+        ]
+
+        if startup_program is None:
+            startup_program = default_startup_program()
+        main_block = loss.block
+
+        self.nrings = 2
+        collective_helper = CollectiveHelper(self.role_maker, self.nrings)
+        collective_helper.update_startup_program(startup_program)
+        p2s = self.create_snapshot_vars(startup_program)
+        self.init_snapshot_vars(startup_program, p2s)
+
+        p2s = self.create_snapshot_vars(main_block.program)
+        with program_guard(main_block.program, startup_program):
+            step = paddle.optimizer.lr.autoincreased_step_counter(begin=1)
+            k_steps = paddle.static.create_global_var(
+                name="k_steps",
+                shape=[1],
+                value=k_steps_value,
+                dtype='int64',
+                persistable=True,
+            )
+
+            begin_step = paddle.static.create_global_var(
+                name="begin_step",
+                shape=[1],
+                value=begin_step_value,
+                dtype='int64',
+                persistable=True,
+            )
+
+            last_step = paddle.static.create_global_var(
+                name="last_step",
+                shape=[1],
+                value=begin_step_value,
+                dtype='int64',
+                persistable=True,
+            )
+
+            def communicate():
+                sub_block = default_main_program().current_block()
+                ring_id = -1
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='c_sync_calc_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    ring_id = (ring_id + 1) % self.nrings
+                    sub_block.append_op(
+                        type='all_reduce',
+                        inputs={'x': [param]},
+                        outputs={'out': [param]},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for ring_id in range(self.nrings):
+                    sub_block.append_op(
+                        type='c_sync_comm_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={
+                            'ring_id': ring_id,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='scale',
+                        inputs={'X': [param]},
+                        outputs={'Out': [param]},
+                        attrs={
+                            'scale': 1.0 / self.role_maker._worker_num(),
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='assign',
+                        inputs={'X': [param]},
+                        outputs={'Out': [snapshot]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                paddle.assign(step, last_step)
+
+            def begin_localsgd():
+                paddle.static.nn.cond(step - last_step == k_steps, communicate)
+
+            paddle.static.nn.cond(
+                step > begin_step, begin_localsgd, communicate
+            )
+        return minimized
+
+
+class AdaptiveLocalSGDOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = ['AMPOptimizer']
+        self.meta_optimizers_black_list = [
+            "LocalSGDOptimizer",
+        ]
+        self.snapshot_key = '@SNAPSHOT'
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if not self.user_defined_strategy.adaptive_localsgd:
+            return False
+
+        if self.role_maker._worker_num() <= 1:
+            return False
+
+        return isinstance(
+            self.inner_opt,
+            (
+                paddle.optimizer.Momentum,
+                paddle.optimizer.sgd.SGD,
+            ),
+        )
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.adaptive_localsgd = False
+        dist_strategy.adaptive_localsgd_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.adaptive_localsgd = True
+        dist_strategy.adaptive_localsgd_configs = {
+            "init_k_steps": 1,
+            "begin_step": 1,
+        }
+
+    def snapshot_name(self, param_name):
+        return param_name + self.snapshot_key
+
+    def create_snapshot_vars(self, program):
+        block = program.global_block()
+
+        non_dist_params = []
+        for param in block.iter_parameters():
+            if not param.is_distributed:
+                non_dist_params.append(param)
+
+        p2s = []
+        for param in non_dist_params:
+            snapshot = block.create_var(
+                name=self.snapshot_name(param.name),
+                shape=param.shape,
+                persistable=True,
+                stop_gradient=True,
+                dtype=param.dtype,
+            )
+            p2s.append([param, snapshot])
+        return p2s
+
+    def init_snapshot_vars(self, startup_program, param2snapshot):
+        with program_guard(startup_program):
+            for param, snapshot in param2snapshot:
+                paddle.assign(param, snapshot)
+
+    def _generate_avg_loss(self, program_block, loss, avg_loss):
+        program_block.append_op(
+            type='all_reduce',
+            inputs={'x': [loss]},
+            outputs={'out': [avg_loss]},
+            attrs={
+                'ring_id': 0,
+                OP_ROLE_KEY: OpRole.Optimize,
+                'reduce_type': paddle.distributed.ReduceOp.SUM,
+            },
+        )
+        program_block.append_op(
+            type='c_sync_calc_stream',
+            inputs={'X': [avg_loss]},
+            outputs={'Out': [avg_loss]},
+            attrs={OP_ROLE_KEY: OpRole.Optimize},
+        )
+
+        program_block.append_op(
+            type='scale',
+            inputs={'X': [avg_loss]},
+            outputs={'Out': [avg_loss]},
+            attrs={
+                'scale': 1.0 / self.role_maker._worker_num(),
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        minimized = self.inner_opt.minimize(
+            loss, startup_program=startup_program
+        )
+
+        init_k_steps = self.user_defined_strategy.adaptive_localsgd_configs[
+            'init_k_steps'
+        ]
+        begin_step_value = self.user_defined_strategy.adaptive_localsgd_configs[
+            'begin_step'
+        ]
+
+        if startup_program is None:
+            startup_program = default_startup_program()
+        main_block = loss.block
+
+        self.nrings = 2
+        collective_helper = CollectiveHelper(self.role_maker, self.nrings)
+        collective_helper.update_startup_program(startup_program)
+        p2s = self.create_snapshot_vars(startup_program)
+        self.init_snapshot_vars(startup_program, p2s)
+
+        p2s = self.create_snapshot_vars(main_block.program)
+        with program_guard(main_block.program, startup_program):
+            step = paddle.optimizer.lr.autoincreased_step_counter(begin=1)
+
+            k_steps = paddle.static.create_global_var(
+                name="k_steps",
+                shape=[1],
+                value=int(init_k_steps),
+                dtype='int64',
+                persistable=True,
+            )
+
+            begin_step = paddle.static.create_global_var(
+                name="begin_step",
+                shape=[1],
+                value=int(begin_step_value),
+                dtype='int64',
+                persistable=True,
+            )
+
+            last_step = paddle.static.create_global_var(
+                name="last_step",
+                shape=[1],
+                value=0,
+                dtype='int64',
+                persistable=True,
+            )
+
+            avg_loss = paddle.static.create_global_var(
+                name="avg_loss",
+                shape=[1],
+                value=float(0),
+                dtype=loss.dtype,
+                persistable=True,
+            )
+
+            lr_0 = paddle.static.create_global_var(
+                name="lr_0",
+                shape=[1],
+                value=float(0),
+                dtype='float32',
+                persistable=True,
+            )
+
+            loss_0 = paddle.static.create_global_var(
+                name="loss_0",
+                shape=[1],
+                value=float(0),
+                dtype='float32',
+                persistable=True,
+            )
+
+            global_lr = self.inner_opt._global_learning_rate()
+
+            def initialize():
+                self._generate_avg_loss(main_block, loss, avg_loss)
+                paddle.assign(avg_loss, loss_0)
+                paddle.assign(global_lr, lr_0)
+
+            paddle.static.nn.cond(step == 1, initialize)
+
+            def communicate():
+                sub_block = default_main_program().current_block()
+                ring_id = -1
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='c_sync_calc_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    ring_id = (ring_id + 1) % self.nrings
+                    sub_block.append_op(
+                        type='all_reduce',
+                        inputs={'x': [param]},
+                        outputs={'out': [param]},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for ring_id in range(self.nrings):
+                    sub_block.append_op(
+                        type='c_sync_comm_stream',
+                        inputs={'X': param},
+                        outputs={'Out': param},
+                        attrs={
+                            'ring_id': ring_id,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+
+                for param, snapshot in p2s:
+                    sub_block.append_op(
+                        type='scale',
+                        inputs={'X': [param]},
+                        outputs={'Out': [param]},
+                        attrs={
+                            'scale': 1.0 / self.role_maker._worker_num(),
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    sub_block.append_op(
+                        type='elementwise_sub',
+                        inputs={'X': [snapshot], 'Y': [param]},
+                        outputs={'Out': [param]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                    sub_block.append_op(
+                        type='assign',
+                        inputs={'X': [param]},
+                        outputs={'Out': [snapshot]},
+                        attrs={OP_ROLE_KEY: OpRole.Optimize},
+                    )
+                paddle.assign(step, last_step)
+
+            def communicate_avg_loss():
+                communicate()
+                self._generate_avg_loss(main_block, loss, avg_loss)
+
+                next_local_steps = paddle.cast(
+                    paddle.ceil(
+                        paddle.sqrt(
+                            lr_0
+                            * avg_loss
+                            / (global_lr * loss_0)
+                            * float(init_k_steps)
+                        )
+                    ),
+                    dtype='int64',
+                )
+                max_local_steps = paddle.full(
+                    shape=[1], dtype='int64', fill_value=16
+                )
+                min_local_steps = paddle.full(
+                    shape=[1], dtype='int64', fill_value=1
+                )
+                next_local_steps = paddle.minimum(
+                    next_local_steps, max_local_steps
+                )
+                next_local_steps = paddle.maximum(
+                    next_local_steps, min_local_steps
+                )
+                paddle.assign(next_local_steps, k_steps)
+
+            def begin_localsgd():
+                paddle.static.nn.cond(
+                    step - last_step == k_steps, communicate_avg_loss
+                )
+
+            paddle.static.nn.cond(
+                step > begin_step, begin_localsgd, communicate
+            )
+
+        return minimized

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

@@ -0,0 +1,106 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.optimizer import Optimizer
+
+__all__ = []
+
+
+class MetaOptimizerBase(Optimizer):
+    def __init__(self, optimizer):
+        self.inner_opt = optimizer
+        self._learning_rate = self.inner_opt._learning_rate
+        self._learning_rate_map = self.inner_opt._learning_rate_map
+        self.meta_optimizers_white_list = []
+        self.meta_optimizers_black_list = []
+
+    def _set_auxiliary_var(self, key, val):
+        super()._set_auxiliary_var(key, val)
+        self.inner_opt._set_auxiliary_var(key, val)
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        self.loss = loss
+        self.role_maker = role_maker
+        self.user_defined_optimizer = user_defined_optimizer
+        self.user_defined_strategy = user_defined_strategy
+
+    def _update_inner_optimizer(self, optimizer):
+        self.inner_opt = optimizer
+
+    def _can_apply(self):
+        return False
+
+    def _is_graph_out(self):
+        return False
+
+    def _can_update(self, optimizer):
+        if str(optimizer.__class__.__name__) in self.meta_optimizers_white_list:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        raise NotImplementedError(
+            f"you should implement disable strategy in {type(self).__name__}"
+        )
+
+    def _enable_strategy(self, dist_strategy, context=None):
+        raise NotImplementedError(
+            f"you should implement enable strategy in {type(self).__name__}"
+        )
+
+    def apply_gradients(self, params_grads):
+        return self.inner_opt.apply_gradients(params_grads=params_grads)
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        return self.inner_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.inner_opt._apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        params_grads = self.backward(
+            loss,
+            startup_program=startup_program,
+            parameter_list=parameter_list,
+            no_grad_set=no_grad_set,
+        )
+
+        optimize_ops = self.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+        return optimize_ops, params_grads
+
+    def minimize(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.minimize_impl(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

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

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

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

@@ -0,0 +1,319 @@
+#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+import paddle
+from paddle.incubate.optimizer import PipelineOptimizer as PO
+
+from .common import (
+    OP_ROLE_KEY,
+    OP_ROLE_VAR_KEY,
+    CollectiveHelper,
+    OpRole,
+    is_backward_op,
+    is_loss_grad_op,
+)
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class PipelineOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = [
+            "RecomputeOptimizer",
+            "AMPOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+        self.global_ring_id = 1
+        self.dp_ring_id = 2
+        self.start_pipeline_ring_id = 20  # Just a magic number
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+        self.micro_batch_size = user_defined_strategy.pipeline_configs[
+            'micro_batch_size'
+        ]
+        self.num_microbatches = user_defined_strategy.pipeline_configs[
+            'accumulate_steps'
+        ]
+        self.schedule_mode = user_defined_strategy.pipeline_configs[
+            'schedule_mode'
+        ]
+        self.use_sharding = user_defined_strategy.sharding
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        # FIXME revise for hybrid parallelism
+        if self.use_sharding:
+            return False
+
+        if self.user_defined_strategy.pipeline:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.pipeline = False
+        dist_strategy.pipeline_configs = {
+            "micro_batch_size": 1,
+            "accumulate_steps": 1,
+            "schedule_mode": "1F1B",
+        }
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.pipeline = True
+        dist_strategy.pipeline_configs = {
+            "micro_batch_size": 1,
+            "accumulate_steps": 1,
+            "schedule_mode": "1F1B",
+        }
+
+    def _broadcast_params(self, ring_id):
+        block = self.startup_program.global_block()
+        param = None
+        for param in block.iter_parameters():
+            if param.is_distributed:
+                continue
+
+            block.append_op(
+                type='broadcast',
+                inputs={'x': param},
+                outputs={'out': param},
+                attrs={
+                    'ring_id': ring_id,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+        if not param:
+            return  # no parameter on this device
+        block.append_op(
+            type='c_sync_comm_stream',
+            inputs={'X': param},
+            outputs={'Out': param},
+            attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Forward},
+        )
+
+    def _get_process_group_info(self):
+        # global ring info
+        self.global_endpoints = self.endpoints
+        self.global_rank = self.rank
+        self.global_nranks = self.nranks
+
+        # data parallel ring info
+        if self.pipeline_num > 1:
+            self.dp_rank = self.rank // self.inner_parallelism
+            self.dp_nranks = self.nranks // self.inner_parallelism
+            start_index = self.rank % self.inner_parallelism
+            self.dp_endpoints = [
+                self.endpoints[start_index + i * self.inner_parallelism]
+                for i in range(self.pipeline_num)
+            ]
+
+    def _init_process_group(self, pipeline_pair, pipeline_ring_map):
+        self._get_process_group_info()
+        collective_helper = CollectiveHelper(self.role_maker, wait_port=False)
+        # Create global ring for all gpus (ring_id = 0)
+        collective_helper._init_communicator(
+            self.startup_program,
+            self.current_endpoint,
+            self.global_endpoints,
+            self.global_rank,
+            self.global_ring_id,
+            True,
+            self.global_ring_id,
+            True,
+        )
+        # Create pipeline rings
+        if self.inner_parallelism > 1:
+            pipeline_id = self.rank // self.inner_parallelism
+            start_index = pipeline_id * self.inner_parallelism
+            for pair in pipeline_pair:
+                pair_key = pair[0] * 1000 + pair[1]
+                ring_id = pipeline_ring_map[pair_key]
+                assert ring_id >= self.start_pipeline_ring_id
+                first_node = pair[0] + start_index
+                second_node = pair[1] + start_index
+                if self.rank != first_node and self.rank != second_node:
+                    collective_helper._init_communicator(
+                        self.startup_program,
+                        None,
+                        None,
+                        None,
+                        None,
+                        False,
+                        self.global_ring_id,
+                        True,
+                    )
+                    continue
+                pipeline_endpoints = [
+                    self.endpoints[first_node],
+                    self.endpoints[second_node],
+                ]
+                pipeline_rank = 0 if self.rank == first_node else 1
+                pipeline_nranks = 2
+                collective_helper._init_communicator(
+                    self.startup_program,
+                    self.current_endpoint,
+                    pipeline_endpoints,
+                    pipeline_rank,
+                    ring_id,
+                    False,
+                    self.global_ring_id,
+                    True,
+                )
+
+        # Create dp rings
+        if self.pipeline_num > 1:
+            collective_helper._init_communicator(
+                self.startup_program,
+                self.current_endpoint,
+                self.dp_endpoints,
+                self.dp_rank,
+                self.dp_ring_id,
+                True,
+                self.global_ring_id,
+                True,
+            )
+            self._broadcast_params(self.dp_ring_id)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self.endpoints = self.role_maker._get_trainer_endpoints()
+        self.current_endpoint = self.endpoints[self.role_maker._worker_index()]
+        self.rank = self.role_maker._worker_index()
+        self.nranks = self.role_maker._worker_num()
+
+        self.wrapped_opt = PO(
+            self.inner_opt, num_microbatches=self.num_microbatches
+        )
+        orig_startup_program = (
+            startup_program
+            if startup_program
+            else paddle.static.default_startup_program()
+        )
+        block = loss.block
+        program = block.program
+
+        program._pipeline_opt = {}
+        program._pipeline_opt['local_rank'] = self.rank
+        program._pipeline_opt['global_ring_id'] = self.global_ring_id
+        program._pipeline_opt['ring_id'] = self.start_pipeline_ring_id
+        program._pipeline_opt['micro_batch_size'] = self.micro_batch_size
+        program._pipeline_opt['schedule_mode'] = self.schedule_mode
+        program._pipeline_opt['use_sharding'] = False
+        program._pipeline_opt['mp_degree'] = 1
+        program._pipeline_opt['mp_rank'] = 0
+        (
+            optimize_ops,
+            params_grads,
+            prog_list,
+            pp_pair,
+            ring_map,
+        ) = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        self.startup_program = orig_startup_program._pipeline_opt[
+            'startup_program'
+        ]
+        self.inner_parallelism = program._pipeline_opt['inner_parallelism']
+        assert self.nranks % self.inner_parallelism == 0
+        assert prog_list
+        self.pipeline_num = len(self.endpoints) // self.inner_parallelism
+
+        self._init_process_group(pp_pair, ring_map)
+
+        self.main_program_list = prog_list
+        self.main_program = program
+        if self.pipeline_num > 1:
+            self._transpile_main_program(loss)
+        return optimize_ops, params_grads
+
+    def _transpile_main_program(self, loss):
+        self._insert_loss_grad_ops(loss, self.pipeline_num)
+        self._insert_allreduce_ops(self.dp_ring_id)
+
+    def _insert_loss_grad_ops(self, loss, pipeline_num):
+        """
+        In order to keep the learning rate consistent in different numbers of
+        training workers, we scale the loss grad by the number of workers
+        """
+        block = self.main_program_list[-1].global_block()
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_loss_grad_op(op):
+                loss_grad_var = block.vars[op.output_arg_names[0]]
+                block._insert_op(
+                    idx + 1,
+                    type='scale',
+                    inputs={'X': loss_grad_var},
+                    outputs={'Out': loss_grad_var},
+                    attrs={
+                        'scale': 1.0 / pipeline_num,
+                        OP_ROLE_KEY: OpRole.Backward,
+                    },
+                )
+
+    def _insert_allreduce_ops(self, ring_id):
+        block = self.main_program._pipeline_opt[
+            'section_program'
+        ].global_block()
+        origin_block = self.main_program.global_block()
+        grad = None
+        processed_param_name = set()
+        first_optimize_op_idx = None
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_backward_op(op) and not first_optimize_op_idx:
+                first_optimize_op_idx = idx + 1
+                # no optimize phase
+                if first_optimize_op_idx == len(block.ops):
+                    return
+            if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
+                op_role_var = op.all_attrs()[OP_ROLE_VAR_KEY]
+                if len(op_role_var) == 0:
+                    continue
+                assert len(op_role_var) % 2 == 0
+                offset = 0
+                for i in range(0, len(op_role_var), 2):
+                    param_name = op_role_var[i]
+                    param = block.vars[op_role_var[i]]
+                    if param_name in processed_param_name:
+                        continue
+                    processed_param_name.add(param_name)
+                    grad_name = op_role_var[i + 1]
+                    if 'MERGED' not in grad_name:
+                        grad_name += '@MERGED'
+                    grad = block.vars[grad_name]
+                    origin_param = origin_block.vars[op_role_var[i]]
+                    if origin_param.is_distributed:
+                        continue
+
+                    block._insert_op(
+                        first_optimize_op_idx + offset,
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )

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

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

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

@@ -0,0 +1,123 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+import copy
+
+import paddle
+from paddle.static.quantization.quanter import (
+    _quant_config_default,
+    quant_aware,
+)
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+
+class QATOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+            "GraphExecutionOptimizer",
+            "RecomputeOptimizer",
+            "GradientMergeOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.qat:
+            return True
+
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.qat = False
+        dist_strategy.qat_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.qat = True
+        dist_strategy.qat_configs = {
+            'channel_wise_abs_max': True,
+            'weight_bits': 8,
+            'activation_bits': 8,
+            'not_quant_pattern': [],
+            'algo': "",
+        }
+
+    def _gen_qat_config(self):
+        # Align the config to auto_parallel quantization pass
+        config = self.user_defined_strategy.qat_configs
+        qat_config = copy.deepcopy(_quant_config_default)
+        qat_config['quantize_op_types'] = [
+            'conv2d',
+            'depthwise_conv2d',
+            'mul',
+            'matmul',
+            'matmul_v2',
+        ]
+        qat_config['weight_quantize_type'] = (
+            'channel_wise_abs_max'
+            if config['channel_wise_abs_max']
+            else 'abs_max'
+        )
+        qat_config['weight_bits'] = config['weight_bits']
+        qat_config['activation_bits'] = config['activation_bits']
+        qat_config['not_quant_pattern'] = list(config['not_quant_pattern'])
+        return qat_config
+
+    def _replace_program(self, main_program, refer_program):
+        main_program._rebuild_from_desc(refer_program.desc)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        optimize_ops, params_grads = self.inner_opt.minimize(
+            loss,
+            startup_program,
+            parameter_list,
+            no_grad_set,
+        )
+        device = paddle.device.get_device()
+        place = paddle.set_device(device)
+        qat_config = self._gen_qat_config()
+        qat_program = quant_aware(
+            loss.block.program, place, config=qat_config, return_program=True
+        )
+        self._replace_program(loss.block.program, qat_program)
+        return optimize_ops, params_grads
+
+    def qat_init(self, place, scope=None, test_program=None):
+        if test_program is not None:
+            qat_config = self._gen_qat_config()
+            qat_program = quant_aware(
+                test_program,
+                place,
+                scope=scope,
+                config=qat_config,
+                for_test=True,
+                return_program=True,
+            )
+            self._replace_program(test_program, qat_program)

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

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

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

@@ -0,0 +1,104 @@
+#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+from paddle.incubate.optimizer import RecomputeOptimizer as RO
+
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class RecomputeOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.wrapped_opt = None
+        # we do not allow meta optimizer to be inner optimizer currently
+        self.meta_optimizers_white_list = [
+            "LarsOptimizer",
+            "LambOptimizer",
+            "DGCOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+
+    def _init_wrapped_opt(self):
+        if self.wrapped_opt is not None:
+            return
+
+        configs = self.user_defined_strategy.recompute_configs
+        self.wrapped_opt = RO(self.inner_opt)
+        self.wrapped_opt._set_checkpoints(list(configs["checkpoints"]))
+        if configs["enable_offload"]:
+            self.wrapped_opt._enable_offload()
+            # TODO(JZ-LIANG) might found a way to infer the checkpoint shape automatically
+            checkpoint_shapes = list(configs["checkpoint_shape"])
+            self.wrapped_opt.checkpoint_shape = checkpoint_shapes
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.recompute:
+            if (
+                len(self.user_defined_strategy.recompute_configs["checkpoints"])
+                == 0
+            ):
+                return False
+            else:
+                return True
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.recompute = False
+        dist_strategy.recompute_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        # we do not support automatically recompute checkpoints currently
+        return
+
+    def backward(
+        self,
+        loss,
+        startup_program=None,
+        parameter_list=None,
+        no_grad_set=None,
+        callbacks=None,
+    ):
+        # maybe inner_opt of other meta optimizer
+        self._init_wrapped_opt()
+        return self.wrapped_opt.backward(
+            loss, startup_program, parameter_list, no_grad_set, callbacks
+        )
+
+    def apply_gradients(self, params_grads):
+        return self.wrapped_opt.apply_gradients(params_grads=params_grads)
+
+    def apply_optimize(self, loss, startup_program, params_grads):
+        return self.wrapped_opt.apply_optimize(
+            loss, startup_program=startup_program, params_grads=params_grads
+        )
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self._init_wrapped_opt()
+        optimize_ops, params_grads = self.wrapped_opt.minimize(
+            loss, startup_program, parameter_list, no_grad_set
+        )
+        return optimize_ops, params_grads

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

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

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

@@ -0,0 +1,273 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+from paddle.distributed.fleet.meta_optimizers.common import (
+    OP_ROLE_KEY,
+    OpRole,
+    is_optimizer_op,
+)
+from paddle.framework import core
+
+__all__ = []
+
+
+class FP16Utils:
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def is_fp16_cast_op(block, op, params):
+        if op.type != "cast":
+            return False
+        if is_optimizer_op(op):
+            return False
+        assert len(op.desc.input_arg_names()) == 1
+        assert len(op.desc.output_arg_names()) == 1
+        input_name, output_name = (
+            op.desc.input_arg_names()[0],
+            op.desc.output_arg_names()[0],
+        )
+        if input_name not in params:
+            return False
+        input_var = block.var(input_name)
+        output_var = block.var(output_name)
+        if (
+            input_var.dtype != core.VarDesc.VarType.FP32
+            or output_var.dtype != core.VarDesc.VarType.FP16
+        ):
+            return False
+        return True
+
+    @staticmethod
+    def is_fp32_cast_op(block, op):
+        if op.type != "cast":
+            return False
+        if not is_optimizer_op(op):
+            return False
+        assert len(op.desc.input_arg_names()) == 1
+        assert len(op.desc.output_arg_names()) == 1
+        input_name, output_name = (
+            op.desc.input_arg_names()[0],
+            op.desc.output_arg_names()[0],
+        )
+        input_var = block.var(input_name)
+        output_var = block.var(output_name)
+        if (
+            input_var.dtype != core.VarDesc.VarType.FP16
+            or output_var.dtype != core.VarDesc.VarType.FP32
+        ):
+            return False
+        return True
+
+    @staticmethod
+    def remove_cast_op(block, params, segment, offset):
+        inserted_op_num = 0
+        for op_idx in reversed(
+            range(offset + segment._start_idx, offset + segment._end_idx)
+        ):
+            op = block.ops[op_idx]
+            if FP16Utils.is_fp16_cast_op(block, op, params):
+                block._remove_op(op_idx, sync=False)
+                inserted_op_num -= 1
+        block._sync_with_cpp()
+        return inserted_op_num
+
+    @staticmethod
+    def prune_fp16(block, shard, reduced_grads_to_param, ring_ids):
+        """
+        1. prune all cast_fp16_to_fp32 ops if the param not belongs to this shard
+        2. revise amp inifine grad checking for sharding
+        """
+        # remove cast
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not FP16Utils.is_fp32_cast_op(block, op):
+                continue
+            output_name = op.desc.output_arg_names()[0]
+            # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+            param_name = output_name.removesuffix("@MERGED").removesuffix(
+                "@GRAD"
+            )
+            if param_name not in shard.global_params:
+                raise ValueError(
+                    "Output 'X' of cast_op must be a grad of"
+                    f"model param, but {output_name} is not a grad"
+                )
+            if output_name in reduced_grads_to_param:
+                continue
+            if shard.has_param(param_name):
+                continue
+            block._remove_op(idx, sync=False)
+            block._remove_var(output_name, sync=False)
+
+        block._sync_with_cpp()
+        update_loss_scaling_op_idx = -1
+        inf_var_name = ''
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if op.type == "update_loss_scaling":
+                update_loss_scaling_op_idx = idx
+                inf_var_name = op.desc.input('FoundInfinite')[0]
+            if op.type in ["check_finite_and_unscale", "update_loss_scaling"]:
+                reversed_x = []
+                reversed_x_paramname = []
+                for input_name in op.desc.input('X'):
+                    # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+                    param_name = input_name.removesuffix(
+                        "@MERGED"
+                    ).removesuffix("@GRAD")
+                    if param_name not in shard.global_params:
+                        raise ValueError(
+                            "Input 'X' of check_finite_and_unscale must"
+                            f"be grads, but {input_name} is not a grad"
+                        )
+                    if shard.has_param(param_name):
+                        reversed_x.append(input_name)
+                        reversed_x_paramname.append(param_name)
+                op.desc.set_input('X', reversed_x)
+                op.desc.set_output('Out', reversed_x)
+
+                # the grad checking should take the all and only param in the current shard
+                to_check_param = set(reversed_x_paramname)
+                should_check_param = set(shard.global_params).intersection(
+                    {
+                        param
+                        for param, worker_idx in shard.global_param2device.items()
+                        if worker_idx == shard.worker_idx
+                    }
+                )
+                assert to_check_param == should_check_param, (
+                    f"amp \
+                    check_finite_and_unscale checking miss [{should_check_param - to_check_param}] and got unexpected [{to_check_param - should_check_param}]"
+                )
+
+        if update_loss_scaling_op_idx == -1:
+            return
+        inf_var = block.var(inf_var_name)
+        inf_var_int32 = block.create_var(
+            name=inf_var_name + "@cast_int32",
+            shape=inf_var.shape,
+            dtype=core.VarDesc.VarType.INT32,
+        )
+
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            type='cast',
+            inputs={'X': inf_var},
+            outputs={'Out': inf_var_int32},
+            attrs={
+                "in_dtype": inf_var.dtype,
+                "out_dtype": inf_var_int32.dtype,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+
+        # allreduce(mp)->allreduce(sharding)->allreduce(pp)
+        for ring_id in ring_ids:
+            if ring_id == -1:
+                continue
+            # this allreduce communication should not overlap with calc
+            block._insert_op_without_sync(
+                update_loss_scaling_op_idx,
+                type='all_reduce',
+                inputs={'x': inf_var_int32},
+                outputs={'out': inf_var_int32},
+                attrs={
+                    'ring_id': ring_id,
+                    'op_type': paddle.distributed.ReduceOp.MAX,
+                    OP_ROLE_KEY: OpRole.Optimize,
+                },
+            )
+            update_loss_scaling_op_idx += 1
+
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            type='cast',
+            inputs={'X': inf_var_int32},
+            outputs={'Out': inf_var},
+            attrs={
+                "in_dtype": inf_var_int32.dtype,
+                "out_dtype": inf_var.dtype,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+        block._sync_with_cpp()
+
+    # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+    @staticmethod
+    def sync_amp_check_nan_inf(block, ring_ids):
+        update_loss_scaling_op_idx = -1
+
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if op.type == "update_loss_scaling":
+                update_loss_scaling_op_idx = idx
+                inf_var_name = op.desc.input('FoundInfinite')[0]
+                break
+
+        # not use amp
+        if update_loss_scaling_op_idx == -1:
+            return
+        # 0. inf_var_int32 = cast(inf_var)
+        # 1. inf_var_int32 = allreduce_max(inf_var_int32)
+        # 3. inf_var = cast(inf_var_int32)
+        inf_var = block.var(inf_var_name)
+        inf_var_int32 = block.create_var(
+            name=inf_var_name + "@cast_int32",
+            shape=inf_var.shape,
+            dtype=core.VarDesc.VarType.INT32,
+        )
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            type='cast',
+            inputs={'X': inf_var},
+            outputs={'Out': inf_var_int32},
+            attrs={
+                "in_dtype": inf_var.dtype,
+                "out_dtype": inf_var_int32.dtype,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+
+        # allreduce(mp)->allreduce(pp)
+        for ring_id in ring_ids:
+            if ring_id == -1:
+                continue
+            block._insert_op_without_sync(
+                update_loss_scaling_op_idx,
+                type='all_reduce',
+                inputs={'x': inf_var_int32},
+                outputs={'out': inf_var_int32},
+                attrs={
+                    'ring_id': ring_id,
+                    'op_type': paddle.distributed.ReduceOp.MAX,
+                    OP_ROLE_KEY: OpRole.Optimize,
+                },
+            )
+            update_loss_scaling_op_idx += 1
+
+        block._insert_op_without_sync(
+            update_loss_scaling_op_idx,
+            type='cast',
+            inputs={'X': inf_var_int32},
+            outputs={'Out': inf_var},
+            attrs={
+                "in_dtype": inf_var_int32.dtype,
+                "out_dtype": inf_var.dtype,
+                OP_ROLE_KEY: OpRole.Optimize,
+            },
+        )
+        update_loss_scaling_op_idx += 1
+        block._sync_with_cpp()

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

@@ -0,0 +1,259 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
+
+__all__ = []
+
+
+class GradientClipHelper:
+    def __init__(self, mp_ring_id):
+        self.mp_ring_id = mp_ring_id
+
+    def _is_gradient_clip_op(self, op):
+        return op.desc.has_attr("op_namescope") and op.desc.attr(
+            "op_namescope"
+        ).startswith("/gradient_clip")
+
+    def prune_gradient_clip(self, block, shard, ring_ids):
+        """
+        prune gradient_clip related ops for params that not belong to cur shard
+        prune: square, reduce_sum, elementwise_mul
+        keep: sum, sqrt, elementwise_max, elementwise_div
+        """
+        deprecated_vars = set()
+        deprecate_op_idx = set()
+        reversed_x_paramname = []
+        global_norm_sum_op_idx = -1
+        for idx, op in enumerate(block.ops):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == "sum":
+                global_norm_sum_op_idx = idx
+                continue
+            deprecate_op = False
+            for input_name in op.desc.input_arg_names():
+                if input_name in deprecated_vars:
+                    deprecate_op = True
+                # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+                param_name = input_name.removesuffix("@MERGED").removesuffix(
+                    "@GRAD"
+                )
+                if shard.is_param(param_name) and not shard.has_param(
+                    param_name
+                ):
+                    deprecate_op = True
+                elif shard.is_param(param_name):
+                    reversed_x_paramname.append(param_name)
+
+            if deprecate_op:
+                deprecate_op_idx.add(idx)
+                for output_name in op.desc.output_arg_names():
+                    if output_name not in op.desc.input_arg_names():
+                        deprecated_vars.add(output_name)
+
+        # NOTE(wangxi): If only have 2 sharding, and 1 param.
+        # sharding 0 will not deprecated_vars, will return, only
+        # sharding 1 will insert allreduce, then hang.
+        if not deprecated_vars and global_norm_sum_op_idx == -1:
+            # got no gradient_clip op
+            return
+
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if idx in deprecate_op_idx:
+                block._remove_op(idx, sync=False)
+                continue
+            if op.type == "sum":
+                reversed_inputs = []
+                global_norm_sum_op_idx = idx
+                for input_name in op.desc.input_arg_names():
+                    if input_name not in deprecated_vars:
+                        reversed_inputs.append(input_name)
+
+                op.desc.set_input("X", reversed_inputs)
+                assert len(op.desc.output_arg_names()) == 1
+                sum_res = op.desc.output_arg_names()[0]
+
+                # NOTE(wangxi): If we have 2 param, but sharding is 4,
+                # then the sum op in some cards will not have input.
+                # So we use fill_constant_op to set `sum_var` to zero,
+                # which does not affect correctness.
+                if len(reversed_inputs) == 0:
+                    sum_var = block.var(sum_res)
+                    namescope = op.attr("op_namescope")
+
+                    block._remove_op(idx, sync=False)
+                    op = block._insert_op_without_sync(
+                        idx,
+                        type='fill_constant',
+                        inputs={},
+                        outputs={'Out': sum_res},
+                        attrs={
+                            'shape': sum_var.shape,
+                            'dtype': sum_var.dtype,
+                            'value': 0.0,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    op._set_attr('op_namescope', namescope)
+
+                # allreduce(mp)->allreduce(sharding)->allreduce(pp)
+                idx_offset = 1
+                for ring_id in ring_ids:
+                    if ring_id == -1:
+                        continue
+                    # this allreduce should not overlap with calc and should be scheduled in calc stream
+                    block._insert_op_without_sync(
+                        idx + idx_offset,
+                        type='all_reduce',
+                        inputs={'x': sum_res},
+                        outputs={'out': sum_res},
+                        attrs={
+                            'ring_id': ring_id,
+                            'op_namescope': "/gradient_clip_model_parallelism",
+                            'reduce_type': paddle.distributed.ReduceOp.Sum,
+                            OP_ROLE_KEY: OpRole.Optimize,
+                        },
+                    )
+                    idx_offset += 1
+
+        # the grad sum here should take the all and only param in the current shard
+        to_check_param = set(reversed_x_paramname)
+        should_check_param = set(shard.global_params).intersection(
+            {
+                param
+                for param, worker_idx in shard.global_param2device.items()
+                if worker_idx == shard.worker_idx
+            }
+        )
+        assert to_check_param == should_check_param, (
+            f"amp check_finite_and_unscale \
+        checking miss [{should_check_param - to_check_param}] and got unexpected [{to_check_param - should_check_param}]"
+        )
+
+        for var_name in deprecated_vars:
+            block._remove_var(var_name, sync=False)
+        block._sync_with_cpp()
+        return
+
+    # TODO (JZ-LIANG) revise this for uniform mixed parallelism
+    def sync_global_norm(self, block, ring_ids, mp_rank):
+        """
+        prune gradient_clip related ops for params that not belong to cur shard
+        prune: square, reduce_sum, elementwise_mul
+        keep: sum, sqrt, elementwise_max, elementwise_div
+        """
+        is_clip_grad_by_global_norm = False
+        for idx, op in list(enumerate(block.ops)):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == 'sum':
+                is_clip_grad_by_global_norm = True
+                break
+        if not is_clip_grad_by_global_norm:
+            # TODO(Yuang Liu): need some extra handles when clip_grad_norm for mp
+            return
+
+        removed_op_idx = set()
+        removed_tmp_var = set()
+        for idx, op in list(enumerate(block.ops)):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == 'sum':
+                break
+            for input_name in op.input_arg_names:
+                input_var = block.var(input_name)
+                # NOTE: when mp_degree > 1, some vars will be split into each mp rank.
+                # However, there still some vars such as Scale, Bias are not split.
+                # Those not be split vars should only be counted once during grad clip
+                # by global norm. Those vars either doesn't have is_distributed attr
+                # or the is_distributed attr has been set as False.
+                # Therefore, we prune those duplicated vars for grad clip.
+                if mp_rank >= 1 and (
+                    not (
+                        hasattr(input_var, 'is_distributed')
+                        and input_var.is_distributed
+                    )
+                ):
+                    removed_op_idx.add(idx)
+                    for output_name in op.output_arg_names:
+                        removed_tmp_var.add(output_name)
+
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if idx in removed_op_idx:
+                block._remove_op(idx, sync=False)
+
+        for var_name in removed_tmp_var:
+            block._remove_var(var_name, sync=False)
+
+        for idx, op in list(enumerate(block.ops)):
+            if not self._is_gradient_clip_op(op):
+                continue
+            if op.type == 'sum':
+                # If mp_rank == 0, no extra handles, just allreduce
+                # If mp_rank >= 1, some extra handles is needed
+                sum_rst_var = block.var(op.output_arg_names[0])
+                if mp_rank >= 1:
+                    reserved_vars = []
+                    for input_name in op.input_arg_names:
+                        if input_name not in removed_tmp_var:
+                            reserved_vars.append(input_name)
+
+                    if len(reserved_vars) > 0:
+                        op.desc.set_input("X", reserved_vars)
+                    else:
+                        # If all input of sum op should be removed, then remove the sum op.
+                        # And set the output's value of sum to 0.
+                        namescope = op.attr("op_namescope")
+                        block._remove_op(idx, sync=False)
+                        fill_constant_op = block._insert_op_without_sync(
+                            idx,
+                            type='fill_constant',
+                            inputs={},
+                            outputs={'Out': sum_rst_var},
+                            attrs={
+                                'shape': sum_rst_var.shape,
+                                'dtype': sum_rst_var.dtype,
+                                'value': 0.0,
+                                OP_ROLE_KEY: OpRole.Optimize,
+                            },
+                        )
+                        fill_constant_op._set_attr('op_namescope', namescope)
+                self._insert_allreduce(block, ring_ids, idx, sum_rst_var)
+                break
+
+    @staticmethod
+    def _insert_allreduce(block, ring_ids, idx, var):
+        for ring_id in ring_ids:
+            if ring_id == -1:
+                continue
+
+            idx = idx + 1
+            block._insert_op_without_sync(
+                idx,
+                type='all_reduce',
+                inputs={'x': var},
+                outputs={'out': var},
+                attrs={
+                    'ring_id': ring_id,
+                    'op_namescope': "/gradient_clip_model_parallelism",
+                    'reduce_type': paddle.distributed.ReduceOp.SUM,
+                    OP_ROLE_KEY: OpRole.Optimize,
+                },
+            )

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

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

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

@@ -0,0 +1,153 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__all__ = []
+
+
+class ProgramDeps:
+    def __init__(self, block, start_vars, end_vars):
+        self._block = block
+        # vars where to start to build the deps
+        self._start_vars = start_vars
+        # vars where to stop to build the deps
+        self._end_vars = end_vars
+        # var name -> op idxs which depends on this var
+        self._var_to_use_op = {}
+        # sub block deps which is a subset of this topo
+        self._sub_block_deps = {}
+        # var name -> op idxs which generate var
+        self._var_to_generate_op = {}
+        self._should_removed_var = set()
+        self._father_block_deps = None
+        self._build_deps()
+
+    def get_sub_block_deps(self, idx):
+        if idx in self._sub_block_deps:
+            return self._sub_block_deps[idx]
+        else:
+            return None
+
+    def get_var_deps(self, var_name):
+        if var_name in self._var_to_use_op:
+            return self._var_to_use_op[var_name]
+        else:
+            return None
+
+    def _build_deps(
+        self,
+    ):
+        for var_name in self._start_vars:
+            self._var_to_use_op[var_name] = []
+            self._var_to_generate_op[var_name] = []
+
+        for idx, op in enumerate(self._block.ops):
+            if op.type in [
+                "c_sync_comm_stream",
+                "c_calc_comm_stream",
+                'all_reduce',
+            ]:
+                continue
+            input_vars = op.desc.input_arg_names()
+            output_vars = op.desc.output_arg_names()
+            deps_reduce = False
+            for input_name in input_vars:
+                if input_name in self._var_to_use_op:
+                    deps_reduce = True
+            if not deps_reduce:
+                continue
+            for input_name in input_vars:
+                if input_name in self._var_to_use_op:
+                    self._var_to_use_op[input_name].append(idx)
+            for output_name in output_vars:
+                if output_name not in self._var_to_use_op:
+                    self._var_to_use_op[output_name] = []
+                if output_name not in self._var_to_generate_op:
+                    self._var_to_generate_op[output_name] = [idx]
+                else:
+                    self._var_to_generate_op[output_name].append(idx)
+            if op.type == "conditional_block":
+                # subblock
+                assert op.desc.has_attr("sub_block")
+                subblock_idx = op.desc.attr("sub_block").id
+                subblock_deps = ProgramDeps(
+                    self._block.program.block(subblock_idx),
+                    op.desc.input_arg_names(),
+                    op.desc.output_arg_names(),
+                )
+                self._sub_block_deps[subblock_idx] = subblock_deps
+                subblock_deps._father_block_deps = self
+
+    def crop_input_var_from_op(self, op_idx, var_name):
+        if var_name in self._var_to_use_op:
+            # update var -> dep_var_op
+            if self._var_to_use_op[var_name] != []:
+                if op_idx not in self._var_to_use_op[var_name]:
+                    raise ValueError(
+                        f"op_idx: {op_idx} is not in self._var_to_use_op[{var_name}], "
+                        f"self._var_to_use_op[{var_name}] is {self._var_to_use_op[var_name]}"
+                    )
+                self._var_to_use_op[var_name].remove(op_idx)
+            # update _should_removed_var
+            if var_name in self._start_vars:
+                self._should_removed_var.discard(var_name)
+            elif (
+                self._var_to_use_op[var_name] == []
+            ):  # no more deps of this var
+                self._should_removed_var.add(var_name)
+            elif (
+                self._var_to_generate_op[var_name][-1]
+                >= self._var_to_use_op[var_name][-1]
+            ):
+                # there are circle in the graph
+                self._should_removed_var.add(var_name)
+            else:  # input_name should not be deleted
+                self._should_removed_var.discard(var_name)
+
+    def crop_output_var_from_op(self, op_idx, var_name):
+        if var_name in self._var_to_generate_op:
+            assert op_idx in self._var_to_generate_op[var_name]
+            self._var_to_generate_op[var_name].remove(op_idx)
+        if self._block.has_var(var_name):
+            if (
+                var_name not in self._var_to_generate_op
+                or self._var_to_generate_op[var_name] == []
+            ):
+                self._block._remove_var(var_name, sync=False)
+
+    def remove_op(self, op_idx, reserved_vars=None):
+        # update deps
+        op = self._block.ops[op_idx]
+        for input_name in op.desc.input_arg_names():
+            if reserved_vars is not None and input_name in reserved_vars:
+                continue
+            self.crop_input_var_from_op(op_idx, input_name)
+        for output_name in op.desc.output_arg_names():
+            if reserved_vars is not None and output_name in reserved_vars:
+                continue
+            self.crop_output_var_from_op(op_idx, output_name)
+        self._block._remove_op(op_idx, sync=False)
+
+    def should_remove_op(self, op_idx):
+        op = self._block.ops[op_idx]
+
+        # NOTE: At present, it is found that the OP without output is
+        # only send_v2 and partial_send op, which will be used in
+        # all device
+        if len(op.desc.output_arg_names()) == 0:
+            return False
+
+        for output_name in op.desc.output_arg_names():
+            if output_name not in self._should_removed_var:
+                return False
+        return True

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

@@ -0,0 +1,175 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+from paddle.distributed.fleet.meta_optimizers.common import is_optimizer_op
+from paddle.distributed.fleet.meta_optimizers.sharding.fp16_helper import (
+    FP16Utils,
+)
+from paddle.distributed.fleet.meta_optimizers.sharding.utils import get_var_size
+
+__all__ = []
+
+
+class Shard:
+    def __init__(
+        self,
+    ):
+        self.global_params = set()
+        self.worker_idx = -1
+        self.worker_num = -1
+        self.global_param2device = {}
+        self.device2global_params = {}
+
+    def setup(self, params_grads, worker_idx, worker_num):
+        # param names of all devices
+        self.global_params = {x[0].name for x in params_grads}
+        # _param(str) -> device_id(int)
+        self.worker_idx = worker_idx
+        self.worker_num = worker_num
+        # global_param2device contains fp32 params and fp16 params
+        # device2global_params only contains fp32 params
+        (
+            self.global_param2device,
+            self.device2global_params,
+        ) = self._split_params(params_grads, worker_idx, worker_num)
+
+    def has_param(self, var_name):
+        return (
+            var_name in self.global_param2device
+            and self._var_device_id(var_name) == self.worker_idx
+        )
+
+    def has_opt_var(self, var_name):
+        return self._var_device_id(var_name) == self.worker_idx
+
+    def has_var(self, var_name):
+        return (
+            self._var_device_id(var_name) == -1
+            or self._var_device_id(var_name) == self.worker_idx
+        )
+
+    def _split_params(self, params_grads, worker_idx, worker_num):
+        param2device = {}
+        total_param_mem = 0.0
+        param2mem = []
+        for param in [x[0] for x in params_grads]:
+            mem = get_var_size(param)
+            total_param_mem += mem
+            param2mem.append((param.name, mem))
+        device2params = {x: [] for x in range(worker_num)}
+        device_idx = 0
+        mem_accu = 0.0
+        for param_name, mem in param2mem:
+            if mem_accu > total_param_mem * 1.0 * (device_idx + 1) / worker_num:
+                device_idx += 1
+            device2params[device_idx].append(param_name)
+            param2device[param_name] = device_idx
+            mem_accu += mem
+        return param2device, device2params
+
+    def _var_device_id(self, var_name):
+        if var_name in self.global_param2device:
+            return self.global_param2device[var_name]
+        for suffix in [
+            "_moment1_0",
+            "_moment2_0",
+            "_beta1_pow_acc_0",
+            "_beta2_pow_acc_0",
+            "_velocity_0",
+        ]:
+            base_name = re.sub(suffix, '', var_name)
+            if base_name in self.global_param2device:
+                return self.global_param2device[base_name]
+        return -1
+
+    def find_broadcast_params(self, block):
+        broadcast_vars = set()
+        fp16_params = set()
+        fp16_to_fp32 = {}
+
+        param_usage = dict.fromkeys(self.global_params, 0)
+        for op in block.ops:
+            if is_optimizer_op(op):
+                continue
+            for input_name in op.desc.input_arg_names():
+                if input_name in self.global_params:
+                    param_usage[input_name] += 1
+
+        for op in block.ops:
+            if not FP16Utils.is_fp16_cast_op(block, op, self.global_params):
+                continue
+            input_name = op.input_arg_names[0]
+            output_name = op.output_arg_names[0]
+            broadcast_vars.add(output_name)
+            fp16_params.add(output_name)
+            fp16_to_fp32[output_name] = input_name
+            param_usage[input_name] -= 1
+            self.global_param2device[output_name] = self.global_param2device[
+                input_name
+            ]
+
+        for param, usage in param_usage.items():
+            if usage > 0:
+                broadcast_vars.add(param)
+        return broadcast_vars
+
+    def device(self, var_name):
+        return self._var_device_id(var_name)
+
+    def is_param(self, var_name):
+        return var_name in self.global_params
+
+    def is_opti_var(self, var_name):
+        if var_name in self.global_params:
+            return True
+        for suffix in [
+            "_moment1_0",
+            "_moment2_0",
+            "_beta1_pow_acc_0",
+            "_beta2_pow_acc_0",
+            "_velocity_0",
+        ]:
+            base_name = re.sub(suffix, '', var_name)
+            if base_name in self.global_params:
+                return True
+        return False
+
+    def filter_grads(self, grads):
+        grads_in_shard = []
+        for grad in grads:
+            param = grad.split("@")[0]
+            if self.has_param(param):
+                grads_in_shard.append(grad)
+        return grads_in_shard
+
+
+class ProgramSegment:
+    def __init__(self, block):
+        self._block = block
+        self._allreduce_vars = []
+        # sub program start idx
+        self._start_idx = -1
+        # sub program end idx
+        self._end_idx = -1
+        # param name to broadcast name
+        self._param2broadcast = {}
+        self._broadcast_vars = []
+        # cast op pairs, fp16 name (str) -> fp32 name (str)
+        self._cast_ops = {}
+        # fill constant vars
+        self._fill_constant_vars = []
+        # parameter mems
+        self._param_mem = 0.0

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

@@ -0,0 +1,41 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_VAR_KEY
+
+__all__ = []
+
+
+class WeightDecayHelper:
+    def __init__(self):
+        pass
+
+    def _is_weight_decay_op(self, op):
+        return op.desc.has_attr("op_namescope") and op.desc.attr(
+            "op_namescope"
+        ).startswith("/regularization")
+
+    def prune_weight_decay(self, block, shard):
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if not self._is_weight_decay_op(op):
+                continue
+            if OP_ROLE_VAR_KEY not in op.attr_names:
+                raise ValueError(
+                    "The Weight Decay op should hold op_role_var attribute"
+                    f"but the {op.type} op does not hold op_role_var"
+                )
+            op_role_var = op.all_attrs()[OP_ROLE_VAR_KEY]
+            if not shard.has_param(op_role_var[0]):
+                block._remove_op(idx, sync=False)
+        block._sync_with_cpp()

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

@@ -0,0 +1,270 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+
+import paddle
+from paddle import static
+
+from .common import (
+    OP_ROLE_KEY,
+    OP_ROLE_VAR_KEY,
+    CollectiveHelper,
+    OpRole,
+    is_backward_op,
+    is_loss_grad_op,
+    is_optimizer_op,
+)
+from .meta_optimizer_base import MetaOptimizerBase
+
+__all__ = []
+
+
+class TensorParallelOptimizer(MetaOptimizerBase):
+    def __init__(self, optimizer):
+        super().__init__(optimizer)
+        self.inner_opt = optimizer
+        self.meta_optimizers_white_list = [
+            "RecomputeOptimizer",
+            "AMPOptimizer",
+            "LarsOptimizer",
+            "LambOptimizer",
+        ]
+        self.meta_optimizers_black_list = []
+        self.mp_ring_id = 0
+        self.global_ring_id = 1
+        self.dp_ring_id = 2
+
+    def _set_basic_info(
+        self, loss, role_maker, user_defined_optimizer, user_defined_strategy
+    ):
+        super()._set_basic_info(
+            loss, role_maker, user_defined_optimizer, user_defined_strategy
+        )
+        self.mp_degree = user_defined_strategy.tensor_parallel_configs[
+            'tensor_parallel_degree'
+        ]
+
+    def _can_apply(self):
+        if not self.role_maker._is_collective:
+            return False
+
+        if self.user_defined_strategy.tensor_parallel:
+            return True
+        return False
+
+    def _disable_strategy(self, dist_strategy):
+        dist_strategy.tensor_parallel = False
+        dist_strategy.tensor_parallel_configs = {}
+
+    def _enable_strategy(self, dist_strategy, context):
+        dist_strategy.tensor_parallel = True
+        dist_strategy.tensor_parallel_configs = {
+            "tensor_parallel_degree": 1,
+        }
+
+    def _broadcast_params(self, ring_id, mp_mode):
+        block = self.startup_program.global_block()
+        param = None
+        for param in block.iter_parameters():
+            if param.is_distributed and mp_mode:
+                continue
+
+            block.append_op(
+                type='broadcast',
+                inputs={'x': param},
+                outputs={'out': param},
+                attrs={
+                    'ring_id': ring_id,
+                    'root': 0,
+                    OP_ROLE_KEY: OpRole.Forward,
+                },
+            )
+
+        if not param:
+            return  # no parameter on this device
+        block.append_op(
+            type='c_sync_comm_stream',
+            inputs={'X': param},
+            outputs={'Out': param},
+            attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Forward},
+        )
+
+    def _get_process_group_info(self):
+        # global ring info
+        self.global_endpoints = self.endpoints
+        self.global_rank = self.rank
+        self.global_nranks = self.nranks
+
+        # model parallel ring info
+        self.mp_rank = self.rank % self.mp_degree
+        self.mp_nranks = self.mp_degree
+        mp_group = self.rank // self.mp_degree
+        self.mp_endpoints = [
+            self.endpoints[i]
+            for i in range(self.global_nranks)
+            if i // self.mp_degree == mp_group
+        ]
+
+        # data parallel ring info
+        if self.nranks > self.mp_degree:
+            self.dp_rank = self.rank // self.mp_degree
+            self.dp_nranks = self.nranks // self.mp_degree
+            start_index = self.rank % self.mp_degree
+            self.dp_endpoints = [
+                self.endpoints[start_index + i * self.mp_degree]
+                for i in range(self.dp_nranks)
+            ]
+
+    def _init_process_group(self):
+        self._get_process_group_info()
+        collective_helper = CollectiveHelper(self.role_maker, wait_port=False)
+
+        # Create global ring for all gpus
+        collective_helper._init_communicator(
+            self.startup_program,
+            self.current_endpoint,
+            self.global_endpoints,
+            self.global_rank,
+            self.global_ring_id,
+            True,
+            self.global_ring_id,
+            True,
+        )
+
+        # Create model parallel ring for all gpus
+        collective_helper._init_communicator(
+            self.startup_program,
+            self.current_endpoint,
+            self.mp_endpoints,
+            self.mp_rank,
+            self.mp_ring_id,
+            True,
+            self.global_ring_id,
+            True,
+        )
+        self._broadcast_params(self.mp_ring_id, mp_mode=True)
+
+        # Create dp rings
+        if self.nranks > self.mp_degree:
+            collective_helper._init_communicator(
+                self.startup_program,
+                self.current_endpoint,
+                self.dp_endpoints,
+                self.dp_rank,
+                self.dp_ring_id,
+                True,
+                self.global_ring_id,
+                True,
+            )
+            self._broadcast_params(self.dp_ring_id, mp_mode=False)
+
+    def minimize_impl(
+        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
+    ):
+        self.endpoints = self.role_maker._get_trainer_endpoints()
+        self.current_endpoint = self.endpoints[self.role_maker._worker_index()]
+        self.startup_program = startup_program
+        if startup_program is None:
+            self.startup_program = static.default_startup_program()
+
+        optimize_ops, params_grads = self.inner_opt.minimize(
+            loss, self.startup_program, parameter_list, no_grad_set
+        )
+
+        self.main_program = loss.block.program
+        self.nranks = len(self.endpoints)
+        self.rank = self.role_maker._worker_index()
+
+        self._init_process_group()
+
+        assert self.nranks % self.mp_degree == 0
+
+        if self.nranks > self.mp_degree:
+            # data parallelism
+            dp_degree = self.nranks // self.mp_degree
+            self._transpile_main_program(loss, dp_degree)
+        return optimize_ops, params_grads
+
+    def _transpile_main_program(self, loss, dp_degree):
+        self._insert_loss_grad_ops(loss, dp_degree)
+        self._insert_allreduce_ops(loss, self.dp_ring_id)
+
+    def _insert_loss_grad_ops(self, loss, dp_degree):
+        """
+        In order to keep the learning rate consistent in different numbers of
+        training workers, we scale the loss grad by the number of workers
+        """
+        block = loss.block
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_loss_grad_op(op):
+                loss_grad_var = block.vars[op.output_arg_names[0]]
+                block._insert_op(
+                    idx + 1,
+                    type='scale',
+                    inputs={'X': loss_grad_var},
+                    outputs={'Out': loss_grad_var},
+                    attrs={
+                        'scale': 1.0 / dp_degree,
+                        OP_ROLE_KEY: OpRole.Backward,
+                    },
+                )
+                break
+
+    def _insert_allreduce_ops(self, loss, ring_id):
+        block = loss.block
+        grad = None
+        for idx, op in reversed(list(enumerate(block.ops))):
+            if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
+                op_role_var = op.attr(OP_ROLE_VAR_KEY)
+                if len(op_role_var) == 0:
+                    continue
+                assert len(op_role_var) % 2 == 0
+                offset = idx
+                for i in range(0, len(op_role_var), 2):
+                    param = block.vars[op_role_var[i]]
+                    grad = block.vars[op_role_var[i + 1]]
+                    if offset == idx:
+                        offset += 1
+                        block._insert_op(
+                            offset,
+                            type='c_sync_calc_stream',
+                            inputs={'X': grad},
+                            outputs={'Out': grad},
+                            attrs={OP_ROLE_KEY: OpRole.Backward},
+                        )
+                        offset += 1
+
+                    block._insert_op(
+                        offset,
+                        type='all_reduce',
+                        inputs={'x': grad},
+                        outputs={'out': grad},
+                        attrs={
+                            'ring_id': ring_id,
+                            'reduce_type': paddle.distributed.ReduceOp.SUM,
+                            OP_ROLE_KEY: OpRole.Backward,
+                        },
+                    )
+
+        if grad is None:
+            return
+
+        for idx, op in list(enumerate(block.ops)):
+            if is_optimizer_op(op):
+                block._insert_op(
+                    idx,
+                    type='c_sync_comm_stream',
+                    inputs={'X': grad},
+                    outputs={'Out': grad},
+                    attrs={'ring_id': ring_id, OP_ROLE_KEY: OpRole.Backward},
+                )
+                break