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

mindspore/train/serialization.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2024 Huawei Technologies Co., Ltd
+# Copyright 2025 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.
@@ -36,13 +36,12 @@ from functools import partial
 import math
 import sys
 import time
-import numpy as np
 from safetensors.numpy import save_file
+import numpy as np
 import google
 
 from mindspore.train.checkpoint_pb2 import Checkpoint
 from mindspore.train.mind_ir_pb2 import ModelProto as mindir_model
-from mindspore.train.print_pb2 import Print
 
 import mindspore
 import mindspore.nn as nn
@@ -52,15 +51,13 @@ from mindspore.log import vlog_print
 from mindspore._checkparam import check_input_data, check_input_dataset
 from mindspore import _checkparam as Validator
 from mindspore.common import dtype as mstype
-from mindspore.common import np_dtype
 from mindspore.common.api import _cell_graph_executor as _executor
 from mindspore.common.api import _JitExecutor
 from mindspore.common.api import _get_parameter_layout
-from mindspore.common.initializer import initializer, One
+from mindspore.common.initializer import initializer
 from mindspore.common.parameter import Parameter, _offload_if_config
 from mindspore.common.tensor import Tensor
 from mindspore._c_expression import TensorPy as Tensor_
-from mindspore.common._utils import is_shape_unknown
 from mindspore.common.file_system import FileSystem, _register_basic_file_system, _register_mindio_file_system
 from mindspore.communication.management import get_rank, get_group_size
 from mindspore.experimental import MapParameter
@@ -76,9 +73,9 @@ from mindspore.parallel.checkpoint_transform import load_distributed_checkpoint
 from mindspore.parallel.checkpoint_transform import merge_sliced_parameter as new_merge_sliced_parameter
 from mindspore.parallel.checkpoint_transform import build_searched_strategy as new_build_searched_strategy
 from mindspore.parallel.transform_safetensors import _fast_safe_open
-from mindspore.train._utils import read_proto, get_parameter_redundancy, _progress_bar, _load_and_transform
+from mindspore.train._utils import get_parameter_redundancy, _progress_bar, _load_and_transform
 from mindspore._c_expression import load_mindir, _encrypt, _decrypt, _is_cipher_file, \
-    split_mindir, split_dynamic_mindir
+    split_mindir, split_dynamic_mindir, _get_snapshot_params
 from mindspore.common.generator import Generator
 
 tensor_to_ms_type = {"Int8": mstype.int8, "UInt8": mstype.uint8, "Int16": mstype.int16, "UInt16": mstype.uint16,
@@ -86,12 +83,9 @@ tensor_to_ms_type = {"Int8": mstype.int8, "UInt8": mstype.uint8, "Int16": mstype
                      "Float16": mstype.float16, "Float32": mstype.float32, "Float64": mstype.float64,
                      "Bool": mstype.bool_, "str": mstype.string, "BFloat16": mstype.bfloat16, "Int4": mstype.qint4x2}
 
-tensor_to_np_type = {"Int8": np.int8, "UInt8": np.uint8, "Int16": np.int16, "UInt16": np.uint16,
-                     "Int32": np.int32, "UInt32": np.uint32, "Int64": np.int64, "UInt64": np.uint64,
-                     "Float16": np.float16, "Float32": np.float32, "Float64": np.float64, "Bool": np.bool_, "str": "U"}
-
-if hasattr(np_dtype, "bfloat16"):
-    tensor_to_np_type["BFloat16"] = np_dtype.bfloat16
+_tensor_to_np_type = {"Int8": np.int8, "UInt8": np.uint8, "Int16": np.int16, "UInt16": np.uint16,
+                      "Int32": np.int32, "UInt32": np.uint32, "Int64": np.int64, "UInt64": np.uint64,
+                      "Float16": np.float16, "Float32": np.float32, "Float64": np.float64, "Bool": np.bool_, "str": "U"}
 
 np_type_convert = {"int32": np.int32, "float32": np.float32, "float16": np.float16, "float64": np.float64}
 
@@ -114,6 +108,21 @@ INT_64_MAX = 9223372036854775807
 cpu_cast = Cast().set_device("CPU")
 
 _ckpt_fs = FileSystem()
+_ckpt_fs_initialized = False
+
+
+def tensor_to_np_type(tensor_type_str):
+    """tensor to numpy type"""
+    if tensor_type_str == "BFloat16":
+        from mindspore.common import np_dtype
+        if not np_dtype.np_dtype_valid(True):
+            raise TypeError(
+                "The Numpy bfloat16 data type is not supported now, please ensure that the current "
+                "Numpy version is not less than the version when the mindspore is compiled, "
+                "and the major versions are same."
+            )
+        return np_dtype.bfloat16
+    return _tensor_to_np_type.get(tensor_type_str)
 
 
 def init_ckpt_file_system(fs: FileSystem):
@@ -123,8 +132,12 @@ def init_ckpt_file_system(fs: FileSystem):
     _register_basic_file_system(fs)
 
 
-# Initialize checkpoint file system
-init_ckpt_file_system(_ckpt_fs)
+def _ensure_ckpt_fs_initialized():
+    """Ensure checkpoint file system is initialized"""
+    global _ckpt_fs_initialized
+    if not _ckpt_fs_initialized:
+        init_ckpt_file_system(_ckpt_fs)
+        _ckpt_fs_initialized = True
 
 
 def _wait_async_process_save_ckpt():
@@ -401,9 +414,6 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
                                                                                 crc_num, crc_check,
                                                                                 ckpt_total_io_time)
                             continue
