mindspore 2.5.0__cp311-cp311-win_amd64.whl → 2.6.0rc1__cp311-cp311-win_amd64.whl

mindspore/parallel/transform_safetensors.py

@@ -16,6 +16,7 @@
 from __future__ import absolute_import
 
 import os
+import sys
 import glob
 import math
 import json
@@ -24,15 +25,17 @@ 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, \
@@ -65,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()
@@ -285,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)
@@ -299,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,
@@ -316,9 +333,10 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
     """
     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():
@@ -352,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:
@@ -362,25 +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)
-                            save_param_name = 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):
-                                    if not choice_out:
-                                        continue
-                                elif isinstance(choice_out, str):
-                                    save_param_name = choice_out
-                                else:
+                                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[save_param_name] = output
+                            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
@@ -399,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
@@ -413,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):
@@ -552,6 +581,7 @@ def _extrace_number(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.
@@ -589,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)
@@ -605,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.
     """
@@ -613,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
@@ -646,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(
@@ -665,22 +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 _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):
+                        max_process_num=64, choice_func=None, split_dst_file=()):
     """
     Merge multiple safetensor files into a unified safetensor file.
 
@@ -692,9 +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): 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.
+        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.
@@ -707,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):
@@ -732,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 = []
@@ -746,37 +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:
-            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)
 
+    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 = [i for i in range(split_num)]
     res = _split_list(res, max_process)
     processes = []
     src_strategy_name = None
@@ -786,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, choice_func))
+            "", dst_dir, "safetensors", None, sub_list, res[i], True, src_strategy_name, choice_func))
         p.start()
         processes.append(p)
     for p in processes:
@@ -801,13 +973,20 @@ def _transform_safetensors_single_semaphore(needed_rank_list_map, all_safetensor
                                             ckpt_prefix, dst_safetensors_dir, output_format,
                                             _transform_param_list, pipe_param_list=None, file_index=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,
-                                      choice_func)
+        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):
@@ -854,22 +1033,13 @@ def _apply_sf_obj_transform_operators(transform_operator_stack, sf_obj, device_n
     return sf_obj
 
 
-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):
+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():
-                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)))
+                total_param[key] = ms.Parameter(ms.Tensor.from_numpy(f.get_tensor(key)))
     return total_param
 
 
@@ -887,12 +1057,15 @@ def _cal_param_name_map_and_param_list(file_list, total_safetensors_dir, json_fi
             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)}.")
+        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)
@@ -907,8 +1080,12 @@ 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.endswith('.json')]
+    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)
@@ -916,14 +1093,16 @@ def _load_parallel_checkpoint(file_info):
     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]
@@ -945,6 +1124,9 @@ def _load_parallel_checkpoint(file_info):
                     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(
@@ -964,19 +1146,29 @@ def _load_parallel_checkpoint(file_info):
             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[:]
-        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 = _process_hyper_params(file_list, total_safetensors_dir, name_map, total_param)
+            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:
         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")