mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/parallel/_tensor.py

@@ -17,7 +17,6 @@ from __future__ import division
 from __future__ import absolute_import
 
 import numpy as np
-
 from mindspore.common.tensor import Tensor
 from mindspore.communication.management import get_rank, get_group_size
 from mindspore._c_expression import TensorTransform
@@ -41,7 +40,7 @@ def _get_tensor_strategy(dev_mat, tensor_map):
         if dim == -1:
             tensor_strategy.append(1)
         else:
-            tensor_strategy.append(dev_mat[-dim-1])
+            tensor_strategy.append(dev_mat[-dim - 1])
     return tensor_strategy
 
 
@@ -198,7 +197,7 @@ def _get_slice_index(dev_mat, tensor_map, opt_shard_group):
     return tensor_slice_index
 
 
-def _load_tensor(tensor, dev_mat, tensor_map):
+def _load_tensor(tensor, dev_mat, tensor_map, rank_id=-1):
     """
     Get the tensor slice of the local device by the device matrix and the tensor map
 
@@ -216,7 +215,10 @@ def _load_tensor(tensor, dev_mat, tensor_map):
         >>> tensor_map = [1, -1]
         >>> tensor_slice = _load_tensor(tensor, dev_mat, tensor_map)
     """
-    rank = get_rank()
+    if rank_id == -1:
+        rank = get_rank()
+    else:
+        rank = rank_id
     tensor_strategy = _get_tensor_strategy(dev_mat, tensor_map)
     tensor_slice_index = _get_tensor_slice_index(dev_mat, tensor_strategy, tensor_map, rank)
     np_tensor = tensor.asnumpy()
@@ -225,7 +227,7 @@ def _load_tensor(tensor, dev_mat, tensor_map):
     return np_tensor_slice
 
 
-def _load_tensor_by_layout(tensor, layout):
+def _load_tensor_by_layout(tensor, layout, rank_id):
     """
     Load tensor by layout.
 
@@ -246,13 +248,13 @@ def _load_tensor_by_layout(tensor, layout):
         raise ValueError("The length of layout must be larger than 5! layout is {}".format(layout))
     dev_mat = layout[0]
     tensor_map = layout[1]
+    if not tensor_map:
+        return tensor
     uniform_split = layout[4]
     group = layout[5]
     if uniform_split == 0:
         raise RuntimeError("The load tensor only support uniform split now")
-    if tensor.size == 1:
-        return tensor
-    tensor_slice = _load_tensor(tensor, dev_mat, tensor_map)
+    tensor_slice = _load_tensor(tensor, dev_mat, tensor_map, rank_id)
     if group:
         # get a totally shard tensor slice for parallel optimizer
         rank = get_rank(group)
@@ -315,7 +317,6 @@ def _reshape_param_data(param_data, dev_mat, tensor_map):
     return Tensor(tensor_slices_new[0])
 
 
-
 def _extract_layout_item(layout_item):
     dev_matrix = layout_item[0]
     tensor_map = layout_item[1]
@@ -541,6 +542,7 @@ def _check_operator(operator):
 
 def _apply_operator(operator_name):
     """apply transform operator"""
+
     def _apply_reshape_operator(numpy_data, reshape_op):
         """
         Apply reshape operator.
@@ -597,8 +599,8 @@ def _apply_operator(operator_name):
             raise ValueError("The slice operator information is wrong.")
         shape_size = len(slice_op[1]) // 3
         begin = slice_op[1][:shape_size]
-        end = slice_op[1][shape_size:shape_size*2]
-        stride = slice_op[1][shape_size*2:]
+        end = slice_op[1][shape_size:shape_size * 2]
+        stride = slice_op[1][shape_size * 2:]
         slice_index = []
         for begin_i, end_i, strides_i in zip(begin, end, stride):
             s = slice(begin_i, end_i, strides_i)
@@ -637,8 +639,8 @@ def _reshape_param_data_with_weight(param_data, dev_mat, field_size):
     for i in range(len(tensor_slices[0][0])):
         tensor_slices_new = np.array(tensor_slices[0][:, i]).reshape(field_size, -1)
         for j in range(1, device_count):