-                        if isinstance(value[2], Tensor) and hasattr(value[2], "slice_num") and value[2].slice_num > 1:
-                            _write_hugeparameter(name, value, f)
-                            continue
 
                         crc_num, ckpt_total_io_time = _write_parameter_bytes_data(name, value, f, enc_key, plain_data,
                                                                                   crc_num, crc_check,
@@ -458,7 +468,7 @@ def _exec_save(ckpt_file_name, data_list, enc_key=None, enc_mode="AES-GCM", map_
                                f"simultaneously modified a file.")
             elif _ckpt_fs.backend != "mindio":
                 os.rename(tmp_name, ckpt_file_name)
-            os.chmod(ckpt_file_name, stat.S_IRUSR)
+                os.chmod(ckpt_file_name, stat.S_IRUSR)
     except BaseException as e:
         logger.critical("Failed to save the checkpoint file %s. Maybe don't have the permission to write files, "
                         "or the disk space is insufficient and so on.", ckpt_file_name)
@@ -546,27 +556,6 @@ def _write_mapparameter(name, value, f, map_param_inc=False):
             break
 
 
-def _write_hugeparameter(name, value, f):
-    """Write huge parameter into protobuf file."""
-    slice_num = value[2].slice_num
-    offset = 0
-    max_size = value[0][0]
-    for param_slice in range(slice_num):
-        checkpoint_list = Checkpoint()
-        param_value = checkpoint_list.value.add()
-        param_value.tag = name
-        param_tensor = param_value.tensor
-        param_tensor.dims.extend(value[0])
-        param_tensor.tensor_type = value[1]
-        param_key = value[3]
-        numpy_data = value[2].asnumpy_of_slice_persistent_data(param_key, param_slice)
-        if offset + numpy_data.shape[0] > max_size:
-            numpy_data = numpy_data[:max_size - offset]
-        param_tensor.tensor_content = numpy_data.tobytes()
-        f.write(checkpoint_list.SerializeToString())
-        offset += numpy_data.shape[0]
-
-
 def _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name, format):
     """Check save_obj and ckpt_file_name for save_checkpoint."""
     if format not in ["safetensors", "ckpt"]:
@@ -718,6 +707,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
           <https://mindspore.cn/tutorials/en/master/beginner/save_load.html#saving-and-loading-the-model-weight>`_
     """
     start_save_time = time.time()
+    _ensure_ckpt_fs_initialized()
     ckpt_file_name = _check_save_obj_and_ckpt_file_name(save_obj, ckpt_file_name, format)
     integrated_save = Validator.check_bool(integrated_save)
     async_save = _check_async_save(async_save)
@@ -767,9 +757,7 @@ def save_checkpoint(save_obj, ckpt_file_name, integrated_save=True,
                 data_list[param["name"]].append(param["data"])
                 continue
             if isinstance(param["data"], list):
-                if param["data"][0] == "persistent_data":
-                    _save_param_list_data(data_list, key, param)
-                elif param["data"][0] == "offload_parameter":
+                if param["data"][0] == "offload_parameter":
                     data_list[key].append("offload_parameter")
                     _save_param_list_data(data_list, key, param)
 
@@ -955,6 +943,8 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
     if not is_parallel_mode:
         save_obj.init_parameters_data()
     param_dict = _convert_cell_param_and_names_to_dict(save_obj, choice_func, is_parallel_mode)
+    enable_ckpt_d2h_sync = os.getenv('MS_ENABLE_D2H_ASYNC') == '1'
+    param_snapshot = _get_snapshot_params() if enable_ckpt_d2h_sync else {}
     for (key, value) in param_dict.items():
         each_param = {"name": key}
         if isinstance(value, MapParameter):
@@ -962,10 +952,7 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
             param_list.append(each_param)
             continue
 
-        if value.data.is_persistent_data():
-            # list save persistent_data: [Tensor, shape, type, param.key]
-            param_data = ["persistent_data", value.data, value.param_info.origin_shape, str(value.dtype), value.key]
-        elif value.data.offload_file_path() != "":
+        if value.data.offload_file_path() != "":
             # list save offload data: [Param, shape, type, param.key]
             param_data = ["offload_parameter"]
             param_tensor = value.data
@@ -980,7 +967,8 @@ def _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_f
             if append_dict and "__exception_save__" in append_dict:
                 param_data = Tensor(Tensor_.move_to(value, "CPU", False))
             else:
-                param_data = Tensor(value.data)
+                # when enable MS_ENABLE_D2H_ASYNC=1, fetch param from sanpshot in priority
+                param_data = param_snapshot.get(key, Tensor(value.data))
 
             # in automatic model parallel scenario, some parameters were split to all the devices,
             # which should be combined before saving
@@ -1004,13 +992,16 @@ def _convert_save_obj_to_param_list(save_obj, integrated_save, append_dict, choi
 
         return _handle_shared_param_for_pipeline_parallel(save_obj)
 
-    return _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_func)
+    if isinstance(save_obj, nn.Cell):
+        return _convert_cell_to_param_list(save_obj, integrated_save, append_dict, choice_func)
+
+    raise TypeError("For 'save_checkpoint', the argument 'save_obj' must be list、dict or nn.cell, "
+                    "but got {}.".format(type(save_obj)))
 
 
 def _save_param_list_data(data_list, key, param):
     """Save persistent data into save_obj."""
     dims = []
-    # persistent_data shape can not be ()
     for dim in param['data'][2]:
         dims.append(dim)
     data_list[key].append(dims)
@@ -1268,11 +1259,7 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                 continue
             data = element.tensor.tensor_content
             data_type = element.tensor.tensor_type
-            np_type = tensor_to_np_type.get(data_type)
             ms_type = tensor_to_ms_type[data_type]
-            if data_type == 'str':
-                str_length = int(len(data) / 4)
-                np_type = np_type + str(str_length)
             param_data_list.append(data)
             if (element_id == len(checkpoint_list.value) - 1) or \
                     (element.tag != checkpoint_list.value[element_id + 1].tag):
@@ -1280,6 +1267,8 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                 param_data_list.clear()
                 dims = element.tensor.dims
                 if data_type == 'str':
+                    str_length = int(len(data) / 4)
+                    np_type = "U" + str(str_length)
                     str_value = np.frombuffer(new_data, np_type)
                     parameter_dict[element.tag] = str(str_value[0])
                 else:
@@ -1288,7 +1277,6 @@ def _load_into_param_dict(ckpt_file_name, parameter_dict, specify_prefix, filter
                     param_data = Tensor_.convert_bytes_to_tensor(new_data, tuple(dims), ms_type)
                     parameter = Parameter(param_data, name=element.tag)
                     parameter_dict[element.tag] = parameter
-                    _offload_if_config(parameter)
 
         logger.info("Loading checkpoint files process is finished.")
         return remove_redundancy
@@ -1386,6 +1374,7 @@ def load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=N
     """
     start_load_time = time.time()
     vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin load checkpoint.")
+    _ensure_ckpt_fs_initialized()
     specify_prefix = _check_prefix(specify_prefix)
     filter_prefix = _check_prefix(filter_prefix)
     dec_key = Validator.check_isinstance('dec_key', dec_key, (type(None), bytes))
@@ -2133,6 +2122,7 @@ def _export(net, file_name, file_format, *inputs, **kwargs):
     if file_format == 'AIR':
         _save_air(net, file_name, *inputs, **kwargs)
     elif file_format == 'ONNX':
+        logger.warning("mindspore.export(file_format='ONNX') will be deleted, please use mindspore.onnx.export()")
         _save_onnx(net, file_name, *inputs, **kwargs)
     elif file_format == 'MINDIR':
         _save_mindir(net, file_name, *inputs, **kwargs)
@@ -2198,6 +2188,11 @@ def _save_onnx(net, file_name, *inputs, **kwargs):
         file_name += ".onnx"
     if os.path.exists(file_name):
         os.chmod(file_name, stat.S_IWUSR)
+    else:
+        dir_path = os.path.dirname(file_name)
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path, mode=0o700, exist_ok=True)
+        os.chmod(dir_path, 0o700)
     with open(file_name, 'wb') as f:
         f.write(onnx_stream)
         os.chmod(file_name, stat.S_IRUSR)
@@ -2477,147 +2472,6 @@ def _save_dataset_to_mindir(model, dataset):
             model.preprocessor.op[-1].offload = op['offload'] if 'offload' in op.keys() else False
 
 
