paddlepaddle--paddle/python/paddle/distributed/fleet/utils/tensor_fusion_helper.py

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

import numpy as np

import paddle
from paddle.framework import (
    _current_expected_place_,
    base as imperative_base,
    core,
)
from paddle.framework.recall_error import check_naninf
from paddle.utils import strtobool

from .log_util import logger


def _share_tensor_ipc_meta(tensor):
    if tensor is None:
        return None

    if paddle.is_compiled_with_xpu():
        return tensor.value().get_tensor()._share_xpu()

    if core.is_compiled_with_cuda() and not core.is_compiled_with_rocm():
        return tensor.value().get_tensor()._share_cuda()
    return None


class HOOK_ACTION:
    ALL_REDUCE = 0
    REDUCE = 1
    REDUCE_SCATTER = 2


alignment = {
    "gpu": 256,
    "npu": 256,
    "xpu": 256,
}

align = {
    paddle.float16: 2,
    paddle.bfloat16: 2,
    paddle.float32: 4,
}


__current_device_type__ = None


def get_current_device_type():
    global __current_device_type__
    if __current_device_type__ is None:
        if paddle.is_compiled_with_cuda():
            device_type = "gpu"
        elif paddle.is_compiled_with_xpu():
            device_type = "xpu"
        else:
            current_device = _current_expected_place_()
            try:
                device_type = current_device.get_device_type()
            except:
                device_type = "unknown"
        assert device_type in alignment.keys(), (
            f"tensor fusion helper now only support {alignment.keys()}, but got device {device_type} instead."
        )
        __current_device_type__ = device_type
    return __current_device_type__


def assign_group_by_size(parameters, group_size=128 * 1024 * 1024):
    is_sparse_gradient = [False] * len(parameters)

    group_indices = core.eager_assign_group_by_size(
        parameters, is_sparse_gradient, [group_size, group_size]
    )

    var_groups = OrderedDict()
    group_msg = []
    for group_idx, indices in enumerate(group_indices):
        group_size = 0
        for index in indices:
            var_groups.setdefault(group_idx, []).append(parameters[index])
            group_size += np.prod(parameters[index].shape)
        dtype = parameters[indices[0]].dtype
        bytes = group_size * core.size_of_dtype(dtype)
        msg = f"group_{group_idx}: {bytes / 1024**2:.4f} MB, dtype: {dtype!s}"
        group_msg.append(msg)

    logger.info(f"Tensor Fusion Group Info:\n{group_msg}\n")
    return var_groups


def get_group_size(parameters, group_size=128 * 1024 * 1024):
    is_sparse_gradient = [False] * len(parameters)

    group_indices = core.eager_assign_group_by_size(
        parameters, is_sparse_gradient, [group_size, group_size]
    )

    opt_states_sizes = []
    for group_idx, indices in enumerate(group_indices):
        group_size = 0
        for index in indices:
            group_size += np.prod(parameters[index].shape)
        dtype = parameters[indices[0]].dtype
        bytes = group_size * core.size_of_dtype(dtype)
        param_size_G = bytes / 1024**3
        opt_states_size_G = param_size_G * 12 / core.size_of_dtype(dtype)
        opt_states_sizes.append(opt_states_size_G)

    return opt_states_sizes


def flatten_dense_tensors(
    parameters,
    use_main_grad=False,
    fuse_param=True,
    warp_buffer=False,
    release_grad=False,
):
    from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_storage import (
        GradStorage,
        ParamStorage,
    )

    _buffer_size = 0
    _param2align = {}
    _param2offset = {}
    dtype = parameters[0].dtype

    for param in parameters:
        assert param.trainable, "param must be trainable..."
        size = np.prod(param.shape) * align[dtype]
        remaining = size % alignment[get_current_device_type()]
        ali = (
            0
            if remaining == 0
            else alignment[get_current_device_type()] - remaining
        )
        align_ = ali // align[dtype]
        _param2offset[param.name] = _buffer_size
        _buffer_size += np.prod(param.shape) + align_
        _param2align[param.name] = align_

    if release_grad:
        return None, _buffer_size, _param2offset

    if fuse_param:
        param_storage = ParamStorage(
            size=_buffer_size, dtype=dtype, device=get_current_device_type()
        )
        param_storage.add_rank_params(parameters, _param2align)

    # process gradient
    grad_dtype = paddle.float32 if use_main_grad else dtype
    grad_storage = GradStorage(
        size=_buffer_size,
        dtype=grad_dtype,
        device=get_current_device_type(),
        destination="0",
        param2align=_param2align,
    )

    for param in parameters:
        grad_storage.add_grad(param, _param2align[param.name])

    if warp_buffer:
        if fuse_param:
            param_storage.warp_buffer()
        grad_storage.warp_buffer()

    outputs = (grad_storage,)
    if fuse_param:
        if not use_main_grad:
            # param_storage --> grad_storage
            param_storage.buffer._copy_gradient_from(grad_storage.buffer)
        else:
            param_storage.buffer.main_grad = grad_storage.buffer
        param_storage.buffer.stop_gradient = False
        outputs = (param_storage, *outputs)

    if release_grad:
        outputs = (*outputs, _buffer_size, _param2offset)

    return outputs


