mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

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

@@ -1,4 +1,4 @@
-# Copyright 2021 Huawei Technologies Co., Ltd
+# Copyright 2021-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.
@@ -30,17 +30,17 @@ import mindspore.common.dtype as mstype
 from mindspore.ops import operations as P
 from mindspore.ops import functional as F
 from mindspore.nn.cell import Cell
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore import log as logger
-from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation
+from mindspore.parallel._utils import _get_parallel_mode
 from mindspore.context import ParallelMode
 from mindspore.log import _LogActionOnce
-from mindspore.nn.transformer.layers import _LayerNorm, _Linear, _check_input_shape, \
+from mindspore.parallel._transformer.layers import _LayerNorm, _Linear, \
     _args_type_validator_check, _valid_type_checks, _valid_value_checks, \
-    _check_shape_equal, _check_past_none_input_none, _check_input_dtype, _check_input_shape_value
-from mindspore.nn.transformer.op_parallel_config import default_dpmp_config, _PipeLineConfig, OpParallelConfig,\
+    _check_past_none_input_none, _check_input_dtype
+from mindspore.parallel._transformer.op_parallel_config import default_dpmp_config, _PipeLineConfig, OpParallelConfig, \
     _Config, _check_config, MoEParallelConfig
-from mindspore.nn.transformer.moe import default_moe_config, MoE, _check_moe_config
+from mindspore.parallel._transformer.moe import default_moe_config, MoE, _check_moe_config
 
 __all__ = [
     "AttentionMask",
@@ -352,9 +352,11 @@ class FeedForward(Cell):
             hidden_size (int): The dimension of the inputs.
             ffn_hidden_size (int): The intermediate hidden size.
             dropout_rate (float): The dropout rate for the second linear's output.
-            hidden_act (str): The activation of the internal feedforward layer. Supports 'relu',
+            hidden_act (str, nn.Cell): The activation of the internal feedforward layer. Supports 'relu',
                 'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
-                'hsigmoid', 'logsigmoid' and so on. Default: gelu.
+                'hsigmoid', 'logsigmoid' and so on. User can provide custom activition to the argument.
+                If user wants to run the net in the parallel mode, the custom activation must also provide
+                the `activation_shard` function. Please see examples. Default: gelu.
             expert_num (int): The number of experts used in Linear. For the case expert_num > 1, BatchMatMul is used
                 and the first dimension in BatchMatMul indicate expert_num. Default: 1.
             expert_group_size (int): The number of tokens in each data parallel group. Default: None. This parameter is
@@ -375,7 +377,7 @@ class FeedForward(Cell):
             [batch * seq_length, hidden_size]`.
 
         Raises:
-            ValueError: `hidden_act` is not a string.
+            TypeError: `hidden_act` is not a string or nn.Cell.
             TypeError: `parallel_config` is not a subclass of OpParallelConfig.
             ValueError: `ffn_hidden_size` is not a multiple of the model parallel way.
             ValueError: `hidden_size` is not a multiple of the model parallel way.
@@ -387,19 +389,51 @@ class FeedForward(Cell):
             >>> import numpy as np
             >>> from mindspore.nn.transformer import FeedForward
             >>> from mindspore import dtype as mstype
-            >>> from mindspore import Tensor
+            >>> from mindspore import Tensor, nn
+            >>> import mindspore.ops as ops
             >>> model = FeedForward(hidden_size=15, ffn_hidden_size=30, dropout_rate=0.1)
             >>> tensor = Tensor(np.ones((2, 20, 15)), mstype.float32)
             >>> output = model(tensor)
             >>> print(output.shape)
             (2, 20, 15)
+            >>> # Example 2 using custom hidden activation
+            >>> class MyActivationNoShard(nn.Cell):
+            ...     def __init__(self):
+            ...         super(MyActivationNoShard, self).__init__()
+            ...         self.add = ops.Add()
+            ...     def construct(self, x):
+            ...         return self.add(x, 0.1)
+            >>> model = FeedForward(hidden_size=15, ffn_hidden_size=30, dropout_rate=0.1,
+            ...                     hidden_act=MyActivationNoShard)
+            >>> tensor = Tensor(np.ones((2, 20, 15)), mstype.float32)
+            >>> output = model(tensor)
+            >>> print(output.shape)
+            (2, 20, 15)
+            >>> # Example 3 using custom hidden activation with activation_shard
+            >>> # If user wantss to run on the SEMI/AUTO parallel mode, the custom activation must provide
+            >>> # a class function named activation_shard. It accepts the argument parallel_config (OpParallelConfig,
+            >>> # MoEParallelConfig) and set the shard for the primitives used in the construct.
+            >>> class MyActivationWithShard(nn.Cell):
+            ...     def __init__(self):
+            ...         super(MyActivationWithShard, self).__init__()
+            ...         self.add = ops.Add()
+            ...     def construct(self, x):
+            ...         return self.add(x, 0.1)
+            ...     def activation_shard(self, parallel_config):
+            ...         self.add.shard(((parallel_config.data_parallel, parallel_config.model_parallel), ()))
+            >>>
+            >>> model = FeedForward(hidden_size=15, ffn_hidden_size=30, dropout_rate=0.1,
+            ...                     hidden_act=MyActivationWithShard)
+            >>> tensor = Tensor(np.ones((2, 20, 15)), mstype.float32)
+            >>> output = model(tensor)
+            >>> print(output.shape)
+            (2, 20, 15)
     """
     @_LogActionOnce(logger=logger, key='FeedForward',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
                                 ffn_hidden_size=Validator.check_positive_int,
                                 dropout_rate=Validator.check_non_negative_float,
-                                hidden_act=_valid_type_checks([str], "FeedForward"),
                                 param_init_type=_valid_value_checks([mstype.float32, mstype.float16],
                                                                     "FeedForward"),
                                 parallel_config=_valid_type_checks([OpParallelConfig, MoEParallelConfig],
@@ -413,7 +447,10 @@ class FeedForward(Cell):
                  param_init_type=mstype.float32,
                  parallel_config=default_dpmp_config):
         super(FeedForward, self).__init__()
-        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
+        if hidden_act is None or not (isinstance(hidden_act, str) or issubclass(hidden_act, nn.Cell)):
+            raise TypeError(f"For FeedForward cell, the hidden_act should str type or nn.Cell type, "
+                            f"but got {hidden_act}.")
+        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,):
             _check_config(parallel_config)
             mp = parallel_config.model_parallel
             if expert_num > 1:
@@ -459,9 +496,9 @@ class FeedForward(Cell):
             else:
                 self.projection.shard(strategy_matmul=((dp, mp), (mp, 1)))
             self.projection.bias.parallel_optimizer = False
-            self.dropout = nn.Dropout(1 - dropout_rate)
-            self.dropout_3d = nn.Dropout(1 - dropout_rate)
-            self.dropout_4d = nn.Dropout(1 - dropout_rate)
+            self.dropout = nn.Dropout(p=dropout_rate)
+            self.dropout_3d = nn.Dropout(p=dropout_rate)
+            self.dropout_4d = nn.Dropout(p=dropout_rate)
             self.cast = P.Cast()
         else:
             _check_config(parallel_config)
@@ -519,16 +556,18 @@ class FeedForward(Cell):
                 self.projection.shard(strategy_matmul=((dp, mp), (mp, 1)),
                                       strategy_bias=((dp, 1), (1,)))
             self.projection.bias.parallel_optimizer = False
-            self.dropout = nn.Dropout(1 - dropout_rate)
+            self.dropout = nn.Dropout(p=dropout_rate)
             self.dropout.dropout.shard(((dp, 1),))
-            self.dropout_3d = nn.Dropout(1 - dropout_rate)
+            self.dropout_3d = nn.Dropout(p=dropout_rate)
             self.dropout_3d.dropout.shard(((dp, 1, 1),))
-            self.dropout_4d = nn.Dropout(1 - dropout_rate)
+            self.dropout_4d = nn.Dropout(p=dropout_rate)
             self.dropout_4d.dropout.shard(((dp, ep, 1, 1),))
             self.cast = P.Cast()
+            # for grouped pairwise exchange alltoall method in pass
+            self.mapping.matmul.add_prim_attr("gpea_label", True)
+            self.projection.matmul.add_prim_attr("gpea_label", True)
 
     def construct(self, x):
-        _check_input_shape(F.shape(x), "x", self.cls_name, [2, 3])
         _check_input_dtype(F.dtype(x), "x", [mstype.float32, mstype.float16], self.cls_name)
         x = self.cast(x, mstype.float16)
         # returned shape is [bs, seq_length, ffn_hidden_size] or [bs * seq_length, ffn_hidden_size]
@@ -601,9 +640,7 @@ class AttentionMask(Cell):
         self.multiply = P.Mul().shard(((parallel_config.data_parallel, 1, 1), (1, 1, 1)))
 
     def construct(self, input_mask):
-        _check_input_shape(F.shape(input_mask), "input_mask", self.cls_name, 2)
         _check_input_dtype(F.dtype(input_mask), "input_mask", [mstype.float32, mstype.float16], self.cls_name)
-        _check_input_shape_value(F.shape(input_mask), 1, "input_mask", self.cls_name, self.seq_length)
         input_mask = P.Cast()(self.not_equal(input_mask, 0), mstype.float16)
         input_shape = P.Shape()(input_mask)
         shape_right = (input_shape[0], 1, input_shape[1])
@@ -698,10 +735,9 @@ class VocabEmbedding(Cell):
                         f"model parallel for the embedding lookup.")
 
     def construct(self, input_ids):
-        _check_input_shape(F.shape(input_ids), "input_ids", self.cls_name, 2)
         _check_input_dtype(F.dtype(input_ids), "input_ids", [mstype.int32], self.cls_name)
         output = self.gather(self.embedding_table, input_ids, 0)
-        return output, self.embedding_table
+        return output, self.embedding_table.value()
 
 
 class MultiHeadAttention(Cell):
@@ -718,7 +754,9 @@ class MultiHeadAttention(Cell):
         if query, key and value tensor is same, then it will be self attention.
 
         Args:
-            batch_size(int): The batch size of the input tensor.
+            batch_size(int): The batch size of the input tensor when do increnmental prediction. Should be a positive
+                value. When do training or prediction, the argument will not work and the user can just pass None to
+                the argument.
             src_seq_length(int): The sequence length of the query vector.
             tgt_seq_length(int): The sequence length of the key and value vector.
             hidden_size(int): The hidden size of the input.
@@ -751,9 +789,9 @@ class MultiHeadAttention(Cell):
             - **value_tensor** (Tensor) - The value vector with shape (batch_size, tgt_seq_length, hidden_size) or
               (batch_size * tgt_seq_length, hidden_size), if the use_past is False or is_first_iteration=True.
               Otherwise, must be (batch_size, 1, hidden_size)
-            - **attention_mask** (Tensor) - The attention mask matrix with shape (batch_size, src_seq_length,
-              tgt_seq_length), if the use_past is False or is_first_iteration=True. Otherwise,
-              must be (batch_size, 1, tgt_seq_length)
+            - **attention_mask** (Tensor) - If the use_past is False or is_first_iteration=True, the attention mask
+              matrix should ba (batch_size, src_seq_length, tgt_seq_length), or None. None means there will be no mask
+              in softmax computation. Otherwise, the mask must be (batch_size, 1, tgt_seq_length)
             - **key_past** (Tensor) - Float16 tensor with shape (batch_size, num_heads, size_per_head, tgt_seq_length).
               The past calculated key vector. Used for incremental prediction when the use_past is True.
               Default None.
@@ -783,7 +821,7 @@ class MultiHeadAttention(Cell):
             >>> from mindspore.nn.transformer import MultiHeadAttention
             >>> from mindspore import dtype as mstype
             >>> from mindspore import Tensor
-            >>> model = MultiHeadAttention(batch_size=2, hidden_size=15, src_seq_length=20, tgt_seq_length=20,
+            >>> model = MultiHeadAttention(batch_size=None, hidden_size=15, src_seq_length=20, tgt_seq_length=20,
             ...                            num_heads=3)
             >>> from_tensor = Tensor(np.ones((2, 20, 15)), mstype.float32)
             >>> to_tensor = Tensor(np.ones((2, 20, 15)), mstype.float16)
@@ -830,8 +868,7 @@ class MultiHeadAttention(Cell):
     """
     @_LogActionOnce(logger=logger, key='MultiHeadAttention',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
-    @_args_type_validator_check(batch_size=Validator.check_positive_int,
-                                hidden_size=Validator.check_positive_int,
+    @_args_type_validator_check(hidden_size=Validator.check_positive_int,
                                 num_heads=Validator.check_positive_int,
                                 src_seq_length=Validator.check_positive_int,
                                 tgt_seq_length=Validator.check_positive_int,
@@ -860,10 +897,13 @@ class MultiHeadAttention(Cell):
                  parallel_config=default_dpmp_config):
         super(MultiHeadAttention, self).__init__()
         self._is_ascend = context.get_context('device_target') in ["Ascend"]
-        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
+        self.dp = parallel_config.data_parallel
+        self.is_parallel_mode = _get_parallel_mode() in (
+            ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
+        if batch_size:
+            Validator.check_positive_int(batch_size)
+        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,):
             _check_config(parallel_config)
-            self.is_parallel_mode = _get_parallel_mode() in (
-                ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
             self.src_seq_length = src_seq_length
             self.tgt_seq_length = tgt_seq_length
             self.hidden_size = hidden_size
@@ -883,11 +923,6 @@ class MultiHeadAttention(Cell):
                                  "'parallel_config.model_parallel', but got the num_heads is {} "
                                  "and the parallel_config.model_parallel  is {}."
                                  .format(num_heads, parallel_config.model_parallel))
-            if self.is_parallel_mode and batch_size % parallel_config.data_parallel != 0:
-                raise ValueError("For 'MultiHeadAttention', the class variable 'batch_size' must be a multiple of "
-                                 "'parallel_config.data_parallel', but got the batch_size is {} "
-                                 "and the parallel_config.data_parallel is {}."
-                                 .format(batch_size, parallel_config.data_parallel))
             self.is_first_iteration = True
             # Output layer
             self.projection = _Linear(in_channels=hidden_size,
@@ -918,8 +953,8 @@ class MultiHeadAttention(Cell):
             # Normalize factor for attention, sqrt(dk) as widely used
             self.scale_factor = Tensor(math.sqrt(math.sqrt(self.size_per_head)))
             self.use_past = use_past
-            self.dropout = nn.Dropout(1 - hidden_dropout_rate)
-            self.prob_dropout = nn.Dropout(1 - attention_dropout_rate)
+            self.dropout = nn.Dropout(p=hidden_dropout_rate)
+            self.prob_dropout = nn.Dropout(p=attention_dropout_rate)
             self.softmax = nn.Softmax().to_float(softmax_compute_type)
             self.softmax_3d = nn.Softmax().to_float(softmax_compute_type)
             self.expand_dims = P.ExpandDims()
@@ -961,8 +996,6 @@ class MultiHeadAttention(Cell):
                 self.mul1 = P.Mul().shard(((1, 1, 1, 1), (1, 1, 1, 1)))
         else:
             _check_config(parallel_config)
-            self.is_parallel_mode = _get_parallel_mode() in (
-                ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
             self.src_seq_length = src_seq_length
             self.tgt_seq_length = tgt_seq_length
             self.hidden_size = hidden_size
@@ -982,11 +1015,6 @@ class MultiHeadAttention(Cell):
                                  "'parallel_config.model_parallel', but got the num_heads is {} "
                                  "and the parallel_config.model_parallel  is {}."
                                  .format(num_heads, parallel_config.model_parallel))
-            if self.is_parallel_mode and batch_size % parallel_config.data_parallel != 0:
-                raise ValueError("For 'MultiHeadAttention', the class variable 'batch_size' must be a multiple of "
-                                 "'parallel_config.data_parallel', but got the batch_size is {} "
-                                 "and the parallel_config.data_parallel is {}."
-                                 .format(batch_size, parallel_config.data_parallel))
             self.is_first_iteration = True
             # Output layer
             self.projection = _Linear(in_channels=hidden_size,
@@ -1026,9 +1054,9 @@ class MultiHeadAttention(Cell):
             # Normalize factor for attention, sqrt(dk) as widely used
             self.scale_factor = Tensor(math.sqrt(math.sqrt(self.size_per_head)))
             self.use_past = use_past
-            self.dropout = nn.Dropout(1 - hidden_dropout_rate)
+            self.dropout = nn.Dropout(p=hidden_dropout_rate)
             self.dropout.dropout.shard(((parallel_config.data_parallel, 1),))
-            self.prob_dropout = nn.Dropout(1 - attention_dropout_rate)
+            self.prob_dropout = nn.Dropout(p=attention_dropout_rate)
             self.prob_dropout.dropout.shard(
                 ((parallel_config.data_parallel, parallel_config.model_parallel, 1, 1),))
             self.softmax = nn.Softmax().to_float(softmax_compute_type)
@@ -1086,10 +1114,12 @@ class MultiHeadAttention(Cell):
                   value_past=None, batch_valid_length=None):
         self._check_inputs(query_tensor, key_tensor, value_tensor, attention_mask, key_past,
                            value_past, batch_valid_length)
-        query_tensor, key_tensor, value_tensor, batch_size, ori_shape = self._convert_to_2d_tensor(query_tensor,
-                                                                                                   key_tensor,
-                                                                                                   value_tensor,
-                                                                                                   attention_mask)
+        ori_shape = F.shape(query_tensor)
+        batch_size = self._get_batch_size_from_query(query_tensor)
+        query_tensor, key_tensor, value_tensor = self._convert_to_2d_tensor(query_tensor,
+                                                                            key_tensor,
+                                                                            value_tensor,
+                                                                            attention_mask)
         ori_dtype = F.dtype(query_tensor)
         query_tensor = F.cast(query_tensor, self.dtype)
         key_tensor = F.cast(key_tensor, self.dtype)
@@ -1102,21 +1132,24 @@ class MultiHeadAttention(Cell):
         query = self.transpose(
             F.reshape(
                 query,
-                (batch_size, -1, self.n_head, self.size_per_head)),
+                (batch_size, self._get_seq_length_under_incremental(self.src_seq_length),
+                 self.n_head, self.size_per_head)),
             (0, 2, 1, 3))
         # the returned shape is [bs, size_per_head, seq_length, num_heads]
         key = self.transpose(
             F.reshape(
-                key, (batch_size, -1, self.n_head, self.size_per_head)),
+                key, (batch_size, self._get_seq_length_under_incremental(self.tgt_seq_length),
+                      self.n_head, self.size_per_head)),
             (0, 2, 3, 1))
         # the returned shape is [bs, num_heads, seq_length, size_per_head]
         value = self.transpose(
             F.reshape(
                 value,
-                (batch_size, -1, self.n_head, self.size_per_head)),
+                (batch_size, self._get_seq_length_under_incremental(self.tgt_seq_length),
+                 self.n_head, self.size_per_head)),
             (0, 2, 1, 3))
         # support input shape is [bs, seq, seq] or [bs, heads, seq, seq]
-        if len(F.shape(attention_mask)) == 3:
+        if attention_mask is not None and len(F.shape(attention_mask)) == 3:
             # expand attention mask from [bs, seq, seq] -> [bs, 1, seq, seq]
             attention_mask = self.expand_dims(attention_mask, 1)
         # key and value for current token(s)
@@ -1167,35 +1200,30 @@ class MultiHeadAttention(Cell):
         output = F.cast(output, ori_dtype)
         return output, layer_present
 
+    def _get_batch_size_from_query(self, query):
+        r"""Get the batch size from query tensor"""
+        # For the incremental prediction, the seq length for the input is 1.
+        if len(F.shape(query)) == 2 and ((self.use_past and self.is_first_iteration) or (not self.use_past)):
+            return F.shape(query)[0] // self.src_seq_length
+        return F.shape(query)[0]
+
+    def _get_seq_length_under_incremental(self, length):
+        r"""Return the length of the tensor.
+            For the incremental prediction, the seq length for the input is 1.
+        """
+        if self.use_past and not self.is_first_iteration:
+            return 1
+        return length
+
     def _check_inputs(self, query_tensor, key_tensor, value_tensor, attention_mask, key_past=None,
                       value_past=None, batch_valid_length=None):
         r"""Check inputs"""
-        if not self.use_past or (self.use_past and self.is_first_iteration):
-            _check_shape_equal(F.shape(query_tensor), "query_tensor", self.cls_name,
-                               [[self.batch_size, self.src_seq_length, self.hidden_size],
-                                [self.batch_size * self.src_seq_length, self.hidden_size]])
-            _check_shape_equal(F.shape(key_tensor), "key_tensor", self.cls_name,
-                               [[self.batch_size, self.tgt_seq_length, self.hidden_size],
-                                [self.batch_size * self.tgt_seq_length, self.hidden_size]])
-            _check_shape_equal(F.shape(value_tensor), "value_tensor", self.cls_name,
-                               [[self.batch_size, self.tgt_seq_length, self.hidden_size],
-                                [self.batch_size * self.tgt_seq_length, self.hidden_size]])
-            _check_shape_equal(F.shape(attention_mask), "attention_mask", self.cls_name,
-                               [self.batch_size, self.src_seq_length, self.tgt_seq_length])
-        else:
-            _check_shape_equal(F.shape(query_tensor), "query_tensor", self.cls_name,
-                               [[self.batch_size, 1, self.hidden_size], [self.batch_size, self.hidden_size]])
-            _check_shape_equal(F.shape(key_tensor), "key_tensor", self.cls_name,
-                               [[self.batch_size, 1, self.hidden_size], [self.batch_size, self.hidden_size]])
-            _check_shape_equal(F.shape(value_tensor), "value_tensor", self.cls_name,
-                               [[self.batch_size, 1, self.hidden_size], [self.batch_size, self.hidden_size]])
-            _check_shape_equal(F.shape(attention_mask), "attention_mask", self.cls_name,
-                               [[self.batch_size, 1, self.tgt_seq_length], [self.batch_size, self.hidden_size]])
-
         _check_input_dtype(F.dtype(query_tensor), "query_tensor", [mstype.float32, mstype.float16], self.cls_name)
         _check_input_dtype(F.dtype(key_tensor), "key_tensor", [mstype.float32, mstype.float16], self.cls_name)
         _check_input_dtype(F.dtype(value_tensor), "value_tensor", [mstype.float32, mstype.float16], self.cls_name)
-        _check_input_dtype(F.dtype(attention_mask), "attention_mask", [mstype.float32, mstype.float16], self.cls_name)
+        if attention_mask is not None:
+            _check_input_dtype(F.dtype(attention_mask), "attention_mask", [mstype.float32, mstype.float16],
+                               self.cls_name)
 
         key_is_tensor = isinstance(key_past, Tensor)
         value_is_tensor = isinstance(value_past, Tensor)
@@ -1210,13 +1238,8 @@ class MultiHeadAttention(Cell):
         _check_past_none_input_none(self.use_past, "batch_valid_length", self.cls_name, None,
                                     batch_valid_length_is_tensor, batch_is_default)
         if self.use_past:
-            _check_shape_equal(F.shape(key_past), "key_past", self.cls_name,
-                               [self.batch_size, self.n_head, self.size_per_head, self.tgt_seq_length])
             _check_input_dtype(F.dtype(key_past), "key_past", [mstype.float16], self.cls_name)
-            _check_shape_equal(F.shape(value_past), "value_past", self.cls_name,
-                               [self.batch_size, self.n_head, self.tgt_seq_length, self.size_per_head])
             _check_input_dtype(F.dtype(value_past), "value_past", [mstype.float16], self.cls_name)
-            _check_shape_equal(F.shape(batch_valid_length), "batch_valid_length", self.cls_name, [self.batch_size])
             _check_input_dtype(F.dtype(batch_valid_length), "batch_valid_length", [mstype.int32], self.cls_name)
         return True
 
@@ -1228,7 +1251,8 @@ class MultiHeadAttention(Cell):
         key_tensor = F.reshape(key_tensor, (-1, key_shape[-1]))
         value_shape = F.shape(value_tensor)
         value_tensor = F.reshape(value_tensor, (-1, value_shape[-1]))
-        return query_tensor, key_tensor, value_tensor, F.shape(attention_mask)[0], query_shape
+
+        return query_tensor, key_tensor, value_tensor
 
     def _merge_heads(self, x):
         """
@@ -1286,30 +1310,31 @@ class MultiHeadAttention(Cell):
         score = self.batch_matmul(query, key)
 
         ori_dtype = P.DType()(score)
-        score = P.Cast()(score, self.softmax_dtype)
+        attention_scores = P.Cast()(score, self.softmax_dtype)
 
         # for input size of (bs, 1) namely the second graph,
         # the shape of attention_mask matrix should be (bs, 1, 1, seq_length)
-        if self.use_past and not self.is_first_iteration:
-            # Calculate the current total token
-            current_index = self.reducesum(F.cast(self.not_equal(self.slice(key, (0, 0, 0, 0),
-                                                                            (F.shape(query)[0], 1, 1, self.seq_length),
-                                                                            (1, 1, 1, 1)),
-                                                                 0), mstype.float32), (1, 2, 3))
-            # Get the precise position index
-            index = self.sub1(F.cast(current_index, mstype.int32), 1)
-            index = F.reshape(index, (-1, 1, 1))
-            # Calculate the attention_mask matrix via the position index
-            attention_mask = F.cast(self.tensor_le(self.range, index), mstype.int32)
-            attention_mask = self.expand_dims(attention_mask, 2)
-
-        # Minus 10000 for the position where masked to exclude them from softmax
-        multiplu_out = self.sub(
-            P.Cast()(F.tuple_to_array((1.0,)), P.DType()(score)),
-            P.Cast()(attention_mask, P.DType()(score)))
-
-        adder = self.mul(multiplu_out, self.multiply_data)
-        attention_scores = self.add(adder, score)
+        if attention_mask is not None:
+            if self.use_past and not self.is_first_iteration:
+                # Calculate the current total token
+                current_index = self.reducesum(F.cast(self.not_equal(self.slice(key, (0, 0, 0, 0),
+                                                                                (F.shape(query)[0], 1, 1,
+                                                                                 self.seq_length),
+                                                                                (1, 1, 1, 1)),
+                                                                     0), mstype.float32), (1, 2, 3))
+                # Get the precise position index
+                index = self.sub1(F.cast(current_index, mstype.int32), 1)
+                index = F.reshape(index, (-1, 1, 1))
+                # Calculate the attention_mask matrix via the position index
+                attention_mask = F.cast(self.tensor_le(self.range, index), mstype.int32)
+                attention_mask = self.expand_dims(attention_mask, 2)
+            # Minus 10000 for the position where masked to exclude them from softmax
+            multiplu_out = self.sub(
+                P.Cast()(F.tuple_to_array((1.0,)), P.DType()(attention_scores)),
+                P.Cast()(attention_mask, P.DType()(attention_scores)))
+
+            adder = self.mul(multiplu_out, self.multiply_data)
+            attention_scores = self.add(adder, attention_scores)
 
         # attention probs
         attention_probs = self._softmax(attention_scores)
@@ -1328,7 +1353,9 @@ class TransformerEncoderLayer(Cell):
         encoder layer, including multihead attention and feedward layer.
 
         Args:
-            batch_size(int): The batch size of the input tensor.
+            batch_size(int): The batch size of the input tensor when do increnmental prediction. Should be a positive
+                value. When do training or prediction, the argument will not work and the user can just pass None to
+                the argument.
             hidden_size(int): The hidden size of the input.
             ffn_hidden_size(int): The hidden size of bottleneck in the feedforward layer.
             num_heads(int): The number of the heads.
@@ -1342,9 +1369,12 @@ class TransformerEncoderLayer(Cell):
                 Should be mstype.float32 or mstype.float16. Default mstype.float32.
             param_init_type(dtype.Number): The parameter initialization type of the module.
                 Should be mstype.float32 or mstype.float16. Default mstype.float32.
-            hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
+            hidden_act (str, nn.Cell): The activation of the internal feedforward layer. Supports 'relu',
                 'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
-                'hsigmoid', 'logsigmoid' and so on. Default: gelu.
+                'hsigmoid', 'logsigmoid' and so on. User can provide custom activition to the argument.
+                If user wants to run the net in the parallel mode, the custom activation must also provide
+                the `activation_shard` function. Please see the examples of the
+                class:`mindspore.nn.transformer.FeedForward`. Default: gelu.
             use_past(bool): Use the past state to compute, used for incremental prediction. For example, if we have two
                 words and want to generate the ten more words. We just need to compute the two words' state only once,
                 and generate the next word one by one. When use_past is True, there are two steps to run the prediction.
@@ -1362,8 +1392,9 @@ class TransformerEncoderLayer(Cell):
             - **x** (Tensor) - Float Tensor, shape should be [batch_size, seq_length, hidden_size] or
               [batch_size * seq_length, hidden_size], if the use_past is False or is_first_iteration=True. Otherwise,
               should be [batch_size, 1, hidden_size]
-            - **input_mask** (Tensor) - Float Tensor, attention mask with shape [batch_size, seq_length, seq_length],
-              if the use_past is False or is_first_iteration=True. Otherwise, should be [batch_size, 1, hidden_size]
+            - **input_mask** (Tensor) - Float Tensor, If the use_past is False or is_first_iteration=True,
+              the attention mask matrix should ba [batch_size, seq_length, seq_length], or None. None means there will
+              be no mask in softmax computation. Otherwise, should be [batch_size, 1, hidden_size]
             - **init_reset** (Tensor) - A bool tensor with shape [1], used to clear the past key parameter and
               past value parameter used in the incremental prediction. Only valid when use_past is True. Default True.
             - **batch_valid_length** (Tensor) - Int32 tensor with shape [batch_size] the past calculated the index.
@@ -1430,14 +1461,12 @@ class TransformerEncoderLayer(Cell):
     """
     @_LogActionOnce(logger=logger, key='TransformerEncoderLayer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
-    @_args_type_validator_check(batch_size=Validator.check_positive_int,
-                                hidden_size=Validator.check_positive_int,
+    @_args_type_validator_check(hidden_size=Validator.check_positive_int,
                                 num_heads=Validator.check_positive_int,
                                 ffn_hidden_size=Validator.check_positive_int,
                                 seq_length=Validator.check_positive_int,
                                 attention_dropout_rate=Validator.check_non_negative_float,
                                 hidden_dropout_rate=Validator.check_non_negative_float,
-                                hidden_act=_valid_type_checks([str], "TransformerEncoderLayer"),
                                 post_layernorm_residual=Validator.check_bool,
                                 layernorm_compute_type=_valid_value_checks([mstype.float32, mstype.float16],
                                                                            "TransformerEncoderLayer"),
@@ -1465,7 +1494,10 @@ class TransformerEncoderLayer(Cell):
                  moe_config=default_moe_config,
                  parallel_config=default_dpmp_config):
         super(TransformerEncoderLayer, self).__init__()
-        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
+        if batch_size or use_past:
+            Validator.check_positive_int(batch_size)
+        self.batch_size = batch_size
+        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,):
             _check_config(parallel_config)
             if num_heads % parallel_config.model_parallel != 0:
                 raise ValueError(
@@ -1488,7 +1520,6 @@ class TransformerEncoderLayer(Cell):
             self.use_past = use_past
             self.seq_length = seq_length
             self.hidden_size = hidden_size
-            self.batch_size = batch_size
             self.layernorm1 = _LayerNorm((hidden_size,)).to_float(layernorm_compute_type)
             self.layernorm2 = _LayerNorm((hidden_size,)).to_float(layernorm_compute_type)
 
@@ -1564,7 +1595,6 @@ class TransformerEncoderLayer(Cell):
             self.use_past = use_past
             self.seq_length = seq_length
             self.hidden_size = hidden_size
-            self.batch_size = batch_size
             self.layernorm1 = _LayerNorm((hidden_size,)).to_float(layernorm_compute_type)
             self.layernorm1.shard(((parallel_config.data_parallel, 1),))
             self.layernorm2 = _LayerNorm((hidden_size,)).to_float(layernorm_compute_type)
@@ -1623,11 +1653,14 @@ class TransformerEncoderLayer(Cell):
             raise RuntimeError(f"The {self.cls_name} only support sharding propagation or "
                                f"semi-auto parallel mode now.")
 
-    def construct(self, x, input_mask, init_reset=True, batch_valid_length=None):
+    def construct(self, x, input_mask=None, init_reset=True, batch_valid_length=None):
         self._check_input(x, input_mask, init_reset, batch_valid_length)
         x_shape = F.shape(x)
         x = F.reshape(x, (-1, x_shape[-1]))
-        input_x = self.layernorm1(x)
+        if self.post_layernorm_residual:
+            input_x = x
+        else:
+            input_x = self.layernorm1(x)
         input_x = F.cast(input_x, self.dtype)
 
         # indicate whether reset saved states
@@ -1636,8 +1669,10 @@ class TransformerEncoderLayer(Cell):
 
         if self.use_past:
             # reset states, init_reset True for reuse and False for reset
-            key_reset = self.assign(self.key_past, self.mul(self.key_past, F.cast(init_reset, self.dtype)))
-            value_reset = self.assign(self.value_past, self.mul(self.value_past, F.cast(init_reset, self.dtype)))
+            self.assign(self.key_past, self.mul(self.key_past, F.cast(init_reset, self.dtype)))
+            key_reset = self.key_past
+            self.assign(self.value_past, self.mul(self.value_past, F.cast(init_reset, self.dtype)))
+            value_reset = self.value_past
             # add dependency for desired execution order
             input_x = F.depend(input_x, key_reset)
             input_x = F.depend(input_x, value_reset)
@@ -1665,8 +1700,10 @@ class TransformerEncoderLayer(Cell):
             # current key and value
             key_present, value_present = layer_present
             # update key and value calculated this step
-            key_update = self.assign(self.key_past, key_present)
-            value_update = self.assign(self.value_past, value_present)
+            self.assign(self.key_past, key_present)
+            key_update = self.key_past
+            self.assign(self.value_past, value_present)
+            value_update = self.value_past
             # add dependency for desired execution order
             key_update = F.depend(key_update, key_reset)
             value_update = F.depend(value_update, value_reset)
@@ -1683,11 +1720,15 @@ class TransformerEncoderLayer(Cell):
 
             if self.post_layernorm_residual:
                 output = self.add_3d(output_x, mlp_logit)
+                output = F.reshape(output, (-1, x_shape[-1]))
+                output = self.layernorm1(output)
+                output = F.reshape(output, x_shape)
             else:
                 output = self.add_3d(x, mlp_logit)
         else:
             if self.post_layernorm_residual:
                 output = self.add(output_x, mlp_logit)
+                output = self.layernorm1(output)
             else:
                 output = self.add(x, mlp_logit)
             output = F.reshape(output, x_shape)
@@ -1698,18 +1739,9 @@ class TransformerEncoderLayer(Cell):
 
     def _check_input(self, x, input_mask, init_reset, batch_valid_length):
         r"""Check inputs"""
-        if not self.use_past or (self.use_past and self.is_first_iteration):
-            _check_shape_equal(F.shape(x), "x", self.cls_name,
-                               [[self.batch_size, self.seq_length, self.hidden_size],
-                                [self.batch_size * self.seq_length, self.hidden_size]])
-            _check_shape_equal(F.shape(input_mask), "input_mask", self.cls_name,
-                               [self.batch_size, self.seq_length, self.seq_length])
-        else:
-            _check_shape_equal(F.shape(x), "x", self.cls_name, [self.batch_size, 1, self.hidden_size])
-            _check_shape_equal(F.shape(input_mask), "input_mask", self.cls_name,
-                               [self.batch_size, 1, self.seq_length])
         _check_input_dtype(F.dtype(x), "x", [mstype.float32, mstype.float16], self.cls_name)
-        _check_input_dtype(F.dtype(input_mask), "input_mask", [mstype.float32, mstype.float16], self.cls_name)
+        if input_mask is not None:
+            _check_input_dtype(F.dtype(input_mask), "input_mask", [mstype.float32, mstype.float16], self.cls_name)
 
         init_reset_is_tensor = isinstance(init_reset, Tensor)
         init_reset_is_default = init_reset is True
@@ -1721,9 +1753,7 @@ class TransformerEncoderLayer(Cell):
                                     batch_valid_length_is_tensor, batch_is_default)
 
         if self.use_past:
-            _check_shape_equal(F.shape(init_reset), "init_reset", self.cls_name, [1])
             _check_input_dtype(F.dtype(init_reset), "init_reset", [mstype.bool_], self.cls_name)
-            _check_shape_equal(F.shape(batch_valid_length), "batch_valid_length", self.cls_name, [self.batch_size])
             _check_input_dtype(F.dtype(batch_valid_length), "batch_valid_length", [mstype.int32], self.cls_name)
         return True
 
@@ -1738,7 +1768,9 @@ class TransformerDecoderLayer(Cell):
             hidden_size(int): The hidden size of the input.
             ffn_hidden_size(int): The hidden size of bottleneck in the feedforward layer.
             num_heads(int): The number of the heads.
-            batch_size(int): The batch size of the input tensor.
+            batch_size(int): The batch size of the input tensor when do increnmental prediction. Should be a positive
+                value. When do training or prediction, the argument will not work and the user can just pass None to
+                the argument.
             src_seq_length(int): The input source sequence length.
             tgt_seq_length(int): The input target sequence length.
             attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1.
@@ -1751,9 +1783,12 @@ class TransformerDecoderLayer(Cell):
                 Should be dtype.float32 or dtype.float16. Default mstype.float32.
             param_init_type(dtype.Number): The parameter initialization type of the module.
                 Should be dtype.float32 or dtype.float16. Default dtype.float32.
-            hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
+            hidden_act (str, nn.Cell): The activation of the internal feedforward layer. Supports 'relu',
                 'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
-                'hsigmoid', 'logsigmoid' and so on. Default: gelu.
+                'hsigmoid', 'logsigmoid' and so on. User can provide custom activition to the argument.
+                If user wants to run the net in the parallel mode, the custom activation must also provide
+                the `activation_shard` function. Please see the examples of the
+                class:`mindspore.nn.transformer.FeedForward`. Default: gelu.
             moe_config(MoEConfig): The configuration of MoE (Mixture of Expert). Default is an instance of MoEConfig
                 with default values. Please see `MoEConfig`.
             parallel_config(OpParallelConfig, MoEParallelConfig): The parallel configure. When MoE is applied,
@@ -1764,13 +1799,13 @@ class TransformerDecoderLayer(Cell):
             - **hidden_stats** (Tensor) - The input tensor with shape [batch_size, tgt_seq_length, hidden_size] or
               [batch_size * tgt_seq_length, hidden_size].
             - **decoder_mask** (Tensor) - The attention mask for decoder with shape [batch_size, src_seq_length,
-              seq_length].
+              seq_length] or None. None means there will be no mask in softmax computation in self attention.
             - **encoder_output** (Tensor) - The output of the encoder with shape [batch_size, seq_length, hidden_size]
               or [batch_size * seq_length, hidden_size].
               Note this args can not be passed by None when the net is in outermost layer. Default None.
             - **memory_mask** (Tensor) - The memory mask of the cross attention with shape [batch, tgt_seq_length,
-              src_seq_length] where tgt_seq_length is the length of the decoder. Note this args can not be passed by
-              None when the net is in outermost layer. Default None.
+              src_seq_length] where tgt_seq_length is the length of the decoder. The user can also pass None. None
+              means there will be no mask in softmax computation in cross attention. Default None.
             - **init_reset** (Tensor) - A bool tensor with shape [1], used to clear the past key parameter and
               past value parameter used in the incremental prediction. Only valid when use_past is True. Default True.
             - **batch_valid_length** (Tensor) - Int32 tensor with shape [batch_size] the past calculated the index.
@@ -1815,15 +1850,13 @@ class TransformerDecoderLayer(Cell):
     """
     @_LogActionOnce(logger=logger, key='TransformerDecoderLayer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
-    @_args_type_validator_check(batch_size=Validator.check_positive_int,
-                                hidden_size=Validator.check_positive_int,
+    @_args_type_validator_check(hidden_size=Validator.check_positive_int,
                                 num_heads=Validator.check_positive_int,
                                 ffn_hidden_size=Validator.check_positive_int,
                                 src_seq_length=Validator.check_positive_int,
                                 tgt_seq_length=Validator.check_positive_int,
                                 attention_dropout_rate=Validator.check_non_negative_float,
                                 hidden_dropout_rate=Validator.check_non_negative_float,
-                                hidden_act=_valid_type_checks([str], "TransformerDecoderLayer"),
                                 post_layernorm_residual=Validator.check_bool,
                                 layernorm_compute_type=_valid_value_checks([mstype.float32, mstype.float16],
                                                                            "TransformerDecoderLayer"),
@@ -1854,7 +1887,9 @@ class TransformerDecoderLayer(Cell):
         _check_moe_config(moe_config, parallel_config)
         self.use_moe = (moe_config.expert_num > 1)
         config_to_attention = parallel_config.dpmp if self.use_moe else parallel_config
-        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
+        if batch_size or use_past:
+            Validator.check_positive_int(batch_size)
+        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,):
             _check_config(parallel_config)
             if num_heads % parallel_config.model_parallel != 0:
                 raise ValueError("For 'TransformerDecoderLayer', the class variable 'num_heads' must be divisibled by "
@@ -2066,8 +2101,10 @@ class TransformerDecoderLayer(Cell):
         value_reset = None
         if self.use_past:
             # reset states, init_reset True for reuse and False for reset
-            key_reset = self.assign(self.key_past, self.mul(self.key_past, F.cast(init_reset, self.dtype)))
-            value_reset = self.assign(self.value_past, self.mul(self.value_past, F.cast(init_reset, self.dtype)))
+            self.assign(self.key_past, self.mul(self.key_past, F.cast(init_reset, self.dtype)))
+            key_reset = self.key_past
+            self.assign(self.value_past, self.mul(self.value_past, F.cast(init_reset, self.dtype)))
+            value_reset = self.value_past
             # add dependency for desired execution order
             input_x = F.depend(input_x, key_reset)
             input_x = F.depend(input_x, value_reset)
@@ -2110,8 +2147,10 @@ class TransformerDecoderLayer(Cell):
             # current key and value
             key_present, value_present = layer_present
             # update key and value calculated this step
-            key_update = self.assign(self.key_past, key_present)
-            value_update = self.assign(self.value_past, value_present)
+            self.assign(self.key_past, key_present)
+            key_update = self.key_past
+            self.assign(self.value_past, value_present)
+            value_update = self.value_past
             # add dependency for desired execution order
             key_update = F.depend(key_update, key_reset)
             value_update = F.depend(value_update, value_reset)
@@ -2143,29 +2182,14 @@ class TransformerDecoderLayer(Cell):
 
     def _check_input(self, hidden_states, attention_mask, encoder_output, memory_mask, init_reset, batch_valid_length):
         r"""Check inputs"""
-        if not self.use_past or (self.use_past and self.is_first_iteration):
-            _check_shape_equal(F.shape(hidden_states), "hidden_states", self.cls_name,
-                               [[self.batch_size, self.tgt_seq_length, self.hidden_size],
-                                [self.batch_size * self.tgt_seq_length, self.hidden_size]])
-            _check_shape_equal(F.shape(attention_mask), "attention_mask", self.cls_name,
-                               [self.batch_size, self.tgt_seq_length, self.tgt_seq_length])
-
-        else:
-            _check_shape_equal(F.shape(hidden_states), "hidden_states", self.cls_name,
-                               [self.batch_size, 1, self.hidden_size])
-            _check_shape_equal(F.shape(attention_mask), "attention_mask", self.cls_name,
-                               [self.batch_size, 1, self.tgt_seq_length])
         _check_input_dtype(F.dtype(hidden_states), "hidden_states", [mstype.float32, mstype.float16], self.cls_name)
-        _check_input_dtype(F.dtype(attention_mask), "attention_mask", [mstype.float32, mstype.float16], self.cls_name)
+        if attention_mask is not None:
+            _check_input_dtype(F.dtype(attention_mask), "attention_mask", [mstype.float32, mstype.float16],
+                               self.cls_name)
         if encoder_output is not None:
-            _check_shape_equal(F.shape(encoder_output), "encoder_output", self.cls_name,
-                               [[self.batch_size, self.src_seq_length, self.hidden_size],
-                                [self.batch_size * self.src_seq_length, self.hidden_size]])
             _check_input_dtype(F.dtype(encoder_output), "encoder_output",
                                [mstype.float32, mstype.float16], self.cls_name)
         if memory_mask is not None:
-            _check_shape_equal(F.shape(memory_mask), "memory_mask", self.cls_name,
-                               [self.batch_size, self.tgt_seq_length, self.src_seq_length])
             _check_input_dtype(F.dtype(memory_mask), "memory_mask",
                                [mstype.float32, mstype.float16], self.cls_name)
 
@@ -2179,9 +2203,7 @@ class TransformerDecoderLayer(Cell):
                                     batch_valid_length_is_tensor, batch_is_default)
 
         if self.use_past:
-            _check_shape_equal(F.shape(init_reset), "init_reset", self.cls_name, [1])
             _check_input_dtype(F.dtype(init_reset), "init_reset", [mstype.bool_], self.cls_name)
-            _check_shape_equal(F.shape(batch_valid_length), "batch_valid_length", self.cls_name, [self.batch_size])
             _check_input_dtype(F.dtype(batch_valid_length), "batch_valid_length", [mstype.int32], self.cls_name)
         return True
 
@@ -2240,7 +2262,9 @@ class TransformerEncoder(Cell):
         attention and feedforward layer.
 
         Args:
-            batch_size(int): The batch size of the input tensor.
+            batch_size(int): The batch size of the input tensor when do increnmental prediction. Should be a positive
+                value. When do training or prediction, the argument will not work and the user can just pass None to
+                the argument.
             num_layers(int): The layers of the `TransformerEncoderLayer`
             hidden_size(int): The hidden size of the input.
             ffn_hidden_size(int): The hidden size of bottleneck in the feedforward layer.
@@ -2248,9 +2272,12 @@ class TransformerEncoder(Cell):
             num_heads(int): The number of the heads.
             attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1.
             hidden_dropout_rate(float): The dropout rate of the final output of the layer. Default: 0.1.
-            hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
+            hidden_act (str, nn.Cell): The activation of the internal feedforward layer. Supports 'relu',
                 'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
-                'hsigmoid', 'logsigmoid' and so on. Default: gelu.
+                'hsigmoid', 'logsigmoid' and so on. User can provide custom activition to the argument.
+                If user wants to run the net in the parallel mode, the custom activation must also provide
+                the `activation_shard` function. Please see the examples of the
+                class:`mindspore.nn.transformer.FeedForward`. Default: gelu.
             post_layernorm_residual(bool): Do residuals adds before the layernorm. Default False.
             layernorm_compute_type(dtype.Number): The computation type of the layernorm.
                 Should be mstype.float32 or mstype.float16. Default mstype.float32.
@@ -2284,7 +2311,9 @@ class TransformerEncoder(Cell):
             - **hidden_states** (Tensor) - Tensor, shape should be [batch_size, seq_length, hidden_size] or
               [batch_size * seq_length, hidden_size], if the use_past is False or is_first_iteration=True. Otherwise,
               should be [batch_size, 1, hidden_size].
-            - **attention_mask** (Tensor) - Tensor, attention mask with shape [batch_size, seq_length, seq_length]
+            - **attention_mask** (Tensor) - Float Tensor, If the use_past is False or is_first_iteration=True,
+              the attention mask matrix should ba [batch_size, seq_length, seq_length], or None. None means there will
+              be no mask in softmax computation. Otherwise, should be [batch_size, 1, hidden_size]
             - **init_reset** (Tensor) - A bool tensor with shape [1], used to clear the past key parameter and
               past value parameter used in the incremental prediction. Only valid when use_past is True. Default True.
             - **batch_valid_length** (Tensor) - Int32 tensor with shape [batch_size] the past calculated the index.
@@ -2361,7 +2390,6 @@ class TransformerEncoder(Cell):
                                 offset=Validator.check_non_negative_int,
                                 attention_dropout_rate=Validator.check_non_negative_float,
                                 hidden_dropout_rate=Validator.check_non_negative_float,
-                                hidden_act=_valid_type_checks([str], "TransformerEncoder"),
                                 post_layernorm_residual=Validator.check_bool,
                                 layernorm_compute_type=_valid_value_checks([mstype.float32, mstype.float16],
                                                                            "TransformerEncoder"),
@@ -2396,7 +2424,7 @@ class TransformerEncoder(Cell):
         _check_moe_config(moe_config, parallel_config)
         self.use_moe = (moe_config.expert_num > 1)
         config_to_layer = parallel_config.moe_parallel_config if self.use_moe else parallel_config.dp_mp_config
-        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
+        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,):
             self.add = P.Add()
             self.aux_loss = Tensor(0.0, mstype.float32)
             self.num_layers = num_layers
@@ -2490,7 +2518,9 @@ class TransformerDecoder(Cell):
 
         Args:
             num_layers(int): The layers of the `TransformerDecoderLayer`.
-            batch_size(int): The batch size of the input tensor.
+            batch_size(int): The batch size of the input tensor when do increnmental prediction. Should be a positive
+                value. When do training or prediction, the argument will not work and the user can just pass None to
+                the argument.
             hidden_size(int): The hidden size of the input.
             ffn_hidden_size(int): The hidden size of bottleneck in the feedforward layer.
             src_seq_length(int): The input source sequence length.
@@ -2505,9 +2535,12 @@ class TransformerDecoder(Cell):
                 Should be mstype.float32 or mstype.float16. Default mstype.float32.
             param_init_type(dtype.Number): The parameter initialization type of the module.
                 Should be mstype.float32 or mstype.float16. Default mstype.float32.
-            hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
+            hidden_act (str, nn.Cell): The activation of the internal feedforward layer. Supports 'relu',
                 'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
-                'hsigmoid', 'logsigmoid' and so on. Default: gelu.
+                'hsigmoid', 'logsigmoid' and so on. User can provide custom activition to the argument.
+                If user wants to run the net in the parallel mode, the custom activation must also provide
+                the `activation_shard` function. Please see the examples of the
+                class:`mindspore.nn.transformer.FeedForward`. Default: gelu.
             lambda_func(function): A function can determine the fusion index,
                 pipeline stages and recompute attribute. If the
                 user wants to determine the pipeline stage and gradient aggregation fusion, the user can pass a
@@ -2528,13 +2561,14 @@ class TransformerDecoder(Cell):
             - **hidden_stats** (Tensor) - The input tensor with shape [batch_size, seq_length, hidden_size] or
               [batch_size * seq_length, hidden_size]
             - **attention_mask** (Tensor) - The attention mask for decoder with shape
-              [batch_size, seq_length, seq_length]
+              [batch_size, seq_length, seq_length] or None. None means there will be no mask in softmax
+              computation in self attention.
             - **encoder_output** (Tensor) - The output of the encoder with shape [batch_size, seq_length, hidden_size]
               or [batch_size * seq_length, hidden_size]. Note this args can not be passed by None when the net is in
               outermost layer. Default None.
             - **memory_mask** (Tensor) - The memory mask of the cross attention with shape [batch, tgt_seq_length,
-              src_seq_length] where tgt_seq_length is the length of the decoder. Note this args can not be passed by
-              None when the net is in outermost layer. Default None.
+              src_seq_length] where tgt_seq_length is the length of the decoder. The user can also pass None. None
+              means there will be no mask in softmax computation in cross attention. Default None.
             - **init_reset** (Tensor) - A bool tensor with shape [1], used to clear the past key parameter and
               past value parameter used in the incremental prediction. Only valid when use_past is True. Default True.
             - **batch_valid_length** (Tensor) - Int32 tensor with shape [batch_size] the past calculated the index.
@@ -2591,7 +2625,6 @@ class TransformerDecoder(Cell):
                                 offset=Validator.check_non_negative_int,
                                 attention_dropout_rate=Validator.check_non_negative_float,
                                 hidden_dropout_rate=Validator.check_non_negative_float,
-                                hidden_act=_valid_type_checks([str], "TransformerDecoder"),
                                 post_layernorm_residual=Validator.check_bool,
                                 layernorm_compute_type=_valid_value_checks([mstype.float32, mstype.float16],
                                                                            "TransformerDecoder"),
@@ -2627,7 +2660,7 @@ class TransformerDecoder(Cell):
         _check_config(parallel_config)
         self.use_moe = (moe_config.expert_num > 1)
         config_to_layer = parallel_config.moe_parallel_config if self.use_moe else parallel_config.dp_mp_config
-        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
+        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,):
             self.add = P.Add()
             self.aux_loss = Tensor(0.0, mstype.float32)
             self.num_layers = num_layers
@@ -2731,12 +2764,14 @@ class Transformer(Cell):
         the residual addition before the layer normalization. And the default hidden act is `gelu`.
         The details can be found in `Attention is all you need <https://arxiv.org/pdf/1706.03762v5.pdf>`_.
 
-        Note:
-            This is an experimental interface that is subject to change or deletion.
+        .. warning::
+            This is an experimental API that is subject to change or deletion.
 
         Args:
             hidden_size(int): The hidden size of the input.
-            batch_size(int): The batch size of the input tensor.
+            batch_size(int): The batch size of the input tensor when do increnmental prediction. Should be a positive
+                value. When do training or prediction, the argument will not work and the user can just pass None to
+                the argument.
             ffn_hidden_size(int): The hidden size of bottleneck in the feedforward layer.
             src_seq_length(int): The seq_length of the encoder's input tensor.
             tgt_seq_length(int): The seq_length of the decoder's input tensor.
@@ -2745,9 +2780,12 @@ class Transformer(Cell):
             num_heads(int): The number of the heads. Default: 2.
             attention_dropout_rate(float): The dropout rate of the attention scores. Default:0.1.
             hidden_dropout_rate(float): The dropout rate of the final output of the layer. Default:0.1.
-            hidden_act(str): The activation of the internal feedforward layer. Supports 'relu',
+            hidden_act (str, nn.Cell): The activation of the internal feedforward layer. Supports 'relu',
                 'relu6', 'tanh', 'gelu', 'fast_gelu', 'elu', 'sigmoid', 'prelu', 'leakyrelu', 'hswish',
-                'hsigmoid', 'logsigmoid' and so on. Default: gelu.
+                'hsigmoid', 'logsigmoid' and so on. User can provide custom activition to the argument.
+                If user wants to run the net in the parallel mode, the custom activation must also provide
+                the `activation_shard` function. Please see the examples of the
+                class:`mindspore.nn.transformer.FeedForward`. Default: gelu.
             post_layernorm_residual(bool): Do residuals adds before the layernorm. Default False.
             layernorm_compute_type(dtype.Number): The computation type of the layernorm.
                 Should be dtype.float32 or dtype.float16. Default dtype.float32.
@@ -2772,15 +2810,17 @@ class Transformer(Cell):
             - **encoder_inputs** (Tensor) - The input tensor with shape [batch_size, seq_length, hidden_size] or
               [batch_size * seq_length, hidden_size].
             - **encoder_masks** (Tensor) - The attention mask for decoder with shape
-              [batch_size, seq_length, seq_length].
+              [batch_size, seq_length, seq_length] or None. None means there will be no mask in softmax computation
+              in self attention of the encoder module.
             - **decoder_inputs** (Tensor) - The output of the encoder with shape [batch_size, seq_length, hidden_size]
               or [batch_size * seq_length, hidden_size], this should be none if the decoder layer is 0.
             - **decoder_masks** (Tensor) - The attention mask for decoder with shape
-              [batch_size, seq_length, seq_length]
+              [batch_size, seq_length, seq_length] or None. None means there will be no mask in softmax computation
+              in self attention of the decoder module.
             - **memory_mask** (Tensor) - The memory mask of the cross attention with shape [batch, tgt_seq_length,
               src_seq_length]
               where tgt_seq_length is the length of the decoder. The output of the encoder with shape [batch_size,
-              seq_length, hidden_size], this should be none if the decoder layer is 0.
+              seq_length, hidden_size], this should be none if the decoder layer is 0 or the user wants no mask.
             - **init_reset** (Tensor) - A bool tensor with shape [1], used to clear the past key parameter and
               past value parameter used in the incremental prediction. Only valid when use_past is True. Default True.
             - **batch_valid_length** (Tensor) - Int32 tensor with shape [batch_size] the past calculated the index.
@@ -2854,7 +2894,6 @@ class Transformer(Cell):
                                 tgt_seq_length=Validator.check_positive_int,
                                 attention_dropout_rate=Validator.check_non_negative_float,
                                 hidden_dropout_rate=Validator.check_non_negative_float,
-                                hidden_act=_valid_type_checks([str], "Transformer"),
                                 post_layernorm_residual=Validator.check_bool,
                                 layernorm_compute_type=_valid_value_checks([mstype.float32, mstype.float16],
                                                                            "Transformer"),
@@ -2884,7 +2923,7 @@ class Transformer(Cell):
                  moe_config=default_moe_config,
                  parallel_config=default_transformer_config):
         super(Transformer, self).__init__()
-        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
+        if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,):
             _check_config(parallel_config)
             self.batch_size = batch_size
             self.hidden_size = hidden_size