mindspore 2.6.0rc1__cp39-cp39-win_amd64.whl → 2.7.0rc1__cp39-cp39-win_amd64.whl

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
@@ -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)
@@ -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/auto_parallel.py

@@ -162,6 +162,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 +264,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):
@@ -410,7 +416,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 +442,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] \* ... \*
@@ -513,9 +519,10 @@ class AutoParallel(Cell):
         if not isinstance(scheduler, str):
             raise TypeError("For 'AutoParallel.pipeline', the argument 'stages' "
                             "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"):
             raise ValueError("For 'AutoParallel.pipeline', the argument "
-                             "'scheduler' must be '1f1b' , 'gpipe' , but got the value : {}."
+                             "'scheduler' must be '1f1b'/'gpipe'/'seqpipe'/'seqvpp'/'seqsmartvpp' ," \
+                             " but got the value : {}."
                              .format(scheduler))
         self._pipeline_stages = stages
         self._pipeline_result_broadcast = output_broadcast

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):

mindspore/parallel/checkpoint_transform.py

@@ -582,6 +582,8 @@ def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
         The number of multiprocess settings is related to the size of the host, and it is not recommended to set it
         too large, otherwise it may cause freezing.
 
+        This function does not support converting remove_redundancy's checkpoint file.
+
     Args:
         src_checkpoints_dir (str): The source checkpoints directory.
         dst_checkpoints_dir (str): The destination checkpoints directory to save the converted checkpoints.
@@ -924,8 +926,8 @@ def set_op_strategy_config(mode="SAVE", path=""):
 
 def build_searched_strategy(strategy_filename):
     """
-    Extract the sharding strategy for each parameter in the network
-    from the strategy file for distributed inference scenarios.
+    Extract the sharding strategy for each parameter in the network from the strategy file
+    for distributed inference scenarios.
 
     Args:
         strategy_filename (str): Name of strategy file.
@@ -1025,8 +1027,10 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
         >>> from mindspore.parallel.auto_parallel import AutoParallel
         >>> from mindspore.nn.utils import no_init_parameters
         >>> from mindspore.common.initializer import initializer, One
+        >>> from mindspore.communication.management import get_group_size
         >>>
         >>> step_per_epoch = 4
+        >>> device_num = get_group_size()
         >>>
         >>> # Define the network structure.
         >>> class Net(nn.Cell):
@@ -1070,7 +1074,7 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
         ...     network = AutoParallel(network, parallel_mode="semi_auto")
         ...     network.save_param_strategy_file(file_path="./train_strategy.ckpt")
         ...     model = ms.Model(network=network, loss_fn=net_loss, optimizer=net_opt)
-        ...     ckpt_config = train.CheckpointConfig(keep_checkpoint_max=1, integrated_save=False)
+        ...     ckpt_config = train.CheckpointConfig(keep_checkpoint_max=1, integrated_save=True)
         ...     global_rank_id = int(os.getenv("RANK_ID"))
         ...     ckpt_path = "./rank_{}_ckpt".format(global_rank_id)
         ...     ckpt_callback = train.ModelCheckpoint(prefix="parallel", directory=ckpt_path, config=ckpt_config)
@@ -1096,10 +1100,10 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
         >>>
         >>> train_net()
         >>> load_model()
-        [[-7.3259363 -7.497216  -7.398196  ... -7.374962  -7.204874  -7.234935 ]
-        [ 3.362938   3.3535435  3.3832688 ...  3.4263954  3.279045   3.3202887]
+        [[-9.62929535e+00, -9.76258755e+00, -9.70192051e+00 ... -9.67151260e+00, -9.71998310e+00, -9.64571190e+00],
+        [-4.63218540e-01, -4.07317460e-01, -3.78161550e-01 ... -3.95918339e-01, -2.87363172e-01, -3.48693460e-01],
         ...
-        [ 1.6067538  1.6244187  1.5384722 ...  1.5449994  1.6195512  1.6176052]]
+        [-4.28075647e+00, -4.36630344e+00, -4.25664043e+00 ... -4.32012939e+00, -4.30337954e+00, -4.27571440e+00]]
     """
     if format not in ['safetensors', 'ckpt'] or output_format not in ['safetensors', 'ckpt']:
         raise ValueError(
@@ -1185,7 +1189,7 @@ def load_distributed_checkpoint(network, checkpoint_filenames=None, predict_stra
     param_not_in_ckpt = []
     for _, param in network.parameters_and_names():
         sliced_params = []
-        if param.name not in rank_list.keys():
+        if param.name not in rank_list:
             param_not_in_strategy.append(param.name)
             continue
         if param.name not in param_total_dict: