mindspore 2.5.0__cp310-cp310-win_amd64.whl → 2.6.0rc1__cp310-cp310-win_amd64.whl

mindspore/parallel/shard.py

@@ -15,17 +15,77 @@
 """shard"""
 
 import copy
+import numpy as np
 import mindspore as ms
 from mindspore import log as logger
 from mindspore._c_expression import Shard_
 
 
-class Layout:
+class _DistributedTensorInfo:
+    """
+    Describe the distributed information of a tensor.
+
+    Args:
+        distributed_info (Union[Layout, DeviceMesh]): The distributed information of a tensor.
+
+    Raises:
+        TypeError: If `distributed_info` is not a Layout type.
+
+    Examples:
+        >>> from mindspore import _DistributedTensorInfo, Layout
+        >>> layout = Layout((2, 2), ("dp", "mp"))
+        >>> src_layout = layout("dp", "mp")
+        >>> distributed_info = _DistributedTensorInfo(src_layout)
+        >>> print(distributed_info.sharding_strategy)
+        [2, 2]
     """
-    Parallel layout describes the detailed sharding information.
 
-    For more detailed information, refer to the file `Higher-order Operator-level Parallelism
-    <https://www.mindspore.cn/docs/en/master/model_train/parallel/advanced_operator_parallel.html>`_.
+    def __init__(self, distributed_info):
+        if isinstance(distributed_info, Layout):
+            self._layout = distributed_info
+            self._distributed_info = distributed_info
+        else:
+            raise TypeError(
+                f"DistributedTensorInfo only supports Layout or DeviceMesh as input, but got {type(distributed_info)}")
+        self._sharding_strategy = None
+
+    @property
+    def layout(self):
+        """return layout of current tensor"""
+        return self._layout
+
+    @property
+    def distributed_info(self):
+        """return the distributed info, it depends on user's input """
+        return self._distributed_info
+
+    @property
+    def sharding_strategy(self):
+        """return the sharding strategy of current tensor"""
+        if self._sharding_strategy is None:
+            layout_info = self._layout.to_dict()
+            device_matrix = layout_info["device_matrix"]
+            tensor_map = layout_info["tensor_map"]
+            sharding_strategy = []
+            for map_value in tensor_map:
+                if isinstance(map_value, (tuple, list)):
+                    shard_size = 1
+                    for value in map_value:
+                        if value != -1:
+                            shard_size *= device_matrix[len(device_matrix) - value - 1]
+                    sharding_strategy.append(shard_size)
+                else:
+                    if map_value != -1:
+                        sharding_strategy.append(device_matrix[len(device_matrix) - map_value - 1])
+                    else:
+                        sharding_strategy.append(1)
+            self._sharding_strategy = sharding_strategy
+        return self._sharding_strategy
+
+
+class Layout:
+    """
+    Topological abstraction describing cluster devices for tensor slice placement on the cluster.
 
     Note:
         - It is valid only in semi auto parallel or auto parallel mode.
@@ -38,28 +98,36 @@ class Layout:
         alias_name (tuple): The alias name for each axis of device_matrix, its length shoits element type is string.
                             When using "interleaved_parallel" as an alias name, the tensor would be split into multiple
                             copies on the corresponding partition dimension on a single card.
+        rank_list (list, optional): Data is allocated to the device according to rank_list. Default: ``None``.
+
     Raises:
         TypeError: `device_matrix` is not a tuple type.
         TypeError: `alias_name` is not a tuple type.
+        TypeError: 'rank_list' is not a list type.
         ValueError: `device_matrix` length is not equal to `alias_name` length.
         TypeError: The element of `device_matrix` is not int type.
         TypeError: The element of `alias_name` is not a str type.
+        TypeError: The element of `rank_list` is not int type.
         ValueError: The element of `alias_name` is an empty str.
         ValueError: The element of `alias_name` is "None".
         ValueError: `alias_name` contains repeated element.
 
+    Supported Platforms:
+        ``Ascend``
+
     Examples:
-        >>> from mindspore import Layout
+        >>> from mindspore.parallel import Layout
         >>> layout = Layout((2, 2, 2), ("dp", "sp", "mp"))
         >>> layout0 = layout("dp", "mp")
         >>> print(layout0.to_dict())
-        {"device_matrix": (2, 2, 2), "tensor_map": (2, 0), "interleaved_parallel": False}
+        {"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.
         >>> layout = Layout((2, 2, 2), ("dp", "sp", "interleaved_parallel"))
         >>> layout1 = layout(("dp", "interleaved_parallel"), "sp")
     """
 
