mindstudio-probe 1.2.2__py3-none-any.whl → 8.1.0__py3-none-any.whl

msprobe/core/config_check/utils/hyperparameter_parser.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2025-2025, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
+import re
+from abc import ABC, abstractmethod
+
+from msprobe.core.config_check.utils.utils import config_checking_print
+from msprobe.core.common.file_utils import FileOpen, load_yaml
+from msprobe.core.common.const import Const, FileCheckConst
+
+
+class Parser(ABC):
+    @abstractmethod
+    def parse(self, file_path: str) -> dict:
+        pass
+
+    def run(self, file_path: str) -> dict:
+        """
+            统一对外调用接口
+        :param file_path: 需解析的文件路径
+        :return:
+        """
+        try:
+            result = self.parse(file_path)
+        except Exception as exc:
+            config_checking_print(f"{self.__class__} parsing error, skip file path: {file_path}, error: {exc}")
+            result = {}
+        return result
+
+
+class ShellParser(Parser):
+    def parse(self, file_path: str) -> dict:
+        """
+        Extracts arguments from bash script used to run a model training.
+        """
+        hyperparameters = {}
+        script_content_list = []
+        with FileOpen(file_path, 'r') as file:
+            for line in file:
+                stripped_line = line.lstrip()
+                if not stripped_line.startswith('#'):
+                    line = line.split('#')[0].rstrip() + '\n'
+                    if line.strip():
+                        script_content_list.append(line)
+        script_content = ''.join(script_content_list)
+
+        command_line = re.search(r'msrun\s[^|]*|torchrun\s[^|]*|python\d? -m torch.distributed.launch\s[^|]*',
+                                 script_content,
+                                 re.DOTALL)
+        if command_line:
+            command_line = command_line.group()
+
+            blocks = re.findall(r'([a-zA-Z0-9_]{1,20}_ARGS)="(.*?)"', script_content, re.DOTALL)
+            block_contents = {}
+            for block_name, block_content in blocks:
+                block_content = block_content.replace('\n', ' ')
+                block_contents[block_name] = block_content
+                command_line = command_line.replace(f"${block_name}", block_content)
+
+            matches = re.findall(r'--([\w-]+)(?:\s+([^\s\\]+))?', command_line)
+            for match in matches:
+                key, value = match
+                args_key = re.match(r'\$\{?(\w+)}?', value)
+                if args_key:
+                    env_vars = re.findall(rf'{args_key.group(1)}=\s*(.+)', script_content)
+                    if env_vars:
+                        value = env_vars[-1]
+                hyperparameters[key] = value if value else True
+
+        return hyperparameters
+
+
+class YamlParser(Parser):
+    hyperparameters = {}
+
+    def parse(self, file_path: str) -> dict:
+        ori_hyper = load_yaml(file_path)
+        self.recursive_parse_parameters(ori_hyper, "")
+        return self.hyperparameters
+
+    def recursive_parse_parameters(self, parameters, prefix):
+        if isinstance(parameters, dict):
+            for key, value in parameters.items():
+                new_prefix = prefix + Const.SEP + key if prefix else key
+                self.recursive_parse_parameters(value, new_prefix)
+        elif isinstance(parameters, list):
+            for value in parameters:
+                self.recursive_parse_parameters(value, prefix)
+        elif isinstance(parameters, (int, str, bool)):
+            self.hyperparameters.update({prefix: parameters})
+
+
+class ParserFactory:
+    __ParserDict = {
+        FileCheckConst.SHELL_SUFFIX: ShellParser(),
+        FileCheckConst.YAML_SUFFIX: YamlParser()
+    }
+
+    def get_parser(self, file_type: str) -> Parser:
+        parser = self.__ParserDict.get(file_type, None)
+        if not parser:
+            raise ValueError(f'Invalid parser type: {file_type}')
+        return parser

