mindspore 2.6.0rc1__cp311-cp311-win_amd64.whl → 2.7.0__cp311-cp311-win_amd64.whl

mindspore/parallel/_ps_context.py

@@ -65,7 +65,6 @@ _set_ps_context_func_map = {
     "enable_ssl": ps_context().set_enable_ssl,
     "client_password": ps_context().set_client_password,
     "server_password": ps_context().set_server_password,
-    "scheduler_manage_port": ps_context().set_scheduler_manage_port,
     "config_file_path": ps_context().set_config_file_path,
 }
 
@@ -78,12 +77,10 @@ _get_ps_context_func_map = {
     "scheduler_ip": ps_context().scheduler_ip,
     "scheduler_port": ps_context().scheduler_port,
     "enable_ssl": ps_context().enable_ssl,
-    "client_password": ps_context().client_password,
-    "server_password": ps_context().server_password,
-    "scheduler_manage_port": ps_context().scheduler_manage_port,
     "config_file_path": ps_context().config_file_path,
 }
 
+
 def _get_ps_mode_rank():
     ps_rank = ps_context().ps_rank_id()
     if ps_rank == -1:
@@ -116,8 +113,9 @@ def _set_ps_context(**kwargs):
                           Only after enable_ps is set True, the environment variables will be effective.
                           Default: ``False``.
         config_file_path (str): Configuration file path used by recovery. Default: ''.
-        scheduler_manage_port (int): scheduler manage port used to scale out/in. Default: 11202.
         enable_ssl (bool): Set PS SSL mode enabled or disabled. Default: ``False``.
+                           When set to False, users need to review and confirm the security of network environment
+                           where the distributed job is located.
         client_password (str): Password to decrypt the secret key stored in the client certificate. Default: ''.
         server_password (str): Password to decrypt the secret key stored in the server certificate. Default: ''.
 
@@ -126,7 +124,7 @@ def _set_ps_context(**kwargs):
 
     Examples:
         >>> import mindspore as ms
-        >>> ms.set_ps_context(enable_ps=True, enable_ssl=True, client_password='123456', server_password='123456')
+        >>> ms.set_ps_context(enable_ps=True, enable_ssl=True, client_password='', server_password='')
     """
     for key, value in kwargs.items():
         if key not in _set_ps_context_func_map:

mindspore/parallel/_tensor.py

@@ -21,9 +21,11 @@ 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
+from mindspore import log as logger
 
 _tensor_transform = TensorTransform.get_instance()
-
+COMM_TENSOR_CELL_CACHE = {}
+RESHARD_OP_MAP_CACHE = {}
 
 def _get_tensor_strategy(dev_mat, tensor_map):
     """
@@ -348,7 +350,7 @@ def _extract_layout_item(layout_item):
     return dev_matrix, tensor_map, opt_shard_step, opt_shard_size
 
 
-def _transform_tensor_by_layout(from_layout, to_layout, device_list, rank_id):
+def _transform_tensor_by_layout(from_layout, to_layout, device_list, rank_id, enable_redist_opt=False):
     """
     Transform tensor from source layout to the destination layout.
 
@@ -362,7 +364,7 @@ def _transform_tensor_by_layout(from_layout, to_layout, device_list, rank_id):
     """
     if not isinstance(from_layout, tuple) or not isinstance(to_layout, tuple):
         raise TypeError("The layout should be tuple! layout is {} and {}".format(from_layout, to_layout))