-            tensor_slices_new = np.concatenate((tensor_slices_new,\
-                                   np.array(tensor_slices[j][:, i]).reshape(field_size, -1)), axis=1)
+            tensor_slices_new = np.concatenate((tensor_slices_new, \
+                                                np.array(tensor_slices[j][:, i]).reshape(field_size, -1)), axis=1)
         tensor_slices_col.append(tensor_slices_new)
     new_tensor = np.array(tensor_slices_col[0]).reshape(-1, 1)
     for i in range(1, len(tensor_slices_col)):

mindspore/parallel/_transformer/layers.py

@@ -366,14 +366,14 @@ class _Linear(Cell):
             is same as `x`. The values of str refer to the function `initializer`. Default: 'normal'.
         bias_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable bias_init parameter. The dtype is
             same as `x`. The values of str refer to the function `initializer`. Default: 'zeros'.
-        has_bias (bool): Specifies whether the layer uses a bias vector. Default: True.
+        has_bias (bool): Specifies whether the layer uses a bias vector. Default: ``True``.
         activation (str): activate function applied to the output of the fully connected layer,
-            eg. 'ReLU'. Default: None.
+            eg. 'ReLU'. Default: ``None``.
         expert_num (int): The number of experts used in this Linear. Here, for the case expert_num > 1, BatchMatMul is
             used and the first dimension in BatchMatMul indicate expert_num. Default: 1.
         outer_batch (int): The replication number of experts. The replication is effective only when MoE is applied.
             Default: 1.
-        expert_group_size (int): The number of tokens in each data parallel group. Default: None.
+        expert_group_size (int): The number of tokens in each data parallel group. Default: ``None``.
         compute_dtype (dtype.Number): The computation type. Default: mstype.float16
     Inputs:
         - **x** (Tensor) - Tensor of shape :math:`(*, in\_channels)`. The `in_channels` in `Args` should be equal
@@ -424,9 +424,11 @@ class _Linear(Cell):
         self.out_channels = out_channels
         if not (isinstance(activation, str) or activation is None or issubclass(activation, nn.Cell)):
             raise TypeError(f"For Linear cell, the activation should str type or nn.Cell type, but got {activation}.")
-        if isinstance(weight_init, Tensor) and (weight_init.ndim != 2 or weight_init.shape[0] != out_channels or
-                                                weight_init.shape[1] != in_channels):
-            raise ValueError("The shape of parameter 'weight_init' is error, please check shape of 'weight_init'.")
+
+        if isinstance(weight_init, Tensor):
+            if weight_init.ndim != 2 or weight_init.shape[0] != out_channels \
+                                or weight_init.shape[1] != in_channels:
+                raise ValueError("The shape of parameter 'weight_init' is error, please check shape of 'weight_init'.")
         weight_shape = [out_channels, in_channels] if transpose_b else [in_channels, out_channels]
         self.expert_num = expert_num
         self.outer_batch = outer_batch

mindspore/parallel/_transformer/loss.py

@@ -139,6 +139,7 @@ class _NLLLoss(Cell):
             self.add = P.Add().shard(((dp, mp), ()))
 
     def construct(self, softmax_result, one_hot_label):
+        """The forward of _NLLLoss"""
         log_softmax_result = self.log(self.add(softmax_result, self.eps_const))
         loss = self.mul(log_softmax_result, one_hot_label)
         loss_unsum = self.neg(loss)

mindspore/parallel/_transformer/moe.py

@@ -51,13 +51,13 @@ class MoEConfig:
                 router) to be added to the entire model loss, which is < 1.0. Default: 0.05.
             num_experts_chosen (int): The number of experts is chosen by each token and it should not be larger
                 than expert_num. Default: 1.
-            expert_group_size (int): The number of tokens in each data parallel group. Default: None. This parameter is
-                effective only when in AUTO_PARALLEL mode, and NOT SHARDING_PROPAGATION.
+            expert_group_size (int): The number of tokens in each data parallel group. Default: ``None``.
+                This parameter is effective only when in AUTO_PARALLEL mode, and NOT SHARDING_PROPAGATION.
             group_wise_a2a (bool): Whether to enable group-wise alltoall communication, which can reduce communication