-def check_checkpoint(ckpt_file_name):
-    """
-    Check whether the checkpoint is valid.
-
-    Note:
-        The interface is deprecated from version 2.5 and will be removed in a future version.
-
-    Args:
-        ckpt_file_name (str): Checkpoint file name.
-
-    Returns:
-        bool, whether the checkpoint is valid.
-
-    Examples:
-        >>> import mindspore as ms
-        >>> ckpt_file_name = "./checkpoint/LeNet5-1_32.ckpt"
-        >>> check_result = ms.check_checkpoint(ckpt_file_name)
-        >>> print(check_result)
-        True
-    """
-    logger.warning("The interface 'mindspore.check_checkpoint' is deprecated from version 2.5 "
-                   "and will be removed in a future version.")
-    if not ckpt_file_name.endswith('.ckpt'):
-        return False
-    checkpoint_list = Checkpoint()
-    with _ckpt_fs.open(ckpt_file_name, *_ckpt_fs.open_args) as f:
-        pb_content = f.read()
-        if pb_content[-17:-10] == b"crc_num":
-            crc_num_bytes = pb_content[-10:]
-            pb_content = pb_content[:-17]
-            crc_num = int.from_bytes(crc_num_bytes, byteorder='big')
-            cal_crc_num = binascii.crc32(pb_content, 0)
-            if cal_crc_num != crc_num:
-                logger.warning("For 'check_checkpoint', the ckpt crc check is failed.")
-                return False
-        try:
-            checkpoint_list.ParseFromString(pb_content)
-        except google.protobuf.message.DecodeError as e:
-            logger.warning("For 'check_checkpoint', the ckpt parse is failed.")
-            logger.warning(e)
-            return False
-        return True
-
-
-def parse_print(print_file_name):
-    """
-    Parse data file generated by :class:`mindspore.ops.Print`.
-
-    Note:
-        The interface is deprecated from version 2.5 and will be removed in a future version.
-
-    Args:
-        print_file_name (str): The file name needs to be parsed.
-
-    Returns:
-        List, element of list is Tensor.
-
-    Raises:
-        ValueError: The print file does not exist or is empty.
-        RuntimeError: Failed to parse the file.
-
-    Examples:
-        >>> import numpy as np
-        >>> import mindspore as ms
-        >>> from mindspore import nn, Tensor, ops
-        >>> ms.set_context(mode=ms.GRAPH_MODE, print_file_path='log.data')
-        >>> class PrintInputTensor(nn.Cell):
-        ...         def __init__(self):
-        ...             super().__init__()
-        ...             self.print = ops.Print()
-        ...
-        ...         def construct(self, input_pra):
-        ...             self.print('print:', input_pra)
-        ...             return input_pra
-        >>> x = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]).astype(np.float32)
-        >>> input_pra = Tensor(x)
-        >>> net = PrintInputTensor()
-        >>> net(input_pra)
-        >>>
-        >>> data = ms.parse_print('./log.data')
-        >>> print(data)
-        ['print:', Tensor(shape=[2, 4], dtype=Float32, value=
-        [[ 1.00000000e+00,  2.00000000e+00,  3.00000000e+00,  4.00000000e+00],
-        [ 5.00000000e+00,  6.00000000e+00,  7.00000000e+00,  8.00000000e+00]])]
-    """
-    logger.warning("The interface 'mindspore.parse_print' is deprecated from version 2.5 "
-                   "and will be removed in a future version.")
-    print_file_path = os.path.realpath(print_file_name)
-
-    if os.path.getsize(print_file_path) == 0:
-        raise ValueError("For 'parse_print', the print file may be empty, please make sure enter the correct "
-                         "'print_file_name'.")
-
-    logger.info("Execute load print process.")
-    print_list = Print()
-
-    try:
-        with open(print_file_path, "rb") as f:
-            pb_content = f.read()
-        print_list.ParseFromString(pb_content)
-    except BaseException as e:
-        logger.critical("Failed to read the print file %s, please check whether the file is "
-                        "correct.", print_file_name)
-        raise ValueError(e.__str__() + "\nFailed to read the print file {}, please check whether "
-                                       "the file is correct.".format(print_file_name)) from e
-
-    tensor_list = []
-
-    try:
-        for print_ in print_list.value:
-            # String type
-            if print_.HasField("desc"):
-                tensor_list.append(print_.desc)
-            elif print_.HasField("tensor"):
-                dims = print_.tensor.dims
-                data_type = print_.tensor.tensor_type
-                data = print_.tensor.tensor_content
-                np_type = tensor_to_np_type.get(data_type)
-                param_data = np.fromstring(data, np_type)
-                ms_type = tensor_to_ms_type.get(data_type)
-                if dims and dims != [0]:
-                    param_value = param_data.reshape(dims)
-                    tensor_list.append(Tensor(param_value, ms_type))
-                # Scalar type
-                else:
-                    data_type_ = data_type.lower()
-                    if 'float' in data_type_:
-                        param_data = float(param_data[0])
-                    elif 'int' in data_type_:
-                        param_data = int(param_data[0])
-                    elif 'bool' in data_type_:
-                        param_data = bool(param_data[0])
-                    tensor_list.append(Tensor(param_data, ms_type))
-
-    except BaseException as e:
-        logger.critical("Failed to load the print file %s.", print_list)
-        raise RuntimeError(e.__str__() + "\nFailed to load the print file {}.".format(print_list)) from e
-
-    return tensor_list
-
-
 def async_ckpt_thread_status():
     """
     Get the status of asynchronous save checkpoint thread.
@@ -2652,170 +2506,132 @@ def _calculation_net_size(net):
     return data_total
 
 
-def _get_mindir_inputs(file_name):
+def _load_file_and_convert_name(path, name_map=None, format="ckpt"):
     """
