mindspore 2.4.10__cp310-cp310-manylinux1_x86_64.whl → 2.5.0__cp310-cp310-manylinux1_x86_64.whl

mindspore/parallel/transform_safetensors.py

@@ -16,16 +16,20 @@
 from __future__ import absolute_import
 
 import os
-import time
 import glob
-import re
 import math
 import json
+import re
 from collections import defaultdict
 
+import time
 import multiprocessing as mp
 import numpy as np
+from safetensors.numpy import save_file, load_file
+from safetensors import safe_open
+
 import mindspore as ms
+from mindspore import log as logger
 from mindspore.parallel._parallel_serialization import _get_device_num_from_strategy, _make_dir, \
     _extract_layout_map, _extract_src_dst_layout_map, _parameter_not_in_local_stage, _extract_pipeline_stage_num, \
     _insert_opt_shard_reshape, _extract_src_dst_layout_map_by_src
@@ -36,70 +40,6 @@ from mindspore.parallel._tensor import _get_tensor_strategy, _construct_from_to_
 from mindspore.parallel._parallel_serialization import _build_searched_strategy, _load_protobuf_strategy, \
     _convert_to_list
 
-from safetensors.numpy import save_file, load_file
-from safetensors import safe_open
-
-
-def _load_and_transform(path, name_map, load_func, transform_func):
-    if load_func is not None:
-        param_dict = load_func(path)
-    else:
-        param_dict = path
-    transform_dict = {}
-    for k, v in param_dict.items():
-        new_name = name_map.get(k, k) if name_map is not None else k
-        transform_dict[new_name] = transform_func(v, new_name)
-    return transform_dict
-
-
-def _transform_tensor_to_numpy(path, name_map=None):
-    return _load_and_transform(path, name_map, ms.load_checkpoint, lambda v, new_name: v.asnumpy())
-
-
-def _transform_numpy_to_tensor(path, name_map=None):
-    return _load_and_transform(path, name_map, load_file, lambda v, new_name: ms.Parameter(v, name=new_name))
-
-
-def _process_file(file_info):
-    cur_ckpt_path, name_map, save_path, file = file_info
-    param_dict_numpy = _transform_tensor_to_numpy(cur_ckpt_path, name_map)
-    safetensors_filename = file.replace(".ckpt", ".safetensors")
-    dst_file = os.path.join(save_path, safetensors_filename)
-    save_file(param_dict_numpy, dst_file)
-
-
-def _process_file_safetensors(file_info):
-    cur_safe_path, name_map, save_path, file = file_info
-    param_dict_tensor = _transform_numpy_to_tensor(cur_safe_path, name_map)
-    ckpt_filename = file.replace(".safetensors", ".ckpt")
-    dst_file = os.path.join(save_path, ckpt_filename)
-    ms.save_checkpoint(param_dict_tensor, dst_file)
-
-
-def _gather_tasks(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
-
-        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 _progress_bar(iterable, total=None):
     """
@@ -134,155 +74,16 @@ def _progress_bar(iterable, total=None):
         print_progress_bar(i)
 
 
-def ckpt_to_safetensors(file_path, save_path=None, name_map=None, file_name_regex=None, processes_num=1):
-    """
-    Converts MindSpore checkpoint files into safetensors 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).
-
-    Note:
-        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.
-        The safetensors format does not support the enc verification function. If ckpt is enabled to save enc
-        verification, an error will be generated when performing the conversion.
-        The safetensors format currently does not support crc verification function. If ckpt contains crc verification
-        information, the crc verification information will be lost after conversion to safetensors.
-
-    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``.
-        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.
-    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'.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-
-    Examples:
-        >>> import mindspore as ms
-        >>> ms.ckpt_to_safetensors("./ckpt_save_path")
-        >>> ms.ckpt_to_safetensors("./ckpt_save_path/rank0/checkpoint_0.ckpt")
-        >>> ms.ckpt_to_safetensors(file_path="./ckpt_save_path/rank0/checkpoint_0.ckpt", save_path="./new_path/")
-        >>> 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(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)
-
-        param_dict_numpy = _transform_tensor_to_numpy(file_path, name_map)
-        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)
-        save_file(param_dict_numpy, dst_file)
-
-
-def _gather_safetensors_tasks(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
-
-        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}")
-
-        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(".safetensors") and (file_name_regex is None or re.findall(file_name_regex, file))
-            )
-    return tasks
-
-
-def safetensors_to_ckpt(file_path, save_path=None, name_map=None, file_name_regex=None, processes_num=1):
-    """
-    Converts safetensors files into MindSpore checkpoint 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).
-
-    Note:
-        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.
-
-    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``.
-        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.
-
-    Raises:
-        ValueError: If the input path is invalid, the save_path is not a directory,
-                    or the file_path does not end with '.safetensors'.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-
-    Examples:
-        >>> import mindspore as ms
-        >>> ms.safetensors_to_ckpt("./safetensors_save_path")
-        >>> ms.safetensors_to_ckpt("./safetensors_save_path/rank0/checkpoint_0.safetensors")
-        >>> ms.safetensors_to_ckpt("./safetensors_save_path/rank0/checkpoint_0.safetensors", "./new_path/")
-        >>> 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)
-
-        param_dict_tensor = _transform_numpy_to_tensor(file_path, name_map)
-        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)
-        ms.save_checkpoint(param_dict_tensor, dst_file)
+def _load_and_transform(path, name_map, load_func, transform_func):
+    if load_func is not None:
+        param_dict = load_func(path)
+    else:
+        param_dict = path
+    transform_dict = {}
+    for k, v in param_dict.items():
+        new_name = name_map.get(k, k) if name_map is not None else k
+        transform_dict[new_name] = transform_func(v, new_name)
+    return transform_dict
 
 
 def _check_transform_safetensors(src_safetensors_dir, ckpt_prefix, src_strategy_file, dst_strategy_file):
@@ -460,7 +261,6 @@ def _transform_safetensors_with_parallel(needed_rank_list_map, all_safetensor_fi
 
         for name, layout in layout_map.items():
             pipe_param_list[layout[6][0]].append(name)
-
     part_list_dict = _distribute_files_by_size(all_safetensor_files_map, needed_rank_list_map, process_num)
     processes = []
     for i in range(process_num):
@@ -512,7 +312,7 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
                                   origin_dst_strategy_list,
                                   ckpt_prefix, dst_safetensors_dir, output_format,
                                   _transform_param_list, pipe_param_list=None, file_index=None, unified_flag=False,
-                                  src_strategy_file=None):
+                                  src_strategy_file=None, choice_func=None):
     """
     Transforms safetensors files to a specified format without using parallel processing.
     """
@@ -567,7 +367,18 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
                                 # param not in ckpt file, check reason
                                 continue
                             output = f.get_tensor(param_name)
-                            saftensor_dict[param_name] = output
+                            save_param_name = param_name
+                            if choice_func is not None:
+                                choice_out = choice_func(param_name)
+                                if isinstance(choice_out, bool):
+                                    if not choice_out:
+                                        continue
+                                elif isinstance(choice_out, str):
+                                    save_param_name = choice_out
+                                else:
+                                    raise ValueError("For 'unified_safetensors', the return value type of the function "
+                                                     f"'choice_func' must be bool or str, but got {type(choice_out)}.")
+                            saftensor_dict[save_param_name] = output
             else:
                 saftensor_dict = load_file(all_safetensor_files_map.get(int(needed_rank)))
             for param_name, param in saftensor_dict.items():
@@ -735,6 +546,12 @@ def transform_safetensors_by_rank(rank_id, safetensor_files_map, save_safetensor
     save_file(transform_param_dict, save_safetensor_file_name)
 
 
+def _extrace_number(file_name):
+    """get file last two number"""
+    number_ls = re.findall(r'\d+', file_name)
+    number_ls = [int(i) for i in number_ls]
+    return number_ls[-2:]
+
 def _collect_safetensor_files(src_safetensors_dir, format='safetensors', file_suffix=None):
     """
     Collects all safetensors files from the specified directory and its subdirectories.
