mindspore 2.4.0__cp311-none-any.whl → 2.4.10__cp311-none-any.whl

mindspore/parallel/_utils.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """Utils of auto parallel"""
 import os
+from time import perf_counter
 from importlib import import_module
 import numpy as np
 import mindspore as ms
@@ -27,7 +28,7 @@ from mindspore.communication._comm_helper import _is_initialized
 from mindspore.parallel._auto_parallel_context import auto_parallel_context
 from mindspore.common.seed import get_seed
 from mindspore._c_expression import GraphExecutor_
-from mindspore.parallel._tensor import _load_tensor_by_layout
+from mindspore.parallel._tensor import _load_tensor_by_layout, _load_tensor_shape_by_layout
 
 SUPPORTED_TUPLE_IN_TUPLE_STRATEGY = ["GroupedMatmul", "FusedInferAttentionScore", "Custom"]
 
@@ -104,31 +105,48 @@ def _need_to_full():
     return not _get_full_batch()
 
 
+class ParallelParamInitProfCtx:
+    """Collect parallel param initialization performance context mgr."""
+
+    def __init__(self, parameter, func_name):
+        self.parameter = parameter
+        self.func_name = func_name
+        self.start_timestamp = None
+
+    def __enter__(self):
+        self.start_timestamp = perf_counter()
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        end_timestamp = perf_counter()
+        duration = end_timestamp - self.start_timestamp
+        if os.getenv("MS_DEV_PARAM_INIT_PROF_COLLECT"):
+            logger.warning(f"{self.func_name}: {self.parameter.name}, shape: {self.parameter.shape}, "
+                           f"sliced: {self.parameter.sliced}, duration: {duration}")
+
+
 def _slice_parameter(parameter, phase, layout):
     """Slice python parameter obj according to the layout."""
-    is_train_phase = phase.startswith('train')
-    is_prefill_phase = phase.startswith('prefill')
-    if layout is not None and parameter.from_ckpt and not is_train_phase:
-        is_opt_shard_group = layout[5]
-        if not parameter.sliced and is_prefill_phase and is_opt_shard_group:
+    if getattr(parameter, "init_param", False):
+        if layout is None:
+            parameter.sliced = True
+            return
+        if not parameter.sliced:
+            rank = get_rank()
+            new_tensor_shape = _load_tensor_shape_by_layout(parameter, layout, rank)
+            parameter.shape = new_tensor_shape
+    else:
+        graph_executor = GraphExecutor_.get_instance()
+        new_param = parameter.init_data(layout, set_sliced=True)
+        parameter = new_param
+        graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
+        if layout is None:
+            parameter.sliced = True
+            return
+        if not parameter.sliced:
             rank = get_rank()
             new_tensor = _load_tensor_by_layout(parameter, layout, rank)
             parameter.set_data(new_tensor, True)
-            return
-        layout_shape = layout[2]
-        parameter.shape = tuple(layout_shape)
-        return
-    graph_executor = GraphExecutor_.get_instance()
-    new_param = parameter.init_data(layout, set_sliced=True)
-    parameter = new_param
-    graph_executor.updata_param_node_default_input(phase, {parameter.name: parameter})
-    if layout is None:
-        parameter.sliced = True
-        return
-    if not parameter.sliced:
-        rank = get_rank()
-        new_tensor = _load_tensor_by_layout(parameter, layout, rank)
-        parameter.set_data(new_tensor, True)
 
 
 def _slice_tensor(tensor, layout, rank_id):

mindspore/parallel/transform_safetensors.py

@@ -32,7 +32,7 @@ from mindspore.parallel._parallel_serialization import _get_device_num_from_stra
 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, \
-    _extract_layout_item, _load_tensor_shape
+    _extract_layout_item, _load_tensor_shape, _apply_operator
 from mindspore.parallel._parallel_serialization import _build_searched_strategy, _load_protobuf_strategy, \
     _convert_to_list
 
@@ -375,12 +375,10 @@ def _transform_stage_safetensors(src_strategy_dict, dst_strategy_dict, ckpt_pref
             if int(needed_rank) not in all_safetensor_files_map:
                 raise ValueError("The safetensor file of rank{} is needed for converting rank{}'s safetensor, "
                                  "but it is missing.".format(needed_rank, rank))
-    if process_num > len(needed_rank_list_map):
+    dst_stage_num = _extract_pipeline_stage_num(dst_strategy_dict)
+    if not (len(needed_rank_list_map) == 1 and dst_stage_num > 1) and process_num > len(needed_rank_list_map):
         ms.log.warning("The value of process_num cannot be greater than that of needed_rank_list_map.")
         process_num = len(needed_rank_list_map)
-    dst_stage_num = _extract_pipeline_stage_num(dst_strategy_dict)
-    if len(needed_rank_list_map) == 1 and dst_stage_num > 1:
-        process_num = dst_stage_num
     _transform_safetensors_with_parallel(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,
@@ -452,18 +450,18 @@ def _transform_safetensors_with_parallel(needed_rank_list_map, all_safetensor_fi
     """
     Transforms safetensors files to a specified format using parallel processing.
     """
-    part_list_dict = _distribute_files_by_size(all_safetensor_files_map, needed_rank_list_map, process_num)
-
     # cal param name for every pipeline, save in pipe_param_list.
     pipe_num = _extract_pipeline_stage_num(dst_strategy_dict)
     pipe_param_list = [None for _ in range(max(pipe_num, process_num))]
     if len(needed_rank_list_map) == 1 and pipe_num > 1:
+        process_num = pipe_num
         pipe_param_list = [[] for _ in range(pipe_num)]
         layout_map = _convert_to_list(dst_strategy_dict)
 
         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):
         p = mp.Process(target=_transform_safetensors_single, args=(
@@ -476,15 +474,74 @@ def _transform_safetensors_with_parallel(needed_rank_list_map, all_safetensor_fi
         p.join()
 
 
+def _count_redundancy_list(rank_num, param_name, redundancy_dict, device_num):
+    """Obtain the specified redundant group."""
+    redundancy_tuple = redundancy_dict.get(param_name)
+    for rank_list in redundancy_tuple:
+        for rank in rank_list:
+            if rank_num % device_num == rank % device_num:
+                return set(rank_list)
+    return set()
+
+
+def _find_remove_redundancy_rank_id(pipe_param_list, single_param_dict, file_dict, saftensor_dict, redundancy_dict,
+                                    needed_rank, device_num):
+    """Find the rank_id under redundant groups."""
+    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)
+        open_file_id = None
+        if single_param_dict.get(param_name) is None:
+            continue
+        for real_rank in single_param_dict[param_name]:
+            for redundancy_rank in redundancy_ranks:
+                if real_rank % device_num == redundancy_rank % device_num:
+                    open_file_id = real_rank
+                    break
+        if open_file_id is not None:
+            output = file_dict[open_file_id].get_tensor(param_name)
+            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]}.")
+
+
 def _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=None, file_index=None, unified_flag=False):
+                                  _transform_param_list, pipe_param_list=None, file_index=None, unified_flag=False,
+                                  src_strategy_file=None):
     """
     Transforms safetensors files to a specified format without using parallel processing.
     """
+    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 = {}
+        device_num = 0
+        for param_name, redundancy in redundancy_dict_tmp.items():
+            if device_num == 0:
+                device_num = max(max(redundancy)) + 1
+            origin_param_name = param_name
+            pipeline_stage = 0
+            if "-" in param_name:
+                pipeline_stage, origin_param_name = param_name.split("-")
+                pipeline_stage = int(pipeline_stage)
+            redundancy_new = tuple(
+                (tuple(x + pipeline_stage * device_num for x in subtuple)) for subtuple in redundancy)
+            redundancy_dict[origin_param_name] = redundancy_new
+        file_dict = {}
+        single_param_dict = {}
+        for file_id, _ in all_safetensor_files_map.items():
+            f = safe_open(all_safetensor_files_map.get(file_id), framework="np")
+            file_dict[file_id] = f
+            for param_name in f.keys():
+                if param_name not in single_param_dict.keys():
+                    single_param_dict[param_name] = {file_id}
+                else:
+                    single_param_dict[param_name].add(file_id)
     src_strategy_list_keys = _convert_to_list(src_strategy_dict).keys() if src_strategy_dict else []
     dst_strategy_list_keys = _convert_to_list(dst_strategy_dict).keys() if dst_strategy_dict else []
     for needed_rank_list_key, transform_rank_list in needed_rank_list_map.items():
@@ -494,19 +551,23 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
         for needed_rank in needed_rank_list:
             if pipe_param_list:
                 saftensor_dict = dict()
-                with safe_open(all_safetensor_files_map.get(int(needed_rank)), framework="np") as f:
-                    if not unified_flag:
-                        all_param_name_set = set(f.keys())
-                        src_param_name_set = set(src_strategy_list_keys)
-                        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))
-                    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)
-                        saftensor_dict[param_name] = output
+                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)
+                else:
+                    with safe_open(all_safetensor_files_map.get(int(needed_rank)), framework="np") as f:
+                        if not unified_flag:
+                            all_param_name_set = set(f.keys())
+                            src_param_name_set = set(src_strategy_list_keys)
+                            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))
+                        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)
+                            saftensor_dict[param_name] = output
             else:
                 saftensor_dict = load_file(all_safetensor_files_map.get(int(needed_rank)))
             for param_name, param in saftensor_dict.items():
@@ -527,7 +588,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)
+                                                                  param_total_dict_keys, src_strategy_file)
             if file_index is not None:
                 save_safetensor_file = f"part{file_index}.{output_format}"
                 save_safetensor_file_dir = dst_safetensors_dir
@@ -674,7 +735,7 @@ def transform_safetensors_by_rank(rank_id, safetensor_files_map, save_safetensor
     save_file(transform_param_dict, save_safetensor_file_name)
 
 
-def _collect_safetensor_files(src_safetensors_dir, format='safetensors'):
+def _collect_safetensor_files(src_safetensors_dir, format='safetensors', file_suffix=None):
     """
     Collects all safetensors files from the specified directory and its subdirectories.
     """
@@ -692,7 +753,10 @@ def _collect_safetensor_files(src_safetensors_dir, format='safetensors'):
                            format(safetensor_dir))
             continue
         rank_id = int(rank_id_str)
-        safetensor_file_name = os.path.join(safetensor_dir, f"*.{format}")
+        if file_suffix is None:
+            safetensor_file_name = os.path.join(safetensor_dir, f"*.{format}")
+        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:
@@ -727,7 +791,7 @@ 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):
+                                   dst_strategy_list, param_total_dict_keys=None, src_strategy_file=None):
     """
     Transform model parallel dimension for distributed safetensor files.
     """
@@ -779,7 +843,7 @@ def _transform_parallel_safetensor(rank_id, param_total_dict, param_attr_dict, s
 
         # when the from_layout is less devices, the safetensor_map for map[device_num] should using map[0]
         device_list = list(range(0, np.prod(from_tensor_layout[0])))
-        if rank_id % device_num not in param_attr_dict[param_name]:
+        if rank_id % device_num not in param_attr_dict[param_name] and src_strategy_file is None:
             raise ValueError("The safetensor of rank {} is missing.".format(rank_id % device_num))
         param_rank_map = _get_needed_rank_transform_operator_map_by_layouts(from_tensor_layout, to_tensor_layout,
                                                                             device_list, rank_id)
@@ -801,7 +865,7 @@ 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):
+def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundancy=True, file_suffix=None):
     """
     Merge multiple safetensor files into a unified safetensor file.
 
@@ -809,6 +873,10 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir):
         src_dir (str): Source weight saving directory.
         src_strategy_file (str): Source weight segmentation strategy file.
         dst_dir (str): Target save directory.
+        merge_with_redundancy (bool, optional): Whether the merged source weight files are de-duplicated and
+            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.
 
     Raises:
         ValueError: If the safetensors file of rank is missing.
@@ -827,8 +895,8 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir):
     _make_dir(dst_dir, "path")
     if os.path.isfile(src_dir):
         raise ValueError("For 'unified_safetensors', the 'src_dir' can not be a file.")
-    all_safetensor_files_map = _collect_safetensor_files(src_dir)
-    all_ckpt_files_map = _collect_safetensor_files(src_dir, format='ckpt')
+    all_safetensor_files_map = _collect_safetensor_files(src_dir, format="safetensors", file_suffix=file_suffix)
+    all_ckpt_files_map = _collect_safetensor_files(src_dir, format="ckpt", file_suffix=file_suffix)
     if all_safetensor_files_map and all_ckpt_files_map:
         raise ValueError("For 'unified_safetensors', the 'src_dir' cannot contain "
                          "both ckpt file and safetensors file simultaneously")
@@ -847,14 +915,21 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir):
     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 = list(layout_map.keys())
+    name_list = []
+    for name in list(params_to_store):
+        if name.startswith("accu_grads"):
+            continue
+        name_list.append(name)
     split_list = _split_list(name_list, split_num)
 
-    all_safetensor_files_map = _collect_safetensor_files(src_dir)
     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)
@@ -878,12 +953,14 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir):
     res = [i for i in range(split_num)]
     res = _split_list(res, max_process)
     processes = []
-
+    src_strategy_name = None
+    if not merge_with_redundancy:
+        src_strategy_name = src_strategy_file
     for i in range(max_process):
         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))
+            "", dst_dir, "safetensors", None, split_list, res[i], True, src_strategy_name))
         p.start()
         processes.append(p)
     for p in processes:
@@ -897,13 +974,13 @@ 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):
+                                            unified_flag=False, src_strategy_file=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)
+                                      _transform_param_list, pipe_param_list[i], i, unified_flag, src_strategy_file)
 
 
 def _split_list(split_list, split_num):
@@ -911,6 +988,45 @@ def _split_list(split_list, split_num):
     return [array.tolist() for array in split_array]
 
 
+def _apply_sf_obj_transform_operators(transform_operator_stack, sf_obj, device_num):
+    """apply safetensors object operators"""
+    if not transform_operator_stack:
+        return sf_obj[:]
+    level = transform_operator_stack[-1][1]
+    level_operators = []
+    while True:
+        if not transform_operator_stack or (level != transform_operator_stack[-1][1]):
+            tmp_tensor_dict = {}
+            if not level_operators:
+                continue
+            op_name = level_operators[0][2][0]
+            for operator_pair in level_operators:
+                rank_id = operator_pair[0]
+                cur_level = operator_pair[1]
+                operator = operator_pair[2]
+                if operator[0] != op_name:
+                    raise ValueError("The operator in the same level should be equal in the transform tensor operator "
+                                     "list, but the find {} and {} in level {}".format(op_name, operator[0], cur_level))
+                if operator[0] != "AllConcat":
+                    sf_obj = _apply_operator(operator[0])(sf_obj, operator)
+                    continue
+                for rank in operator[1][:-1]:
+                    if rank % device_num not in sf_obj:
+                        raise ValueError("The checkpoint file of rank {} is missing.".format(rank % device_num))
+                allgather_list = [sf_obj for _ in operator[1][:-1]]
+                tmp_tensor_dict[rank_id % device_num] = _apply_operator(operator[0])(allgather_list, operator)
+            if op_name == "AllConcat":
+                for rank, value in tmp_tensor_dict.items():
+                    sf_obj = value
+            level_operators.clear()
+        if not transform_operator_stack:
+            break
+        operator_pair = transform_operator_stack.pop()
+        level = operator_pair[1]
+        level_operators.append(operator_pair)
+    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."""
@@ -930,7 +1046,9 @@ def _load_parallel_checkpoint(total_safetensors_dir, dst_strategy_file, net=None
         param_list = param_name_map.keys()
 
     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:
         if param_name not in param_name_map:
             continue
@@ -939,19 +1057,54 @@ def _load_parallel_checkpoint(total_safetensors_dir, dst_strategy_file, net=None
             if param_name not in f.keys():
                 continue
             sf_obj = f.get_slice(param_name)
-        param_dict = dict()
-        param_dict[param_name] = sf_obj
 
+        tensor_shape = sf_obj.get_shape()
+        from_dev_matrix = [1]
+        from_tensor_map = [-1] * len(tensor_shape)
+        from_opt_shard_step = 0
+        from_opt_shard_size = 0
         if dst_strategy_list is not None:
             if param_name not in dst_strategy_list:
                 continue
-            slice_op, shape = _get_slice(rank_id, sf_obj, param_name, dst_strategy_list)
+            to_dev_matrix_origin, to_tensor_map_origin, to_opt_shard_step, to_opt_shard_size = _extract_layout_item(
+                dst_strategy_list.get(param_name))
+
+            device_num = np.prod(from_dev_matrix)
+            param_strategy = _get_tensor_strategy(from_dev_matrix, from_tensor_map)
+            origin_tensor_shape = ()
+            for i, item in enumerate(tensor_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],)
+
+            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(
+                to_dev_matrix_origin, to_tensor_map_origin, to_opt_shard_step, to_opt_shard_size, origin_tensor_shape)
+            # Convert tensor layout to same device num
+            from_tensor_layout, to_tensor_layout = _construct_from_to_tensor_layout(from_full_tensor_shape,
+                                                                                    from_dev_matrix,
+                                                                                    from_tensor_map,
+                                                                                    to_full_tensor_shape,
+                                                                                    to_dev_matrix, to_tensor_map)
+
+            # when the from_layout is less devices, the safetensor_map for map[device_num] should using map[0]
+            device_list = list(range(0, np.prod(from_tensor_layout[0])))
+            param_rank_map = _get_needed_rank_transform_operator_map_by_layouts(from_tensor_layout, to_tensor_layout,
+                                                                                device_list, local_rank_id)
+
+            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)
+            transform_operator_stack = _generate_transform_operator_stack(param_rank_map, local_rank_id)
+
+            slice_param = _apply_sf_obj_transform_operators(transform_operator_stack, sf_obj, device_num)
         else:
-            slice_op, shape = slice(None, None, None), None
-        slice_param = sf_obj[slice_op]
-        if shape is not None:
-            slice_param = slice_param.reshape(shape)
+            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:

mindspore/profiler/profiling.py

@@ -31,7 +31,7 @@ from mindspore.context import get_auto_parallel_context
 from mindspore.communication.management import GlobalComm, get_rank, get_group_size, get_local_rank
 import mindspore._c_expression as c_expression
 import mindspore._c_dataengine as cde
-from mindspore._c_expression import _framework_profiler_enable_mi
+from mindspore._c_expression import _framework_profiler_enable_mi, _framework_profiler_disable_mi
 from mindspore.profiler.common.exceptions.exceptions import ProfilerFileNotFoundException, \
     ProfilerIOException, ProfilerException, ProfilerRawFileException, ProfilerParamTypeErrorException
 from mindspore.profiler.common.exceptions.exceptions import ProfilerPathErrorException
@@ -824,6 +824,10 @@ class Profiler:
             self._ascend_profiler.stop()
 
             self._stop_time = int(time.time() * 10000000)
+
+        if self._profile_framework:
+            _framework_profiler_disable_mi()
+
         ProfilerInfo.set_profiling_stop_time(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
         self._init_profiler_info()
         ProfilerInfo.set_diff_time(self._start_time - self._monotonic_time)

mindspore/run_check/_check_version.py

@@ -258,7 +258,7 @@ class AscendEnvChecker(EnvChecker):
 
     def __init__(self, library_path):
         self.library_path = library_path
-        self.version = ["7.3", "7.5"]
+        self.version = ["7.5", "7.6"]
 
         # env
         self.path = os.getenv("PATH")
@@ -266,8 +266,10 @@ class AscendEnvChecker(EnvChecker):
         self.ld_lib_path = os.getenv("LD_LIBRARY_PATH")
         self.ascend_opp_path = os.getenv("ASCEND_OPP_PATH")
         self.ascend_aicpu_path = os.getenv("ASCEND_AICPU_PATH")
-        self.compiler_version = self.ascend_opp_path.split("opp")[0] + "compiler/version.info"
-
+        if self.ascend_opp_path is not None:
+            self.compiler_version = self.ascend_opp_path.split("opp")[0] + "compiler/version.info"
+        else:
+            self.compiler_version = ""
         # check content
         self.path_check = "/compiler/ccec_compiler/bin"
         self.python_path_check = "opp/built-in/op_impl/ai_core/tbe"
@@ -285,7 +287,7 @@ class AscendEnvChecker(EnvChecker):
     def check_custom_version(self):
         """custom op version check"""
         custom_ascendc_soc_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)),