-    def __init__(self, device_matrix, alias_name):
+    def __init__(self, device_matrix, alias_name, rank_list=None):
         if not isinstance(device_matrix, tuple):
             raise TypeError(f'device_matrix must be tuple type, but got:{type(device_matrix)}')
         if not isinstance(alias_name, tuple):
@@ -85,6 +153,20 @@ class Layout:
         self._device_shape = device_matrix
         self._alias_name = alias_name
         self._tensor_map = None
+        self._rank_list = list(range(np.prod(np.array(self._device_shape))))
+        if rank_list is not None:
+            if not isinstance(rank_list, list):
+                raise TypeError(f"The rank_list should be a list, but got {type(rank_list).__name__}.")
+            for in_ele in rank_list:
+                if not isinstance(in_ele, int):
+                    raise TypeError(f"The element of rank_list should be int, but got {type(in_ele).__name__}.")
+            if len(np.array(rank_list).shape) != 1:
+                raise ValueError(
+                    f"The rank_list should be a 1-D list, but got {len(np.array(rank_list).shape)}-D list.")
+            if len(rank_list) != np.prod(np.array(self._device_shape)):
+                raise ValueError(f"The length of rank_list should be equal to the product of device_matrix, "
+                                 f"but got {len(rank_list)} and {np.prod(np.array(self._device_shape))}.")
+            self._rank_list = rank_list
 
     def __call__(self, *tensor_map):
         self._tensor_map = ()
@@ -125,8 +207,8 @@ class Layout:
             raise ValueError("The tensor_map of layout is None")
         interleaved_parallel = "interleaved_parallel" in self._alias_name
         return {"device_matrix": self._device_shape, "tensor_map": self._tensor_map,
-                "interleaved_parallel": interleaved_parallel, "alias_name": self._alias_name}
-
+                "interleaved_parallel": interleaved_parallel, "alias_name": self._alias_name,
+                "rank_list": self._rank_list}
 
 
 class Shard(Shard_):
