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

msprobe/pytorch/debugger/precision_debugger.py

@@ -13,36 +13,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import namedtuple
+from torch.utils.data import dataloader
 
-import torch
-from msprobe.core.common.const import Const, FileCheckConst, MsgConst
+from msprobe.core.common.const import Const, MsgConst
 from msprobe.core.common.exceptions import MsprobeException
-from msprobe.core.common.file_utils import FileChecker
-from msprobe.core.common.utils import get_real_step_or_rank
+from msprobe.core.common.utils import check_token_range
+from msprobe.core.debugger.precision_debugger import BasePrecisionDebugger
 from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.common.utils import check_save_param
+from msprobe.pytorch.common.utils import check_save_param, is_torch_nn_module
 from msprobe.pytorch.debugger.debugger_config import DebuggerConfig
 from msprobe.pytorch.dump.module_dump.module_dump import ModuleDumper
 from msprobe.pytorch.grad_probe.grad_monitor import GradientMonitor
-from msprobe.pytorch.pt_config import parse_json_config
-from msprobe.pytorch.service import Service
-from torch.utils.data import dataloader
-
-ConfigParameters = namedtuple("ConfigParameters", ["config_path", "task",
-                                                   "dump_path", "level", "model"])
-
+from msprobe.pytorch.pytorch_service import PytorchService
+from msprobe.pytorch.pt_config import parse_task_config
 
-class PrecisionDebugger:
-    _instance = None
-    tasks_not_need_debugger = [Const.GRAD_PROBE]
 
-    def __new__(cls, *args, **kwargs):
-        if cls._instance is None:
-            cls._instance = super(PrecisionDebugger, cls).__new__(cls)
-            cls._instance.config = None
-            cls._instance.enable_dataloader = False
-        return cls._instance
+class PrecisionDebugger(BasePrecisionDebugger):
 
     def __init__(
         self,
@@ -53,90 +39,48 @@ class PrecisionDebugger:
         model=None,
         step=None
     ):
-        if not hasattr(self, "initialized"):
-            config_params = ConfigParameters(config_path,
-                                             task,
-                                             dump_path,
-                                             level,
-                                             model)
-            self.check_input_params(config_params)
-
-            self.initialized = True
-            self.model = model
-            common_config, task_config = parse_json_config(config_path, task)
-            self.task = task if task else common_config.task
-            if self.task == Const.GRAD_PROBE:
-                self.gm = GradientMonitor(common_config, task_config)
-                return
-            if step is not None:
-                common_config.step = get_real_step_or_rank(step, Const.STEP)
-            self.config = DebuggerConfig(
-                common_config, task_config, task, dump_path, level
-            )
-            self.service = Service(self.config)
-            self.module_dumper = ModuleDumper(self.service)
-            self.enable_dataloader = self.config.enable_dataloader
-            if self.enable_dataloader:
-                logger.warning_on_rank_0("The enable_dataloader feature will be deprecated in the future.")
-                dataloader._BaseDataLoaderIter.__next__ = iter_tracer(dataloader._BaseDataLoaderIter.__next__)
+        if self.initialized:
+            return
+        super().__init__(config_path, task, dump_path, level, step)
+        self.model = model
+        if self.task == Const.GRAD_PROBE:
+            self.gm = GradientMonitor(self.common_config, self.task_config)
+            return
+        self.config = DebuggerConfig(
+            self.common_config, self.task_config, task, dump_path, level
+        )
+        self.service = PytorchService(self.config)
+        self.module_dumper = ModuleDumper(self.service)
+        self.ori_customer_func = {}
+        self.enable_dataloader = self.config.enable_dataloader
+        self.param_warning()
 
     @property
     def instance(self):
         return self._instance
 
     @staticmethod
-    def check_input_params(args):
-        if args.config_path is not None:
-            if not isinstance(args.config_path, str):
-                raise MsprobeException(
-                    MsprobeException.INVALID_PARAM_ERROR, f"config_path must be a string")
-            file_checker = FileChecker(
-                file_path=args.config_path, path_type=FileCheckConst.FILE, file_type=FileCheckConst.JSON_SUFFIX)
-            file_checker.common_check()
-
-        if args.task is not None and args.task not in Const.TASK_LIST:
-            raise MsprobeException(
-                MsprobeException.INVALID_PARAM_ERROR, f"task must be one of {Const.TASK_LIST}")
-
-        if args.dump_path is not None:
-            if not isinstance(args.dump_path, str):
-                raise MsprobeException(
-                    MsprobeException.INVALID_PARAM_ERROR, f"dump_path must be a string")
-
-        if args.level is not None and args.level not in Const.LEVEL_LIST:
-            raise MsprobeException(
-                MsprobeException.INVALID_PARAM_ERROR, f"level must be one of {Const.LEVEL_LIST}")
-
-        if args.model is not None:
-            logger.warning_on_rank_0(
-                "The 'model' parameter in the PrecisionDebugger will be deprecated in the future."
-                "It is recommended to pass the 'model' parameter in the start interface instead."
-            )
+    def get_task_config(task, json_config):
+        return parse_task_config(task, json_config)
 
     @classmethod