@@ -758,12 +575,9 @@ def _collect_safetensor_files(src_safetensors_dir, format='safetensors', file_su
         else:
             safetensor_file_name = os.path.join(safetensor_dir, f"*{file_suffix}.{format}")
         rank_ckpts = glob.glob(safetensor_file_name)
-        rank_ckpts.sort()
-        for safetensor_file in rank_ckpts:
-            if not os.path.isfile(safetensor_file):
-                ms.log.warning("{} is not a safetensor file.".format(safetensor_file))
-                continue
-            all_safetensor_files_map[rank_id] = safetensor_file
+        rank_ckpts.sort(key=_extrace_number)
+        if rank_ckpts:
+            all_safetensor_files_map[rank_id] = rank_ckpts[-1]
     return all_safetensor_files_map
 
 
@@ -865,7 +679,8 @@ def _transform_parallel_safetensor(rank_id, param_total_dict, param_attr_dict, s
     return transform_param_dict
 
 
-def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundancy=True, file_suffix=None):
+def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundancy=True, file_suffix=None,
+                        max_process_num=64, choice_func=None):
     """
     Merge multiple safetensor files into a unified safetensor file.
 
@@ -877,6 +692,9 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
             saved safetensors files. Default: ``True``, indicating that the merged source weight files are complete.
         file_suffix (str, optional): Specify the filename suffix for merging safetensors files. Default: ``None``,
             meaning all safetensors files in the source weight directory will be merged.
+        max_process_num (int): Maximum number of processes. Default: 64.
+        choice_func (callable): A callable function used to filter parameters or modify parameter names.
+            The return value of the function must be of type str (string) or bool (boolean). Default: None.
 
     Raises:
         ValueError: If the safetensors file of rank is missing.
@@ -943,13 +761,21 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
     param_name_dict = dict()
     for index, part_list in enumerate(split_list):
         for name in part_list:
-            param_name_dict[name] = f"part{index}.safetensors"
+            save_param_name = name
+            if choice_func is not None:
+                choice_out = choice_func(name)
+                if isinstance(choice_out, bool):
+                    if not choice_out:
+                        continue
+                elif isinstance(choice_out, str):
+                    save_param_name = choice_out
+            param_name_dict[save_param_name] = f"part{index}.safetensors"
     json_str = json.dumps(param_name_dict, indent=4)
     map_file = os.path.join(dst_dir, "param_name_map.json")
     with open(map_file, 'w') as f:
         f.write(json_str)
 
-    max_process = min(split_num, 100)
+    max_process = min(split_num, max_process_num)
     res = [i for i in range(split_num)]
     res = _split_list(res, max_process)
     processes = []
@@ -960,7 +786,7 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
         p = mp.Process(target=_transform_safetensors_single_semaphore, args=(
             needed_rank_list_map, all_safetensor_files_map, src_stage_device_num, dst_stage_device_num,
             src_strategy_dict, None, origin_src_strategy_list, origin_dst_strategy_list,
-            "", dst_dir, "safetensors", None, split_list, res[i], True, src_strategy_name))
+            "", dst_dir, "safetensors", None, split_list, res[i], True, src_strategy_name, choice_func))
         p.start()
         processes.append(p)
     for p in processes:
@@ -974,13 +800,14 @@ def _transform_safetensors_single_semaphore(needed_rank_list_map, all_safetensor
                                             origin_dst_strategy_list,
                                             ckpt_prefix, dst_safetensors_dir, output_format,
                                             _transform_param_list, pipe_param_list=None, file_index=None,
-                                            unified_flag=False, src_strategy_file=None):
+                                            unified_flag=False, src_strategy_file=None, choice_func=None):
     for i in file_index:
         _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map, src_stage_device_num,
                                       dst_stage_device_num, src_strategy_dict, dst_strategy_dict,
                                       origin_src_strategy_list,
                                       origin_dst_strategy_list, ckpt_prefix, dst_safetensors_dir, output_format,
-                                      _transform_param_list, pipe_param_list[i], i, unified_flag, src_strategy_file)
+                                      _transform_param_list, pipe_param_list[i], i, unified_flag, src_strategy_file,
+                                      choice_func)
 
 
 def _split_list(split_list, split_num):
@@ -1027,24 +854,64 @@ def _apply_sf_obj_transform_operators(transform_operator_stack, sf_obj, device_n
     return sf_obj
 
 
-def _load_parallel_checkpoint(total_safetensors_dir, dst_strategy_file, net=None, dst_safetensors_dir=None,
-                              rank_id=None):
-    """load parallel safetensors by merged file."""
-    file_list = os.listdir(total_safetensors_dir)
-    json_files = [file for file in file_list if file.endswith('.json')]
-    if len(json_files) != 1:
-        raise ValueError(f"For 'load_parallel_checkpoint', the number of json files in 'total_safetensors_dir' "
-                         f"must be 1, but got {len(json_files)}.")
-    param_name_json = os.path.join(total_safetensors_dir, json_files[0])
-    with open(param_name_json, 'r') as f:
-        param_name_map = json.load(f)
+def _check_name_map_value_is_str(value):
+    """check input is bool"""
+    if not isinstance(value, str):
+        raise ValueError(
+            f"For 'load_distributed_checkpoint', the value of name_map must be str, but got {type(value)}.")
+
+
+def _process_hyper_params(file_list, total_safetensors_dir, name_map, total_param):
+    """process hyper params"""
+    if 'hyper_param.safetensors' in file_list:
+        hyper_parameter_file_name = os.path.join(total_safetensors_dir, "hyper_param.safetensors")
+        with safe_open(hyper_parameter_file_name, framework="np") as f:
+            for key in f.keys():
+                cur_param_name = name_map.get(key) if name_map is not None and key in name_map else key
+                _check_name_map_value_is_str(cur_param_name)
+                total_param[cur_param_name] = ms.Parameter(ms.Tensor.from_numpy(f.get_tensor(key)))
+    return total_param
+
+
+def _cal_param_name_map_and_param_list(file_list, total_safetensors_dir, json_files, dst_strategy_file, rank_id):
+    """calculate param_name_map and param_list"""
+    if len(file_list) == 1:
+        logger.info("There is only one weight file in the directory, which will be automatically mapped.")
+        file_name = os.path.join(total_safetensors_dir, file_list[0])
+        is_file = os.path.isfile(file_name)
+        if not is_file:
+            raise ValueError(f"For 'load_parallel_checkpoint', weight files must be included "
+                             f"in the `unified_safetensors_dir`.")
+        with safe_open(file_name, framework="np") as f:
+            keys = f.keys()
+            values = len(keys) * [file_list[0]]
+            param_name_map = dict(zip(keys, values))
+    else:
+        if len(json_files) != 1:
+            raise ValueError(f"For 'load_parallel_checkpoint', the number of json files in 'total_safetensors_dir' "
+                             f"must be 1, but got {len(json_files)}.")
+        param_name_json = os.path.join(total_safetensors_dir, json_files[0])
+        with open(param_name_json, 'r') as f:
+            param_name_map = json.load(f)
+
     if dst_strategy_file is not None:
         _, dst_strategy_list = _extract_src_dst_layout_map(rank_id, None, dst_strategy_file)
         param_list = dst_strategy_list.keys()
     else:
         dst_strategy_list = None
         param_list = param_name_map.keys()
+    return param_name_map, param_list, dst_strategy_list
 
+
+def _load_parallel_checkpoint(file_info):
+    """load parallel safetensors by merged file."""
+    total_safetensors_dir, dst_strategy_file, net, dst_safetensors_dir, \
+    rank_id, output_format, name_map, return_param_dict = file_info
+    file_list = os.listdir(total_safetensors_dir)
+    json_files = [file for file in file_list if file.endswith('.json')]
+    param_name_map, param_list, dst_strategy_list = _cal_param_name_map_and_param_list(file_list, total_safetensors_dir,
+                                                                                       json_files, dst_strategy_file,
+                                                                                       rank_id)
     total_param = dict()
     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
@@ -1102,20 +969,19 @@ def _load_parallel_checkpoint(total_safetensors_dir, dst_strategy_file, net=None
             slice_param = _apply_sf_obj_transform_operators(transform_operator_stack, sf_obj, device_num)
         else:
             slice_param = sf_obj[:]
+        cur_param_name = name_map.get(param_name) if name_map is not None and param_name in name_map else param_name
+        _check_name_map_value_is_str(cur_param_name)
+        total_param[cur_param_name] = ms.Parameter(ms.Tensor.from_numpy(slice_param))
 
-        total_param[param_name] = ms.Parameter(slice_param)
-
-    if 'hyper_param.safetensors' in file_list:
-        hyper_parameter_file_name = os.path.join(total_safetensors_dir, "hyper_param.safetensors")
-        with safe_open(hyper_parameter_file_name, framework="np") as f:
-            for key in f.keys():
-                total_param[key] = ms.Parameter(f.get_tensor(key))
+    total_param = _process_hyper_params(file_list, total_safetensors_dir, name_map, total_param)
     if net is not None:
-        param_not_load, ckpt_not_load = ms.load_param_into_net(net, total_param)
-        return param_not_load, ckpt_not_load
+        if not return_param_dict:
+            param_not_load, ckpt_not_load = ms.load_param_into_net(net, total_param)
+            return param_not_load, ckpt_not_load
+        return total_param
     _make_dir(os.path.join(dst_safetensors_dir, f"rank_{rank_id}"), "path")
-    ms.save_checkpoint(total_param, os.path.join(dst_safetensors_dir, f"rank_{rank_id}", f"net.safetensors"),
-                       format='safetensors')
+    ms.save_checkpoint(total_param, os.path.join(dst_safetensors_dir, f"rank_{rank_id}", f"net.{output_format}"),
+                       format=output_format)
     return None
 
 
@@ -1143,4 +1009,4 @@ def _get_slice(rank_id, sf_obj, param_name, dst_strategy_list):
 
 
 __all__ = ["_transform_safetensors", "transform_safetensors_by_stage",
-           "transform_safetensors_by_rank", "ckpt_to_safetensors", "safetensors_to_ckpt", "unified_safetensors"]
+           "transform_safetensors_by_rank", "unified_safetensors"]

mindspore/profiler/__init__.py

@@ -20,9 +20,22 @@ Users can visualize the results using the MindInsight tool.
 Now, Profiler supports AICORE operator, AICPU operator, HostCPU operator, memory,
 correspondence, cluster, etc data analysis.
 """
-__all__ = ["Profiler", "EnvProfiler", "ProfilerLevel", "DynamicProfilerMonitor"]
+__all__ = [
+    "tensor_board_trace_handler",
+    "schedule",
+    "Profiler",
+    "EnvProfiler",
+    "ProfilerLevel",
+    "ProfilerActivity",
+    "AicoreMetrics",
+    "DynamicProfilerMonitor",
+    "mstx"
+]
 
-from mindspore.profiler.profiling import Profiler
-from mindspore.profiler.envprofiling import EnvProfiler
-from mindspore.profiler.profiling import ProfilerLevel
+from mindspore.profiler.mstx import Mstx as mstx
+from mindspore.profiler.profiler import Profiler
+from mindspore.profiler.profiler import tensor_board_trace_handler
+from mindspore.profiler.schedule import Schedule as schedule
+from mindspore.profiler.envprofiler import EnvProfiler
+from mindspore.profiler.common.constant import ProfilerLevel, ProfilerActivity, AicoreMetrics
 from mindspore.profiler.dynamic_profiler import DynamicProfilerMonitor