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

msprobe/pytorch/debugger/precision_debugger.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,7 +21,9 @@ 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.pytorch.common.log import logger
+from msprobe.pytorch.common.utils import check_save_param
 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
@@ -49,7 +51,7 @@ class PrecisionDebugger:
         dump_path=None,
         level=None,
         model=None,
-        step=None,
+        step=None
     ):
         if not hasattr(self, "initialized"):
             config_params = ConfigParameters(config_path,
@@ -59,7 +61,6 @@ class PrecisionDebugger:
                                              model)
             self.check_input_params(config_params)
 
-            self.api_origin = False
             self.initialized = True
             self.model = model
             common_config, task_config = parse_json_config(config_path, task)
@@ -67,12 +68,13 @@ class PrecisionDebugger:
             if self.task == Const.GRAD_PROBE:
                 self.gm = GradientMonitor(common_config, task_config)
                 return
-            if step:
+            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.")
@@ -105,9 +107,11 @@ class PrecisionDebugger:
             raise MsprobeException(
                 MsprobeException.INVALID_PARAM_ERROR, f"level must be one of {Const.LEVEL_LIST}")
 
-        if args.model is not None and not isinstance(args.model, torch.nn.Module):
-            raise MsprobeException(
-                MsprobeException.INVALID_PARAM_ERROR, f"model must be a torch.nn.Module")
+        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."
+            )
 
     @classmethod
     def start(cls, model=None):
@@ -120,15 +124,12 @@ class PrecisionDebugger:
         if instance.enable_dataloader:
             logger.warning_on_rank_0("DataLoader is enabled, start() skipped.")
         else:
-            instance.service.start(instance.model, instance.api_origin)
-            instance.api_origin = False
+            instance.service.start(instance.model)
 
-    # 指定代码段dump前反向结束符，之后的计算过程数据将被忽略，无法被dump
     @classmethod
     def forward_backward_dump_end(cls):
         instance = cls._instance
-        instance.service.forward_backward_dump_end()
-        instance.api_origin = True
+        instance.stop()
 
     @classmethod
     def stop(cls):
@@ -158,6 +159,49 @@ class PrecisionDebugger:
             return
         cls._instance.gm.monitor(model)
 
+    @classmethod
+    def save(cls, variable, name, save_backward=True):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if instance.task not in [Const.TENSOR, Const.STATISTICS] or instance.config.level != Const.LEVEL_DEBUG:
+            return
+        try:
+            check_save_param(variable, name, save_backward)
+        except ValueError:
+            return
+        instance.service.save(variable, name, save_backward)
+
+
+def module_dump(module, dump_name):
+    if not isinstance(module, torch.nn.Module):
+        raise MsprobeException(
+            MsprobeException.INVALID_PARAM_ERROR,
+            f"the module argument in module_dump must be a torch.nn.Module subclass"
+        )
+    if not isinstance(dump_name, str):
+        raise MsprobeException(
+            MsprobeException.INVALID_PARAM_ERROR,
+            f"the dump_name argument in module_dump must be a str type"
+        )
+    instance = PrecisionDebugger._instance
+    if not instance:
+        raise MsprobeException(
+            MsprobeException.INTERFACE_USAGE_ERROR,
+            f"PrecisionDebugger must be instantiated before using module_dump interface"
+        )
+    instance.module_dumper.start_module_dump(module, dump_name)
+
+
+def module_dump_end():
+    instance = PrecisionDebugger._instance
+    if not instance:
+        raise MsprobeException(
+            MsprobeException.INTERFACE_USAGE_ERROR,
+            f"PrecisionDebugger must be instantiated before using module_dump_end interface"
+        )
+    instance.module_dumper.stop_module_dump()
+
 
 def iter_tracer(func):
     def func_wrapper(*args, **kwargs):

msprobe/pytorch/dump/module_dump/module_dump.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2024-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 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'
+
+
+class ModuleDumper:
+    def __init__(self, service):
+        self.service = service
+        self.hook_handle_list = []
+
+    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()
+
+    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
+        )
+
+        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])

