mindspore 2.7.0__cp311-cp311-win_amd64.whl → 2.7.1__cp311-cp311-win_amd64.whl

mindspore/parallel/distributed/distributed_data_parallel.py

@@ -0,0 +1,393 @@
+# Copyright 2025 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+""" Distributed data parallel wrapper. """
+from __future__ import absolute_import
+
+__all__ = ["DistributedDataParallel"]
+
+import itertools
+from contextlib import contextmanager
+from typing import Optional
+import mindspore.nn as nn
+import mindspore.log as logger
+from mindspore import Tensor, mint
+from mindspore.common import dtype as mstype
+from mindspore.mint.distributed import get_world_size
+from mindspore.communication import GlobalComm
+from mindspore.common.api import _pynative_executor
+from mindspore.mint.distributed import broadcast, get_global_rank
+from mindspore.parallel.distributed.flatten_grad_buffer import FlattenGradBuffer
+from mindspore._c_expression import Reducer, _find_unused_parameters
+
+
+def get_data_parallel_group():
+    """get default global data parallel group"""
+    return GlobalComm.WORLD_COMM_GROUP
+
+
+def get_data_parallel_world_size(group):
+    """get group world size"""
+    return get_world_size(group)
+
+
+def _find_tensors(obj):
+    if isinstance(obj, Tensor):
+        return [obj]
+    if isinstance(obj, (list, tuple)):
+        return itertools.chain.from_iterable(map(_find_tensors, obj))
+    if isinstance(obj, dict):
+        return itertools.chain.from_iterable(map(_find_tensors, obj.values()))
+
+    return []
+
+
+class DistributedDataParallel(nn.Cell):
+    """
+    DistributedDataParallel wrapper. DistributedDataParallel allocates contiguous memory buffer for gradients.
+    Parameters' gradients will be combined into multiple buckets which are the unit to conduct all-reduce
+    communication among data parallel group to overlap communication latency.
+
+    .. warning::
+        - The method is currently only supported in PyNative mode.
+        - This is an experimental interface, may be changed or canceled in the future.
+
+    Args:
+        module (nn.Cell): the module to be wrapped with DDP.
+        init_sync (bool, optional): whether to sync params from rank0 of process_group when init. Default: ``True``.
+        process_group (str, optional): the comm group of data prallel. Default: ``None``.
+        bucket_cap_mb (int, optional): size of bucket in MB, default is 25MB if not set. Default: ``None``.
+        find_unused_parameters (bool, optional): whether to find unused params in the bucket. Default: ``False``.
+        average_in_collective (bool, optional): True means allreduce sum within DP group firstly then scaling with
+            dp size. Otherwise scaling local rank grad first and then allreduce sum. Default: ``False``.
+        static_graph (bool, optional): Indicate whether it is a static network. When it is a static network, the
+            parameter `find_unused_parameters` will be ignored, and unused parameters will be searched for in the
+            first step. Bucket reconstruction will be performed in execution order before the second step to achieve
+            better performance. Default: ``False``.
+        reducer_mode (str, optional): the backend to be used, could be "CppReducer" for cpp backend or "PythonReducer"
+            for Python backend. Default: ``"CppReducer"``.
+
+    Returns:
+        Model wrapped with DistributedDataParallel.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            - When enabling recomputation or gradient freezing, the model should be wrapped by
+              `DistributedDataParallel` at the outermost layer.
+            - Before running the following examples, you need to configure the communication environment variables.
+              For Ascend devices, it is recommended to use the msrun startup method
+              without any third-party or configuration file dependencies. For detailed information, refer to
+              `msrun launch <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_ .
+
+        >>> from mindspore.parallel.distributed import DistributedDataParallel
+        >>> from mindspore.mint.optim import AdamW
+        >>> from mindspore import Parameter, Tensor, ops, nn
+        >>> import mindspore as ms
+        >>> from mindspore.communication import init
+        >>> from mindspore.mint.distributed.distributed import init_process_group
+        >>> ms.set_context(mode=ms.PYNATIVE_MODE)
+        >>> init_process_group()
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = DistributedDataParallel(module=net,
+        ...                              bucket_cap_mb=None,
+        ...                              average_in_collective=True,
+        ...                              static_graph=True)
+        >>> optimizer = AdamW(net.trainable_params(), 1e-4)
+        >>> loss_fn = nn.CrossEntropyLoss()
+        >>>
+        >>> def forward_fn(data, target):
+        ...     logits = net(data)
+        ...     loss = loss_fn(logits, target)
+        ...     return loss, logits
+        >>>
+        >>> grad_fn = ms.value_and_grad(forward_fn, None, net.trainable_params(), has_aux=True)
+        >>>
+        >>> # Create the dataset taking MNIST as an example. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
+        >>> dataset = create_dataset()
+        >>> for epoch in range(1):
+        ...     step = 0
+        ...     for image, label in dataset:
+        ...         (loss_value, _), grads = grad_fn(image, label)
+        ...         optimizer(grads)
+        ...         net.zero_grad()
+        ...         step += 1
+        ...         print("epoch: %s, step: %s, loss is %.15f" % (epoch, step, loss_value))
+    """
+
+    def __init__(self, module, init_sync=True, process_group=None, bucket_cap_mb: Optional[int] = None,
+                 find_unused_parameters=False, average_in_collective: bool = False, static_graph=False,
+                 reducer_mode="CppReducer"):
+        super(DistributedDataParallel, self).__init__(auto_prefix=False)
+        self.init_sync = init_sync
+        self.bucket_cap_mb = bucket_cap_mb
+        self.average_in_collective = average_in_collective
+        self.grad_reduce_in_fp32 = False
+        self.process_group = process_group if process_group else get_data_parallel_group()
+        self.static_graph = static_graph
+        self.find_unused_parameters = find_unused_parameters
+
+        self.module = module
+        self.param_to_buffer = {}
+        self.has_buckets_grad_sync = False
+
+        # default is 25MB for each buck
+        if bucket_cap_mb is None:
+            bucket_cap_mb = 25
+        self.bucket_bytes_cap = int(bucket_cap_mb * 1024 * 1024)
+
+        # grads sync with allreduce comm
+        self.sync_enabled = True
+        self.reducer_mode = reducer_mode # "CppReducer" or "PythonReducer"
+        self.buffers = []
+        self.has_mark_unused_param = False
+
+        bucketed_params = []
+        self.skipped_params = []
+        for _, param in self.module.parameters_and_names():
+            if not param.requires_grad:
+                self.skipped_params.append(param)
+                continue
+            param.grad = None
+            param.main_grad = None
+            bucketed_params.append(param)
+            if self.average_in_collective:
+                # allreduce to add grads, then to scale grads with dp size
+                self.gradient_scaling_factor = 1.0
+            else:
+                # scale grads with dp size locally, then allreduce to add grads
+                data_parallel_world_size = get_data_parallel_world_size(self.process_group)
+                self.gradient_scaling_factor = 1.0 / data_parallel_world_size
+        self.bucketed_params = bucketed_params
+
+        if self.reducer_mode == "CppReducer":
+            self.reducer = Reducer(self.bucketed_params,
+                                   self.process_group,
+                                   bucket_cap_mb,
+                                   self.grad_reduce_in_fp32,
+                                   average_in_collective,
+                                   static_graph,
+                                   find_unused_parameters)
+            if self.init_sync:
+                self.broadcast_coalesced()
+            return
+        # allocate buffer for trained params
+        self.buffers = self.allocate_buffers_for_parameters(
+            self.bucketed_params,
+            group=self.process_group,
+            gradient_scaling_factor=self.gradient_scaling_factor,
+        )
+        if self.init_sync:
+            self.broadcast_coalesced()
+
+        # register hook for bucket grad reduce
+        self._register_hook_for_params()
+
+        # bucket rebuilding
+        self.rebuilt_params_ = []
+        self.buffer_iterations = 0
+        self.has_bucket_rebuilt = False
+        self.buffer_issued = 0
+        self.triggered_once = False
+
+    def _group_params_by_dtype(self, input_params):
+        param_and_grad_dtype_to_params = {}
+        # group all params by parameter's data type and their gradient's data type.
+        for param in input_params:
+            param_dtype = param.dtype
+            grad_dtype = mstype.float32 if self.grad_reduce_in_fp32 else param.dtype
+            if (param_dtype, grad_dtype) not in param_and_grad_dtype_to_params:
+                param_and_grad_dtype_to_params[(param_dtype, grad_dtype)] = []
+            param_and_grad_dtype_to_params[(param_dtype, grad_dtype)].append(param)
+        return param_and_grad_dtype_to_params
+
+    def allocate_buffers_for_parameters(self, input_params, group, gradient_scaling_factor):
+        """allocate buffers for parameters in different dtype group."""
+        param_and_grad_dtype_to_params = self._group_params_by_dtype(input_params)
+
+        buffers = []
+        # allocate buffer for each group separately
+        for (param_dtype, grad_dtype,), params in param_and_grad_dtype_to_params.items():
+            buffers.append(
+                FlattenGradBuffer(
+                    average_in_collective=self.average_in_collective,
+                    param_dtype=param_dtype,
+                    grad_dtype=grad_dtype,
+                    params=params,
+                    data_parallel_group=group,
+                    bucket_size=self.bucket_bytes_cap,
+                    gradient_scaling_factor=gradient_scaling_factor,
+                    ddp_handle=self,
+                )
+            )
+            for param in params:
+                self.param_to_buffer[param] = buffers[-1]
+        logger.debug("allocate buffers for parameters: %s", buffers)
+        return buffers
+
+    def final_grad_reduce(self):
+        """trigger final grad reduction"""
+        logger.debug("trigger ddp final grad reduce, %d, %d", self.static_graph, len(self.unused_param))
+        if self._should_rebuild_buckets():
+            for param in self.unused_param:
+                self.rebuilt_params_.append(param)
+        for buffer in self.buffers:
+            buffer.final_grad_reduce()
+            buffer.issued = 0
+        self.buffer_issued = 0
+
+    def _register_hook_for_params(self):
+        """register backward hook for each params."""
+        for param in self.module.get_parameters():
+            if param.requires_grad:
+                param.register_hook(self._make_param_hook(param))
+
+    def _post_forward(self, output):
+        """prepare for backward (e.g. find unused params) if needed"""
+        if self.reducer_mode == "CppReducer":
+            if _pynative_executor.grad_flag() and self.sync_enabled:
+                self.reducer.prepare_for_backward(list(_find_tensors(output)))
+        else:
+            unused_param_idx = []
+            if self.static_graph and not self.triggered_once:
+                self.triggered_once = True
+                self.find_unused_parameters = False
+                unused_param_idx = _find_unused_parameters(list(_find_tensors(output)), self.bucketed_params)
+            elif self.find_unused_parameters:
+                unused_param_idx = _find_unused_parameters(list(_find_tensors(output)), self.bucketed_params)
+            self.unused_param = [self.bucketed_params[idx] for idx in unused_param_idx]
+            self.unused_param_name = [param.name for param in self.unused_param]
+            self.has_mark_unused_param = False
+
+    def _pre_forward(self):
+        """pre-process of forward pass to allocate buffer for parameters."""
+        if self.reducer_mode == "CppReducer":
+            if _pynative_executor.grad_flag() and self.sync_enabled:
+                self.reducer.prepare_for_forward()
+                self.reducer.rebuild_buckets()
+            return
+        if self.rebuilt_params_ and self._should_rebuild_buckets():
+            for i in self.rebuilt_params_:
+                i.old_grad = i.grad
+
+            self.buffers = self.allocate_buffers_for_parameters(
+                self.rebuilt_params_,
+                group=self.process_group,
+                gradient_scaling_factor=self.gradient_scaling_factor,
+            )
+            for buffer in self.buffers:
+                buffer.sync_enabled = self.sync_enabled
+
+            for i in self.rebuilt_params_:
+                i.grad.copy_(i.old_grad)
+                i.old_grad = None
+
+            logger.debug("register unused param: %s", self.rebuilt_params_)
+            self.has_bucket_rebuilt = True
+            self.rebuilt_params_ = []
+
+    def construct(self, *inputs, **inputs_dict):
+        """construct for DistributedDataParallel."""
+        self._pre_forward()
+        output = self.module(*inputs, **inputs_dict)
+        self._post_forward(output)
+        return output
+
+    def zero_grad(self):
+        """DPP will accumulate grads automatically, it will zero grads when call zero_grad() manually."""
+        if self.reducer_mode == "CppReducer":
+            self.reducer.zero_grad()
+        else:
+            for buffer in self.buffers:
+                buffer.reset()
+
+    def _enable_sync(self, enable):
+        """enable grad buffer sync or not."""
+        for buffer in self.buffers:
+            buffer.sync_enabled = enable
+        self.sync_enabled = enable
+
+    @contextmanager
+    def no_sync(self):
+        """Context manager helper function. When enabled, no grad allreduce synchronization will be executed."""
+        self._enable_sync(False)
+        try:
+            yield
+        finally:
+            self._enable_sync(True)
+
+    def _should_rebuild_buckets(self):
+        if self.static_graph and not self.has_bucket_rebuilt:
+            return True
+        return False
+
+    def _make_param_hook(self, param):
+        """make closure function as the param hook."""
+        def param_hook(grad):
+            if not self.has_mark_unused_param:
+                for cur_param in self.unused_param:
+                    buffer = self.param_to_buffer[cur_param]
+                    logger.debug("register unused param: %s", cur_param)
+                    buffer.register_grad_ready(cur_param)
+                self.has_mark_unused_param = True
+            elif param.name in self.unused_param_name:
+                logger.debug("unused param already registered: %s", param)
+                return param.grad
+
+            logger.debug("register normal param: %s", param)
+            buffer = self.param_to_buffer[param]
+            param.grad.add_(grad)
+            buffer.register_grad_ready(param)
+            if self._should_rebuild_buckets():
+                self.rebuilt_params_.append(param)
+            return param.grad
+
+        return param_hook
+
+    def broadcast_coalesced(self):
+        """broadcast params from rank 0"""
+        if self.reducer_mode == "CppReducer":
+            buckets = [[self.bucketed_params[idx] for idx in bucket] for bucket in self.reducer.bucket_indices]
+        else:
+            buckets = [bucket.params_list for buffer in self.buffers for bucket in buffer.buckets]
+        if self.skipped_params:
+            param_and_grad_dtype_to_params = self._group_params_by_dtype(self.skipped_params)
+            for params_list in param_and_grad_dtype_to_params.values():
+                buckets.append(params_list)
+
+        def finish(rate_limiter):
+            for _ in rate_limiter:
+                handle, coalesced, params = rate_limiter.pop(0)
+                handle.wait()
+                ptr = 0
+                for param in params:
+                    param.view(-1).copy_(coalesced[ptr:ptr + param.numel()])
+                    ptr += param.numel()
+
+        rate_limiter = []
+        for params in buckets:
+            flat_tensors = [t.view(-1) for t in params]
+            coalesced = mint.cat(flat_tensors)
+            global_rank = get_global_rank(self.process_group, 0)
+            handle = broadcast(coalesced, src=global_rank, group=self.process_group, async_op=True)
+            rate_limiter.append((handle, coalesced, params))
+
+            if len(rate_limiter) >= 2:
+                finish(rate_limiter)
+        finish(rate_limiter)

