mindspore 2.4.10__cp310-cp310-win_amd64.whl → 2.6.0rc1__cp310-cp310-win_amd64.whl

mindspore/parallel/transform_safetensors.py

@@ -16,19 +16,26 @@
 from __future__ import absolute_import
 
 import os
-import time
+import sys
 import glob
-import re
 import math
 import json
+import re
 from collections import defaultdict
 
+import time
 import multiprocessing as mp
+import psutil
 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.log import vlog_print
 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
+    _insert_opt_shard_reshape, _extract_src_dst_layout_map_by_src, _insert_expand_layout_reshape
 from mindspore.parallel._tensor import _get_tensor_strategy, _construct_from_to_tensor_layout, \
     _get_needed_rank_transform_operator_map_by_layouts, \
     _generate_transform_operator_stack, _apply_tensor_transform_operators, _construct_tensor_layout_for_opt_shard, \
@@ -36,70 +43,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):
     """
@@ -125,6 +68,7 @@ def _progress_bar(iterable, total=None):
         elapsed_time_str = time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
         remaining_time_str = time.strftime("%H:%M:%S", time.gmtime(remaining_time))
 
+        sys.stdout.reconfigure(encoding="utf-8")
         print(f'\r{percent}%|{bar}|[{elapsed_time_str}<{remaining_time_str}]', end='')
         if iteration == total:
             print()
@@ -134,155 +78,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 +265,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):
@@ -485,8 +289,9 @@ def _count_redundancy_list(rank_num, param_name, redundancy_dict, device_num):
 
 
 def _find_remove_redundancy_rank_id(pipe_param_list, single_param_dict, file_dict, saftensor_dict, redundancy_dict,
-                                    needed_rank, device_num):
+                                    needed_rank, device_num, choice_func):
     """Find the rank_id under redundant groups."""
+    io_time = 0
     for param_name in pipe_param_list:
         rank_num = int(needed_rank)
         redundancy_ranks = _count_redundancy_list(rank_num, param_name, redundancy_dict, device_num)
@@ -499,11 +304,23 @@ def _find_remove_redundancy_rank_id(pipe_param_list, single_param_dict, file_dic
                     open_file_id = real_rank
                     break
         if open_file_id is not None:
-            output = file_dict[open_file_id].get_tensor(param_name)
+            start_time = time.time()
+            output = file_dict[open_file_id].get_slice(param_name)
+            end_time = time.time()
+            cost_time = end_time - start_time
+            io_time += cost_time
+            if choice_func is not None:
+                choice_out = choice_func(param_name)
+                if isinstance(choice_out, bool) and not choice_out:
+                    continue
+                if not isinstance(choice_out, (bool, str)):
+                    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[param_name] = output
         else:
             raise ValueError(f"For _transform_safetensors_single, {param_name} should be in "
                              f"{redundancy_ranks}, but in {single_param_dict[param_name]}.")
+    return io_time
 
 
 def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map, src_stage_device_num,
@@ -512,13 +329,14 @@ 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.
     """
+    io_cost_time = 0
     if src_strategy_file is not None:
         from mindspore.train._utils import get_parameter_redundancy
-        redundancy_dict_tmp = get_parameter_redundancy(src_strategy_file)
+        redundancy_dict_tmp = get_parameter_redundancy(src_strategy_file, initial_rank=0)
         redundancy_dict = {}
         device_num = 0
         for param_name, redundancy in redundancy_dict_tmp.items():
@@ -552,8 +370,10 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
             if pipe_param_list:
                 saftensor_dict = dict()
                 if src_strategy_file is not None:
-                    _find_remove_redundancy_rank_id(pipe_param_list, single_param_dict, file_dict, saftensor_dict,
-                                                    redundancy_dict, needed_rank, device_num)
+                    io_time = _find_remove_redundancy_rank_id(pipe_param_list, single_param_dict, file_dict,
+                                                              saftensor_dict, redundancy_dict, needed_rank,
+                                                              device_num, choice_func)
+                    io_cost_time += io_time
                 else:
                     with safe_open(all_safetensor_files_map.get(int(needed_rank)), framework="np") as f:
                         if not unified_flag:
@@ -562,14 +382,32 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
                             dst_param_name_set = set(dst_strategy_list_keys)
                             hyper_param_set = all_param_name_set - (src_param_name_set & dst_param_name_set)
                             pipe_param_list.extend(list(hyper_param_set))
+                        io_time = 0
                         for param_name in pipe_param_list:
                             if param_name not in f.keys():
                                 # param not in ckpt file, check reason
                                 continue
-                            output = f.get_tensor(param_name)
+                            start_time = time.time()
+                            output = f.get_slice(param_name)
+                            end_time = time.time()
+                            cost_time = end_time - start_time
+                            io_time += cost_time
+                            io_cost_time += io_time
+                            if choice_func is not None:
+                                choice_out = choice_func(param_name)
+                                if isinstance(choice_out, bool) and not choice_out:
+                                    continue
+                                if not isinstance(choice_out, (bool, str)):
+                                    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[param_name] = output
             else:
+                start_time = time.time()
                 saftensor_dict = load_file(all_safetensor_files_map.get(int(needed_rank)))
+                end_time = time.time()
+                cost_time = end_time - start_time
+                io_cost_time += cost_time
+
             for param_name, param in saftensor_dict.items():
                 src_rank = int(needed_rank) % src_stage_device_num
                 param_total_dict[param_name][src_rank] = param
@@ -588,7 +426,7 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
             local_rank_id = transform_rank % dst_stage_device_num
             transform_param_dict = _transform_parallel_safetensor(local_rank_id, param_total_dict,
                                                                   param_attr_dict, src_strategy_list, dst_strategy_list,
-                                                                  param_total_dict_keys, src_strategy_file)
+                                                                  param_total_dict_keys, src_strategy_file, choice_func)
             if file_index is not None:
                 save_safetensor_file = f"part{file_index}.{output_format}"
                 save_safetensor_file_dir = dst_safetensors_dir
@@ -602,15 +440,17 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
             if _transform_param_list is not None:
                 _transform_param_list.append({save_file_name: transform_param_dict})
             else:
-                if output_format == "safetensors":
-                    save_file(transform_param_dict, save_file_name)
-                else:
-                    transform_param_dict = _load_and_transform(transform_param_dict, None, None,
-                                                               transform_func=lambda v, name: ms.Parameter(v,
-                                                                                                           name=name))
-                    ms.save_checkpoint(transform_param_dict, save_file_name)
+                if transform_param_dict:
+                    if output_format == "safetensors":
+                        save_file(transform_param_dict, save_file_name)
+                    else:
+                        transform_param_dict = _load_and_transform(transform_param_dict,
+                                                                   None, None, transform_func=
+                                                                   lambda v, name: ms.Parameter(v, name=name))
+                        ms.save_checkpoint(transform_param_dict, save_file_name)
             del param_total_dict_keys
         del param_total_dict
+    return io_cost_time
 
 
 def _save_final_safetensors(_transform_param_list, output_format):
@@ -735,6 +575,13 @@ 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 +605,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
 
 
@@ -775,7 +619,7 @@ def _find_needed_ranks(src_strategy_dict, dst_strategy_dict):
     dst_stage_device_num = _get_device_num_from_strategy(dst_strategy_dict)
     dst_stage_num = _extract_pipeline_stage_num(dst_strategy_dict)
     dst_device_num = dst_stage_device_num * dst_stage_num
-    for rank in _progress_bar(range(dst_device_num)):
+    for rank in range(dst_device_num):
         needed_rank_list = ms.rank_list_for_transform(rank, src_strategy_dict, dst_strategy_dict)
         needed_rank_list_key = "-".join([str(r) for r in needed_rank_list])
         needed_rank_list_map[needed_rank_list_key].append(rank)
@@ -791,7 +635,8 @@ def load_file_by_param_name(filename, parme_name_list):
 
 
 def _transform_parallel_safetensor(rank_id, param_total_dict, param_attr_dict, src_strategy_list,
-                                   dst_strategy_list, param_total_dict_keys=None, src_strategy_file=None):
+                                   dst_strategy_list, param_total_dict_keys=None, src_strategy_file=None,
+                                   choice_func=None):
     """
     Transform model parallel dimension for distributed safetensor files.
     """
