mindspore 2.0.0a0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/parallel/_offload_context.py

@@ -0,0 +1,185 @@
+# Copyright 2023 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.
+# ============================================================================
+"""Context of offload"""
+from __future__ import absolute_import
+from __future__ import division
+
+import threading
+
+from mindspore._c_expression import OffloadContext
+from mindspore._checkparam import args_type_check
+from mindspore import _checkparam as Validator
+
+
+class _OffloadConfig:
+    """
+    The key of the Offload Config.
+    """
+    ENABLE_OFFLOAD = "enable_offload"
+    OFFLOAD_PARAM = "offload_param"
+    OFFLOAD_PATH = "offload_path"
+    OFFLOAD_CHECKPOINT = "offload_checkpoint"
+    OFFLOAD_DDR_SIZE = "offload_ddr_size"
+    OFFLOAD_DISK_SIZE = "offload_disk_size"
+    ENABLE_AIO = "enable_aio"
+    AIO_BLOCK_SIZE = "aio_block_size"
+    AIO_QUEUE_DEPTH = "aio_queue_depth"
+    ENABLE_PINNED_MEM = "enable_pinned_mem"
+
+
+class _OffloadContext:
+    """
+    _OffloadContext is the configuration for offload.
+
+    Note:
+        Create a context through instantiating Context object is not recommended.
+        Should use offload_context() to get the context since Context is singleton.
+    """
+    _instance = None
+    _instance_lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance_lock.acquire()
+            cls._instance = object.__new__(cls)
+            cls._instance_lock.release()
+        return cls._instance
+
+    def __init__(self):
+        self._context_handle = OffloadContext.get_instance()
+
+    def check_context_handle(self):
+        """
+        Check context handle.
+
+        Raises:
+            ValueError: If the context handle is none.
+        """
+        if self._context_handle is None:
+            raise ValueError("Context handle is none in context!!!")
+
+    def set_offload_config(self, offload_config):
+        """Set offfload context"""
+        self.check_context_handle()
+        enable_offload = _OffloadConfig.ENABLE_OFFLOAD
+        offload_param = _OffloadConfig.OFFLOAD_PARAM
+        offload_path = _OffloadConfig.OFFLOAD_PATH
+        offload_checkpoint = _OffloadConfig.OFFLOAD_CHECKPOINT
+        offload_ddr_size = _OffloadConfig.OFFLOAD_DDR_SIZE
+        offload_disk_size = _OffloadConfig.OFFLOAD_DISK_SIZE
+        enable_aio = _OffloadConfig.ENABLE_AIO
+        aio_block_size = _OffloadConfig.AIO_BLOCK_SIZE
+        aio_queue_depth = _OffloadConfig.AIO_QUEUE_DEPTH
+        enable_pinned_mem = _OffloadConfig.ENABLE_PINNED_MEM
+
+        for config_name in offload_config:
+            unknown_config = []
+            if config_name not in [enable_offload, offload_param, offload_path, offload_checkpoint,
+                                   offload_ddr_size, offload_disk_size, enable_aio, aio_block_size,
+                                   aio_queue_depth, enable_pinned_mem]:
+                unknown_config.append(config_name)
+
+            if unknown_config:
+                raise ValueError("Unknown config: {}".format(unknown_config))
+
+        if enable_offload in offload_config:
+            Validator.check_bool(
+                offload_config[enable_offload], enable_offload, enable_offload)
+            self._context_handle.set_enable_offload(
+                offload_config[enable_offload])
+
+        if offload_param in offload_config:
+            Validator.check_string(
+                offload_config[offload_param].lower(), ["cpu", "disk"])
+            self._context_handle.set_offload_param(
+                offload_config[offload_param].lower())
+
+        if offload_path in offload_config:
+            if not isinstance(offload_config[offload_path], str):
+                raise TypeError("For 'set_offload_path', "
+                                "the argument 'offload_path' must be str, but got the type : {}."
+                                .format(type(offload_config[offload_path])))
+            self._context_handle.set_offload_path(
+                offload_config[offload_path])
+        if offload_checkpoint in offload_config:
+            Validator.check_string(
+                offload_config[offload_checkpoint].lower(), ["cpu", "disk"])
+            self._context_handle.set_offload_checkpoint(
+                offload_config[offload_checkpoint].lower())
+
+        if offload_ddr_size in offload_config:
+            Validator.check_positive_int(offload_config[offload_ddr_size])
+            self._context_handle.set_offload_ddr_size(
+                offload_config[offload_ddr_size])
+
+        if offload_disk_size in offload_config:
+            Validator.check_positive_int(offload_config[offload_disk_size])
+            self._context_handle.set_offload_disk_size(
+                offload_config[offload_disk_size])
+        if enable_aio in offload_config:
+            Validator.check_bool(
+                offload_config[enable_aio], enable_aio, enable_aio)
+            self._context_handle.set_enable_aio(
+                offload_config[enable_aio])
+        if aio_block_size in offload_config:
+            Validator.check_positive_int(offload_config[aio_block_size])
+            self._context_handle.set_aio_block_size(
+                offload_config[aio_block_size])
+        if aio_queue_depth in offload_config:
+            Validator.check_positive_int(offload_config[aio_queue_depth])
+            self._context_handle.set_aio_queue_depth(
+                offload_config[aio_queue_depth])
+        if enable_pinned_mem in offload_config:
+            Validator.check_bool(
+                offload_config[enable_pinned_mem], enable_pinned_mem, enable_pinned_mem)
+            self._context_handle.set_enable_pinned_mem(
+                offload_config[enable_pinned_mem])
+
+    def offload_config(self):
+        """Get config of offload"""
+        self.check_context_handle()
+        offload_config = {
+            _OffloadConfig.ENABLE_OFFLOAD: self._context_handle.enable_offload(),
+            _OffloadConfig.OFFLOAD_PARAM: self._context_handle.offload_param(),
+            _OffloadConfig.OFFLOAD_PATH: self._context_handle.offload_path(),
+            _OffloadConfig.OFFLOAD_CHECKPOINT: self._context_handle.offload_checkpoint(),
+            _OffloadConfig.OFFLOAD_DDR_SIZE: self._context_handle.offload_ddr_size(),
+            _OffloadConfig.OFFLOAD_DISK_SIZE: self._context_handle.offload_disk_size(),
+            _OffloadConfig.ENABLE_AIO: self._context_handle.enable_aio(),
+            _OffloadConfig.AIO_BLOCK_SIZE: self._context_handle.aio_block_size(),
+            _OffloadConfig.AIO_QUEUE_DEPTH: self._context_handle.aio_queue_depth(),
+            _OffloadConfig.ENABLE_PINNED_MEM: self._context_handle.enable_pinned_mem()
+        }
+        return offload_config
+
+
+_OFFLOAD_CONTEXT = None
+
+
+def offload_context():
+    """Get offload_context. if it is not created, create a new one."""
+    global _OFFLOAD_CONTEXT
+    if _OFFLOAD_CONTEXT is None:
+        _OFFLOAD_CONTEXT = _OffloadContext()
+    return _OFFLOAD_CONTEXT
+
+
+@args_type_check(offload_config=dict)
+def _set_offload_context(offload_config):
+    offload_context().set_offload_config(offload_config)
+
+
+def _get_offload_context():
+    return offload_context().offload_config()

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, \
@@ -81,7 +82,7 @@ def _convert_to_layout(param_name, tensor_layout):
     return strategy
 
 
