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

msprobe/core/debugger/precision_debugger.py

@@ -0,0 +1,143 @@
+# 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 os
+
+from msprobe.core.common.const import Const, FileCheckConst, MsgConst
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.common.file_utils import FileChecker, load_json
+from msprobe.core.common.utils import get_real_step_or_rank, check_init_step
+from msprobe.core.common_config import CommonConfig
+
+
+class BasePrecisionDebugger:
+    _instance = None
+    tasks_not_need_debugger = [Const.GRAD_PROBE]
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(BasePrecisionDebugger, cls).__new__(cls)
+            cls._instance.config = None
+            cls._instance.enable_dataloader = False
+            cls._instance.initialized = False
+            cls.service = None
+            cls.first_start = False
+        return cls._instance
+
+    def __init__(
+            self,
+            config_path=None,
+            task=None,
+            dump_path=None,
+            level=None,
+            step=None
+    ):
+        if self.initialized:
+            return
+        self.initialized = True
+        self.check_input_params(config_path, task, dump_path, level)
+        self.common_config, self.task_config = self.parse_config_path(config_path, task)
+        self.task = self.common_config.task
+        if step is not None:
+            self.common_config.step = get_real_step_or_rank(step, Const.STEP)
+
+    @staticmethod
+    def get_task_config(task, json_config):
+        raise NotImplementedError("Subclass must implment get_task_config")
+
+    @staticmethod
+    def check_input_params(config_path, task, dump_path, level):
+        if not config_path:
+            config_path = os.path.join(os.path.dirname(__file__), "../../config.json")
+        if config_path is not None:
+            if not isinstance(config_path, str):
+                raise MsprobeException(
+                    MsprobeException.INVALID_PARAM_ERROR, f"config_path must be a string")
+            file_checker = FileChecker(
+                file_path=config_path, path_type=FileCheckConst.FILE, file_type=FileCheckConst.JSON_SUFFIX)
+            file_checker.common_check()
+
+        if task is not None and task not in Const.TASK_LIST:
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"task must be one of {Const.TASK_LIST}")
+
+        if dump_path is not None:
+            if not isinstance(dump_path, str):
+                raise MsprobeException(
+                    MsprobeException.INVALID_PARAM_ERROR, f"dump_path must be a string")
+
+        if level is not None and level not in Const.LEVEL_LIST:
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"level must be one of {Const.LEVEL_LIST}")
+
+    @classmethod
+    def get_instance(cls):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if instance.task in BasePrecisionDebugger.tasks_not_need_debugger:
+            instance = None
+        return instance
+
+    @classmethod
+    def forward_backward_dump_end(cls):
+        instance = cls._instance
+        instance.stop()
+
+    @classmethod
+    def set_init_step(cls, step):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        check_init_step(step)
+        instance.service.init_step = step
+        instance.service.loop = 0
+
+    @classmethod
+    def register_custom_api(cls, module, api, api_prefix=None):
+        if not api_prefix:
+            api_prefix = getattr(module, "__name__", "Custom")
+        if not isinstance(api_prefix, str):
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, "api_prefix must be string")
+        if not hasattr(module, api):
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"module {str(module)} does not have {api}")
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        instance.service.register_custom_api(module, api, api_prefix)
+
+    @classmethod
+    def restore_custom_api(cls, module, api):
+        if not hasattr(module, api):
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR, f"module {str(module)} does not have {api}")
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        instance.service.restore_custom_api(module, api)
+
+    def parse_config_path(self, json_file_path, task):
+        if not json_file_path:
+            json_file_path = os.path.join(os.path.dirname(__file__), "../../config.json")
+        json_config = load_json(json_file_path)
+        common_config = CommonConfig(json_config)
+        if task:
+            task_config = self.get_task_config(task, json_config)
+        else:
+            if not common_config.task:
+                common_config.task = Const.STATISTICS
+            task_config = self.get_task_config(common_config.task, json_config)
+        return common_config, task_config

msprobe/core/grad_probe/constant.py

@@ -31,6 +31,7 @@ class GradConst:
     STEP = "step"
     BOUNDS = "bounds"
     OUTPUT_PATH = "output_path"
+    TIME_STAMP = "time_stamp"
 
     # level const
     LEVEL = "level"
@@ -51,7 +52,7 @@ class GradConst:
     BOUNDS_MINIMUM = -2**63
     BOUNDS_MAXIMUM = 2**63 - 1
 
-    # file safty
+    # file safety
     DATA_DIR_AUTHORITY = 0o750
     DATA_FILE_AUTHORITY = 0o640
     DIRECTORY_LENGTH = 4096

msprobe/core/grad_probe/grad_compare.py

@@ -112,7 +112,7 @@ class GradComparator:
             result.append([key] + value)
         result_csv_path = os.path.join(output_dir, "similarities.csv")
         if os.path.exists(result_csv_path):