@@ -799,7 +644,10 @@ def _transform_parallel_safetensor(rank_id, param_total_dict, param_attr_dict, s
     device_num = -1
     param_total_dict_keys = list(param_total_dict.keys()) if param_total_dict_keys is None else param_total_dict_keys
     for param_name in param_total_dict_keys:
-        tensor_shape = list(param_total_dict[param_name].values())[0].shape
+        if str(type(list(param_total_dict[param_name].values())[0])) == "<class 'builtins.PySafeSlice'>":
+            tensor_shape = list(param_total_dict[param_name].values())[0].get_shape()
+        else:
+            tensor_shape = list(param_total_dict[param_name].values())[0].shape
         from_dev_matrix = [1]
         from_tensor_map = [-1] * len(tensor_shape)
         from_opt_shard_step = 0
@@ -832,6 +680,9 @@ def _transform_parallel_safetensor(rank_id, param_total_dict, param_attr_dict, s
                 continue
             origin_tensor_shape += (item * param_strategy[i],)
 
+        has_layout_from = any(isinstance(i, (list, tuple)) for i in from_tensor_map)
+        has_layout_to = any(isinstance(i, (list, tuple)) for i in to_tensor_map_origin)
+
         from_dev_matrix, from_tensor_map, from_full_tensor_shape = _construct_tensor_layout_for_opt_shard(
             from_dev_matrix, from_tensor_map, from_opt_shard_step, from_opt_shard_size, origin_tensor_shape)
         to_dev_matrix, to_tensor_map, to_full_tensor_shape = _construct_tensor_layout_for_opt_shard(
@@ -851,21 +702,132 @@ def _transform_parallel_safetensor(rank_id, param_total_dict, param_attr_dict, s
         from_info_tuple = (from_opt_shard_size, from_dev_matrix, from_tensor_map, from_full_tensor_shape)
         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)
+        _insert_expand_layout_reshape(param_rank_map, from_info_tuple, to_info_tuple, has_layout_from, has_layout_to)
         transform_operator_stack = _generate_transform_operator_stack(param_rank_map, rank_id)
         param_total_dict_copy = param_total_dict[param_name].copy()
         _apply_tensor_transform_operators(transform_operator_stack, param_total_dict_copy, device_num)
-
+        if choice_func is not None:
+            choice_out = choice_func(param_name)
+            if isinstance(choice_out, str):
+                param_name = choice_out
         transform_param_dict[param_name] = param_total_dict_copy[rank_id % device_num]
+        if str(type(transform_param_dict[param_name])) == "<class 'builtins.PySafeSlice'>":
+            transform_param_dict[param_name] = transform_param_dict[param_name][:]
 
     # Handle those parameter like learning_rate, global_step which not in strategy_file.
     for param_name in param_total_dict_keys:
+        if choice_func is not None:
+            choice_out = choice_func(param_name)
+            if isinstance(choice_out, str):
+                continue
         if param_name not in transform_param_dict:
             transform_para = param_total_dict[param_name][rank_id % device_num]
+            if str(type(transform_para)) == "<class 'builtins.PySafeSlice'>":
+                transform_para = transform_para[:]
             transform_param_dict[param_name] = transform_para
     return transform_param_dict
 
 
-def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundancy=True, file_suffix=None):
+def _cal_param_size(shape, dtype):
+    """cal param size by dtype and shape"""
+    dtype_size = {
+        "BOOL": 1,
+        "U8": 1,
+        "I8": 1,
+        "F8_E5M2": 1,
+        "F8_E4M3": 1,
+        "I16": 2,
+        "U16": 2,
+        "I32": 4,
+        "U32": 4,
+        "I64": 8,
+        "U64": 8,
+        "F16": 2,
+        "BF16": 2,
+        "F32": 4,
+        "F64": 8,
+    }
+    num_elements = math.prod(shape)
+    element_size = dtype_size.get(dtype, 4)
+    total_bytes = num_elements * element_size
+    return total_bytes
+
+
+def _split_weight_dict(weights, num_groups):
+    """split weights by num"""
+    sorted_items = sorted(weights.items(), key=lambda x: -x[1])
+    groups = [[] for _ in range(num_groups)]
+    total_bytes = [0] * num_groups
+    for weight_name, byte_size in sorted_items:
+        min_index = total_bytes.index(min(total_bytes))
+        groups[min_index].append(weight_name)
+        total_bytes[min_index] += byte_size
+
+    return groups
+
+
+def _save_hyper_param(split_dst_file, all_safetensor_files_map, name_list, dst_dir):
+    """save hyper param"""
+    if not split_dst_file or (split_dst_file and split_dst_file[0] == 1):
+        with safe_open(all_safetensor_files_map.get(0), framework="np") as f:
+            all_key = f.keys()
+            hyper_parameter = set(all_key) - set(name_list)
+            if hyper_parameter:
+                hyper_dict = {}
+                for key in hyper_parameter:
+                    hyper_dict[key] = f.get_tensor(key)
+                save_file(hyper_dict, os.path.join(dst_dir, "hyper_param.safetensors"))
+
+
+def _save_parameter_map_json(split_list, choice_func, split_dst_file, dst_dir, param_total_size):
+    """save parameter map json file"""
+    param_name_dict = dict()
+    for index, part_list in enumerate(split_list):
+        for name in part_list:
+            save_param_name = name
+            if choice_func is not None:
+                choice_out = choice_func(name)
+                if isinstance(choice_out, str):
+                    save_param_name = choice_out
+            if save_param_name == -1:
+                break
+            param_name_dict[save_param_name] = f"part{index}.safetensors"
+    output_dict = {"metadata": {"total_size": param_total_size}, "weight_map": param_name_dict}
+    if not split_dst_file or (split_dst_file and split_dst_file[0] == 1):
+        json_str = json.dumps(output_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)
+
+
+def _get_dst_shape(param_name, param_shape, src_strategy_list):
+    """get dst shape by strategy"""
+    from_dev_matrix = [1]
+    from_tensor_map = [-1] * len(param_shape)
+    from_opt_shard_size = 0
+    if src_strategy_list is not None:
+        from_dev_matrix, from_tensor_map, _, from_opt_shard_size = _extract_layout_item(
+            src_strategy_list.get(param_name))
+    to_dev_matrix_origin = [1]
+    to_tensor_map_origin = [-1] * len(param_shape)
+    to_opt_shard_step = 0
+    to_opt_shard_size = 0
+
+    param_strategy = _get_tensor_strategy(from_dev_matrix, from_tensor_map)
+    origin_tensor_shape = ()
+    for i, item in enumerate(param_shape):
+        if i == 0 and from_opt_shard_size > 0:
+            origin_tensor_shape += (item * param_strategy[i] * from_opt_shard_size,)
+            continue
+        origin_tensor_shape += (item * param_strategy[i],)
+
+    _, _, to_full_tensor_shape = _construct_tensor_layout_for_opt_shard(
+        to_dev_matrix_origin, to_tensor_map_origin, to_opt_shard_step, to_opt_shard_size, origin_tensor_shape)
+    return to_full_tensor_shape
+
+
+def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundancy=True, file_suffix=None,
+                        max_process_num=64, choice_func=None, split_dst_file=()):
     """
     Merge multiple safetensor files into a unified safetensor file.
 
@@ -877,6 +839,14 @@ 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, optional): Maximum number of processes. Default: ``64``.
+        choice_func (callable, optional): 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``.
+        split_dst_file (tuple, optional) - A parameter used to manually split a task into multiple subtasks for
+            execution, represented as a tuple containing two elements. The first element indicates the number of
+            the current subtask, and the second element indicates the total number of tasks. This parameter supports
+            splitting and executing tasks multiple times on a single machine, and also supports executing different
+            subtasks on multiple machines respectively. Default: ``()``.
 
     Raises:
         ValueError: If the safetensors file of rank is missing.
@@ -889,8 +859,12 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
         >>> src_dir = "/usr/safetensors/llama31B/4p_safetensors/"
         >>> src_strategy_file = "/usr/safetensors/llama31B/strategy_4p.ckpt"
         >>> dst_dir = "/usr/safetensors/llama31B/merge_llama31B_4p/"
-        >>> ms.unified_safetensors(src_dir, src_strategy_file, dst_dir)
+        >>> ms.parallel.unified_safetensors(src_dir, src_strategy_file, dst_dir)
     """
+    pid = os.getpid()
+    total_cores = os.cpu_count()
+    all_cores = set(range(total_cores))
+    os.sched_setaffinity(pid, all_cores)
     _check_transform_safetensors(src_dir, "", src_strategy_file, None)
     _make_dir(dst_dir, "path")
     if os.path.isfile(src_dir):
@@ -914,13 +888,11 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
                                  "but it is missing.".format(needed_rank, rank))
     layout_map = _convert_to_list(src_strategy_dict)
 
-    total_size = 0
     actual_params = set()
     for _, file_name in all_safetensor_files_map.items():
-        total_size += os.path.getsize(file_name) / 1024 / 1024 / 1024
         with safe_open(file_name, framework="np") as f:
             actual_params.update(f.keys())
-    split_num = math.ceil(total_size / 3)
+
     params_to_store = actual_params & set(layout_map.keys())
 
     name_list = []
@@ -928,29 +900,55 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
         if name.startswith("accu_grads"):
             continue
         name_list.append(name)
-    split_list = _split_list(name_list, split_num)
-
-    with safe_open(all_safetensor_files_map.get(0), framework="np") as f:
-        all_key = f.keys()
-        hyper_parameter = set(all_key) - set(name_list)
-        if hyper_parameter:
-            hyper_dict = {}
-            for key in hyper_parameter:
-                hyper_dict[key] = f.get_tensor(key)
-            save_file(hyper_dict, os.path.join(dst_dir, "hyper_param.safetensors"))
-
-    # save parameter map json
-    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"
-    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)
-    res = [i for i in range(split_num)]
+
+    param_size_dict = {}
+    param_total_size = 0
+    for _, file_name in all_safetensor_files_map.items():
+        with safe_open(file_name, framework="np") as f:
+            for k in f.keys():
+                if k in name_list:
+                    py_slice = f.get_slice(k)
+                    param_total_size += _cal_param_size(py_slice.get_shape(), py_slice.get_dtype())
+                    param_dst_shape = _get_dst_shape(k, py_slice.get_shape(), origin_src_strategy_list)
+                    # Convert the shape of np.int32 type to int type to prevent overflow in subsequent calculations.
+                    param_dst_shape = [int(item) for item in param_dst_shape]
+                    if choice_func is not None:
+                        choice_out = choice_func(k)
+                        if isinstance(choice_out, bool):
+                            if not choice_out:
+                                continue
+                    if k not in param_size_dict:
+                        param_size_dict[k] = _cal_param_size(param_dst_shape, py_slice.get_dtype())
+    split_num = math.ceil(sum(param_size_dict.values()) / 1024 / 1024 / 1024 / 3)
+    split_num = min(split_num, len(name_list))
+    split_list = _split_weight_dict(param_size_dict, split_num)
+
+    if split_dst_file:
+        current_machine_num = split_dst_file[0]
+        total_machine_num = split_dst_file[1]
+        n = len(split_list)
+        avg_length = n // total_machine_num
+        remainder = n % total_machine_num
+        start_index = (avg_length * (current_machine_num - 1)) + min(current_machine_num - 1, remainder)
+        end_index = start_index + avg_length + (1 if current_machine_num <= remainder else 0)
+        sub_list = []
+        for i in range(len(split_list)):
+            if start_index <= i < end_index:
+                sub_list.append(split_list[i])
+            else:
+                sub_list.append([-1])
+    else:
+        sub_list = split_list
+
+    _save_hyper_param(split_dst_file, all_safetensor_files_map, name_list, dst_dir)
+    _save_parameter_map_json(split_list, choice_func, split_dst_file, dst_dir, param_total_size)
+
+    if split_dst_file:
+        split_num = end_index - start_index
+        res = list(range(start_index, end_index))
+    else:
+        res = [i for i in range(split_num)]
+    max_process = min(split_num, max_process_num)
     res = _split_list(res, max_process)
     processes = []
     src_strategy_name = None
