mindspore 2.2.10__cp38-none-any.whl → 2.2.14__cp38-none-any.whl

mindspore/parallel/_parallel_serialization.py

@@ -243,6 +243,33 @@ def _extract_pipeline_stage_num(strategy_file):
     return pipeline_stage_num
 
 
+def _extract_src_dst_layout_map_by_src(src_strategy_file=None, dst_strategy_file=None):
+    """Extract strategy list by src strategy"""
+    src_layout_map = _extract_layout_map(src_strategy_file)
+    dst_layout_map = _extract_layout_map(dst_strategy_file)
+    if dst_layout_map is None:
+        return src_layout_map, dst_layout_map
+    for param_name in list(dst_layout_map.keys()):
+        if param_name in src_layout_map.keys():
+            continue
+        dst_layout_map.pop(param_name)
+    stage_id = 0
+    if src_strategy_file[-5:] == ".json":
+        with open(src_strategy_file, 'r') as f:
+            json_content = json.load(f)
+        strategy_items = json_content.get("parallel_strategy_item")
+        if not strategy_items:
+            raise ValueError("The strategy file {} if empty.".format(src_strategy_file))
+        stage_id = strategy_items.get(list(strategy_items.keys())[0]).get('stage')
+    else:
+        src_parallel_strategy_map = _load_protobuf_strategy(src_strategy_file)
+        strategy_items = src_parallel_strategy_map.parallel_strategy_item
+        if not strategy_items:
+            raise ValueError("The strategy file {} if empty.".format(src_strategy_file))
+        stage_id = strategy_items[0].parallel_strategys.stage
+    return src_layout_map, dst_layout_map, stage_id
+
+
 def _extract_src_dst_layout_map(rank_id, src_strategy_file=None, dst_strategy_file=None):
     """Extract strategy list"""
     src_layout_map = _extract_layout_map(src_strategy_file)
