mindstudio-probe 1.0.4__py3-none-any.whl → 1.1.0__py3-none-any.whl

msprobe/pytorch/free_benchmark/result_handlers/base_handler.py

@@ -1,8 +1,23 @@
+# Copyright (c) 2024-2024, 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 math
 from abc import ABC, abstractmethod
 from typing import Any, Optional, Tuple
-import numpy as np
 
+import numpy as np
 import torch
 from msprobe.core.common.const import Const
 from msprobe.pytorch.free_benchmark import logger
@@ -35,7 +50,9 @@ class FuzzHandler(ABC):
             origin_ouput = origin_ouput.values
             perturbed_output = perturbed_output.values
         if hasattr(perturbed_output, "dtype"):
-            abs_tol = ThresholdConfig.ABS_TOL_VALUE_DICT.get(perturbed_output.dtype, FuzzThreshold.F32_THD)
+            abs_tol = ThresholdConfig.ABS_TOL_VALUE_DICT.get(
+                perturbed_output.dtype, FuzzThreshold.F32_THD
+            )
         else:
             abs_tol = FuzzThreshold.F32_THD
         return (
@@ -53,16 +70,22 @@ class FuzzHandler(ABC):
         :return origin_output_chunks: 切块后原始输出列表
         :return perturbed_output_chunks: 切块后扰动后输出列表
         """
-        single_output_mem = origin_output.element_size() * origin_output.nelement() / Const.ONE_MB
+        single_output_mem = (
+            origin_output.element_size() * origin_output.nelement() / Const.ONE_MB
+        )
         if single_output_mem == 0 or origin_output.ndim == 0:
             return [origin_output], [perturbed_output]
         # 张量大小和批数之间的关系：chunks_exp=math.log(M,2)-4, chunks=2**chunks_exp (M为对比张量数据大小[Mb])
         chunks_exp = int(math.log(single_output_mem, 2)) - 4
-        chunks = 2 ** chunks_exp
+        chunks = 2**chunks_exp
         chunks = max(chunks, 1)
         chunks = min(chunks, ThresholdConfig.TENSOR_SPLIT_MAX_CHUNK)
-        origin_output_chunks = TorchC.tensor_split(TorchC.reshape(origin_output, (-1,)), chunks)
-        perturbed_output_chunks = TorchC.tensor_split(TorchC.reshape(perturbed_output, (-1,)), chunks)
+        origin_output_chunks = TorchC.tensor_split(
+            TorchC.reshape(origin_output, (-1,)), chunks
+        )
+        perturbed_output_chunks = TorchC.tensor_split(
+            TorchC.reshape(perturbed_output, (-1,)), chunks
+        )
         return origin_output_chunks, perturbed_output_chunks
 
     @staticmethod
@@ -80,14 +103,16 @@ class FuzzHandler(ABC):
         pass
 
     def get_ratio_from_specific_norm(
-            self, origin_output, perturbed_output, norm_type, abs_tol
+        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 ThresholdConfig.COMP_CONSISTENT
 
     def calculate_error(self, origin_output, perturbed_output, abs_tol):
-        origin_output_chunks, perturbed_output_chunks = self.tensor_split_for_error_calculate(origin_output, perturbed_output)
+        origin_output_chunks, perturbed_output_chunks = (
+            self.tensor_split_for_error_calculate(origin_output, perturbed_output)
+        )
         norm1 = -np.inf
         norm2 = -np.inf
         norm3 = np.inf
@@ -95,11 +120,25 @@ class FuzzHandler(ABC):
             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)
-            norm_values = TorchC.stack([TorchC.max(ratio_tensor1), TorchC.max(ratio_tensor2)])
+            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,
+            )
+            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))
@@ -126,13 +165,13 @@ class FuzzHandler(ABC):
         if self.params.fuzz_stage == Const.BACKWARD:
             abs_tol = ThresholdConfig.BACKWARD_OUTPUT_LOWER_BOUND
         else:
-            abs_tol = abs_tol ** 0.5
+            abs_tol = abs_tol**0.5
         return self.get_ratio_from_specific_norm(
             origin_output, perturbed_output, norm_type, abs_tol
         )
 
     def npu_compare(
-            self, origin_output, perturbed_output
+        self, origin_output, perturbed_output
     ) -> Tuple[bool, Optional[float]]:
 
         if isinstance(perturbed_output, int):
@@ -189,7 +228,7 @@ class FuzzHandler(ABC):
                         max_fuzz_ratio if ratio is None else max(max_fuzz_ratio, ratio)
                     )
                     data_params.is_consistent = (
-                            is_consistent and data_params.is_consistent
+                        is_consistent and data_params.is_consistent
                     )
                     if not is_consistent and data_params.grad_unequal_flag:
                         self.unequal_rows.append(

msprobe/pytorch/free_benchmark/result_handlers/check_handler.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, 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 typing import Any
 
 from msprobe.pytorch.free_benchmark import logger

msprobe/pytorch/free_benchmark/result_handlers/fix_handler.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, 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 typing import Any
 
 from msprobe.pytorch.free_benchmark.common.params import DataParams

msprobe/pytorch/free_benchmark/result_handlers/handler_factory.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, 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 msprobe.pytorch.free_benchmark import FreeBenchmarkException
 from msprobe.pytorch.free_benchmark.common.constant import PreheatConfig
 from msprobe.pytorch.free_benchmark.common.enums import HandlerType

msprobe/pytorch/free_benchmark/result_handlers/preheat_handler.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, 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 math
 from typing import Any
 
@@ -118,8 +133,10 @@ class PreheatHandler(FuzzHandler):
         """
         # 每一步样本数
         total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-        sample_count_per_step = self._get_sample_count_per_step()
         need_sample_set = set()
+        if total_count == 0:
+            return need_sample_set
+        sample_count_per_step = self._get_sample_count_per_step()
         prehead_step = self.params.preheat_config.get("preheat_step")
         for i in range(1, sample_count_per_step + 1):
             count = (prehead_step * (i - 1) + self.params.step) % total_count
@@ -136,9 +153,7 @@ class PreheatHandler(FuzzHandler):
 
     def _adjust_threshold_for_dtype(self, dtype_str, compare_result):
         con_ratio = [ratio for ratio, is_consistent in compare_result if is_consistent]
-        incon_ratio = [
-            ratio for ratio, is_consistent in compare_result if not is_consistent
-        ]
+        incon_ratio = [ratio for ratio, is_consistent in compare_result if not is_consistent]
         old_thd = preheat_counter.get_api_thd(self.pure_name, dtype_str)
         new_thd = old_thd
         # 正例负例都存在

msprobe/pytorch/function_factory.py

@@ -1,4 +1,18 @@
-from msprobe.pytorch.common.utils import logger
+# Copyright (c) 2024-2024, 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 msprobe.pytorch.bench_functions.apply_adam_w import npu_apply_adam_w
 from msprobe.pytorch.bench_functions.confusion_transpose import npu_confusion_transpose, \
     npu_confusion_transpose_backward
@@ -12,7 +26,8 @@ from msprobe.pytorch.bench_functions.rms_norm import npu_rms_norm, npu_rms_norm_
 from msprobe.pytorch.bench_functions.rotary_mul import npu_rotary_mul, npu_rotary_mul_backward
 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, swish_grad, swish
+from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward
+from msprobe.pytorch.common.utils import logger
 
 
 class Register(dict):

msprobe/pytorch/functional/module_dump.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2024-2024, 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
+import torch.nn as nn
+from msprobe.core.common.const import Const
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.data_dump.scope import BaseScope
+from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.debugger.precision_debugger import PrecisionDebugger
+from msprobe.pytorch.hook_module.api_registry import api_register
+from msprobe.pytorch.service import torch_version_above_or_equal_2
+
+hook_handle_list = []
+
+
+def module_dump(module, dump_name):
+    if not isinstance(module, nn.Module):
+        logger.error("The parameter module in module_dump must be a Module subclass.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if not isinstance(dump_name, str):
+        logger.error("The parameter dump_name in module_dump must be a str type.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+
+    api_register.api_originality()
+    register_hook(module, dump_name)
+
+
+def module_dump_end():
+    api_register.api_modularity()
+    remove_hook()
+    hook_handle_list.clear()
+
+
+def register_hook(module, dump_name):
+    prefix = BaseScope.Module_Type_Module + Const.SEP + dump_name + Const.SEP + module.__class__.__name__ + Const.SEP
+
+    pdg = PrecisionDebugger()
+    _, forward_hook, backward_hook, forward_hook_torch_version_below_2 = \
+        pdg.service.build_hook(BaseScope.Module_Type_Module, prefix)
+
+    if torch_version_above_or_equal_2:
+        forward_hook_handle = module.register_forward_hook(forward_hook, with_kwargs=True)
+        hook_handle_list.append(forward_hook_handle)
+    else:
+        pdg.service.check_register_full_backward_hook(module)
+        full_backward_hook_handle = module.register_full_backward_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
+        forward_hook_handle = module.register_forward_hook(forward_hook_torch_version_below_2)
+        hook_handle_list.extend([full_backward_hook_handle, forward_hook_handle])
+    pdg.service.check_register_full_backward_hook(module)
+    full_backward_hook_handle = module.register_full_backward_hook(backward_hook)
+
+    forward_pre_hook_handle = module.register_forward_pre_hook(
+        pdg.service.module_processor.node_hook(prefix + Const.FORWARD, Const.START))
+    forward_hook_handle = module.register_forward_hook(
+        pdg.service.module_processor.node_hook(prefix + Const.FORWARD, Const.STOP))
+    hook_handle_list.extend([full_backward_hook_handle, forward_pre_hook_handle, forward_hook_handle])
+
+    if torch_version_above_or_equal_2:
+        backward_pre_hook_handle = module.register_full_backward_pre_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.START))
+        pdg.service.check_register_full_backward_hook(module)
+        full_backward_hook_handle = module.register_full_backward_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
+        hook_handle_list.extend([backward_pre_hook_handle, full_backward_hook_handle])
+
+
+def remove_hook():
+    for hook_handle in hook_handle_list:
+        if isinstance(hook_handle, torch.utils.hooks.RemovableHandle):
+            hook_handle.remove()

msprobe/pytorch/grad_probe/grad_stat_csv.py

@@ -100,7 +100,7 @@ class CSV_max(CsvItem):
 
 
 @register_csv_item(GradConst.MIN)
-class CSV_max(CsvItem):
+class CSV_min(CsvItem):
     def generate_csv_header(csv_header_input):
         return ["min"]
 
@@ -110,7 +110,7 @@ class CSV_max(CsvItem):
 
 
 @register_csv_item(GradConst.NORM)
-class CSV_max(CsvItem):
+class CSV_norm(CsvItem):
     def generate_csv_header(csv_header_input):
         return ["norm"]
 

msprobe/pytorch/hook_module/__init__.py

@@ -1 +1,16 @@
-from .wrap_functional import remove_dropout
+# Copyright (c) 2024-2024, 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 .wrap_functional import remove_dropout

msprobe/pytorch/hook_module/api_registry.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2023. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, 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
 #
@@ -13,7 +12,6 @@
 # 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
 import torch.distributed as dist
@@ -107,7 +105,14 @@ class ApiRegistry:
         if not is_gpu:
             self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
 
-    def initialize_hook(self, hook):
+    def initialize_hook(self, hook, online_run_ut=False):
+        """
+        initialize_hook
+        Args:
+            hook (_type_): initialize_hook
+            online_run_ut (bool): default False, whether online run_ut or not.
+                If online_run_ut is True, the hook will not wrap the aten ops.
+        """
         self.store_ori_attr(torch.Tensor, get_tensor_ops(), self.tensor_ori_attr)
         wrap_tensor.wrap_tensor_ops_and_bind(hook)
         for attr_name in dir(wrap_tensor.HOOKTensor):
@@ -137,7 +142,7 @@ class ApiRegistry:
                     self.npu_distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP,
                                                                             attr_name)
 
