mindspore 2.5.0__cp39-cp39-win_amd64.whl → 2.6.0rc1__cp39-cp39-win_amd64.whl

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

mindspore/parallel/nn/parallel_grad_reducer.py

@@ -0,0 +1,169 @@
+# 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.
+# ============================================================================
+"""parallel serialization"""
+from __future__ import absolute_import
+
+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
+from mindspore.common.sparse_tensor import Tensor
+from mindspore.common.api import jit
+from mindspore.common.parameter import Parameter
+from mindspore.nn.layer import Identity
+from mindspore.parallel._utils import _get_enable_parallel_optimizer
+
+__all__ = ['PipelineGradReducer']
+
+
+grad_scale = C.MultitypeFuncGraph("grad_scale")
+shard_grad_scale = C.MultitypeFuncGraph("shard_grad_scale")
+reciprocal = P.Reciprocal()
+
+
+@grad_scale.register("Tensor", "Tensor", "Tensor")
+def tensor_grad_scale_pipeline(scale, grad, accu_grad):
+    accu_grad = F.depend(accu_grad, grad)
+    new_grad = accu_grad * reciprocal(scale)
+    accu_grad = F.depend(accu_grad, new_grad)
+    zeros = F.tensor_mul(accu_grad, 0.0)
+    new_grad = F.depend(new_grad, F.assign(accu_grad, zeros))
+    return new_grad
+
+
+@shard_grad_scale.register("Tensor", "Tensor", "Tensor")
+def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
+    new_grad = grad * reciprocal(scale)
+    accu_grad = F.depend(accu_grad, new_grad)
+    new_grad = F.depend(new_grad, F.assign(accu_grad, F.zeros_like(accu_grad)))
+    return new_grad
+
+
+class PipelineGradReducer(Cell):
+    """
+    Functional training scenarios for gradient statute and accumulation of pipeline parallel.
+
+    Args:
+        parameters (list): the parameters to be updated.
+        scale_sense (float, optional): the scale sense of the gradient. Default: 1.0.
+        opt_shard(bool, optional): if use parallel optimizer, set opt_shard True. Default: ``None``.
+
+    Raise:
+        RuntimeError: If the mode is not graph mode.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/en/master/parallel/rank_table.html>`_
+            for more details.
+
+            This example should be run with multiple devices.
+
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import nn, ops, Tensor
+        >>> from mindspore.communication import init
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> ms.reset_auto_parallel_context()
+        >>> init()
+        >>> ms.set_seed(1)
+        >>>
+        >>> class Network(nn.Cell):
+        ...     def __init__(self, in_features, out_features, sens=1.0):
+        ...         super().__init__()
+        ...         self.layer1 = nn.Dense(in_features, 16)
+        ...         self.relu1 = nn.ReLU()
+        ...         self.layer2 = nn.Dense(16, 16)
+        ...         self.relu2 = nn.ReLU()
+        ...         self.layer3 = nn.Dense(16, out_features)
+        ...
+        ...     def construct(self, x):
+        ...         x = self.layer1(x)
+        ...         x = self.relu1(x)
+        ...         x = self.layer2(x)
+        ...         x = self.relu2(x)
+        ...         logits = self.layer3(x)
+        ...         return logits
+        >>>
+        >>> size, in_features, out_features = 16, 32, 10
+        >>> net = Network(in_features, out_features)
+        >>> net.layer1.pipeline_stage = 0
+        >>> net.relu1.pipeline_stage = 0
+        >>> net.layer2.pipeline_stage = 0
+        >>> net.relu2.pipeline_stage = 1
+        >>> net.layer3.pipeline_stage = 1
+        >>> loss_fn = nn.CrossEntropyLoss()
+        >>> optimizer = nn.SGD(net.trainable_params(), 1e-2)
+        >>> net_with_loss = nn.Pipeline(nn.WithLossCell(net, loss_fn), 2)
+        >>> net_with_loss.set_train()
+        >>> def forward_fn(inputs, target):
+        ...     loss = net_with_loss(inputs, target)
+        ...     return loss
+        >>>
+        >>> grad_fn = ops.value_and_grad(forward_fn, None, net_with_loss.trainable_params())
+        >>> pp_grad_reducer = nn.PipelineGradReducer(optimizer.parameters)
+        >>>
+        >>> @ms.jit
+        >>> def train_one_step(inputs, target):
+        ...     loss, grads = grad_fn(inputs, target)
+        ...     grads = pp_grad_reducer(grads)
+        ...     optimizer(grads)
+        ...     return loss, grads
+        >>>
+        >>> parallel_net = AutoParallel(train_one_step, parallel_mode="semi_auto")
+        >>> parallel_net.pipeline(stages=2)
+        >>> inputs = Tensor(np.ones([size, in_features]).astype(np.float32))
+        >>> label = Tensor(np.ones([size, out_features]).astype(np.float32))
+        >>> loss, _ = train_one_step(inputs, label)
+        >>> print(loss)
+        46.36721
+    """
+    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)
+        self.scale_sense = Parameter(scale_sense, name='scale_sense')
+        self.hyper_map = C.HyperMap()
+        if opt_shard is None:
+            self.opt_shard = _get_enable_parallel_optimizer()
+        else:
+            self.opt_shard = opt_shard
+
+    @jit
+    def construct(self, grads):
+        new_grads = None
+        if self.opt_shard:
+            grads = self.grad_reducer(grads)
+            new_grads = self.hyper_map(F.partial(shard_grad_scale, self.scale_sense * self.degree),
+                                       grads, self.accu_grads)
+        else:
+            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/parameter_broadcast.py