msprobe/core/config_check/utils/utils.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
+import os
+import re
+import hashlib
+
+from msprobe.core.common.framework_adapter import FmkAdp
+from msprobe.core.common.log import logger
+
+
+def merge_keys(dir_0, dir_1):
+    output_list = list(dir_0.keys())
+    output_list.extend(list(dir_1.keys()))
+    return set(output_list)
+
+
+def compare_dict(bench_dict, cmp_dict):
+    result = []
+    for key in set(bench_dict.keys()) | set(cmp_dict.keys()):
+        if key in bench_dict and key in cmp_dict:
+            if bench_dict[key] != cmp_dict[key]:
+                result.append(f"{key}: {bench_dict[key]} -> {cmp_dict[key]}")
+        elif key in bench_dict:
+            result.append(f"{key}: [deleted] -> {bench_dict[key]}")
+        else:
+            result.append(f"{key}: [added] -> {cmp_dict[key]}")
+    return result
+
+
+def config_checking_print(msg):
+    logger.info(f"[config checking log] {msg}")
+
+
+def tensor_to_hash(tensor):
+    """Compute the hash value of a tensor"""
+    tensor_bytes = tensor.clone().detach().cpu().numpy().tobytes()
+    return bytes_hash(tensor_bytes)
+
+
+def get_tensor_features(tensor):
+    features = {
+        "max": FmkAdp.tensor_max(tensor),
+        "min": FmkAdp.tensor_min(tensor),
+        "mean": FmkAdp.tensor_mean(tensor),
+        "norm": FmkAdp.tensor_norm(tensor),
+    }
+
+    return features
+
+
+def compare_dicts(dict1, dict2, path=''):
+    deleted = []
+    added = []
+    changed = []
+    result = {}
+
+    for key in dict1:
+        if key not in dict2:
+            deleted.append(f"[Deleted]: {path + key}")
+            result[key] = "[deleted]"
+        else:
+            if isinstance(dict1[key], dict) and isinstance(dict2[key], dict):
+                sub_deleted, sub_added, sub_changed, sub_result = compare_dicts(
+                    dict1[key], dict2[key], path + key + '/')
+                deleted.extend(sub_deleted)
+                added.extend(sub_added)
+                changed.extend(sub_changed)
+                if sub_result:
+                    result[key] = sub_result
+            elif dict1[key] != dict2[key]:
+                changed.append(f"[Changed]: {path + key} : {dict1[key]} -> {dict2[key]}")
+                result[key] = f"[changed]: {dict1[key]} -> {dict2[key]}"
+    for key in dict2:
+        if key not in dict1:
+            added.append(f"[Added]: {path + key}")
+            result[key] = "[added]"
+    return deleted, added, changed, result
+
+
+def bytes_hash(obj: bytes):
+    hex_dig = hashlib.sha256(obj).hexdigest()
+    short_hash = int(hex_dig, 16) % (2 ** 16)
+    return short_hash
+
+
+def update_dict(ori_dict, new_dict):
+    for key, value in new_dict.items():
+        if key in ori_dict and ori_dict[key] != value:
+            if "values" in ori_dict.keys():
+                ori_dict[key]["values"].append(new_dict[key])
+            else:
+                ori_dict[key] = {"description": "duplicate_value", "values": [ori_dict[key], new_dict[key]]}
+        else:
+            ori_dict[key] = value

msprobe/core/data_dump/api_registry.py