-        if torch_version_above_2:
+        if torch_version_above_2 and not online_run_ut:
             self.store_ori_attr(torch.ops.aten, get_aten_ops(), self.aten_ori_attr)
             wrap_aten.wrap_aten_ops_and_bind(hook)
             for attr_name in dir(wrap_aten.HOOKAtenOP):

msprobe/pytorch/hook_module/hook_module.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2019-2020. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, 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
 #
@@ -13,7 +12,6 @@
 # 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 threading
@@ -58,12 +56,12 @@ class HOOKModule(nn.Module):
                 self.register_forward_hook(forward_hook)
             self.register_backward_hook(backward_hook)
 
-    def __call__(self, *input, **kwargs):
+    def __call__(self, *args, **kwargs):
         changed = False
         if not self.stop_hook:
             HOOKModule.inner_stop_hook[self.current_thread] = True
             changed = True
-        result = self._call_func(*input, **kwargs)
+        result = self._call_func(*args, **kwargs)
         if changed:
             HOOKModule.inner_stop_hook[self.current_thread] = False
         return result
@@ -72,28 +70,28 @@ class HOOKModule(nn.Module):
     def reset_module_stats(cls):
         cls.module_count = {}
 
-    def _call_func(self, *input, **kwargs):
+    def _call_func(self, *args, **kwargs):
         full_backward_hooks, non_full_backward_hooks = [], []
         if len(self._backward_hooks) > 0:
             full_backward_hooks, non_full_backward_hooks = self._get_backward_hooks()
         for hook in self._forward_pre_hooks.values():
