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

msprobe/pytorch/compare/pt_compare.py

@@ -14,19 +14,29 @@
 # limitations under the License.
 
 import os.path
+
 import torch
+
 from msprobe.core.common.const import FileCheckConst
-from msprobe.pytorch.common.log import logger
 from msprobe.core.common.exceptions import FileCheckException
-from msprobe.core.compare.acc_compare import Comparator
-from msprobe.core.common.utils import check_configuration_param, check_compare_param, \
-    CompareException, set_dump_path, get_dump_mode
 from msprobe.core.common.file_utils import FileChecker, create_directory, load_yaml
+from msprobe.core.common.utils import CompareException, check_compare_param, check_configuration_param, get_dump_mode, \
+    set_dump_path
+from msprobe.core.compare.acc_compare import Comparator, ModeConfig
+from msprobe.core.compare.utils import set_stack_json_path
+from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.common.utils import load_pt
 
 
-class PTComparator (Comparator):
-    def __init__(self, data_mapping=None):
+class PTComparator(Comparator):
+    def __init__(self, mode_config, data_mapping=None):
+        super().__init__(mode_config)
+
+        self.stack_mode = mode_config.stack_mode
+        self.auto_analyze = mode_config.auto_analyze
+        self.fuzzy_match = mode_config.fuzzy_match
+        self.dump_mode = mode_config.dump_mode
+
         self.frame_name = PTComparator.__name__
         self.data_mapping = data_mapping
         if isinstance(self.data_mapping, str) or self.data_mapping is None:
@@ -37,23 +47,24 @@ class PTComparator (Comparator):
             raise TypeError(f"The type of parameter `data_mapping` must be dict, str or None, but got "
                             f"{type(self.data_mapping)}")
 
-    def load_mapping_file(self, mapping_file):
+    @staticmethod
+    def load_mapping_file(mapping_file):
         if isinstance(mapping_file, str):
             mapping_dict = load_yaml(mapping_file)
         else:
             mapping_dict = {}
         return mapping_dict
-    
+
     def read_npy_data(self, dir_path, file_name):
         if not file_name:
             return None
         data_path = os.path.join(dir_path, file_name)
         path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE,
-                                FileCheckConst.PT_SUFFIX, False)
+                                   FileCheckConst.PT_SUFFIX, False)
         data_path = path_checker.common_check()
         try:
-            data_value = load_pt(data_path,
-                                 to_cpu=True).detach()  # detach because numpy can not process gradient information
+            # detach because numpy can not process gradient information
+            data_value = load_pt(data_path, to_cpu=True).detach()
         except RuntimeError as e:
             # 这里捕获 load_pt 中抛出的异常
             logger.error(f"Failed to load the .pt file at {data_path}.")
@@ -65,20 +76,29 @@ class PTComparator (Comparator):
         if data_value.dtype == torch.bfloat16:
             data_value = data_value.to(torch.float32)
         data_value = data_value.numpy()
-        return data_value  
-    
-    
-def compare(input_param, output_path, stack_mode=False, auto_analyze=True, fuzzy_match=False, **kwargs):
+        return data_value
+
+
+def compare(input_param, output_path, **kwargs):
     try:
+        auto_analyze = kwargs.get('auto_analyze', True)
+        fuzzy_match = kwargs.get('fuzzy_match', False)
+        data_mapping = kwargs.get('data_mapping', None)
+        suffix = kwargs.get('suffix', '')
+
         set_dump_path(input_param)
         dump_mode = get_dump_mode(input_param)
+        if "stack_json_path" in input_param:
+            stack_mode = kwargs.get('stack_mode', False)
+        else:
+            stack_mode = set_stack_json_path(input_param)  # set stack_mode and set "stack_json_path" in input_param
         check_configuration_param(stack_mode, auto_analyze, fuzzy_match, input_param.get('is_print_compare_log', True))
         create_directory(output_path)
-        check_compare_param(input_param, output_path, dump_mode)
-        data_mapping = kwargs.get('data_mapping', None)
+        check_compare_param(input_param, output_path, dump_mode, stack_mode)
     except (CompareException, FileCheckException) as error:
         logger.error('Compare failed. Please check the arguments and do it again!')
         raise CompareException(error.code) from error
-    pt_comparator = PTComparator(data_mapping)
-    pt_comparator.compare_core(input_param, output_path, stack_mode=stack_mode,
-                 auto_analyze=auto_analyze, fuzzy_match=fuzzy_match, dump_mode=dump_mode)
+
+    mode_config = ModeConfig(stack_mode, auto_analyze, fuzzy_match, dump_mode)
+    pt_comparator = PTComparator(mode_config, data_mapping)
+    pt_comparator.compare_core(input_param, output_path, suffix=suffix)

msprobe/pytorch/debugger/debugger_config.py

@@ -34,6 +34,7 @@ class DebuggerConfig:
         self.summary_mode = task_config.summary_mode if task_config.summary_mode else Const.STATISTICS
         self.overflow_nums = task_config.overflow_nums if task_config.overflow_nums else 1
         self.framework = Const.PT_FRAMEWORK