@@ -144,18 +226,6 @@ class Shard(Shard_):
         self.level = None
 
     def __call__(self, fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascend", level=0):
-        parallel_mode = ms.context.get_auto_parallel_context("parallel_mode")
-        if parallel_mode not in ("auto_parallel", "semi_auto_parallel"):
-            raise AssertionError(
-                f"Cell shard only supports auto parallel and semi auto parallel.")
-        if ms.context.get_context("device_target") not in ("Ascend", "GPU"):
-            raise AssertionError(
-                f"'Shard' now only supports 'Ascend' and 'GPU'")
-        if parallel_mode == "auto_parallel" and \
-            ms.context.get_auto_parallel_context("search_mode") != "sharding_propagation":
-            raise AssertionError(f"'search_mode' must be 'sharding_propagation' for 'Shard' when the "
-                                 f"'parallel_mode' is 'auto_parallel.'")
-
         if not isinstance(in_strategy, tuple):
             raise TypeError(
                 f"For 'Shard', the 'in_strategy' should be a tuple, but got {type(in_strategy).__name__}.")
@@ -184,7 +254,8 @@ 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:
+        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'.")
@@ -196,13 +267,13 @@ class Shard(Shard_):
 
         if isinstance(fn, ms.nn.Cell):
             for param in fn.trainable_params():
-                param.is_in_shard = True
+                param.param_info.is_in_pynative_shard = True
 
         # Set parameter layout to corresponding parameter
         self._set_param_layout_into_parameter(fn, parameter_plan)
 
         def shard_fn(*args):
-            @ms.common.jit(hash_args=fn)
+            @ms.common.jit(hash_args=fn, backend="ms_backend")
             def after_shard(*args):
                 return shard_(fn, in_strategy, out_strategy, device, level)(*args)
 
@@ -293,7 +364,7 @@ class Shard(Shard_):
         for stra in strategy:
             if not isinstance(stra, (tuple, Layout)):
                 raise TypeError(
-                    f"The '{log_info}' should be a tuple(tuple(int)) or tuple(mindspore.Layout), "
+                    f"The '{log_info}' should be a tuple(tuple(int)) or tuple(mindspore.parallel.Layout), "
                     f"but got {type(stra).__name__}")
             if isinstance(stra, Layout):
                 strategy_set.add("layout")
@@ -315,7 +386,7 @@ class Shard(Shard_):
             for in_ele in layout:
                 if not isinstance(in_ele, Layout):
                     raise TypeError(f"The {log_info} item should be a object of class Layout.")
-                layout_value += (in_ele.to_dict(),)
+                layout_value += ({k: v for k, v in in_ele.to_dict().items() if k != "rank_list"},)
         return layout_value
 
     def _check_tuple_strategy(self, dim_strategy):
@@ -326,8 +397,8 @@ class Shard(Shard_):
 
 def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascend", level=0):
     """
-    Defining the input and output layouts of this cell and the parallel strategies of remaining ops will be
-    generated by sharding propagation. In PyNative mode, use this method to specify a Cell for distributed
+    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,
     strategy for others will be set by sharding propagation.
     in_strategy and out_strategy define the input and output layout respectively.
@@ -337,33 +408,37 @@ 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 ms.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".
-        If the input contain Parameter, its strategy should be set in `in_strategy`.
+        - 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".
+        - If the input contain Parameter, its strategy should be set in `in_strategy`.
+        - This method currently does not support dynamic shapes.
 
     Args:
         fn (Union[Cell, Function]): Function to be executed in parallel.
-                                    Its arguments and return value must be Tensor or Parameter.
+                                    Its arguments and return value must be Tensor.
                                     If `fn` is a Cell with parameters, `fn` needs to be an instantiated object,
                                     otherwise its arguments cannot be accessed.
         in_strategy (tuple): Define the layout of inputs, each element of the tuple should be a tuple(int) or
-                             tuple(mindspore.Layout).
+                             tuple(mindspore.parallel.Layout).
                              Tuple defines the layout of the corresponding input.
-        out_strategy (Union[tuple, None]): Define the layout of outputs similar with `in_strategy`.
-                                           It is not in use right now. Default: ``None`` .
-        parameter_plan (Union[dict, None]): Define the layout for the specified parameters. Each element in dict
+        out_strategy (Union[tuple, None], optional): Define the layout of outputs similar with `in_strategy`.
+                                           Default: ``None`` .
+        parameter_plan (Union[dict, None], optional): Define the layout for the specified parameters.
+                                            Each element in dict
                                             defines the layout of the parameter like "param_name: layout".
                                             The key is a parameter name of type 'str'.
-                                            The value is a 1-D integer tuple or a 1-D mindspore.Layout tuple,
+                                            The value is a 1-D integer tuple or a 1-D mindspore.parallel.Layout tuple,
                                             indicating the corresponding layout.
                                             If the parameter name is incorrect or the corresponding parameter
-                                            has been set, the parameter setting will be ignored.
+                                            has been set, the parameter setting will be ignored. Supported
+                                            only when `fn` is a Cell with parameters.
                                             Default: ``None`` .
-        device (string): Select a certain `device` target. It is not in use right now.
-                         Support ["CPU", "GPU", "Ascend"]. Default: ``"Ascend"`` .
-        level (int): 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`` .
+        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.
@@ -373,26 +448,28 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
         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.Layout).
-        TypeError: If any element in `out_strategy` is not a tuple(int) or tuple(mindspore.Layout).
+        TypeError: If any element in `in_strategy` is not a tuple(int) or tuple(mindspore.parallel.Layout).
+        TypeError: If any element in `out_strategy` is not a tuple(int) or tuple(mindspore.parallel.Layout).
         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.Layout).
+        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`` ``GPU``
+        ``Ascend``
 
     Examples:
         >>> import numpy as np
         >>> import mindspore as ms
-        >>> from mindspore import Tensor, nn
+        >>> from mindspore import Tensor, nn, ops
         >>> from mindspore.communication import init
+        >>> from mindspore.parallel import shard
+        >>> from mindspore.parallel import Layout
+        >>> from mindspore.nn.utils import no_init_parameters
+        >>> from mindspore.parallel.auto_parallel import AutoParallel
         >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> init()
-        >>> ms.set_auto_parallel_context(parallel_mode="auto_parallel", search_mode="sharding_propagation",
-        ...                              device_num=8)
         >>>
         >>> # Case 1: cell uses functional
         >>> class BasicBlock(nn.Cell):
@@ -404,7 +481,7 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
         >>>             x = ops.abs(x)
         >>>             return x + y
         >>>         # shard a function with tuple(int) strategies
-        >>>         self.shard_my_add = ms.shard(my_add, in_strategy=((2, 2), (1, 4)), out_strategy=((4, 1),))
+        >>>         self.shard_my_add = shard(my_add, in_strategy=((2, 2), (1, 4)), out_strategy=((4, 1),))
         >>>
         >>>     def construct(self, x, u):
         >>>         x = self.gelu(x)
@@ -432,7 +509,7 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
         >>>         super(Net, self).__init__()
         >>>         # setting cell sharding strategy and parameter_plan by tuple(int)
         >>>         self.layer_net1 = NetForward()
-        >>>         self.layer_net1_shard = ms.shard(self.layer_net1, in_strategy=((4, 2), (2, 1)),
+        >>>         self.layer_net1_shard = shard(self.layer_net1, in_strategy=((4, 2), (2, 1)),
         ...                                          parameter_plan={"self.layer_net1.block1.weight": (4, 1)})
         >>>
         >>>         # setting cell sharding strategy and parameter_plan by tuple(ms.Layout)
@@ -440,7 +517,7 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
         >>>         layout = Layout((4, 2, 1), ("dp", "mp", "sp"))
         >>>         in_layout = (layout("dp", "mp"), layout("mp", "sp"))
         >>>         param_layout = layout("dp", "sp")
-        >>>         self.layer_net2_shard = ms.shard(self.layer_net2, in_strategy=in_layout,
+        >>>         self.layer_net2_shard = shard(self.layer_net2, in_strategy=in_layout,
         ...                                          parameter_plan={"self.layer_net2.block2.weight": param_layout})
         >>>         self.flatten = nn.Flatten()
         >>>         self.layer1 = nn.Dense(64, 64)
@@ -458,26 +535,25 @@ def shard(fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascen
         >>>         x = self.matmul(x, Tensor(np.ones(shape=(32, 32)), dtype=ms.float32))
         >>>         return x
         >>>
-        >>> net = Net()
+        >>> with no_init_parameters():
+        >>>     net = Net()
         >>> x = Tensor(np.ones(shape=(64, 1, 8, 8)), dtype=ms.float32)
         >>> y = Tensor(np.ones(shape=(64, 1, 8, 8)), dtype=ms.float32)
-        >>> net(x, y)
+        >>> parallel_net = AutoParallel(net, parallel_mode='sharding_propagation')
+        >>> parallel_net(x, y)
         >>>
         >>> # Case 2: function uses functional sharding
         >>> def test_shard(x, y):
         ...     return x + y
         >>> x = Tensor(np.ones(shape=(32, 10)), dtype=ms.float32)
         >>> y = Tensor(np.ones(shape=(32, 10)), dtype=ms.float32)
-        >>> output = ms.shard(test_shard, in_strategy=((4, 2), (4, 2)))(x, y)
+        >>> output = shard(test_shard, in_strategy=((4, 2), (4, 2)))(x, y)
         >>> print(output.shape)
         (32, 10)
 
-    Tutorial Examples:
-        - `Functional Operator Sharding
-          <https://www.mindspore.cn/docs/en/master/model_train/parallel/shard_function_parallel.html>`_
-        - `mindspore.Layout
-          <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.Layout.html>`_
     """
+    if ms.communication.management.get_group_size() == 1:
+        return fn
     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.")