-            result_input, result_kwargs = hook(self, input, kwargs)
-            if result_input is not None:
-                if not isinstance(result_input, tuple):
-                    result_input = (result_input,)
-                input = result_input
+            result_args, result_kwargs = hook(self, args, kwargs)
+            if result_args is not None:
+                if not isinstance(result_args, tuple):
+                    result_args = (result_args,)
+                args = result_args
             if result_kwargs is not None:
                 kwargs = result_kwargs
         bw_hook = None
         if len(full_backward_hooks) > 0:
             bw_hook = full_hooks.BackwardHook(self, full_backward_hooks)
-            input = bw_hook.setup_input_hook(input)
+            args = bw_hook.setup_input_hook(args)
         if torch._C._get_tracing_state():
-            result = self._slow_forward(*input, **kwargs)
+            result = self._slow_forward(*args, **kwargs)
         else:
-            result = self.forward(*input, **kwargs)
+            result = self.forward(*args, **kwargs)
         for hook in self._forward_hooks.values():
-            hook_result = hook(self, input, kwargs, result)
+            hook_result = hook(self, args, kwargs, result)
             if hook_result is not None:
                 result = hook_result
         if bw_hook:
@@ -116,5 +114,5 @@ class HOOKModule(nn.Module):
                     wrapper = functools.partial(hook, self)
                     functools.update_wrapper(wrapper, hook)
                     grad_fn.register_hook(wrapper)