@@ -960,7 +958,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, sub_list, res[i], True, src_strategy_name, choice_func))
         p.start()
         processes.append(p)
     for p in processes:
@@ -974,13 +972,21 @@ 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):
+    """transform safetensors single semaphore"""
+    total_io_cost_time = 0
     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)
+        io_cost_time = _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, choice_func)
+        while psutil.virtual_memory().percent > 50:
+            time.sleep(1)
+        total_io_cost_time += io_cost_time
+    vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+               f"Unified safetensors io cost time:{total_io_cost_time}.")
 
 
 def _split_list(split_list, split_num):
@@ -1027,36 +1033,76 @@ 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 _process_hyper_params(file_list, total_safetensors_dir, 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():
+                total_param[key] = 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 not json_files:
+            raise ValueError(
+                f"For 'load_parallel_checkpoint', there must be a JSON file named 'param_name_map.json' in "
+                f"the 'total_safetensors_dir'.")
+        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 "weight_map" in param_name_map:
+                param_name_map = param_name_map["weight_map"]
+
     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
+    pid = os.getpid()
+    total_cores = os.cpu_count()
+    all_cores = set(range(total_cores))
+    os.sched_setaffinity(pid, all_cores)
+    file_list = os.listdir(total_safetensors_dir)
+    json_files = [file for file in file_list if file == "param_name_map.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
     local_rank_id = rank_id % dst_stage_device_num
-    for param_name in param_list:
+    total_io_cost_time = 0
+    for param_name in _progress_bar(param_list):
         if param_name not in param_name_map:
             continue
         file_name = os.path.join(total_safetensors_dir, param_name_map[param_name])
         with safe_open(file_name, framework="np") as f:
-            if param_name not in f.keys():
+            cur_param_name = name_map.get(param_name) if name_map is not None and param_name in name_map else param_name
+            if cur_param_name not in f.keys():
                 continue
-            sf_obj = f.get_slice(param_name)
+            sf_obj = f.get_slice(cur_param_name)
 
         tensor_shape = sf_obj.get_shape()
         from_dev_matrix = [1]
@@ -1078,6 +1124,9 @@ def _load_parallel_checkpoint(total_safetensors_dir, dst_strategy_file, net=None
                     continue
                 origin_tensor_shape += (item * param_strategy[i],)
 
+            has_layout_from = any(isinstance(i, (list, tuple)) for i in from_tensor_map)
+            has_layout_to = any(isinstance(i, (list, tuple)) for i in to_tensor_map_origin)
+
             from_dev_matrix, from_tensor_map, from_full_tensor_shape = _construct_tensor_layout_for_opt_shard(
                 from_dev_matrix, from_tensor_map, from_opt_shard_step, from_opt_shard_size, origin_tensor_shape)
             to_dev_matrix, to_tensor_map, to_full_tensor_shape = _construct_tensor_layout_for_opt_shard(
@@ -1097,25 +1146,34 @@ def _load_parallel_checkpoint(total_safetensors_dir, dst_strategy_file, net=None
             from_info_tuple = (from_opt_shard_size, from_dev_matrix, from_tensor_map, from_full_tensor_shape)
             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)
+            _insert_expand_layout_reshape(param_rank_map, from_info_tuple, to_info_tuple,
+                                          has_layout_from, has_layout_to)
             transform_operator_stack = _generate_transform_operator_stack(param_rank_map, local_rank_id)
-
+            start_time = time.time()
             slice_param = _apply_sf_obj_transform_operators(transform_operator_stack, sf_obj, device_num)
+            end_time = time.time()
+            cost_time = end_time - start_time
+            total_io_cost_time += cost_time
         else:
+            start_time = time.time()
             slice_param = sf_obj[:]
-
-        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))
+            end_time = time.time()
+            cost_time = end_time - start_time
+            total_io_cost_time += cost_time
+        total_param[param_name] = ms.Parameter(ms.Tensor.from_numpy(slice_param))
+    vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+               f"load distributed safetensors io cost time:{total_io_cost_time}.")
+    total_param = _process_hyper_params(file_list, total_safetensors_dir, 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:
+            logger.info("start load param into net...")
+            param_not_load, ckpt_not_load = ms.load_param_into_net(net, total_param)
+            logger.info("load param into net is end...")
+            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 +1201,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"]