@@ -0,0 +1,239 @@
+# Copyright (c) 2025-2025, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
+import inspect
+from typing import Dict, Any, Optional, Callable, Union, List, Tuple
+
+from msprobe.core.common.const import Const
+from msprobe.core.common.file_utils import load_yaml
+from msprobe.core.common.log import logger
+
+
+def _get_attr(module, attr_name):
+    if Const.SEP in attr_name:
+        sub_module_name, sub_attr = attr_name.rsplit(Const.SEP, 1)
+        sub_module = getattr(module, sub_module_name, None)
+        attr = getattr(sub_module, sub_attr, None)
+    else:
+        attr = getattr(module, attr_name, None)
+    return attr
+
+
+class ApiWrapper:
+    def __init__(
+        self, api_types: Dict[str, Dict[str, Any]],
+        api_list_paths: Union[str, List[str], Tuple[str]],
+        backlist: Union[List[str], Tuple[str]] = None
+    ):
+        self.api_types = api_types
+        if not isinstance(api_list_paths, (list, tuple)):
+            api_list_paths = [api_list_paths] * len(self.api_types)
+        elif len(api_list_paths) != len(self.api_types):
+            raise RuntimeError("The number of api_list_paths must be equal to the number of frameworks in 'api_types', "
+                               "when api_list_paths is a list or tuple.")
+        self.api_list_paths = api_list_paths
+        self.backlist = backlist if backlist else []
+        self.api_names = self._get_api_names()
+        self.wrapped_api_functions = dict()
+
+    @staticmethod
+    def deal_with_self_kwargs(api_name, api_func, args, kwargs):
+        if kwargs and 'self' in kwargs:
+            func_params = None
+            try:
+                func_params = inspect.signature(api_func).parameters
+            except Exception:
+                if api_name in Const.API_WITH_SELF_ARG:
+                    func_params = inspect.signature(Const.API_WITH_SELF_ARG.get(api_name)).parameters
+            if func_params is None:
+                return False, args, kwargs
+
+            for name, param in func_params.items():
+                if name == 'self' and param.kind == inspect.Parameter.KEYWORD_ONLY:
+                    return False, args, kwargs
+            args_ = list(args)
+            names_and_values = []
+            self_index = 0
+            for i, item in enumerate(func_params.items()):
+                names_and_values.append((item[0], item[1].default))
+                if item[0] == 'self':
+                    self_index = i
+                    break
+            for i in range(len(args), self_index + 1):
+                if names_and_values[i][0] in kwargs:
+                    args_.append(kwargs.pop(names_and_values[i][0]))
+                else:
+                    args_.append(names_and_values[i][1])
+            args = tuple(args_)
+
+        return True, args, kwargs
+
+    def wrap_api(
+        self, api_templates, hook_build_func: Optional[Callable]
+    ):
+        api_types_num = sum([len(v) for v in self.api_types.values()])
+        if not isinstance(api_templates, (list, tuple)):
+            api_templates = [api_templates] * api_types_num
+        elif len(api_templates) != api_types_num:
+            raise RuntimeError("The number of api_templates must be equal to the number of api_types, "
+                               "when api_templates is a list or tuple.")
+
+        self.wrapped_api_functions.clear()
+        index = 0
+        for framework, api_types in self.api_types.items():
+            wrapped_functions_in_framework = dict()
+            for api_type, api_modules in api_types.items():
+                wrapped_functions = dict()
+                name_prefix = Const.API_DATA_PREFIX.get(framework, {}).get(api_type, "API")
+                api_template = api_templates[index]
+                index += 1
+                for api_name in self.api_names.get(framework, {}).get(api_type, []):
+                    ori_api = _get_attr(api_modules[0], api_name)
+                    if callable(ori_api):
+                        def wrap_api_func(api_name, api_func, prefix, hook_build_func, api_template):
+                            def api_function(*args, **kwargs):
+                                api_name_with_prefix = prefix + Const.SEP + str(api_name.split(Const.SEP)[-1])
+                                enable_wrap, args, kwargs = self.deal_with_self_kwargs(api_name_with_prefix,
+                                                                                       api_func, args, kwargs)
+                                if not enable_wrap:
+                                    logger.warning(f'Cannot collect precision data of {api_name_with_prefix}. '
+                                                   'It may be fixed by passing the value of "self" '
+                                                   'as a positional argument instead of a keyword argument. ')
+                                    return api_func(*args, **kwargs)
+                                return api_template(api_name, api_func, prefix, hook_build_func)(*args, **kwargs)
+                            api_function.__name__ = api_name
+                            return api_function
+                        wrapped_functions[api_name] = wrap_api_func(api_name, ori_api, name_prefix,
+                                                                    hook_build_func, api_template)
+                wrapped_functions_in_framework[api_type] = wrapped_functions
+            self.wrapped_api_functions[framework] = wrapped_functions_in_framework
+        return self.wrapped_api_functions
+
+    def _get_api_names(self):
+        api_names = dict()
+
+        for index, framework in enumerate(self.api_types.keys()):
+            api_list = load_yaml(self.api_list_paths[index])
+            valid_names = dict()
+            for api_type, api_modules in self.api_types.get(framework, {}).items():
+                key_in_file = Const.SUPPORT_API_DICT_KEY_MAP.get(framework, {}).get(api_type)
+                api_from_file = api_list.get(key_in_file, [])
+                names = set()
+                for api_name in api_from_file:
+                    if f'{key_in_file}.{api_name}' in self.backlist:
+                        continue
+                    target_attr = api_name
+                    target_module = api_modules[0]
+                    if Const.SEP in api_name:
+                        sub_module_name, target_attr = api_name.rsplit(Const.SEP, 1)
+                        target_module = getattr(api_modules[0], sub_module_name, None)
+                    if target_module and target_attr in dir(target_module):
+                        names.add(api_name)
+                valid_names[api_type] = names
+            api_names[framework] = valid_names
+
+        return api_names
+
+
+class ApiRegistry:
+    """
+    Base class for api registry.
+    """
+
+    def __init__(self, api_types, inner_used_api, supported_api_list_path, api_templates, backlist=None):
+        self.ori_api_attr = dict()
+        self.wrapped_api_attr = dict()
+        self.inner_used_ori_attr = dict()
+        self.inner_used_wrapped_attr = dict()
+        self.api_types = api_types
+        self.inner_used_api = inner_used_api
+        self.supported_api_list_path = supported_api_list_path
+        self.api_templates = api_templates
+        self.backlist = backlist if backlist else []
+        self.all_api_registered = False
+
+    @staticmethod
+    def store_ori_attr(ori_api_group, api_list, api_ori_attr):
+        for api in api_list:
+            api_ori_attr[api] = _get_attr(ori_api_group, api)
+
+    @staticmethod
+    def set_api_attr(api_group, attr_dict):
+        for api, api_attr in attr_dict.items():
+            if Const.SEP in api:
+                sub_module_name, sub_op = api.rsplit(Const.SEP, 1)
+                sub_module = getattr(api_group, sub_module_name, None)
+                if sub_module is not None:
+                    setattr(sub_module, sub_op, api_attr)
+            else:
+                setattr(api_group, api, api_attr)
+
+    @staticmethod
+    def register_custom_api(module, api_name, api_prefix, hook_build_func, api_template):
+        def wrap_api_func(api_name, api_func, prefix, hook_build_func, api_template):
+            def api_function(*args, **kwargs):
+                return api_template(api_name, api_func, prefix, hook_build_func)(*args, **kwargs)
+
+            api_function.__name__ = api_name
+            return api_function
+
+        setattr(module, api_name,
+                wrap_api_func(api_name, getattr(module, api_name), api_prefix, hook_build_func, api_template))
+
+    def register_all_api(self):
+        self.all_api_registered = True
+        for framework, api_types in self.api_types.items():
+            for api_type, api_modules in api_types.items():
+                api_type_with_framework = framework + Const.SEP + api_type
+                for module in api_modules[1]:
+                    self.set_api_attr(module, self.wrapped_api_attr.get(api_type_with_framework, {}))
+
+    def register_inner_used_api(self):
+        for api_type in self.inner_used_api.keys():
+            self.set_api_attr(self.inner_used_api.get(api_type)[0], self.inner_used_wrapped_attr.get(api_type, {}))
+
+    def restore_all_api(self):
+        self.all_api_registered = False
+        for framework, api_types in self.api_types.items():
+            for api_type, api_modules in api_types.items():
+                api_type_with_framework = framework + Const.SEP + api_type
+                for module in api_modules[1]:
+                    self.set_api_attr(module, self.ori_api_attr.get(api_type_with_framework, {}))
+
+    def restore_inner_used_api(self):
+        for api_type in self.inner_used_api.keys():
+            self.set_api_attr(self.inner_used_api.get(api_type)[0], self.inner_used_ori_attr.get(api_type, {}))
+
+    def initialize_hook(self, hook_build_func):
+        api_wrapper = ApiWrapper(self.api_types, self.supported_api_list_path, self.backlist)
+        wrapped_api_functions = api_wrapper.wrap_api(self.api_templates, hook_build_func)
+
+        for framework, api_types in self.api_types.items():
+            for api_type, api_modules in api_types.items():
+                ori_attr = dict()
+                self.store_ori_attr(api_modules[0], api_wrapper.api_names.get(framework).get(api_type), ori_attr)
+                api_type_with_framework = framework + Const.SEP + api_type
+                self.ori_api_attr[api_type_with_framework] = ori_attr
+                self.wrapped_api_attr[api_type_with_framework] = wrapped_api_functions.get(framework).get(api_type)
+
+        for inner_used_api_type, inner_used_api_list in self.inner_used_api.items():
+            ori_attr = dict()
+            wrapped_attr = dict()
+            for api_name in inner_used_api_list[1:]:
+                if self.ori_api_attr.get(inner_used_api_type, {}).get(api_name):
+                    ori_attr[api_name] = self.ori_api_attr.get(inner_used_api_type).get(api_name)
+                    wrapped_attr[api_name] = self.wrapped_api_attr.get(inner_used_api_type).get(api_name)
+            self.inner_used_ori_attr[inner_used_api_type] = ori_attr
+            self.inner_used_wrapped_attr[inner_used_api_type] = wrapped_attr

