mindspore 1.10.0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/parallel/_utils.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ============================================================================
 """Utils of auto parallel"""
-from __future__ import absolute_import
 from importlib import import_module
 import numpy as np
 import mindspore as ms
@@ -44,6 +43,20 @@ def _is_in_auto_parallel_mode():
     return _get_parallel_mode() in [ms.ParallelMode.SEMI_AUTO_PARALLEL, ms.ParallelMode.AUTO_PARALLEL]
 
 
+def _is_in_data_parallel_mode():
+    return _get_parallel_mode() == ms.ParallelMode.DATA_PARALLEL
+
+
+def _is_in_hybrid_parallel_mode():
+    return _get_parallel_mode() == ms.ParallelMode.HYBRID_PARALLEL
+
+
+def _is_pynative_parallel():
+    parallel_mode = context.get_auto_parallel_context('parallel_mode')
+    return context.get_context('mode') == context.PYNATIVE_MODE and parallel_mode in (
+        context.ParallelMode.SEMI_AUTO_PARALLEL, context.ParallelMode.AUTO_PARALLEL)
+
+
 def _get_full_batch():
     """Get whether to use full_batch."""
     return auto_parallel_context().get_full_batch()
@@ -79,16 +92,25 @@ def _need_to_full():
 
 def _slice_parameter(parameter, phase, layout):
     """Slice python parameter obj according to the layout."""
-    if not parameter.sliced:
-        graph_executor = GraphExecutor_.get_instance()
-        new_param = parameter.init_data(layout, set_sliced=True)
-        parameter = new_param
-        graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
-    else:
+    graph_executor = GraphExecutor_.get_instance()
+    new_param = parameter.init_data(layout, set_sliced=True)
+    parameter = new_param
+    graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
+    if not parameter.sliced and (layout is not None):
         new_tensor = _load_tensor_by_layout(parameter, layout)
         parameter.set_data(new_tensor, True)
 
 
+def _init_optimizer_state(parameter, phase):
+    """init optimizer state"""
+    if not parameter.has_init:
+        return
+    graph_executor = GraphExecutor_.get_instance()
+    new_param = parameter.init_data()
+    parameter = new_param
+    graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
+
+
 def _to_full_shapes(shapes, device_num):
     """Expanding batch dimension according to device_num, adapt to mindspore minddata graph solution."""
     new_shapes = []
@@ -381,3 +403,21 @@ def _infer_rank_list(train_map, predict_map=None):
         else:
             ret[param_name] = (rank_list, False)
     return ret
+
+
+def _grads_divided_by_device_num_if_recomputation(grads):
+    """
+    If in pynative parallel and full_batch is True, divide grads by device num to ensure that the gradients is correct.
+    """
+    if not grads or not _is_pynative_parallel() or not _get_full_batch():
+        return grads
+
+    device_num = Tensor(_get_device_num(), grads[0].dtype)
+    logger.info(f"In PyNative mode, when parallel mode is in "
+                f"({context.ParallelMode.SEMI_AUTO_PARALLEL}, {context.ParallelMode.AUTO_PARALLEL}) and "
+                f"full_batch is Ture, the gradients will be automatically divided by device_num({device_num}).")
+    new_grads = ()
+    for grad in grads:
+        new_grads += (grad / device_num,)
+
+    return new_grads

mindspore/parallel/algo_parameter_config.py

@@ -14,6 +14,8 @@
 # ============================================================================
 """Configuration of parameters for strategy-searching algorithm in auto_parallel"""
 from __future__ import absolute_import
+from __future__ import division
+
 import threading
 
 from mindspore._c_expression import CostModelContext
@@ -21,6 +23,8 @@ from mindspore._checkparam import args_type_check
 
 __all__ = ["get_algo_parameters", "reset_algo_parameters", "set_algo_parameters"]
 
+_PARAMETER_CONFIG = None
+
 
 class _AlgoParameterConfig:
     """
@@ -225,7 +229,7 @@ def set_algo_parameters(**kwargs):
     """
     Set parameters in the algorithm for parallel strategy searching. See a typical use in
     `test_auto_parallel_resnet.py
-    <https://gitee.com/mindspore/mindspore/blob/r1.10/tests/ut/python/parallel/test_auto_parallel_resnet.py>`_.
+    <https://gitee.com/mindspore/mindspore/blob/r2.0/tests/ut/python/parallel/test_auto_parallel_resnet.py>`_.
 
     Note:
         The attribute name is required. This interface works ONLY in AUTO_PARALLEL mode.