msprobe/pytorch/dump/module_dump/module_processer.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2024-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 msprobe.core.common.const import Const
+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
+
+torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
+
+
+def checkpoint_without_early_stop(*args, **kwargs):
+    with set_checkpoint_early_stop(False):
+        return origin_checkpoint(*args, **kwargs)
+
+
+def replace_checkpoint():
+    torch.utils.checkpoint.checkpoint = checkpoint_without_early_stop
+
+
+class ModuleProcesser:
+    module_count = {}
+    module_stack = []
+    api_parent_node = ""
+    module_node = {}
+
+    def __init__(self, scope):
+        self.scope = scope if isinstance(scope, (ModuleRangeScope, MixRangeScope)) else None
+        BackwardHook.setup_input_hook = ModuleProcesser.clone_return_value(BackwardHook.setup_input_hook)
+        BackwardHook.setup_output_hook = ModuleProcesser.clone_return_value(BackwardHook.setup_output_hook)
+        replace_checkpoint()
+
+    @staticmethod
+    def clone_return_value(func):
+        @wraps(func)
+        def clone_return_value_func(*args, **kwargs):
+            result = func(*args, **kwargs)
+            return ModuleProcesser.clone_if_tensor(result)
+
+        return clone_return_value_func
+
+    @staticmethod
+    def clone_if_tensor(result):
+        if isinstance(result, torch.Tensor):
+            return result.clone()
+        elif type(result) is tuple:
+            return tuple(ModuleProcesser.clone_if_tensor(x) for x in result)
+        elif type(result) is list:
+            return list(ModuleProcesser.clone_if_tensor(x) for x in result)
+        elif type(result) is dict:
+            return {k: ModuleProcesser.clone_if_tensor(v) for k, v in result.items()}
+        else:
+            return result
+
+    @staticmethod
+    def module_count_func(module_name):
+        if module_name not in ModuleProcesser.module_count:
+            ModuleProcesser.module_count[module_name] = 0
+        else:
+            ModuleProcesser.module_count[module_name] += 1
+        return ModuleProcesser.module_count[module_name]
+
+    @staticmethod
+    def has_register_backward_hook(module):
+        return hasattr(module, '_backward_hooks') and \
+            len(module._backward_hooks) > 0 and \
+            module._is_full_backward_hook is False
+
+    @staticmethod
+    def get_modules_and_names(models):
+        modules_and_names_with_index = {}
+        if isinstance(models, (list, tuple)):
+            for index, model in enumerate(models):
+                modules_and_names_with_index[str(index)] = model.named_modules()
+        else:
+            modules_and_names_with_index["-1"] = models.named_modules()
+        return modules_and_names_with_index
+
+    @classmethod
+    def reset_module_stats(cls):
+        cls.module_count = {}
+        cls.module_stack = []
+        cls.api_parent_node = ""
+        cls.module_node = {}
+
+    def register_module_hook(self, models, build_hook):
+        logger.info_on_rank_0("The init dump is enabled, and the module dump function will not be available.")
+        modules_and_names_with_index = self.get_modules_and_names(models)
+        for index, modules_and_names in modules_and_names_with_index.items():
+            model = models if index == "-1" else models[int(index)]
+            for name, module in modules_and_names:
+                if module == model:
+                    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)
+                pre_forward_hook, forward_hook, backward_hook, forward_hook_torch_version_below_2 = build_hook(
+                    BaseScope.Module_Type_Module,
+                    prefix_name
+                )
+
+                if self.has_register_backward_hook(module):
+                    logger.warning(
+                        f"The {prefix_name[:-1]} 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:
+                    module.register_forward_hook(forward_hook, with_kwargs=True)
+                else:
+                    if not self.has_register_backward_hook(module):
+                        module.register_full_backward_hook(self.node_hook(prefix_name + Const.BACKWARD, Const.STOP))
+                    module.register_forward_hook(forward_hook_torch_version_below_2)
+                if not self.has_register_backward_hook(module):
+                    module.register_full_backward_hook(backward_hook)
+
+                module.register_forward_pre_hook(self.node_hook(prefix_name + Const.FORWARD, Const.START))
+                module.register_forward_hook(self.node_hook(prefix_name + Const.FORWARD, Const.STOP))
+                if torch_version_above_or_equal_2 and not self.has_register_backward_hook(module):
+                    module.register_full_backward_pre_hook(self.node_hook(prefix_name + Const.BACKWARD, Const.START))
+                    module.register_full_backward_hook(self.node_hook(prefix_name + Const.BACKWARD, Const.STOP))
+
+    def node_hook(self, name_prefix, start_or_stop, **kwargs):
+
+        def pre_hook(module, input, output=None):
+            try:
+                index = ModuleProcesser.module_count_func(name_prefix)
+            except IndexError as e:
+                index = None
+                pass
+            full_name = name_prefix + Const.SEP + str(index)
+            if not hasattr(module, "mindstudio_reserved_name") or not module.mindstudio_reserved_name:
+                module.mindstudio_reserved_name = []
+            module.mindstudio_reserved_name.append(full_name)
+            if self.module_stack:
+                ModuleProcesser.module_node[full_name] = self.module_stack[-1]
+            else:
+                ModuleProcesser.module_node[full_name] = None
+
+            ModuleProcesser.module_stack.append(full_name)
+            if self.module_stack:
+                ModuleProcesser.api_parent_node = self.module_stack[-1]
+            if self.scope:
+                self.scope.begin_module(full_name)
+
+        def end_hook(module, input, output=None):
+            if self.module_stack:
+                ModuleProcesser.module_stack.pop()
+            if self.module_stack:
+                ModuleProcesser.api_parent_node = self.module_stack[-1]
+            else:
+                ModuleProcesser.api_parent_node = None
+            if not hasattr(module, "mindstudio_reserved_name") or not module.mindstudio_reserved_name:
+                raise RuntimeError(f"module reserve name is None when pop")
+            current_name = module.mindstudio_reserved_name.pop()
+            if self.scope:
+                self.scope.end_module(current_name)
+
+        def backward_hook(module, input, output=None):
+            try:
+                index = ModuleProcesser.module_count_func(name_prefix)
+            except IndexError as e:
+                index = None
+                pass
+            full_name = name_prefix + Const.SEP + str(index)
+            if not hasattr(module, "mindstudio_reserved_name") or not module.mindstudio_reserved_name:
+                module.mindstudio_reserved_name = []
+            module.mindstudio_reserved_name.append(full_name)
+            forward_full_name = replace_last_occurrence(full_name, Const.BACKWARD, Const.FORWARD)
+            ModuleProcesser.module_node[full_name] = replace_last_occurrence(
+                ModuleProcesser.module_node.get(forward_full_name), Const.FORWARD, Const.BACKWARD)
+            ModuleProcesser.api_parent_node = None
+            if self.scope:
+                self.scope.begin_module(full_name)
+
+        if torch_version_above_or_equal_2:
+            if Const.START in start_or_stop:
+                return pre_hook
+            else:
+                return end_hook
+        else:
+            if Const.FORWARD in name_prefix and Const.START in start_or_stop:
+                return pre_hook
+            elif Const.BACKWARD in name_prefix:
+                return backward_hook
+            else:
+                return end_hook

msprobe/pytorch/free_benchmark/common/params.py

@@ -39,7 +39,6 @@ class DataParams:
     origin_func: Optional[Callable] = None
     api_type: Optional[str] = None
     fuzz_stage: Optional[str] = None
-    grad_unequal_flag: Optional[bool] = True
 
 
 @dataclass
@@ -127,6 +126,8 @@ def make_unequal_row(
     )
     if isinstance(ratio, float):
         row.max_rel = ratio - 1
+    if isinstance(ratio, str):
+        row.max_rel = ratio
     origin_tensor = data_params.original_result
     perturbed_tensor = data_params.perturbed_result
     if index is not None:

msprobe/pytorch/free_benchmark/common/utils.py

@@ -124,6 +124,7 @@ class TorchC:
     abs = torch._C._VariableFunctionsClass.abs
     where = torch._C._VariableFunctionsClass.where
     div = torch._C._VariableFunctionsClass.div
+    mul = torch._C._VariableFunctionsClass.mul
     max = torch._C._VariableFunctionsClass.max
     min = torch._C._VariableFunctionsClass.min
     gt = torch._C._VariableFunctionsClass.gt
@@ -138,3 +139,5 @@ class TorchC:
     tensor_split = torch._C._VariableFunctionsClass.tensor_split
     stack = torch._C._VariableFunctionsClass.stack
     reshape = torch._C._VariableFunctionsClass.reshape
+    nan_to_num = torch._C._VariableFunctionsClass.nan_to_num
+    aminmax = torch._C._VariableFunctionsClass.aminmax

msprobe/pytorch/free_benchmark/compare/grad_saver.py

@@ -82,13 +82,11 @@ class GradSaver:
         data_params = DataParams()
         data_params.original_result = origin_grad
         data_params.perturbed_result = perturbed_grad
-        data_params.grad_unequal_flag = False
         data_params.valid_input_index = index
         try:
             handler.handle(data_params)
             if not data_params.is_consistent:
                 self.is_compare = False
-                data_params.grad_unequal_flag = True
                 data_params.is_consistent = True
                 data_params.perturbed_result = self.perturbed_grad_input
                 data_params.original_result = self.origin_grad_input

msprobe/pytorch/free_benchmark/result_handlers/base_handler.py

@@ -89,12 +89,6 @@ class FuzzHandler(ABC):
         )
         return origin_output_chunks, perturbed_output_chunks
 
