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

mindspore/parallel/strategy.py

@@ -0,0 +1,336 @@
+# Copyright 2025 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Checkpoint strategy info"""
+from __future__ import absolute_import
+
+__all__ = ["get_strategy_metadata", "get_current_strategy_metadata", "enable_save_strategy_online", \
+           "clear_strategy_metadata"]
+
+from itertools import chain
+from typing import Sequence, Union, Tuple, List, Dict
+from types import SimpleNamespace
+
+import numpy as np
+
+from mindspore import log as logger
+from mindspore._c_expression import StrategyInfo
+from mindspore._c_expression import StrategyLayout
+from mindspore.parallel.shard import Layout
+
+LayoutInfo = Tuple[Layout, str, str]
+StrOrTuple = Union[str, Tuple["StrOrTuple", ...], List["StrOrTuple"]]
+
+
+def get_strategy_metadata(network, rank_id=None) -> Dict[int, Dict[str, List[LayoutInfo]]]:
+    """
+    Get all params strategy info or specific rank strategy info in this cell.
+    For more information on layouts, please refer to: :class:`mindspore.parallel.Layout`.
+
+    Args:
+        network (str): The network name.
+        rank_id (int, optional): The rank id of the process on which this cell will be launched.
+            Defaults to ``None``, which means strategy metadata for all ranks will be returned.
+
+    Returns:
+        Dict. A dictionary containing the parameter slicing strategies for either all ranks or a specific rank.
+        The key is `rank_id`, and the value is the slicing strategy for all parameters on that rank.
+        Within each rank's strategy, the key is the parameter name, and the value is the slicing strategy.
+        If a `rank_id` is specified, the dictionary returns the strategy information for that specific rank.
+        Otherwise, it returns the strategy information for all ranks in the network. If not supported, returns None.
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import nn
+        >>> from mindspore.communication import init
+        >>> from mindspore.nn.utils import no_init_parameters
+        >>> from mindspore.parallel.auto_parallel import AutoParallel
+        >>> from mindspore.train import Model
+        >>> from mindspore.parallel.strategy import get_strategy_metadata, get_current_strategy_metadata,
+        ...     enable_save_strategy_online, clear_strategy_metadata
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> init()
+        >>> ms.set_seed(1)
+        >>>
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> with no_init_parameters():
+        ...     net = LeNet5()
+        ...     optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
+        >>>
+        >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
+        >>> train_net = AutoParallel(net, parallel_mode="semi_auto")
+        >>> model = Model(network=train_net, loss_fn=loss, optimizer=optim, metrics=None)
+        >>>
+        >>> # Create the dataset taking MNIST as an example. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
+        >>> dataset = create_dataset()
+        >>>
+        >>> enable_save_strategy_online()
+        >>> model.train(2, dataset)
+        >>>
+        >>> global_info = get_strategy_metadata(network=model.train_network)
+        >>> rank0_info = get_strategy_metadata(network=model.train_network, rank_id=0)
+        >>> local_info = get_current_strategy_metadata(network=model.train_network)
+        >>> clear_strategy_metadata()
+    """
+    return _NetStrategyInfo(network, global_layout=None, local_layout=None).get_rank_layout(rank_id)
+
+
+def get_current_strategy_metadata(network) -> Dict[int, Dict[str, List[LayoutInfo]]]:
+    """
+    Get parameters dictionary of cur rank of the network.
+
+    Args:
+        network(str): The network name.
+
+    Returns:
+        Dict. The key is 0 (representing the local rank), and the value is the slicing strategy for all parameters.
+        The key within the value represents the parameter name, and the value is the corresponding slicing strategy \
+        for that parameter. If not supported, returns None.
+    """
+    return _NetStrategyInfo(network, global_layout=None, local_layout=None).get_local_rank_layout()
+
+
+def enable_save_strategy_online():
+    """
+    Enable save strategy metadata online.
+    """
+    strategy_layout_handle = StrategyLayout.get_instance()
+    if strategy_layout_handle is None:
+        raise ValueError("Strategy layout handle is none in parallel_strategy_checkpoint!!!")
+    strategy_layout_handle.enable_save_strategy_online()
+
+
+def clear_strategy_metadata():
+    """Clear all saved strategy metadata on the C++ side."""
+    strategy_layout_handle = StrategyLayout.get_instance()
+    if strategy_layout_handle is None:
+        raise ValueError("Strategy layout handle is none in parallel_strategy_checkpoint!!!")
+    return strategy_layout_handle.clear_strategy_metadata()
+
+
+class _NetStrategyInfo:
+    """
+    Describe the strategy information of a network.
+    """
+
+    def __init__(self, network, global_layout=None, local_layout=None):
+        self._network = network
+        self._compile_phase = network.compile_phase
+        if global_layout is None or local_layout is None:
+            layout_handle = self._get_layout_handle()
+            global_layout = layout_handle.global_network_layout()
+            local_layout = layout_handle.local_network_layout()
+        self._raw_global_layout = global_layout
+        self._raw_local_layout = local_layout
+
+    @staticmethod
+    def _get_layout_handle():
+        """Get strategy handle"""
+        layout_handle = StrategyLayout.get_instance()
+        if layout_handle is None:
+            raise ValueError("Strategy layout handle is none in parallel_strategy_checkpoint!!!")
+        return layout_handle
+
+    def get_rank_layout(self, rank_id=None):
+        """Get params of the network, global rank or special rank, interface."""
+        raw_global_layout = self._get_valid_layout(self._compile_phase, self._raw_global_layout)
+        if raw_global_layout is None:
+            return None
+        global_layout = self._extract_layout_metadata(raw_global_layout)
+        if rank_id is not None:
+            cur_rank_layout = {rank_id: global_layout[rank_id]}
+            self._layout_to_string(cur_rank_layout)
+            return cur_rank_layout
+        self._layout_to_string(global_layout)
+        return global_layout
+
+    def get_local_rank_layout(self):
+        """Get local rank params of the network, {param_name: param_info[layout]}."""
+        raw_local_layout = self._get_valid_layout(self._compile_phase, self._raw_local_layout)
+        if raw_local_layout is None:
+            return None
+        local_layout = self._extract_layout_metadata(raw_local_layout)
+        self._layout_to_string(local_layout)
+        return local_layout
+
+    @staticmethod
+    def _get_valid_layout(phase, layout_dict):
+        """Helper: Validate and extract layout by phase."""
+        if not phase:
+            return None
+        layout = layout_dict.get(phase)
+        if not layout or all(not v for v in layout.values()):
+            return None
+        return layout
+
+    def _extract_layout_metadata(self, layout: Dict[int, Dict[str, StrategyInfo]]) -> Dict:
+        """Return new layout of special network."""
+        new_layout = {}
+        for rank_id, param_dict in layout.items():
+            new_param_info = {}
+            for param_name, param_info in param_dict.items():
+                new_param_layout = self._layout_process(param_info)
+                new_param_info[param_name] = new_param_layout
+            new_layout[rank_id] = new_param_info
+        return new_layout
+
+    def _layout_process(self, stra_layout):
+        """
+        Return the layout list, stra_layout is one of params_info of cur_rank.
+        """
+        new_dev_mat, counter, new_tensor_map, full_opt_shard = self._get_dev_mat_for_opt_shard(
+            stra_layout.opt_weight_shard_size, stra_layout.dev_matrix, stra_layout.tensor_map)
+        alphabet = 'abcdefghijklmnopqrstuvwxyz'
+        alias_name = [alphabet[i] for i in range(len(new_dev_mat))]
+        if stra_layout.opt_weight_shard_size == 0:
+            new_tensor_map = tuple(tuple(alias_name[len(alias_name) - idx - 1] if idx != -1 else "None" for idx in sub)
+                                   for sub in new_tensor_map)
+        else:
+            info = SimpleNamespace(
+                new_dev_mat=new_dev_mat,
+                new_tensor_map=new_tensor_map,
+                full_opt_shard=full_opt_shard,
+                counter=counter,
+                alias_name=alias_name
+            )
+            new_tensor_map = self._get_tensor_map_for_opt_shard(info)
+        new_tensor_map = self._compact_tensor_map(new_tensor_map)
+        new_dev_mat = tuple(new_dev_mat)
+        alias_name = tuple(alias_name)
+        layout = Layout(new_dev_mat, alias_name, stra_layout.rank_list)
+        final_layout = layout(*new_tensor_map)
+        logger.debug("The final layout is %s", final_layout.to_dict())
+        cur_param_list = [final_layout, stra_layout.tensor_type, stra_layout.tensor_shape]
+        return cur_param_list
+
+    def _get_dev_mat_for_opt_shard(self, opt_shard, dev_mat, tensor_map):
+        """generate device matrix for opt shard scenario"""
+        if opt_shard == 0:
+            return dev_mat, -1, tensor_map, True
+        used_dev_num = self._calc_used_dev_num(dev_mat, tensor_map)
+        total_dev_num = int(np.prod(np.array(dev_mat)))
+        if opt_shard == -1 or used_dev_num * opt_shard == total_dev_num:
+            return dev_mat, -1, tensor_map, True
+        remain_dev_num = total_dev_num // (used_dev_num * opt_shard)
+        used_dev_mat_mask = self._get_used_dev_mat(dev_mat, tensor_map)
+        info = SimpleNamespace(
+            dev_mat=dev_mat,
+            tensor_map=tensor_map,
+            counter=-1,
+            real_remain_dev_num=1,
+            remain_dev_num=remain_dev_num
+        )
+        for axis, value in enumerate(dev_mat):
+            if used_dev_mat_mask[axis]:
+                continue
+            info.counter = axis
+            if info.real_remain_dev_num == info.remain_dev_num:
+                return dev_mat, axis, tensor_map, False
+            if info.real_remain_dev_num < info.remain_dev_num:
+                info.real_remain_dev_num *= value
+                continue
+            # info.real_remain_dev_num > info.remain_dev_num，split axis.
+            return self._split_dev_dim(info)
+        if info.real_remain_dev_num == info.remain_dev_num:
+            return dev_mat, info.counter, tensor_map, False
+        return self._split_dev_dim(info)
+
+    def _get_tensor_map_for_opt_shard(self, info: SimpleNamespace):
+        """generate tensor map for opt shard scenario"""
+
+        def idx_to_alias(idx):
+            return "None" if idx == -1 else info.alias_name[len(info.alias_name) - idx - 1]
+
+        def entry_to_alias(entry):
+            if isinstance(entry, (list, tuple)):
+                return tuple(idx_to_alias(i) for i in entry)
+            return idx_to_alias(entry)
+
+        used_dev_mat = self._get_used_dev_mat(info.new_dev_mat, info.new_tensor_map)
+        if info.full_opt_shard:
+            unused_idx = [len(used_dev_mat) - i - 1 for i, used in enumerate(used_dev_mat) if not used]
+        else:
+            unused_idx = [len(used_dev_mat) - i - 1 for i, used in enumerate(used_dev_mat) if
+                          not used and i > info.counter]
+        first_entry = info.new_tensor_map[0]
+        first_list = list(first_entry) if isinstance(first_entry, (list, tuple)) else [first_entry]
+        new_first_list = [dim for dim in first_list + unused_idx if dim != -1]
+        first_alias_list = [idx_to_alias(i) for i in new_first_list] or ["None"]
+        first_alias = first_alias_list[0] if len(first_alias_list) == 1 else tuple(first_alias_list)
+        rest_alias = [entry_to_alias(entry) for entry in info.new_tensor_map[1:]]
+        new_tensor_map = tuple([first_alias] + rest_alias)
+        return new_tensor_map
+
+    @staticmethod
+    def _split_dev_dim(info: SimpleNamespace):
+        """Split the counter dimension of dev_mat and adjust tensor_map."""
+        dev_mat = info.dev_mat
+        counter = info.counter
+        splitted_dev_value = dev_mat[counter]
+        new_dev_mat_value_first = info.remain_dev_num // (info.real_remain_dev_num // splitted_dev_value)
+        new_dev_mat_value_second = splitted_dev_value // new_dev_mat_value_first
+        new_dev_mat = dev_mat[:counter] + [new_dev_mat_value_first, new_dev_mat_value_second] + dev_mat[counter + 1:]
+        flag = len(new_dev_mat) - 1 - counter
+        new_tensor_map = [[v if v < flag or v == -1 else v + 1 for v in sub] for sub in info.tensor_map]
+        return new_dev_mat, counter, new_tensor_map, False
+
+    @staticmethod
+    def _calc_used_dev_num(dev_mat, tensor_map):
+        """Count the total number of device nums that have been used."""
+        idx_flat = [idx for idx in chain.from_iterable(tensor_map) if idx != -1]
+        if not idx_flat:
+            return 1
+        prod_list = [dev_mat[len(dev_mat) - idx - 1] for idx in idx_flat]
+        return int(np.prod(prod_list))
+
+    @staticmethod
+    def _get_used_dev_mat(dev_mat, tensor_map) -> List[bool]:
+        """List that records whether the device ID is being used or not."""
+        used = set()
+        for elem in tensor_map:
+            if isinstance(elem, (list, tuple)):
+                used.update(i for i in elem if i != -1)
+            elif elem != -1:
+                used.add(elem)
+        return [(len(dev_mat) - i - 1) in used for i in range(len(dev_mat))]
+
+    @staticmethod
+    def _compact_tensor_map(alias_map: Sequence[StrOrTuple]) -> Tuple[StrOrTuple, ...]:
+        """Extend tensor map of 'None'."""
+
+        def _compress(elem: StrOrTuple) -> StrOrTuple:
+            if isinstance(elem, (list, tuple)):
+                compressed = tuple(_compress(e) for e in elem)
+                if len(compressed) == 1:
+                    return compressed[0]
+                if all(x == 'None' for x in compressed):
+                    return 'None'
+                return compressed
+            return elem
+
+        return tuple(_compress(e) for e in alias_map)
+
+    @staticmethod
+    def _layout_to_string(layout_info):
+        """Print layout info."""
+        for rank_id, param_layout in layout_info.items():
+            logger.info("rank_id=%s", rank_id)
+            for param_name, cur_param_list in param_layout.items():
+                final_layout, param_type, global_shape = cur_param_list
+                logger.info("param_name=%s: [param_layout=%s, param_type=%s, global_shape=%s]",
+                            param_name, final_layout.to_dict(), param_type, global_shape)
+            logger.info("\n")