mindspore/parallel/checkpoint_transform.py

@@ -0,0 +1,329 @@
+# Copyright 2022 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.
+# ============================================================================
+"""Transform distributed checkpoint"""
+from __future__ import absolute_import
+
+import os
+import glob
+import copy
+from collections import defaultdict
+import numpy as np
+import mindspore as ms
+from mindspore.parallel._parallel_serialization import _rank_list_for_transform_parallel_checkpoint, \
+    _transform_parallel_checkpoint, _get_device_num_from_strategy, _make_dir, \
+    _extract_layout_map, _extract_src_dst_layout_map, _parameter_not_in_local_stage, _extract_pipeline_stage_num, \
+    _merge_protobuf_strategy, _merge_json_strategy
+
+
+__all__ = ["merge_pipeline_strategys", "rank_list_for_transform", "transform_checkpoint_by_rank",
+           "transform_checkpoints"]
+
+
+def merge_pipeline_strategys(src_strategy_dirs, dst_strategy_file):
+    """
+    Merge parallel strategy between all pipeline stages in pipeline parallel mode.
+    For more details about converting distributed Checkpoint, please refer to
+    `Distributed Resilience Training and
+    Inference <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/resilience_train_and_predict.html>`_.
+
+    Note:
+        Strategy file of each pipeline stage should be included in src_strategy_dirs.
+
+    Args:
+        src_strategy_dirs (str): The directory of strategy files including all pipeline stage which is saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'
+        dst_strategy_file (str): The file merged strategy to save.
+
+    Raises:
+        NotADirectoryError: `src_strategy_dirs` is not a directory.
+
+    Examples:
+        >>> # src_strategy_dir/stra0.ckpt, src_strategy_dir/stra1.ckpt ... src_strategy_dir/stra127.ckpt
+        >>> merge_pipeline_strategys("./src_strategy_dir", "./dst_strategy.ckpt")
+
+    """
+    dst_strategy_dir, _ = os.path.split(dst_strategy_file)
+    if not os.path.exists(dst_strategy_dir):
+        _make_dir(dst_strategy_dir, "path")
+    if not os.path.isdir(src_strategy_dirs):
+        raise NotADirectoryError("src_strategy_dirs {} is not a directory.".format(src_strategy_dirs))
+    src_strategy_files_protobuf = glob.glob(os.path.join(src_strategy_dirs, "*.ckpt"))
+    src_strategy_files_json = glob.glob(os.path.join(src_strategy_dirs, "*.json"))
+    if src_strategy_files_protobuf and src_strategy_files_json:
+        raise ValueError("The strategys format should be all '.ckpt' or all '.json'")
+    is_protobuf = len(src_strategy_files_protobuf) > 0
+    if is_protobuf:
+        _merge_protobuf_strategy(src_strategy_files_protobuf, dst_strategy_file)
+    else:
+        _merge_json_strategy(src_strategy_files_json, dst_strategy_file)
+
+
+
+def rank_list_for_transform(rank_id, src_strategy_file=None, dst_strategy_file=None):
+    """
+    List of original distributed checkpoint rank index for obtaining the target checkpoint of a rank_id
+    during the distributed checkpoint conversion. For more details about converting distributed Checkpoint,
+    please refer to `Distributed Resilience Training and
+    Inference <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/resilience_train_and_predict.html>`_.
+
+    Args:
+        rank_id (int): The rank of which distributed checkpoint needs to be obtained after conversion.
+        src_strategy_file (str): Name of source sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'src_strategy_file' is None, it means that the source sharding strategy is
+                                 without any sharing for each parameter. Default:None.
+        dst_strategy_file (str): Name of destination sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'dst_strategy_file' is None, it means that the destination sharding strategy
+                                 is without any sharing for each parameter. Default:None.
+
+    Returns:
+        List, the rank list required for converting the distributed checkpoint of rank_id.
+
+    Raises:
+        ValueError: `src_strategy_file` or dst_strategy_file is incorrect.
+        TypeError: `src_strategy_file` or dst_strategy_file is not a string.
+        TypeError: `rank_id` is not a int.
+
+    Examples:
+        >>> rank_id = 0
+        >>> rank_list = rank_list_for_transform(rank_id, "./src_strategy.ckpt", "./dst_strategy.ckpt")
+        >>> checkpoint_files_map = {}
+        >>> for rank in rank_list:
+        ...     checkpoint_files_map[rank] = "./pangu{}-100_2.ckpt".format(rank)
+
+    """
+    if not isinstance(rank_id, int):
+        raise TypeError("The rank_id should be a int.")
+    if src_strategy_file is None:
+        return [0]
+    src_strategy_list, dst_strategy_list = _extract_src_dst_layout_map(rank_id, src_strategy_file, dst_strategy_file)
+    src_stage_device_num = np.prod(src_strategy_list.get(list(src_strategy_list.keys())[0])[0]) if src_strategy_list \
+                                                                                                   is not None else 1
+    dst_stage_device_num = np.prod(dst_strategy_list.get(list(dst_strategy_list.keys())[0])[0]) if dst_strategy_list \
+                                                                                                   is not None else 1
+
+    if not src_strategy_list:
+        raise ValueError("The src_strategy_file is empty.")
+    local_rank_id = rank_id % dst_stage_device_num if dst_stage_device_num > 1 else rank_id
+    needed_rank_list_in_local_stage = _rank_list_for_transform_parallel_checkpoint(local_rank_id,
+                                                                                   src_strategy_list, dst_strategy_list)
+    result_set = set()
+    handled_pipeline_stage = []
+    for _, layout in src_strategy_list.items():
+        for src_pipeline_stage_id in layout[6]:
+            if src_pipeline_stage_id in handled_pipeline_stage:
+                continue
+            src_rank_id_start = src_pipeline_stage_id * src_stage_device_num
+            result_set.update([src_rank_id_start + rank for rank in needed_rank_list_in_local_stage])
+            handled_pipeline_stage.append(src_pipeline_stage_id)
+    result_list = list(result_set)
+    result_list.sort(reverse=True)
+    return result_list
+
+
+def transform_checkpoint_by_rank(rank_id, checkpoint_files_map, save_checkpoint_file_name,
+                                 src_strategy_file=None, dst_strategy_file=None):
+    """
+    Transform distributed checkpoint from source sharding strategy to destination sharding strategy by rank
+    for a network. For more details about converting distributed Checkpoint, please refer to
+    `Distributed Resilience Training and
+    Inference <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/resilience_train_and_predict.html>`_.
+
+    Args:
+        rank_id (int): The rank of which distributed checkpoint needs to be obtained after conversion.
+        checkpoint_files_map (dict): The checkpoint files map whose key is the rank id and the value is
+                                     the checkpoint file name.
+        save_checkpoint_file_name (str): The file name to save the converted checkpoint.
+        src_strategy_file (str): Name of source sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'src_strategy_file' is None, it means that the source sharding strategy is
+                                 without any sharing for each parameter. Default:None.
+        dst_strategy_file (str): Name of destination sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'dst_strategy_file' is None, it means that the destination sharding strategy
+                                 is without any sharing for each parameter. Default:None.
+
+    Raises:
+        ValueError: `src_strategy_file` or `dst_strategy_file` is incorrect.
+        ValueError: item in `checkpoint_files_map` is incorrect.
+        ValueError: `save_checkpoint_file_name` is not end with ".ckpt".
+        TypeError: `checkpoint_files_map` is not a dict.
+        TypeError: `src_strategy_file` or `dst_strategy_file` is not a string.
+        TypeError: `rank_id` is not a int.
+        TypeError: `save_checkpoint_file_name` is not a string.
+
+    Examples:
+        >>> dst_device_num = 8
+        >>> for rank_id in range(dst_device_num):
+        ...     rank_list = rank_list_for_transform(rank_id, "./src_strategy.ckpt", "./dst_strategy.ckpt")
+        ...     checkpoint_files_map = {}
+        ...     for rank in rank_list:
+        ...         checkpoint_files_map[rank] = "./origin_checkpoint_rank{}/pangu{}-100_2.ckpt".format(rank)
+        ...     save_checkpoint_file_name = "./new_checkpoint_rank{}/pangu{}-100_2.ckpt".format(rank_id)
+        ...     transform_checkpoint_by_rank(rank_id, checkpoint_files_map, save_checkpoint_file_name,
+        ...                                  "./src_strategy.ckpt", "./dst_strategy.ckpt")
+
+    """
+    if not isinstance(checkpoint_files_map, dict):
+        raise TypeError("The checkpoint_files_map should be a dict.")
+    if not isinstance(rank_id, int):
+        raise TypeError("The rank_id should be a int.")
+    if not isinstance(save_checkpoint_file_name, str):
+        raise TypeError("The save_checkpoint_file_name should be a str.")
+    if save_checkpoint_file_name[-5:] != ".ckpt":
+        raise ValueError("The save_checkpoint_file_name {} should end with .ckpt".format(save_checkpoint_file_name))
+    if dst_strategy_file and os.path.dirname(dst_strategy_file) and not os.path.exists(
+            os.path.dirname(dst_strategy_file)):
+        raise ValueError("The director of dst_strategy_file: {} is not exists.".
+                         format(os.path.dirname(dst_strategy_file)))
+    for rank, local_file in checkpoint_files_map.items():
+        if not os.path.exists(local_file):
+            raise ValueError("Checkpoint file {} in rank {} not exits: ".format(local_file, rank))
+    param_total_dict = defaultdict(dict)
+    param_attr_dict = defaultdict(dict)
+    src_strategy_list, dst_strategy_list = _extract_src_dst_layout_map(rank_id, src_strategy_file, dst_strategy_file)
+    # src rank => local rank inside pipeline stage
+    src_stage_device_num = np.prod(src_strategy_list.get(list(src_strategy_list.keys())[0])[0]) if src_strategy_list \
+                                                                                                   is not None else 1
+    dst_stage_device_num = np.prod(dst_strategy_list.get(list(dst_strategy_list.keys())[0])[0]) if dst_strategy_list \
+                                                                                                   is not None else 1
+    origin_dst_strategy_list = _extract_layout_map(dst_strategy_file)
+    origin_src_strategy_list = _extract_layout_map(src_strategy_file)
+    for rank, file_name in checkpoint_files_map.items():
+        ckpt_dict = ms.load_checkpoint(file_name)
+        for param_name, param in ckpt_dict.items():
+            # cut the parameter not in the pipeline stage.
+            if _parameter_not_in_local_stage(param_name, origin_src_strategy_list, src_strategy_list) \
+                    and _parameter_not_in_local_stage(param_name, origin_dst_strategy_list, dst_strategy_list):
+                continue
+            src_rank = rank % src_stage_device_num
+            param_total_dict[param_name][src_rank] = param.data.asnumpy()
+            param_attr_dict[param_name][src_rank] = (param.requires_grad, param.layerwise_parallel)
+    local_rank_id = rank_id % dst_stage_device_num
+    transform_param_list = _transform_parallel_checkpoint(local_rank_id, param_total_dict,
+                                                          param_attr_dict, src_strategy_list, dst_strategy_list)
+    ms.save_checkpoint(transform_param_list, save_checkpoint_file_name)
+
+
+def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix, src_strategy_file=None,
+                          dst_strategy_file=None):
+    """
+    Transform distributed checkpoint from source sharding strategy to destination sharding strategy for a rank.
+    For more details about converting distributed Checkpoint, please refer to
+    `Distributed Resilience Training and
+    Inference <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/resilience_train_and_predict.html>`_.
+
+    Note:
+        The `src_checkpoints_dir` directory structure should be organized like "src_checkpoints_dir/rank_0/a.ckpt", the
+        rank number should be set to a subdirectory and the checkpoint file is stored in this subdirectory. If multiple
+        files exist in a rank directory, the last file in the lexicgraphic order would be selected.
+
+    Args:
+        src_checkpoints_dir (str): The source checkpoints directory.
+        dst_checkpoints_dir (str): The destination checkpoints directory to save the converted checkpoints.
+        ckpt_prefix (str): The destination checkpoint name prefix.
+        src_strategy_file (str): Name of source sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'src_strategy_file' is None, it means that the source sharding strategy is
+                                 without any sharing for each parameter. Default:None.
+        dst_strategy_file (str): Name of destination sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'dst_strategy_file' is None, it means that the destination sharding strategy
+                                 is without any sharing for each parameter. Default:None.
+
+    Raises:
+        ValueError: `src_strategy_file` or `dst_strategy_file` is incorrect.
+        NotADirectoryError: `src_checkpoints_dir` or `dst_checkpoints_dir` is not a directory.
+        ValueError: The checkpoint file is missing in `src_checkpoints_dir`.
+        TypeError: `src_strategy_file` or `dst_strategy_file` is not a string.
+
+    Examples:
+        >>> transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, "dst_checkpoint",
+        ...                       "./src_strategy.ckpt", "./dst_strategy.ckpt")
+
+    """
+    if not os.path.isdir(src_checkpoints_dir):
+        raise NotADirectoryError("src_checkpoints_dir {} is not a directory.".format(src_checkpoints_dir))
+    _make_dir(dst_checkpoints_dir, "path")
+    if not isinstance(ckpt_prefix, str):
+        raise TypeError("The ckpt_prefix should be a str.")
+    checkpoints_rank_dir_list = os.path.join(src_checkpoints_dir, "rank_[0-9]*")
+    all_checkpoint_files_map = {}
+    for checkpoint_dir in glob.glob(checkpoints_rank_dir_list):
+        if not os.path.isdir(checkpoint_dir):
+            ms.log.warning("{} is not a directory.".format(checkpoint_dir))
+            continue
+        rank_id_str = checkpoint_dir.split('rank_')[-1]
+        if not rank_id_str.isdigit():
+            ms.log.warning("{} is not a expected directory, the directory should end with rank_0/rank_1.....".
+                           format(checkpoint_dir))
+            continue
+        rank_id = int(rank_id_str)
+        checkpoint_file_name = os.path.join(checkpoint_dir, "*.ckpt")
+        rank_ckpts = glob.glob(checkpoint_file_name)
+        rank_ckpts.sort()
+        for checkpoint_file in rank_ckpts:
+            if not os.path.isfile(checkpoint_file):
+                ms.log.warning("{} is not a checkpoint file.".format(checkpoint_file))
+                continue
+            all_checkpoint_files_map[rank_id] = checkpoint_file
+
+    needed_rank_list_map = defaultdict(list)
+    dst_stage_device_num = _get_device_num_from_strategy(dst_strategy_file)
+    src_stage_device_num = _get_device_num_from_strategy(src_strategy_file)
+    dst_stage_num = _extract_pipeline_stage_num(dst_strategy_file)
+    dst_device_num = dst_stage_device_num * dst_stage_num
+    origin_src_strategy_list = _extract_layout_map(src_strategy_file)
+    origin_dst_strategy_list = _extract_layout_map(dst_strategy_file)
+    for rank in range(dst_device_num):
+        needed_rank_list = rank_list_for_transform(rank, src_strategy_file, dst_strategy_file)
+        for needed_rank in needed_rank_list:
+            if needed_rank not in all_checkpoint_files_map:
+                raise ValueError("The checkpoint file of rank{} is needed for converting rank{}'s checkpoint, "
+                                 "but it is missing.".format(needed_rank, rank))
+        needed_rank_list_key = "-".join([str(r) for r in needed_rank_list])
+        needed_rank_list_map[needed_rank_list_key].append(rank)
+    for needed_rank_list_key, transform_rank_list in needed_rank_list_map.items():
+        param_total_dict = defaultdict(dict)
+        param_attr_dict = defaultdict(dict)
+        needed_rank_list = needed_rank_list_key.split("-")
+        for needed_rank in needed_rank_list:
+            ckpt_dict = ms.load_checkpoint(all_checkpoint_files_map.get(int(needed_rank)))
+            for param_name, param in ckpt_dict.items():
+                src_rank = int(needed_rank) % src_stage_device_num
+                param_total_dict[param_name][src_rank] = param.data.asnumpy()
+                param_attr_dict[param_name][src_rank] = (param.requires_grad, param.layerwise_parallel)
+        for transform_rank in transform_rank_list:
+            param_total_dict_copy = copy.deepcopy(param_total_dict)
+            src_strategy_list, dst_strategy_list = _extract_src_dst_layout_map(transform_rank, src_strategy_file,
+                                                                               dst_strategy_file)
+            # cut the parameter not in the pipeline stage.
+            for param in list(param_total_dict_copy.keys()):
+                if _parameter_not_in_local_stage(param, origin_src_strategy_list, src_strategy_list) \
+                        and _parameter_not_in_local_stage(param, origin_dst_strategy_list, dst_strategy_list):
+                    param_total_dict_copy.pop(param)
+
+            local_rank_id = transform_rank % dst_stage_device_num
+            transform_param_list = _transform_parallel_checkpoint(local_rank_id, param_total_dict_copy,
+                                                                  param_attr_dict, src_strategy_list, dst_strategy_list)
+            save_checkpoint_file = "{}{}.ckpt".format(ckpt_prefix, transform_rank)
+            save_checkpoint_file_dir = os.path.join(dst_checkpoints_dir, "rank_{}".format(transform_rank))
+            if not os.path.exists(save_checkpoint_file_dir):
+                _make_dir(save_checkpoint_file_dir, "path")
+            save_checkpoint_file_name = os.path.join(save_checkpoint_file_dir, save_checkpoint_file)
+            ms.save_checkpoint(transform_param_list, save_checkpoint_file_name)
+            del param_total_dict_copy
+        del param_total_dict