-    return _tensor_transform.transform_tensor_sharding(from_layout, to_layout, device_list, rank_id)
+    return _tensor_transform.transform_tensor_sharding(from_layout, to_layout, device_list, enable_redist_opt, rank_id)
 
 
 def _construct_from_to_tensor_layout(from_full_tensor_shape, from_dev_matrix,
@@ -587,13 +589,15 @@ def _get_needed_rank_list_by_layouts(from_tensor_layout, to_tensor_layout, devic
     return result_list
 
 
-def _get_needed_rank_transform_operator_map_by_layouts(from_tensor_layout, to_tensor_layout, device_list, self_rank):
+def _get_needed_rank_transform_operator_map_by_layouts(from_tensor_layout, to_tensor_layout, device_list, self_rank,
+                                                       enable_redist_opt=False):
     """
     AllGather op: {op_name, group_ranks + axis}
     """
     stack = []
     index = 0
-    transform_operators = _transform_tensor_by_layout(from_tensor_layout, to_tensor_layout, device_list, self_rank)
+    transform_operators = _transform_tensor_by_layout(from_tensor_layout, to_tensor_layout, device_list, self_rank,
+                                                      enable_redist_opt)
     result_map = {self_rank: transform_operators}
     for operators in transform_operators:
         op_name = operators[0]
@@ -606,7 +610,7 @@ def _get_needed_rank_transform_operator_map_by_layouts(from_tensor_layout, to_te
         for rank in group_info[1]:
             if rank not in result_map:
                 new_transform_operators = _transform_tensor_by_layout(from_tensor_layout, to_tensor_layout,
-                                                                      device_list, rank)
+                                                                      device_list, rank, enable_redist_opt)
                 result_map[rank] = new_transform_operators
                 index = 0
                 for operators in new_transform_operators:
@@ -710,8 +714,6 @@ def _apply_operator(operator_name):
         Returns:
             The data of tensor after apply operator.
         """
-        if str(type(numpy_data)) == "<class 'builtins.PySafeSlice'>":
-            numpy_data = numpy_data[:]
         if not isinstance(numpy_data, np.ndarray):
             raise TypeError("The data should be a numpy.ndarray.")
         _check_operator(reshape_op)
@@ -732,10 +734,7 @@ def _apply_operator(operator_name):
             raise TypeError("The data_list should be a list.")
         new_numpy_data_list = []
         for numpy_data in numpy_data_list:
-            if str(type(numpy_data)) == "<class 'builtins.PySafeSlice'>":
-                new_numpy_data_list.append(numpy_data[:])
-            else:
-                new_numpy_data_list.append(numpy_data)
+            new_numpy_data_list.append(numpy_data)
         numpy_data_list = new_numpy_data_list
         _check_operator(allgather_op)
         concat_group = allgather_op[1][:-1]
@@ -896,3 +895,159 @@ def _chunk_shape(np_tensor, strategy, depth):
         output.extend(
             _chunk_shape(ret_, strategy[len(strategy) - depth + 1:len(strategy)], depth - 1))
     return output
+
+
+def _infer_pp_op_map(from_layout, to_layout, self_rank):
+    """
+    get the ops map for merging pp stages
+    """
+    from_rank_list = from_layout[3]
+    to_rank_list = to_layout[3]
+    from_dev_num_in_stage = len(from_rank_list)
+    current_rank_stage_id = self_rank // from_dev_num_in_stage
+    diff_rank_id = [
+        rank_id for rank_id in to_rank_list if rank_id not in from_rank_list]
+    end_stage = from_dev_num_in_stage * (current_rank_stage_id + 1)
+    start_stage = from_dev_num_in_stage * current_rank_stage_id
+    rank_pos_in_stage = list(range(start_stage, end_stage)).index(self_rank)
+    root_idx = from_rank_list[rank_pos_in_stage]
+    broadcast_rank_list = [root_idx]
+    while rank_pos_in_stage < len(diff_rank_id):
+        broadcast_rank_list.append(diff_rank_id[rank_pos_in_stage])
+        rank_pos_in_stage += from_dev_num_in_stage
+    broadcast_rank_list.sort()
+    broadcast_map = {rank_id: [('Broadcast', root_idx, broadcast_rank_list)] for rank_id in broadcast_rank_list}
+    return broadcast_map
+
+
+def _get_pipeline_operator_map(from_layout, to_layout, self_rank):
+    """
+    If src_pp_stages is greater than dst_pp_stages, the weights of the corresponding cards need to
+    be communicated via broadcast to swap. Need to communicate src rank0's 01 to src rank2,
+    so that rank2 holds param0's data. Similarly, communicate rank1's 02 to rank3
+    rank0 01           01 11
+    rank1 02           02 12
+    pp2 ------->  pp1
+    rank2 11           03 13
+    rank3 12           04 14
+
+    Args:
+        from_layout (tuple): Use tuple to present layout
+          (device_matrix(list), tensor_map(list), global_shape(list), rank_list(list))
+        to_layout (tuple): Use tuple to present layout
+          (device_matrix(list), tensor_map(list), global_shape(list), rank_list(list))
+        self_rank (int): rank_id
+    """
+    if len(from_layout[3]) < len(to_layout[3]):
+        logger.debug(f"from {from_layout} to {to_layout} need to broadcast data across pp stages")
+        comm_tensor_cache_key = (
+            f"{from_layout[0]}, {from_layout[1]}, {from_layout[2]}, {from_layout[3]}"
+            f" -> "
+            f"{to_layout[0]}, {to_layout[1]}, {from_layout[2]}, {to_layout[3]}")
+        global COMM_TENSOR_CELL_CACHE
+        if comm_tensor_cache_key not in COMM_TENSOR_CELL_CACHE:
+            logger.debug(f"comm_tensor_cache_key is {comm_tensor_cache_key}, not match cache")
+            broadcast_map = _infer_pp_op_map(from_layout, to_layout, self_rank)
+            broadcast_op_map_dict = {rank_id: broadcast_map for rank_id in broadcast_map.keys()}
+            COMM_TENSOR_CELL_CACHE[comm_tensor_cache_key] = broadcast_op_map_dict
+        else:
+            comm_tensor_cache_key_rank_list = COMM_TENSOR_CELL_CACHE[comm_tensor_cache_key]
+            if self_rank in comm_tensor_cache_key_rank_list:
+                logger.debug(f"comm_tensor_cache_key is {comm_tensor_cache_key}, match cache")
+                broadcast_map = comm_tensor_cache_key_rank_list[self_rank]
+            else:
+                logger.debug(f"comm_tensor_cache_key is {comm_tensor_cache_key}, but rank {self_rank} not match cache")
+                broadcast_map = _infer_pp_op_map(from_layout, to_layout, self_rank)
+                for rank_id in broadcast_map.keys():
+                    COMM_TENSOR_CELL_CACHE[comm_tensor_cache_key][rank_id] = broadcast_map
+        return broadcast_map
+    logger.debug(f"from {from_layout} to {to_layout} no need to broadcast data across pp stages")
+    return {}
+
+
+def _is_multi_shard(in_tensor_map):
+    """
+    whether the input tensor map is in multi shard
+    """
+    for tensor_map in in_tensor_map:
+        if isinstance(tensor_map, (list, tuple)) and len(tensor_map) > 1:
+            return True
+    return False
+
+
+def _insert_expand_layout_reshape(param_rank_map, from_info_tuple, to_info_tuple,
+                                  insert_from_reshape, insert_to_reshape):
+    """ insert layout expand op reshape """
+    from_dev_matrix = from_info_tuple[0]
+    from_tensor_map = from_info_tuple[1]
+    from_full_tensor_shape = from_info_tuple[2]
+    to_dev_matrix_origin = to_info_tuple[0]
+    to_tensor_map_origin = to_info_tuple[1]
+    origin_tensor_shape = to_info_tuple[2]
+    for param_rank, _ in param_rank_map.items():
+        if insert_from_reshape:
+            from_slice_tensor_shape = ()
+            from_tensor_strategy = _get_tensor_strategy(from_dev_matrix, from_tensor_map)
+            for i, item in enumerate(from_full_tensor_shape):
+                from_slice_tensor_shape += (item // from_tensor_strategy[i],)
+            param_rank_map.get(param_rank).insert(0, ('Reshape', list(from_slice_tensor_shape)))
+        if insert_to_reshape:
+            to_tensor_strategy = _get_tensor_strategy(to_dev_matrix_origin, to_tensor_map_origin)
+            to_slice_tensor_shape = ()
+            for i, item in enumerate(origin_tensor_shape):
+                to_slice_tensor_shape += (item // to_tensor_strategy[i],)
+            param_rank_map.get(param_rank).append(('Reshape', list(to_slice_tensor_shape)))
+
+
+def _infer_reshard_op_map(from_layout, to_layout, self_rank):
+    """infer reshard op map"""
+    from_layout_without_rank_list = from_layout[:-1]
+    to_layout_without_rank_list = to_layout[:-1]
+    if _is_multi_shard(from_layout[1]):
+        # ((2, 1), 1) --> (2, 1, 1) expand tensormap
+        new_layout = _expand_layout(from_layout[0], from_layout[1], from_layout[2])
+        from_layout_without_rank_list = (new_layout[0], new_layout[1], new_layout[2])
+    if _is_multi_shard(to_layout[1]):
+        new_layout = _expand_layout(to_layout[0], to_layout[1], to_layout[2])
+        to_layout_without_rank_list = (new_layout[0], new_layout[1], new_layout[2])
+    operator_map = _get_needed_rank_transform_operator_map_by_layouts(from_layout_without_rank_list,
+                                                                      to_layout_without_rank_list,
+                                                                      from_layout[3], self_rank,
+                                                                      True)
+    new_to_layout_info = to_layout[:-1]
+    _insert_expand_layout_reshape(operator_map, from_layout_without_rank_list, new_to_layout_info,
+                                  _is_multi_shard(from_layout[1]), _is_multi_shard(to_layout[1]))
+    return operator_map
+
+
+def _get_resharding_operator_map(from_layout, to_layout, self_rank):
+    """
+        Args:
+        from_layout (tuple): Use tuple to present layout
+          (device_matrix(list), tensor_map(list), global_shape(list), rank_list(list))
+        to_layout (tuple): Use tuple to present layout
+          (device_matrix(list), tensor_map(list), global_shape(list), rank_list(list))
+        self_rank (int): rank_id
+    """
+    reshard_op_cache_key = (
+        f"{from_layout[0]}, {from_layout[1]}, {from_layout[2]}, {from_layout[3]}"
+        f" -> "
+        f"{to_layout[0]}, {to_layout[1]}, {from_layout[2]}, {to_layout[3]}")
+    global RESHARD_OP_MAP_CACHE
+    if reshard_op_cache_key not in RESHARD_OP_MAP_CACHE:
+        operator_map = _infer_reshard_op_map(from_layout, to_layout, self_rank)
+        op_map_dict = {rank_id: operator_map for rank_id in operator_map}
+        RESHARD_OP_MAP_CACHE[reshard_op_cache_key] = op_map_dict
+        logger.debug(f"reshard_op_cache_key is {reshard_op_cache_key}, not match cache")
+    else:
+        cache_rank_list_dict = RESHARD_OP_MAP_CACHE[reshard_op_cache_key]
+        if self_rank in cache_rank_list_dict:
+            operator_map = cache_rank_list_dict[self_rank]
+            logger.debug(f"reshard_op_cache_key is {reshard_op_cache_key}, match cache")
+        else:
+            logger.debug(f"reshard_op_cache_key is {reshard_op_cache_key}, "
+                         f"but rank {self_rank} is not match cache")
+            operator_map = _infer_reshard_op_map(from_layout, to_layout, self_rank)
+            for rank_id in operator_map:
+                RESHARD_OP_MAP_CACHE[reshard_op_cache_key][rank_id] = operator_map
+    return operator_map

mindspore/parallel/_transformer/moe.py

@@ -111,7 +111,7 @@ def _check_moe_config(moe_config=None, parallel_config=None):
     """
     if not isinstance(moe_config, MoEConfig):
         raise TypeError(f"'moe_config' must be an instance of MoEConfig, but got {type(moe_config).__name__}.")
-    use_moe = (moe_config.expert_num > 1)
+    use_moe = moe_config.expert_num > 1
     if use_moe is False:
         return
     if moe_config.expert_num % parallel_config.expert_parallel != 0:

mindspore/parallel/_transformer/transformer.py

@@ -1365,7 +1365,12 @@ class MultiHeadAttention(Cell):
 class TransformerEncoderLayer(Cell):
     r"""
         Transformer Encoder Layer. This is an implementation of the single layer of the transformer
-        encoder layer, including multihead attention and feedward layer.
+        encoder layer, mainly including Multi-Head Attention, Feed Forward, Add and LayerNorm layer.
+
+        The TransformerEncoderLayer structure is shown in the following figure:
+
+        .. image:: ../images/TransformerEncoderLayer.png
+            :align: center
 
         Args:
             batch_size(int): The batch size of the input tensor when do increnmental prediction. Should be a positive
@@ -1448,7 +1453,7 @@ class TransformerEncoderLayer(Cell):
             >>> # When use use_past=True, it includes two steps to implement the incremental prediction.
             >>> # Step 1: set is_first_iteration=True, and input the full sequence length's state.
             >>> batch_valid_length = Tensor(np.ones((2,)), mstype.int32)
-            >>> init_reset = Tensor([True], mstype.bool_)
+            >>> init_reset = Tensor([True], mstype.bool)
             >>> # Set is_first_iteration=True to generate the full memory states
             >>> model = TransformerEncoderLayer(batch_size=2, hidden_size=8, ffn_hidden_size=64, seq_length=16,
             ...                                 num_heads=2, use_past=True)
@@ -1462,7 +1467,7 @@ class TransformerEncoderLayer(Cell):
             (2, 2, 16, 4)
             >>> encoder_input_value = Tensor(np.ones((2, 1, 8)), mstype.float32)
             >>> encoder_input_mask = Tensor(np.ones((2, 1, 16)), mstype.float16)
-            >>> init_reset = Tensor([False], mstype.bool_)
+            >>> init_reset = Tensor([False], mstype.bool)
             >>> # Step 2: set is_first_iteration=False, and pass the single word to run the prediction rather than
             >>> # the full sequence.
             >>> model.add_flags_recursive(is_first_iteration=False)
@@ -1532,7 +1537,7 @@ class TransformerEncoderLayer(Cell):
                     "and parallel_config. model_parallel is {}."
                     .format(ffn_hidden_size, parallel_config.model_parallel))
             _check_moe_config(moe_config, parallel_config)
-            self.use_moe = (moe_config.expert_num > 1)
+            self.use_moe = moe_config.expert_num > 1
             self.use_past = use_past
             self.seq_length = seq_length
             self.hidden_size = hidden_size
@@ -1607,7 +1612,7 @@ class TransformerEncoderLayer(Cell):
                     "and parallel_config. model_parallel is {}."
                     .format(ffn_hidden_size, parallel_config.model_parallel))
             _check_moe_config(moe_config, parallel_config)
-            self.use_moe = (moe_config.expert_num > 1)
+            self.use_moe = moe_config.expert_num > 1
             self.use_past = use_past
             self.seq_length = seq_length
             self.hidden_size = hidden_size
@@ -1902,7 +1907,7 @@ class TransformerDecoderLayer(Cell):
                  parallel_config=default_dpmp_config):
         super(TransformerDecoderLayer, self).__init__()
         _check_moe_config(moe_config, parallel_config)
-        self.use_moe = (moe_config.expert_num > 1)
+        self.use_moe = moe_config.expert_num > 1
         config_to_attention = parallel_config.dpmp if self.use_moe else parallel_config
         if batch_size or use_past:
             Validator.check_positive_int(batch_size)
@@ -2370,7 +2375,7 @@ class TransformerEncoder(Cell):
             >>> # When use use_past=True, it includes two steps to implement the incremental prediction.
             >>> # Step 1: set is_first_iteration=True, and input the full sequence length's state.
             >>> batch_valid_length = Tensor(np.ones((2,)), mstype.int32)
-            >>> init_reset = Tensor([True], mstype.bool_)
+            >>> init_reset = Tensor([True], mstype.bool)
             >>> # Set is_first_iteration=True to generate the full memory states
             >>> model = TransformerEncoder(batch_size=2, hidden_size=8, ffn_hidden_size=64, seq_length=16,
             ...                            num_heads=2, num_layers=2, use_past=True)
@@ -2384,7 +2389,7 @@ class TransformerEncoder(Cell):
             (2, 2, 16, 4)
             >>> encoder_input_value = Tensor(np.ones((2, 1, 8)), mstype.float32)
             >>> encoder_input_mask = Tensor(np.ones((2, 1, 16)), mstype.float16)
-            >>> init_reset = Tensor([False], mstype.bool_)
+            >>> init_reset = Tensor([False], mstype.bool)
             >>> # Step 2: set is_first_iteration=False, and pass the single word to run the prediction rather than
             >>> # the full sequence.
             >>> model.add_flags_recursive(is_first_iteration=False)
@@ -2440,7 +2445,7 @@ class TransformerEncoder(Cell):
         super(TransformerEncoder, self).__init__()
         _check_config(parallel_config)
         _check_moe_config(moe_config, parallel_config)
-        self.use_moe = (moe_config.expert_num > 1)
+        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,):
             self.add = P.Add()
@@ -2677,7 +2682,7 @@ class TransformerDecoder(Cell):
         super(TransformerDecoder, self).__init__()
         _check_moe_config(moe_config, parallel_config)
         _check_config(parallel_config)
-        self.use_moe = (moe_config.expert_num > 1)
+        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,):
             self.add = P.Add()
@@ -2959,7 +2964,7 @@ class Transformer(Cell):
             if not lambda_func:
                 lambda_func = _get_lambda_func(total_layer=encoder_layers + decoder_layers)
             _check_moe_config(moe_config, parallel_config)
-            self.use_moe = (moe_config.expert_num > 1)
+            self.use_moe = moe_config.expert_num > 1
             self.add = P.Add()
             self.aux_loss = Tensor(0.0, mstype.float32)
             if encoder_layers > 0:
@@ -3026,7 +3031,7 @@ class Transformer(Cell):
             if not lambda_func:
                 lambda_func = _get_lambda_func(total_layer=encoder_layers + decoder_layers)
             _check_moe_config(moe_config, parallel_config)
-            self.use_moe = (moe_config.expert_num > 1)
+            self.use_moe = moe_config.expert_num > 1
             self.add = P.Add().shard(((), ()))
             self.aux_loss = Tensor(0.0, mstype.float32)
             if encoder_layers > 0:

mindspore/parallel/_utils.py

@@ -21,7 +21,7 @@ import mindspore as ms
 from mindspore import context, log as logger
 from mindspore._c_expression import reset_op_id, reset_op_id_with_offset
 from mindspore.common.tensor import Tensor
-from mindspore.common.dtype import dtype_to_nptype
+from mindspore.common.dtype import _dtype_to_nptype
 from mindspore.common import dtype as mstype
 from mindspore.communication.management import get_group_size, get_rank
 from mindspore.communication._comm_helper import _is_initialized
@@ -156,7 +156,7 @@ def _is_in_auto_parallel_mode():
 
 
 def _is_parallel_mode():
-    if not _is_initialized() or context.get_context('mode') == context.PYNATIVE_MODE:
+    if not _is_initialized():
         return False
     if os.getenv("RUN_MODE") != "predict":
         return False
@@ -173,12 +173,6 @@ def _is_in_hybrid_parallel_mode():
     return _get_parallel_mode() == ms.ParallelMode.HYBRID_PARALLEL
 
 
-def _is_pynative_parallel():
-    parallel_mode = context.get_auto_parallel_context('parallel_mode')
-    return context.get_context('mode') == context.PYNATIVE_MODE and parallel_mode in (
-        context.ParallelMode.SEMI_AUTO_PARALLEL, context.ParallelMode.AUTO_PARALLEL)
-
-
 def _get_full_batch():
     """Get whether to use full_batch."""
     return auto_parallel_context().get_full_batch()
@@ -452,7 +446,7 @@ def _to_full_tensor(elem, global_device_num, global_rank, scaling_sens=None):
                     batchsize_per_device = item
                 else:
                     new_shape += (item,)
-            new_tensor_numpy = np.zeros(new_shape, dtype_to_nptype(type_))
+            new_tensor_numpy = np.zeros(new_shape, _dtype_to_nptype(type_))  # pylint:disable=protected-access
             start = stage_rank * batchsize_per_device
             new_tensor_numpy[start: start + batchsize_per_device] = data.asnumpy()
         else:
@@ -466,7 +460,7 @@ def _to_full_tensor(elem, global_device_num, global_rank, scaling_sens=None):
                 end = (stage_rank % dataset_strategy[index][i] + 1) * item
                 s = slice(start, end, 1)
                 slice_index += (s,)
-            new_tensor_numpy = np.zeros(new_shape, dtype_to_nptype(type_))
+            new_tensor_numpy = np.zeros(new_shape, _dtype_to_nptype(type_))  # pylint:disable=protected-access
             new_tensor_numpy[slice_index] = data.asnumpy()
         new_tensor = Tensor(new_tensor_numpy, dtype=type_)
         lst.append(new_tensor)
@@ -773,7 +767,7 @@ 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 _is_pynative_parallel() or not _get_full_batch():
+    if not _get_full_batch():
         return grads
 
     device_num = _get_device_num()

mindspore/parallel/auto_parallel.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """Cell of auto parallel"""
 import os
+from mindspore import jit
 from mindspore.nn.cell import Cell
 from mindspore.parallel.shard import Layout
 from mindspore.communication.management import get_rank, get_group_size
@@ -162,6 +163,10 @@ class AutoParallel(Cell):
         super(AutoParallel, self).__init__(auto_prefix=False)
         self.network = network
 
+        if parallel_mode not in ["semi_auto", "sharding_propagation", "recursive_programming"]:
+            raise ValueError("the argument 'parallel_mode' must be one of ['semi_auto', 'sharding_propagation'," \
+                  " 'recursive_programming'], but got the value : {} .".format(parallel_mode))
+
         self._parallel_mode = parallel_mode
 
         self._global_rank = get_rank()
@@ -260,8 +265,10 @@ class AutoParallel(Cell):
         self._save_strategy_file_path = file_path
 
     def disable_strategy_file_only_for_trainable_params(self):
-        """By default, MindSpore only loads and saves trainable parameters. This API enables the loading and saving of
-        non-trainable parameters as well."""
+        """
+        By default, MindSpore only loads and saves trainable parameters. This API enables the loading and saving of
+        non-trainable parameters as well.
+        """
         self._only_trainable_params = False
 
     def save_operator_strategy_file(self, file_path):
@@ -275,7 +282,8 @@ class AutoParallel(Cell):
         Note:
             - It only works when `parallel_mode=sharding_propagation`.
             - When performing distributed training, users can first save the strategy using dryrun on a single device
-              and then load strategy to perform distributed training.
+              and then load strategy to perform distributed training. Note that only the first device of each node will
+              save the strategy file, so the simulated rank id specified by Dryrun must be divisible by 8.
 
         Args:
             file_path (str): Path to save parallel strategy json, must be an absolute path.
@@ -410,7 +418,7 @@ class AutoParallel(Cell):
                 raise ValueError("For 'AutoParallel.dataset_strategy', the argument "
                                  "'config' must be 'full_batch' or 'data_parallel', but got the value : {}."
                                  .format(config))
-            self._full_batch = (config == "full_batch")
+            self._full_batch = config == "full_batch"
             self._dataset_strategy_config = config
             return
         if not isinstance(config, tuple):
@@ -436,9 +444,9 @@ class AutoParallel(Cell):
         Args:
             shard_size (int, optional): Set the optimizer weight shard group size if you want to specific the
                 maximum group size across devices when the parallel optimizer is
-                enabled. The numerical range can be (0, device_num]. Default value
+                enabled. The numerical range can be (0, device_num] or -1. Default value
                 is -1, which means the optimizer weight shard group size will
-                the data parallel group of each parameter. Default  ``-1``.
+                the data parallel group of each parameter.
             threshold (int, optional): Set the threshold of parallel optimizer. When parallel optimizer is
                 enabled, parameters with size smaller than this threshold will not be
                 sharded across the devices. Parameter size = shape[0] \* ... \*
@@ -505,17 +513,18 @@ class AutoParallel(Cell):
             raise ValueError("For 'AutoParallel.pipeline', the argument 'stages' "
                              "must be larger than zero, but got value: {}.".format(stages))
         if not isinstance(output_broadcast, bool):
-            raise TypeError("For 'AutoParallel.pipeline', the argument 'stages' "
+            raise TypeError("For 'AutoParallel.pipeline', the argument 'output_broadcast' "
                             "must be bool type, but got the type : {}.".format(type(output_broadcast)))
         if not isinstance(interleave, bool):
-            raise TypeError("For 'AutoParallel.pipeline', the argument 'stages' "
+            raise TypeError("For 'AutoParallel.pipeline', the argument 'interleave' "
                             "must be bool type, but got the type : {}.".format(type(interleave)))
         if not isinstance(scheduler, str):
-            raise TypeError("For 'AutoParallel.pipeline', the argument 'stages' "
+            raise TypeError("For 'AutoParallel.pipeline', the argument 'scheduler' "
                             "must be str type, but got the type : {}.".format(type(scheduler)))
-        if scheduler not in ("1f1b", "gpipe"):
+        if scheduler not in ("1f1b", "gpipe", "seqpipe", "seqvpp", "seqsmartvpp", "zero_bubble_v"):
             raise ValueError("For 'AutoParallel.pipeline', the argument "
-                             "'scheduler' must be '1f1b' , 'gpipe' , but got the value : {}."
+                             "'scheduler' must be '1f1b'/'gpipe'/'seqpipe'/'seqvpp'/'seqsmartvpp'/'zero_bubble_v' ," \
+                             " but got the value : {}."
                              .format(scheduler))
         self._pipeline_stages = stages
         self._pipeline_result_broadcast = output_broadcast
@@ -658,8 +667,11 @@ class AutoParallel(Cell):
                 - recomputation_communication_overlap (bool): Enable overlap between recompute ops and communication ops
                   if True.
                   Default: False.
-                - grad_matmul_communication_overlap (bool): Enable overlap between dw matmul and
-                  tensor parallel communication ops if True. Default: False.
+                - grad_matmul_communication_overlap (bool, str): When set to ``True``, it indicates that overlap
+                  between dw matmul and tensor parallel communication is enabled. When set to ``False``, it indicates
+                  that this feature is disabled. When set to str, it only optimizes the specified communication
+                  operator types, with operators separated by ``,``. For example, "AlltoAll,AlltoAllV" indicates that
+                  only ``AlltoAll`` and ``AlltoAllV`` are optimized. Default: ``False``.
                 - grad_fa_allgather_overlap (bool): Enable overlap between duplicated allgather by recomputing
                   in sequence parallel and flashattentionscoregrad ops if True. Default: False.
                 - enable_communication_fusion (bool): Enable communication fusion to optimize the number of
@@ -674,7 +686,9 @@ class AutoParallel(Cell):
                   and optimizer parallel allgather communication if True. Currently, do not support
                   `O2 <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.JitConfig.html>`_
                   Default: False.
-                - computation_communication_fusion_level (int): Enable the fusion between compute and communicate.
+                - computation_communication_fusion_level (int): Enable the fusion between compute and communicate,
+                  which fuses communication tasks and computing tasks, allows for partial pipelining and parallel
+                  execution of these tasks during operation, thereby enhancing performance.
                   Default: ``0``. Note: This function must be used with Ascend Training Solution 24.0.RC2 or later.
                   This is an experimental configuration, may be changed or canceled in the future.
 
@@ -685,6 +699,12 @@ class AutoParallel(Cell):
                   - 2: Apply fusion to backward nodes.
 
                   - 3: Apply fusion to all nodes.
+
+                  .. warning::
+                      After setting ``export MS_ENABLE_LCCL=on``, the fusion operator based on memory semantics will be
+                      used. Please note that this operator is still in an experimental stage and may be changed or
+                      removed in the future.
+
                 - dataset_broadcast_opt_level (int): Optimize the scenario that the dataset repeated reading. Only
                   support O0/O1 jit level. It doesn't work in O2 mode. Default: ``0``.
 
@@ -728,5 +748,6 @@ class AutoParallel(Cell):
         self._transformer_opt_config = file_path
         ctx.ascend_config['parallel_speed_up_json_path'] = file_path
 
+    @jit
     def construct(self, *args, **kwargs):
         return self.network(*args, **kwargs)

mindspore/parallel/checkpoint_convert.py

@@ -15,10 +15,10 @@
 """Convert distributed checkpoint"""
 from __future__ import absolute_import
 
-from mindspore.parallel.checkpoint_transform import rank_list_for_transform, transform_checkpoint_by_rank, transform_checkpoints
+__all__ = ["rank_list_for_convert", "convert_checkpoint_by_rank", "convert_checkpoints"]
 
-__all__ = ["rank_list_for_convert",
-           "convert_checkpoint_by_rank", "convert_checkpoints"]
+from mindspore.parallel.checkpoint_transform import rank_list_for_transform, transform_checkpoint_by_rank, \
+    transform_checkpoints
 
 
 def rank_list_for_convert(rank_id, src_strategy_file=None, dst_strategy_file=None):