msprobe/core/data_dump/data_collector.py

@@ -41,7 +41,7 @@ class DataCollector:
         self.backward_module_names = {}
         self.optimizer_status = ""
         self.optimizer_status_first_start = {Const.OPTIMIZER: True, Const.CLIP_GRAD: True}
-        atexit.register(self.write_json)
+        atexit.register(self.write_json_at_exit)
 
     @property
     def dump_data_dir(self):
@@ -78,6 +78,11 @@ class DataCollector:
     def write_json(self):
         self.data_writer.write_json()
 
+    def write_json_at_exit(self):
+        if self.config.async_dump and self.config.task == Const.TENSOR:
+            self.data_processor.dump_async_data()
+        self.data_writer.write_json()
+
     def update_data(self, name, data_info):
         msg = f"msprobe is collecting data on {name}."
         if self.config.task == Const.OVERFLOW_CHECK:
@@ -89,6 +94,10 @@ class DataCollector:
         logger.debug(msg)
         self.data_writer.update_data(data_info)
 
+    def call_stack_collect(self, name):
+        stack_info = self.data_processor.analyze_api_call_stack(name)
+        self.data_writer.update_stack(name, stack_info)
+
     def forward_input_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
         if self.config.task == Const.FREE_BENCHMARK:
             backward_name = name.replace(Const.FORWARD, Const.BACKWARD)