-            logger.warning(f"{result_csv_path} will be recoverd")
+            logger.warning(f"{result_csv_path} will be deleted")
             remove_path(result_csv_path)
         write_csv(result, result_csv_path)
 
@@ -121,7 +121,7 @@ class GradComparator:
         similarities = {}
         logger.info(f"{len(steps)} steps will be compared")
         grad_weight_order = cls._get_grad_weight_order(path1, path2)
-        for step in tqdm(steps, desc="culculate similarities (by step)"):
+        for step in tqdm(steps, desc="calculate similarities (by step)"):
             grad_files = cls._get_matched_grad_files(path1, path2, step)
             same_count_summary = 0
             total_count_summary = 0

msprobe/core/grad_probe/utils.py

@@ -82,7 +82,7 @@ class ListCache(list):
         if len(self) == 0:
             return
         if not self._output_file:
-            logger.warning("dumpfile path is not setted")
+            logger.warning("dumpfile path is not set.")
         write_csv(self, self._output_file)
         logger.info(f"write {len(self)} items to {self._output_file}.")
         self.clear()

msprobe/core/hook_manager.py

@@ -0,0 +1,242 @@
+# Copyright (c) 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.
+
+
+from abc import ABC, abstractmethod
+import os
+
+from msprobe.core.common.runtime import Runtime
+from msprobe.core.common.utils import Const
+from msprobe.core.data_dump.data_processor.base import (ModuleBackwardInputsOutputs, ModuleForwardInputsOutputs)
+
+
+class HookSet:
+    def __init__(self, forward_hook=None, forward_pre_hook=None, backward_hook=None, backward_pre_hook=None):
+        self.forward_hook = forward_hook
+        self.forward_pre_hook = forward_pre_hook
+        self.backward_hook = backward_hook
+        self.backward_pre_hook = backward_pre_hook
+
+
+class BaseHookManager(ABC):
+    inner_switch = False
+    hook_handle_dict = {}
+    params_grad_info = {}
+
+    def __init__(self, data_collector, config, attl_manager=None):
+        self.data_collector = data_collector
+        self.config = config
+        self.attl_manager = attl_manager
+
+    @property
+    def _pid(self):
+        return os.getpid()
+
+    @property
+    @abstractmethod
+    def _is_recompute(self):
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _no_grad_context():
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _add_count(name):
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _process_kwargs_and_output(module, hook_type, kwargs_or_output, output_or_kwargs):
+        pass
+
+    @staticmethod
+    def _clear_input_kwargs(module):
+        if hasattr(module, 'msprobe_input_kwargs'):
+            del module.msprobe_input_kwargs
+
+    @abstractmethod
+    def build_hook(self):
+        pass
+
+    @abstractmethod
+    def _get_params_dict(self, module):
+        pass
+
+    @abstractmethod
+    def _need_exchange(self, module):
+        pass
+
+    def _register_param_hook(self, name, module, params_dict):
+        ori_name = name.rsplit(Const.SEP, 2)[0]
+        grad_name = ori_name + Const.SEP + Const.PARAMS_GRAD
+        # 首次执行前向hook时，添加params_grad_name属性，并注册参数hook
+        setattr(module, 'params_grad_name', grad_name)
+         # data_mode为forward时，不注册参数hook
+        if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
+            for param_name, param in params_dict.items():
+                if param.requires_grad:
+                    name = ori_name + Const.SEP + param_name
+                    old_handle = BaseHookManager.hook_handle_dict.get(name)
+                    if old_handle and hasattr(old_handle, "remove"):
+                        old_handle.remove()
+                    handle = param.register_hook(self._build_grad_hook(module, ori_name, param_name))
+                    BaseHookManager.hook_handle_dict[name] = handle
+
+    def _init_params_grad_info(self, module, params_dict):
+        '''
+        初始化参数梯度信息, 在前向hook结束后, 将参数梯度信息写入cache_data中用于占位
+        '''
+        if not params_dict:
+            return
+        if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
+            grad_name = module.params_grad_name if hasattr(module, 'params_grad_name') else None
+            # 判断是否已经在cache_data中进行了占位, 若没有则先写入cache_data中
+            if not BaseHookManager.params_grad_info.get(grad_name):
+                data_info = {grad_name: {key: [None] for key, value in params_dict.items() if value.requires_grad}}
+                # 当模块中的参数有requires_grad属性为True时，才会进行梯度计算，此时才需要占位
+                if data_info.get(grad_name):
+                    # 将grad_name的data_info先写入cache_data中, 梯度计算后再更新
+                    self.data_collector.handle_data(grad_name, data_info,
+                                                    flush=self.data_collector.data_processor.is_terminated)
+                # 记录当前模块的参数梯度信息已占位
+                BaseHookManager.params_grad_info[grad_name] = True
+
+    def _should_execute_hook(self, hook_type, module, is_forward):
+        is_module_hook = hook_type == Const.MODULE
+        if is_module_hook and not Runtime.is_running:
+            return False
+        elif not is_module_hook and is_forward and not Runtime.is_running:
+            return False
+        elif not is_module_hook and not is_forward and not module.forward_data_collected:
+            return False
+        if BaseHookManager.inner_switch:
+            return False
+        if not self.data_collector or self.data_collector.data_processor.is_terminated:
+            return False
+        return True
+
+    def _build_grad_hook(self, module, ori_name, param_name):
+        def hook_fn(grad):
+            if not self._should_execute_hook(Const.MODULE, module, False):
+                return
+            BaseHookManager.inner_switch = True
+            self.data_collector.params_data_collect(ori_name, param_name, self._pid, grad)
+            BaseHookManager.inner_switch = False
+            return
+        return hook_fn
+
+    def _build_forward_pre_hook(self, hook_type, full_name, api_name):
+        def forward_pre_hook(module, args, kwargs=None):
+            if hook_type == Const.MODULE:
+                return
+            if not self._should_execute_hook(hook_type, module, True):
+                return
+            if kwargs is None:
+                kwargs = module.msprobe_input_kwargs if hasattr(module, 'msprobe_input_kwargs') else {}
+            with self._no_grad_context():
+                BaseHookManager.inner_switch = False
+                module.forward_data_collected = True
+                self._add_count(api_name)
+                module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
+                self.data_collector.update_api_or_module_name(full_name)
+                if getattr(self.config, "online_run_ut", False):
+                    BaseHookManager.inner_switch = False
+                    return
+                self.data_collector.forward_input_data_collect(
+                    full_name,
+                    module,
+                    self._pid,
+                    module_input_output,
+                    self._is_recompute
+                )
+                BaseHookManager.inner_switch = False
+        return forward_pre_hook
+
+    def _build_forward_hook(self, hook_type, full_name):
+        def forward_hook(module, args, kwargs_or_output, output_or_kwargs=None):
+            if not self._should_execute_hook(hook_type, module, True):
+                self._clear_input_kwargs(module)
+                return None
+            kwargs, output = self._process_kwargs_and_output(module, hook_type, kwargs_or_output, output_or_kwargs)
+            BaseHookManager.inner_switch = True
+            self.data_collector.update_api_or_module_name(full_name)
+            module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+            with self._no_grad_context():
+                if getattr(self.config, "online_run_ut", False):
+                    if self.data_collector.scope and not self.data_collector.scope.check(full_name):
+                        return None
+                    if self.attl_manager:
+                        self.attl_manager.attl_send(full_name, args, kwargs, output)
+                    BaseHookManager.inner_switch = False
+                    return None
+                if hook_type == Const.MODULE:
+                    params_dict = self._get_params_dict(module)
+                    setattr(module_input_output, Const.PARAMS, params_dict)
+                    if params_dict:
+                        self._register_param_hook(full_name, module, params_dict)
+                    self.data_collector.update_api_or_module_name(full_name)
+                    self.data_collector.forward_data_collect(
+                        full_name,
+                        module,
+                        self._pid,
+                        module_input_output,
+                        self._is_recompute
+                    )
+                    self._init_params_grad_info(module, params_dict)
+                else:
+                    self.data_collector.forward_output_data_collect(
+                        full_name,
+                        module,
+                        self._pid,
+                        module_input_output,
+                        self._is_recompute
+                    )
+                self._clear_input_kwargs(module)
+
+                if self.data_collector.if_return_forward_new_output():
+                    forward_new_output = self.data_collector.get_forward_new_output()
+                    BaseHookManager.inner_switch = False
+                    return forward_new_output
+
+                BaseHookManager.inner_switch = False
+                return output
+        return forward_hook
+
+    def _build_backward_hook(self, hook_type, full_name):
+        def backward_hook(module, grad_input, grad_output):
+            if not self._should_execute_hook(hook_type, module, False):
+                return
+            BaseHookManager.inner_switch = True
+            self.data_collector.update_api_or_module_name(full_name)
+            if getattr(self.config, "online_run_ut", False):
+                BaseHookManager.inner_switch = False
+                return
+            need_exchange = self._need_exchange(module) if hook_type == Const.MODULE else True
+            if need_exchange:
+                module_input_output = ModuleBackwardInputsOutputs(grad_input=grad_output, grad_output=grad_input)
+            else:
+                module_input_output = ModuleBackwardInputsOutputs(grad_input=grad_input, grad_output=grad_output)
+            self.data_collector.backward_data_collect(
+                    full_name,
+                    module,
+                    self._pid,
+                    module_input_output,
+                    self._is_recompute
+                )
+            BaseHookManager.inner_switch = False
+        return backward_hook