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

msprobe/pytorch/functional/dump_module.py

@@ -1,39 +1,39 @@
-import torch.nn as nn
-from msprobe.pytorch.common.log import logger
-from msprobe.core.common.const import Const
-from msprobe.pytorch.hook_module.api_registry import api_register
-from msprobe.pytorch.debugger.precision_debugger import PrecisionDebugger
-from msprobe.core.common.exceptions import MsaccException
-from msprobe.core.data_dump.scope import BaseScope
-
-module_count = {}
-
-
-def module_dump(module, dump_name):
-    if not isinstance(module, nn.Module):
-        logger.error("The parameter:module in module_dump is not a Module subclass.")
-        raise MsaccException(MsaccException.INVALID_PARAM_ERROR)
-    if not isinstance(dump_name, str):
-        logger.error("The parameter:dump_name in module_dump is not a str type.")
-        raise MsaccException(MsaccException.INVALID_PARAM_ERROR)
-    api_register.api_originality()
-    if dump_name not in module_count:
-        module_count[dump_name] = 0
-    else:
-        module_count[dump_name] += 1
-    dump_name = dump_name + Const.SEP + str(module_count.get(dump_name)) + Const.SEP
-
-    pdg = PrecisionDebugger()
-    _, forward_hook, backward_hook = pdg.service.build_hook(BaseScope.Module_Type_Module, dump_name)
-    module.register_forward_hook(forward_hook, with_kwargs=True)
-    module.register_full_backward_hook(backward_hook)
-
-    module.register_forward_pre_hook(pdg.service.module_processor.node_hook(dump_name + Const.FORWARD, Const.START))
-    module.register_forward_hook(pdg.service.module_processor.node_hook(dump_name + Const.FORWARD, Const.STOP))
-    module.register_full_backward_pre_hook(
-        pdg.service.module_processor.node_hook(dump_name + Const.BACKWARD, Const.START))
-    module.register_full_backward_hook(pdg.service.module_processor.node_hook(dump_name + Const.BACKWARD, Const.STOP))
-
-
-def module_dump_end():
-    api_register.api_modularity()
+import torch.nn as nn
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.const import Const
+from msprobe.pytorch.hook_module.api_registry import api_register
+from msprobe.pytorch.debugger.precision_debugger import PrecisionDebugger
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.data_dump.scope import BaseScope
+
+module_count = {}
+
+
+def module_dump(module, dump_name):
+    if not isinstance(module, nn.Module):
+        logger.error("The parameter:module in module_dump is not a Module subclass.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if not isinstance(dump_name, str):
+        logger.error("The parameter:dump_name in module_dump is not a str type.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    api_register.api_originality()
+    if dump_name not in module_count:
+        module_count[dump_name] = 0
+    else:
+        module_count[dump_name] += 1
+    dump_name = dump_name + Const.SEP + str(module_count.get(dump_name)) + Const.SEP
+
+    pdg = PrecisionDebugger()
+    _, forward_hook, backward_hook, _ = pdg.service.build_hook(BaseScope.Module_Type_Module, dump_name)
+    module.register_forward_hook(forward_hook, with_kwargs=True)
+    module.register_full_backward_hook(backward_hook)
+
+    module.register_forward_pre_hook(pdg.service.module_processor.node_hook(dump_name + Const.FORWARD, Const.START))
+    module.register_forward_hook(pdg.service.module_processor.node_hook(dump_name + Const.FORWARD, Const.STOP))
+    module.register_full_backward_pre_hook(
+        pdg.service.module_processor.node_hook(dump_name + Const.BACKWARD, Const.START))
+    module.register_full_backward_hook(pdg.service.module_processor.node_hook(dump_name + Const.BACKWARD, Const.STOP))
+
+
+def module_dump_end():
+    api_register.api_modularity()

msprobe/pytorch/grad_probe/grad_monitor.py

@@ -0,0 +1,91 @@
+import os
+from collections import defaultdict
+
+import torch
+if int(torch.__version__.split('.')[0]) >= 2:
+    from torch.optim.optimizer import register_optimizer_step_pre_hook
+from msprobe.pytorch.grad_probe.grad_stat_csv import GradStatCsv
+from msprobe.core.grad_probe.utils import check_numeral_list_ascend, data_in_list_target
+from msprobe.core.grad_probe.constant import level_adp
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.file_utils import remove_path, save_npy, write_csv, create_directory
+from msprobe.pytorch.common.utils import get_rank_id, print_rank_0
+
+
+class GradientMonitor:
+
+    def __init__(self, common_config, task_config):
+        level = task_config.grad_level
+        if level not in level_adp:
+            raise Exception(f"level is valid, not in {level_adp.keys()}")
+        self._level_adp = level_adp[level]
+        self._param_list = task_config.param_list
+        self._target_ranks = common_config.rank
+        logger.info(f"target rank {self._target_ranks}")
+        self._target_step = common_config.step
+        logger.info(f"target step {self._target_step}")
+        self._bounds = task_config.bounds
+        check_numeral_list_ascend(self._bounds)
+        self._output_path = common_config.dump_path
+        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")
+        self._step = -1
+        self._param2name = defaultdict(str)
+
+    @property
+    def output_path(self):
+        return self._output_path
+
+    @staticmethod
+    def save_grad_direction(param_name, grad, save_path):
+        if not os.path.exists(save_path):
+            create_directory(save_path)
+        param_grad = grad.clone().detach()
+        is_positive = param_grad > 0
+        save_filepath = os.path.join(save_path, f"{param_name}.npy")
+        save_npy(is_positive.cpu().numpy(), save_filepath)
+
+    def monitor(self, model):
+        print_rank_0("> parameter names:")
+        for name, param in model.named_parameters():
+            self._param2name[param] = name
+            print_rank_0(f"\t{name}")
+        setattr(self, "_rank", get_rank_id())
+        if torch.distributed.is_initialized() and not data_in_list_target(getattr(self, "_rank"), self._target_ranks):
+            return
+        self._hook_optimizer()
+
+    def _hook_optimizer(self):
+        def optimizer_pre_step_hook(optimizer, args, kargs):
+            self._step += 1
+            logger.info(f"grad_probe: optimizer step {self._step}")
+            if not data_in_list_target(self._step, self._target_step):
+                return
+            output_lines = []
+            for param, param_name in self._param2name.items():
+                if not data_in_list_target(param_name, self._param_list):
+                    continue
+                grad = param.main_grad if hasattr(param, "main_grad") else param.grad
+                if grad is None:
+                    logger.info(f"grad is None: {param_name}")
+                    continue
+                grad_info = GradStatCsv.generate_csv_line(param_name, self._level_adp, grad, self._bounds)
+                output_lines.append(grad_info)
+                if self._level_adp["have_grad_direction"]:
+                    GradientMonitor.save_grad_direction(param_name, grad,
+                                                    f'{self._output_path}/rank{self._rank}/step{self._step}')
+            output_dirpath = os.path.join(self._output_path, f"rank{getattr(self, '_rank')}")
+            if not os.path.isdir(output_dirpath):
+                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")
+                remove_path(output_path)
+            header_result = GradStatCsv.generate_csv_header(self._level_adp, self._bounds)
+            output_lines.insert(0, header_result)
+            write_csv(output_lines, output_path)
+            logger.info(f"write grad data to {output_path}")
+        if int(torch.__version__.split('.')[0]) >= 2:
+            register_optimizer_step_pre_hook(optimizer_pre_step_hook)

msprobe/pytorch/grad_probe/grad_stat_csv.py

@@ -0,0 +1,129 @@
+from abc import ABC, abstractmethod
+from collections import namedtuple
+import hashlib
+import torch
+from msprobe.core.grad_probe.constant import GradConst
+
+CSV_header_input = namedtuple("CSV_header_input", ["bounds"])
+CSV_content_input = namedtuple("CSV_content_input", ["grad", "bounds"])
+
+
+class GradStatCsv:
+    csv = {}
+
+    @staticmethod
+    def generate_csv_header(level, bounds):
+        header = ["param_name"]
+        for key in level["header"]:
+            csv_header_input = CSV_header_input(bounds=bounds)
+            header.extend(GradStatCsv.csv[key].generate_csv_header(csv_header_input))
+        return header
+
+    @staticmethod
+    def generate_csv_line(param_name, level, grad, bounds):
+        line = [param_name]
+        for key in level["header"]:
+            csv_content_input = CSV_content_input(grad=grad, bounds=bounds)
+            line.extend(GradStatCsv.csv[key].generate_csv_content(csv_content_input))
+        return line
+
+
+def register_csv_item(key, cls=None):
+    if cls is None:
+        # 无参数时，返回装饰器函数
+        return lambda cls: register_csv_item(key, cls)
+    GradStatCsv.csv[key] = cls
+    return cls
+
+
+class CsvItem(ABC):
+    @abstractmethod
+    def generate_csv_header(csv_header_input):
+        pass
+
+    @abstractmethod
+    def generate_csv_content(csv_content_input):
+        pass
+
+
+@register_csv_item(GradConst.MD5)
+class CSV_md5(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["MD5"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        tensor_bytes = grad.cpu().detach().float().numpy().tobytes()
+        md5_hash = hashlib.md5(tensor_bytes)
+        return [md5_hash.hexdigest()]
+
+
+@register_csv_item(GradConst.DISTRIBUTION)
+class CSV_distribution(CsvItem):
+    def generate_csv_header(csv_header_input):
+        bounds = csv_header_input.bounds
+        intervals = []
+        if bounds:
+            intervals.append(f"(-inf, {bounds[0]}]")
+            for i in range(1, len(bounds)):
+                intervals.append(f"({bounds[i-1]}, {bounds[i]}]")
+        if intervals:
+            intervals.append(f"({bounds[-1]}, inf)")
+        intervals.append("=0")
+    
+        return intervals
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        bounds = csv_content_input.bounds
+        grad = grad.cpu().detach()
+        if grad.dtype == torch.bfloat16:
+            grad = grad.to(torch.float32)
+        element_num = grad.numel()
+        grad_equal_0_num = (grad == 0).sum().item()
+        bound = torch.Tensor(bounds)
+        bucketsize_result = torch.bucketize(grad, bound)
+        interval_nums = [(bucketsize_result == i).sum().item() for i in range(len(bound) + 1)]
+        interval_nums.append(grad_equal_0_num)
+        return_list = [x / element_num if element_num != 0 else 0 for x in interval_nums]
+        return return_list
+
+
+@register_csv_item(GradConst.MAX)
+class CSV_max(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["max"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [torch.max(grad).cpu().detach().float().numpy().tolist()]
+
+
+@register_csv_item(GradConst.MIN)
+class CSV_max(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["min"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [torch.min(grad).cpu().detach().float().numpy().tolist()]
+
+
+@register_csv_item(GradConst.NORM)
+class CSV_max(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["norm"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [torch.norm(grad).cpu().detach().float().numpy().tolist()]
+    
+
+@register_csv_item(GradConst.SHAPE)
+class CSV_shape(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["shape"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [list(grad.shape)]

msprobe/pytorch/hook_module/api_registry.py

@@ -1,161 +1,161 @@
-#!/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");
-# 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.distributed as dist
-
-from msprobe.pytorch.hook_module import wrap_torch, wrap_functional, wrap_tensor, wrap_vf, wrap_distributed, wrap_aten
-from msprobe.pytorch.hook_module.wrap_aten import get_aten_ops
-from msprobe.pytorch.hook_module.wrap_distributed import get_distributed_ops
-from msprobe.pytorch.hook_module.wrap_functional import get_functional_ops
-from msprobe.pytorch.hook_module.wrap_tensor import get_tensor_ops
-from msprobe.pytorch.hook_module.wrap_torch import get_torch_ops
-from msprobe.pytorch.hook_module.wrap_vf import get_vf_ops
-from msprobe.pytorch.common.utils import torch_without_guard_version, npu_distributed_api, is_gpu
-from msprobe.core.common.const import Const
-
-torch_version_above_2 = torch.__version__.split('+')[0] > '2.0'
-
-if not is_gpu:
-    import torch_npu
-    from . import wrap_npu_custom
-    from .wrap_npu_custom import get_npu_ops
-
-
-class ApiRegistry:
-    def __init__(self):
-        self.tensor_ori_attr = {}
-        self.torch_ori_attr = {}
-        self.functional_ori_attr = {}
-        self.distributed_ori_attr = {}
-        self.npu_distributed_ori_attr = {}
-        self.vf_ori_attr = {}
-        self.aten_ori_attr = {}
-        self.torch_npu_ori_attr = {}
-
-        self.tensor_hook_attr = {}
-        self.torch_hook_attr = {}
-        self.functional_hook_attr = {}
-        self.distributed_hook_attr = {}
-        self.npu_distributed_hook_attr = {}
-        self.vf_hook_attr = {}
-        self.aten_hook_attr = {}
-        self.torch_npu_hook_attr = {}
-
-    @staticmethod
-    def store_ori_attr(ori_api_group, api_list, api_ori_attr):
-        for api in api_list:
-            if '.' in api:
-                sub_module_name, sub_op = api.rsplit('.', 1)
-                sub_module = getattr(ori_api_group, sub_module_name)
-                api_ori_attr[api] = getattr(sub_module, sub_op)
-            else:
-                api_ori_attr[api] = getattr(ori_api_group, api)
-
-    @staticmethod
-    def set_api_attr(api_group, attr_dict):
-        for api, api_attr in attr_dict.items():
-            if '.' in api:
-                sub_module_name, sub_op = api.rsplit('.', 1)
-                sub_module = getattr(api_group, sub_module_name, None)
-                if sub_module is not None:
-                    setattr(sub_module, sub_op, api_attr)
-            else:
-                setattr(api_group, api, api_attr)
-
-    def api_modularity(self):
-        self.set_api_attr(torch.Tensor, self.tensor_hook_attr)
-        self.set_api_attr(torch, self.torch_hook_attr)
-        self.set_api_attr(torch.nn.functional, self.functional_hook_attr)
-        self.set_api_attr(dist, self.distributed_hook_attr)
-        self.set_api_attr(dist.distributed_c10d, self.distributed_hook_attr)
-        if not is_gpu and not torch_without_guard_version:
-            self.set_api_attr(torch_npu.distributed, self.npu_distributed_hook_attr)
-            self.set_api_attr(torch_npu.distributed.distributed_c10d, self.npu_distributed_hook_attr)
-        if torch_version_above_2:
-            self.set_api_attr(torch.ops.aten, self.aten_hook_attr)
-        self.set_api_attr(torch._VF, self.vf_hook_attr)
-        if not is_gpu:
-            self.set_api_attr(torch_npu, self.torch_npu_hook_attr)
-
-    def api_originality(self):
-        self.set_api_attr(torch.Tensor, self.tensor_ori_attr)
-        self.set_api_attr(torch, self.torch_ori_attr)
-        self.set_api_attr(torch.nn.functional, self.functional_ori_attr)
-        self.set_api_attr(dist, self.distributed_ori_attr)
-        self.set_api_attr(dist.distributed_c10d, self.distributed_ori_attr)
-        if not is_gpu and not torch_without_guard_version:
-            self.set_api_attr(torch_npu.distributed, self.npu_distributed_ori_attr)
-            self.set_api_attr(torch_npu.distributed.distributed_c10d, self.npu_distributed_ori_attr)
-        if torch_version_above_2:
-            self.set_api_attr(torch.ops.aten, self.aten_ori_attr)
-        self.set_api_attr(torch._VF, self.vf_ori_attr)
-        if not is_gpu:
-            self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
-
-    def initialize_hook(self, hook):
-        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):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.tensor_hook_attr[attr_name[5:]] = getattr(wrap_tensor.HOOKTensor, attr_name)
-
-        self.store_ori_attr(torch, get_torch_ops(), self.torch_ori_attr)
-        wrap_torch.wrap_torch_ops_and_bind(hook)
-        for attr_name in dir(wrap_torch.HOOKTorchOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.torch_hook_attr[attr_name[5:]] = getattr(wrap_torch.HOOKTorchOP, attr_name)
-
-        self.store_ori_attr(torch.nn.functional, get_functional_ops(), self.functional_ori_attr)
-        wrap_functional.wrap_functional_ops_and_bind(hook)
-        for attr_name in dir(wrap_functional.HOOKFunctionalOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.functional_hook_attr[attr_name[5:]] = getattr(wrap_functional.HOOKFunctionalOP, attr_name)
-
-        self.store_ori_attr(dist, get_distributed_ops(), self.distributed_ori_attr)
-        wrap_distributed.wrap_distributed_ops_and_bind(hook)
-        if not is_gpu and not torch_without_guard_version:
-            self.store_ori_attr(torch_npu.distributed, npu_distributed_api, self.npu_distributed_ori_attr)
-        for attr_name in dir(wrap_distributed.HOOKDistributedOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP, attr_name)
-                if not is_gpu and not torch_without_guard_version and attr_name[5:] in npu_distributed_api:
-                    self.npu_distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP,
-                                                                            attr_name)
-
-        if torch_version_above_2:
-            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):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    self.aten_hook_attr[attr_name[5:]] = getattr(wrap_aten.HOOKAtenOP, attr_name)
-
-        self.store_ori_attr(torch._VF, get_vf_ops(), self.vf_ori_attr)
-        wrap_vf.wrap_vf_ops_and_bind(hook)
-        for attr_name in dir(wrap_vf.HOOKVfOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.vf_hook_attr[attr_name[5:]] = getattr(wrap_vf.HOOKVfOP, attr_name)
-
-        if not is_gpu:
-            self.store_ori_attr(torch_npu, get_npu_ops(), self.torch_npu_ori_attr)
-            wrap_npu_custom.wrap_npu_ops_and_bind(hook)
-            for attr_name in dir(wrap_npu_custom.HOOKNpuOP):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    self.torch_npu_hook_attr[attr_name[5:]] = getattr(wrap_npu_custom.HOOKNpuOP, attr_name)
-
-
-api_register = ApiRegistry()
+#!/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");
+# 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.distributed as dist
+
+from msprobe.pytorch.hook_module import wrap_torch, wrap_functional, wrap_tensor, wrap_vf, wrap_distributed, wrap_aten
+from msprobe.pytorch.hook_module.wrap_aten import get_aten_ops
+from msprobe.pytorch.hook_module.wrap_distributed import get_distributed_ops
+from msprobe.pytorch.hook_module.wrap_functional import get_functional_ops
+from msprobe.pytorch.hook_module.wrap_tensor import get_tensor_ops
+from msprobe.pytorch.hook_module.wrap_torch import get_torch_ops
+from msprobe.pytorch.hook_module.wrap_vf import get_vf_ops
+from msprobe.pytorch.common.utils import torch_without_guard_version, npu_distributed_api, is_gpu
+from msprobe.core.common.const import Const
+
+torch_version_above_2 = torch.__version__.split('+')[0] > '2.0'
+
+if not is_gpu:
+    import torch_npu
+    from . import wrap_npu_custom
+    from .wrap_npu_custom import get_npu_ops
+
+
+class ApiRegistry:
+    def __init__(self):
+        self.tensor_ori_attr = {}
+        self.torch_ori_attr = {}
+        self.functional_ori_attr = {}
+        self.distributed_ori_attr = {}
+        self.npu_distributed_ori_attr = {}
+        self.vf_ori_attr = {}
+        self.aten_ori_attr = {}
+        self.torch_npu_ori_attr = {}
+
+        self.tensor_hook_attr = {}
+        self.torch_hook_attr = {}
+        self.functional_hook_attr = {}
+        self.distributed_hook_attr = {}
+        self.npu_distributed_hook_attr = {}
+        self.vf_hook_attr = {}
+        self.aten_hook_attr = {}
+        self.torch_npu_hook_attr = {}
+
+    @staticmethod
+    def store_ori_attr(ori_api_group, api_list, api_ori_attr):
+        for api in api_list:
+            if '.' in api:
+                sub_module_name, sub_op = api.rsplit('.', 1)
+                sub_module = getattr(ori_api_group, sub_module_name)
+                api_ori_attr[api] = getattr(sub_module, sub_op)
+            else:
+                api_ori_attr[api] = getattr(ori_api_group, api)
+
+    @staticmethod
+    def set_api_attr(api_group, attr_dict):
+        for api, api_attr in attr_dict.items():
+            if '.' in api:
+                sub_module_name, sub_op = api.rsplit('.', 1)
+                sub_module = getattr(api_group, sub_module_name, None)
+                if sub_module is not None:
+                    setattr(sub_module, sub_op, api_attr)
+            else:
+                setattr(api_group, api, api_attr)
+
+    def api_modularity(self):
+        self.set_api_attr(torch.Tensor, self.tensor_hook_attr)
+        self.set_api_attr(torch, self.torch_hook_attr)
+        self.set_api_attr(torch.nn.functional, self.functional_hook_attr)
+        self.set_api_attr(dist, self.distributed_hook_attr)
+        self.set_api_attr(dist.distributed_c10d, self.distributed_hook_attr)
+        if not is_gpu and not torch_without_guard_version:
+            self.set_api_attr(torch_npu.distributed, self.npu_distributed_hook_attr)
+            self.set_api_attr(torch_npu.distributed.distributed_c10d, self.npu_distributed_hook_attr)
+        if torch_version_above_2:
+            self.set_api_attr(torch.ops.aten, self.aten_hook_attr)
+        self.set_api_attr(torch._VF, self.vf_hook_attr)
+        if not is_gpu:
+            self.set_api_attr(torch_npu, self.torch_npu_hook_attr)
+
+    def api_originality(self):
+        self.set_api_attr(torch.Tensor, self.tensor_ori_attr)
+        self.set_api_attr(torch, self.torch_ori_attr)
+        self.set_api_attr(torch.nn.functional, self.functional_ori_attr)
+        self.set_api_attr(dist, self.distributed_ori_attr)
+        self.set_api_attr(dist.distributed_c10d, self.distributed_ori_attr)
+        if not is_gpu and not torch_without_guard_version:
+            self.set_api_attr(torch_npu.distributed, self.npu_distributed_ori_attr)
+            self.set_api_attr(torch_npu.distributed.distributed_c10d, self.npu_distributed_ori_attr)
+        if torch_version_above_2:
+            self.set_api_attr(torch.ops.aten, self.aten_ori_attr)
+        self.set_api_attr(torch._VF, self.vf_ori_attr)
+        if not is_gpu:
+            self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
+
+    def initialize_hook(self, hook):
+        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):
+            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                self.tensor_hook_attr[attr_name[5:]] = getattr(wrap_tensor.HOOKTensor, attr_name)
+
+        self.store_ori_attr(torch, get_torch_ops(), self.torch_ori_attr)
+        wrap_torch.wrap_torch_ops_and_bind(hook)
+        for attr_name in dir(wrap_torch.HOOKTorchOP):
+            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                self.torch_hook_attr[attr_name[5:]] = getattr(wrap_torch.HOOKTorchOP, attr_name)
+
+        self.store_ori_attr(torch.nn.functional, get_functional_ops(), self.functional_ori_attr)
+        wrap_functional.wrap_functional_ops_and_bind(hook)
+        for attr_name in dir(wrap_functional.HOOKFunctionalOP):
+            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                self.functional_hook_attr[attr_name[5:]] = getattr(wrap_functional.HOOKFunctionalOP, attr_name)
+
+        self.store_ori_attr(dist, get_distributed_ops(), self.distributed_ori_attr)
+        wrap_distributed.wrap_distributed_ops_and_bind(hook)
+        if not is_gpu and not torch_without_guard_version:
+            self.store_ori_attr(torch_npu.distributed, npu_distributed_api, self.npu_distributed_ori_attr)
+        for attr_name in dir(wrap_distributed.HOOKDistributedOP):
+            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                self.distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP, attr_name)
+                if not is_gpu and not torch_without_guard_version and attr_name[5:] in npu_distributed_api:
+                    self.npu_distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP,
+                                                                            attr_name)
+
+        if torch_version_above_2:
+            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):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    self.aten_hook_attr[attr_name[5:]] = getattr(wrap_aten.HOOKAtenOP, attr_name)
+
+        self.store_ori_attr(torch._VF, get_vf_ops(), self.vf_ori_attr)
+        wrap_vf.wrap_vf_ops_and_bind(hook)
+        for attr_name in dir(wrap_vf.HOOKVfOP):
+            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                self.vf_hook_attr[attr_name[5:]] = getattr(wrap_vf.HOOKVfOP, attr_name)
+
+        if not is_gpu:
+            self.store_ori_attr(torch_npu, get_npu_ops(), self.torch_npu_ori_attr)
+            wrap_npu_custom.wrap_npu_ops_and_bind(hook)
+            for attr_name in dir(wrap_npu_custom.HOOKNpuOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    self.torch_npu_hook_attr[attr_name[5:]] = getattr(wrap_npu_custom.HOOKNpuOP, attr_name)
+
+
+api_register = ApiRegistry()