-    @staticmethod
-    def convert_overflow_ratio_to_consistent(ratio):
-        if math.isnan(ratio) or math.isinf(ratio):
-            return ThresholdConfig.COMP_CONSISTENT
-        return ratio
-
     @abstractmethod
     def get_threshold(self, dtype):
         pass
@@ -107,10 +101,10 @@ class FuzzHandler(ABC):
         self, origin_output, perturbed_output, norm_type, abs_tol
     ):
         if norm_type == NormType.ENDLESS_NORM:
-            return self.calculate_error(origin_output, perturbed_output, abs_tol)
+            return self.calculate_max_ratio(origin_output, perturbed_output, abs_tol)
         return ThresholdConfig.COMP_CONSISTENT
 
-    def calculate_error(self, origin_output, perturbed_output, abs_tol):
+    def calculate_max_ratio(self, origin_output, perturbed_output, abs_tol):
         origin_output_chunks, perturbed_output_chunks = (
             self.tensor_split_for_error_calculate(origin_output, perturbed_output)
         )
@@ -122,42 +116,30 @@ class FuzzHandler(ABC):
             raise FreeBenchmarkException(
                 FreeBenchmarkException.OutputIndexError, err_msg
             )
-        norm1 = -np.inf
-        norm2 = -np.inf
-        norm3 = np.inf
+
+        max_ratio = ThresholdConfig.COMP_CONSISTENT
         for i, chunk_origin in enumerate(origin_output_chunks):
             if chunk_origin.nelement() == 0:
                 break
             chunk_perturbed = perturbed_output_chunks[i]
