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

mindspore/parallel/_parallel_serialization.py

@@ -16,6 +16,7 @@
 from __future__ import absolute_import
 
 import os
+import json
 import numpy as np
 import mindspore as ms
 from mindspore.parallel._tensor import _get_tensor_strategy, _construct_from_to_tensor_layout, \
@@ -39,7 +40,15 @@ def _convert_to_list(strategy):
             field_size = int(layout.field)
             shard_stride = int(layout.opt_weight_shard_step)
             shard_size = int(layout.opt_weight_shard_size)
-            train_map[param_name] = [dev_mat, tensor_map, param_split_shape, field_size, shard_stride, shard_size]
+            pipeline_stage = 0
+            origin_param_name = param_name
+            if "-" in param_name:
+                pipeline_stage, origin_param_name = param_name.split("-")
+            if origin_param_name not in train_map:
+                train_map[origin_param_name] = [dev_mat, tensor_map, param_split_shape, field_size, shard_stride,
+                                                shard_size, [int(pipeline_stage)]]
+            else:
+                train_map.get(origin_param_name)[6].append(int(pipeline_stage))
         except BaseException as e:
             raise ValueError(f"{e.__str__()}. Convert layout strategy to list "
                              f"failed, please make sure that strategy matches the node_strategy.proto, you can "
@@ -73,8 +82,8 @@ def _convert_to_layout(param_name, tensor_layout):
     return strategy
 
 
-def _build_searched_strategy(strategy_filename):
-    """build searched strategy"""
+def _check_strategy_file(strategy_filename):
+    """load parallel strategy file"""
     if not isinstance(strategy_filename, str):
         raise TypeError(f"For 'build_searched_strategy', the argument 'strategy_filename' should be string, "
                         f"but got {type(strategy_filename)}.")
@@ -86,12 +95,25 @@ def _build_searched_strategy(strategy_filename):
     if os.path.getsize(strategy_filename) == 0:
         raise ValueError(f"For 'build_searched_strategy', the strategy file {strategy_filename} should not "
                          f"be empty. Please check whether the 'strategy_filename' is correct.")
-    parallel_strategy_map = ms.train.node_strategy_pb2.ParallelStrategyMap()
 
+
+def _load_protobuf_strategy(strategy_filename):
+    """load strategy from protobuf file"""
+    parallel_strategy_map = ms.train.node_strategy_pb2.ParallelStrategyMap()
     with open(strategy_filename, 'rb') as f:
         pb_content = f.read()
-    parallel_strategy_map.ParseFromString(pb_content)
+    try:
+        parallel_strategy_map.ParseFromString(pb_content)
+    except BaseException as e:
+        raise TypeError("The strategy file type should be one of json or protobuf. "
+                        "When the file name extension is not '.json', "
+                        "the file is considered as a protobuf file.") from e
+    return parallel_strategy_map
 
+
+def _build_protobuf_strategy(strategy_filename):
+    """build strategy from protobuf file"""
+    parallel_strategy_map = _load_protobuf_strategy(strategy_filename)
     layout_items = parallel_strategy_map.parallel_layout_item
     if not layout_items:
         raise ValueError(f"For 'build_searched_strategy', the strategy file {strategy_filename} has no sliced "
@@ -102,10 +124,143 @@ def _build_searched_strategy(strategy_filename):
         parameter_name = layout_item.param_name
         layout = layout_item.parallel_layouts
         strategy[parameter_name] = layout
+    return strategy
+
 
+def _build_json_strategy(strategy_filename):
+    """build strategy from json file"""
+    with open(strategy_filename, 'r') as f:
+        json_content = json.load(f)
+    layout_items = json_content.get("parallel_layout_item")
+    strategy = {}
+    for parameter_name, layout_item in layout_items.items():
+        layout = ms.train.node_strategy_pb2.ParallelLayouts()
+        layout.field = layout_item.get("field")
+        layout.opt_weight_shard_size = layout_item.get("opt_weight_shard_size")
+        layout.opt_weight_shard_step = layout_item.get("opt_weight_shard_step")
+        dev_matrix = layout.dev_matrix.add()
+        for item in layout_item.get("dev_matrix"):
+            dev_matrix.dim.append(item)
+        tensor_map = layout.tensor_map.add()
+        for item in layout_item.get("tensor_map"):
+            tensor_map.dim.append(item)
+        param_split_shape = layout.param_split_shape.add()
+        if "param_split_shape" in layout_item:
+            for item in layout_item.get("param_split_shape"):
+                param_split_shape.dim.append(item)
+        indices_offset = layout.indices_offset.add()
+        if "indices_offset" in layout_item:
+            for item in layout_item.get("indices_offset"):
+                indices_offset.dim.append(item)
+        strategy[parameter_name] = layout
     return strategy
 
 
+def _build_searched_strategy(strategy_filename):
+    """build searched strategy"""
+    _check_strategy_file(strategy_filename)
+    if strategy_filename[-5:] != ".json":
+        return _build_protobuf_strategy(strategy_filename)
+    return _build_json_strategy(strategy_filename)
+
+
+def _merge_protobuf_strategy(src_strategy_files, dst_strategy_file):
+    """merge protobuf strategy"""
+    dst_parallel_strategy_map = ms.train.node_strategy_pb2.ParallelStrategyMap()
+    merged_stage = []
+    for src_strategy_file in src_strategy_files:
+        src_parallel_strategy_map = _load_protobuf_strategy(src_strategy_file)
+        strategy_items = src_parallel_strategy_map.parallel_strategy_item
+        layout_items = src_parallel_strategy_map.parallel_layout_item
+        if not strategy_items or not layout_items:
+            raise ValueError("The strategy file {} is empty".format(src_strategy_file))
+        pipeline_stage = strategy_items[0].parallel_strategys.stage
+        if pipeline_stage in merged_stage:
+            continue
+        for layout_item in layout_items:
+            layout_item.param_name = "-".join([str(pipeline_stage), layout_item.param_name])
+        dst_parallel_strategy_map.parallel_strategy_item.extend(strategy_items)
+        dst_parallel_strategy_map.parallel_layout_item.extend(layout_items)
+        merged_stage.append(pipeline_stage)
+    dst_parallel_strategy_map.current_stage = 1
+    with open(dst_strategy_file, "wb") as f:
+        f.write(dst_parallel_strategy_map.SerializeToString())
+
+
+def _merge_json_strategy(src_strategy_files, dst_strategy_file):
+    """merge protobuf strategy"""
+    dst_parallel_strategy_map = {"current_stage": 1, "parallel_strategy_item": {}, "parallel_layout_item": {}}
+    merged_stage = []
+    for src_strategy_file in src_strategy_files:
+        with open(src_strategy_file, 'r') as f:
+            json_content = json.load(f)
+        layout_items = json_content.get("parallel_layout_item")
+        strategy_items = json_content.get("parallel_strategy_item")
+        if not strategy_items or not layout_items:
+            raise ValueError("The strategy file {} is empty".format(src_strategy_file))
+        pipeline_stage = strategy_items.get(list(strategy_items.keys())[0]).get('stage')
+        if pipeline_stage in merged_stage:
+            continue
+        for param_name, layout_item in layout_items.items():
+            new_layout_item = {}
+            new_param_name = "-".join([str(pipeline_stage), param_name])
+            new_layout_item[new_param_name] = layout_item
+            dst_parallel_strategy_map.get("parallel_layout_item").update(new_layout_item)
+        dst_parallel_strategy_map.get("parallel_strategy_item").update(strategy_items)
+        merged_stage.append(pipeline_stage)
+    with open(dst_strategy_file, "w") as f:
+        json.dump(dst_parallel_strategy_map, f)
+
+
+def _parameter_not_in_local_stage(param_name, origin_strategy_list, strategy_list):
+    """parameter whether in the local stage"""
+    if origin_strategy_list is None or strategy_list is None:
+        return True
+    return param_name in origin_strategy_list and param_name not in strategy_list
+
+
+def _extract_layout_map(strategy_file):
+    """Extract layout map"""
+    layout_map = None
+    if strategy_file is not None:
+        src_strategy = _build_searched_strategy(strategy_file)
+        layout_map = _convert_to_list(src_strategy)
+    return layout_map
+
+
+def _extract_pipeline_stage_num(strategy_file):
+    """extract pipeline stage num"""
+    pipeline_stage_num = 1
+    if strategy_file is not None:
+        src_strategy = _build_searched_strategy(strategy_file)
+        layout_map = _convert_to_list(src_strategy)
+        pipeline_stage_set = set()
+        for _, layout in layout_map.items():
+            pipeline_stage_set.update(layout[6])
+        pipeline_stage_num = len(pipeline_stage_set)
+        if list(pipeline_stage_set) != list(range(pipeline_stage_num)):
+            raise ValueError("The strategy file for pipeline parallel dose not contains all stages.")
+    return pipeline_stage_num
+
+
+def _extract_src_dst_layout_map(rank_id, src_strategy_file=None, dst_strategy_file=None):
+    """Extract strategy list"""
+    src_layout_map = _extract_layout_map(src_strategy_file)
+    dst_layout_map = _extract_layout_map(dst_strategy_file)
+    if dst_layout_map is None:
+        return src_layout_map, dst_layout_map
+    dst_stage_device_num = np.prod(dst_layout_map.get(list(dst_layout_map.keys())[0])[0])
+    dst_stage_id = rank_id // dst_stage_device_num
+    # cut the source and destination layout, remain the parameter in the dst_stage
+    for param_name in list(dst_layout_map.keys()):
+        if dst_stage_id in dst_layout_map.get(param_name)[6]:
+            continue
+        dst_layout_map.pop(param_name)
+        if src_layout_map is not None and param_name in src_layout_map:
+            src_layout_map.pop(param_name)
+    return src_layout_map, dst_layout_map
+
+
 def _restore_group_info_list(group_info_file_name):
     """restore group info"""
     parallel_group_map = ms.train.node_strategy_pb2.ParallelGroupMap()
@@ -122,6 +277,7 @@ def _restore_group_info_list(group_info_file_name):
 
 
 def _get_device_num_from_strategy(strategy_file=None):
+    """Get device num from strategy file"""
     if strategy_file is None:
         return 1
     src_strategy = _build_searched_strategy(strategy_file)
@@ -130,23 +286,14 @@ def _get_device_num_from_strategy(strategy_file=None):
     return np.prod(device_mat)
 
 
-def _rank_list_for_transform_parallel_checkpoint(rank_id, src_strategy_file=None, dst_strategy_file=None):
+def _rank_list_for_transform_parallel_checkpoint(rank_id, src_strategy_list, dst_strategy_list):
     """
     Get the needed rank list for transform model parallel dim of checkpoint.
     """
-    if src_strategy_file is None:
-        return [rank_id]
-    src_strategy = _build_searched_strategy(src_strategy_file)
-    src_strategy_list = _convert_to_list(src_strategy)
-    if not src_strategy_list:
-        raise ValueError("The src_strategy_file is empty.")
-    if dst_strategy_file is not None:
-        dst_strategy = _build_searched_strategy(dst_strategy_file)
-        dst_strategy_list = _convert_to_list(dst_strategy)
     result_list = set()
     handled_layout = []
     for param_name, _ in src_strategy_list.items():
-        if dst_strategy_file is not None and param_name not in dst_strategy_list:
+        if dst_strategy_list is not None and param_name not in dst_strategy_list:
             continue
         from_dev_matrix, from_tensor_map, from_opt_shard_step, from_opt_shard_size = _extract_layout_item(
             src_strategy_list.get(param_name))
@@ -156,7 +303,7 @@ def _rank_list_for_transform_parallel_checkpoint(rank_id, src_strategy_file=None
         to_tensor_map = [-1] * len(fake_tensor_shape)
         to_opt_shard_step = 0
         to_opt_shard_size = 0
-        if dst_strategy_file is not None:
+        if dst_strategy_list is not None:
             to_dev_matrix, to_tensor_map, to_opt_shard_step, to_opt_shard_size = _extract_layout_item(
                 dst_strategy_list.get(param_name))
         handled_key = (from_dev_matrix, from_tensor_map, from_opt_shard_step, from_opt_shard_size,
@@ -188,18 +335,10 @@ def _rank_list_for_transform_parallel_checkpoint(rank_id, src_strategy_file=None
     return list(result_list)
 
 
-def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, src_strategy_file=None,
-                                   dst_strategy_file=None):
+def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, src_strategy_list, dst_strategy_list):
     """
     Transform model parallel dimension for distributed checkpoint files.
     """
-    device_num = rank_id + 1
-    if src_strategy_file is not None:
-        src_strategy = _build_searched_strategy(src_strategy_file)
-        src_strategy_list = _convert_to_list(src_strategy)
-    if dst_strategy_file is not None:
-        dst_strategy = _build_searched_strategy(dst_strategy_file)
-        dst_strategy_list = _convert_to_list(dst_strategy)
     transform_param_dict = {}
     for param_name, _ in param_total_dict.items():
         tensor_shape = list(param_total_dict[param_name].values())[0].shape
@@ -207,7 +346,7 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
         from_tensor_map = [-1] * len(tensor_shape)
         from_opt_shard_step = 0
         from_opt_shard_size = 0
-        if src_strategy_file is not None:
+        if src_strategy_list is not None:
             if param_name not in src_strategy_list:
                 ms.log.warning("The parameter {} is not in src_strategy.".format(param_name))
                 continue
@@ -217,7 +356,7 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
         to_tensor_map_origin = [-1] * len(tensor_shape)
         to_opt_shard_step = 0
         to_opt_shard_size = 0
-        if dst_strategy_file is not None:
+        if dst_strategy_list is not None:
             if param_name not in dst_strategy_list:
                 ms.log.warning("The parameter {} is not in dst_strategy.".format(param_name))
                 continue

mindspore/parallel/_ps_context.py

@@ -40,7 +40,7 @@ def _need_reset_device_target_for_ps(target):
     For Ascend backend, the card can't be occupied by multiple processes in distributed traning,
     so we need to reset the device target for some roles.
     '''
-    is_server = (_get_ps_context("ms_role") in ["MS_PSERVER", "MS_SERVER", "MS_SCHED"])
+    is_server = (os.getenv('MS_ROLE') in ["MS_PSERVER", "MS_SERVER", "MS_SCHED"])
     return is_server and target == "Ascend"
 
 
@@ -184,10 +184,6 @@ def _reinsert_hash_table_size(new_name, cur_name, cache_vocab_size, embedding_si
     ps_context().reinsert_hash_table_size(new_name, cur_name, cache_vocab_size, embedding_size)
 
 
-def _insert_weight_init_info(name, global_seed, op_seed):
-    ps_context().insert_weight_init_info(name, global_seed, op_seed)
-
-
 def _insert_accumu_init_info(name, init_val):
     ps_context().insert_accumu_init_info(name, init_val)
 
@@ -210,6 +206,14 @@ def _cache_enable():
     return ps_context().cache_enable()
 
 
+def _set_cache_size(cache_size):
+    ps_context().set_cache_size(cache_size)
+
+
+def _set_sparse_format(sparse_format):
+    ps_context().set_sparse_format(sparse_format)
+
+
 def _set_rank_id(rank_id):
     ps_context().set_rank_id(rank_id)
 

mindspore/parallel/_recovery_context.py

@@ -14,7 +14,7 @@
 # ============================================================================
 """Context for recovery"""
 
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore._c_expression import RecoveryContext
 
 RECOVERY_CONTEXT = None

mindspore/parallel/_tensor.py

@@ -15,7 +15,9 @@
 """load tensor and combine tensor"""
 from __future__ import division
 from __future__ import absolute_import
+
 import numpy as np
+
 from mindspore.common.tensor import Tensor
 from mindspore.communication.management import get_rank, get_group_size
 from mindspore._c_expression import TensorTransform
@@ -173,20 +175,26 @@ def _chunk_tensor_by_strategy(np_tensor, strategy):
     return _chunk_tensor(np_tensor, strategy, len(strategy))
 
 
-def _get_slice_index(dev_mat, tensor_map):
+def _get_slice_index(dev_mat, tensor_map, opt_shard_group):
     """
     Get the slice index for current slice.
 
     Args:
         dev_mat (list): The device matrix of devices.
         tensor_map (list): The split strategy of tensor.
+        opt_shard_group(string): The group of optimizer shard
 
     Returns:
         Integer, the slice index for slice on this device.
     """
     rank = get_rank()
+    dev_num = get_group_size()
     tensor_strategy = _get_tensor_strategy(dev_mat, tensor_map)
     tensor_slice_index = _get_tensor_slice_index(dev_mat, tensor_strategy, tensor_map, rank)
+    if opt_shard_group:
+        tensor_slice_index += dev_num
+        opt_rank = get_rank(opt_shard_group)
+        tensor_slice_index += opt_rank
     return tensor_slice_index
 
 

mindspore/{nn/transformer → parallel/_transformer}/__init__.py

@@ -15,17 +15,17 @@
 """
 NOTE:
     Transformer Networks.
-    This is an experimental interface that is subject to change or deletion.
+    These are experimental APIs that are subject to change or deletion.
 """
 from __future__ import absolute_import
 
-from mindspore.nn.transformer.transformer import AttentionMask, VocabEmbedding, MultiHeadAttention, \
+from mindspore.parallel._transformer.transformer import AttentionMask, VocabEmbedding, MultiHeadAttention, \
     FeedForward, TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer, \
     Transformer, TransformerOpParallelConfig, EmbeddingOpParallelConfig, TransformerRecomputeConfig
-from mindspore.nn.transformer.moe import MoEConfig
-from mindspore.nn.transformer.layers import FixedSparseAttention
-from mindspore.nn.transformer.loss import CrossEntropyLoss
-from mindspore.nn.transformer.op_parallel_config import OpParallelConfig
+from mindspore.parallel._transformer.moe import MoEConfig
+from mindspore.parallel._transformer.layers import FixedSparseAttention
+from mindspore.parallel._transformer.loss import CrossEntropyLoss
+from mindspore.parallel._transformer.op_parallel_config import OpParallelConfig
 
 __all__ = []
 __all__.extend(transformer.__all__)

mindspore/{nn/transformer → parallel/_transformer}/layers.py

@@ -33,11 +33,12 @@ from mindspore._extends import cell_attr_register
 from mindspore.nn.cell import Cell
 from mindspore.nn.layer.activation import get_activation
 from mindspore.ops import functional as F
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.ops.primitive import constexpr
 from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation
 from mindspore.context import ParallelMode
-from mindspore.nn.transformer.op_parallel_config import default_dpmp_config, OpParallelConfig
+from mindspore.parallel._transformer.op_parallel_config import default_dpmp_config, OpParallelConfig, MoEParallelConfig
+from mindspore import log as logger
 
 __all__ = [
     "FixedSparseAttention"
@@ -154,6 +155,29 @@ class _LayerInputCheck:
                              f"but got {input_shape[dim]}")
         return True
 
+    @staticmethod
+    def check_shape_equal_without_batch(input_shape, param_name, func_name, target_shape):
+        """
+        Check the input shape's is equal to the expected shape, the value on 0-th is viewed as batch, and the
+        batch size will not be checked.
+        """
+        length, hidden = target_shape
+        if isinstance(input_shape, tuple):
+            input_shape = list(input_shape)
+        _LayerInputCheck.check_shape_length(input_shape, param_name, func_name,
+                                            [len(target_shape), len(target_shape) + 1])
+        if input_shape[-1] != hidden:
+            raise ValueError(f"For {func_name}, the last dimension of {param_name} shape must be {hidden},"
+                             f"but got the last dimension {input_shape[-1]} in {input_shape}.")
+        if input_shape[0] == 0:
+            raise ValueError(f"For {func_name}, the first dimension of {param_name} shape greater than 0,"
+                             f"but got the first dimension {input_shape[0]} in {input_shape}.")
+        if len(input_shape) == 2 and input_shape[0] % length != 0:
+            raise ValueError(f"For {func_name}, the first dimension of {param_name} shape should be divisible "
+                             f"by {length}, "
+                             f"but got the first dimension {input_shape[0]} in {input_shape}.")
+        return True
+
 
 @constexpr
 def _check_past_none_input_none(use_past, param_name, func_name, default_value, is_tensor, is_default):
@@ -175,23 +199,6 @@ def _check_input_dtype(input_dtype, param_name, allow_dtypes, cls_name):
     Validator.check_type_name(param_name, input_dtype, allow_dtypes, cls_name)
 
 
-@constexpr
-def _check_input_shape(input_shape, param_name, func_name, target_len):
-    # check the input length
-    _LayerInputCheck.check_shape_length(input_shape, param_name, func_name, target_len)
-
-
-@constexpr
-def _check_shape_equal(input_shape, param_name, func_name, target_shape):
-    # check the input length
-    _LayerInputCheck.check_shape_equal(input_shape, param_name, func_name, target_shape)
-
-
-@constexpr
-def _check_input_shape_value(input_shape, dim, param_name, cls_name, target_value):
-    _LayerInputCheck.check_shape_value_on_axis(input_shape, dim, param_name, cls_name, target_value)
-
-
 class _Dropout(nn.Cell):
     r"""
         A Dropout Implements with P.DropoutGenMask and  P.DropoutDoMask for parallel training.
@@ -361,7 +368,7 @@ class _Linear(Cell):
             same as `x`. The values of str refer to the function `initializer`. Default: 'zeros'.
         has_bias (bool): Specifies whether the layer uses a bias vector. Default: True.
         activation (str): activate function applied to the output of the fully connected layer,
-            eg. 'ReLU'.Default: None.
+            eg. 'ReLU'. Default: None.
         expert_num (int): The number of experts used in this Linear. Here, for the case expert_num > 1, BatchMatMul is
             used and the first dimension in BatchMatMul indicate expert_num. Default: 1.
         outer_batch (int): The replication number of experts. The replication is effective only when MoE is applied.
@@ -392,7 +399,6 @@ class _Linear(Cell):
     @_args_type_validator_check(in_channels=Validator.check_positive_int,
                                 out_channels=Validator.check_positive_int,
                                 has_bias=Validator.check_bool,
-                                activation=_valid_type_checks([type(None), str], "Linear"),
                                 transpose_b=Validator.check_bool,
                                 expert_num=Validator.check_positive_int,
                                 outer_batch=Validator.check_positive_int,
@@ -416,7 +422,9 @@ class _Linear(Cell):
         super(_Linear, self).__init__()
         self.in_channels = in_channels
         self.out_channels = out_channels
-        if isinstance(weight_init, Tensor) and (weight_init.ndim != 2 or weight_init.shape[0] != out_channels or \
+        if not (isinstance(activation, str) or activation is None or issubclass(activation, nn.Cell)):
+            raise TypeError(f"For Linear cell, the activation should str type or nn.Cell type, but got {activation}.")
+        if isinstance(weight_init, Tensor) and (weight_init.ndim != 2 or weight_init.shape[0] != out_channels or
                                                 weight_init.shape[1] != in_channels):
             raise ValueError("The shape of parameter 'weight_init' is error, please check shape of 'weight_init'.")
         weight_shape = [out_channels, in_channels] if transpose_b else [in_channels, out_channels]
@@ -449,7 +457,10 @@ class _Linear(Cell):
             self.bias.parallel_optimizer = False
             self.bias_add = P.Add()
         self.act_name = activation
-        self.activation = get_activation(activation) if isinstance(activation, str) else activation
+        if callable(activation):
+            self.activation = activation()
+        else:
+            self.activation = get_activation(activation) if isinstance(activation, str) else activation
         self.activation_flag = self.activation is not None
         self.dtype = compute_dtype
         self.cast = P.Cast()
@@ -491,7 +502,7 @@ class _Linear(Cell):
         self.matmul.shard(strategy_matmul)
         if self.has_bias:
             self.bias_add.shard(strategy_bias)
-        if self.activation_flag:
+        if self.activation_flag and isinstance(self.act_name, str):
             # some operations has many primitives, need to manually set the shard
             if self.act_name.lower() == "leakyrelu":
                 self.activation.select_op.shard((strategy_activation[0], strategy_activation[0]))
@@ -506,7 +517,26 @@ class _Linear(Cell):
                                  "or auto parallel mode.")
             else:
                 getattr(self.activation, self.act_name).shard(strategy_activation)
-
+        elif self.activation_flag and isinstance(self.activation, Cell):
+            if hasattr(self.activation, 'activation_shard') and strategy_activation:
+                shard_tuple = strategy_activation[0]
+                if len(shard_tuple) == 2:
+                    parallel_config = OpParallelConfig(data_parallel=shard_tuple[0],
+                                                       model_parallel=shard_tuple[1])
+                elif len(shard_tuple) == 4:
+                    parallel_config = MoEParallelConfig(data_parallel=shard_tuple[0],
+                                                        expert_parallel=shard_tuple[1],
+                                                        model_parallel=shard_tuple[2])
+                else:
+                    raise ValueError("The user-defined activation function currently only supports the case where the "
+                                     "input policy is 2 or 4, so that relevant policies can be extracted from it."
+                                     "To avoid this error, you need to add the function of extracting "
+                                     "'ParallelConfig' or 'OpParallelConfig' for the incoming strategy_activation ")
+                self.activation.activation_shard(parallel_config)
+            else:
+                logger.warning(f"The user passed the custom defined activation function {self.activation_flag}. "
+                               f"If the user want to enable shard for the activation cell, "
+                               f"the user should set the shard for each primitives in the cell.")
         return self
 
 
@@ -540,13 +570,13 @@ class FixedSparseAttention(nn.Cell):
             default args.
 
     Inputs:
-        - **q** (Tensor) - Tensor query (:class:`mstype.fp16` [batch_size, seq_length, hidden_size]): Sequence of
+        - **q** (Tensor) - Tensor query ( `mstype.fp16` [batch_size, seq_length, hidden_size]): Sequence of
           queries to query the context.
-        - **k** (Tensor) - Tensor key (:class:`mstype.fp16` [batch_size, seq_length, hidden_size]): Sequence of
+        - **k** (Tensor) - Tensor key ( `mstype.fp16` [batch_size, seq_length, hidden_size]): Sequence of
           queries to query the context.
-        - **v** (Tensor) - Tensor value (:class:`mstype.fp16` [batch size, sequence length, Embedding Size]):
+        - **v** (Tensor) - Tensor value ( `mstype.fp16` [batch size, sequence length, Embedding Size]):
           Sequence of queries to query the context.
-        - **attention_mask** (Tensor) - Float Tensor the mask of (:class:`mstype.fp32`, :class:`mstype.fp16`
+        - **attention_mask** (Tensor) - Float Tensor the mask of ( `mstype.fp32`,  `mstype.fp16`
           [batch_size, seq_length, seq_length]): Lower triangular matrix to pass masked information.
 
     Outputs:
@@ -654,17 +684,9 @@ class FixedSparseAttention(nn.Cell):
         self.slice1 = P.StridedSlice().shard(((dp, 1, 1),))
 
     def construct(self, q, k, v, attention_mask):
-        _check_shape_equal(F.shape(q), "q", self.cls_name,
-                           [self.batch_size, self.seq_length, self.hidden_size])
         _check_input_dtype(F.dtype(q), "q", [mstype.float16], self.cls_name)
-        _check_shape_equal(F.shape(k), "k", self.cls_name,
-                           [self.batch_size, self.seq_length, self.hidden_size])
         _check_input_dtype(F.dtype(k), "k", [mstype.float16], self.cls_name)
-        _check_shape_equal(F.shape(v), "v", self.cls_name,
-                           [self.batch_size, self.seq_length, self.hidden_size])
         _check_input_dtype(F.dtype(v), "v", [mstype.float16], self.cls_name)
-        _check_shape_equal(F.shape(attention_mask), "attention_mask", self.cls_name,
-                           [self.batch_size, self.seq_length, self.seq_length])
         _check_input_dtype(F.dtype(attention_mask), "attention_mask", [mstype.float32, mstype.float16], self.cls_name)
 
         q, k, v = self._transpose_inputs(q, k, v)

mindspore/{nn/transformer → parallel/_transformer}/loss.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 # ============================================================================
 """
-Parallel Loss for the Parallel Training
-This is an experimental interface that is subject to change or deletion.
+Parallel Loss for the Parallel Training.
+These are experimental APIs that are subject to change or deletion.
 """
 from __future__ import absolute_import
 
@@ -30,8 +30,8 @@ from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_device_num, _get_pipeline_stages
 from mindspore.log import _LogActionOnce
 from mindspore import log as logger
-from mindspore.nn.transformer.layers import _check_input_dtype, _check_input_shape
-from mindspore.nn.transformer.op_parallel_config import default_dpmp_config, OpParallelConfig
+from mindspore.parallel._transformer.layers import _check_input_dtype
+from mindspore.parallel._transformer.op_parallel_config import default_dpmp_config, OpParallelConfig
 
 __all__ = ["CrossEntropyLoss"]
 
@@ -247,7 +247,4 @@ class CrossEntropyLoss(Cell):
         _check_input_dtype(F.dtype(logits), "logits", [mstype.float32, mstype.float16], self.cls_name)
         _check_input_dtype(F.dtype(label), "label", [mstype.int32], self.cls_name)
         _check_input_dtype(F.dtype(input_mask), "input_mask", [mstype.float32], self.cls_name)
-        _check_input_shape(F.shape(logits), "logits", self.cls_name, 2)
-        _check_input_shape(F.shape(label), "label", self.cls_name, 1)
-        _check_input_shape(F.shape(input_mask), "input_mask", self.cls_name, 1)
         return True