-                time by converting part of inter communication into intra communication. Default: False. This parameter
-                is effective only when model parallel > 1 and data_parallel equal to expert parallel.
+                time by converting part of inter communication into intra communication. Default: ``False``.
+                This parameter is effective only when model parallel > 1 and data_parallel equal to expert parallel.
             comp_comm_parallel (bool): Whether to enable ffn compute and communication parallel, which can reduce pure
-                communicattion time by splitting and overlapping compute and communication. Default: False.
+                communicattion time by splitting and overlapping compute and communication. Default: ``False``.
             comp_comm_parallel_degree (int): The split number of compute and communication. The larger the numbers,
                 the more overlap there will be but will consume more memory. Default: 2. This parameter is effective
                 only when comp_comm_parallel enable.
@@ -273,7 +273,7 @@ class MoE(Cell):
         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)
+                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),
@@ -330,7 +330,7 @@ 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)
+            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),

mindspore/parallel/_transformer/op_parallel_config.py

@@ -147,10 +147,11 @@ class _PipeLineConfig(_Config):
             >>> config=_PipeLineConfig(pipeline_stage=1, micro_batch_num=1)
     """
 
-    def __init__(self, pipeline_stage=1, micro_batch_num=1):
+    def __init__(self, pipeline_stage=1, micro_batch_num=1, pipeline_segment=1):
         Validator.check_positive_int(pipeline_stage, "pipeline_stage")
         Validator.check_positive_int(micro_batch_num, "micro_batch_num")
         self.pipeline_stage = pipeline_stage
+        self.pipeline_segment = pipeline_segment
         self.micro_batch_num = micro_batch_num
 
     @property
@@ -163,6 +164,16 @@ class _PipeLineConfig(_Config):
         self._pipeline_stage = value
         context.set_auto_parallel_context(pipeline_stages=value)
 
+    @property
+    def pipeline_segment(self):
+        return self._pipeline_segment
+
+    @pipeline_segment.setter
+    def pipeline_segment(self, value):
+        Validator.check_positive_int(value, "pipeline_segment")
+        self._pipeline_segment = value
+        context.set_auto_parallel_context(pipeline_segments=value)
+
     @property
     def micro_batch_num(self):
         return self._micro_batch_num

mindspore/parallel/_transformer/transformer.py

@@ -70,7 +70,7 @@ class EmbeddingOpParallelConfig(_Config):
             vocab_emb_dp(bool): Shard embedding in model parallel or data parallel. If True, the embedding lookup
                 will be a data parallel style training and model_parallel value will be ignored.  If false, the
                 embedding table will be sharded into n parts at the 0-th dimension row slice of the embedding table,
-                where the n is the model parallel way determined by this parameter. Default: True
+                where the n is the model parallel way determined by this parameter. Default: ``True``
 
         Supported Platforms:
             ``Ascend`` ``GPU``
@@ -120,13 +120,13 @@ class TransformerRecomputeConfig(_Config):
         TransformerRecomputeConfig for the setting recompute attributes for encoder/decoder layers.
 
         Args:
-            recompute (bool): Enable recomputation of the transformer block or not. Default: False.
+            recompute (bool): Enable recomputation of the transformer block or not. Default: ``False``.
             parallel_optimizer_comm_recompute (bool): Specifies whether the communication operator allgathers
                 introduced by optimizer shard are recomputed in auto parallel or semi auto parallel mode.
-                Default: False.
+                Default: ``False``.
             mp_comm_recompute (bool): Specifies whether the model parallel communication operators
-                in the cell are recomputed in auto parallel or semi auto parallel mode. Default: True.
-            recompute_slice_activation (bool): Slice the cell output which would remains in memory. Default: False.
+                in the cell are recomputed in auto parallel or semi auto parallel mode. Default: ``True``.
+            recompute_slice_activation (bool): Slice the cell output which would remains in memory. Default: ``False``.
 
         Supported Platforms:
             ``Ascend`` ``GPU``
@@ -214,7 +214,7 @@ class TransformerOpParallelConfig(_Config):
             gradient_aggregation_group (int): The fusion group size of the optimizer state sharding. Default: 4.
             recompute (Union[TransformerRecomputeConfig, bool]): The configuration of recomputation for
                 the transformer block. Default: An instance of TransformerRecomputeConfig with default values.
-            vocab_emb_dp (bool): Shard embedding in model parallel or data parallel. Default: True.
+            vocab_emb_dp (bool): Shard embedding in model parallel or data parallel. Default: ``True``.
 
         Supported Platforms:
             ``Ascend`` ``GPU``
@@ -226,7 +226,8 @@ class TransformerOpParallelConfig(_Config):
             >>> config=TransformerOpParallelConfig(data_parallel=1, model_parallel=1, recompute=recompute_config)
     """
 