@@ -19,7 +19,10 @@ __all__ = ["parameter_broadcast"]
 
 import numpy as np
 import mindspore as ms
-from mindspore.communication import get_rank, create_group, get_group_size
+from mindspore.communication import create_group, get_group_size
+from mindspore.parallel._utils import _get_auto_parallel_net, _parallel_mode_map, _check_rank
+# disable pylint too broad Exception
+# pylint: disable=W0212
 
 
 def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
@@ -34,7 +37,8 @@ def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
         layout (Dict): Parameter layout dictionary. Come from
             :func:`mindspore.nn.Cell.parameter_layout_dict`
             or read from file(for example: "strategy.ckpt" saved by using the
-            `strategy_ckpt_config` parameter of :func:`mindspore.set_auto_parallel_context`).
+            `strategy_ckpt_config` parameter of
+            :func:`mindspore.parallel.auto_parallel.AutoParallel.save_param_strategy_file` ).
             The key is param name, the value is the layout of this parameter.
         cur_rank (int, optional): current rank id. Default: ``0``.
         initial_rank (int, optional): Start rank id for each pipeline. Default: ``0``.
@@ -45,6 +49,9 @@ def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
         ValueError: Parameter name in `layout` can not be found in
             :func:`mindspore.nn.Cell.parameters_dict`.
 
+    Supported Platforms:
+        ``Ascend``
+
     Examples:
         >>> import os
         >>> import mindspore as ms
@@ -53,11 +60,11 @@ def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
         >>> from mindspore.communication import init
         >>> from mindspore.common.initializer import initializer
         >>> from mindspore.train import Model
-        >>> from mindspore.parallel.parameter_broadcast import parameter_broadcast
         >>> from mindspore.train.serialization import load_checkpoint, load_param_into_net
+        >>> from mindspore.parallel.auto_parallel import AutoParallel
+        >>> from mindspore.parallel import parameter_broadcast
         >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> ms.runtime.set_memory(max_size="28GB")
-        >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL)
         >>> init()
         >>> ms.set_seed(1)
         >>> class Network(nn.Cell):
@@ -90,7 +97,8 @@ def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
         >>> dataset = create_dataset()
         >>> optim = nn.SGD(net.trainable_params(), 1e-2)
         >>> loss = nn.CrossEntropyLoss()
-        >>> model = Model(net, loss_fn=loss, optimizer=optim)
+        >>> parallel_net = AutoParallel(net)
+        >>> model = Model(parallel_net, loss_fn=loss, optimizer=optim)
         >>> model.train(1, dataset)
         >>> ms.save_checkpoint(net, "./simple.ckpt", False)
         >>> layout = model.train_network.parameter_layout_dict
@@ -104,17 +112,20 @@ def parameter_broadcast(net, layout, cur_rank=0, initial_rank=0):
         ...         print("step end, cur step num: ", cb_params.cur_step_num, flush=True)
         >>> model.train(1, dataset, callbacks=[LossCallBack()])
     """
-    if not layout:
+    if not layout or get_group_size() <= 1:
         return
     from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
     from mindspore.nn.wrap.cell_wrapper import AllreduceGraph
-    origin_parallel_mode = ms.get_auto_parallel_context("parallel_mode")
-    if origin_parallel_mode not in ("semi_auto_parallel", "auto_parallel"):
-        return
-    if cur_rank != get_rank():
-        raise ValueError(f"For parameter broadcast, the cur_rank: {cur_rank} is wrong.")
-    if initial_rank % (get_group_size() / ms.get_auto_parallel_context("pipeline_stages")) != 0:
-        raise ValueError(f"For parameter broadcast, the initial_rank: {initial_rank} is wrong.")
+    origin_parallel_mode = ""
+    pipeline_stages = 1
+    parallel_net = _get_auto_parallel_net(net)
+    if type(parallel_net).__name__ == 'AutoParallel':
+        origin_parallel_mode = _parallel_mode_map(parallel_net._parallel_mode)
+        pipeline_stages = parallel_net._pipeline_stages
+    else:
+        origin_parallel_mode = ms.get_auto_parallel_context("parallel_mode")
+        pipeline_stages = ms.get_auto_parallel_context("pipeline_stages")
+    _check_rank(cur_rank, initial_rank, pipeline_stages)
     param_redundancy = get_parameter_redundancy(layout, initial_rank)
     if not param_redundancy:
         return