-                self._maybe_warn_non_full_backward_hook(input, result, grad_fn)
+                self._maybe_warn_non_full_backward_hook(args, result, grad_fn)
         return result

msprobe/pytorch/hook_module/utils.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2024-2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, 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
 #
@@ -13,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 
 import os
 from msprobe.core.common.file_utils import load_yaml

msprobe/pytorch/hook_module/wrap_aten.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2023-2023. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, 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
 #
@@ -13,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 
 import os
 import torch
@@ -24,6 +22,7 @@ from msprobe.core.common.const import Const
 from msprobe.core.common.file_utils import load_yaml
 from msprobe.pytorch.function_factory import npu_custom_grad_functions
 
+
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
 ops = load_yaml(yaml_path)
@@ -50,6 +49,8 @@ class AtenOPTemplate(HOOKModule):
     def __init__(self, op, hook, need_hook=True):
         if isinstance(op, torch._ops.OpOverloadPacket):
             op_name_ = op._qualified_op_name.split("::")[-1]
+        elif isinstance(op, str):
+            op_name_ = str(op)
         else:
             op_name_ = op.name().split("::")[-1]
             overload_name = op._overloadname
@@ -76,13 +77,13 @@ class AtenOPTemplate(HOOKModule):
 
 
 class AtenOPPacketTemplate():
-    def __init__(self, opPacket, hook):
-        self.opPacket = opPacket
+    def __init__(self, op_packet, hook):
+        self.op_packet = op_packet
         self.hook = hook
 
     def __getattr__(self, key):
         try:
-            attr = getattr(self.opPacket, key)
+            attr = getattr(self.op_packet, key)
         except AttributeError as e:
             raise AttributeError(f"AtenOPPacketTemplate or OpOverloadPacket does not have attribute '{key}'.") from e
         if isinstance(attr, torch._ops.OpOverload):
@@ -92,10 +93,10 @@ class AtenOPPacketTemplate():
 
     @torch_device_guard
     def __call__(self, *args, **kwargs):
-        return AtenOPTemplate(self.opPacket, self.hook)(*args, **kwargs)
+        return AtenOPTemplate(self.op_packet, self.hook)(*args, **kwargs)
 
     def overloads(self):
-        return self.opPacket.overloads()
+        return self.op_packet.overloads()
 
 
 def wrap_aten_op(op, hook):

msprobe/pytorch/hook_module/wrap_distributed.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2023. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, 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
 #
@@ -13,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""
 
 import os
 from functools import wraps
@@ -23,6 +21,7 @@ from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.common.utils import torch_device_guard
 from msprobe.core.common.const import Const
 from msprobe.core.common.file_utils import load_yaml
+from msprobe.core.common.inplace_op_checker import InplaceOpChecker
 
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
@@ -50,7 +49,7 @@ class DistributedOPTemplate(HOOKModule):
         self.op_name_ = op_name
         self.prefix_op_name_ = "Distributed" + Const.SEP + str(op_name) + Const.SEP
         super().__init__(build_hook)
-        if not self.stop_hook and self.op_name_ in Const.INPLACE_LIST:
+        if not self.stop_hook and InplaceOpChecker.check(self.op_name_, InplaceOpChecker.OP_DISTRIBUTED):
             self.op_is_inplace = True
 
     @torch_device_guard