-    def start(cls, model=None):
-        instance = cls._instance
-        if not instance:
-            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
-        if instance.task in PrecisionDebugger.tasks_not_need_debugger:
+    def start(cls, model=None, token_range=None):
+        instance = cls.get_instance()
+        if instance is None:
             return
-        instance.config.check_model(instance, model)
+
+        check_token_range(token_range)
+        instance.config.check_model(instance, model, token_range)
+
         if instance.enable_dataloader:
             logger.warning_on_rank_0("DataLoader is enabled, start() skipped.")
         else:
-            instance.service.start(instance.model)
-
-    @classmethod
-    def forward_backward_dump_end(cls):
-        instance = cls._instance
-        instance.stop()
+            instance.service.start(instance.model, token_range)
 
     @classmethod
     def stop(cls):
-        instance = cls._instance
-        if not instance:
-            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
-        if instance.task in PrecisionDebugger.tasks_not_need_debugger:
+        instance = cls.get_instance()
+        if instance is None:
             return
         if instance.enable_dataloader:
             logger.warning_on_rank_0("DataLoader is enabled, stop() skipped.")
@@ -145,9 +89,8 @@ class PrecisionDebugger:
 
     @classmethod
     def step(cls):
-        if not cls._instance:
-            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
-        if cls._instance.task in PrecisionDebugger.tasks_not_need_debugger:
+        instance = cls.get_instance()
+        if instance is None:
             return
         cls._instance.service.step()
 
@@ -172,12 +115,23 @@ class PrecisionDebugger:
             return
         instance.service.save(variable, name, save_backward)
 
+    def param_warning(self):
+        if self.model is not None:
+            logger.warning_on_rank_0(
+                "The 'model' parameter in the PrecisionDebugger will be deprecated in the future."
+                "It is recommended to pass the 'model' parameter in the start interface instead."
+            )
+        if self.enable_dataloader:
+            logger.warning_on_rank_0("The enable_dataloader feature will be deprecated in the future.")
+            dataloader._BaseDataLoaderIter.__next__ = iter_tracer(dataloader._BaseDataLoaderIter.__next__)
+
 
 def module_dump(module, dump_name):
