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

msprobe/pytorch/common/utils.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -28,7 +28,7 @@ from msprobe.core.common.exceptions import DistributedNotInitializedError
 from msprobe.core.common.file_utils import (FileCheckConst, change_mode,
                                             check_file_or_directory_path, check_path_before_create, FileOpen)
 from msprobe.core.common.log import logger
-from msprobe.core.common.utils import check_seed_all
+from msprobe.core.common.utils import check_seed_all, is_save_variable_valid
 from packaging import version
 
 try:
@@ -57,7 +57,7 @@ def parameter_adapter(func):
 
     @wraps(func)
     def inner(self, *args, **kwargs):
-        if self.op_name_ == "__getitem__" and len(args) > 1 and isinstance(args[1], torch.Tensor):
+        if self.api_name == "__getitem__" and len(args) > 1 and isinstance(args[1], torch.Tensor):
             input_tensor = args[0]
             indices = args[1]
             if indices.dtype == torch.uint8:
@@ -77,7 +77,7 @@ def parameter_adapter(func):
                 else:
                     res = [input_tensor[tensor_index] for tensor_index in indices]
                     return getattr(torch._C._VariableFunctionsClass, "stack")(res, 0)
-        if self.op_name_ == "__eq__" and len(args) > 1 and args[1] is None:
+        if self.api_name == "__eq__" and len(args) > 1 and args[1] is None:
             return False
         return func(self, *args, **kwargs)
 
@@ -261,6 +261,10 @@ class Const:
     NPU = 'NPU'
     DISTRIBUTED = 'Distributed'
 
