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

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}"

mindspore/parallel/function/reshard_func.py

@@ -42,11 +42,12 @@ def reshard(tensor, layout):
     can check :class:`mindspore.parallel.Layout` for reference.
 
     Note:
-        - In the Graph mode, this function can set the sharding propagation strategy of a tensor.
-          For those tensor do not manually be set, their strategies are decided by the sharding
-          strategy propagation algorithm automatically.
-        - In PyNative mode, you can use this method to arrange tensors in a cell (that is, cells
-          that use Cell.shard/F.shard in PyNative mode) that is executed in parallel in graph mode.
+        In the Graph mode, this function can set the sharding propagation strategy of a tensor.
+        For those tensor do not manually be set, their strategies are decided by the sharding
+        strategy propagation algorithm automatically.
+
+    .. warning::
+        The method is currently not supported in PyNative mode.
 
     Args:
         tensor (Tensor): The tensor to be set the sharding strategy.

mindspore/parallel/nn/parallel_cell_wrapper.py

@@ -28,7 +28,8 @@ from mindspore import log as logger
 
 class PipelineCell(Cell):
     """
-    Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training.
+    Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training,
+    and specify the segment info.
 
     Note:
         micro_size must be greater or equal to pipeline stages.
@@ -37,6 +38,8 @@ class PipelineCell(Cell):
         network (Cell): The target network to wrap.
         micro_size (int): MicroBatch size.
         stage_config (dict, optional): The stage configuration for each cell's execution in pipeline parallel.
+        segment_config (dict, optional): The segment configuration for each cell's execution in pipeline parallel.
+            Default ``None``.
 
     Supported Platforms:
         ``Ascend``
@@ -48,7 +51,7 @@ class PipelineCell(Cell):
         >>> net = LeNet5()
         >>> net = nn.PipelineCell(net, 4, stage_config={"cell_name_0": 0, "cell_name_1": 1})
     """
-    def __init__(self, network, micro_size, stage_config=None):
+    def __init__(self, network, micro_size, stage_config=None, segment_config=None):
         super(PipelineCell, self).__init__(auto_prefix=False)
         self.network = network
         self.micro_inputs = nn.CellList()
@@ -104,6 +107,37 @@ class PipelineCell(Cell):
                     logger.warning(cell_name)
                 raise KeyError("For 'PipelineCell', the argument 'stage_config' : {} is not "
                                "found in 'network' : {}".format(config_dict, network))
+        if segment_config is None:
+            return
+        self._config_segment(segment_config)
+
+
+    def _config_segment(self, segment_config):
+        """
+        Config segment num for cell.
+        """
+        config_dict = segment_config.copy()
+
+        for cell_name, cell in self.network.cells_and_names():
+            if cell_name in segment_config:
+                setattr(cell, "pipeline_segment", segment_config[cell_name])
+                del config_dict[cell_name]
+        if str(self.network) in segment_config:
+            setattr(self.network, "pipeline_segment", segment_config[str(self.network)])
+            del config_dict[str(self.network)]
+        # if there are any config elements left, print them
+        if config_dict:
+            for config_cell_name, config_segment_num in config_dict.items():
+                logger.error("pipeline_cell segment_config set pipeline_segment fail!")
+                logger.warning("config cell name:" + str(config_cell_name) +
+                               " config segment num:" + str(config_segment_num))
+            logger.warning("network:" + str(self.network))
+            logger.warning("cell name available:")
+            for cell_name, _ in self.network.cells_and_names():
+                logger.warning(cell_name)
+            raise KeyError("For 'PipelineCell', the argument 'segment_config' : {} is not "
+                           "found in 'network' : {}".format(config_dict, self.network))
+
 
     def construct(self, *args, **kwargs):
         ret = None
@@ -119,7 +153,8 @@ class PipelineCell(Cell):
 
 class Pipeline(PipelineCell):
     """
-    Specify the number of micro_batch for pipeline parallelism and the division rules for stage.
+    Specify the number of micro_batch for pipeline parallelism and the division rules for stage,
+    and specify the segment info.
 
     Note:
         micro_size must be greater or equal to pipeline stages.
@@ -128,6 +163,8 @@ class Pipeline(PipelineCell):
         network (Cell): The target network to wrap.
         micro_size (int): MicroBatch size.
         stage_config (dict, optional): Stage configuration for cell's execution in pipeline parallel. Default ``None``.
+        segment_config (dict, optional): The segment configuration for each cell's execution in pipeline parallel.
+            Default ``None``.
 
     Raises:
         TypeError: The type of `net` is not cell.

mindspore/parallel/nn/parallel_grad_reducer.py

@@ -17,7 +17,6 @@ from __future__ import absolute_import
 
 __all__ = ['PipelineGradReducer']
 
-from mindspore import context
 from mindspore.nn.cell import Cell
 from mindspore.ops import functional as F, composite as C, operations as P
 import mindspore.common.dtype as mstype
@@ -140,7 +139,6 @@ class PipelineGradReducer(Cell):
     """
     def __init__(self, parameters, scale_sense=1.0, opt_shard=None):
         super(PipelineGradReducer, self).__init__(auto_prefix=False)