def bw_hook_func(buffer, param):
    @paddle.autograd.no_grad()
    def fused_comm(*_):
        buffer.add_grad(param)

    return fused_comm


class ShardingGradView:
    def __init__(
        self,
        param,
        param_buffer,
        grad_buffer,
        index,
        padded_size,
        sharding_degree,
        rank,
        use_main_grad=False,
        release_grad=False,
    ):
        self._param = param
        self._param_buffer = param_buffer
        self._grad_buffer = grad_buffer
        self._index = index
        self._padded_size = padded_size
        self._sharding_degree = sharding_degree
        self._rank = rank
        self._use_main_grad = use_main_grad
        self._release_grad = release_grad
        shard_size = param_buffer._numel() // sharding_degree
        rank_begin = max(rank, 0) * shard_size
        rank_end = rank_begin + shard_size

        param_begin = max(self._index, rank_begin)
        param_end = min(self._index + self._padded_size, rank_end)
        self._param_begin = param_begin
        self._param_end = param_end
        self._rank_begin = rank_begin

        self._slice_grad = None

        if not self._release_grad:
            self._link_grad_to_buffer()

        # share param buffer
        self._share_param_buffer()

    def _get_padding(self):
        if self._param_begin < self._param_end and self._slice_grad is not None:
            padding_start = self._index + self._param._numel()
            padding_end = self._index + self._padded_size
            padding_start = max(self._param_begin, padding_start)
            padding_end = min(self._param_end, padding_end)

            if padding_start >= padding_end:
                return None

            padding = padding_end - padding_start
            grad_numel = self._slice_grad._numel()
            assert grad_numel >= padding, f"{grad_numel} vs {padding}"
            padding_grad = self._slice_grad._slice(
                grad_numel - padding, grad_numel
            )
            return padding_grad
        else:
            return None

    def _slice_grad_from_buffer(self):
        assert self._grad_buffer is not None
        if self._param_begin < self._param_end:
            self._slice_grad = self._grad_buffer._slice(
                self._param_begin, self._param_end
            )
        tmp_grad = self._grad_buffer._slice(
            self._index, self._index + self._param._numel()
        )
        return tmp_grad

    def _link_grad_to_buffer(self):
        tmp_grad = self._slice_grad_from_buffer()
        tmp_grad.get_tensor()._set_dims(self._param.shape)
        if not self._use_main_grad:
            self._param._copy_gradient_from(tmp_grad)
        else:
            self._param.main_grad = tmp_grad

    def _share_param_buffer(self):
        param_shape = self._param.shape
        stop_gradient = self._param.stop_gradient
        self._param.stop_gradient = True
        self._param.flatten_()
        paddle.assign(
            self._param,
            self._param_buffer._slice(
                self._index, self._index + self._param._numel()
            ),
        )
        self._param.get_tensor()._set_dims(param_shape)
        self._param.stop_gradient = stop_gradient
        self._param_buffer._slice(
            self._index, self._index + self._param._numel()
        )._share_buffer_to(self._param)

    def fill_slice_param(self, slice_param):
        slice_begin = self._param_begin
        slice_end = self._param_end
        if slice_param._is_initialized():
            assert self._param_buffer._is_shared_buffer_with(slice_param)
            assert len(slice_param.shape) == 1
            assert slice_param.shape[0] == (slice_end - slice_begin)
        slice_begin = self._param_begin
        slice_end = self._param_end
        slice_buffer = self._param_buffer._slice(slice_begin, slice_end)
        slice_buffer._share_buffer_to(slice_param)
        slice_param.get_tensor()._set_dims([slice_end - slice_begin])

    def assign_slice_grad(self, slice_param):
        assert self._param_buffer._is_shared_buffer_with(self._param)
        slice_grad = self._slice_grad
        if slice_grad is None:
            return
        self.fill_slice_param(slice_param)
        if hasattr(self._param, "main_grad"):
            if not hasattr(slice_param, "main_grad"):
                slice_param.main_grad = slice_grad
            else:
                assert slice_param.main_grad is slice_grad
        elif slice_grad is not None:
            if slice_param.grad is None:
                slice_param._copy_gradient_from(slice_grad)
            else:
                assert slice_param.grad._is_shared_buffer_with(slice_grad)

    def _clear_param_buffer(self):
        self._param._clear_to_zero_allocation()
        self._param_buffer._clear_to_zero_allocation()

    def _reset_param_buffer(self, new_param_storage):
        new_param = paddle.empty_like(self._param)
        new_param._share_buffer_to(self._param)
        new_param_storage._share_buffer_to(self._param_buffer)
        self._share_param_buffer()

    def _clear_grad_buffer(self):
        if self._slice_grad is not None:
            self._slice_grad._clear_dataptr()
            self._slice_grad = None

        if self._grad_buffer is not None:
            self._grad_buffer._clear_dataptr()
            self._grad_buffer = None

    def _reset_grad_buffer(self, slice_grad_buffer):
        self._clear_grad_buffer()
        self._grad_buffer = slice_grad_buffer
        if self._param_begin < self._param_end:
            self._slice_grad = self._grad_buffer._slice(
                self._param_begin - self._rank_begin,
                self._param_end - self._rank_begin,
            )

    @property
    def has_effective_slice_param(self):
        return self._param_begin < self._param_end