-    Get MindIR file's inputs.
-
-    Note:
-        1. Parsing encrypted MindIR file is not supported.
-        2. Parsing dynamic shape MindIR file is not supported.
+    Load file, during load convert name by name_map.
 
     Args:
-        file_name (str): MindIR file name.
+        path (str): The file path.
+        name_map (dict): Convert the name of parameter by name_map.
+        format (str): The format of the file. Option: 'ckpt', 'safetensors'
 
     Returns:
-        Tensor, list(Tensor), the input of MindIR file.
-
-    Raises:
-        TypeError: If the parameter file_name is not `str`.
-        RuntimeError: MindIR's input is not tensor type or has no dims.
-
-    Examples:
-        >>> input_tensor = get_mindir_inputs("lenet.mindir")
+        Dict, key is parameter name, value is a Parameter or string.
     """
-    Validator.check_file_name_by_regular(file_name)
-    file_name = os.path.realpath(file_name)
-    model = read_proto(file_name)
-    input_tensor = []
-
-    for ele_input in model.graph.input:
-        input_shape = []
-        if not hasattr(ele_input, "tensor") or not hasattr(ele_input.tensor[0], "dims"):
-            raise RuntimeError("MindIR's inputs has no tensor or tensor has no dims, please check MindIR file.")
-
-        for ele_shape in ele_input.tensor[0].dims:
-            input_shape.append(ele_shape)
-        if is_shape_unknown(input_shape):
-            raise RuntimeError(f"MindIR input's shape is: {input_shape}, dynamic shape is not supported.")
-
-        mindir_type = ele_input.tensor[0].data_type
-        if mindir_type not in mindir_to_tensor_type:
-            raise RuntimeError(f"MindIR input's type: {mindir_type} is not supported.")
-
-        input_type = mindir_to_tensor_type.get(mindir_type)
-        input_tensor.append(Tensor(shape=input_shape, dtype=input_type, init=One()))
-
-    if not input_tensor:
-        logger.warning("The MindIR model has no input, return None.")
-        return None
-    return input_tensor[0] if len(input_tensor) == 1 else input_tensor
-
-
-def convert_model(mindir_file, convert_file, file_format):
-    """
-    Convert mindir model to other format model. The current version only supports conversion to ONNX models.
-
-    Note:
-        The interface is deprecated from version 2.5 and will be removed in a future version.
-
-    Args:
-        mindir_file (str): MindIR file name.
-        convert_file (str): Convert model file name.
-        file_format (str): Convert model's format, current version only supports "ONNX".
-
-    Raises:
-        TypeError: If the parameter `mindir_file` is not `str`.
-        TypeError: If the parameter `convert_file` is not `str`.
-        ValueError: If the parameter `file_format` is not "ONNX".
-
-    Examples:
-        >>> import mindspore as ms
-        >>> ms.convert_model("lenet.mindir", "lenet.onnx", "ONNX")
-    """
-    logger.warning("The interface 'mindspore.train.serialization.convert_model' is deprecated from version 2.5 "
-                   "and will be removed in a future version.")
-    Validator.check_file_name_by_regular(mindir_file)
-    Validator.check_file_name_by_regular(convert_file)
-    if file_format != "ONNX":
-        raise ValueError(f"For 'convert_model', 'file_format' must be 'ONNX', but got {file_format}.")
-    net_input = _get_mindir_inputs(mindir_file)
-    graph = load(mindir_file)
-    net = nn.GraphCell(graph)
-    if isinstance(net_input, Tensor):
-        export(net, net_input, file_name=convert_file, file_format=file_format)
-    else:
-        export(net, *net_input, file_name=convert_file, file_format=file_format)
-
-
-def _load_ckpt_to_new_name_map(path, name_map=None):
-    return _load_and_transform(path, name_map, mindspore.load_checkpoint, None)
-
+    if name_map is not None:
+        load_func = partial(mindspore.load_checkpoint, format=format)
+        return _load_and_transform(path, name_map, load_func)
 
-def _load_sf_to_new_name_map(path, name_map=None):
-    load_func = partial(mindspore.load_checkpoint, format="safetensors")
-    return _load_and_transform(path, name_map, load_func, None)
+    return mindspore.load_checkpoint(path, format=format)
 
 
 def _process_file(file_info):
-    cur_ckpt_path, name_map, save_path, file = file_info
-    if name_map is not None:
-        param_dict = _load_ckpt_to_new_name_map(cur_ckpt_path, name_map)
+    """Rrocess file."""
+    cur_path, name_map, save_path, file, dst_format = file_info
+    if dst_format == "safetensors":
+        param_dict = _load_file_and_convert_name(cur_path, name_map, format="ckpt")
+        safetensors_filename = file.replace(".ckpt", ".safetensors")
+        dst_file = os.path.join(save_path, safetensors_filename)
+        mindspore.save_checkpoint(param_dict, dst_file, format='safetensors')
     else:
-        param_dict = mindspore.load_checkpoint(cur_ckpt_path)
-    safetensors_filename = file.replace(".ckpt", ".safetensors")
-    dst_file = os.path.join(save_path, safetensors_filename)
-    mindspore.save_checkpoint(param_dict, dst_file, format='safetensors')
+        param_dict = _load_file_and_convert_name(cur_path, name_map, format="safetensors")
+        ckpt_filename = file.replace(".safetensors", ".ckpt")
+        dst_file = os.path.join(save_path, ckpt_filename)
+        mindspore.save_checkpoint(param_dict, dst_file)
 
 
-def _process_file_safetensors(file_info):
-    cur_safe_path, name_map, save_path, file = file_info
-    if name_map is not None:
-        param_dict = _load_sf_to_new_name_map(cur_safe_path, name_map)
+def _gather_all_tasks(file_path, save_path, file_name_regex, name_map, dst_format="ckpt"):
+    """gather transform rank together"""
+    if dst_format == "ckpt":
+        cur_file_suffix = ".safetensors"
     else:
-        param_dict = mindspore.load_checkpoint(cur_safe_path, format="safetensors")
-    ckpt_filename = file.replace(".safetensors", ".ckpt")
-    dst_file = os.path.join(save_path, ckpt_filename)
-    mindspore.save_checkpoint(param_dict, dst_file)
+        cur_file_suffix = ".ckpt"
 
-
-def _gather_safetensors_tasks(file_path, save_path, file_name_regex, name_map):
-    """gather transform rank together"""
-    tasks = []
+    tasks_list = []
     for root, dirs, _ in os.walk(file_path):
         if root != file_path:
             continue
 
         rank_dirs = [d for d in dirs if d.startswith('rank')]
         if not rank_dirs:
-            raise ValueError(
-                f"For 'safetensors_to_ckpt', no directories starting with 'rank' found in {file_path}")
+            if dst_format == "safetensors":
+                raise ValueError(
+                    f"For 'ckpt_to_safetensors', no directories starting with 'rank' found in {file_path}.")
+            if dst_format == "ckpt":
+                raise ValueError(
+                    f"For 'safetensors_to_ckpt', no directories starting with 'rank' found in {file_path}.")
+
+            raise ValueError(f"For '_gather_all_tasks', error args 'format': {dst_format}.")
 
         for rank_dir in rank_dirs:
             rank_dir_path = os.path.join(root, rank_dir)
-            dst_root = os.path.join(save_path,
-                                    os.path.relpath(rank_dir_path, file_path)) if save_path else rank_dir_path
+            if save_path is not None:
+                dst_root = os.path.join(save_path, os.path.relpath(rank_dir_path, file_path))
+            else:
+                dst_root = rank_dir_path
+
             os.makedirs(dst_root, exist_ok=True)
-            tasks.extend(
-                (os.path.join(rank_dir_path, file), name_map, dst_root, file)
-                for file in os.listdir(rank_dir_path)
-                if file.endswith(".safetensors") and (file_name_regex is None or re.findall(file_name_regex, file))
-            )
-    return tasks
 
+            for file in os.listdir(rank_dir_path):
+                if file.endswith(cur_file_suffix) and (file_name_regex is None or re.search(file_name_regex, file)):
+                    tasks_list.append((os.path.join(rank_dir_path, file), name_map, dst_root, file, dst_format))
 
-def _gather_tasks_covert(file_path, save_path, file_name_regex, name_map):
-    """gather transform rank together"""
-    tasks = []
-    for root, dirs, _ in os.walk(file_path):
-        if root != file_path:
-            continue
+    return tasks_list
 
-        rank_dirs = [d for d in dirs if d.startswith('rank')]
-        if not rank_dirs:
-            raise ValueError(
-                f"For 'ckpt_to_safetensors', no directories starting with 'rank' found in {file_path}")
 
-        for rank_dir in rank_dirs:
-            rank_dir_path = os.path.join(root, rank_dir)
-            dst_root = os.path.join(save_path,
-                                    os.path.relpath(rank_dir_path, file_path)) if save_path else rank_dir_path
-            os.makedirs(dst_root, exist_ok=True)
-            tasks.extend(
-                (os.path.join(rank_dir_path, file), name_map, dst_root, file)
-                for file in os.listdir(rank_dir_path)
-                if file.endswith(".ckpt") and (file_name_regex is None or re.findall(file_name_regex, file))
-            )
-    return tasks
+def _convert_checkpoint_file(file_path, save_path=None, name_map=None, file_name_regex=None,
+                             processes_num=1, dst_format="safetensors"):
+    """
+    Converts MindSpore checkpoint files format and saves them to `save_path`.
+    Safetensors is a reliable and portable machine learning model storage format introduced by Huggingface,
+    used for securely storing Tensors with fast speed (zero copy).
+
+    Args:
+        file_path (str): Path to the directory containing checkpoint files or a single checkpoint file (.ckpt).
+        save_path (str, optional): Directory path where safetensors files will be saved. Default: ``None``.
+        name_map (dict, optional): Dictionary mapping original parameter names to new names. Default: ``None``.
+        file_name_regex (str, optional): Regular expression used to match the file that needs to be converted.
+                                       Default: ``None``.
+        processes_num (int, optional): Number of processes to use for parallel processing. Default: 1.
+        dst_format (str): dst file format. Default: "safetensors".
+    """
+    if dst_format == "safetensors":
+        src_format = "ckpt"
+        src_file_suffix = ".ckpt"
+        dst_file_suffix = ".safetensors"
+        func_name = "ckpt_to_safetensors"
+    else:
+        src_format = "safetensors"
+        src_file_suffix = ".safetensors"
+        dst_file_suffix = ".ckpt"
+        func_name = "safetensors_to_ckpt"
+    is_dir = os.path.isdir(file_path)
+    is_file = os.path.isfile(file_path)
+    if not is_dir and not is_file:
+        raise ValueError(f"For {func_name}, the input path must be a valid path or file, but got {file_path}")
+    if save_path and os.path.splitext(save_path)[1]:
+        raise ValueError(f"For {func_name}, the save_path must be a directory, but got '{save_path}'")
+    if name_map is not None and not isinstance(name_map, dict):
+        raise ValueError(
+            f"For {func_name}, the type of 'name_map' must be a directory, but got '{type(name_map)}'")
+
+    if is_dir:
+        tasks_list = _gather_all_tasks(file_path, save_path, file_name_regex, name_map, dst_format=dst_format)
+        with mp.Pool(processes=processes_num) as pool:
+            list(_progress_bar(pool.imap(_process_file, tasks_list), total=len(tasks_list)))
+    elif is_file:
+        if not file_path.endswith(src_file_suffix):
+            raise ValueError(f"For {func_name}, the input file must be a {src_file_suffix} file, but got {file_path}")
+        if file_name_regex is not None and not re.findall(file_name_regex, file_path):
+            raise ValueError(f"For {func_name}, the input file does not match the regular expression.")
+        if save_path and not os.path.exists(save_path):
+            os.makedirs(save_path, exist_ok=True)
+
+        param_dict = _load_file_and_convert_name(file_path, name_map, format=src_format)
+
+        file_filename = os.path.basename(file_path).replace(src_file_suffix, dst_file_suffix)
+        dst_file = os.path.join(save_path if save_path else os.path.dirname(file_path), file_filename)
+        mindspore.save_checkpoint(param_dict, dst_file, format=dst_format)
 
 
 def ckpt_to_safetensors(file_path, save_path=None, name_map=None, file_name_regex=None, processes_num=1):
@@ -2834,11 +2650,11 @@ def ckpt_to_safetensors(file_path, save_path=None, name_map=None, file_name_rege
 
     Args:
         file_path (str): Path to the directory containing checkpoint files or a single checkpoint file (.ckpt).
-        save_path (str, optional): Directory path where safetensors files will be saved. Defaults: ``None``.
-        name_map (dict, optional): Dictionary mapping original parameter names to new names. Defaults: ``None``.
+        save_path (str, optional): Directory path where safetensors files will be saved. Default: ``None``.
+        name_map (dict, optional): Dictionary mapping original parameter names to new names. Default: ``None``.
         file_name_regex (str, optional): Regular expression used to match the file that needs to be converted.
-                                   Defaults: ``None``.
-        processes_num (int, optional): Number of processes to use for parallel processing. Defaults: 1.
+                                   Default: ``None``.
+        processes_num (int, optional): Number of processes to use for parallel processing. Default: 1.
     Raises:
         ValueError: If the input path is invalid or the save_path is not a directory,
                     or the file_path does not end with '.ckpt'.
@@ -2854,36 +2670,8 @@ def ckpt_to_safetensors(file_path, save_path=None, name_map=None, file_name_rege
         >>> namemap = {"lin.weight":"new_name"}
         >>> ms.ckpt_to_safetensors("./ckpt_save_path/rank0/checkpoint_0.ckpt", "./new_path/", namemap)
     """