-            ratio_tensor1 = TorchC.where(
-                TorchC.abs(chunk_perturbed) > abs_tol,
-                TorchC.div(
-                    TorchC.clamp(chunk_origin, min=abs_tol),
-                    TorchC.clamp(chunk_perturbed, min=abs_tol),
-                ),
-                1,
-            )
-            ratio_tensor2 = TorchC.where(
-                TorchC.abs(chunk_origin) > abs_tol,
-                TorchC.div(
-                    TorchC.clamp(chunk_perturbed, min=abs_tol),
-                    TorchC.clamp(chunk_origin, min=abs_tol),
-                ),
-                1,
+            # 如果乘积最小值 < 极小值乘积的负值，认为存在非极小值符号相反的情况
+            if TorchC.lt(
+                TorchC.min(TorchC.mul(chunk_origin, chunk_perturbed)), -(abs_tol**2)
+            ):
+                return ThresholdConfig.SYMBOL_FLIPPING
+            # 求A/B B/A的比值前，将值限制在大于极小值范围内
+            clamp_origin = TorchC.clamp(TorchC.abs(chunk_origin), min=abs_tol)
+            clamp_perturbed = TorchC.clamp(TorchC.abs(chunk_perturbed), min=abs_tol)
+            # 对于计算结果为nan的情况，认为两者没有差异
+            ratio_tensor = TorchC.nan_to_num(
+                TorchC.div(clamp_origin, clamp_perturbed),
+                nan=ThresholdConfig.COMP_CONSISTENT,
             )
-            norm_values = TorchC.stack(
-                [TorchC.max(ratio_tensor1), TorchC.max(ratio_tensor2)]
-            )
-            max_ratio1, max_ratio2 = norm_values.tolist()
-            norm1 = max(norm1, self.convert_overflow_ratio_to_consistent(max_ratio1))
-            norm2 = max(norm2, self.convert_overflow_ratio_to_consistent(max_ratio2))
-            norm3 = min(norm3, self.convert_overflow_ratio_to_consistent(max_ratio1))
-
-        if norm3 < 0:
-            ratio = ThresholdConfig.SYMBOL_FLIPPING
-        else:
-            ratio = max(norm1, norm2)
-        return ratio
+            # 求A/B 和 B/A比值最大值，其中 B/A的最大值为 A/B的最小值的倒数
+            min_ratio, max_ratio = TorchC.stack([*TorchC.aminmax(ratio_tensor)]).tolist()
+            min_ratio_reciprocal = np.inf if min_ratio == 0 else 1 / min_ratio
+            max_ratio = max(max_ratio, min_ratio_reciprocal)
+        return max_ratio
 
     def ratio_calculate(self, origin_output, perturbed_output, norm_type) -> float:
         try:
@@ -220,10 +202,12 @@ class FuzzHandler(ABC):
                 )
                 npu_consistent = is_consistent
                 max_fuzz_ratio = (
-                    max_fuzz_ratio if ratio is None else max(max_fuzz_ratio, ratio)
+                    max_fuzz_ratio
+                    if not isinstance(ratio, (int, float))
+                    else max(max_fuzz_ratio, ratio)
                 )