-                                              "../lib/plugin/ascend/custom_ascendc_ops/op_impl/ai_core/tbe/kernel",
+                                              "../lib/plugin/ascend/custom_ascendc_910b/op_impl/ai_core/tbe/kernel",
                                               MSContext.get_instance().get_ascend_soc_version())
         if not os.path.exists(custom_ascendc_soc_dir):
             logger.debug(f"The path {custom_ascendc_soc_dir} of the custom ascend c operator does not exist.")
@@ -296,7 +298,7 @@ class AscendEnvChecker(EnvChecker):
 
         cur_version = self._read_version(self.compiler_version)
         custom_version_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
-                                           "../lib/plugin/ascend/custom_ascendc_ops/version.info")
+                                           "../lib/plugin/ascend/custom_ascendc_910/version.info")
         with open(custom_version_path, 'r') as f:
             all_info = f.readlines()
             for line in all_info:
@@ -370,12 +372,21 @@ class AscendEnvChecker(EnvChecker):
         curr_path = os.path.realpath(os.path.dirname(__file__))
         cust_aicpu_path = os.path.realpath(os.path.join(curr_path, "../lib/plugin/ascend/custom_aicpu_ops"))
         cust_aicore_path = os.path.realpath(os.path.join(curr_path, "../lib/plugin/ascend/custom_aicore_ops"))
-        cust_ascendc_path = os.path.realpath(os.path.join(curr_path, "../lib/plugin/ascend/custom_ascendc_ops"))
+        cust_ascendc_ascend910_path = os.path.realpath(
+            os.path.join(curr_path, "../lib/plugin/ascend/custom_ascendc_910"))
+        cust_ascendc_ascend910b_path = os.path.realpath(
+            os.path.join(curr_path, "../lib/plugin/ascend/custom_ascendc_910b"))
         if os.getenv('ASCEND_CUSTOM_OPP_PATH'):
             os.environ['ASCEND_CUSTOM_OPP_PATH'] = os.environ['ASCEND_CUSTOM_OPP_PATH'] + ":" + \
-                                                   cust_ascendc_path + ":" + cust_aicore_path + ":" + cust_aicpu_path
+                                                   cust_ascendc_ascend910_path + ":" + cust_ascendc_ascend910b_path + \
+                                                   ":" + cust_aicore_path + ":" + cust_aicpu_path
         else:
-            os.environ['ASCEND_CUSTOM_OPP_PATH'] = cust_ascendc_path + ":" + cust_aicore_path + ":" + cust_aicpu_path
+            os.environ['ASCEND_CUSTOM_OPP_PATH'] = cust_ascendc_ascend910_path + ":" + \
+                                                   cust_ascendc_ascend910b_path + ":" + cust_aicore_path + ":" + \
+                                                   cust_aicpu_path
+        # Ignore ge infer missing error. To be removed after infers are completed.
+        os.environ['FAST_IGNORE_INFER_ERROR'] = "1"
+        os.environ['IGNORE_INFER_ERROR'] = "1"
         plugin_dir = os.path.dirname(self.library_path)
         akg_dir = os.path.join(plugin_dir, "ascend")
         AscendEnvChecker._concat_variable('LD_LIBRARY_PATH', akg_dir)
@@ -398,7 +409,7 @@ class AscendEnvChecker(EnvChecker):
                 "Environment Variable PYTHONPATH is set. For details, refer to the installation guidelines: "
                 "https://www.mindspore.cn/install")
 
-        if self.ld_lib_path is None or not self.ld_lib_path_check_fwk in self.ld_lib_path:
+        if self.ld_lib_path is None or self.ld_lib_path_check_fwk not in self.ld_lib_path:
             logger.warning("Can not find driver so(need by mindspore-ascend). Please check whether the "
                            "Environment Variable LD_LIBRARY_PATH is set. For details, refer to the installation "
                            "guidelines: https://www.mindspore.cn/install")