-    def __init__(self, data_parallel=1, model_parallel=1, expert_parallel=1, pipeline_stage=1, micro_batch_num=1,
+    def __init__(self, data_parallel=1, model_parallel=1, expert_parallel=1, pipeline_stage=1, pipeline_segment=1,
+                 micro_batch_num=1,
                  recompute=default_transformer_recompute_config,
                  optimizer_shard=False, gradient_aggregation_group=4, vocab_emb_dp=True):
         self.recompute = recompute
@@ -234,7 +235,8 @@ class TransformerOpParallelConfig(_Config):
         self.gradient_aggregation_group = gradient_aggregation_group
         self._embed_dp_mp_config = EmbeddingOpParallelConfig(data_parallel=data_parallel, model_parallel=model_parallel,
                                                              vocab_emb_dp=vocab_emb_dp)
-        self._pp_config = _PipeLineConfig(pipeline_stage=pipeline_stage, micro_batch_num=micro_batch_num)
+        self._pp_config = _PipeLineConfig(pipeline_stage=pipeline_stage, micro_batch_num=micro_batch_num,
+                                          pipeline_segment=pipeline_segment)
         self._moe_config = MoEParallelConfig(data_parallel=data_parallel, model_parallel=model_parallel,
                                              expert_parallel=expert_parallel)
 
@@ -309,6 +311,14 @@ class TransformerOpParallelConfig(_Config):
     def pipeline_stage(self, value):
         self._pp_config.pipeline_stage = value
 
+    @property
+    def pipeline_segment(self):
+        return self._pp_config.pipeline_segment
+
+    @pipeline_segment.setter
+    def pipeline_segment(self, value):
+        self._pp_config.pipeline_segment = value
+
     @property
     def optimizer_shard(self):
         return self._optimizer_shard
@@ -359,8 +369,8 @@ class FeedForward(Cell):
                 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
-                effective only when in AUTO_PARALLEL mode, and NOT SHARDING_PROPAGATION.
+            expert_group_size (int): The number of tokens in each data parallel group. Default: ``None``.
+                This parameter is effective only when in AUTO_PARALLEL mode, and NOT SHARDING_PROPAGATION.
             param_init_type (dtype.Number): The parameter initialization type. Should be mstype.float32 or
                 mstype.float16. Default: mstype.float32.
             parallel_config (OpParallelConfig, MoEParallelConfig): The config of parallel setting, see
@@ -429,6 +439,7 @@ class FeedForward(Cell):
             >>> 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,
@@ -622,6 +633,7 @@ class AttentionMask(Cell):
               [1. 1. 1. 0]
               [0. 0. 0. 0]]]
     """
+
     @_LogActionOnce(logger=logger, key='AttentionMask',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(seq_length=Validator.check_positive_int,
@@ -710,6 +722,7 @@ class VocabEmbedding(Cell):
             >>> print(table.shape)
             (30, 30)
     """