-    is_dir = os.path.isdir(file_path)
-    is_file = os.path.isfile(file_path)
-    if not is_dir and not is_file:
-        raise ValueError(f"For 'ckpt_to_safetensors', the input path must be a valid path or file, but got {file_path}")
-    if save_path and os.path.splitext(save_path)[1]:
-        raise ValueError(f"For 'ckpt_to_safetensors', the save_path must be a directory, but got '{save_path}'")
-    if name_map is not None and not isinstance(name_map, dict):
-        raise ValueError(
-            f"For 'ckpt_to_safetensors', the type of 'name_map' must be a directory, but got '{type(name_map)}'")
-
-    if is_dir:
-        tasks = _gather_tasks_covert(file_path, save_path, file_name_regex, name_map)
-        with mp.Pool(processes=processes_num) as pool:
-            list(_progress_bar(pool.imap(_process_file, tasks), total=len(tasks)))
-    elif is_file:
-        if not file_path.endswith(".ckpt"):
-            raise ValueError(f"For 'ckpt_to_safetensors', the input file must be a .ckpt file, but got {file_path}")
-        if file_name_regex is not None and not re.findall(file_name_regex, file_path):
-            raise ValueError(f"For 'ckpt_to_safetensors', the input file does not match the regular expression.")
-        if save_path and not os.path.exists(save_path):
-            os.makedirs(save_path, exist_ok=True)
-
-        if name_map is not None:
-            param_dict = _load_ckpt_to_new_name_map(file_path, name_map)
-        else:
-            param_dict = mindspore.load_checkpoint(file_path)
-
-        safetensors_filename = os.path.basename(file_path).replace(".ckpt", ".safetensors")
-        dst_file = os.path.join(save_path if save_path else os.path.dirname(file_path), safetensors_filename)
-        mindspore.save_checkpoint(param_dict, dst_file, format='safetensors')
+    _convert_checkpoint_file(file_path, save_path, name_map,
+                             file_name_regex, processes_num, "safetensors")
 
 
 def safetensors_to_ckpt(file_path, save_path=None, name_map=None, file_name_regex=None, processes_num=1):
