mindspore 2.4.10__cp311-cp311-win_amd64.whl → 2.6.0rc1__cp311-cp311-win_amd64.whl

mindspore/parallel/function/reshard_func.py

@@ -0,0 +1,259 @@
+# Copyright 2023 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.
+# ============================================================================
+"""Defines parameter operators with functional form."""
+import mindspore as ms
+from mindspore import context, ops
+from mindspore import log as logger
+from mindspore.ops import operations as P
+from mindspore.ops._primitive_cache import _get_cache_prim
+from mindspore.common.tensor import Tensor
+from mindspore.communication.management import get_group_size, get_rank
+from mindspore.parallel.shard import Layout, _DistributedTensorInfo
+from mindspore.parallel._auto_parallel_context import _get_all_auto_parallel_context, _recover_auto_parallel_context
+from mindspore.ops.primitive import constexpr
+
+
+REDIST_CELL_CACHE = {}
+COMM_TENSOR_CELL_CACHE = {}
+
+
+@constexpr
+def group_size():
+    """ Return the device number in the Cell's construct method. """
+    return get_group_size()
+
+
+# pylint: disable=W0212
+def reshard(tensor, layout):
+    r"""
+    Converting a tensor from one distributed arrangement to another distributed arrangement.
+    The given layout must be type mindspore.parallel.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.
+
+    Args:
+        tensor (Tensor): The tensor to be set the sharding strategy.
+        layout (Layout): The layout to shard the tensor precisely, including the device
+                         arrangement (device_matrix) and the alias for the device matrix
+                         (alias_name).
+
+    Returns:
+        Tensor. The mathematically equivalent of the input tensor.
+
+    Raises:
+        TypeError: Reshard takes in Tensor type as the first input param, but got: `type(tensor)`.
+        TypeError: Reshard only support type mindspore.parallel.Layout but got: `type(layout)`.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start-up
+            <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 8 devices.
+
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import ops, nn, Tensor, context, Layout
+        >>> from mindspore.parallel.function import reshard
+        >>> from mindspore.nn.utils import no_init_parameters
+        >>> from mindspore.parallel.auto_parallel import AutoParallel
+        >>> from mindspore.communication import init
+        >>> context.set_context(mode=ms.GRAPH_MODE)
+        >>> init()
+        >>> class Network(nn.Cell):
+        ...     def __init__(self):
+        ...         super().__init__()
+        ...         self.matmul = ops.MatMul()
+        ...         self.relu = ops.ReLU()
+        ...     def construct(self, x, layout):
+        ...         x = self.relu(x)
+        ...         x_reshard = reshard(x, layout)
+        ...         y = Tensor(np.ones(shape=(128, 128)), dtype=ms.float32)
+        ...         x = self.matmul(x_reshard, y)
+        ...         return x
+        >>> layout = Layout((4, 2), ("dp", "mp"))
+        >>> input_layout = layout("dp", "mp")
+        >>> with no_init_parameters():
+        ...     net = Network()
+        >>> parallel_net = AutoParallel(net, parallel_mode='sharding_propagation')
+        >>> tensor = Tensor(np.ones(shape=(128, 128)), dtype=ms.float32)
+        >>> out = parallel_net(tensor, input_layout)
+    """
+    if group_size() == 1:
+        return tensor
+    if not isinstance(tensor, Tensor):
+        raise TypeError(f"Reshard takes in Tensor type as the first input param, but got: {type(tensor)}.")
+    if not isinstance(layout, Layout):
+        raise TypeError(f"Reshard only support type mindspore.parallel.Layout, but got: {type(layout)}.")
+
+    def layout_to_tuple(layout):
+        layout_dict = layout.to_dict()
+        tensor_map = layout_dict["tensor_map"]
+        device_matrix_rev = layout_dict["device_matrix"][::-1]
+        axis_stgy = ()
+        for ind in tensor_map:
+            if ind == -1:
+                axis_stgy += (1,)
+            else:
+                axis_stgy += (device_matrix_rev[ind],)
+        return axis_stgy
+
+    in_strategy = layout_to_tuple(layout)
+    _reshard = _get_cache_prim(P.Reshard)(in_layout=(layout,), out_layout=(layout,), in_strategy=(in_strategy,))
+    return _reshard(tensor)
+
+
+def _redistribute(tensor, dst_dtensor_info):
+    """
+    Redistribute the tensor from the source sharding strategy to the destination sharding strategy.
+
+    Args:
+        tensor (Tensor): The source tensor.
+        dst_dtensor_info (_DistributedTensorInfo): The destination sharding strategy.
+
+    Returns:
+        Tensor, value is same as the source tensor, but the sharding strategy is the destination sharding strategy.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor, Layout, _DistributedTensorInfo
+        >>>
+        >>> init()
+        >>> layout = Layout((2, 1), ("dp", "mp"))
+        >>> src_layout = layout("dp", "mp")
+        >>> distributed_info = _DistributedTensorInfo(src_layout)
+        >>> x = Tensor(np.ones([2, 2]).astype(np.float32))
+        >>> out = x.redistribute(distributed_info)
+        >>> print(out)
+        [[1. 1.]]
+    """
+    from mindspore.parallel._cell_wrapper import RedistributionCell, _insert_virtual_pp_dim
+    if not isinstance(dst_dtensor_info, _DistributedTensorInfo):
+        raise TypeError(
+            "dst_dtensor_info should be _DistributedTensorInfo type, but got {}".format(type(dst_dtensor_info)))
+    run_mode = context.get_context("mode")
+    context.set_context(mode=context.GRAPH_MODE)
+    og_auto_parallel_context, pp_config = _get_all_auto_parallel_context()
+    context.reset_auto_parallel_context()
+    tensor_data = tensor
+    all_reduce_data = False
+    # If src_pp_stages is less than or equal to dst_pp_stages, the parameters of each pp stage of src can be
+    # directly swapped to the corresponding card of dst
+    # rank0 01 11           01
+    # rank1 02 12           02
+    #  pp1   ------>  pp2
+    # rank2 03 13           11
+    # rank3 04 14           12
+    # if dtensor info is None, return the all 1 strategy as from dtensor info
+    if tensor._dtensor_info is None:
+        all_dev_num = get_group_size()
+        dev_mat = Layout((all_dev_num,), ("replica",))
+        tensor_map = ["None"] * len(tensor.shape)
+        layout = dev_mat(*tensor_map)
+        tensor._dtensor_info = _DistributedTensorInfo(layout)
+    src_layout_info = tensor._dtensor_info.layout.to_dict()
+    dst_layout_info = dst_dtensor_info.layout.to_dict()
+    if len(tensor._dtensor_info.layout.to_dict()["rank_list"]) < len(dst_dtensor_info.layout.to_dict()["rank_list"]):
+        # If src_pp_stages is greater than dst_pp_stages, the weights of the corresponding cards need to
+        # be communicated via AllReduce to swap. Need to communicate src rank0's 01 to src rank2,
+        # so that rank2 holds param0's data. Similarly, communicate rank1's 02 to rank3
+        # rank0 01           01 11
+        # rank1 02           02 12
+        # pp2 ------->  pp1
+        # rank2 11           03 13
+        # rank3 12           04 14
+        from mindspore.parallel._cell_wrapper import CommTensorDataForPP
+        if get_rank() in dst_dtensor_info.layout.to_dict()["rank_list"]:
+            comm_tensor_cache_key = (
+                f"{src_layout_info['device_matrix']}, {src_layout_info['tensor_map']}, {src_layout_info['rank_list']}"
+                f" -> "
+                f"{dst_layout_info['device_matrix']}, {dst_layout_info['tensor_map']}, {dst_layout_info['rank_list']}")
+            global COMM_TENSOR_CELL_CACHE
+            if comm_tensor_cache_key not in COMM_TENSOR_CELL_CACHE:
+                comm_tensor_data_func = CommTensorDataForPP(tensor._dtensor_info, dst_dtensor_info)
+                COMM_TENSOR_CELL_CACHE[comm_tensor_cache_key] = comm_tensor_data_func
+                logger.debug(f"comm_tensor_cache_key is {comm_tensor_cache_key}, not match cache")
+            else:
+                comm_tensor_data_func = COMM_TENSOR_CELL_CACHE[comm_tensor_cache_key]
+                logger.debug(f"comm_tensor_cache_key is {comm_tensor_cache_key}, match cache")
+            if not comm_tensor_data_func._current_rank_has_data:
+                new_tensor_shape = tuple([tensor_data.shape[i] // tensor._dtensor_info.sharding_strategy[i]
+                                          for i in range(len(tensor.shape))])
+                tensor_data = comm_tensor_data_func.comm_data(ops.zeros(new_tensor_shape, tensor.dtype))
+            else:
+                tensor_data = comm_tensor_data_func.comm_data(tensor)
+            all_reduce_data = True
+    ms.communication.comm_func.barrier()
+    if src_layout_info['device_matrix'] == dst_layout_info['device_matrix'] and src_layout_info['tensor_map'] == \
+            dst_layout_info['tensor_map']:
+        return tensor_data
+    dataset_strategy = (_insert_virtual_pp_dim(tensor._dtensor_info.layout),)
+    if get_rank() not in tensor._dtensor_info.layout.to_dict()["rank_list"] and not all_reduce_data:
+        dataset_strategy = "full_batch"
+    context.set_auto_parallel_context(dataset_strategy=dataset_strategy,
+                                      parallel_mode="semi_auto_parallel", device_num=get_group_size())
+    global REDIST_CELL_CACHE
+    redist_cache_key = (f"{src_layout_info['device_matrix']}, {src_layout_info['tensor_map']} -> "
+                        f"{dst_layout_info['device_matrix']}, {dst_layout_info['tensor_map']}")
+    if redist_cache_key in REDIST_CELL_CACHE.keys():
+        logger.debug(f"redist_cache_key is {redist_cache_key}, match cache")
+        redist_func = REDIST_CELL_CACHE[redist_cache_key]
+    else:
+        logger.debug(f"redist_cache_key is {redist_cache_key}, not match cache")
+        redist_func = RedistributionCell(tensor._dtensor_info.layout, dst_dtensor_info.layout)
+        REDIST_CELL_CACHE[redist_cache_key] = redist_func
+    redist_func.set_train(True)
+    redist_tensor_data = redist_func(tensor_data)
+    context.reset_auto_parallel_context()
+    _recover_auto_parallel_context(og_auto_parallel_context, pp_config)
+    context.set_context(mode=run_mode)
+    redist_tensor_data._dtensor_info = dst_dtensor_info
+    return redist_tensor_data
+
+
+__all__ = [
+    'reshard'
+]
+
+__all__.sort()