+
     @_LogActionOnce(logger=logger, key='VocabEmbedding',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(vocab_size=Validator.check_positive_int,
@@ -866,6 +879,7 @@ class MultiHeadAttention(Cell):
             >>> print(past[1].shape)
             (2, 3, 20, 5)
     """
+
     @_LogActionOnce(logger=logger, key='MultiHeadAttention',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
@@ -1203,7 +1217,8 @@ class MultiHeadAttention(Cell):
     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)):
+        incr_infer = self.use_past and self.is_first_iteration
+        if len(F.shape(query)) == 2 and ((incr_infer) or (not self.use_past)):
             return F.shape(query)[0] // self.src_seq_length
         return F.shape(query)[0]
 
@@ -1459,6 +1474,7 @@ class TransformerEncoderLayer(Cell):
             >>> print(past[1].shape)
             (2, 2, 16, 4)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerEncoderLayer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
@@ -1848,6 +1864,7 @@ class TransformerDecoderLayer(Cell):
             >>> print(past[3].shape)
             (2, 2, 20, 32)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerDecoderLayer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(hidden_size=Validator.check_positive_int,
@@ -2292,7 +2309,7 @@ class TransformerEncoder(Cell):
                 represents the transformer block, `layer_id(int)` means the layer index for the current module, counts
                 from zero, `offset(int)` means the layer_index needs an offset, if there are other modules in the net.
                 The default setting for the pipeline is: `(layer_id + offset) // (layers / pipeline_stage)`.
-                Default: None.
+                Default: ``None``.
             offset(int): The initial layer index for the `encoder`. Used for setting the fusion id and stage id, to not
                 overlap with the encoder layer. Default 0.
             use_past(bool): Use the past state to compute, used for incremental prediction. For example, if we have two
@@ -2301,7 +2318,7 @@ class TransformerEncoder(Cell):
                 In the first step, set the is_first_iteration to be True by
                 `model.add_flags_recursive(is_first_iteration=True)`, and pass the full inputs. Then, set the
                 is_first_iteration to be False by `model.add_flags_recursive(is_first_iteration=False)`. At this moment,
-                pass the single step's input tensor, and loop it. Default: False.
+                pass the single step's input tensor, and loop it. Default: ``False``.
             moe_config(MoEConfig): The configuration of MoE (Mixture of Expert). Default is an instance of MoEConfig
                 with default values. Please see `MoEConfig`.
             parallel_config(TransformerOpParallelConfig): The parallel configure. Default `default_transformer_config`,
@@ -2379,6 +2396,7 @@ class TransformerEncoder(Cell):
             >>> print(past[0][1].shape)
             (2, 2, 16, 4)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerEncoder',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(batch_size=Validator.check_positive_int,
@@ -2548,7 +2566,7 @@ class TransformerDecoder(Cell):
                 represents the transformer block, `layer_id(int)` means the layer index for the current module, counts
                 from zero, `offset(int)` means the layer_index needs an offset, if there are other modules in the net.
                 The default setting for the pipeline is: `(layer_id + offset) // (layers / pipeline_stage)`.
-                Default: None.
+                Default: ``None``.
             use_past(bool): Use the past state to compute, used for incremental prediction. Default False.
             offset(int): The initial layer index for the `decoder`. Used for setting the fusion id and stage id, to not
                 overlap with the encoder layer. Default 0.
@@ -2613,6 +2631,7 @@ class TransformerDecoder(Cell):
             >>> print(past[0][3].shape)
             (2, 2, 20, 32)
     """
+
     @_LogActionOnce(logger=logger, key='TransformerDecoder',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(batch_size=Validator.check_positive_int,
@@ -2882,6 +2901,7 @@ class Transformer(Cell):
             >>> print(de_past[0][3].shape)
             (2, 2, 20, 32)
     """
+
     @_LogActionOnce(logger=logger, key='Transformer',
                     no_warning=_get_parallel_mode() in (ParallelMode.STAND_ALONE,))
     @_args_type_validator_check(batch_size=Validator.check_positive_int,

mindspore/parallel/_utils.py

@@ -96,11 +96,21 @@ def _slice_parameter(parameter, phase, layout):
     new_param = parameter.init_data(layout, set_sliced=True)
     parameter = new_param
     graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
-    if not parameter.sliced and (layout is not None):
-        new_tensor = _load_tensor_by_layout(parameter, layout)
+    if layout is None:
+        parameter.sliced = True
+        return
+    if not parameter.sliced:
+        rank = get_rank()
+        new_tensor = _load_tensor_by_layout(parameter, layout, rank)
         parameter.set_data(new_tensor, True)
 
 
+def _slice_tensor(tensor, layout, rank_id):
+    """Slice python tensor obj according to the layout."""
+    new_tensor = _load_tensor_by_layout(tensor, layout, rank_id)
+    return new_tensor
+
+
 def _init_optimizer_state(parameter, phase):
     """init optimizer state"""
     if not parameter.has_init:
@@ -127,14 +137,17 @@ def _to_full_shapes(shapes, device_num):
                                  "dataset strategy item size {}".format(len(shape), len(dataset_strategy[index])))
             new_shape = ()
             for i, item in enumerate(shape):
-                new_shape += (item * dataset_strategy[index][i],)
+                if item > 0:
+                    new_shape += (item * dataset_strategy[index][i],)  # static shape
+                else:
+                    new_shape += (item,)  # dynamic shape
             new_shapes.append(new_shape)
         return new_shapes
     for shape in shapes:
         new_shape = ()
         for i, item in enumerate(shape):
-            if i == 0:
-                new_shape += (item * device_num,)
+            if i == 0 and item > 0:
+                new_shape += (item * device_num,)  # only for static shape
             else:
                 new_shape += (item,)
         new_shapes.append(new_shape)
@@ -192,7 +205,7 @@ def _to_full_tensor(elem, global_device_num, global_rank, scaling_sens=None):
                 slice_index += (s,)
             new_tensor_numpy = np.zeros(new_shape, dtype_to_nptype(type_))
             new_tensor_numpy[slice_index] = data.asnumpy()
-        new_tensor = Tensor(new_tensor_numpy)
+        new_tensor = Tensor(new_tensor_numpy, dtype=type_)
         lst.append(new_tensor)
     if scaling_sens:
         lst.append(Tensor(scaling_sens, mstype.float32))
@@ -325,7 +338,7 @@ def _parallel_predict_check():
         dataset_strategy = context.get_auto_parallel_context("dataset_strategy")
         is_shard_dataset_mp = (dataset_strategy and dataset_strategy not in ("data_parallel", "full_batch"))
         if not context.get_auto_parallel_context("full_batch") and not is_shard_dataset_mp:
-            raise RuntimeError('Model prediction only supports full batch dataset. Please set "full_batch" with True.')
+            logger.warning('Using non full-batch dataset in model prediction may lead to incorrect data.')
 
 
 def _check_similar_layout(tensor_layout1, tensor_layout2):
@@ -335,7 +348,7 @@ def _check_similar_layout(tensor_layout1, tensor_layout2):
     for i in tensor_layout1[1]:
         if i == -1:
             continue
-        if tensor_layout1[0][-1-i] != tensor_layout2[0][-1-i]:
+        if tensor_layout1[0][-1 - i] != tensor_layout2[0][-1 - i]:
             return False
     return True
 
@@ -353,7 +366,7 @@ def _remove_repeated_slices(tensor_layout):
     tensor_map = tensor_layout[1]
     for dim in range(len(dev_mat)):
         if dim not in tensor_map:
-            dev_mat[-1-dim] = 1
+            dev_mat[-1 - dim] = 1
     new_tensor_layout[0] = dev_mat
     return new_tensor_layout
 
@@ -409,15 +422,24 @@ def _grads_divided_by_device_num_if_recomputation(grads):
     """
     If in pynative parallel and full_batch is True, divide grads by device num to ensure that the gradients is correct.
     """
-    if not grads or not _is_pynative_parallel() or not _get_full_batch():
+    if not _is_pynative_parallel() or not _get_full_batch():
         return grads
 
-    device_num = Tensor(_get_device_num(), grads[0].dtype)
+    device_num = _get_device_num()
     logger.info(f"In PyNative mode, when parallel mode is in "
                 f"({context.ParallelMode.SEMI_AUTO_PARALLEL}, {context.ParallelMode.AUTO_PARALLEL}) and "
                 f"full_batch is Ture, the gradients will be automatically divided by device_num({device_num}).")
-    new_grads = ()
-    for grad in grads:
-        new_grads += (grad / device_num,)
 
+    if not isinstance(grads, (tuple, Tensor)):
+        raise ValueError(f"The type of grads must be either Tuple[Tensor] or Tensor, but got {type(grads)}.")
+
+    if isinstance(grads, tuple):
+        new_grads = ()
+        if grads:
+            device_num_tensor = Tensor(device_num, grads[0].dtype)
+            for grad in grads:
+                new_grads += (grad / device_num_tensor,)
+    else:
+        device_num_tensor = Tensor(device_num, grads.dtype)
+        new_grads = grads / device_num_tensor
     return new_grads