sharding_optimizer.py

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

import os

from paddle.fluid import core
from paddle.fluid.optimizer import PipelineOptimizer
from paddle.static import (
    create_global_var,
    default_startup_program,
    device_guard,
)
from paddle.utils import unique_name

from ..utils.log_util import logger
from .common import (
    OP_ROLE_KEY,
    OP_ROLE_VAR_KEY,
    CollectiveHelper,
    OpRole,
    is_backward_op,
    is_optimizer_op,
    is_update_op,
)
from .meta_optimizer_base import MetaOptimizerBase
from .sharding import utils
from .sharding.fp16_helper import FP16Utils
from .sharding.gradient_clip_helper import GradientClipHelper
from .sharding.offload_helper import OffloadHelper
from .sharding.prune import ProgramDeps
from .sharding.shard import ProgramSegment, Shard
from .sharding.utils import (
    get_first_optimize_op_idx,
    get_grad_device,
    get_var_size,
    insert_allreduce_ops,
    insert_broadcast_ops,
    insert_cast_ops,
    insert_fill_constant_ops,
    insert_reduce_ops,
    insert_scale_loss_grad_ops,
    insert_sync_calc_op,
    insert_sync_comm_ops,
)
from .sharding.weight_decay_helper import WeightDecayHelper

__all__ = []