-                data_params.is_consistent = is_consistent and data_params.is_consistent
-                if not is_consistent and data_params.grad_unequal_flag:
+                data_params.is_consistent = is_consistent
+                if not is_consistent:
                     self.unequal_rows.append(
                         make_unequal_row(data_params, self.params, ratio=ratio)
                     )
@@ -235,12 +219,12 @@ class FuzzHandler(ABC):
                     )
                     npu_consistent = npu_consistent and is_consistent
                     max_fuzz_ratio = (
-                        max_fuzz_ratio if ratio is None else max(max_fuzz_ratio, ratio)
-                    )
-                    data_params.is_consistent = (
-                        is_consistent and data_params.is_consistent
+                        max_fuzz_ratio
+                        if not isinstance(ratio, (int, float))
+                        else max(max_fuzz_ratio, ratio)
                     )
-                    if not is_consistent and data_params.grad_unequal_flag:
+                    data_params.is_consistent = is_consistent
+                    if not is_consistent:
                         self.unequal_rows.append(
                             make_unequal_row(
                                 data_params, self.params, ratio=ratio, index=index_

msprobe/pytorch/free_benchmark/result_handlers/preheat_handler.py

@@ -75,10 +75,6 @@ class PreheatHandler(FuzzHandler):
         if self.params.preheat_config.get("preheat_step") <= self.params.step:
             return data_params.original_result
 
-        if not data_params.grad_unequal_flag:
-            data_params.grad_unequal_flag = True
-            data_params.is_consistent = False
-            return data_params.original_result
         preheat_counter.add_api_called_time(self.pure_name)
 
         if not self._is_take_a_sample():

msprobe/pytorch/function_factory.py

@@ -27,6 +27,11 @@ from msprobe.pytorch.bench_functions.rotary_mul import npu_rotary_mul, npu_rotar
 from msprobe.pytorch.bench_functions.scaled_mask_softmax import npu_scaled_masked_softmax, \
     npu_scaled_masked_softmax_backward
 from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward
+from msprobe.pytorch.bench_functions.apply_adam import npu_apply_adam
+from msprobe.pytorch.bench_functions.group_norm_silu import npu_group_norm_silu
+from msprobe.pytorch.bench_functions.mish import npu_mish
+from msprobe.pytorch.bench_functions.moe_gating_top_k_softmax import npu_moe_gating_top_k_softmax
+from msprobe.pytorch.bench_functions.sort_v2 import npu_sort_v2
 from msprobe.pytorch.common.utils import logger
 
 
@@ -79,7 +84,8 @@ class Register(dict):
 npu_custom_functions = Register()
 npu_custom_functions([
     npu_apply_adam_w, npu_confusion_transpose, npu_fast_gelu, npu_layer_norm_eval, npu_linear, npu_fusion_attention,
-    npu_rms_norm, npu_rotary_mul, npu_scaled_masked_softmax, npu_swiglu, gpu_fusion_attention
+    npu_rms_norm, npu_rotary_mul, npu_scaled_masked_softmax, npu_swiglu, gpu_fusion_attention, npu_apply_adam,
+    npu_group_norm_silu, npu_mish, npu_moe_gating_top_k_softmax, npu_sort_v2
 ])
 
 # register for npu custom backward bench functions

msprobe/pytorch/hook_module/__init__.py

@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .wrap_functional import remove_dropout
+from msprobe.pytorch.common.utils import remove_dropout