@@ -341,6 +368,7 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
     Transform model parallel dimension for distributed checkpoint files.
     """
     transform_param_dict = {}
+    device_num = -1
     for param_name, _ in param_total_dict.items():
         tensor_shape = list(param_total_dict[param_name].values())[0].shape
         from_dev_matrix = [1]
@@ -394,14 +422,18 @@ def _transform_parallel_checkpoint(rank_id, param_total_dict, param_attr_dict, s
         to_info_tuple = (to_opt_shard_size, to_dev_matrix_origin, to_tensor_map_origin, origin_tensor_shape)
         _insert_opt_shard_reshape(param_rank_map, from_info_tuple, to_info_tuple)
         transform_operator_stack = _generate_transform_operator_stack(param_rank_map, rank_id)
-        _apply_tensor_transform_operators(transform_operator_stack, param_total_dict[param_name], device_num)
-        transform_tensor = ms.Tensor(param_total_dict[param_name][rank_id % device_num])
+        param_total_dict_copy = param_total_dict[param_name].copy()
+        _apply_tensor_transform_operators(transform_operator_stack, param_total_dict_copy, device_num)
+        transform_tensor = ms.Tensor(param_total_dict_copy[rank_id % device_num])
         requires_grad = param_attr_dict[param_name][rank_id % device_num][0]
         layerwise_parallel = param_attr_dict[param_name][rank_id % device_num][1]
         transform_para = ms.Parameter(transform_tensor, param_name, requires_grad, layerwise_parallel)
         if param_type_dict[param_name][rank_id % device_num] == "BFloat16":
             transform_para.set_dtype(ms.bfloat16)
         transform_param_dict[param_name] = transform_para
+    if device_num < 0:
+        raise ValueError("None of the parameters in checkpoint file are in either src strategy or "
+                         "dst strategy. Please check correctness of strategy files.")
 
     # Handle those parameter like learning_rate, global_step which not in strategy_file.
     for param_name, _ in param_total_dict.items():

mindspore/parallel/_tensor.py

@@ -223,7 +223,9 @@ def _load_tensor(tensor, dev_mat, tensor_map, rank_id=-1):
     tensor_strategy = _get_tensor_strategy(dev_mat, tensor_map)
     tensor_slice_index = _get_tensor_slice_index(dev_mat, tensor_strategy, tensor_map, rank)
     if tensor.dtype == mstype.bfloat16:
-        tensor = tensor.float()
+        from mindspore.ops.operations import Cast
+        cpu_cast = Cast().set_device("CPU")
+        tensor = cpu_cast(tensor, mstype.float32)
     np_tensor = tensor.asnumpy()
     np_tensor_list = _chunk_tensor_by_strategy(np_tensor, tensor_strategy)
     np_tensor_slice = np_tensor_list[int(tensor_slice_index)]

mindspore/parallel/_transformer/transformer.py

@@ -805,14 +805,14 @@ class MultiHeadAttention(Cell):
             - **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).
+            - **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.
-            - **value_past** (Tensor) - Float16 tensor with shape
+            - **value_past** (Tensor) - float16 tensor with shape
               (batch_size, num_heads, tgt_seq_length, size_per_head).
               The past calculated value vector. Used for incremental prediction when the use_past is True.
               Default None.
-            - **batch_valid_length** (Tensor) - Int32 tensor with shape (batch_size,) the past calculated the index.
+            - **batch_valid_length** (Tensor) - int32 tensor with shape (batch_size,) the past calculated the index.
               Used for incremental prediction when the use_past is True. Default None.
 
         Outputs:
@@ -1412,7 +1412,7 @@ class TransformerEncoderLayer(Cell):
               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.
+            - **batch_valid_length** (Tensor) - int32 tensor with shape [batch_size] the past calculated the index.
               Used for incremental prediction when the use_past is True. Default None.
 
         Outputs:
@@ -1824,7 +1824,7 @@ class TransformerDecoderLayer(Cell):
               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.
+            - **batch_valid_length** (Tensor) - int32 tensor with shape [batch_size] the past calculated the index.
               Used for incremental prediction when the use_past is True. Default None.
 
         Outputs:
@@ -2333,7 +2333,7 @@ class TransformerEncoder(Cell):
               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.
+            - **batch_valid_length** (Tensor) - int32 tensor with shape [batch_size] the past calculated the index.
               Used for incremental prediction when the use_past is True. Default None.
 
         Outputs:
@@ -2589,7 +2589,7 @@ class TransformerDecoder(Cell):
               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.
+            - **batch_valid_length** (Tensor) - int32 tensor with shape [batch_size] the past calculated the index.
               Used for incremental prediction when the use_past is True. Default None.
 
         Outputs:
@@ -2842,7 +2842,7 @@ class Transformer(Cell):
               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.
+            - **batch_valid_length** (Tensor) - int32 tensor with shape [batch_size] the past calculated the index.
               Used for incremental prediction when the use_past is True. Default None.
 
         Outputs:

mindspore/parallel/checkpoint_transform.py

@@ -25,11 +25,11 @@ from mindspore.common import dtype as mstype
 from mindspore.parallel._parallel_serialization import _rank_list_for_transform_parallel_checkpoint, \
     _transform_parallel_checkpoint, _get_device_num_from_strategy, _make_dir, \
     _extract_layout_map, _extract_src_dst_layout_map, _parameter_not_in_local_stage, _extract_pipeline_stage_num, \
-    _merge_protobuf_strategy, _merge_json_strategy
+    _merge_protobuf_strategy, _merge_json_strategy, _extract_src_dst_layout_map_by_src
 
 
 __all__ = ["merge_pipeline_strategys", "rank_list_for_transform", "transform_checkpoint_by_rank",
-           "transform_checkpoints"]
+           "transform_checkpoints", "load_segmented_checkpoints"]
 
 
 def merge_pipeline_strategys(src_strategy_dirs, dst_strategy_file):
@@ -71,7 +71,6 @@ def merge_pipeline_strategys(src_strategy_dirs, dst_strategy_file):
         _merge_json_strategy(src_strategy_files_json, dst_strategy_file)
 
 
-
 def rank_list_for_transform(rank_id, src_strategy_file=None, dst_strategy_file=None):
     """
     List of original distributed checkpoint rank index for obtaining the target checkpoint of a rank_id