@@ -118,9 +127,16 @@ class DataCollector:
         self.set_is_recomputable(data_info, is_recompute)
         if self.config.level == Const.LEVEL_L2:
             return
-        self.data_writer.update_stack(self.data_processor.analyze_api_call_stack(name))
+        self.call_stack_collect(name)
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
+    def forward_data_collect_only_tensor(self, name, module, pid, module_input_output):
+        if not self.check_scope_and_pid(self.scope, name, pid):
+            return
+
+        self.data_processor.analyze_forward(name, module, module_input_output)
+
+
     def forward_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
@@ -130,9 +146,15 @@ class DataCollector:
         if self.config.task != Const.STRUCTURE:
             data_info = self.data_processor.analyze_forward(name, module, module_input_output)
         self.set_is_recomputable(data_info, is_recompute)
-        self.data_writer.update_stack(self.data_processor.analyze_api_call_stack(name))
+        self.call_stack_collect(name)
         self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
+    def backward_data_collect_only_tensor(self, name, module, pid, module_input_output, is_recompute=None):
+        if not self.check_scope_and_pid(self.scope, name, pid):
+            return
+
+        self.data_processor.analyze_backward(name, module, module_input_output)
+
     def backward_data_collect(self, name, module, pid, module_input_output, is_recompute=None):
         self.update_construct(name)
         if not self.check_scope_and_pid(self.scope, name, pid):
@@ -180,7 +202,10 @@ class DataCollector:
                     self.optimizer_status_first_start[self.optimizer_status] = False
                 self.data_writer.update_construct({name: self.optimizer_status})
             else:
-                self.data_writer.update_construct({name: self.module_processor.api_parent_node})
+                if self.config.level == Const.LEVEL_MIX and \
+                  not (name.startswith(Const.MODULE) or name.startswith(Const.CELL)):
+                    self.data_writer.update_construct({name: self.module_processor.api_parent_node})
+
             self.data_writer.update_construct(self.module_processor.module_node)
 
     def handle_data(self, name, data_info, flush=False):
@@ -204,6 +229,7 @@ class DataCollector:
 
     def params_data_collect(self, name, param_name, pid, data):
         grad_name = name + Const.SEP + Const.PARAMS_GRAD
+        self.update_api_or_module_name(grad_name)
         # 校验scope和pid，以及当前name是否有过反向计算
         if not self.check_scope_and_pid(self.scope, name, pid) and not self.backward_module_names.get(name):
             # 如果没有反向计算，则需要清除之前占位写入的grad数据
@@ -213,18 +239,19 @@ class DataCollector:
         data_info = self.data_processor.analyze_params(grad_name, param_name, data)
         self.handle_data(grad_name, data_info, flush=self.data_processor.is_terminated)
 
-    def fill_stack_tensor_data(self):
-        self.data_writer.fill_stack_tensor_data()
 
     def debug_data_collect_forward(self, variable, name_with_count):
 
         data_info = self.data_processor.analyze_debug_forward(variable, name_with_count)
-        self.data_writer.update_debug({name_with_count: data_info})
+        name_with_count_category = name_with_count + Const.SEP + Const.DEBUG
+        self.data_writer.update_debug({name_with_count_category: data_info})
 
     def debug_data_collect_backward(self, variable, grad_name_with_count):
         # prepare all None nested data structure
         all_none_data_info = self.data_processor.analyze_element_to_all_none(variable)
-        self.data_writer.update_debug({grad_name_with_count: all_none_data_info})
+        grad_name_with_count_category = grad_name_with_count + Const.SEP + Const.DEBUG
+        self.data_writer.update_debug({grad_name_with_count_category: all_none_data_info})
 
         # register tensor backward hook
-        self.data_processor.analyze_debug_backward(variable, grad_name_with_count, self.data_writer.cache_debug['data'])
+        self.data_processor.analyze_debug_backward(variable, grad_name_with_count_category,
+                                                   self.data_writer.cache_debug['data'])