-    if not isinstance(module, torch.nn.Module):
+    if not is_torch_nn_module(module):
         raise MsprobeException(
             MsprobeException.INVALID_PARAM_ERROR,
-            f"the module argument in module_dump must be a torch.nn.Module subclass"
+            f"the module argument in module_dump must be a torch.nn.Module type, "
+            f"but currently there is an unsupported {type(module)} type."
         )
     if not isinstance(dump_name, str):
         raise MsprobeException(

msprobe/pytorch/dump/module_dump/hook_wrapper.py

@@ -0,0 +1,93 @@
+# 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.
+
+from functools import wraps
+
+import torch
+from torch.utils.hooks import BackwardHook
+
+from msprobe.core.common.const import Const
+from msprobe.core.common.decorator import recursion_depth_decorator
+from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.common.utils import is_float8_tensor
+
+
+def wrap_setup_backward_hook(func):
+    def requires_clone(tensor):
+        return isinstance(tensor, torch.Tensor) and not is_float8_tensor(tensor) and \
+            tensor.requires_grad and torch.is_grad_enabled()
+
+    @recursion_depth_decorator("Dump: wrap_setup_backward_hook.parse_tensor", max_depth=Const.DUMP_MAX_DEPTH)
+    def parse_tensor(item, tensor_list):
+        if requires_clone(item):
+            tensor_list.append(item)
+        elif isinstance(item, (list, tuple)):
+            for value in item:
+                parse_tensor(value, tensor_list)
+        elif isinstance(item, dict):
+            for value in item.values():
+                parse_tensor(value, tensor_list)
+
+    @recursion_depth_decorator("Dump: wrap_setup_backward_hook.rebuild_args", max_depth=Const.DUMP_MAX_DEPTH)
+    def rebuild_args(item, tensor_iter):
+        if requires_clone(item):
+            result = next(tensor_iter)
+            if hasattr(result, "_base") and result._base is not None:
+                if torch._C._autograd._get_creation_meta(result) != torch._C._autograd.CreationMeta(0):
+                    torch._C._autograd._set_creation_meta(result, torch._C._autograd.CreationMeta(0))
+            return result    
+        if isinstance(item, list):
+            for index, value in enumerate(item):
+                item[index] = rebuild_args(value, tensor_iter)
+            return item
+        if isinstance(item, dict):
+            for key, value in item.items():
+                item[key] = rebuild_args(value, tensor_iter)
+            return item
+        if isinstance(item, tuple):
+            if hasattr(item, '_fields'):
+                return type(item)(*[rebuild_args(i, tensor_iter) for i in item])
+            return type(item)([rebuild_args(i, tensor_iter) for i in item])
+        return item
+
+    @wraps(func)
+    def wrap_setup_hook_func(*args, **kwargs):
+        if len(args) < 2:
+            return func(*args, **kwargs)
+
+        actual_args = args[1]
+
+        tensor_list = []
+
+        parse_tensor(actual_args, tensor_list)
+
+        new_args = args[0], tuple(tensor_list)
+        hooked_tensors = func(*new_args, **kwargs)
+
+        tensor_iter = iter(hooked_tensors)
+        try:
+            new_data = rebuild_args(actual_args, tensor_iter)
+        except Exception as e:
+            logger.debug(f"Unsupported data in setup input/output hook. The detail info: {e}")
+            new_data = actual_args
+
+        return new_data
+
+    return wrap_setup_hook_func
+
+
+def wrap_setup_input_output_hook():
+    BackwardHook.setup_input_hook = wrap_setup_backward_hook(BackwardHook.setup_input_hook)
+    BackwardHook.setup_output_hook = wrap_setup_backward_hook(BackwardHook.setup_output_hook)

msprobe/pytorch/dump/module_dump/module_dump.py

@@ -13,74 +13,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
-from msprobe.core.common.const import Const
-from msprobe.core.data_dump.scope import BaseScope
 from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.hook_module.api_registry import api_register
-
-torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
+from msprobe.pytorch.dump.module_dump.module_processer import ModuleProcesser
+from msprobe.pytorch.hook_module.api_register import get_api_register
 
 
 class ModuleDumper:
     def __init__(self, service):
         self.service = service
-        self.hook_handle_list = []
+        self.api_register = get_api_register()
 
     def start_module_dump(self, module, dump_name):
-        api_register.api_originality()
-        self.register_hook(module, dump_name)
-
-    def stop_module_dump(self):
-        api_register.api_modularity()
-        for hook_handle in self.hook_handle_list:
-            if isinstance(hook_handle, torch.utils.hooks.RemovableHandle):
-                hook_handle.remove()
-        self.hook_handle_list.clear()
+        if hasattr(module, 'msprobe_hook') and not hasattr(module, 'msprobe_module_dump'):
+            logger.info_on_rank_0("The init dump is enabled, and the module dump function will not be available.")
+            return
 
-    def register_hook(self, module, dump_name):
-        prefix_name = (
-                BaseScope.Module_Type_Module + Const.SEP +
-                dump_name + Const.SEP +
-                module.__class__.__name__ + Const.SEP
-        )
-        module_processor = self.service.module_processor
-        _, forward_hook, backward_hook, forward_hook_torch_version_below_2 = self.service.build_hook(
-            BaseScope.Module_Type_Module,
-            prefix_name
-        )
+        ModuleProcesser.enable_module_dump = True
+        self.api_register.restore_all_api()
+        if not hasattr(module, 'msprobe_module_dump'):
+            self.service.module_processor.register_module_hook(module, self.service.build_hook,
+                                                               recursive=False, module_names=[dump_name])
+            setattr(module, 'msprobe_module_dump', True)
 
-        if module_processor.has_register_backward_hook(module):
-            logger.warning(
-                f"The {dump_name} module has registered deprecated register_backward_hook,"
-                f"which may cause abnormal data dump. The backward data dump for this module will be skipped."
-            )
-        if torch_version_above_or_equal_2:
-            forward_hook_handle = module.register_forward_hook(forward_hook, with_kwargs=True)
-        else:
-            if not module_processor.has_register_backward_hook(module):
-                backward_hook_handle = module.register_full_backward_hook(
-                    module_processor.node_hook(prefix_name + Const.BACKWARD, Const.STOP)
-                )
-                self.hook_handle_list.append(backward_hook_handle)
-            forward_hook_handle = module.register_forward_hook(forward_hook_torch_version_below_2)
-        self.hook_handle_list.append(forward_hook_handle)
-        if not module_processor.has_register_backward_hook(module):
-            backward_hook_handle = module.register_full_backward_hook(backward_hook)
-            self.hook_handle_list.append(backward_hook_handle)
-
-        forward_pre_hook_handle = module.register_forward_pre_hook(
-            module_processor.node_hook(prefix_name + Const.FORWARD, Const.START)
-        )
-        forward_hook_handle = module.register_forward_hook(
-            module_processor.node_hook(prefix_name + Const.FORWARD, Const.STOP)
-        )
-        self.hook_handle_list.extend([forward_pre_hook_handle, forward_hook_handle])
-        if torch_version_above_or_equal_2 and not module_processor.has_register_backward_hook(module):
-            backward_pre_hook_handle = module.register_full_backward_pre_hook(
-                module_processor.node_hook(prefix_name + Const.BACKWARD, Const.START)
-            )
-            backward_hook_handle = module.register_full_backward_hook(
-                module_processor.node_hook(prefix_name + Const.BACKWARD, Const.STOP)
-            )
-            self.hook_handle_list.extend([backward_pre_hook_handle, backward_hook_handle])
+    def stop_module_dump(self):
+        ModuleProcesser.enable_module_dump = False
+        self.api_register.register_all_api()