-def _load_strategy_file(strategy_filename):
+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, "
@@ -94,18 +95,25 @@ def _load_strategy_file(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_searched_strategy(strategy_filename):
-    """build searched strategy"""
-    parallel_strategy_map = _load_strategy_file(strategy_filename)
-
+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 "
@@ -116,10 +124,94 @@ 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:

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

@@ -175,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,11 @@ 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, MoEParallelConfig
+from mindspore.parallel._transformer.op_parallel_config import default_dpmp_config, OpParallelConfig, MoEParallelConfig
 from mindspore import log as logger
 
 __all__ = [
@@ -161,7 +161,6 @@ class _LayerInputCheck:
         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.
         """
-        target_shape = target_shape
         length, hidden = target_shape
         if isinstance(input_shape, tuple):
             input_shape = list(input_shape)
@@ -200,28 +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)
-
-
-@constexpr
-def _check_shape_equal_without_batch(input_shape, param_name, func_name, target_shape):
-    _LayerInputCheck.check_shape_equal_without_batch(input_shape, param_name, func_name, target_shape)
-
-
 class _Dropout(nn.Cell):
     r"""
         A Dropout Implements with P.DropoutGenMask and  P.DropoutDoMask for parallel training.
@@ -593,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:
@@ -707,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

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

@@ -13,26 +13,27 @@
 # limitations under the License.
 # ============================================================================
 """
-Note: Mixture of Expert (MoE) structure. This is an experimental interface that is subject to change or deletion.
+Note:
+Mixture of Expert (MoE) structure.
+These are experimental APIs that are subject to change or deletion.
 """
 from __future__ import absolute_import
 from __future__ import division
 
-import math
 import numpy as np
 
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 import mindspore.communication.management as D
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
-from mindspore.ops.primitive import constexpr
+from mindspore.ops.primitive import _primexpr
 from mindspore.nn.cell import Cell
 from mindspore.nn.layer import Dense
 from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation
-from mindspore.nn.transformer.op_parallel_config import default_moeparallel_config
+from mindspore.parallel._transformer.op_parallel_config import default_moeparallel_config
 
 __all__ = [
     "MoEConfig"]
@@ -132,9 +133,11 @@ def _check_moe_config(moe_config=None, parallel_config=None):
                          f"should be less than device_num: {device_num}.")
 
 
-@constexpr
+@_primexpr
 def calculate_expert_capacity(k, tokens_per_group, capacity_factor, expert_dim):
-    return math.ceil(k * tokens_per_group * capacity_factor / expert_dim)
+    res = k * tokens_per_group * capacity_factor / expert_dim
+    res_int = int(res)
+    return res_int if res < 0 or res == res_int else res_int + 1
 
 
 class MoE(Cell):
@@ -187,7 +190,7 @@ class MoE(Cell):
             self.group_wise_a2a = moe_config.group_wise_a2a
             if not (self.mp > 1 and self.dp == self.ep):
                 self.group_wise_a2a = False
-            from mindspore.nn.transformer import FeedForward
+            from mindspore.parallel._transformer import FeedForward
 
             self.ffn = FeedForward(hidden_size=hidden_size,
                                    ffn_hidden_size=ffn_hidden_size,
@@ -232,7 +235,7 @@ class MoE(Cell):
             self.group_wise_a2a = moe_config.group_wise_a2a
             if not (self.mp > 1 and self.dp == self.ep):
                 self.group_wise_a2a = False
-            from mindspore.nn.transformer import FeedForward
+            from mindspore.parallel._transformer import FeedForward
 
             self.ffn = FeedForward(hidden_size=hidden_size,
                                    ffn_hidden_size=ffn_hidden_size,
@@ -269,8 +272,8 @@ class MoE(Cell):
         pad_size = 0
         if self.group_wise_a2a:
             # If capacity can't div by mp, pad for mp shard.
-            if capacity%self.mp != 0:
-                pad_size = self.mp-(capacity%self.mp)
+            if capacity % self.mp != 0:
+                pad_size = self.mp-(capacity % self.mp)
             if pad_size != 0:
                 capacity += pad_size
                 pad_tensor = self.stride_slice_dp(expert_input, (0, 0, 0, 0),
@@ -326,8 +329,8 @@ class MoE(Cell):
         """
         # Pad capacity for comp_comm_parallel_degree split.
         pad_size = 0
-        if capacity%self.comp_comm_parallel_degree != 0:
-            pad_size = self.comp_comm_parallel_degree-(capacity%self.comp_comm_parallel_degree)
+        if capacity % self.comp_comm_parallel_degree != 0:
+            pad_size = self.comp_comm_parallel_degree-(capacity % self.comp_comm_parallel_degree)
             capacity += pad_size
             pad_tensor = self.stride_slice_dp(expert_input, (0, 0, 0, 0),
                                               (self.expert_dim, self.dp_group, pad_size, self.hidden_size),
@@ -646,9 +649,10 @@ class TopkRouter(Cell):
                                                     self.on_value, self.off_value))
             accum_combine_tensor = self.add2(accum_combine_tensor, combine_tensor)
 
-        # expert weights normalization
-        combine_tensor_sum = self.reduce_sum_keep2(self.reduce_sum_keep2(accum_combine_tensor, -1), -2)
-        accum_combine_tensor = self.div2(accum_combine_tensor, self.add4(combine_tensor_sum, 1e-9))
+        # expert weights normalization when k > 1
+        if self.num_experts_chosen > 1:
+            combine_tensor_sum = self.reduce_sum_keep2(self.reduce_sum_keep2(accum_combine_tensor, -1), -2)
+            accum_combine_tensor = self.div2(accum_combine_tensor, self.add4(combine_tensor_sum, 1e-9))
         # dispatch_tensor is of boolean type. Here, using NotEqual instead of Cast, for that 'Cast to bool' has
         # bad performance
         dispatch_tensor = self.not_equal(accum_combine_tensor, 0.0)

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

@@ -13,12 +13,12 @@
 # limitations under the License.
 # ============================================================================
 """
-Parallel Config for the Parallel Training
-This is an experimental interface that is subject to change and/or deletion.
+Parallel Config for the Parallel Training.
+These are experimental APIs that are subject to change or deletion.
 """
 from __future__ import absolute_import
 
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore import context
 import mindspore.communication.management as D
 from mindspore.context import ParallelMode