def build_reduce_scatter_buffer(
    parameters,
    sharding_degree,
    rank,
    use_main_grad=False,
    release_grad=False,
    init_slice_param=False,
    slice_params={},
):
    total_buffer_size = 0
    param2index = {}
    dtype = parameters[0].dtype

    def get_padded_size(param):
        size = np.prod(param.shape)
        align_size = alignment[get_current_device_type()] // align[dtype]
        align_size = align_size * sharding_degree
        padded_size = ((size + align_size - 1) // align_size) * align_size
        return padded_size

    for param in parameters:
        assert param.trainable, "param must be trainable..."
        param2index[param.name] = total_buffer_size
        total_buffer_size += get_padded_size(param)

    grad_dtype = paddle.float32 if use_main_grad else dtype

    param_buffer = paddle.zeros(shape=[total_buffer_size], dtype=dtype)
    param_buffer_ipc_meta = _share_tensor_ipc_meta(param_buffer)
    grad_buffer = (
        paddle.zeros(shape=[total_buffer_size], dtype=grad_dtype)
        if not release_grad
        else None
    )

    sharding_grad_view = {}
    for param in parameters:
        padded_size = get_padded_size(param)
        grad_view = ShardingGradView(
            param,
            param_buffer,
            grad_buffer,
            param2index[param.name],
            padded_size,
            sharding_degree,
            rank,
            use_main_grad,
            release_grad,
        )
        if init_slice_param and grad_view.has_effective_slice_param:
            assert param.name in slice_params
            grad_view.fill_slice_param(slice_params[param.name])
        # hack main_grad
        sharding_grad_view[param.name] = grad_view
    return (
        sharding_grad_view,
        total_buffer_size,
        param_buffer,
        grad_buffer,
        param_buffer_ipc_meta,
    )


def get_grad_address(param, use_main_grad):
    addr = None
    if use_main_grad:
        if param.main_grad is not None:
            addr = param.main_grad.data_ptr()
    else:
        if (param.grad is not None) and param.grad._is_initialized():
            addr = param.grad.data_ptr()
    return addr


class FusedCommBuffer:
    class Status(enum.Enum):
        """Status of this bucket, Only useful when param allgather overlap is enabled"""

        # Parameters are sharded between processes
        SHARDED = enum.auto()
        # Asynchronous communication is in progress
        SYNCING = enum.auto()
        # Parameters are ready to use
        READY = enum.auto()

    def __init__(
        self,
        id,
        params,
        comm_group,
        acc_steps=1,
        act=None,
        dst=-1,
        use_main_grad=None,
        fuse_param=False,
        scale_after_comm=True,
        release_grads=False,
        use_reduce_avg=False,
        free_grads_in_comm=False,
        init_slice_param=False,
        slice_params={},
    ):
        self._id = id
        self._params = params
        self._acc_steps = acc_steps
        self._comm_group = comm_group
        self._scale_after_comm = scale_after_comm
        self._fuse_param = fuse_param
        self._release_grads = release_grads
        self._use_reduce_avg = use_reduce_avg
        self._free_grads_in_comm = free_grads_in_comm
        self._log_message_printed = False

        self.status = FusedCommBuffer.Status.READY
        self.sync_param_task = None

        if self._free_grads_in_comm:
            assert acc_steps == 1, (
                f"No need to use free_grads_in_comm when acc_steps `{acc_steps}` != 1"
            )
            assert act == HOOK_ACTION.REDUCE_SCATTER, (
                "Currently, only support reduce_scatter"
            )
            assert release_grads, "Currently, only support release_grads"

        assert not (self._fuse_param and self._release_grads), (
            "It's not supported when using fuse_param and release_grad at the same time."
        )

        self.use_main_grad = (
            use_main_grad
            if use_main_grad is not None
            else hasattr(self._params[0], "main_grad")
        )

        self._task = None
        self._dtype = (
            paddle.float32 if self.use_main_grad else self._params[0].dtype
        )
        self._params_step_dict = {}
        self._params_checked_in = 0
        self._grads_to_addr = {}

        self._param_buffer_meta_tensor = None

        self._act = act
        if self._act == HOOK_ACTION.ALL_REDUCE:
            assert dst == -1
        elif self._act == HOOK_ACTION.REDUCE_SCATTER:
            assert dst == -1
        elif self._act == HOOK_ACTION.REDUCE:
            assert dst != -1
        else:
            raise ValueError(
                "The act should be allreduce for dp or reduce for sharding."
            )
        self._dst = dst

        self._init_step_dict()
        if self._act != HOOK_ACTION.REDUCE_SCATTER:
            if self._fuse_param:
                self.param_storage, self.grad_storage = flatten_dense_tensors(
                    self._params,
                    use_main_grad=use_main_grad,
                    fuse_param=True,
                    warp_buffer=True,
                )
                self.param_storage = self.param_storage.buffer
                self.grad_storage = self.grad_storage.buffer
            elif self._release_grads:
                self.param_storage = None
                (
                    grad_storage,
                    self.buffer_size,
                    self.param2offset,
                ) = flatten_dense_tensors(
                    self._params,
                    use_main_grad=self.use_main_grad,
                    fuse_param=False,
                    warp_buffer=False,
                    release_grad=True,
                )
                self.grad_storage = (
                    None if grad_storage is None else grad_storage.buffer
                )
            else:
                self.param_storage = None
                self.grad_storage = flatten_dense_tensors(
                    self._params,
                    use_main_grad=self.use_main_grad,
                    fuse_param=False,
                    warp_buffer=False,
                )[0].buffer
        else:
            assert not self._fuse_param, "not supported"
            (
                self._sharding_param_grad_view,
                self.buffer_size,
                self.param_storage,
                self.grad_storage,
                _,
            ) = build_reduce_scatter_buffer(
                self._params,
                self._comm_group.nranks,
                self._comm_group.rank,
                use_main_grad=self.use_main_grad,
                release_grad=self._release_grads,
                init_slice_param=init_slice_param,
                slice_params=slice_params,
            )
            # hack, for parameter sync in dygraph sharding optimizer after step
            self._params[0].comm_buffer_ref = weakref.ref(self)
            self._param_buffer_meta_tensor = self.param_storage
        if not self._release_grads:
            self._record_addr()

    def _refresh_param_buffer_ipc_meta(self):
        if self._param_buffer_meta_tensor is None:
            return None
        return _share_tensor_ipc_meta(self._param_buffer_meta_tensor)

    @property
    def param_buffer_ipc_meta(self):
        return self._refresh_param_buffer_ipc_meta()

    def _record_addr(self):
        for param in self._params:
            self._grads_to_addr[param.name] = get_grad_address(
                param, self.use_main_grad
            )

    def _clear_param_storage(self):
        self.param_storage._clear_to_zero_allocation()
        for param in self._params:
            self._sharding_param_grad_view[param.name]._clear_param_buffer()

    def _reset_param_storage(self):
        new_param_storage = paddle.empty_like(self.param_storage)
        new_param_storage._share_buffer_to(self.param_storage)
        for param in self._params:
            grad_view = self._sharding_param_grad_view[param.name]
            grad_view._reset_param_buffer(new_param_storage)

    def _clear_grad_storage(self):
        self.grad_storage._clear_dataptr()
        self.grad_storage = None
        if self._act == HOOK_ACTION.REDUCE_SCATTER:
            for param in self._params:
                self._sharding_param_grad_view[param.name]._clear_grad_buffer()

    def _reset_grad_storage(self, slice_grad_buffer):
        self._clear_grad_storage()
        for param in self._params:
            self._sharding_param_grad_view[param.name]._reset_grad_buffer(
                slice_grad_buffer
            )
        self.grad_storage = slice_grad_buffer

    def _init_step_dict(self):
        for p in self._params:
            self._params_step_dict[p.name] = 0

    def _copy_grad_to_buffer(self, param):
        if self._params_step_dict[param.name] > 0:
            return

        if self.grad_storage is None:
            assert self._params_step_dict[param.name] == 0

            self.grad_storage = paddle.zeros(
                [self.buffer_size], dtype=self._dtype
            )

        if self._act == HOOK_ACTION.REDUCE_SCATTER:
            self._sharding_param_grad_view[
                param.name
            ]._grad_buffer = self.grad_storage
            tmp_var = self._sharding_param_grad_view[
                param.name
            ]._slice_grad_from_buffer()
        else:
            grad_end = self.param2offset[param.name] + np.prod(param.shape)
            assert grad_end <= self.buffer_size
            tmp_var = self.grad_storage._slice(
                self.param2offset[param.name], grad_end
            )

        grad_var = param.main_grad if self.use_main_grad else param.grad

        if grad_var is not None:
            grad_var.stop_gradient = True
            grad_var.flatten_()
            tmp_var.add_(grad_var)
            grad_var._clear()

        tmp_var.get_tensor()._set_dims(param.shape)
        if self.use_main_grad:
            if not self._free_grads_in_comm:
                param.main_grad = tmp_var
                param.main_grad.name = "main_grad@" + param.name
        else:
            if not self._free_grads_in_comm:
                param._copy_gradient_from(tmp_var)

        # record address for the following `acc_steps - 1` steps.
        self._grads_to_addr[param.name] = get_grad_address(
            param, self.use_main_grad
        )

    def _reset_params_checked_in(self):
        self._task = None
        self._init_step_dict()
        self._params_checked_in = 0

    @property
    def _all_params_checked_in(self):
        return (
            len(self._params) == self._params_checked_in
            and len(self._params_step_dict) == 0
        )

    def add_grad(self, param, use_comm=True):
        assert param.name in self._params_step_dict

        if not self._release_grads or self._params_step_dict[param.name] > 0:
            current_ptr = get_grad_address(param, self.use_main_grad)
            if self._grads_to_addr[param.name] != current_ptr:
                error_message = f"The address of the grad/main_grad of param {param.name} has been changed during training, which is not allowed for dp/sharding overlap with pp. This may be caused by some non-inplace operations on the grad/main_grad. Here are some examples: 1. The grad/main_grad of the param is changed by other operations, such as: clear_grad; 2. Using non-inplace operations on the grad/main_grad, such as: add, sub, mul, div, etc."
                logger.error(error_message)
                raise ValueError(error_message)
        else:
            # When release_grads is enabled, fusing of gradients only happen
            # in the 0-th gradient accumulation step, and remain unchanged for
            # the following `acc_steps - 1` steps.
            self._copy_grad_to_buffer(param)

        self._params_step_dict[param.name] += 1

        if self._params_step_dict[param.name] == self._acc_steps:
            self._params_checked_in += 1
            self._params_step_dict.pop(param.name)

        if self._all_params_checked_in and use_comm:
            self.comm_grads()

    @imperative_base.no_grad
    def assign_slice_grad(self, param, slice_param):
        assert self._act == HOOK_ACTION.REDUCE_SCATTER
        assert param.name in self._sharding_param_grad_view
        grad_view = self._sharding_param_grad_view[param.name]
        grad_view.assign_slice_grad(slice_param)

    @imperative_base.no_grad
    def sync_params(self, sync=True, param2task={}):
        if not self.need_reduce_scale_sync():
            return
        assert self._act == HOOK_ACTION.REDUCE_SCATTER
        full_buffer = self.param_storage
        group = self._comm_group
        shard_size = full_buffer._numel() // group.nranks

        begin = shard_size * max(group.rank, 0)
        end = begin + shard_size
        slice_buffer = full_buffer._slice(begin, end)
        if group.nranks == 1:
            return
        if sync:
            # default sync_op is False, so we need to wait here.
            # this will call distributed_py.cc in paddle. In distributed_py.cc, there defines two all gather function, their parameters are different.
            group.process_group.all_gather(slice_buffer, full_buffer).wait()
        else:
            # default sync_op is False, so we don't need to to set sync_op = false here.
            task = group.process_group.all_gather(slice_buffer, full_buffer)

            self.sync_param_task = task
            for param in self.params:
                assert param.name not in param2task
                param2task[param.name] = task

    @property
    def params(self):
        return self._params

    @imperative_base.no_grad
    def comm_grads(self):
        assert self._all_params_checked_in, (
            "Not all params checked in."
            f"Parameter number: {len(self._params)}, Check-in number: {self._params_checked_in}"
        )
        self._comm_grads()

    def need_reduce_scale_sync(self):
        stop_gradient_values = [param.stop_gradient for param in self.params]
        if all(stop_gradient_values):
            return False
        else:
            if any(stop_gradient_values) and not self._log_message_printed:
                logger.info(
                    "There is at least one parameter whose stop_gradient attribute is True"
                )
            self._log_message_printed = True
            return True

    @imperative_base.no_grad
    def _comm_grads(self):
        if not self.need_reduce_scale_sync():
            return

        reduce_op = (
            paddle.distributed.ReduceOp.AVG
            if self._use_reduce_avg
            else paddle.distributed.ReduceOp.SUM
        )
        # scale will be skipped when reduce_avg comm operation is enabled.
        if not self._scale_after_comm and not self._use_reduce_avg:
            scale_factor = 1.0 / self._comm_group.nranks
            self.grad_storage.scale_(scale_factor)

        need_check = strtobool(os.getenv('FLAGS_pp_check_naninf', '0'))
        if need_check:
            err_msg = check_naninf(self.grad_storage)
            if err_msg is not None:
                raise ValueError(
                    f"{err_msg}. Tensor contains inf or nan values at rank {paddle.distributed.get_rank()} before gradient communication"
                )

        if self._act == HOOK_ACTION.ALL_REDUCE:
            task = paddle.distributed.all_reduce(
                self.grad_storage,
                op=reduce_op,
                group=self._comm_group,
                sync_op=False,
            )

        elif self._act == HOOK_ACTION.REDUCE:
            task = paddle.distributed.reduce(
                self.grad_storage,
                dst=self._dst,
                op=reduce_op,
                group=self._comm_group,
                sync_op=False,
            )

        elif self._act == HOOK_ACTION.REDUCE_SCATTER:
            # In align mode, we scale the grad in advance, so we need a SUM head
            if paddle.distributed.in_auto_parallel_align_mode():
                reduce_op = paddle.distributed.ReduceOp.SUM
            shard_size = self.grad_storage._numel() // self._comm_group.nranks
            begin = shard_size * max(self._comm_group.rank, 0)
            end = begin + shard_size
            reduce_scattered = (
                paddle.empty_like(self.grad_storage._slice(begin, end))
                if self._free_grads_in_comm
                else self.grad_storage._slice(begin, end)
            )
            task = paddle.distributed.reduce_scatter(
                reduce_scattered,
                self.grad_storage,
                op=reduce_op,
                group=self._comm_group,
                sync_op=False,
            )

            if self._free_grads_in_comm:
                self._reset_grad_storage(reduce_scattered)

        self._task = task

    @imperative_base.no_grad
    def scale_grads(self):
        if self.need_reduce_scale_sync():
            if self._comm_group.nranks == 1 and self._task is None:
                self._reset_params_checked_in()
                return
            assert self._task is not None, "Task is not initialized."
            self._task.wait()

            # scale will be skipped when use reduce_avg comm operation
            if self._scale_after_comm and not self._use_reduce_avg:
                scale_factor = 1.0 / self._comm_group.nranks
                self.grad_storage.scale_(scale_factor)

        self._reset_params_checked_in()


def obtain_storage(
    parameters,
    use_main_grad=False,
    clip=True,
    dist=False,
    fuse_param=True,
    comm_overlap=False,
    act=None,
    comm_group=None,
    dst=-1,
    acc_steps=1,
    scale_after_comm=False,
    use_reduce_avg=False,
    group_size=256 * 1024 * 1024,
):
    if len(parameters) < 1:
        return [], []

    var_groups = assign_group_by_size(parameters, group_size=group_size)
    storage = []
    buffers = []
    for group_idx, parameters in var_groups.items():
        comm_buffer = FusedCommBuffer(
            group_idx,
            parameters,
            comm_group=comm_group,
            acc_steps=acc_steps,
            act=act,
            dst=dst,
            use_main_grad=use_main_grad,
            fuse_param=fuse_param,
            scale_after_comm=scale_after_comm,
            use_reduce_avg=use_reduce_avg,
        )
        if fuse_param:
            param_buffer = comm_buffer.param_storage
            param_buffer.need_clip = clip
            param_buffer.is_distributed = dist
            storage.append(param_buffer)
        if comm_overlap:
            for param in parameters:
                param._register_backward_hook(bw_hook_func(comm_buffer, param))
            buffers.append(comm_buffer)

    return storage, buffers


def filter_params(params, is_fp32, is_distributed, need_clip):
    params = list(
        filter(
            lambda x: (
                x.is_distributed if is_distributed else (not x.is_distributed)
            ),
            params,
        )
    )
    params = list(
        filter(
            lambda x: (
                getattr(x, 'need_clip', True)
                if need_clip
                else (not getattr(x, 'need_clip', True))
            ),
            params,
        )
    )
    params = list(
        filter(
            lambda x: (
                x.dtype == paddle.float32
                if is_fp32
                else x.dtype != paddle.float32
            ),
            params,
        )
    )
    dtype = None
    for p in params:
        if dtype is None:
            dtype = p.dtype
        else:
            assert dtype == p.dtype

    return params, dtype


def _fused_parameters_impl(
    parameters,
    use_main_grad=False,
    fuse_param=True,
    comm_overlap=False,
    comm_group=None,
    act=None,
    dst=-1,
    acc_step=1,
    scale_after_comm=False,
    apply_decay_param_fun=None,
    use_reduce_avg=False,
    group_size=256 * 1024 * 1024,
):
    param_groups = []
    attrs = []

    is_fp32 = [True, False]
    is_distributed = [True, False]
    need_clip = [True, False]

    no_fp32_dtype = None
    for fp32, dist, clip in itertools.product(
        is_fp32, is_distributed, need_clip
    ):
        params, dtype = filter_params(parameters, fp32, dist, clip)
        if not fp32:
            if no_fp32_dtype is None:
                no_fp32_dtype = dtype
            elif dtype is not None:
                assert no_fp32_dtype == dtype
        attrs.append([dtype, dist, clip])
        param_groups.append(params)

    decay_fused = []
    all_fused = []
    all_buffers = []
    for params, attr in zip(param_groups, attrs):
        decay_params = []
        other_params = []

        for param in params:
            if apply_decay_param_fun is not None and apply_decay_param_fun(
                param.name
            ):
                decay_params.append(param)
            else:
                other_params.append(param)

        is_distributed = attr[1]
        need_clip = attr[2]
        decay, decay_buffers = obtain_storage(
            decay_params,
            use_main_grad=use_main_grad,
            clip=need_clip,
            dist=is_distributed,
            fuse_param=fuse_param,
            comm_overlap=comm_overlap,
            act=act,
            comm_group=comm_group,
            dst=dst,
            acc_steps=acc_step,
            scale_after_comm=scale_after_comm,
            use_reduce_avg=use_reduce_avg,
            group_size=group_size,
        )
        other, other_buffers = obtain_storage(
            other_params,
            fuse_param=fuse_param,
            comm_overlap=comm_overlap,
            use_main_grad=use_main_grad,
            clip=need_clip,
            dist=is_distributed,
            act=act,
            comm_group=comm_group,
            dst=dst,
            acc_steps=acc_step,
            scale_after_comm=scale_after_comm,
            use_reduce_avg=use_reduce_avg,
            group_size=group_size,
        )
        decay_fused += decay
        all_fused += decay
        all_fused += other
        all_buffers += decay_buffers
        all_buffers += other_buffers

    return decay_fused, all_fused, all_buffers


def fused_parameters(
    parameters,
    use_main_grad=False,
    fuse_param=True,
    comm_overlap=False,
    comm_group=None,
    act=None,
    dst=-1,
    acc_step=1,
    scale_after_comm=False,
    group_params=False,
    apply_decay_param_fun=None,
    use_reduce_avg=False,
    group_size=256 * 1024 * 1024,
):
    """
    Fuse gradients. Fuse parameters if be enabled. Prepare for comm overlap if be enabled.
    :param parameters: all parameters to be fused.
    :param use_main_grad: does the gradient use main grad or not
    :param comm_overlap: enable comm overlap or not
    :param comm_group: the comm group for comm overlap
    :param act: the comm operation, could be chosen from reduce and allreduce
    :param dst: the dst for comm overlap
    :param acc_step: acc steps, using for comm overlap
    :param fuse_param: fuse param or not
    :param scale_after_comm: if enable comm overlap, specify the location of grad scale
    :param group_params: the format of the input parameters is param group
    :param apply_decay_param_fun: the function to filter decay param
    :param use_reduce_avg: use reduce_avg comm operation instead of scale and reduce_sum
    :param group_size: the size of each group, default is 256MB
    :return: param storage if fused, comm buffers if comm overlap, param groups if use group params
    """
    if act is None:
        act = HOOK_ACTION.REDUCE

    if comm_overlap:
        if comm_group is None:
            assert act == HOOK_ACTION.ALL_REDUCE, (
                "Only allreduce action can use default comm group"
            )
            comm_group = paddle.distributed.collective._get_default_group()
    if act == HOOK_ACTION.REDUCE:
        assert dst != -1
    elif act == HOOK_ACTION.ALL_REDUCE:
        dst = -1

    if group_params:
        updated_parameters = []
        comm_buffers = []
        for idx, group_param in enumerate(parameters):
            assert isinstance(group_param, dict), (
                "For group params, each group should be a dictionary."
            )
            assert 'params' in group_param.keys(), (
                "For group params, each group should have parameters."
            )
            real_param = group_param['params']
            (
                group_decay_fused,
                group_all_fused,
                group_all_buffers,
            ) = _fused_parameters_impl(
                real_param,
                use_main_grad=use_main_grad,
                fuse_param=fuse_param,
                comm_overlap=comm_overlap,
                comm_group=comm_group,
                act=act,
                dst=dst,
                acc_step=acc_step,
                scale_after_comm=scale_after_comm,
                apply_decay_param_fun=apply_decay_param_fun,
                use_reduce_avg=use_reduce_avg,
                group_size=group_size,
            )
            if comm_overlap:
                comm_buffers.extend(group_all_buffers)
            for fused_tensor in group_all_fused:
                fused_tensor.optimize_attr = real_param[0].optimize_attr
            group_param['params'] = group_all_fused
            updated_parameters.append(group_param)
        return updated_parameters, comm_buffers
    else:
        decay_fused, all_fused, all_buffers = _fused_parameters_impl(
            parameters,
            use_main_grad=use_main_grad,
            fuse_param=fuse_param,
            comm_overlap=comm_overlap,
            comm_group=comm_group,
            act=act,
            dst=dst,
            acc_step=acc_step,
            scale_after_comm=scale_after_comm,
            apply_decay_param_fun=apply_decay_param_fun,
            use_reduce_avg=use_reduce_avg,
            group_size=group_size,
        )

        return decay_fused, all_fused, all_buffers