mindspore/parallel/nn/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+# ============================================================================
+"""
+Interfaces for parallel-related functionality
+"""
+from __future__ import absolute_import
+
+from mindspore.parallel.nn.parallel_grad_reducer import PipelineGradReducer
+from mindspore.parallel.nn.parallel_cell_wrapper import PipelineCell, Pipeline, MicroBatchInterleaved, GradAccumulation
+
+__all__ = []
+__all__.extend(parallel_grad_reducer.__all__)
+__all__.extend(parallel_cell_wrapper.__all__)

mindspore/parallel/nn/parallel_cell_wrapper.py

@@ -0,0 +1,263 @@
+# 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.
+# ============================================================================
+"""Cell_wrapper."""
+from __future__ import absolute_import
+from __future__ import division
+
+from mindspore import nn
+from mindspore.ops import operations as P
+from mindspore.nn.cell import Cell
+from mindspore.nn.wrap.cell_wrapper import _MicroBatch
+from mindspore import log as logger
+
+
+__all__ = ['PipelineCell', 'Pipeline', 'MicroBatchInterleaved', 'GradAccumulation']
+
+
+class PipelineCell(Cell):
+    """
+    Slice MiniBatch into finer-grained MicroBatch for use in pipeline-parallel training.
+
+    Note:
+        micro_size must be greater or equal to pipeline stages.
+
+    Args:
+        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.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore.nn as nn
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> 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):
+        super(PipelineCell, self).__init__(auto_prefix=False)
+        self.network = network
+        self.micro_inputs = nn.CellList()
+        self.micro_size = micro_size
+        self.add_list = []
+        if not isinstance(network, Cell):
+            raise TypeError("For 'PipelineCell', the argument 'network' must cell type, "
+                            "but got the type : {}.".format(type(network)))
+        if not isinstance(micro_size, int):
+            raise TypeError("For 'PipelineCell', the argument 'micro_size' must be integer, "
+                            "but got the type : {}.".format(type(micro_size)))
+        if micro_size <= 0:
+            raise ValueError("For 'PipelineCell', the argument 'micro_size' must be large than 0, "
+                             "but got {}.".format(micro_size))
+        for i in range(micro_size):
+            micro_input = _MicroBatch(micro_size)
+            self.micro_inputs.append(micro_input)
+            self.add = P.Add().add_prim_attr("pipeline_end", i)
+            self.add_list.append(self.add)
+        self._get_attr_from_cell(network)
+
+        # prase stage_config
+        config_dict = {}
+        if stage_config is not None:
+            for cell_name, stage_num in stage_config.items():
+                config_cell_name = cell_name
+                config_stage_num = stage_num
+                config_dict[config_cell_name] = config_stage_num
+
+        # set cell.stage_config
+            for cell_name, cell in self.network.cells_and_names():
+                for config_cell_name, config_stage_num in config_dict.copy().items():
+                    if not cell_name or not config_cell_name:
+                        continue
+                    if cell_name == config_cell_name:
+                        setattr(cell, "pipeline_stage", config_stage_num)
+                        del config_dict[config_cell_name]
+
+            for config_cell_name, config_stage_num in config_dict.copy().items():
+                if str(network) == config_cell_name:
+                    setattr(network, "pipeline_stage", config_stage_num)
+                    del config_dict[config_cell_name]
+
+            # if there are any config elements left, print them
+            if config_dict:
+                for config_cell_name, config_stage_num in config_dict.items():
+                    logger.error("pipeline_cell stage_config set pipeline_stage fail!")
+                    logger.warning("config cell name:" + str(config_cell_name) +
+                                   " config stage num:" + str(config_stage_num))
+                logger.warning("network:" + str(self.network))
+                logger.warning("cell name available:")
+                for cell_name, cell in self.network.cells_and_names():
+                    logger.warning(cell_name)
+                raise KeyError("For 'PipelineCell', the argument 'stage_config' : {} is not "
+                               "found in 'network' : {}".format(config_dict, network))
+
+    def construct(self, *inputs):
+        ret = None
+        for i in range(self.micro_size):
+            micro_input = self.micro_inputs[i](i, *inputs)
+            output = self.network(*micro_input)
+            if ret is not None:
+                ret = self.add_list[i](ret, output)
+            else:
+                ret = output
+        return ret
+
+
+class Pipeline(PipelineCell):
+    """
+    Specify the number of micro_batch for pipeline parallelism and the division rules for stage.
+
+    Note:
+        micro_size must be greater or equal to pipeline stages.
+
+    Args:
+        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``.
+
+    Raises:
+        TypeError: The type of `net` is not cell.
+        TypeError: If the type of `micro_size` is not int.
+        ValueError: When `micro_size` is 0 or negative value.
+        KeyError: `dict` cell name matching exception,
+            there are remaining configuration items after traversing all `cell` under the current net.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore.parallel.nn import Pipeline
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = Pipeline(net, 4, stage_config={"cell_name_0": 0, "cell_name_1": 1})
+    """
+
+
+class MicroBatchInterleaved(Cell):
+    """
+    Implement the static graph parallel multi-copy splitting function to enable concurrent computation
+    and communication.
+    Application scenario: When there is model parallelism in semi-automatic mode
+    and network, if the first slice data is calculating forward, the second slice data will execute the
+    communication operators at the same time, to achieve the performance acceleration of communication and computing
+    concurrency.
+
+    Args:
+        network (Cell): The target network to wrap.
+        interleave_num (int, optional): split num of batch size. Default: ``2`` .
+
+    Inputs:
+        tuple[Tensor]. It's the same with the input of the `network` .
+
+    Outputs:
+        The wrapped input. The output of the input `network` should be a Tensor.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore.nn as nn
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = nn.MicroBatchInterleaved(net, 2)
+    """
+    def __init__(self, network, interleave_num=2):
+        super(MicroBatchInterleaved, self).__init__(auto_prefix=False)
+        if not isinstance(interleave_num, int):
+            raise TypeError("For 'MicroBatchInterleaved', the argument 'interleave_num' must be integer, "
+                            "but got the type : {}.".format(type(interleave_num)))
+        if interleave_num <= 0:
+            raise ValueError("For 'MicroBatchInterleaved', the argument 'interleave_num' must be large than 0, "
+                             "but got {}.".format(interleave_num))
+        self.network = network
+        self.interleave_num = interleave_num
+        self.interleave_inputs = nn.CellList()
+        self.add = P.Add().add_prim_attr("micro_interleaved_add_flag", True)
+        for _ in range(interleave_num):
+            interleave_data = _MicroBatch(interleave_num)
+            interleave_data.strided_slice.add_prim_attr("strided_slice_flag", True)
+            interleave_data.strided_slice.add_prim_attr("interleave_num", interleave_num)
+            self.interleave_inputs.append(interleave_data)
+        self._get_attr_from_cell(network)
+
+    def construct(self, *inputs):
+        output = 0.0
+        for i in range(self.interleave_num):
+            interleave_input = self.interleave_inputs[i](i, *inputs)
+            output = self.add(output, self.network(*interleave_input))
+        return output
+
+
+class GradAccumulation(Cell):
+    """
+    Implementation of parallel gradient accumulation for static graphs.
+
+    Args:
+        network (Cell): The target network to wrap.
+        micro_size (int): MicroBatch size.
+
+    Raises:
+        TypeError: The type of `network` is not cell.
+        TypeError: If the type of `micro_size` is not int.
+        ValueError: When `micro_size` is 0 or negative value.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore.parallel.nn import GradAccumulation
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = GradAccumulation(net, 4)
+    """
+    def __init__(self, network, micro_size):
+        super(GradAccumulation, self).__init__(auto_prefix=False)
+        self.network = network
+        self.micro_inputs = nn.CellList()
+        self.micro_size = micro_size
+        self.add_list = []
+        if not isinstance(network, Cell):
+            raise TypeError("For 'GradAccumulation', the argument 'network' must cell type, "
+                            "but got the type : {}.".format(type(network)))
+        if not isinstance(micro_size, int):
+            raise TypeError("For 'GradAccumulation', the argument 'micro_size' must be integer, "
+                            "but got the type : {}.".format(type(micro_size)))
+        if micro_size <= 0:
+            raise ValueError("For 'GradAccumulation', the argument 'micro_size' must be large than 0, "
+                             "but got {}.".format(micro_size))
+        for i in range(micro_size):
+            micro_input = _MicroBatch(micro_size)
+            micro_input.strided_slice.add_prim_attr("grad_accu_num", micro_size)
+            self.micro_inputs.append(micro_input)
+            self.add = P.Add().add_prim_attr("forward_end", i)
+            self.add_list.append(self.add)
+        self._get_attr_from_cell(network)
+
+    def construct(self, *inputs):
+        ret = None
+        for i in range(self.micro_size):
+            micro_input = self.micro_inputs[i](i, *inputs)
+            output = self.network(*micro_input)
+            if ret is not None:
+                ret = self.add_list[i](ret, output)
+            else:
+                ret = output
+        return ret