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

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.set_context(max_device_memory="28GB")
-        >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL)
+        >>> ms.runtime.set_memory(max_size="28GB")
         >>> 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

mindspore/parallel/shard.py

@@ -15,14 +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 _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]
+    """
+
+    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:
     """
-    Parallel layout describes the detailed sharding information.
+    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.
@@ -35,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):
@@ -82,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 = ()
@@ -122,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_):
@@ -141,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__}.")
@@ -181,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'.")
@@ -193,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)
 
@@ -290,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")
@@ -312,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):
@@ -323,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.
@@ -334,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.
@@ -370,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):
@@ -401,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)
@@ -429,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)
@@ -437,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)
@@ -455,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.")