@@ -222,48 +221,63 @@ def transform_checkpoint_by_rank(rank_id, checkpoint_files_map, save_checkpoint_
     ms.save_checkpoint(transform_param_list, save_checkpoint_file_name)
 
 
-def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix, src_strategy_file=None,
-                          dst_strategy_file=None):
-    """
-    Transform distributed checkpoint from source sharding strategy to destination sharding strategy for a rank.
-    For more details about converting distributed Checkpoint, please refer to
-    `Model Transformation <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/model_transformation.html>`_.
-
-    Note:
-        The `src_checkpoints_dir` directory structure should be organized like "src_checkpoints_dir/rank_0/a.ckpt", the
-        rank number should be set to a subdirectory and the checkpoint file is stored in this subdirectory. If multiple
-        files exist in a rank directory, the last file in the lexicgraphic order would be selected.
-
-    Args:
-        src_checkpoints_dir (str): The source checkpoints directory.
-        dst_checkpoints_dir (str): The destination checkpoints directory to save the converted checkpoints.
-        ckpt_prefix (str): The destination checkpoint name prefix.
-        src_strategy_file (str): Name of source sharding strategy file which saved by
-                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
-                                 when the 'src_strategy_file' is None, it means that the source sharding strategy is
-                                 without any sharing for each parameter. Default:None.
-        dst_strategy_file (str): Name of destination sharding strategy file which saved by
-                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
-                                 when the 'dst_strategy_file' is None, it means that the destination sharding strategy
-                                 is without any sharing for each parameter. Default:None.
-
-    Raises:
-        ValueError: `src_strategy_file` or `dst_strategy_file` is incorrect.
-        NotADirectoryError: `src_checkpoints_dir` or `dst_checkpoints_dir` is not a directory.
-        ValueError: The checkpoint file is missing in `src_checkpoints_dir`.
-        TypeError: `src_strategy_file` or `dst_strategy_file` is not a string.
-
-    Examples:
-        >>> import mindspore as ms
-        >>> ms.transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, "dst_checkpoint",
-        ...                       "./src_strategy.ckpt", "./dst_strategy.ckpt")
-
-    """
-    if not os.path.isdir(src_checkpoints_dir):
-        raise NotADirectoryError("src_checkpoints_dir {} is not a directory.".format(src_checkpoints_dir))
-    _make_dir(dst_checkpoints_dir, "path")
-    if not isinstance(ckpt_prefix, str):
-        raise TypeError("The ckpt_prefix should be a str.")
+def _transform_checkpoint_by_stage(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix, src_strategy_file,
+                                   dst_strategy_file=None):
+    """Transform checkpoint for stage in src_strategy_file"""
+    param_total_dict = defaultdict(dict)
+    param_attr_dict = defaultdict(dict)
+    param_type_dict = defaultdict(dict)
+    src_strategy_list, dst_strategy_list, stage_id = _extract_src_dst_layout_map_by_src(src_strategy_file, \
+                                                                                             dst_strategy_file)
+    src_stage_device_num = np.prod(src_strategy_list.get(list(src_strategy_list.keys())[0])[0]) if src_strategy_list \
+                                                                                                   is not None else 1
+    dst_stage_device_num = np.prod(dst_strategy_list.get(list(dst_strategy_list.keys())[0])[0]) if dst_strategy_list \
+                                                                                                   is not None else 1
+    origin_dst_strategy_list = _extract_layout_map(dst_strategy_file)
+    origin_src_strategy_list = _extract_layout_map(src_strategy_file)
+    checkpoint_files_map = {}
+    src_rank_id_start = stage_id * src_stage_device_num
+    for local_rank in range(src_stage_device_num):
+        rank_id = src_rank_id_start + local_rank
+        checkpoint_file_name = os.path.join(src_checkpoints_dir, "rank_{}".format(rank_id), "*.ckpt")
+        rank_ckpts = glob.glob(checkpoint_file_name)
+        rank_ckpts.sort()
+        for checkpoint_file in rank_ckpts:
+            if not os.path.isfile(checkpoint_file):
+                ms.log.warning("{} is not a checkpoint file.".format(checkpoint_file))
+                continue
+            checkpoint_files_map[rank_id] = checkpoint_file
+    for rank, local_file in checkpoint_files_map.items():
+        if not os.path.exists(local_file):
+            raise ValueError("Checkpoint file {} in rank {} not exits: ".format(local_file, rank))
+    for rank, file_name in checkpoint_files_map.items():
+        ckpt_dict = ms.load_checkpoint(file_name)
+        for param_name, param in ckpt_dict.items():
+            # cut the parameter not in the pipeline stage.
+            if _parameter_not_in_local_stage(param_name, origin_src_strategy_list, src_strategy_list) \
+                    and _parameter_not_in_local_stage(param_name, origin_dst_strategy_list, dst_strategy_list):
+                continue
+            src_rank = rank % src_stage_device_num
+            param_type_dict[param_name][src_rank] = str(param.data.dtype)
+            if param.data.dtype == mstype.bfloat16:
+                param.set_dtype(mstype.float32)
+            param_total_dict[param_name][src_rank] = param.data.asnumpy()
+            param_attr_dict[param_name][src_rank] = (param.requires_grad, param.layerwise_parallel)
+    for local_rank_id in range(dst_stage_device_num):
+        transform_param_list = _transform_parallel_checkpoint(local_rank_id, param_total_dict,
+                                                              param_attr_dict, src_strategy_list, dst_strategy_list,
+                                                              param_type_dict)
+        save_checkpoint_file = "{}{}_part{}.ckpt".format(ckpt_prefix, local_rank_id, stage_id)
+        save_checkpoint_file_dir = os.path.join(dst_checkpoints_dir, "rank_{}".format(local_rank_id))
+        if not os.path.exists(save_checkpoint_file_dir):
+            _make_dir(save_checkpoint_file_dir, "path")
+        save_checkpoint_file_name = os.path.join(save_checkpoint_file_dir, save_checkpoint_file)
+        ms.save_checkpoint(transform_param_list, save_checkpoint_file_name)
+
+
+def _transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix, src_strategy_file=None,
+                           dst_strategy_file=None):
+    """Transform checkpoints for all stages in src_strategy_file"""
     checkpoints_rank_dir_list = os.path.join(src_checkpoints_dir, "rank_[0-9]*")
     all_checkpoint_files_map = {}
     for checkpoint_dir in glob.glob(checkpoints_rank_dir_list):