mindspore/parallel/transform_safetensors.py

@@ -15,6 +15,7 @@
 """Transform distributed safetensors"""
 from __future__ import absolute_import
 
+import copy
 import os
 import sys
 import glob
@@ -37,7 +38,6 @@ from mindspore import log as logger
 from mindspore.log import vlog_print
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
-from mindspore.common import np_dtype
 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_expand_layout_reshape
@@ -51,8 +51,6 @@ from mindspore.common import dtype as mstype
 
 safetensors_to_mstype = {'Int4': mstype.qint4x2}
 
-np.bfloat16 = np_dtype.bfloat16
-
 MAX_HEADER_SIZE = 100 * 1000 * 1000
 
 dtype_size = {
@@ -71,6 +69,7 @@ dtype_size = {
     "F64": 8,
 }
 np_dtype_size = {
+    "bool": 1,
     "bool_": 1,
     "uint8": 1,
     "int8": 1,
@@ -96,7 +95,6 @@ numpy_dtype = {
     "I64": np.int64,
     "U64": np.uint64,
     "F16": np.float16,
-    "BF16": np.bfloat16,  # no bf16
     "F32": np.float32,
     "F64": np.float64,
 }
@@ -215,6 +213,16 @@ class PySafeSlice:
 
     @property
     def dtype(self):
+        """Get dtype by numpy_dtype"""
+        if self.info["dtype"] == "BF16":
+            from mindspore.common import np_dtype
+            if not np_dtype.np_dtype_valid(True):
+                raise TypeError(
+                    "The Numpy bfloat16 data type is not supported now, please ensure that the current "
+                    "Numpy version is not less than the version when the mindspore is compiled, "
+                    "and the major versions are same."
+                )
+            return np_dtype.bfloat16
         return numpy_dtype[self.info["dtype"]]
 
     @property
@@ -690,6 +698,8 @@ def _transform_safetensors_single(needed_rank_list_map, all_safetensor_files_map
             else:
                 if transform_param_dict:
                     if output_format == "safetensors":
+                        if meta_data and "remove_redundancy" in meta_data:
+                            meta_data["remove_redundancy"] = "False"
                         _save_file_atomically(transform_param_dict, save_file_name, metadata=meta_data)
                     else:
                         transform_param_dict = _load_and_transform(transform_param_dict, None, None,
@@ -759,6 +769,11 @@ def transform_safetensors_by_stage(src_safetensors_dir, dst_safetensors_dir, ckp
             param_type_dict[param_name][src_rank] = str(param.data.dtype)
             param_total_dict[param_name][src_rank] = param
             param_attr_dict[param_name][src_rank] = (True, False)
+
+    ckpt_prefix = os.path.basename(ckpt_prefix)
+    if '..' in ckpt_prefix or '/' in ckpt_prefix or '\\' in ckpt_prefix:
+        raise ValueError(f"Invalid ckpt_prefix: {ckpt_prefix}. Must not contain path traversal characters.")
+
     for local_rank_id in range(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,
@@ -776,6 +791,7 @@ def transform_safetensors_by_rank(rank_id, safetensor_files_map, save_safetensor
     """
     Transform distributed checkpoint from source sharding strategy to destination sharding strategy by rank.
     """
+    save_safetensor_file_name = os.path.abspath(save_safetensor_file_name)
     if not isinstance(safetensor_files_map, dict):
         raise TypeError("The safetensor_files_map should be a dict.")
     if not isinstance(rank_id, int):
@@ -823,11 +839,84 @@ def transform_safetensors_by_rank(rank_id, safetensor_files_map, save_safetensor
     _save_file_atomically(transform_param_dict, save_safetensor_file_name, metadata={"format": "ms"})
 
 
-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 _extract_numbers(s):
+    """Extract all numbers from a string and convert them to integers."""
+    return [int(num) for num in re.findall(r'\d+', s)]
+
+
+def _extract_last_two_numbers(file_name):
+    """Get the last two numbers from a filename."""
+    all_numbers = _extract_numbers(file_name)
+    return all_numbers[-2:]
+
+
+def _find_shortest_file(matched_files, rank_ckpts, new_file_suffix, file_suffix):
+    """Find the shortest file from a list of matched files."""
+    min_length = min(len(os.path.basename(ckpt)) for ckpt in matched_files)
+    shortest_files = [ckpt for ckpt in matched_files if len(os.path.basename(ckpt)) == min_length]
+    if len(shortest_files) == 1:
+        return shortest_files[0]
+    raise ValueError(f"Multiple files with suffix '{file_suffix}' found in {rank_ckpts}. Following MindSpore naming "
+                     f"rules, searched for files ending with '{new_file_suffix}' but found multiple "
+                     f"files {matched_files}. Then searched for the shortest filename, but found multiple shortest "
+                     f"files {shortest_files}. Please set file_suffix to the longest common suffix of all files.")
+
+
+def _get_matched_file(matched, rank_ckpts, new_file_suffix, file_suffix):
+    """Get the file from a list of matched files."""
+    if len(matched) == 1:
+        return matched[0]
+    if len(matched) > 1:
+        return _find_shortest_file(matched, rank_ckpts, new_file_suffix, file_suffix)
+    raise ValueError(f"Multiple files with suffix '{file_suffix}' found in {rank_ckpts}. Following MindSpore naming "
+                     f"rules, searched for files ending with '{new_file_suffix}' but found zero files. "
+                     f"Please set file_suffix to the longest common suffix of all files.")
+
+
+def _find_most_matching_file(rank_ckpts, file_suffix, format):
+    """Finds the most matching checkpoint file based on the file_suffix."""
+    if file_suffix is None:
+        rank_ckpts.sort(key=_extract_last_two_numbers)
+        return rank_ckpts[-1]
+
+    new_file_suffix = file_suffix
+    pattern1 = rf'^_(\d+)-(\d+)_(\d+)$'
+    matches1 = re.search(pattern1, file_suffix)
+    pattern2 = rf'^(\d+)-(\d+)_(\d+)$'
+    matches2 = re.search(pattern2, file_suffix)
+    # Pattern matching for _{task_id}-{epoch}_{step} format (e.g., _1-10_100 or 1-10_100)
+    if matches1 is not None or matches2 is not None:
+        if matches2 is not None:
+            new_file_suffix = "_" + new_file_suffix
+        matched = [ckpt for ckpt in rank_ckpts if ckpt.endswith(f"{new_file_suffix}.{format}") and
+                   not ckpt.endswith(f"rank{new_file_suffix}.{format}")]
+        return _get_matched_file(matched, rank_ckpts, new_file_suffix, file_suffix)
+
+    pattern3 = rf'^-(\d+)_(\d+)$'
+    matches3 = re.search(pattern3, file_suffix)
+    pattern4 = rf'^(\d+)_(\d+)$'
+    matches4 = re.search(pattern4, file_suffix)
+    # Pattern matching for -{epoch}_{step} format (e.g., -10_100 or 10_100)
+    if matches3 is not None or matches4 is not None:
+        if matches4 is not None:
+            new_file_suffix = "-" + new_file_suffix
+        matched = [ckpt for ckpt in rank_ckpts if ckpt.endswith(f"{new_file_suffix}.{format}")]
+        return _get_matched_file(matched, rank_ckpts, new_file_suffix, file_suffix)
+
+    pattern5 = rf'^_(\d+)$'
+    matches5 = re.search(pattern5, file_suffix)
+    pattern6 = rf'^(\d+)$'
+    matches6 = re.search(pattern6, file_suffix)
+    # Pattern matching for _{step} format (e.g., _100 or 100)
+    if matches5 is not None or matches6 is not None:
+        if matches6 is not None:
+            new_file_suffix = "_" + new_file_suffix
+        matched = [ckpt for ckpt in rank_ckpts if ckpt.endswith(f"{new_file_suffix}.{format}")]
+        return _get_matched_file(matched, rank_ckpts, new_file_suffix, file_suffix)
+
+    raise ValueError(f"Multiple {format} files ending with '{file_suffix}' found in {rank_ckpts}. "
+                     f"Cannot determine which file is the intended one. "
+                     f"Please set file_suffix to the longest common suffix.")
 
 
 def _collect_safetensor_files(src_safetensors_dir, format='safetensors', file_suffix=None):
@@ -838,6 +927,9 @@ def _collect_safetensor_files(src_safetensors_dir, format='safetensors', file_su
         return {0: src_safetensors_dir}
     safetensors_rank_dir_list = os.path.join(src_safetensors_dir, "rank_[0-9]*")
     all_safetensor_files_map = {}
+    multiple_files_found_flag = False
+    multiple_files_list = None
+    chosen_file = None
     for safetensor_dir in glob.glob(safetensors_rank_dir_list):
         if not os.path.isdir(safetensor_dir):
             ms.log.warning("{} is not a directory.".format(safetensor_dir))
@@ -853,9 +945,23 @@ 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(key=_extrace_number)
-        if rank_ckpts:
-            all_safetensor_files_map[rank_id] = rank_ckpts[-1]
+        if len(rank_ckpts) > 1:
+            all_safetensor_files_map[rank_id] = _find_most_matching_file(rank_ckpts, file_suffix, format)
+            if not multiple_files_found_flag:
+                multiple_files_found_flag = True
+                multiple_files_list = copy.deepcopy(rank_ckpts)
+                chosen_file = all_safetensor_files_map[rank_id]
+        elif rank_ckpts:
+            all_safetensor_files_map[rank_id] = rank_ckpts[0]
+        elif file_suffix is not None:
+            raise ValueError(f"No safetensors files found in directory '{safetensor_dir}' "
+                             f"with suffix '{file_suffix}' and format '{format}'. "
+                             f"Please verify the directory contains the expected files. "
+                             f"Recommend setting file_suffix to the longest common suffix.")
+    if file_suffix is not None and multiple_files_found_flag:
+        logger.warning(f"When unified_safetensors files with file_suffix `{file_suffix}`, multiple files were found. "
+                       f"Showing one list: {multiple_files_list}; selected `{chosen_file}` from it. "
+                       f"Please check whether the file_suffix is set correctly.")
     return all_safetensor_files_map
 
 
@@ -972,7 +1078,7 @@ def _transform_parallel_safetensor(rank_id, param_total_dict, param_attr_dict, s
 def _cal_param_size(shape, dtype):
     """cal param size by dtype and shape"""
     num_elements = math.prod(shape)
-    element_size = np_dtype_size.get(dtype, 4)
+    element_size = np_dtype_size.get(str(dtype), 4)
     total_bytes = num_elements * element_size
     return total_bytes
 
@@ -1135,7 +1241,7 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
     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, format="safetensors", file_suffix=file_suffix)
-    all_ckpt_files_map = _collect_safetensor_files(src_dir, format="ckpt", file_suffix=file_suffix)
+    all_ckpt_files_map = _collect_safetensor_files(src_dir, format="ckpt")
     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")
@@ -1173,11 +1279,6 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
         with _fast_safe_open(file_name, framework="np") as f:
             for k in f.keys():
                 if k in name_list:
-                    py_slice = f.get_tensor(k)
-                    param_total_size += _cal_param_size(py_slice.shape, py_slice.dtype)
-                    param_dst_shape = _get_dst_shape(k, py_slice.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):
@@ -1185,7 +1286,13 @@ def unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundan
                                 name_list.remove(k)
                                 continue
                     if k not in param_size_dict:
-                        param_size_dict[k] = _cal_param_size(param_dst_shape, py_slice.dtype)
+                        py_slice = f.get_tensor(k)
+                        param_dst_shape = _get_dst_shape(k, py_slice.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]
+                        param_size = _cal_param_size(param_dst_shape, py_slice.dtype)
+                        param_total_size += param_size
+                        param_size_dict[k] = param_size
     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)

mindspore/profiler/analysis/viewer/ascend_kernel_details_viewer.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ============================================================================
 """Ascend kernel details viewer"""
-import csv
 import os
 from decimal import Decimal
 
@@ -25,8 +24,7 @@ from mindspore.profiler.common.constant import (
     ProfilerActivity
 )
 from mindspore.profiler.common.log import ProfilerLogger
-from mindspore.profiler.common.path_manager import PathManager
-
+from mindspore.profiler.common.file_manager import FileManager
 from mindspore import log as logger
 
 
@@ -110,12 +108,15 @@ class AscendKernelDetailsViewer(BaseViewer):
         Write data to csv file.
         """
         self._logger.info("Kernel details saved start")
-        PathManager.check_directory_path_writeable(os.path.dirname(self._save_path))
-        with open(self._save_path, "w", newline="", encoding="utf-8") as csvfile:
-            writer = csv.writer(csvfile)
-            writer.writerow(self.kernel_details_headers)
-            for row in self.op_summary:
-                writer.writerow([row[field] for field in self.op_summary_headers])
+        csv_data = []
+        for row in self.op_summary:
+            csv_row = [row[field] for field in self.op_summary_headers]
+            csv_data.append(csv_row)
+        FileManager.create_csv_file(
+            file_path=self._save_path,
+            data=csv_data,
+            headers=self.kernel_details_headers
+        )
         self._logger.info("Kernel details saved done")
 
     def _update_headers(self):
@@ -247,4 +248,7 @@ def _get_step_id_by_ts(ts: Decimal, step_events_dict: dict):
         if st <= ts <= et:
             return step_id
 
+    if step_events_dict:
+        return list(step_events_dict.keys())[-1]
+
     return None

mindspore/profiler/analysis/viewer/ascend_op_memory_viewer.py

@@ -330,7 +330,7 @@ class AscendOpMemoryViewer:
                 res.append(self._combine_alloc_and_free_event(alloc_event, free_event))
                 alloc_event, free_event = None, None
             elif alloc_event is None and free_event:
-                self._logger.error("Alloc event is None, but free event is not None")
+                self._logger.warning("Alloc event is None, but free event is not None")
 
         if alloc_event:
             res.append(self._combine_alloc_and_free_event(alloc_event))

mindspore/profiler/analysis/viewer/ms_minddata_viewer.py

@@ -500,7 +500,7 @@ class BottleneckAnalyzer:
             in_op_id, out_q = self._get_non_inline_child_recur(op_id), self.queue_utilization_pct[op_id]
             # This is a leaf node since input queue does not exist and output queue exists
             if in_op_id == self.op_id_not_exist and out_q != self.queue_usage_not_exist:
-                if out_q < self._THRESHOLDS['_LEAF_OUTPUT_QUEUE_EMPTY_FREQ_PCT_MAXIMUM']:
+                if out_q <= self._THRESHOLDS['_LEAF_OUTPUT_QUEUE_EMPTY_FREQ_PCT_MAXIMUM']:
                     queue_usage_analysis.append(self._format_leaf_node_suggestion(op_id, out_q))
             # This is device_queue op
             elif self.op_names[op_id] == "DeviceQueue" and in_op_id != self.op_id_not_exist:

mindspore/profiler/common/constant.py

@@ -226,3 +226,8 @@ class HostSystem(Enum):
     DISK = "disk"
     NETWORK = "network"
     OSRT = "osrt"
+
+
+class MsprofModeName:
+    """msprof mode name"""
+    MSPROF_DYNAMIC_ENV = "PROFILING_MODE"