-        self._check_mode()
         self.accu_grads = parameters.clone(prefix="accu_grads", init="zeros")
         self.grad_reducer = Identity()
         self.degree = Tensor(1, mstype.float32)
@@ -162,9 +160,3 @@ class PipelineGradReducer(Cell):
             accu_grads = self.grad_reducer(self.accu_grads)
             new_grads = self.hyper_map(F.partial(grad_scale, self.scale_sense * self.degree), grads, accu_grads)
         return new_grads
-
-    def _check_mode(self):
-        """check parallel mode"""
-        mode = context.get_context('mode')
-        if mode != context.GRAPH_MODE:
-            raise RuntimeError(f"PipelineGradReducer only support graph mode, but get {mode}")

mindspore/parallel/shard.py

@@ -253,13 +253,6 @@ class Shard(Shard_):
                            "will be overwritten as False.")
             ms.set_algo_parameters(fully_use_devices=False)
 
-        if ms.context.get_auto_parallel_context("full_batch_is_set") is False and \
-            ms.context.get_context("mode") == ms.context.PYNATIVE_MODE:
-            logger.warning("When calling the shard interface, "
-                           "'dataset_strategy' or 'full_batch' is not manually set by the user, "
-                           "and the 'dataset_strategy' will be set to 'full_batch'.")
-            ms.context.set_auto_parallel_context(dataset_strategy="full_batch")
-
         if self._is_attrs_has_been_set(fn, in_strategy, out_strategy, device, level):
             return self.shard_fn
         shard_ = Shard()
@@ -394,11 +387,10 @@ class Shard(Shard_):
                 f"The tuple strategy for each dimension should be tuple(int).")
 
 
-def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascend", level=0):
+def shard(fn, in_strategy, out_strategy=None, parameter_plan=None):
     """
     Specify the input and output slicing strategy for a Cell or function.
-    In PyNative mode, use this method to specify a Cell for distributed
-    execution in graph mode. In Graph mode, use this method to specify distribution strategy for a Cell,
+    In Graph mode, use this method to specify distribution strategy for a Cell,
     strategy for others will be set by sharding propagation.
     in_strategy and out_strategy define the input and output layout respectively.
     in_strategy/out_strategy should be a tuple, each element of which corresponds to the desired layout of
@@ -410,7 +402,9 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
         - It is valid only in semi auto parallel or auto parallel mode.
           In other parallel modes, strategies set here will be ignored.
         - If the input contain Parameter, its strategy should be set in `in_strategy`.
-        - This method currently does not support dynamic shapes.
+
+    .. warning::
+        The method is currently not supported in PyNative mode.
 
     Args:
         fn (Union[Cell, Function]): Function to be executed in parallel.
@@ -432,19 +426,12 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
                                             has been set, the parameter setting will be ignored. Supported
                                             only when `fn` is a Cell with parameters.
                                             Default: ``None`` .
-        device (str, optional): Select a certain `device` target. It is not in use right now.
-                                Support ["CPU", "GPU", "Ascend"]. Default: ``"Ascend"`` .
-        level (int, optional): Option for parallel strategy infer algorithm, namely the object function,
-            maximize computation
-            over communication ratio, maximize speed performance, minimize memory usage etc. It is not in
-            use right now. Support [0, 1, 2]. Default: ``0`` .
 
     Returns:
         Function, return the function that will be executed under auto parallel process.
 
     Raises:
         AssertionError: If parallel mode is not "auto_parallel" nor "semi_auto_parallel".
-        AssertionError: If device_target it not "Ascend" or "GPU".
         TypeError: If `in_strategy` is not a tuple.
         TypeError: If `out_strategy` is not a tuple or None.
         TypeError: If any element in `in_strategy` is not a tuple(int) or tuple(mindspore.parallel.Layout).
@@ -452,8 +439,6 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
         TypeError: If `parameter_plan` is not a dict or None.
         TypeError: If any key in `parameter_plan` is not a str.
         TypeError: If any value in `parameter_plan` is not a tuple(int) or a tuple(mindspore.parallel.Layout).
-        TypeError: If `device` is not a str.
-        TypeError: If `level` is not an integer.
 
     Supported Platforms:
         ``Ascend``
@@ -556,4 +541,5 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
     if not isinstance(fn, (ms.nn.Cell)):
         logger.warning("'fn' is not a mindspore.nn.Cell, and its definition cannot involve Parameter; "
                        "otherwise, the result may be incorrect.")
-    return Shard()(fn, in_strategy, out_strategy, parameter_plan, device, level)
+
+    return Shard()(fn, in_strategy, out_strategy, parameter_plan)