+    HIFLOAT8_TYPE = "torch_npu.HiFloat8Tensor"
+    FLOAT8_E5M2_TYPE = "torch.float8_e5m2"
+    FLOAT8_E4M3FN_TYPE = "torch.float8_e4m3fn"
+
     RAISE_PRECISION = {
         torch.float16: torch.float32,
         torch.bfloat16: torch.float32,
@@ -419,7 +423,11 @@ def is_recomputation():
         bool: True if in the re-computation phase, False otherwise.
     """
     backward_function_indices = []
-    call_stack = inspect.stack()
+    try:
+        call_stack = inspect.stack()
+    except Exception as e:
+        logger.warning(f"Failed to capture stack trace, recomputation validation may be incorrect, error info: {e}.")
+        return False
 
     # Identify the function 'backward' is being executed within the 'torch/_tensor.py' file.
     for frame_info in call_stack:
@@ -449,9 +457,11 @@ def is_recomputation():
 
 def check_save_param(variable, name, save_backward):
     # try catch this api to skip invalid call
-    if not isinstance(variable, (list, dict, torch.Tensor, int, float, str)):
+    valid_data_types = tuple([torch.Tensor, int, float, str])
+    if not is_save_variable_valid(variable, valid_data_types):
+        valid_data_types_with_nested_types = valid_data_types + (dict, tuple, list)
         logger.warning("PrecisionDebugger.save variable type not valid, "
-                       "should be one of list, dict, torch.Tensor, int, float or string. "
+                       f"should be one of {valid_data_types_with_nested_types}"
                        "Skip current save process.")
         raise ValueError
     if not isinstance(name, str):
@@ -473,3 +483,15 @@ def replace_last_occurrence(text, old, new):
     if index != -1:
         return text[:index] + text[index:].replace(old, new, 1)
     return text
+
+
+def is_hifloat8_tensor(tensor):
+    if not is_gpu and hasattr(torch_npu, "HiFloat8Tensor") and isinstance(tensor, torch_npu.HiFloat8Tensor):
+        return True
+    return False
+
+
+def is_float8_tensor(tensor):
+    if str(tensor.dtype) in [Const.FLOAT8_E5M2_TYPE, Const.FLOAT8_E4M3FN_TYPE]:
+        return True
+    return is_hifloat8_tensor(tensor)

msprobe/pytorch/debugger/precision_debugger.py

@@ -19,7 +19,7 @@ import torch
 from msprobe.core.common.const import Const, FileCheckConst, 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 get_real_step_or_rank, check_init_step
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.common.utils import check_save_param
 from msprobe.pytorch.debugger.debugger_config import DebuggerConfig
@@ -172,6 +172,15 @@ class PrecisionDebugger:
             return
         instance.service.save(variable, name, save_backward)
 
+    @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
+
 
 def module_dump(module, dump_name):
     if not isinstance(module, torch.nn.Module):

msprobe/pytorch/dump/module_dump/module_dump.py

@@ -17,7 +17,7 @@ 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
+from msprobe.pytorch.hook_module.api_register import get_api_register
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
@@ -26,13 +26,14 @@ 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.api_register.restore_all_api()
         self.register_hook(module, dump_name)
 
     def stop_module_dump(self):
-        api_register.api_modularity()
+        self.api_register.register_all_api()
         for hook_handle in self.hook_handle_list:
             if isinstance(hook_handle, torch.utils.hooks.RemovableHandle):
                 hook_handle.remove()

msprobe/pytorch/dump/module_dump/module_processer.py

@@ -16,15 +16,17 @@
 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.core.data_dump.scope import BaseScope, ModuleRangeScope, MixRangeScope
 from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.common.utils import replace_last_occurrence
-from torch.utils.checkpoint import checkpoint as origin_checkpoint
-from torch.utils.checkpoint import set_checkpoint_early_stop
-from torch.utils.hooks import BackwardHook
+from msprobe.pytorch.common.utils import replace_last_occurrence, is_float8_tensor
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
+if torch_version_above_or_equal_2:
+    from torch.utils.checkpoint import checkpoint as origin_checkpoint, set_checkpoint_early_stop
 
 
 def checkpoint_without_early_stop(*args, **kwargs):
@@ -33,7 +35,8 @@ def checkpoint_without_early_stop(*args, **kwargs):
 
 
 def replace_checkpoint():
-    torch.utils.checkpoint.checkpoint = checkpoint_without_early_stop
+    if torch_version_above_or_equal_2:
+        torch.utils.checkpoint.checkpoint = checkpoint_without_early_stop
 
 
 class ModuleProcesser:
@@ -58,8 +61,9 @@ class ModuleProcesser:
         return clone_return_value_func
 
     @staticmethod
+    @recursion_depth_decorator("ModuleDump: ModuleProcesser.clone_if_tensor", max_depth=Const.DUMP_MAX_DEPTH)
     def clone_if_tensor(result):
-        if isinstance(result, torch.Tensor):
+        if isinstance(result, torch.Tensor) and not is_float8_tensor(result):
             return result.clone()
         elif type(result) is tuple:
             return tuple(ModuleProcesser.clone_if_tensor(x) for x in result)
@@ -109,6 +113,8 @@ class ModuleProcesser:
             for name, module in modules_and_names:
                 if module == model:
                     continue
+                if module.__class__.__name__ == "FullyShardedDataParallel":
+                    continue
                 module_index = (index + Const.SEP) if index != "-1" else ""
                 prefix_name = (BaseScope.Module_Type_Module + Const.SEP + module_index +
                                name + Const.SEP + module.__class__.__name__ + Const.SEP)

msprobe/pytorch/free_benchmark/common/utils.py

@@ -16,7 +16,7 @@
 
 import torch
 from msprobe.core.common.exceptions import FreeBenchmarkException
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark.common.enums import DeviceType
 
 

msprobe/pytorch/free_benchmark/compare/single_benchmark.py

@@ -16,7 +16,7 @@
 import math
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
 from msprobe.pytorch.free_benchmark.common.utils import TorchC

msprobe/pytorch/free_benchmark/perturbed_layers/npu/add_noise.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode
@@ -95,13 +95,13 @@ class AddNoiseLayer(NpuBaseLayer):
         except Exception:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"when calculate maximun value, tensor is changed to float32."
+                f"when calculating the maximum value, the tensor is changed to float32."
             )
             max_val = TorchC.max(TorchC.abs(tensor_obj.to(torch.float32))).item()
         if max_val < abs_tol:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"Maximun value is less than the minimun threshold. Cancel add noise."
+                f"maximum value is less than the minimum threshold. Cancel adding noise."
             )
             return False
         return True

msprobe/pytorch/free_benchmark/perturbed_layers/npu/bit_noise.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode
@@ -100,13 +100,13 @@ class BitNoiseLayer(NpuBaseLayer):
         except Exception:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"when calculate maximun value, tensor is changed to float32."
+                f"when calculate the maximum value, the tensor is changed to float32."
             )
             max_val = TorchC.max(TorchC.abs(tensor_obj.to(torch.float32))).item()
         if max_val < abs_tol:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.api_name}, "
-                f"Maximun value is less than the minimun threshold. Cancel add noise."
+                f"maximum value is less than the minimum threshold. Cancel adding noise."
             )
             return False
         return True

msprobe/pytorch/free_benchmark/perturbed_layers/npu/change_value.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import torch
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode
 from msprobe.pytorch.free_benchmark.common.params import DataParams

msprobe/pytorch/free_benchmark/perturbed_layers/npu/improve_precision.py

@@ -15,7 +15,7 @@
 
 import torch
 from msprobe.core.common.const import Const
-from msprobe.core.common.utils import recursion_depth_decorator
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.free_benchmark import logger
 from msprobe.pytorch.free_benchmark.common.constant import CommonField
 from msprobe.pytorch.free_benchmark.common.enums import PerturbationMode

msprobe/pytorch/free_benchmark/result_handlers/check_handler.py

@@ -49,6 +49,6 @@ class CheckerHandler(FuzzHandler):
         except Exception as e:
             logger.warning_on_rank_0(
                 f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when campare the result exception raise {e}"
+                f"when comparing the results, an exception is raised: {e}"
             )
         return data_params.original_result

msprobe/pytorch/function_factory.py

@@ -70,7 +70,7 @@ class Register(dict):
 
         def add_register_item(key, value):
             if key in self._dict:
-                logger.warning(f"{value.__name__} has been registered before, so we will overriden it.")
+                logger.warning(f"{value.__name__} has been registered before, so we will override it.")
             self[key] = value
             return value
 

msprobe/pytorch/grad_probe/grad_monitor.py

@@ -46,7 +46,7 @@ class GradientMonitor:
         if not os.path.exists(self._output_path):
             create_directory(self._output_path)
         else:
-            logger.warning(f"the file in {self._output_path} will be recoverd")
+            logger.warning(f"the file in {self._output_path} will be deleted")
         self._step = -1
         self._param2name = defaultdict(str)
 
@@ -97,7 +97,7 @@ class GradientMonitor:
                 create_directory(output_dirpath)
             output_path = os.path.join(output_dirpath, f"grad_summary_{self._step}.csv")
             if os.path.exists(output_path):
-                logger.warning(f"{output_path} will be recoverd")
+                logger.warning(f"{output_path} will be deleted")
                 remove_path(output_path)
             header_result = GradStatCsv.generate_csv_header(self._level_adp, self._bounds)
             output_lines.insert(0, header_result)

msprobe/pytorch/hook_module/api_register.py

@@ -0,0 +1,131 @@
+# 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 functools
+import os
+
+import torch
+import torch.distributed as dist
+
+from msprobe.core.common.const import Const
+from msprobe.core.data_dump.api_registry import ApiRegistry
+from msprobe.pytorch.common.utils import (
+    torch_without_guard_version, is_gpu, torch_device_guard, parameter_adapter
+)
+from msprobe.pytorch.function_factory import npu_custom_functions
+from msprobe.pytorch.hook_module.hook_module import HOOKModule
+
+
+torch_version_above_2 = torch.__version__.split('+')[0] > '2.0'
+
+_api_types = {
+    Const.PT_FRAMEWORK: {
+        Const.PT_API_TYPE_FUNCTIONAL: (torch.nn.functional, (torch.nn.functional,)),
+        Const.PT_API_TYPE_TENSOR: (torch.Tensor, (torch.Tensor,)),
+        Const.PT_API_TYPE_TORCH: (torch, (torch,)),
+        Const.PT_API_TYPE_VF: (torch._C._VariableFunctionsClass, (torch._VF,)),
+        Const.PT_API_TYPE_DIST: (dist, (dist, dist.distributed_c10d))
+    }
+}
+if not is_gpu:
+    import torch_npu
+    if torch_without_guard_version:
+        _api_types.get(Const.PT_FRAMEWORK).update(
+            {
+                Const.PT_API_TYPE_NPU: (torch.ops.npu, (torch_npu, torch.ops.npu))
+            }
+        )
+    else:
+        _api_types.get(Const.PT_FRAMEWORK).update(
+            {Const.PT_API_TYPE_NPU: (torch_npu._C._VariableFunctionsClass, (torch_npu,))}
+        )
+        _api_types.get(Const.PT_FRAMEWORK).update(
+            {
+                Const.PT_API_TYPE_NPU_DIST: (torch_npu.distributed, (torch_npu.distributed,
+                                                                     torch_npu.distributed.distributed_c10d))
+            }
+        )
+
+_inner_used_api = {}
+_supported_api_list_path = (os.path.join(os.path.dirname(os.path.realpath(__file__)), Const.SUPPORT_API_FILE_NAME),)
+_cuda_func_mapping = {"npu_fusion_attention": "gpu_fusion_attention"}
+
+
+@parameter_adapter
+def tensor_module_forward(module, *args, **kwargs):
+    return module.api_func(*args, **kwargs)
+
+
+def dist_module_forward(module, *args, **kwargs):
+    handle = module.api_func(*args, **kwargs)
+    if kwargs.get("async_op") or module.api_name in ["isend", "irecv"]:
+        if handle and hasattr(handle, 'wait'):
+            handle.wait()
+    if module.api_name == "batch_isend_irecv":
+        if isinstance(handle, list):
+            for req in handle:
+                req.wait()
+    return handle
+
+
+def npu_module_forward(module, *args, **kwargs):
+    if not module.need_hook:
+        if module.api_name not in npu_custom_functions:
+            raise Exception(f'There is not bench function {module.api_name}')
+        if module.device == Const.CUDA_LOWERCASE:
+            module.api_name = _cuda_func_mapping.get(module.api_name, module.api_name)
+        if module.device in [Const.CUDA_LOWERCASE, Const.CPU_LOWERCASE]:
+            return npu_custom_functions[module.api_name](*args, **kwargs)
+    return module.api_func(*args, **kwargs)
+
+
+forward_methods = {
+    "Tensor": tensor_module_forward,
+    "Distributed": dist_module_forward,
+    "NPU": npu_module_forward
+}
+
+
+class ApiTemplate(HOOKModule):
+    def __init__(self, api_name, api_func, prefix, hook_build_func, need_hook=True, device=Const.CPU_LOWERCASE):
+        self.api_name = api_name
+        self.api_func = api_func
+        self.prefix = prefix
+        self.prefix_api_name = prefix + Const.SEP + str(api_name.split(Const.SEP)[-1]) + Const.SEP
+        self.need_hook = need_hook
+        self.device = device
+        if self.need_hook:
+            super().__init__(hook_build_func)
+        if prefix == Const.DIST_API_TYPE_PREFIX:
+            self.op_is_distributed = True
+
+    @torch_device_guard
+    def forward(self, *args, **kwargs):
+        exec_func = forward_methods.get(self.prefix)
+        exec_func = functools.partial(exec_func, self) if exec_func else self.api_func
+        return exec_func(*args, **kwargs)
+
+
+api_register = None
+
+
+def get_api_register(return_new=False):
+    if return_new:
+        return ApiRegistry(_api_types, _inner_used_api, _supported_api_list_path, ApiTemplate)
+
+    global api_register
+    if api_register is None:
+        api_register = ApiRegistry(_api_types, _inner_used_api, _supported_api_list_path, ApiTemplate)
+    return api_register

msprobe/pytorch/hook_module/hook_module.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -21,6 +21,8 @@ import torch
 import torch.nn as nn
 import torch.utils.hooks as full_hooks
 
+from msprobe.pytorch.common.utils import is_float8_tensor
+
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
 
@@ -28,28 +30,27 @@ class HOOKModule(nn.Module):
     module_count = defaultdict(int)
     inner_stop_hook = {}
 
-    def __init__(self, build_hook) -> None:
+    def __init__(self, hook_build_func) -> None:
         super(HOOKModule, self).__init__()
         self.has_overflow = False
-        self.prefix = ""
         self.current_thread = threading.current_thread().ident
         if self.current_thread not in HOOKModule.inner_stop_hook:
             HOOKModule.inner_stop_hook[self.current_thread] = False
         self.stop_hook = HOOKModule.inner_stop_hook.get(self.current_thread, False)
 
         if not self.stop_hook:
-            if hasattr(self, "prefix_op_name_"):
-                self.prefix = self.prefix_op_name_
-
             self.forward_data_collected = False
-            forward_pre_hook, forward_hook, backward_hook, _ = build_hook(self.prefix)
-            if torch_version_above_or_equal_2:
-                self.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
-                self.register_forward_hook(forward_hook, with_kwargs=True)
-            else:
-                self.register_forward_pre_hook(forward_pre_hook)
-                self.register_forward_hook(forward_hook)
-            self.register_backward_hook(backward_hook)
+
+            prefix = self.prefix_api_name if hasattr(self, "prefix_api_name") else ""
+            if callable(hook_build_func):
+                forward_pre_hook, forward_hook, backward_hook, _ = hook_build_func(prefix)
+                if torch_version_above_or_equal_2:
+                    self.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
+                    self.register_forward_hook(forward_hook, with_kwargs=True)
+                else:
+                    self.register_forward_pre_hook(forward_pre_hook)
+                    self.register_forward_hook(forward_hook)
+                self.register_backward_hook(backward_hook)
 
     def __call__(self, *args, **kwargs):
         changed = False
@@ -111,6 +112,10 @@ class HOOKModule(nn.Module):
                         return result
                 else:
                     return result
+
+            if is_float8_tensor(var) or not (var.requires_grad and torch.is_grad_enabled()):
+                return result
+
             grad_fn = var.grad_fn
             if grad_fn is not None:
                 for hook in non_full_backward_hooks:

msprobe/pytorch/hook_module/register_optimizer_hook.py

@@ -32,8 +32,9 @@ def register_optimizer_hook(data_collector):
     def patch_clip_grad(func):
         def wrapper(*args, **kwargs):
             data_collector.optimizer_status = Const.CLIP_GRAD
-            func(*args, **kwargs)
+            result = func(*args, **kwargs)
             data_collector.optimizer_status = Const.END_PREFIX + Const.CLIP_GRAD
+            return result
 
         return wrapper