class ShardingOptimizer(MetaOptimizerBase):
    """Sharding Optimizer."""

    def __init__(self, optimizer):
        super().__init__(optimizer)
        self.inner_opt = optimizer
        self.meta_optimizers_white_list = [
            "RecomputeOptimizer",
            "AMPOptimizer",
            "LarsOptimizer",
            "LambOptimizer",
            "ASPOptimizer",
            # "ModelParallelOptimizer",
            # "PipelineOptimizer",
        ]
        self.meta_optimizers_black_list = []
        self._main_program = None
        self._startup_program = None
        self._segments = []
        # params and fp16 params is for broadcast
        self._params = set()
        self._broadcast_vars = set()
        # reduced grads to param name
        self._reduced_grads_to_param = {}
        self._shard = Shard()
        self._verbose = False

        # use sharding as outer parallelism (e.g. inner:Megatron & outer sharding)
        self.mp_degree = 1

    def _can_apply(self):
        if not self.role_maker._is_collective:
            return False
        if self.role_maker._worker_num() <= 1:
            return False
        return self.user_defined_strategy.sharding

    def _disable_strategy(self, dist_strategy):
        dist_strategy.sharding = False
        dist_strategy.sharding_configs = {}

    def _enable_strategy(self, dist_strategy, context):
        dist_strategy.sharding = True
        dist_strategy.sharding_configs = {"segment_broadcast_MB": 32}

    def _get_sharding_segment_strategy(self):
        """get
        self._sharding_segment_strategy
        1. if by_size:    self._broadcast_MB
        2. if by_anchors: self._sharding_segment_anchors
                          self._backward_remain_anchors
                          self._forward_remain_anchors
        """
        strategy = self.user_defined_strategy
        sharding_configs = strategy.sharding_configs
        segment_strategy = str(sharding_configs["sharding_segment_strategy"])

        if segment_strategy == "segment_broadcast_MB":
            self._broadcast_MB = sharding_configs["segment_broadcast_MB"]
            assert (
                self._broadcast_MB > 0
            ), "segment size should larger than zero !"
        elif segment_strategy == "segment_anchors":
            self._sharding_segment_anchors = sharding_configs["segment_anchors"]
            assert (
                len(self._sharding_segment_anchors) > 0
            ), "you should set the sharding segment anchors !"
            self._backward_remain_anchors = self._sharding_segment_anchors[:]
            self._forward_remain_anchors = []
        else:
            raise NotImplementedError(
                "the sharding segment strategy [{}] is not implemented".format(
                    str(segment_strategy)
                )
            )
        self._sharding_segment_strategy = segment_strategy

    def _get_hybrid_degree(self):
        """get
        self.hybrid_dp
        self.sharding_degree
        self.mp_degree
        self.pp_degree
        self.dp_degree
        """
        strategy = self.user_defined_strategy
        sharding_configs = strategy.sharding_configs

        # parallelism
        sharding_degree = int(sharding_configs["sharding_degree"])
        mp_degree = int(sharding_configs["mp_degree"])
        pp_degree = int(sharding_configs["pp_degree"])
        dp_degree = int(sharding_configs['dp_degree'])
        global_world_size = self.role_maker._worker_num()

        assert sharding_degree > 0, "sharding degree must be larger than zero"
        # pipeline setting
        # TODO (JZ-LIANG) should revise here for support mix parallelism with pipeline
        if pp_degree > 1:
            assert strategy.pipeline is True

        if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None):
            assert pp_degree == 2, (
                "For manually set pipeline, only " "pp_degree = 2 is supported."
            )
            assert (
                global_world_size == mp_degree * sharding_degree * dp_degree
            ), "global work size [{}], mp_degree [{}], sharding_degree [{}], dp_degree [{}].".format(
                global_world_size, mp_degree, sharding_degree, dp_degree
            )
        else:
            assert (
                global_world_size
                == mp_degree * sharding_degree * pp_degree * dp_degree
            ), "global work size [{}], mp_degree [{}], sharding_degree [{}], pp_degree [{}], dp_degree [{}].".format(
                global_world_size,
                mp_degree,
                sharding_degree,
                pp_degree,
                dp_degree,
            )

        # FIXME (JZ-LIANG) deprecated hybrid_dp
        if sharding_configs["hybrid_dp"]:
            logger.warning(
                "[hybrid_dp] API setting is deprecated. Now when "
                "dp_degree >= 2, its will be in hybrid dp mode automatically"
            )
            assert dp_degree >= 1

        self.hybrid_dp = True if dp_degree > 1 else False
        self.sharding_degree = sharding_degree
        self.mp_degree = mp_degree
        self.pp_degree = pp_degree
        self.dp_degree = dp_degree

    def _get_hybrid_dp_mode(self):
        """get
        self.hybrid_dp_mode = 'pp_hybrid_dp' or 'sharding_hybrid_dp'
        self.gradient_merge_mode = 'pp_gm' or 'sharding_gm'
        self._gradient_merge_acc_step
        self.pp_allreduce_in_optimize
        self._optimizer_sharding
        """
        strategy = self.user_defined_strategy
        sharding_configs = strategy.sharding_configs

        # NOTE (JZ-LIANG)
        # There 2 kind of modes for gradient-merge and hybrid-dp in mixed parallelism [sharding] and [pipeline].
        # We distinguish this two modes since the gm/hybrid-dp related allreduce should be insert in different place
        # according different mode to have best performance:
        # sharding: communication within node, and therefore should insert within backward segment
        #           to overlap with bw calc, conduct every micro step.
        # pipeline: communication across nodes, and therefore should insert in update segment,
        #           conduct just once per global step.
        dp_mode = None
        # dp here is the pure dp as the outest parallelism
        if self.hybrid_dp:
            if self.pp_degree > 1:
                dp_mode = "pp_hybrid_dp"
            else:
                assert self.sharding_degree > 1, (
                    "by now we only support five kind of hybrid dp: sharding_hybrid_dp, "
                    "mp_sharding_hybrid_dp, pp_hybrid_dp, mp_sharding_pp_hybrid_dp, sharding_pp_hybrid_dp."
                )
                dp_mode = "sharding_hybrid_dp"

        # gradient merge
        gm_mode = None
        gm_acc_step = int(sharding_configs["gradient_merge_acc_step"])
        if self.pp_degree <= 1:
            gm_mode = "sharding_gm"
            self._grad2merged_grad = {}
        else:
            gm_mode = "pp_gm"
            gm_acc_step = strategy.pipeline_configs['accumulate_steps']
            gradient_scale_configs = strategy.gradient_scale_configs
            assert gradient_scale_configs['scale_strategy'] == 'avg', (
                'For pipeline mode, the '
                'gradient scale mode should '
                'be "avg", but got {}'.format(
                    gradient_scale_configs['scale_strategy']
                )
            )
            # Note (Yuang Liu): this avg_loss flag determines where to do the average op for grad merge.
            # If True, will do sum firstly for gradient merge, then do scale by gm_acc_step.
            # If False, will scale loss by gm_acc_step first, then do sum for gradient merge.
            self.scale_gradient = gradient_scale_configs['scale_gradient']
        if gm_acc_step > 1:
            logger.info(
                "Gradient merge in [{}], acc step = [{}]".format(
                    gm_mode, gm_acc_step
                )
            )

        optimizer_sharding = False
        # TODO(wangxi): need support dp_as_opt_sharding with sharding
        #               need support without pp in future
        if (
            self.sharding_degree == 1
            and self.dp_degree > 1
            and sharding_configs['_dp_as_optimizer_sharding']
            and self.pp_degree > 1
        ):
            optimizer_sharding = True

        self.hybrid_dp_mode = dp_mode
        self.gradient_merge_mode = gm_mode
        self._gradient_merge_acc_step = gm_acc_step
        self._optimizer_sharding = optimizer_sharding

        # this feature is design for ascend, and should NOT be used in GPU training
        self.pp_allreduce_in_optimize = sharding_configs[
            "pp_allreduce_in_optimize"
        ]

    def _inner_opt_minimize(
        self, loss, startup_program, parameter_list, no_grad_set
    ):
        pipeline_configs = self.user_defined_strategy.pipeline_configs

        if self.inner_opt is None:
            raise ValueError(
                "self.inner_opt of ShardingOptimizer should not be None."
            )

        if self.pp_degree > 1:
            pp_optimizer = PipelineOptimizer(
                self.inner_opt, self._gradient_merge_acc_step
            )
            self._pp_optimizer = pp_optimizer

            global_rank = self.role_maker._worker_index()
            schedule_mode = pipeline_configs['schedule_mode']

            pipeline_opt = {
                'schedule_mode': schedule_mode,
                'micro_batch_size': pipeline_configs['micro_batch_size'],
                'local_rank': self.pp_rank,
                'global_rank': global_rank,
                'use_sharding': True,
                # TODO (JZ-LIANG) should revise here for support mix parallelism with pipeline
                'ring_id': 20,
                'global_ring_id': 3,
                'mp_degree': self.mp_degree,
                'mp_rank': global_rank % self.mp_degree,
                'scale_gradient': self.scale_gradient,
            }
            main_program = loss.block.program
            main_program._pipeline_opt = pipeline_opt

            (
                optimize_ops,
                params_grads,
                program_list,
                self.pipeline_pair,
                self.pp_ring_map,
            ) = pp_optimizer.minimize(
                loss, startup_program, parameter_list, no_grad_set
            )
            assert self.pp_degree == len(program_list)
        else:
            optimize_ops, params_grads = self.inner_opt.minimize(
                loss, startup_program, parameter_list, no_grad_set
            )

        if startup_program is None:
            startup_program = default_startup_program()

        if self.pp_degree > 1:
            startup_program = startup_program._pipeline_opt['startup_program']
            print("pp_rank:", self.pp_rank)
            if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None):
                main_program = program_list[
                    int(os.getenv("PADDLE_MANUAL_PIPELINE_STAGE"))
                ]
            else:
                main_program = program_list[self.pp_rank]
            with open("main_%d" % self.role_maker._worker_index(), 'w') as f:
                f.writelines(str(main_program))
            main_block = main_program.global_block()
            new_params_grads = []
            for param, grad in params_grads:
                if main_block.has_var(param.name):
                    new_params_grads.append((param, grad))
            params_grads = new_params_grads
        else:
            main_block = loss.block

        startup_block = startup_program.global_block()
        self._main_program = main_block.program
        self._startup_program = startup_program

        if self.pp_degree > 1:
            pp_optimizer._rename_gradient_var_name(main_block)
            with open("main_%d" % self.role_maker._worker_index(), 'w') as f:
                f.writelines(str(main_program))

        return optimize_ops, params_grads

    def _apply_sharding_pass(self, params_grads):
        if self.sharding_degree == 1:
            return

        main_block = self._main_program.global_block()
        startup_block = self._startup_program.global_block()

        # step1: build shard
        self._build_shard(
            params_grads, self.sharding_rank, self.sharding_degree
        )

        # step2: split_program
        self._split_program(main_block)

        # step3: add broadcast and reduce ops
        self._add_broadcast_allreduce(main_block)
        main_block._sync_with_cpp()
        startup_block._sync_with_cpp()

        # step4: remove unneeded ops and vars from block
        self._prune_main_program(
            main_block,
            self._shard,
            [self.mp_ring_id, self.sharding_ring_id, self.pp_ring_id],
        )
        self._prune_startup_program(startup_block, self._shard)

    def _apply_opt_sharding_pass(self, params_grads):
        """outer dp as optimizer sharding"""
        if self._optimizer_sharding is False:
            return

        main_block = self._main_program.global_block()
        startup_block = self._startup_program.global_block()

        # step1: build shard
        self._build_shard(params_grads, self.dp_rank, self.dp_degree)

        # NOTE(wangxi): prune_main_program will prune cast if not add this
        for param, grad in params_grads:
            self._reduced_grads_to_param[grad.name] = param.name

        # step4: remove unneeded ops and vars from block
        self._prune_main_program(
            main_block,
            self._shard,
            [self.mp_ring_id, self.pp_ring_id, self.dp_ring_id],
        )
        self._prune_startup_program(startup_block, self._shard)

    def _insert_allreduce_for_pp(self, params_grads):
        if self.pp_degree == 1:
            return

        strategy = self.user_defined_strategy
        sharding_configs = strategy.sharding_configs

        main_block = self._main_program.global_block()
        startup_block = self._startup_program.global_block()

        # sharding-pp related logic
        # pp_optimizer._rename_gradient_var_name(main_block)
        # crop ops
        if self.sharding_degree > 1:
            for idx, op in reversed(list(enumerate(main_block.ops))):
                if is_update_op(op):
                    op_role_var = op.attr('op_role_var')
                    param_name = op_role_var[0]
                    if not self._shard.has_param(param_name):
                        main_block._remove_op(idx)

            for idx, op in reversed(list(enumerate(main_block.ops))):
                if op.type != 'cast':
                    continue
                in_name = op.input_arg_names[0]
                if in_name not in self._params:
                    continue
                # if self._shard.has_param(param_name): continue
                if in_name not in main_block.vars:
                    main_block._remove_op(idx)

        if self._optimizer_sharding:
            # TODO(wangxi): support fp16_allreduce with optimizer sharding
            strategy.fp16_allreduce = False

        shard = self._shard if self._optimizer_sharding else None
        accumulated_grad_names = self._pp_optimizer._accumulate_gradients(
            main_block, strategy=strategy, shard=shard
        )

        len_of_ops = len(main_block.ops)
        if self.scale_gradient:
            self._avg_grad_merge_after_sum(main_block, accumulated_grad_names)
        first_optimize_op_index = get_first_optimize_op_idx(main_block)

        if self.pp_allreduce_in_optimize:
            logger.info(
                f"Pipeline Persistable grad is {accumulated_grad_names}"
            )
            # FIXME(wangxi): accumulated_grad get from pipeline is not
            #  include sharding's param@BroadCast grad when
            #  pp_allreduce_in_optimize
            accumulated_grad_names = insert_reduce_ops(
                main_block,
                first_optimize_op_index,
                self.sharding_ring_id,
                accumulated_grad_names,
                self._shard,
                core.op_proto_and_checker_maker.OpRole.Optimize,
                use_calc_stream=True,
                rank=self.sharding_rank,
            )

            logger.info(f"PP-Sharding grad is {accumulated_grad_names}")
            first_optimize_op_index += len(main_block.ops) - len_of_ops
            len_of_ops = len(main_block.ops)

        if self._optimizer_sharding:
            accumulated_grad_names = utils.insert_reduce_ops(
                main_block,
                first_optimize_op_index,
                self.dp_ring_id,
                accumulated_grad_names,
                self._shard,
                OpRole.Optimize,
                use_calc_stream=True,
                rank=self.dp_rank,
                strategy=strategy,
            )
            logger.info(f"Optimizer grad in this rank {accumulated_grad_names}")
            first_optimize_op_index += len(main_block.ops) - len_of_ops
            len_of_ops = len(main_block.ops)

            # NOTE(wangxi): we fused after optimize_cast
            optimize_cast = sharding_configs['optimize_cast']
            optimizer_param = utils.insert_broadcast_param_ops(
                main_block,
                len_of_ops,
                self.dp_ring_id,
                [x[0].name for x in params_grads],
                self._shard,
                OpRole.Optimize,
                use_calc_stream=True,
                rank=self.dp_rank,
                strategy=None if optimize_cast else strategy,
            )
            logger.info(f"Optimizer param in this rank {optimizer_param}")
            if not strategy.fuse_grad_merge and not optimize_cast:
                assert len(accumulated_grad_names) == len(optimizer_param)
        elif self.hybrid_dp and self.hybrid_dp_mode == "pp_hybrid_dp":
            insert_allreduce_ops(
                main_block,
                first_optimize_op_index,
                self.dp_ring_id,
                accumulated_grad_names,
                core.op_proto_and_checker_maker.OpRole.Optimize,
                use_calc_stream=True,
                user_defined_strategy=strategy,
            )
            first_optimize_op_index += len(main_block.ops) - len_of_ops
            len_of_ops = len(main_block.ops)

        # FIXME(wangxi): if fp16_allreduce, put cast fp16->fp32 to there?

    def _avg_grad_merge_after_sum(self, main_block, accumulated_grad_names):
        if (
            self.user_defined_strategy.amp
            and self.user_defined_strategy.amp_configs[
                'use_dynamic_loss_scaling'
            ]
        ):
            # For AMP, if using dynamic loss scaling the avg
            # operation can be simple done by modify the LossScaling op.
            for idx, op in enumerate(main_block.ops):
                if op.type == 'check_finite_and_unscale':
                    loss_scale_name = op.input('Scale')[0]
                    loss_scaling_var = main_block.var(loss_scale_name)
                    loss_scale_tmp_var_name = loss_scale_name + '@TMP'
                    loss_scale_tmp_var = main_block.create_var(
                        name=loss_scale_tmp_var_name,
                        shape=loss_scaling_var.shape,
                        dtype=loss_scaling_var.dtype,
                    )
                    main_block._insert_op_without_sync(
                        idx,
                        type='scale',
                        inputs={'X': loss_scaling_var},
                        outputs={'Out': loss_scale_tmp_var},
                        attrs={
                            'scale': self._gradient_merge_acc_step,
                            'bias': 0.0,
                            'bias_after_scale': False,
                            OP_ROLE_KEY: OpRole.Optimize,
                        },
                    )
                    op._rename_input(loss_scale_name, loss_scale_tmp_var_name)
                    break
        else:
            # For pp, do the avg operation for gradient merge after merging
            # the gradient to meet the logic for gradient merge under pure dp.
            tmp_first_opt_idx = None
            for idx, op in enumerate(main_block.ops):
                if is_optimizer_op(op) and op.type != 'c_sync_comm_stream':
                    tmp_first_opt_idx = idx
                    break
            assert (
                tmp_first_opt_idx is not None
            ), 'Occurs some errors, no optimize ops'
            for grad in accumulated_grad_names:
                main_block._insert_op_without_sync(
                    tmp_first_opt_idx,
                    type='scale',
                    inputs={'X': grad},
                    outputs={'Out': grad},
                    attrs={
                        'scale': 1.0 / self._gradient_merge_acc_step,
                        'bias': 0.0,
                        'bias_after_scale': False,
                        OP_ROLE_KEY: OpRole.Optimize,
                    },
                )

    def _adapt_amp_clip_without_sharding(self):
        # if not use sharding, adapt amp/clip, for remain parallelism.
        # cast --> amp --> clip --> opt
        if self.sharding_degree > 1:
            return
        if self._optimizer_sharding:
            return

        main_block = self._main_program.global_block()
        startup_block = self._startup_program.global_block()

        # amp inf_var & clip global_norm_var

        rings = [self.mp_ring_id, self.pp_ring_id]
        # FIXME(wangxi): some problem with NPU found_finite, need sync with DP
        if core.is_compiled_with_custom_device('npu'):
            rings += [self.dp_ring_id]
        FP16Utils.sync_amp_check_nan_inf(main_block, rings)

        gradientclip_helper = GradientClipHelper(None)
        gradientclip_helper.sync_global_norm(
            main_block, [self.mp_ring_id, self.pp_ring_id], self.mp_rank
        )

    def _insert_loss_grad_scale_op(self):
        main_block = self._main_program.global_block()

        # step6: loss div dp_degree
        global_dp_degree = self.sharding_degree * self.dp_degree
        assert int(global_dp_degree) == global_dp_degree
        if global_dp_degree > 1:
            insert_scale_loss_grad_ops(main_block, scale=global_dp_degree)

        main_block._sync_with_cpp()

    def _apply_optimize_offload_pass(self, params_grads):
        strategy = self.user_defined_strategy
        sharding_configs = strategy.sharding_configs
        main_block = self._main_program.global_block()
        startup_block = self._startup_program.global_block()

        mp_ring_id = self.mp_ring_id if self.mp_degree > 1 else None
        dp_ring_id = self.dp_ring_id if self.dp_degree > 1 else None
        offload_helper = OffloadHelper(
            mp_ring_id=mp_ring_id, dp_ring_id=dp_ring_id
        )

        # optimize offload should be enable while gradient merge is enable and
        # acc_step is quite large (e.g. >> 100). Since its memcpy could not be
        # overlap with calc, otherwise it will slower down training severely.
        if sharding_configs["optimize_offload"]:
            logger.info("Sharding with optimize offload !")
            offload_helper.offload(main_block, startup_block)
            # The optimize_cast is already included in offload_fp32param
            offload_helper.offload_fp32param(main_block, startup_block)
        elif sharding_configs['optimize_cast']:
            logger.info("Sharding with optimize cast !")
            # NOTE(wangxi): optimize_cast will persist fp16 param, it
            # will take more memory, but will be faster. Trade space for time.
            if self._optimizer_sharding:
                offload_helper.opt_sharding_cast_fp32param(
                    main_block, startup_block, [x[0].name for x in params_grads]
                )
                # NOTE(wangxi): fused after optimize_cast
                utils.fuse_opt_broadcast_param_ops(
                    main_block, dp_ring_id, self._shard, strategy=strategy
                )
            else:
                offload_helper.cast_fp32param_in_optimize(
                    main_block, startup_block
                )

    def _dump_program_for_debug(self):
        main_block = self._main_program.global_block()
        startup_block = self._startup_program.global_block()
        with open(
            "start_sharding_%d" % self.role_maker._worker_index(), 'w'
        ) as f:
            f.writelines(str(startup_block.program))
        with open(
            "main_sharding_%d" % self.role_maker._worker_index(), 'w'
        ) as f:
            f.writelines(str(main_block.program))

    def minimize_impl(
        self, loss, startup_program=None, parameter_list=None, no_grad_set=None
    ):
        # TODO: (JZ-LIANG) support multiple comm in future
        # self._nrings = self.user_defined_strategy.nccl_comm_num
        self._nrings_sharding = 1
        self._nrings_dp = 1

        self._get_sharding_segment_strategy()
        self._get_hybrid_degree()
        self._get_hybrid_dp_mode()

        # config sharding & dp groups
        self._build_groups()

        # inner optimize minimize
        optimize_ops, params_grads = self._inner_opt_minimize(
            loss, startup_program, parameter_list, no_grad_set
        )

        self._init_comm()

        self._apply_sharding_pass(params_grads)

        self._apply_opt_sharding_pass(params_grads)

        self._insert_allreduce_for_pp(params_grads)

        self._adapt_amp_clip_without_sharding()

        # loss div dp_degree
        self._insert_loss_grad_scale_op()

        # apply optimize offload or optimize cast
        self._apply_optimize_offload_pass(params_grads)

        # step6: (optional) sharding gradient merge
        self._sharding_gradient_merge()

        # # check op dependecy
        # FIXME (JZ-LIANG) enable checking in future.
        # check_broadcast(main_block)
        # check_allreduce_sum(main_block, self._shard, self.sharding_ring_id,
        #                     self.dp_ring_id)

        # NOTE(JZ-LIANG) ensure in both sharding_hybrid_dp & pp_hybrid_dp
        # init param broadcast should be called after startup pruning
        self._initialization_broadcast()

        # NOTE(wangxi): if param is not persistable, program.clone will
        #  failed, so we remove no persistable param, recreate param as a var
        self._recreate_not_persist_param_as_var()

        self._dump_program_for_debug()

        # GPU need to wait server ready, GPU and NPU is Layered connection
        if not core.is_compiled_with_custom_device('npu'):
            self._wait()
        return optimize_ops, params_grads

    def _init_pair_comm(self, pair, ring_id):
        pp_group_endpoints = [
            self.pp_group_endpoints[pair[0]],
            self.pp_group_endpoints[pair[1]],
        ]
        pp_rank = 0 if self.pp_rank == pair[0] else 1
        if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None) is None:
            self._collective_helper._init_communicator(
                self._startup_program,
                self.current_endpoint,
                pp_group_endpoints,
                pp_rank,
                ring_id,
                False,
                sync=False,
            )

    def _init_npu_pipeline_comm(self, startup_block):
        assert (self.pp_degree % 2) == 0

        max_ring_id = -1
        my_pair = []
        for pair in self.pipeline_pair:
            pair_key = pair[0] * 1000 + pair[1]
            ring_id = self.pp_ring_map[pair_key]
            max_ring_id = max(max_ring_id, ring_id)
            logger.info(f"pp pair:{pair}, ring_id: {ring_id}")

            if self.pp_rank in pair:
                my_pair.append(pair)

        # for example: self.pp_rank=2, self.pp_degree=4
        send_to_next_pair = (
            self.pp_rank,
            (self.pp_rank + 1) % self.pp_degree,
        )  # 2->3
        recv_from_next_pair = (
            (self.pp_rank + 1) % self.pp_degree,
            self.pp_rank,
        )  # 3->2
        recv_from_prev_pair = (
            (self.pp_rank - 1 + self.pp_degree) % self.pp_degree,
            self.pp_rank,
        )  # 1->2
        send_to_prev_pair = (
            self.pp_rank,
            (self.pp_rank - 1 + self.pp_degree) % self.pp_degree,
        )  # 2->1

        even = (self.pp_rank % 2) == 0

        # 1. even send to next, odd recv from prev, 0->1, 2->3
        pair = send_to_next_pair if even else recv_from_prev_pair
        ring_id = self.pp_ring_map[pair[0] * 1000 + pair[1]]
        self._init_pair_comm(pair, ring_id)
        my_pair.remove(pair)
        logger.info(f"pair0(even->odd): pp pair:{pair}, ring_id: {ring_id}")

        # 2. even recv from next, odd send to prev, 1->0, 3->2
        pair = recv_from_next_pair if even else send_to_prev_pair
        ring_id = self.pp_ring_map[pair[0] * 1000 + pair[1]]
        self._init_pair_comm(pair, ring_id)
        my_pair.remove(pair)
        logger.info(f"pair1(even<-odd): pp pair:{pair}, ring_id: {ring_id}")

        # if pp_degree is 2, only need pair(0->1, 1->0)
        if self.pp_degree > 2:
            # 3. odd send to next, even recv from prev, 1->2, 3->0
            pair = send_to_next_pair if not even else recv_from_prev_pair
            ring_id = self.pp_ring_map.get(
                pair[0] * 1000 + pair[1], max_ring_id + 1
            )  # 3->0 not in pp_ring_map
            self._init_pair_comm(pair, ring_id)
            if self.pp_rank != 0 and self.pp_rank != self.pp_degree - 1:
                my_pair.remove(pair)
            logger.info(
                "pair2(odd->even): pp pair:{}, ring_id: {}".format(
                    pair, ring_id
                )
            )

            # 4. odd recv from next, even send to prev, 2->1, 0->3
            pair = recv_from_next_pair if not even else send_to_prev_pair
            ring_id = self.pp_ring_map.get(
                pair[0] * 1000 + pair[1], max_ring_id + 2
            )  # 0->3 not in pp_ring_map
            self._init_pair_comm(pair, ring_id)
            if self.pp_rank != 0 and self.pp_rank != self.pp_degree - 1:
                my_pair.remove(pair)
            logger.info(
                "pair3(odd<-even): pp pair:{}, ring_id: {}".format(
                    pair, ring_id
                )
            )

        assert len(my_pair) == 0, (
            "Current pipeline does not support cross stage communication, "
            "please check unexpected pair {}".format(my_pair)
        )

    def _init_pipeline_comm(self, startup_block):
        # TODO (JZ-LIANG) to unify pp_rank_ and pp_rank
        if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None) is None:
            self._collective_helper._init_communicator(
                self._startup_program,
                self.current_endpoint,
                self.pp_group_endpoints,
                self.pp_rank,
                self.pp_ring_id,
                False,
                sync=False,
            )

        if core.is_compiled_with_custom_device('npu'):
            self._init_npu_pipeline_comm(startup_block)
            return

        # GPU
        for pair in self.pipeline_pair:
            pair_key = pair[0] * 1000 + pair[1]
            ring_id = self.pp_ring_map[pair_key]
            logger.info(f"pp pair:{pair}, ring_id: {ring_id}")
            if self.pp_rank in pair:
                self._init_pair_comm(pair, ring_id)

    def _init_comm(self):
        # sync var
        startup_block = self._startup_program.global_block()

        # mp ring
        if self.mp_degree > 1:
            self._collective_helper._init_communicator(
                self._startup_program,
                self.current_endpoint,
                self.mp_group_endpoints,
                self.mp_rank,
                self.mp_ring_id,
                False,
                sync=False,
            )

        # sharding ring
        if self.sharding_degree > 1:
            self._collective_helper._init_communicator(
                self._startup_program,
                self.current_endpoint,
                self.sharding_group_endpoints,
                self.sharding_rank,
                self.sharding_ring_id,
                False,
                sync=False,
            )

        # pp ring
        if self.pp_degree > 1:
            self._init_pipeline_comm(startup_block)

        # pure dp ring
        if self.dp_degree > 1:
            self._collective_helper._init_communicator(
                self._startup_program,
                self.current_endpoint,
                self.dp_group_endpoints,
                self.dp_rank,
                self.dp_ring_id,
                False,
                sync=False,
            )

        startup_block._sync_with_cpp()

    def _build_shard(self, params_grads, shard_rank, shard_size):
        # step 2: split params
        self._params = {x[0].name for x in params_grads}
        self._shard.setup(params_grads, shard_rank, shard_size)

        # step 3: get broadcast vars
        self._broadcast_vars = self._shard.find_broadcast_params(
            self._main_program.global_block()
        )

    def _wait(
        self,
    ):
        endpoints = self.global_endpoints[:]
        current_endpoint = endpoints[self.global_rank]
        if self.global_rank == 0:
            self._collective_helper._wait(current_endpoint, endpoints)

    def collect_segment(self, segment, op_idx, block):
        segment._start_idx = op_idx + 1
        self._segments.insert(0, segment)
        new_segment = ProgramSegment(block)
        new_segment._end_idx = op_idx + 1

        return new_segment

    def _split_program(self, block):
        for op_idx, op in reversed(list(enumerate(block.ops))):
            if int(op.attr('op_role')) != int(OpRole.Optimize):
                last_backward_op_idx = op_idx + 1
                break

        var2broadcast_time = {}
        segment = ProgramSegment(block)
        segment._end_idx = last_backward_op_idx
        for op_idx in reversed(range(last_backward_op_idx)):
            op = block.ops[op_idx]
            assert int(op.attr('op_role')) != int(OpRole.Optimize)
            if self._sharding_segment_strategy == "segment_broadcast_MB":
                if segment._param_mem >= self._broadcast_MB:
                    segment = self.collect_segment(segment, op_idx, block)

            elif self._sharding_segment_strategy == "segment_anchors":
                if int(op.attr('op_role')) == int(OpRole.Backward):
                    for input_name in op.desc.input_arg_names():

                        # NOTE (JZ-LIANG) naive rule to support amp, if amp change, should modify here accordingly
                        if self.user_defined_strategy.amp:
                            if ".cast_fp16@GRAD" not in input_name:
                                continue
                            else:
                                input_name = input_name[
                                    : input_name.find(".cast_fp16@GRAD")
                                ]

                        if input_name in self._backward_remain_anchors:
                            segment = self.collect_segment(
                                segment, op_idx, block
                            )
                            assert (
                                input_name not in self._forward_remain_anchors
                            ), "segment anchor [{}] met twice !".format(
                                input_name
                            )
                            self._backward_remain_anchors.remove(input_name)
                            self._forward_remain_anchors.append(input_name)
                elif int(op.attr('op_role')) == int(OpRole.Forward):
                    for output_name in op.desc.output_arg_names():
                        if output_name in self._forward_remain_anchors:
                            segment = self.collect_segment(
                                segment, op_idx, block
                            )
                            self._forward_remain_anchors.remove(output_name)

            # find broadcast vars
            for input_name in op.desc.input_arg_names():
                if input_name not in self._broadcast_vars:
                    continue
                if input_name in segment._param2broadcast:
                    # skip broadcast because it reuse the old broadcast var
                    broadcast_name = segment._param2broadcast[input_name]
                    if input_name != broadcast_name:
                        op._rename_input(input_name, broadcast_name)
                    continue
                if self._shard.has_param(input_name):
                    broadcast_var_name = input_name
                else:
                    broadcast_var_name = unique_name.generate(
                        input_name + "@BroadCast"
                    )
                    segment._fill_constant_vars.append(broadcast_var_name)

                # (JZ-LIANG) should use Param base name ?
                broadcast_var_base_name = input_name
                if "subprog" in broadcast_var_base_name:
                    # remove suffix
                    broadcast_var_base_name = broadcast_var_base_name[
                        : broadcast_var_base_name.find(".subprog")
                    ]

                var2broadcast_time[broadcast_var_base_name] = (
                    var2broadcast_time.get(broadcast_var_base_name, 0) + 1
                )

                segment._param2broadcast[input_name] = broadcast_var_name
                segment._broadcast_vars.append(
                    (broadcast_var_name, self._shard.device(input_name))
                )
                segment._param_mem += get_var_size(
                    self._main_program.global_block().var(input_name)
                )

            # find reduce vars
            if self.pp_degree > 1 and self.pp_allreduce_in_optimize:
                # place pipeline gradient allreduce in optimize
                pass
            else:
                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:
                        assert len(op_role_var) % 2 == 0
                        for i in range(0, len(op_role_var), 2):
                            param, reduced_grad = (
                                op_role_var[i],
                                op_role_var[i + 1],
                            )
                            segment._allreduce_vars.append(reduced_grad)
                            assert (
                                reduced_grad not in self._reduced_grads_to_param
                            )
                            self._reduced_grads_to_param[reduced_grad] = param

            # find cast op
            if FP16Utils.is_fp16_cast_op(block, op, self._params):
                fp32_param = op.desc.input_arg_names()[0]
                fp16_param = op.desc.output_arg_names()[0]
                if self._shard.has_param(fp32_param):
                    segment._cast_ops[fp16_param] = fp32_param

        if segment._param_mem > 0:
            segment._start_idx = 0
            self._segments.insert(0, segment)

        if self._sharding_segment_strategy == "segment_anchors":
            assert (
                len(self._forward_remain_anchors) == 0
            ), f"remain anchors {self._forward_remain_anchors}"
            assert (
                len(self._backward_remain_anchors) == 0
            ), f"remain anchors {self._backward_remain_anchors}"

        if self._verbose:
            for varname in sorted(
                var2broadcast_time, key=var2broadcast_time.get, reverse=True
            ):
                logger.info(
                    "Sharding broadcast: [{}] times [{}]".format(
                        var2broadcast_time[varname], varname
                    )
                )
            for idx_ in range(len(self._segments)):
                logger.info(f"segment [{idx_}] :")
                logger.info(
                    "start op: [{}]  [{}]".format(
                        block.ops[self._segments[idx_]._start_idx].desc.type(),
                        block.ops[
                            self._segments[idx_]._start_idx
                        ].desc.input_arg_names(),
                    )
                )
                logger.info(
                    "end   op: [{}]  [{}]".format(
                        block.ops[self._segments[idx_]._end_idx].desc.type(),
                        block.ops[
                            self._segments[idx_]._end_idx
                        ].desc.input_arg_names(),
                    )
                )
        return

    def _prune_main_program(self, block, shard, rings):
        """
        calculate deps from allredce op to optimize op,
        remove ops and vars not needed in this worker

        1. prune regularization (weight decay)
        2. prune cast_fp32_to_fp16; update amp_infine_checking
        3. prune gradient_clip related; update global_norm_sum
        4. prune optimizer op + param + gradient

        """
        weightdecay_helper = WeightDecayHelper()
        weightdecay_helper.prune_weight_decay(block, shard)

        # FIXME(wangxi): mp should prune duplicated param_grads
        # NOTE (JZ-LIANG) the sync of FoundInfinite should among one entire Model Parallelism
        # group. and each Data Parallelism group should have its own sync of FoundInfinite
        # amp could use global group for sync
        FP16Utils.prune_fp16(block, shard, self._reduced_grads_to_param, rings)

        # clipbyglobalnorm should only use the Model paramllelism group (mp-sharding-pp)
        gradientclip_helper = GradientClipHelper(None)
        gradientclip_helper.prune_gradient_clip(block, shard, rings)

        # build prog deps
        reduced_grads = []
        for idx, op in enumerate(block.ops):
            input_names = op.desc.input_arg_names()
            output_names = op.desc.output_arg_names()
            # FIXME(wangxi): need use grads, pipeline grad is @GRAD@MERGE
            if (
                op.type == "c_allreduce_sum"
                and op.attr('use_model_parallel') is False
            ):
                assert len(output_names) == 1
                output_name = output_names[0]
                reduced_grads.append(output_name)

        # prune optimizer state and param
        pruned_opti_vars = []
        for var_name in list(block.vars.keys()):
            if shard.is_opti_var(var_name) and not shard.has_opt_var(var_name):
                pruned_opti_vars.append(var_name)
        program_deps = ProgramDeps(block, reduced_grads, pruned_opti_vars)

        # Init
        for var_name in program_deps._end_vars:
            program_deps._should_removed_var.add(var_name)

        # Prune
        for idx, op in reversed(list(enumerate(block.ops))):
            if op.type in [
                "c_allreduce_sum",
                "c_sync_comm_stream",
                "c_calc_comm_stream",
                "c_gen_nccl_id",
                "c_gen_bkcl_id",
                "c_comm_init",
                'send_v2',
                'recv_v2',
            ]:
                pass
            elif op.type == "conditional_block":
                assert op.desc.has_attr("sub_block")
                subblock_idx = op.desc.attr("sub_block").id
                subblock_deps = program_deps.get_sub_block_deps(subblock_idx)
                # only prune amp subblock
                if subblock_deps is None or not self._is_amp_subblock(op):
                    continue
                # init
                reversed_output_vars = []
                for output_name in op.desc.output("Out"):
                    if output_name in program_deps._should_removed_var:
                        subblock_deps._should_removed_var.add(output_name)
                        program_deps.crop_output_var_from_op(idx, output_name)
                    else:
                        reversed_output_vars.append(output_name)
                # prune
                for sub_op_idx, _ in reversed(
                    list(enumerate(subblock_deps._block.ops))
                ):
                    if subblock_deps.should_remove_op(sub_op_idx):
                        subblock_deps.remove_op(sub_op_idx)
                reversed_input_vars = []
                for input_name in op.desc.input('Input'):
                    if input_name not in subblock_deps._should_removed_var:
                        reversed_input_vars.append(input_name)
                    else:
                        program_deps.crop_input_var_from_op(idx, input_name)
                op.desc.set_input('Input', reversed_input_vars)
                op.desc.set_output('Out', reversed_output_vars)
            else:
                # if all outputs of this op are in _should_removed_var
                # _should_removed_var: opt state not cur shard
                if program_deps.should_remove_op(idx):
                    # NOTE(wangxi): need reserve all param in optimizer_sharding
                    reserved_vars = (
                        self._params if self._optimizer_sharding else None
                    )
                    program_deps.remove_op(idx, reserved_vars)

        # NOTE (JZ-LIANG) revise and unify logic here
        # sharding support fp16_allreduce logic
        block._sync_with_cpp()
        for idx, op in reversed(list(enumerate(block.ops))):
            if op.type == 'concat' and is_optimizer_op(op):
                # remove inputs that not on this card
                reserved_x = []
                for var_name in op.desc.input("X"):
                    if block.has_var(var_name):
                        reserved_x.append(var_name)
                op.desc.set_input('X', reserved_x)
        block._sync_with_cpp()
        return

    def _add_broadcast_allreduce(self, block):
        """
        add broadcast allreduce op
        if enable gradient_merge, insert related ops

        if combined with pipeline(grad accumulate),
        the grad allreduce should be done in optimize role
        """
        if len(self._segments) < 1:
            return
        # sharding
        if self.pp_degree > 1 and self.pp_allreduce_in_optimize:
            for idx in range(len(self._segments)):
                assert len(self._segments[idx]._allreduce_vars) == 0

        # NOTE (JZ-LIANG) revise and unify logic here
        # fix the _end_idx for segments[-1] if pp is used.
        new_end_idx = self._segments[-1]._end_idx
        for idx in range(
            self._segments[-1]._end_idx - 1,
            self._segments[-1]._start_idx - 1,
            -1,
        ):
            op = block.ops[idx]
            if op.type == "fill_constant" or op.type == "sum":
                if "MERGED" in op.output_arg_names[0]:
                    new_end_idx = idx + 1
            elif op.type == "cast":
                if "@TMP" in op.output_arg_names[0]:
                    new_end_idx = idx + 1
        self._segments[-1]._end_idx = new_end_idx

        if self._segments[-1]._allreduce_vars:
            shard_allredue_vars = self._shard.filter_grads(
                self._segments[-1]._allreduce_vars
            )
            if (
                self.gradient_merge_mode != "sharding_gm"
                or self._gradient_merge_acc_step <= 1
            ):
                if (
                    self.hybrid_dp
                    and self.hybrid_dp_mode == "sharding_hybrid_dp"
                    and len(shard_allredue_vars) >= 1
                ):
                    insert_sync_comm_ops(
                        block,
                        self._segments[-1]._end_idx,
                        self.dp_ring_id,
                        shard_allredue_vars,
                    )
                    insert_allreduce_ops(
                        block,
                        self._segments[-1]._end_idx,
                        self.dp_ring_id,
                        shard_allredue_vars,
                        user_defined_strategy=self.user_defined_strategy,
                    )
            # gradient merge
            elif (
                self.gradient_merge_mode == "sharding_gm"
                and self._gradient_merge_acc_step > 1
            ):
                self.create_persistable_gradients_and_insert_merge_ops(
                    block,
                    self._startup_program.global_block(),
                    self._segments[-1]._end_idx,
                    shard_allredue_vars,
                    self._shard,
                )

            insert_sync_comm_ops(
                block,
                self._segments[-1]._end_idx,
                self.sharding_ring_id,
                self._segments[-1]._allreduce_vars,
            )
            # allreduce --> reduce
            insert_reduce_ops(
                block,
                self._segments[-1]._end_idx,
                self.sharding_ring_id,
                self._segments[-1]._allreduce_vars,
                self._shard,
                op_role=OpRole.Backward,
                use_calc_stream=False,
            )

        for idx, segment in reversed(list(enumerate(self._segments))):
            allreduce_vars = (
                self._segments[idx - 1]._allreduce_vars if idx > 0 else []
            )
            broadcast_vars = (
                self._segments[idx + 1]._broadcast_vars
                if idx < len(self._segments) - 1
                else []
            )
            fill_constant_vars = (
                self._segments[idx + 2]._fill_constant_vars
                if idx < len(self._segments) - 2
                else []
            )
            cast_ops = (
                self._segments[idx + 2]._cast_ops
                if idx < len(self._segments) - 2
                else {}
            )

            for op_idx in reversed(range(segment._start_idx, segment._end_idx)):
                op = block.ops[op_idx]
                for input_name in op.desc.input_arg_names():
                    if (
                        input_name in segment._param2broadcast
                        and input_name != segment._param2broadcast[input_name]
                    ):
                        op._rename_input(
                            input_name, segment._param2broadcast[input_name]
                        )

            for param_name, broadcast_name in segment._param2broadcast.items():
                if param_name != broadcast_name:
                    block.create_var(
                        name=broadcast_name,
                        shape=self._main_program.global_block()
                        .var(param_name)
                        .shape,
                        dtype=self._main_program.global_block()
                        .var(param_name)
                        .dtype,
                        persistable=False,
                    )

            # step1: remove cast ops
            block._sync_with_cpp()
            segment._end_idx += FP16Utils.remove_cast_op(
                block, self._params, segment, 0
            )

            # step2: add Sync ops
            shard_allredue_vars = self._shard.filter_grads(allreduce_vars)

            if (
                self.gradient_merge_mode != "sharding_gm"
                or self._gradient_merge_acc_step <= 1
            ):
                if (
                    self.hybrid_dp
                    and self.hybrid_dp_mode == "sharding_hybrid_dp"
                    and len(shard_allredue_vars) >= 1
                ):
                    insert_sync_comm_ops(
                        block,
                        segment._end_idx,
                        self.dp_ring_id,
                        shard_allredue_vars,
                    )

                    broad_cast_vars = [x[0] for x in broadcast_vars]
                    if len(broad_cast_vars) > 0:
                        insert_sync_comm_ops(
                            block,
                            segment._end_idx,
                            self.sharding_ring_id,
                            broad_cast_vars,
                        )
                else:
                    comm_dep_vars = allreduce_vars + [
                        x[0] for x in broadcast_vars
                    ]
                    if len(comm_dep_vars) > 0:
                        insert_sync_comm_ops(
                            block,
                            segment._end_idx,
                            self.sharding_ring_id,
                            comm_dep_vars,
                        )
            # gradient merge
            elif (
                self.gradient_merge_mode == "sharding_gm"
                and self._gradient_merge_acc_step > 1
            ):
                broad_cast_vars = [x[0] for x in broadcast_vars]
                if len(broad_cast_vars) > 0:
                    insert_sync_comm_ops(
                        block,
                        segment._end_idx,
                        self.sharding_ring_id,
                        broad_cast_vars,
                    )

            calc_dep_vars = (
                fill_constant_vars
                + [k for k, v in cast_ops.items()]
                + self._segments[idx]._allreduce_vars
            )

            if len(calc_dep_vars) > 0:
                insert_sync_calc_op(
                    block, segment._end_idx, [calc_dep_vars[-1]]
                )

            # step3: insert `fill_constant` ops
            insert_fill_constant_ops(
                block, segment._end_idx, fill_constant_vars
            )

            # step4: add `cast` ops
            insert_cast_ops(block, segment._end_idx, cast_ops)

            # step5: add broadcast ops
            # gradient merge
            if (
                self.gradient_merge_mode == "sharding_gm"
                and self._gradient_merge_acc_step > 1
            ):
                self.create_persistable_gradients_and_insert_merge_ops(
                    block,
                    self._startup_program.global_block(),
                    segment._start_idx,
                    shard_allredue_vars,
                    self._shard,
                )

            insert_broadcast_ops(
                block, segment._start_idx, self.sharding_ring_id, broadcast_vars
            )

            # step6: add all_reduce ops
            # dp
            if (
                self.gradient_merge_mode != "sharding_gm"
                or self._gradient_merge_acc_step <= 1
            ):
                if (
                    self.hybrid_dp
                    and self.hybrid_dp_mode == "sharding_hybrid_dp"
                    and len(shard_allredue_vars) >= 1
                ):
                    insert_allreduce_ops(
                        block,
                        segment._start_idx,
                        self.dp_ring_id,
                        shard_allredue_vars,
                        user_defined_strategy=self.user_defined_strategy,
                    )
                    insert_sync_comm_ops(
                        block,
                        segment._start_idx,
                        self.sharding_ring_id,
                        allreduce_vars,
                    )
            # gradient merge
            elif (
                self.gradient_merge_mode == "sharding_gm"
                and self._gradient_merge_acc_step > 1
            ):
                insert_sync_comm_ops(
                    block,
                    segment._start_idx,
                    self.sharding_ring_id,
                    allreduce_vars,
                )
            # sharding
            # allreduce --> reduce
            # TODO temp change
            if len(allreduce_vars) > 0:
                insert_reduce_ops(
                    block,
                    segment._start_idx,
                    self.sharding_ring_id,
                    allreduce_vars,
                    self._shard,
                    op_role=OpRole.Backward,
                    use_calc_stream=False,
                )

            block._sync_with_cpp()

        if self._segments[0]._broadcast_vars:
            broadcast_vars = [x[0] for x in self._segments[0]._broadcast_vars]
            insert_sync_comm_ops(
                block,
                self._segments[0]._start_idx,
                self.sharding_ring_id,
                broadcast_vars,
            )
            insert_broadcast_ops(
                block,
                self._segments[0]._start_idx,
                self.sharding_ring_id,
                self._segments[0]._broadcast_vars,
            )

        fill_constant_vars = []
        for x in self._segments[:2]:
            fill_constant_vars += x._fill_constant_vars

        # Join
        cast_ops = {}
        for x in self._segments[:2]:
            for k, v in x._cast_ops.items():
                cast_ops[k] = v

        calc_deps_vars = fill_constant_vars + [k for k, v in cast_ops.items()]
        if fill_constant_vars or cast_ops:
            insert_sync_calc_op(
                block, self._segments[0]._start_idx, [calc_deps_vars[-1]]
            )

        if fill_constant_vars:
            insert_fill_constant_ops(
                block, self._segments[0]._start_idx, fill_constant_vars
            )

        if cast_ops:
            insert_cast_ops(block, self._segments[0]._start_idx, cast_ops)

        return

    def _prune_startup_program(self, block, shard):
        for idx, op in reversed(list(enumerate(block.ops))):
            for output_name in op.desc.output_arg_names():
                if shard.has_var(output_name):
                    continue
                if self._optimizer_sharding and shard.is_param(output_name):
                    continue
                # TODO why do we remove op, when only one var is removed
                block._remove_op(idx, sync=False)
                break

        for var_name in list(block.vars.keys()):
            if shard.has_var(var_name):
                continue
            if self._optimizer_sharding and shard.is_param(var_name):
                continue
            block._remove_var(var_name, sync=False)
        block._sync_with_cpp()

    def _build_groups(self):
        """
        pre-assign ring ids
            mp: 0
            sharding: 1
            pure-dp: 2
            global: 3
            pp: 4
            pp-pair: >= 20
        if one parallelism is not enable: -1
        and only support parallelism hierarchy: mp --> sharding --> pp --> dp
        """
        # step 1: initialize nccl
        self.global_word_size = self.role_maker._worker_num()
        self.global_rank = self.role_maker._worker_index()
        self.global_endpoints = self.role_maker._get_trainer_endpoints()
        self.current_endpoint = self.global_endpoints[self.global_rank]
        self._collective_helper = CollectiveHelper(
            self.role_maker, nrings=self._nrings_sharding
        )
        assert (
            self.global_word_size % self.mp_degree == 0
        ), "global_word_size: {} should be divisible to the mp_degree: {}".format(
            self.global_word_size, self.mp_degree
        )
        assert (
            self.global_word_size % self.sharding_degree == 0
        ), "global_word_size: {} should be divisible to the sharding_degree: {}".format(
            self.global_word_size, self.sharding_degree
        )
        assert (
            self.global_word_size % self.pp_degree == 0
        ), "global_word_size: {} should be divisible to the pp_degree: {}".format(
            self.global_word_size, self.pp_degree
        )
        assert (
            self.global_word_size % self.dp_degree == 0
        ), "global_word_size: {} should be divisible to the dp_degree: {}".format(
            self.global_word_size, self.dp_degree
        )

        # mp group
        if self.mp_degree > 1:
            self.mp_ring_id = 0
            self.mp_rank = self.global_rank % self.mp_degree
            self.mp_group_id = self.global_rank // self.mp_degree
            self.mp_group_endpoints = [
                ep
                for idx, ep in enumerate(self.global_endpoints)
                if idx // self.mp_degree == self.mp_group_id
            ]
            assert self.current_endpoint in self.mp_group_endpoints
            assert (
                len(self.mp_group_endpoints) == self.mp_degree
            ), "num of mp worker in group is [{}], but mp group size is [{}]".format(
                len(self.mp_group_endpoints), self.mp_degree
            )
        else:
            self.mp_degree = 1
            self.mp_ring_id = -1
            self.mp_rank = -1
            self.mp_group_id = -1
            self.mp_group_endpoints = []

        # sharding
        if self.sharding_degree > 1:
            self.sharding_ring_id = 1
            self.sharding_rank = (
                self.global_rank // self.mp_degree
            ) % self.sharding_degree
            self.sharding_group_id = self.global_rank // (
                self.mp_degree * self.sharding_degree
            )
            # mp + sharding + ...
            if self.mp_degree > 1:
                self.sharding_group_endpoints = [
                    ep
                    for idx, ep in enumerate(self.global_endpoints)
                    if (idx // (self.mp_degree * self.sharding_degree))
                    == self.sharding_group_id
                    and idx % self.mp_degree == self.mp_rank
                ]
            # sharding + ...
            else:
                self.sharding_group_endpoints = [
                    ep
                    for idx, ep in enumerate(self.global_endpoints)
                    if (idx // (self.mp_degree * self.sharding_degree))
                    == self.sharding_group_id
                ]
            assert self.current_endpoint in self.sharding_group_endpoints
        else:
            self.sharding_degree = 1
            self.sharding_ring_id = -1
            self.sharding_rank = -1
            self.sharding_group_id = -1
            self.sharding_group_endpoints = []

        # pp
        if self.pp_degree > 1:
            self.pp_pair_ring_id = 20
            # pipeline global ring_id set to 4 for sharding0, mp1, dp2, global3
            self.pp_ring_id = 4
            self.pp_rank = (
                self.global_rank
                // (self.sharding_degree * self.mp_degree)
                % self.pp_degree
            )
            # (NOTE): Already adjust for (outter-pure) dp
            self.pp_group_id = self.global_rank // (
                self.mp_degree * self.sharding_degree * self.pp_degree
            )
            pp_first_stage_idx = self.global_rank % (
                self.sharding_degree * self.mp_degree
            ) + self.pp_group_id * (
                self.mp_degree * self.sharding_degree * self.pp_degree
            )
            pp_stage_offset = self.sharding_degree * self.mp_degree
            self.pp_group_endpoints = []
            for i in range(self.pp_degree):
                self.pp_group_endpoints.append(
                    self.global_endpoints[
                        pp_first_stage_idx + pp_stage_offset * i
                    ]
                )
            assert self.current_endpoint in self.pp_group_endpoints
        else:
            self.pp_ring_id = -1
            self.pp_degree = 1
            self.pp_pair_ring_id = -1
            self.pp_rank = -1
            self.pp_group_id = -1
            self.pp_group_endpoints = []

        # outter-pure-dp group
        # NOTE (JZ-LIANG) support outter-pure-dp to scale the throughput in 3D parallelism
        # e.g. mp-sharding-pp-dp
        # sharding-hybrid-dp as one senario of outter-pure-dp
        local_pp_degree = self.pp_degree
        if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None):
            assert self.pp_degree == 2, (
                "For manually set pipeline, only " "pp_degree = 2 is supported."
            )
            assert (
                self.global_word_size
                == self.mp_degree * self.sharding_degree * self.dp_degree
            ), "global work size [{}], mp_degree [{}], sharding_degree [{}], dp_degree [{}].".format(
                self.global_word_size,
                self.mp_degree,
                self.sharding_degree,
                self.dp_degree,
            )
            local_pp_degree = 1
        else:
            assert (
                self.global_word_size
                == self.mp_degree
                * self.sharding_degree
                * self.pp_degree
                * self.dp_degree
            ), "mp_degree: [{}], sharding_degree: [{}], pp_degree: [{}], dp_degree: [{}]; BUT global nrank: [{}]".format(
                self.mp_degree,
                self.sharding_degree,
                self.pp_degree,
                self.dp_degree,
                self.global_word_size,
            )

        if self.dp_degree > 1:
            self.dp_ring_id = 2
            self.dp_rank = self.global_rank // (
                self.sharding_degree * self.mp_degree * local_pp_degree
            )
            dp_first_rank_idx = self.global_rank % (
                self.sharding_degree * self.mp_degree * local_pp_degree
            )
            dp_offset = self.sharding_degree * self.mp_degree * local_pp_degree
            self.dp_group_endpoints = []
            for i in range(self.dp_degree):
                self.dp_group_endpoints.append(
                    self.global_endpoints[dp_first_rank_idx + dp_offset * i]
                )
            assert self.current_endpoint in self.dp_group_endpoints
            logger.info("Hybrid DP mode turn on !")
        else:
            self.dp_ring_id = -1
            self.dp_rank = -1
            self.dp_group_endpoints = []

        # global group
        # use for gen_nccl_comm_sync, amp check nan inf, clip by global norm
        # NOTE (JZ-LIANG) when use global ring for calc global norm and dp_degree > 1, the allreduce result should be devided by dp_degree
        self.global_ring_id = 3

        logger.info(f"global word size: {self.global_word_size}")
        logger.info(f"global rank: {self.global_rank}")
        logger.info(f"global endpoints: {self.global_endpoints}")
        logger.info(f"global ring id: {self.global_ring_id}")
        logger.info("#####" * 6)

        logger.info(f"mp group size: {self.mp_degree}")
        logger.info(f"mp rank: {self.mp_rank}")
        logger.info(f"mp group id: {self.mp_group_id}")
        logger.info(f"mp group endpoints: {self.mp_group_endpoints}")
        logger.info(f"mp ring id: {self.mp_ring_id}")
        logger.info("#####" * 6)

        logger.info(f"sharding group size: {self.sharding_degree}")
        logger.info(f"sharding rank: {self.sharding_rank}")
        logger.info(f"sharding group id: {self.sharding_group_id}")
        logger.info(
            f"sharding group endpoints: {self.sharding_group_endpoints}"
        )
        logger.info(f"sharding ring id: {self.sharding_ring_id}")
        logger.info("#####" * 6)

        logger.info(f"pp group size: {self.pp_degree}")
        logger.info(f"pp rank: {self.pp_rank}")
        logger.info(f"pp group id: {self.pp_group_id}")
        logger.info(f"pp group endpoints: {self.pp_group_endpoints}")
        logger.info(f"pp ring id: {self.pp_ring_id}")
        logger.info("#####" * 6)

        logger.info(f"pure dp group size: {self.dp_degree}")
        logger.info(f"pure dp rank: {self.dp_rank}")
        logger.info(f"pure dp group endpoints: {self.dp_group_endpoints}")
        logger.info(f"pure dp ring id: {self.dp_ring_id}")
        logger.info("#####" * 6)

        return

    def _recreate_not_persist_param_as_var(self):
        def recreate_not_persist_param_as_var(program):
            block = program.global_block()
            params = block.all_parameters()
            for param in params:
                if param.persistable:
                    continue

                name = param.name
                shape = param.shape
                dtype = param.dtype
                type = param.type
                lod_level = param.lod_level
                stop_gradient = param.stop_gradient
                trainable = param.trainable
                optimize_attr = param.optimize_attr
                regularizer = param.regularizer
                have_dist_attr = False
                is_distributed = False
                if hasattr(param, 'is_distributed'):
                    have_dist_attr = True
                    is_distributed = param.is_distributed

                block._remove_var(name, sync=False)
                var = block.create_var(
                    name=name,
                    shape=shape,
                    dtype=dtype,
                    type=type,
                    lod_level=lod_level,
                    stop_gradient=stop_gradient,
                    trainable=trainable,
                    persistable=False,
                )
                if have_dist_attr:
                    var.is_distributed = is_distributed

            block._sync_with_cpp()

        recreate_not_persist_param_as_var(self._startup_program)
        recreate_not_persist_param_as_var(self._main_program)

    def _initialization_broadcast(self):
        """
        this funtion is to ensure the initialization between dp group to be
        identical when hybrid-dp is used, and the initialization of
        not distributed param between mp group to be identical.
        """
        if self.dp_degree <= 1 and self.mp_degree <= 1:
            return

        startup_block = self._startup_program.global_block()

        params = startup_block.all_parameters()
        params_name = []
        not_dist_param_name = set()

        for param in params:
            params_name.append(param.name)
            if not hasattr(param, 'is_distributed') or not param.is_distributed:
                not_dist_param_name.add(param.name)

        # offload and optimize_cast will insert broadcast op
        broadcast_params = set()
        for op in startup_block.ops:
            if op.type == 'c_broadcast':
                broadcast_params.add(op.desc.output_arg_names()[0])

        for param in params_name:
            if param in broadcast_params:
                continue

            rings = []
            # need sync not distributed param in mp group
            if self.mp_degree > 1 and param in not_dist_param_name:
                rings.append(self.mp_ring_id)
            if self.dp_degree > 1:
                rings.append(self.dp_ring_id)

            for ring in rings:
                startup_block.append_op(
                    type='c_broadcast',
                    inputs={'X': param},
                    outputs={'Out': param},
                    attrs={
                        'ring_id': ring,
                        'root': 0,
                        'use_calc_stream': True,
                        OP_ROLE_KEY: OpRole.Forward,
                    },
                )

        startup_block._sync_with_cpp()

    # sharding gradient merge
    def create_persistable_gradients_and_insert_merge_ops(
        self, main_block, startup_block, insert_idx, grad_names, shard
    ):

        for grad_name in grad_names:
            assert (
                get_grad_device(grad_name, shard) == shard.worker_idx
            ), "try to merge gradient not belong to current shard: [{}]".format(
                grad_name
            )
            persistable_grad_name = grad_name + '@GradiantMerge'
            assert (
                grad_name not in self._grad2merged_grad
            ), "grad [{}] already in grad2merged_grad, maybe you meet sharing weight case !".format(
                grad_name
            )
            self._grad2merged_grad[grad_name] = persistable_grad_name
            grad_var = main_block.var(grad_name)
            # create var
            gradient_merge_var = main_block.create_var(
                name=persistable_grad_name,
                shape=grad_var.shape,
                dtype=grad_var.dtype,
                persistable=True,
            )
            startup_gradient_merge_var = startup_block.create_var(
                name=persistable_grad_name,
                shape=grad_var.shape,
                dtype=grad_var.dtype,
                persistable=True,
            )

            # merge gradient
            main_block._insert_op_without_sync(
                insert_idx,
                type="elementwise_add",
                inputs={'X': grad_name, 'Y': gradient_merge_var},
                outputs={'Out': gradient_merge_var},
                attrs={
                    'axis': -1,
                    'use_mkldnn': False,
                    OP_ROLE_KEY: OpRole.Backward,
                },
            )

            # startup initialization
            startup_block.append_op(
                type="fill_constant",
                outputs={"Out": startup_gradient_merge_var},
                attrs={
                    "shape": grad_var.shape,
                    "dtype": grad_var.dtype,
                    "value": float(0),
                },
            )

        main_block._sync_with_cpp()
        startup_block._sync_with_cpp()

    def _create_gm_cond(self, main_block):
        # Add const var
        acc_step_var = create_global_var(
            name="gradient_merge_acc_step",
            shape=[1],
            value=int(self._gradient_merge_acc_step),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )

        zero_var = create_global_var(
            name="gradient_merge_zero",
            shape=[1],
            value=int(0),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )

        # Add step var & cond var
        current_step_var = create_global_var(
            name="gradient_merge_current_step",
            shape=[1],
            value=int(0),
            dtype='int32',
            persistable=True,
            force_cpu=True,
        )

        cond_var = main_block.create_var(
            name="gradient_merge_cond", shape=[1], dtype='bool'
        )

        with device_guard("cpu"):
            # step_var = (step_var + 1) % k_step
            main_block.append_op(
                type='increment',
                inputs={'X': [current_step_var]},
                outputs={'Out': [current_step_var]},
                attrs={'step': float(1), OP_ROLE_KEY: OpRole.Optimize},
            )

            main_block.append_op(
                type='elementwise_mod',
                inputs={'X': current_step_var, 'Y': acc_step_var},
                outputs={'Out': current_step_var},
                attrs={
                    'axis': -1,
                    OP_ROLE_KEY: OpRole.Optimize,
                    'use_mkldnn': False,
                },
            )

            # cond_var = (step_var == 0)
            main_block.append_op(
                type='equal',
                inputs={'X': current_step_var, 'Y': zero_var},
                outputs={'Out': cond_var},
                attrs={OP_ROLE_KEY: OpRole.Optimize},
            )
        # paddle.static.Print(current_step_var, message="in FWBW last conditional")
        return cond_var

    def _true_apply_gradient(self):
        """
        allreduce grad@gradientmerge in dp group
        grad@gradientmerge / acc_step
        re-create all optimize ops of origin main block and rename them
            cast(backward)
            amp
            clip
            opt
        # fill constant grad@gradientmerge

        """
        # current conditional block
        main_block = self._main_program.global_block()
        cur_block_idx = self._main_program.current_block_idx
        cur_block = self._main_program.current_block()
        self.cond_block = self._main_program.current_block()

        # cur_block's forward_block & backward_block is itself
        cur_block._set_forward_block_idx(cur_block_idx)

        # allreduce grad@gradientmerge
        if self.hybrid_dp:
            assert (
                self.dp_ring_id >= 0
            ), "dp_ring_id should larger than 0 when in sharding&DP mode"
            for grad, merged_grad in self._grad2merged_grad.items():
                merged_grad_var = main_block.var(merged_grad)
                cur_block.append_op(
                    type='c_allreduce_sum',
                    inputs={'X': merged_grad_var},
                    outputs={'Out': merged_grad_var},
                    attrs={
                        'ring_id': self.dp_ring_id,
                        'use_calc_stream': True,
                        OP_ROLE_KEY: OpRole.Optimize,
                    },
                )

        # grad@gradientmerge / acc_step
        for grad, merged_grad in self._grad2merged_grad.items():
            # grad /= k_steps
            merged_grad_var = main_block.var(merged_grad)
            cur_block.append_op(
                type='scale',
                inputs={'X': merged_grad_var},
                outputs={'Out': merged_grad_var},
                attrs={
                    'scale': 1.0 / float(self._gradient_merge_acc_step),
                    'bias': 0.0,
                    'bias_after_scale': False,
                    OP_ROLE_KEY: OpRole.Optimize,
                },
            )

        # re-create optimize ops
        already_moved_var_names = []
        for op_desc in self.original_optimize_ops_desc:
            new_op_desc = cur_block.desc.append_op()
            new_op_desc.copy_from(op_desc)

            for input_name in new_op_desc.input_arg_names():
                if input_name in self._grad2merged_grad:
                    new_op_desc._rename_input(
                        input_name, self._grad2merged_grad[input_name]
                    )

            for output_name in new_op_desc.output_arg_names():
                if output_name in self._grad2merged_grad:
                    new_op_desc._rename_output(
                        output_name, self._grad2merged_grad[output_name]
                    )

                # move non temp optimize vars from block0 to cond block
                if (
                    output_name not in already_moved_var_names
                    and output_name not in self._grad2merged_grad.keys()
                ):
                    var_ = self._main_program.global_block().var(output_name)
                    if not var_.persistable:
                        # move
                        name_ = var_.name
                        shape_ = var_.shape
                        type_ = var_.dtype
                        self._main_program.global_block()._remove_var(
                            var_.name, sync=False
                        )
                        self.cond_block.create_var(
                            name=name_,
                            shape=shape_,
                            dtype=type_,
                            persistable=False,
                        )
                        already_moved_var_names.append(name_)

        self._main_program.global_block()._sync_with_cpp()
        cur_block._sync_with_cpp()

        # fill zero to grad@gradientmerge
        for grad, merged_grad in self._grad2merged_grad.items():
            merged_grad_var = main_block.var(merged_grad)
            cur_block.append_op(
                type='fill_constant',
                outputs={'Out': merged_grad_var},
                attrs={
                    "shape": merged_grad_var.shape,
                    "dtype": merged_grad_var.dtype,
                    "value": float(0),
                    OP_ROLE_KEY: OpRole.Optimize,
                },
            )

        # lr_var = main_block.var("gradient_merge_current_step")
        # paddle.static.Print(lr_var, message="in OPTIMIZE last conditional")

    def _sharding_gradient_merge(self):
        """
        copy all optimize ops in origin main block
        remove all optimize ops in origin main block
        create cond block

        """
        if (
            self.gradient_merge_mode != "sharding_gm"
            or self._gradient_merge_acc_step <= 1
        ):
            return

        main_block = self._main_program.global_block()
        # copy original optimize ops to temp ops desc list
        # remove them from block 0
        tmp_copy_block = self._main_program._create_block()

        self.original_optimize_ops_desc = []
        for op_idx, op in reversed(list(enumerate(main_block.ops))):
            if int(op.attr('op_role')) != int(OpRole.Optimize):
                continue
            else:
                tmp_op_desc = tmp_copy_block.desc.append_op()
                tmp_op_desc.copy_from(op.desc)
                self.original_optimize_ops_desc.append(tmp_op_desc)
                main_block._remove_op(op_idx, sync=False)
        tmp_copy_block._sync_with_cpp()
        self.original_optimize_ops_desc = list(
            reversed(self.original_optimize_ops_desc)
        )

        # back to block 0
        self._main_program._rollback()

        # create cond vars and ops at the end of block 0
        cond = self._create_gm_cond(main_block)

        # create cond block
        cond_block = self._main_program._create_block()
        self._true_apply_gradient()

        # back to block 0
        self._main_program._rollback()

        # cond op
        step_scope = self._main_program.global_block().create_var(
            type=core.VarDesc.VarType.STEP_SCOPES
        )
        conditional_block_op = self._main_program.global_block().append_op(
            type='conditional_block',
            inputs={
                'Cond': cond,
                'Input': [],
            },
            outputs={'Out': [], 'Scope': [step_scope]},
            attrs={
                'sub_block': cond_block,
                'is_scalar_condition': True,
            },
        )