@@ -2898,11 +2686,11 @@ def safetensors_to_ckpt(file_path, save_path=None, name_map=None, file_name_rege
 
     Args:
         file_path (str): Path to the directory containing safetensors files or a single safetensors file (.safetensors).
-        save_path (str, optional): Directory path where checkpoint files will be saved. Defaults: ``None``.
-        name_map (dict, optional): Dictionary mapping original parameter names to new names. Defaults: ``None``.
+        save_path (str, optional): Directory path where checkpoint files will be saved. Default: ``None``.
+        name_map (dict, optional): Dictionary mapping original parameter names to new names. Default: ``None``.
         file_name_regex (str, optional): Regular expression used to match the file that needs to be converted.
-                                   Defaults: ``None``.
-        processes_num (int, optional): Number of processes to use for parallel processing. Defaults: 1.
+                                   Default: ``None``.
+        processes_num (int, optional): Number of processes to use for parallel processing. Default: 1.
 
     Raises:
         ValueError: If the input path is invalid, the save_path is not a directory,
@@ -2919,37 +2707,8 @@ def safetensors_to_ckpt(file_path, save_path=None, name_map=None, file_name_rege
         >>> namemap = {"lin.weight":"new_name"}
         >>> ms.safetensors_to_ckpt("./safetensors_save_path/rank0/checkpoint_0.safetensors", "./new_path/", namemap)
     """
-    is_dir = os.path.isdir(file_path)
-    is_file = os.path.isfile(file_path)
-    if not is_dir and not is_file:
-        raise ValueError(f"For 'safetensors_to_ckpt', the input path must be a valid path or file, but got {file_path}")
-    if save_path and os.path.splitext(save_path)[1]:
-        raise ValueError(f"For 'safetensors_to_ckpt', the save_path must be a directory, but got '{save_path}'")
-    if name_map is not None and not isinstance(name_map, dict):
-        raise ValueError(
-            f"For 'safetensors_to_ckpt', the type of 'name_map' must be a directory, but got '{type(name_map)}'")
-
-    if is_dir:
-        tasks = _gather_safetensors_tasks(file_path, save_path, file_name_regex, name_map)
-        with mp.Pool(processes=processes_num) as pool:
-            list(_progress_bar(pool.imap(_process_file_safetensors, tasks), total=len(tasks)))
-    elif is_file:
-        if not file_path.endswith(".safetensors"):
-            raise ValueError(
-                f"For 'safetensors_to_ckpt', the input file must be a .safetensors file, but got {file_path}")
-        if file_name_regex is not None and not re.findall(file_name_regex, file_path):
-            raise ValueError(f"For 'safetensors_to_ckpt', the input file does not match the regular expression.")
-        if save_path and not os.path.exists(save_path):
-            os.makedirs(save_path, exist_ok=True)
-
-        if name_map is not None:
-            param_dict = _load_sf_to_new_name_map(file_path, name_map)
-        else:
-            param_dict = mindspore.load_checkpoint(file_path, format="safetensors")
-
-        ckpt_filename = os.path.basename(file_path).replace(".safetensors", ".ckpt")
-        dst_file = os.path.join(save_path if save_path else os.path.dirname(file_path), ckpt_filename)
-        mindspore.save_checkpoint(param_dict, dst_file)
+    _convert_checkpoint_file(file_path, save_path, name_map,
+                             file_name_regex, processes_num, "ckpt")
 
 
 def restore_group_info_list(group_info_file_name):