@@ -336,3 +350,135 @@ def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
             ms.save_checkpoint(transform_param_list, save_checkpoint_file_name)
             del param_total_dict_copy
         del param_total_dict
+
+
+def transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix, src_strategy_file=None,
+                          dst_strategy_file=None):
+    """
+    Transform distributed checkpoint from source sharding strategy to destination sharding strategy for a rank.
+    For more details about converting distributed Checkpoint, please refer to
+    `Model Transformation <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/model_transformation.html>`_.
+
+    Note:
+        The `src_checkpoints_dir` directory structure should be organized like "src_checkpoints_dir/rank_0/a.ckpt", the
+        rank number should be set to a subdirectory and the checkpoint file is stored in this subdirectory. If multiple
+        files exist in a rank directory, the last file in the lexicgraphic order would be selected.
+
+    Args:
+        src_checkpoints_dir (str): The source checkpoints directory.
+        dst_checkpoints_dir (str): The destination checkpoints directory to save the converted checkpoints.
+        ckpt_prefix (str): The destination checkpoint name prefix.
+        src_strategy_file (str): Name of source sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'src_strategy_file' is None, it means that the source sharding strategy is
+                                 without any sharing for each parameter. Default:None.
+        dst_strategy_file (str): Name of destination sharding strategy file which saved by
+                                 'mindspore.set_auto_parallel_context(strategy_ckpt_save_file)'.
+                                 when the 'dst_strategy_file' is None, it means that the destination sharding strategy
+                                 is without any sharing for each parameter. Default:None.
+
+    Raises:
+        ValueError: `src_strategy_file` or `dst_strategy_file` is incorrect.
+        NotADirectoryError: `src_checkpoints_dir` or `dst_checkpoints_dir` is not a directory.
+        ValueError: The checkpoint file is missing in `src_checkpoints_dir`.
+        TypeError: `src_strategy_file` or `dst_strategy_file` is not a string.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> ms.transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, "dst_checkpoint",
+        ...                       "./src_strategy.ckpt", "./dst_strategy.ckpt")
+
+    """
+    if not os.path.isdir(src_checkpoints_dir):
+        raise NotADirectoryError("src_checkpoints_dir {} is not a directory.".format(src_checkpoints_dir))
+    _make_dir(dst_checkpoints_dir, "path")
+    if not isinstance(ckpt_prefix, str):
+        raise TypeError("The ckpt_prefix should be a str.")
+    if src_strategy_file and os.path.dirname(src_strategy_file) and not os.path.exists(
+            os.path.dirname(src_strategy_file)):
+        raise ValueError("The director of src_strategy_file: {} is not exists.".
+                         format(os.path.dirname(src_strategy_file)))
+    if dst_strategy_file and os.path.dirname(dst_strategy_file) and not os.path.exists(
+            os.path.dirname(dst_strategy_file)):
+        raise ValueError("The director of dst_strategy_file: {} is not exists.".
+                         format(os.path.dirname(dst_strategy_file)))
+    src_layout_map = _extract_layout_map(src_strategy_file)
+    dst_layout_map = _extract_layout_map(dst_strategy_file)
+    pipeline_stage_num = _extract_pipeline_stage_num(src_strategy_file)
+    if src_layout_map:
+        src_param_keys = {param_name for param_name in src_layout_map if not param_name.startswith("accu_grads")}
+    if dst_layout_map:
+        dst_param_keys = {param_name for param_name in dst_layout_map if not param_name.startswith("accu_grads")}
+    if src_layout_map and dst_layout_map and pipeline_stage_num == 1 \
+        and src_param_keys.issubset(dst_param_keys) and len(src_param_keys) < len(dst_param_keys):
+        dst_stage_num = _extract_pipeline_stage_num(dst_strategy_file)
+        if dst_stage_num > 1:
+            raise NotImplementedError("When using unmerged src strategy, dst strategy doesn't \
+                                       support strategy with pipeline parallel.")
+        _transform_checkpoint_by_stage(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
+                                       src_strategy_file, dst_strategy_file)
+    else:
+        _transform_checkpoints(src_checkpoints_dir, dst_checkpoints_dir, ckpt_prefix,
+                               src_strategy_file, dst_strategy_file)
+
+
+def load_segmented_checkpoints(ckpt_file_dir, net=None, strict_load=False, filter_prefix=None,
+                               dec_key=None, dec_mode="AES-GCM", specify_prefix=None, choice_func=None):
+    """
+    Load checkpoint info from a specified file. If the specified ckpt_file_dir path contains multiple
+    checkpoint files, all checkpoint files will be loaded one by one and the combined dictionary will be return.
+
+    Note:
+        - `specify_prefix` and `filter_prefix` do not affect each other.
+        - If none of the parameters are loaded from checkpoint file, it will throw ValueError.
+        - `specify_prefix` and `filter_prefix` are in the process of being deprecated,
+          `choice_func` is recommended instead.
+          And using either of those two args will override `choice_func` at the same time.
+
+    Args:
+        ckpt_file_dir (str): Checkpoint file directory.
+        net (Cell): The network where the parameters will be loaded. Default: ``None`` .
+        strict_load (bool): Whether to strict load the parameter into net. If ``False`` , it will load parameter
+                            into net when parameter name's suffix in checkpoint file is the same as the
+                            parameter in the network. When the types are inconsistent perform type conversion
+                            on the parameters of the same type, such as float32 to float16. Default: ``False`` .
+        filter_prefix (Union[str, list[str], tuple[str]]): Deprecated(see `choice_func`). Parameters starting with the
+            filter_prefix will not be loaded. Default: ``None`` .
+        dec_key (Union[None, bytes]): Byte type key used for decryption. If the value is ``None`` , the decryption
+                                      is not required. Default: ``None`` .
+        dec_mode (str): This parameter is valid only when dec_key is not set to ``None`` . Specifies the decryption
+                        mode, currently supports ``"AES-GCM"`` and ``"AES-CBC"`` and ``"SM4-CBC"`` .
+                        Default: ``"AES-GCM"`` .
+        specify_prefix (Union[str, list[str], tuple[str]]): Deprecated(see `choice_func`). Parameters starting with the
+            specify_prefix will be loaded. Default: ``None`` .
+        choice_func (Union[None, function]) : Input value of the function is a Parameter name of type string,
+            and the return value is a bool. If returns ``True`` , the Parameter
+            that matches the custom condition will be loaded. If returns ``False`` , the Parameter that
+            matches the custom condition will be removed. Default: ``None`` .
+
+    Returns:
+        Dict, key is parameter name, value is a Parameter or string. When the `append_dict` parameter of
+        :func:`mindspore.save_checkpoint` and the `append_info` parameter of :class:`mindspore.train.CheckpointConfig`
+        are used to save the checkpoint, `append_dict` and `append_info` are dict types, and their value are string,
+        then the return value obtained by loading checkpoint is string, and in other cases the return value is
+        Parameter.
+
+    Raises:
+        TypeError: Input ckpt_file_dir is not a string.
+        ValueError: Checkpoint file directory doesn't exist. Or it's not a directory
+        ValueError: Checkpoint file's format is incorrect.
+        ValueError: Parameter's dict is None after load checkpoint file.
+        TypeError: The type of `specify_prefix` or `filter_prefix` is incorrect.
+    """
+    if not isinstance(ckpt_file_dir, str):
+        raise TypeError("The ckpt_file_dir should be a str.")
+    if not os.path.isdir(ckpt_file_dir):
+        raise ValueError("The dst_strategy_file: {} doesn't exists. Or it's not a directory.".
+                         format(os.path.dirname(ckpt_file_dir)))
+    checkpoint_file_name = os.path.join(ckpt_file_dir, "*.ckpt")
+    rank_ckpts = glob.glob(checkpoint_file_name)
+    parameter_dict = {}
+    for checkpoint_file in rank_ckpts:
+        parameter_dict.update(ms.load_checkpoint(checkpoint_file, net, strict_load, filter_prefix, dec_key,
+                                                 dec_mode, specify_prefix, choice_func))
+    return parameter_dict

