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

mindspore/parallel/nn/parallel_cell_wrapper.py

@@ -17,6 +17,8 @@
 from __future__ import absolute_import
 from __future__ import division
 
+__all__ = ['PipelineCell', 'Pipeline', 'MicroBatchInterleaved', 'GradAccumulation']
+
 from mindspore import nn
 from mindspore.ops import operations as P
 from mindspore.nn.cell import Cell
@@ -24,12 +26,10 @@ 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.
+    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.
@@ -38,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``
@@ -49,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()
@@ -101,15 +103,46 @@ class PipelineCell(Cell):
                                    " 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():
+                for cell_name, _ 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):
+        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
         for i in range(self.micro_size):
-            micro_input = self.micro_inputs[i](i, *inputs)
+            micro_input = self.micro_inputs[i](i, *args, **kwargs)
             output = self.network(*micro_input)
             if ret is not None:
                 ret = self.add_list[i](ret, output)
@@ -120,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.
@@ -129,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.
@@ -197,10 +233,10 @@ class MicroBatchInterleaved(Cell):
             self.interleave_inputs.append(interleave_data)
         self._get_attr_from_cell(network)
 
-    def construct(self, *inputs):
+    def construct(self, *args, **kwargs):
         output = 0.0
         for i in range(self.interleave_num):
-            interleave_input = self.interleave_inputs[i](i, *inputs)
+            interleave_input = self.interleave_inputs[i](i, *args, **kwargs)
             output = self.add(output, self.network(*interleave_input))
         return output
 
@@ -251,10 +287,10 @@ class GradAccumulation(Cell):
             self.add_list.append(self.add)
         self._get_attr_from_cell(network)
 
-    def construct(self, *inputs):
+    def construct(self, *args, **kwargs):
         ret = None
         for i in range(self.micro_size):
-            micro_input = self.micro_inputs[i](i, *inputs)
+            micro_input = self.micro_inputs[i](i, *args, **kwargs)
             output = self.network(*micro_input)
             if ret is not None:
                 ret = self.add_list[i](ret, output)

mindspore/parallel/nn/parallel_grad_reducer.py

@@ -15,7 +15,8 @@
 """parallel serialization"""
 from __future__ import absolute_import
 
-from mindspore import context
+__all__ = ['PipelineGradReducer']
+
 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
@@ -25,8 +26,6 @@ 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")
@@ -81,6 +80,7 @@ class PipelineGradReducer(Cell):
         >>> import mindspore as ms
         >>> from mindspore import nn, ops, Tensor
         >>> from mindspore.communication import init
+        >>> from mindspore.parallel.auto_parallel import AutoParallel
         >>>
         >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> ms.reset_auto_parallel_context()
@@ -113,7 +113,7 @@ class PipelineGradReducer(Cell):
         >>> 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 = nn.PipelineCell(nn.WithLossCell(net, loss_fn), 2)
         >>> net_with_loss.set_train()
         >>> def forward_fn(inputs, target):
         ...     loss = net_with_loss(inputs, target)
@@ -135,11 +135,10 @@ class PipelineGradReducer(Cell):
         >>> label = Tensor(np.ones([size, out_features]).astype(np.float32))
         >>> loss, _ = train_one_step(inputs, label)
         >>> print(loss)
-        46.36721
+        46.304886
     """
     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)
@@ -151,19 +150,13 @@ class PipelineGradReducer(Cell):
             self.opt_shard = opt_shard
 
     @jit
-    def construct(self, grads):
+    def construct(self, *args, **kwargs):
         new_grads = None
         if self.opt_shard:
-            grads = self.grad_reducer(grads)
+            grads = self.grad_reducer(*args, **kwargs)
             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/shard.py

@@ -121,8 +121,7 @@ class Layout:
         >>> layout0 = layout("dp", "mp")
         >>> print(layout0.to_dict())
         {"device_matrix": (2, 2, 2), "tensor_map": (2, 0), "interleaved_parallel": False,
-        'alias_name': {'dp', 'sp', 'mp'}, "rank_list": [0, 1, 2, 3, 4, 1, 6, 7]}
-        >>> # Total device num is 4, but split the tensor in local device into two copies.
+        'alias_name': {'dp', 'sp', 'mp'}, "rank_list": [0, 1, 2, 3, 4, 5, 6, 7]}
         >>> layout = Layout((2, 2, 2), ("dp", "sp", "interleaved_parallel"))
         >>> layout1 = layout(("dp", "interleaved_parallel"), "sp")
     """
@@ -254,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()
@@ -395,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
@@ -408,10 +399,12 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
     The parallel strategies of remaining operators are derived from the strategy specified by the input and output.
 
     Note:
-        - If shard is called, the parallel mode in `set_auto_parallel_context` (parallel_mode) will be set to
-          "auto_parallel" and the search mode (search_mode) to "sharding_propagation".
+        - 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.
@@ -433,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).
@@ -453,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``
@@ -557,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)