mindspore/parallel/shard.py

@@ -0,0 +1,229 @@
+# Copyright 2022 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.
+# ============================================================================
+"""pynative shard"""
+
+import mindspore as ms
+from mindspore import log as logger
+from mindspore._c_expression import Shard_
+
+
+class Shard(Shard_):
+    """Shard operation"""
+
+    def __init__(self):
+        """Initialize Shard."""
+        Shard_.__init__(self, 'Shard')
+        self.shard_fn = None
+        self.fn = None
+        self.in_strategy = None
+        self.out_strategy = None
+        self.parameter_plan = None
+        self.device = None
+        self.level = None
+
+    def __call__(self, fn, in_strategy, out_strategy=None, parameter_plan=None, device="Ascend", level=0):
+        if ms.context.get_context("mode") != ms.context.PYNATIVE_MODE or \
+                ms.context.get_auto_parallel_context("parallel_mode") not in ["auto_parallel"]:
+            raise AssertionError(f"Cell shard only supports auto parallel under PyNative mode.")
+        if ms.context.get_context("device_target") not in ["Ascend", "GPU"]:
+            raise AssertionError(f"'Shard' now only supports 'Ascend' and 'GPU'")
+        if ms.context.get_auto_parallel_context("search_mode") != "sharding_propagation":
+            raise AssertionError(f"'search_mode' must be 'sharding_propagation' for 'Shard'")
+        if not isinstance(in_strategy, tuple):
+            raise TypeError(f"For 'Shard', the 'in_strategy' should be a tuple, but got {type(in_strategy).__name__}")
+        if not isinstance(out_strategy, (type(None), tuple)):
+            raise TypeError(f"For 'Shard', the 'out_strategy' should be None or tuple, "
+                            f"but got {type(out_strategy).__name__}")
+
+        if not isinstance(device, str):
+            raise TypeError(f"For 'Shard', the 'device' should be a string, "
+                            f"but got {type(device).__name__}")
+        if not isinstance(level, int):
+            raise TypeError(f"For 'Shard', the 'level' should be an integer, "
+                            f"but got {type(level).__name__}")
+
+        if ms.get_algo_parameters("fully_use_devices") is True:
+            logger.warning("After calling 'shard', the environment variable 'fully_use_devices' "
+                           "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:
+            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()
+
+        if isinstance(fn, ms.nn.Cell):
+            for param in fn.trainable_params():
+                param.is_in_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)
+            def after_shard(*args):
+                return shard_(fn, in_strategy, out_strategy, device, level)(*args)
+
+            return after_shard(*args)
+
+        self.shard_fn = shard_fn
+        self.fn = fn
+        self.in_strategy = in_strategy
+        self.out_strategy = out_strategy
+        self.device = device
+        self.level = level
+        return self.shard_fn
+
+    @staticmethod
+    def _search_parameter_by_name(param_name: str, net):
+        param_name = param_name.replace("self.", "")
+        for param in net.trainable_params():
+            if param.name == param_name:
+                return param
+        return None
+
+    @staticmethod
+    def _check_layout_is_valid(param_name, param_shape, param_strategy):
+        if len(param_strategy) != len(param_shape):
+            raise ValueError(f"For {param_name}, the length of param_strategy: {len(param_strategy)}, "
+                             f"is not equal to param_shape len: {len(param_shape)}.")
+        for i, _ in enumerate(param_strategy):
+            if param_shape[i] % param_strategy[i] != 0:
+                raise ValueError(f"For '{param_name}', the param_shape is {param_shape} and "
+                                 f"the setting param_strategy is {param_strategy}. "
+                                 f"The param_shape[{i}]: {param_shape[i]} cannot be divisible by "
+                                 f"param_strategy[{i}]: {param_strategy[i]}.")
+
+    def _set_param_layout_into_parameter(self, fn, parameter_plan):
+        """ Set param_strategy into parameter if fn is a Cell and parameter_plan is a dict."""
+        if parameter_plan is None:
+            return
+        if isinstance(parameter_plan, dict):
+            if not isinstance(fn, ms.nn.Cell):
+                raise TypeError(f"If parameter_plan is set, type of fn must be mindspore.nn.Cell, but got {type(fn)}")
+            for k in parameter_plan.keys():
+                v = parameter_plan[k]
+                if not isinstance(k, str) or not isinstance(v, tuple):
+                    raise TypeError(f"For 'Shard', the type of each key and value in 'parameter_plan' must be str and "
+                                    f"tuple, but got {type(k).__name__} and {type(v).__name__}")
+        else:
+            raise TypeError(f"For 'Shard', the 'parameter_plan' should be a dict or None, "
+                            f"but got {type(parameter_plan).__name__}")
+
+        for param_name in parameter_plan.keys():
+            param_strategy = parameter_plan[param_name]
+            param = self._search_parameter_by_name(param_name, fn)
+            if param is None:
+                logger.warning(f"{param_name} is not exist, ignored its setting.")
+                continue
+
+            self._check_layout_is_valid(param_name, param.shape, param_strategy)
+            if param.param_info.param_strategy:
+                logger.warning(f"The layout of parameter '{param_name}' "
+                               f"has been set to {param.param_info.param_strategy}, "
+                               f"current setting {param_strategy} will be ignored.")
+            param.param_info.param_strategy = param_strategy
+
+    def _is_attrs_has_been_set(self, fn, in_strategy, out_strategy, device, level):
+        return self.shard_fn is not None and self.fn == fn and self.in_strategy == in_strategy and \
+               self.out_strategy == out_strategy and self.device == device and self.level == level
+
+
+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 execution in graph mode.
+    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
+    this input/output, and None represents data_parallel,
+    which can refer to the description of :func:`mindspore.ops.Primitive.shard`.
+    The parallel strategies of remaining operators are derived from the strategy specified by the input and output.
+
+    Note:
+        You need to set the execution mode to PyNative mode,
+        set the parallel mode in `set_auto_parallel_context` to "auto_parallel"
+        and the search mode to "sharding_propagation".
+        If the input contain Parameter, its strategy should be set in `in_strategy`.
+        For more details about shard, please refer to `Functional Operator Sharding
+        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/pynative_shard_function_parallel.html>`_.
+
+    Args:
+        fn (Union[Cell, Function]): Function to be executed in parallel.
+                                    Its arguments and return value must be Tensor or Parameter.
+                                    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 or None.
+                             Tuple defines the layout of the corresponding input
+                             and None represents a data parallel strategy.
+        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
+                                            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, indicating the corresponding layout.
+                                            If the parameter name is incorrect or the corresponding parameter
+                                            has been set, the parameter setting will be ignored.
+                                            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".
+
+    Returns:
+        Function, return the function that will be executed under auto parallel process.
+
+    Raises:
+        AssertionError: If execute mode is not PYNATIVE_MODE.
+        AssertionError: If parallel mode is not "auto_parallel".
+        AssertionError: If search_mode it not "sharding_propagation".
+        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 `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.
+        TypeError: If `device` is not a str.
+        TypeError: If `level` is not an integer.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor, set_context, set_auto_parallel_context, shard, PYNATIVE_MODE
+        >>> from mindspore.communication import init
+        >>> set_context(mode=PYNATIVE_MODE)
+        >>> init()
+        >>> set_auto_parallel_context(parallel_mode="auto_parallel", search_mode="sharding_propagation",
+        ...                           device_num=2)
+        >>> def test_shard(x, y):
+        ...     return x + y
+        >>> x = Tensor(np.ones(shape=(32, 10)))
+        >>> y = Tensor(np.ones(shape=(32, 10)))
+        >>> output = shard(test_shard, in_strategy=((2, 1), (2, 1)))(x, y)
+        >>> print(output.shape)
+        (32, 10)
+    """
+    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)