mindspore/profiler/parser/ascend_cluster_generator.py

@@ -0,0 +1,111 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""trace step time analyse model"""
+import csv
+import fnmatch
+import json
+import logging
+import os
+import stat
+
+import numpy as np
+from mindspore.profiler.common.exceptions.exceptions import ProfilerIOException
+
+
+def find_files(directory, pattern):
+    """Find files with feature 'pattern' from the directory"""
+    file_list = []
+    for root, _, files in os.walk(directory):
+        files.sort(key=lambda x: os.path.getctime(os.path.join(directory, x)))
+        for basename in files:
+            if fnmatch.fnmatch(basename, pattern):
+                filename = os.path.join(root, basename)
+                file_list.append(filename)
+    return file_list
+
+
+class AscendClusterGenerator:
+    """Generate step trace time data from msprof*.json"""
+
+    def __init__(self, source_path):
+        self.root_path = source_path
+        self.msprof_data = np.array([])
+        self.step_trace_time = {'Step': None, 'Computing': 0, 'comunNotOverlp': 0, 'Overlapped': 0, 'Communication': 0,
+                                'Free': 0, 'Stage': 0, 'Bubble': 0, 'comunNotOverlpRec': 0}
+        self.msprof_data_df = np.dtype([('name', object), ('ts', float), ('dur', float)])
+        self.trace_step_time_df = np.dtype(
+            [('Step', int), ('Computing', float), ('comunNotOverlp', float), ('Communication', float), ('Free', float),
+             ('Stage', float), ('Bubble', float), ('comunNotOverlpRec', float)])
+        self.title = ['Step', 'Computing', 'Communication(Not Overlapped)', 'Overlapped', 'Communication', 'Free',
+                      'Stage', 'Bubble', 'Communication(Not Overlapped and Exclude Receive)']
+
+    def parse(self):
+        """
+        Analyse msprof json generate cluster data.
+        """
+        self.read_msprof()
+        if self.msprof_data.size < 1:
+            return
+        self.step_trace_time['Computing'] = np.sum(self.msprof_data[self.msprof_data['name'] == 'Computing']['dur'])
+        self.step_trace_time['comunNotOverlp'] = np.sum(
+            self.msprof_data[self.msprof_data['name'] == 'Communication(Not Overlapped)']['dur'])
+        self.step_trace_time['Communication'] = np.sum(
+            self.msprof_data[self.msprof_data['name'] == 'Communication']['dur'])
+        self.step_trace_time['Free'] = np.sum(self.msprof_data[self.msprof_data['name'] == 'Free']['dur'])
+        self.step_trace_time['Bubble'] = np.sum(
+            self.msprof_data[np.char.find(self.msprof_data['name'].astype('str'), '/Receive-op')]['dur'])
+
+        self.step_trace_time['Overlapped'] = self.step_trace_time['Communication'] - self.step_trace_time[
+            'comunNotOverlp']
+        self.step_trace_time['Stage'] = np.max(self.msprof_data['ts'] + self.msprof_data['dur']) - np.min(
+            self.msprof_data['ts']) - self.step_trace_time['Bubble']
+        self.step_trace_time['comunNotOverlpRec'] = self.step_trace_time['comunNotOverlp'] - self.step_trace_time[
+            'Bubble']
+
+    def read_msprof(self):
+        """
+        read msprof json information into memory.
+        """
+        msprof_data = []
+        for file in find_files(self.root_path, "msprof_*.json"):
+            with open(file) as jsonfile:
+                for row in json.load(jsonfile):
+                    if row.get('name') in ['Computing', 'Communication', 'Communication(Not Overlapped)',
+                                           'Free'] or row.get('name').find('/Receive-op'):
+                        name = row.get('name', '')
+                        ts = row.get('ts', 0)
+                        dur = row.get('dur', 0)
+                        msprof_data.append(tuple([name, ts, dur]))
+        self.msprof_data = np.array(msprof_data, dtype=self.msprof_data_df)
+
+    def write(self, step_trace_time_path):
+        """
+        Write the step trace time csv.
+
+        Args:
+            step_trace_time_path(str): step_trace_time.csv path.
+        """
+        try:
+            with os.fdopen(os.open(step_trace_time_path,
+                                   os.O_WRONLY | os.O_CREAT | os.O_TRUNC, stat.S_IWUSR | stat.S_IRUSR),
+                           'w') as step_trace_time:
+                writer = csv.writer(step_trace_time)
+                writer.writerow(self.title)
+                writer.writerow([v for _, v in self.step_trace_time.items()])
+        except (IOError, OSError) as err:
+            logging.critical('Error occurred when write step trace time file: %s', err)
+            raise ProfilerIOException() from err
+        if os.path.exists(step_trace_time_path):
+            os.chmod(step_trace_time_path, stat.S_IREAD | stat.S_IWRITE)