mindspore/parallel/distributed/flatten_grad_buffer.py

@@ -0,0 +1,295 @@
+# Copyright 2025 Huawei Technologies Co., Ltd
+#
+# 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.
+# ============================================================================
+""" Param and grad buffer, bucket implemenatrion. """
+from __future__ import absolute_import
+
+__all__ = ["Bucket", "FlattenGradBuffer"]
+
+from enum import Enum
+import numpy as np
+from mindspore import mint, Tensor
+from mindspore.common.initializer import Zero
+from mindspore.communication.management import get_group_size
+import mindspore.communication.comm_func as comm_func
+
+
+class BufferType(Enum):
+    PARAM = 0
+    GRAD = 1
+
+
+MEM_ALIGN_SIZE = 512
+ALIGN_BYTES = 32
+MIN_BUCKET_SIZE = int(1 * 1024 * 1024)
+DEFAULT_BUCKET_SIZE = int(25 * 1024 * 1024)
+
+
+class Bucket:
+    """
+    Bucket to track a subset of parameters and gradients in the buffer. Bucket records the parameters
+    whose gradient has already been computed. It also provide functionality to synchronize gradients among
+    data parallel group when all parameters' graidents have been computed.
+
+    Args:
+        average_in_collective (bool): Scaling grads before/after AllReduce, True: scaling after AllReduce.
+        params (List(Parameters)): Parameters belongs to this bucket.
+        grad_data (Tensor): A section of buffers' gradient data, coressponding to parameters in this bucket.
+        offset (int): Start index in the buffer.
+        numel_unpadded (int): Number of unpadded elements in bucket.
+        data_parallel_group (str): Data parallel group name.
+        data_parallel_world_size (int): Data parallel group size.
+        gradient_scaling_factor (float): Work with average_in_collective, it is 1.0 when average_in_collective
+        true else 1.0/dp
+    """
+
+    def __init__(self, average_in_collective, params, grad_data, offset, numel_unpadded, data_parallel_group,
+                 data_parallel_world_size, gradient_scaling_factor):
+        self.average_in_collective = average_in_collective
+        self.params_list = params
+        self.params = set(params)
+        self.params_grad_ready = set()
+        self.grad_data = grad_data
+        self.grad_data_numel = self.grad_data.numel()
+        self.offset = offset
+        self.numel_unpadded = numel_unpadded
+        self.data_parallel_group = data_parallel_group
+        self.data_parallel_world_size = data_parallel_world_size
+        self.gradient_scaling_factor = gradient_scaling_factor
+
+        if self.data_parallel_world_size > 1:
+            self.grad_reducer = comm_func.all_reduce
+
+        self.reset()
+
+    def inplace_reduce_dp(self, src):
+        """conduct all-reduce/reduce-scatter on src tensor and inplace update result into target."""
+        self.communication_result, self.communication_handle = self.grad_reducer(
+            src, "sum", self.data_parallel_group, async_op=True
+        )
+
+    def reset(self):
+        """reset bucket for the next iteration."""
+        self.params_grad_ready = set()
+        self.is_reduce_issued = False
+        self.communication_handle = None
+        self.communication_result = None
+
+    def issue_grad_reduce(self):
+        """issue grad reduce for the local grad data view."""
+        if self.is_reduce_issued:
+            raise RuntimeError("The bucket reduce is already issued")
+
+        if self.gradient_scaling_factor != 1.0:
+            self.grad_data.copy_(mint.mul(self.grad_data, self.gradient_scaling_factor))
+
+        if self.data_parallel_world_size > 1:
+            self.inplace_reduce_dp(self.grad_data)
+
+        self.is_reduce_issued = True
+
+    def final_grad_reduce(self):
+        """finalize grad reduce for the local grad data view."""
+        start_idx = 0
+        end_idx = self.grad_data_numel
+        target = self.grad_data[start_idx:end_idx]
+
+        if not self.is_reduce_issued:
+            raise RuntimeError(
+                f"The bucket reduce has not been issued "
+                f"with only {len(self.params_grad_ready)}/{len(self.params)} params ready"
+            )
+
+        if self.data_parallel_world_size > 1:
+            self.communication_handle.wait()
+            target.copy_(self.communication_result)
+            self.communication_result = None
+            if self.average_in_collective:
+                target.copy_(mint.div(target, self.data_parallel_world_size))
+
+    def register_grad_ready(self, param):
+        """register grad ready and issue bucket grad reduce when the bucket is ready."""
+        if param not in self.params:
+            raise ValueError("The param to be registered is not in the bucket")
+
+        if param in self.params_grad_ready:
+            raise ValueError(f"The param {param} is already registered")
+
+        self.params_grad_ready.add(param)
+        if len(self.params_grad_ready) == len(self.params):
+            self.issue_grad_reduce()
+            return True
+
+        return False
+
+    def __repr__(self):
+        return f"Bucket (offset={self.offset}, param_lens={len(self.params)})"
+
+
+class FlattenGradBuffer:
+    """
+    Allocate contiguous memory buffer for given parameters and corresponding gradients. Breaking
+    up parameters and gradients buffer into small buckets, which is the unit for all-reduce/reduce-scatter
+    communication during back-propagation.
+
+    Args:
+        average_in_collective (bool): Scaling grads before/after AllReduce, True: scaling after AllReduce.
+        param_dtype (mindspore.dtype): The parameters' datatype.
+        grad_dtype (mindspore.dtype): The gradients' datatype.
+        params (List(Parameters)): Parameters belongs to this buffer.
+        data_parallel_group (str): Data parallel group name.
+        bucket_size (int): Bucket size threshold used to partition bucekts.
+        gradient_scaling_factor (float):
+    """
+
+    def __init__(self, average_in_collective, param_dtype, grad_dtype, params, data_parallel_group,
+                 bucket_size, gradient_scaling_factor, ddp_handle):
+        super(FlattenGradBuffer, self).__init__()
+        self.param_dtype = param_dtype
+        self.grad_dtype = grad_dtype
+        self.data_parallel_group = data_parallel_group
+        self.data_parallel_world_size = get_group_size(group=self.data_parallel_group)
+        self.gradient_scaling_factor = gradient_scaling_factor
+        self.average_in_collective = average_in_collective
+
+        self.buckets = []
+        self.param_index_map = {}
+        self.param_to_bucket = {}
+        self.sync_enabled = True
+        self.issued = 0
+        self.ddp_handle = ddp_handle
+
+        buckets_metadata = self.calc_partition_metadata(bucket_size, params)
+        self.instantiate_buckets(buckets_metadata, params)
+
+    def calc_partition_metadata(self, bucket_size, params):
+        """calc bucket partition metadata"""
+        # helper func
+        def _need_new_bucket(bucket_numel, bucket_id):
+            target_bucket_size = bucket_size
+            if bucket_id == 0 and bucket_size == DEFAULT_BUCKET_SIZE:
+                target_bucket_size = MIN_BUCKET_SIZE
+            return (
+                bucket_size is not None
+                and bucket_numel != 0
+                and bucket_numel >= target_bucket_size
+            )
+
+        def _build_bucket():
+            nonlocal buckets_metadata, bucket_start_index, bucket_params, bucket_id
+            bucket_end_index = data_start_index
+            buckets_metadata.append(
+                (bucket_start_index, bucket_end_index, bucket_params)
+            )
+            bucket_start_index = bucket_end_index
+            bucket_id = bucket_id + 1
+            bucket_params = []
+
+        param_data_list = []
+        buckets_metadata = []
+        data_start_index = 0
+        data_end_index = 0
+        bucket_id = 0
+        bucket_start_index = 0
+        bucket_params = []
+        for param in params[::]:  # traverse from the beginning
+            last_bucket_numel = data_start_index - bucket_start_index
+            if _need_new_bucket(last_bucket_numel, bucket_id):
+                _build_bucket()
+            data_end_index = data_start_index + param.numel()
+            bucket_params.append(param)
+            param_data_list.append(param)
+            self.param_index_map[param] = (data_start_index, data_end_index, bucket_id)
+            data_start_index = data_end_index
+
+        # add bucket for the last few params which do not reach the bucket_size threshold
+        if data_start_index - bucket_start_index > 0:
+            bucket_end_index = data_start_index
+            buckets_metadata.append(
+                (bucket_start_index, bucket_end_index, bucket_params)
+            )
+            data_start_index = bucket_end_index
+
+        # allocate contiguous memory for parameters and gradients
+        self.numel = data_start_index
+        self.grad_data = Tensor(shape=(self.numel), dtype=self.grad_dtype, init=Zero())
+        self.grad_data.init_data()
+        self.numel_unpadded = 0
+        return buckets_metadata
+
+    def instantiate_buckets(self, buckets_metadata, params):
+        """build bucket instance according to partition metadata"""
+        for bucket_start_index, bucket_end_index, bucket_params in buckets_metadata:
+            local_grad_data = self.grad_data[bucket_start_index:bucket_end_index]
+            self.numel_unpadded += bucket_end_index - bucket_start_index
+            bucket = Bucket(
+                average_in_collective=self.average_in_collective,
+                params=bucket_params,
+                grad_data=local_grad_data,
+                offset=bucket_start_index,
+                numel_unpadded=bucket_end_index - bucket_start_index,
+                data_parallel_group=self.data_parallel_group,
+                data_parallel_world_size=self.data_parallel_world_size,
+                gradient_scaling_factor=self.gradient_scaling_factor,
+            )
+            self.buckets.append(bucket)
+            for param in bucket_params:
+                self.param_to_bucket[param] = bucket
+
+        for param in params:
+            data_start_index, _, _ = self.param_index_map[param]
+            param.grad = self._get_buffer_slice(
+                param.shape, data_start_index, BufferType.GRAD
+            )
+
+    def _get_buffer_slice(self, shape, start_index, buffer_type):
+        """get the buffer view with the same shape"""
+        end_index = start_index + int(np.prod(shape))
+        if start_index < 0 or end_index > self.numel:
+            raise ValueError("index out of range")
+        if buffer_type == BufferType.GRAD:
+            buffer_tensor = self.grad_data[start_index:end_index]
+        else:
+            raise TypeError("Invalid buffer type for _get_buffer_slice.")
+        buffer_tensor = buffer_tensor.view(shape)
+        return buffer_tensor
+
+    def reset(self):
+        """reset buffer for the next iteration."""
+        self.grad_data.zero_()
+        for bucket in self.buckets:
+            bucket.reset()
+        self.sync_enabled = True
+
+    def final_grad_reduce(self):
+        """finalize grad reduce for each bucket"""
+        for bucket in self.buckets:
+            bucket.final_grad_reduce()
+
+    def register_grad_ready(self, param):
+        """register ready grad in its buckets"""
+        if self.sync_enabled:
+            bucket = self.param_to_bucket[param]
+            if bucket.register_grad_ready(param):
+                self.issued += 1
+            if self.issued == len(self.buckets):
+                self.ddp_handle.buffer_issued += 1
+                if self.ddp_handle.buffer_issued == len(self.ddp_handle.buffers):
+                    self.ddp_handle.final_grad_reduce()
+
+    def __repr__(self):
+        param_index_with_name = {
+            param.name: index for (param, index) in self.param_index_map.items()
+        }
+        return f"Buffer has buckets: \n {self.buckets} \n and param_index_map: \n {param_index_with_name}"