+        self.async_dump = common_config.async_dump if common_config.async_dump else False
 
         if self.level == Const.LEVEL_L2:
             self.is_backward_kernel_dump = False
@@ -74,29 +75,43 @@ class DebuggerConfig:
         if not self.dump_path:
             raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
                                    f"The dump_path not found.")
+        if not isinstance(self.async_dump, bool):
+            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                                   f"The parameters async_dump should be bool.")
 
     def check(self):
         self.check_kwargs()
         return True
 
     def check_model(self, instance, start_model):
-        if self.level not in ["L0", "mix"]:
+        if self.level not in [Const.LEVEL_L0, Const.LEVEL_MIX]:
             if instance.model is not None or start_model is not None:
-                logger.warning_on_rank_0(
+                logger.info_on_rank_0(
                     f"The current level is not L0 or mix level, so the model parameters will not be used.")
             return
-        if start_model is None:
-            if instance.model is None:
-                logger.error_on_rank_0(
-                    f"For level {self.level}, PrecisionDebugger or start interface must receive a 'model' argument.")
-                raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, f"missing the parameter 'model'")
-            return
-        if isinstance(start_model, torch.nn.Module):
-            instance.model = start_model
+        if start_model is None and instance.model is None:
+            logger.error_on_rank_0(
+                f"For level {self.level}, PrecisionDebugger or start interface must receive a 'model' parameter.")
+            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, f"missing the parameter 'model'")
+        
+        instance.model = start_model if start_model is not None else instance.model
+        if isinstance(instance.model, torch.nn.Module):
+            return 
+
+        error_model = None
+        if isinstance(instance.model, (list, tuple)):
+            for model in instance.model:
+                if not isinstance(model, torch.nn.Module):
+                    error_model = model
+                    break
         else:
-            logger.error_on_rank_0(f"The 'model' parameter of start must be a torch.nn.Module type.")
+            error_model = instance.model
+
+        if error_model is not None:
+            error_info = (f"The 'model' parameter must be a torch.nn.Moudle or list[torch.nn.Moudle] "
+                          f"type, currently there is a {type(error_model)} type.")
             raise MsprobeException(
-                MsprobeException.INVALID_PARAM_ERROR, f"model must be a torch.nn.Module")
+                MsprobeException.INVALID_PARAM_ERROR, error_info)
 
     def _check_and_adjust_config_with_l2(self):
         if self.scope:

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");
@@ -22,6 +22,7 @@ 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.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 +50,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 +60,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 +67,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 +106,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 +123,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):
@@ -159,6 +159,36 @@ class PrecisionDebugger:
         cls._instance.gm.monitor(model)
 
 
+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):
         debugger_instance = PrecisionDebugger.instance

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/{module_processer.py → dump/module_dump/module_processer.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");
@@ -17,12 +17,24 @@ from functools import wraps
 
 import torch
 from msprobe.core.common.const import Const
-from msprobe.core.data_dump.scope import ModuleRangeScope, MixRangeScope
+from msprobe.core.data_dump.scope import BaseScope, ModuleRangeScope, MixRangeScope
+from msprobe.pytorch.common.log import logger
+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 = []
@@ -34,6 +46,7 @@ class ModuleProcesser:
         BackwardHook.setup_input_hook = ModuleProcesser.clone_return_value(BackwardHook.setup_input_hook)
         BackwardHook.setup_output_hook = ModuleProcesser.clone_return_value(BackwardHook.setup_output_hook)
         BackwardHook.setup_output_hook = ModuleProcesser.filter_tensor_and_tuple(BackwardHook.setup_output_hook)
+        replace_checkpoint()
 
     @staticmethod
     def filter_tensor_and_tuple(func):
@@ -63,7 +76,7 @@ class ModuleProcesser:
             return ModuleProcesser.clone_if_tensor(result)
 
         return clone_return_value_func
-    
+
     @staticmethod
     def clone_if_tensor(result):
         if isinstance(result, torch.Tensor):
@@ -85,6 +98,22 @@ class ModuleProcesser:
             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 = {}
@@ -92,6 +121,42 @@ class ModuleProcesser:
         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):
@@ -100,7 +165,10 @@ class ModuleProcesser:
             except IndexError as e:
                 index = None
                 pass
-            module.mindstudio_reserved_name = full_name = name_prefix + Const.SEP + str(index)
+            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:
@@ -119,8 +187,11 @@ class ModuleProcesser:
                 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(module.mindstudio_reserved_name)
+                self.scope.end_module(current_name)
 
         def backward_hook(module, input, output=None):
             try:
@@ -128,7 +199,10 @@ class ModuleProcesser:
             except IndexError as e:
                 index = None
                 pass
-            module.mindstudio_reserved_name = full_name = name_prefix + Const.SEP + str(index)
+            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 = full_name.replace(Const.BACKWARD, Const.FORWARD)
             ModuleProcesser.module_node[full_name] = ModuleProcesser.module_node[forward_full_name].replace(
                 Const.FORWARD, Const.BACKWARD) if ModuleProcesser.module_node[forward_full_name] else None

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