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

msprobe/mindspore/compare/ms_graph_compare.py

@@ -1,12 +1,27 @@
+# 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 copy
-import csv
 import glob
 import os
+import re
 
 import numpy as np
 import pandas as pd
-from msprobe.core.common.const import CompareConst, GraphMode, Const, FileCheckConst
-from msprobe.core.common.file_utils import FileOpen, check_path_before_create, change_mode, load_npy
+from msprobe.core.common.const import CompareConst, GraphMode, Const
+from msprobe.core.common.file_utils import load_npy, read_csv, save_excel
 from msprobe.core.common.log import logger
 from msprobe.core.common.utils import add_time_with_xlsx, CompareException
 from msprobe.core.compare.multiprocessing_compute import _ms_graph_handle_multi_process, check_accuracy
@@ -14,7 +29,7 @@ from msprobe.core.compare.npy_compare import npy_data_check, statistics_data_che
 from msprobe.mindspore.common.utils import convert_to_int, list_lowest_level_directories
 
 
-class row_data:
+class RowData:
     def __init__(self, mode):
         self.basic_data = copy.deepcopy(CompareConst.MS_GRAPH_BASE)
         self.npy_data = copy.deepcopy(CompareConst.MS_GRAPH_NPY)
@@ -28,17 +43,34 @@ class row_data:
         return self.data
 
 
+def get_name_dict(name: str) -> dict:
+    compare_pattern = re.compile(r'^([^.]+)\.([^.]+)\.([^.]+)\.([^.]+)\.(\d+(?:\.\d+)*)\.'
+                                 r'((?:in|out)put(?:\.\d+)*)\.([^.]+)\.([^.]+)\.npy$')
+    match = compare_pattern.match(name)
+    if match:
+        return {'op_type': match.group(1),
+                'op_name': match.group(2),
+                'task_id': match.group(3),
+                'stream_id': match.group(4),
+                'timestamp': match.group(5).split(Const.SEP)[0],
+                'input_output_index': match.group(6),
+                'slot': match.group(7),
+                'format': match.group(8)}
+    return {}
+
+
 def npy_data_read(data_path, npy_file_list, mapping_dict):
     data_list = []
+    compare_key_elements = ['op_name', 'task_id', 'input_output_index', 'slot']
     for data in npy_file_list:
         if data in mapping_dict:
-            split_list = mapping_dict[data].split(Const.SEP)
+            name_dict = get_name_dict(mapping_dict[data])
         else:
-            split_list = data.split(Const.SEP)
-        if len(split_list) < 7:
+            name_dict = get_name_dict(data)
+        if not name_dict:
             continue
-        compare_key = f"{split_list[1]}.{split_list[2]}.{split_list[3]}.{split_list[5]}.{split_list[6]}"
-        timestamp = convert_to_int(split_list[4])
+        compare_key = Const.SEP.join([name_dict.get(element) for element in compare_key_elements])
+        timestamp = convert_to_int(name_dict.get('timestamp'))
 
         data_list.append([os.path.join(data_path, data), compare_key, timestamp])
     return data_list
@@ -47,17 +79,18 @@ def npy_data_read(data_path, npy_file_list, mapping_dict):
 def statistic_data_read(statistic_file_list, statistic_file_path):
     data_list = []
     statistic_data_list = []
-    header_index = {'Data Type': None, 'Shape': None, 'Max Value': None, 'Min Value': None,
-                    'Avg Value': None, 'L2Norm Value': None}
+    header_index = {
+        'Data Type': None, 'Shape': None, 'Max Value': None,
+        'Min Value': None, 'Avg Value': None, 'L2Norm Value': None
+    }
     for statistic_file in statistic_file_list:
-        with FileOpen(statistic_file, "r") as f:
-            csv_reader = csv.reader(f, delimiter=",")
-            header = next(csv_reader)
-            for key in header_index.keys():
-                for index, value in enumerate(header):
-                    if key == value:
-                        header_index[key] = index
-            statistic_data_list.extend([row for row in csv_reader])
+        content = read_csv(statistic_file, as_pd=False)
+        header = content[0]
+        for key in header_index.keys():
+            for index, value in enumerate(header):
+                if key == value:
+                    header_index[key] = index
+        statistic_data_list.extend(content[1:])
 
     for key in header_index.keys():
         if header_index[key] is None:
@@ -65,8 +98,9 @@ def statistic_data_read(statistic_file_list, statistic_file_path):
 
     for data in statistic_data_list:
         compare_key = f"{data[1]}.{data[2]}.{data[3]}.{data[5]}"
+        op_name = f"{compare_key} {statistic_file_path}"
         timestamp = int(data[4])
-        result_data = [statistic_file_path, compare_key, timestamp]
+        result_data = [op_name, compare_key, timestamp]
         for key in header_index.keys():
             if header_index[key] is None:
                 result_data.append(np.nan)
@@ -94,11 +128,9 @@ def generate_data_name(data_path):
     mapping_dict = {}
     if mapping_exist:
         for mapping_file in mapping_file_list:
-            with FileOpen(mapping_file, "r") as f:
-                csv_reader = csv.reader(f, delimiter=",")
-                header = next(csv_reader)
-                for row in csv_reader:
-                    mapping_dict[row[0]] = row[1]
+            content = read_csv(mapping_file, False)
+            for row in content[1:]:
+                mapping_dict[row[0]] = row[1]
 
     if npy_exist:
         data_list = npy_data_read(data_path, npy_file_list, mapping_dict)
@@ -133,7 +165,7 @@ class GraphMSComparator:
     def compare_ops(compare_result_db, mode):
 
         def npy_mode_compute(row):
-            result_dict = row_data(GraphMode.NPY_MODE)()
+            result_dict = RowData(GraphMode.NPY_MODE)()
 
             def process_npy_file(file_path, name_prefix, result):
                 if os.path.exists(file_path):
@@ -168,7 +200,7 @@ class GraphMSComparator:
             return pd.Series(result_dict)
 
         def statistic_mode_compute(row):
-            result_dict = row_data('STATISTIC')()
+            result_dict = RowData('STATISTIC')()
 
             def update_result_dict(result, rows, prefix):
                 result[f'{prefix} Name'] = rows[f'{prefix} Name']
@@ -195,24 +227,30 @@ class GraphMSComparator:
                     result_dict[CompareConst.NPU_NORM] - result_dict[CompareConst.BENCH_NORM])
                 result_dict[CompareConst.MAX_RELATIVE_ERR] = result_dict[CompareConst.MAX_DIFF] / result_dict[
                     CompareConst.BENCH_MAX] if result_dict[CompareConst.BENCH_MAX] > 0 else 0
-                result_dict[CompareConst.MAX_RELATIVE_ERR] = str(result_dict[CompareConst.MAX_RELATIVE_ERR] * 100) + "%"
+                if not np.isnan(result_dict[CompareConst.MAX_RELATIVE_ERR]):
+                    result_dict[CompareConst.MAX_RELATIVE_ERR] = str(
+                        result_dict[CompareConst.MAX_RELATIVE_ERR] * 100) + "%"
                 result_dict[CompareConst.MIN_RELATIVE_ERR] = result_dict[CompareConst.MIN_DIFF] / result_dict[
                     CompareConst.BENCH_MIN] if result_dict[CompareConst.BENCH_MIN] > 0 else 0
-                result_dict[CompareConst.MIN_RELATIVE_ERR] = str(result_dict[CompareConst.MIN_RELATIVE_ERR] * 100) + "%"
+                if not np.isnan(result_dict[CompareConst.MIN_RELATIVE_ERR]):
+                    result_dict[CompareConst.MIN_RELATIVE_ERR] = \
+                        str(result_dict[CompareConst.MIN_RELATIVE_ERR] * 100) + "%"
                 result_dict[CompareConst.MEAN_RELATIVE_ERR] = result_dict[CompareConst.MEAN_DIFF] / result_dict[
                     CompareConst.BENCH_MEAN] if result_dict[CompareConst.BENCH_MEAN] > 0 else 0
-                result_dict[CompareConst.MEAN_RELATIVE_ERR] = str(
-                    result_dict[CompareConst.MEAN_RELATIVE_ERR] * 100) + "%"
+                if not np.isnan(result_dict[CompareConst.MEAN_RELATIVE_ERR]):
+                    result_dict[CompareConst.MEAN_RELATIVE_ERR] = str(
+                        result_dict[CompareConst.MEAN_RELATIVE_ERR] * 100) + "%"
                 result_dict[CompareConst.NORM_RELATIVE_ERR] = result_dict[CompareConst.NORM_DIFF] / result_dict[
                     CompareConst.BENCH_NORM] if result_dict[CompareConst.BENCH_NORM] > 0 else 0
-                result_dict[CompareConst.NORM_RELATIVE_ERR] = str(
-                    result_dict[CompareConst.NORM_RELATIVE_ERR] * 100) + "%"
+                if not np.isnan(result_dict[CompareConst.NORM_RELATIVE_ERR]):
+                    result_dict[CompareConst.NORM_RELATIVE_ERR] = str(
+                        result_dict[CompareConst.NORM_RELATIVE_ERR] * 100) + "%"
                 magnitude_diff = result_dict[CompareConst.MAX_DIFF] / (
                         max(result_dict[CompareConst.NPU_MAX], result_dict[CompareConst.BENCH_MAX]) + 1e-10)
-                if magnitude_diff > CompareConst.MAGNITUDE:
-                    result_dict[CompareConst.ACCURACY] = 'No'
-                else:
-                    result_dict[CompareConst.ACCURACY] = 'Yes'
+                if np.isnan(result_dict[CompareConst.NPU_MAX]) and np.isnan(result_dict[CompareConst.BENCH_MAX]):
+                    magnitude_diff = 0
+                result_dict[CompareConst.ACCURACY] = CompareConst.YES if \
+                    magnitude_diff <= CompareConst.MAGNITUDE else CompareConst.NO
 
             return pd.Series(result_dict)
 
@@ -235,14 +273,24 @@ class GraphMSComparator:
                 is_empty = True
             if is_empty or not mode:
                 continue
-            compare_result_df = self._do_multi_process(compare_result_df, mode)
+            compare_result_df = self.do_multi_process(compare_result_df, mode)
             compare_result_name = add_time_with_xlsx(f"compare_result_{str(rank_id)}_{str(step_id)}")
             compare_result_path = os.path.join(os.path.realpath(self.output_path), f"{compare_result_name}")
-            check_path_before_create(compare_result_path)
-            compare_result_df.to_excel(compare_result_path, index=False)
-            change_mode(compare_result_path, FileCheckConst.DATA_FILE_AUTHORITY)
+            self.to_excel(compare_result_df, compare_result_path)
             logger.info(f"Compare rank: {rank_id} step: {step_id} finish. Compare result: {compare_result_path}.")
 
+    def to_excel(self, compare_result_df: pd.DataFrame, compare_result_path: str, slice_num=0, need_slice=False) -> int:
+        size = len(compare_result_df)
+        # sheet size cannot be larger than 1048576
+        if size < CompareConst.MAX_EXCEL_LENGTH:
+            compare_result_path = compare_result_path.replace('.xlsx', f'_slice_{slice_num}.xlsx') if \
+                need_slice else compare_result_path
+            save_excel(compare_result_path, compare_result_df)
+            return slice_num + 1
+        else:
+            slice_num = self.to_excel(compare_result_df.iloc[0: size // 2], compare_result_path, slice_num, True)
+            return self.to_excel(compare_result_df.iloc[size // 2:], compare_result_path, slice_num, True)
+
     def compare_process(self, rank_id, step_id):
         # generate data_path
         npu_data_path_list = self.npu_rank_step_dict.get((rank_id, step_id))
@@ -251,8 +299,8 @@ class GraphMSComparator:
             return [], ''
 
         # generate file name
-        npu_mode = 'ERROR_MODE'
-        bench_mode = 'ERROR_MODE'
+        npu_mode = GraphMode.ERROR_MODE
+        bench_mode = GraphMode.ERROR_MODE
         npu_data_list = []
         bench_data_list = []
         for npu_data_path in npu_data_path_list:
@@ -262,7 +310,7 @@ class GraphMSComparator:
             bench_mode, data_list = generate_data_name(bench_data_path)
             bench_data_list.extend(data_list)
 
-        if npu_mode == "ERROR_MODE" or bench_mode == "ERROR_MODE":
+        if npu_mode == GraphMode.ERROR_MODE or bench_mode == GraphMode.ERROR_MODE:
             logger.warning(f"Data_path {npu_data_path} or {bench_data_path} is not exist.")
             return [], ''
         if npu_mode != bench_mode:
@@ -286,11 +334,13 @@ class GraphMSComparator:
                                                   CompareConst.BENCH_NORM])
 
             npu_float_type = [CompareConst.NPU_MAX, CompareConst.NPU_MIN, CompareConst.NPU_MEAN, CompareConst.NPU_NORM]
-            npu_data_df[npu_float_type] = npu_data_df[npu_float_type].astype(np.float32)
+            npu_data_df[npu_float_type] = npu_data_df[npu_float_type].astype(float)
 
-            bench_float_type = [CompareConst.BENCH_MAX, CompareConst.BENCH_MIN, CompareConst.BENCH_MEAN,
-                                CompareConst.BENCH_NORM]
-            bench_data_df[bench_float_type] = bench_data_df[bench_float_type].astype(np.float32)
+            bench_float_type = [
+                CompareConst.BENCH_MAX, CompareConst.BENCH_MIN,
+                CompareConst.BENCH_MEAN, CompareConst.BENCH_NORM
+            ]
+            bench_data_df[bench_float_type] = bench_data_df[bench_float_type].astype(float)
 
         npu_data_df['Local Index'] = npu_data_df.sort_values('TimeStamp').groupby('Compare Key').cumcount()
         bench_data_df['Local Index'] = bench_data_df.sort_values('TimeStamp').groupby('Compare Key').cumcount()
@@ -339,7 +389,7 @@ class GraphMSComparator:
                 rank_step_path_dict[rank_step_key] = [dir_path]
         return dict(sorted(rank_step_path_dict.items()))
 
-    def _do_multi_process(self, result_df, mode):
+    def do_multi_process(self, result_df, mode):
         try:
             result_df = _ms_graph_handle_multi_process(self.compare_ops, result_df, mode)
         except ValueError as e:

msprobe/mindspore/debugger/debugger_config.py

@@ -1,9 +1,24 @@
+# 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 os
 
 from msprobe.core.common.const import Const
+from msprobe.core.common.file_utils import create_directory
 from msprobe.mindspore.common.const import Const as MsConst
 from msprobe.mindspore.common.const import FreeBenchmarkConst
-from msprobe.core.common.file_utils import create_directory
 
 
 class DebuggerConfig:
@@ -18,7 +33,7 @@ class DebuggerConfig:
         self.level_ori = common_config.level
         self.list = [] if not task_config.list else task_config.list
         self.scope = [] if not task_config.scope else task_config.scope
-        self.data_mode = [] if not task_config.data_mode else task_config.data_mode
+        self.data_mode = [Const.ALL] if not task_config.data_mode else task_config.data_mode
         self.file_format = task_config.file_format
         self.overflow_nums = 1 if not task_config.overflow_nums else task_config.overflow_nums
         self.check_mode = task_config.check_mode
@@ -37,6 +52,9 @@ class DebuggerConfig:
                self.pert_type != FreeBenchmarkConst.DEFAULT_PERT_TYPE:
                 raise ValueError("pert_mode must be improve_precision or empty when handler_type is fix, "
                                  f"but got {self.pert_type}.")
+            if self.stage == Const.BACKWARD and self.handler_type == FreeBenchmarkConst.FIX:
+                raise ValueError("handler_type must be check or empty when fuzz_stage is backward, "
+                                 f"but got {self.handler_type}.")
             self.dump_level = FreeBenchmarkConst.DEFAULT_DUMP_LEVEL
 
     def check(self):
@@ -51,16 +69,4 @@ class DebuggerConfig:
             self.file_format = "npy"
         if not self.check_mode:
             self.check_mode = "all"
-        self._check_rank()
-        self._check_step()
         return True
-
-    def _check_rank(self):
-        for rank_id in self.rank:
-            if not isinstance(rank_id, int) or rank_id < 0:
-                raise ValueError(f"rank {self.rank} must be a positive integer.")
-
-    def _check_step(self):
-        for s in self.step:
-            if not isinstance(s, int) or s < 0:
-                raise ValueError(f"step element {s} must be a positive integer.")

msprobe/mindspore/debugger/precision_debugger.py

@@ -1,17 +1,34 @@
+# 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 os
+from collections import defaultdict
 
 import mindspore as ms
 from mindspore._c_expression import MSContext
 
-from msprobe.mindspore.service import Service
-from msprobe.mindspore.ms_config import parse_json_config
-from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
-from msprobe.mindspore.task_handler_factory import TaskHandlerFactory
-from msprobe.core.common.const import Const
+from msprobe.core.common.const import Const, MsgConst
+from msprobe.mindspore.cell_processor import CellProcessor
 from msprobe.mindspore.common.const import Const as MsConst
-from msprobe.mindspore.runtime import Runtime
-
+from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
+from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 from msprobe.mindspore.grad_probe.grad_monitor import GradientMonitor
+from msprobe.mindspore.ms_config import parse_json_config
+from msprobe.mindspore.runtime import Runtime
+from msprobe.mindspore.service import Service
+from msprobe.mindspore.task_handler_factory import TaskHandlerFactory
 
 
 class PrecisionDebugger:
@@ -65,11 +82,11 @@ class PrecisionDebugger:
     def start(cls, model=None):
         instance = cls._instance
         if not instance:
-            raise Exception("No instance of PrecisionDebugger found.")
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
         if instance.task in PrecisionDebugger.task_not_need_service:
             return
 
-        instance.config.execution_mode = instance._get_execution_mode()
+        instance.config.execution_mode = cls._get_execution_mode()
         if cls._need_service():
             if not instance.service:
                 instance.service = Service(instance.config)
@@ -82,11 +99,21 @@ class PrecisionDebugger:
         instance.first_start = True
         Runtime.is_running = True
 
+    @classmethod
+    def forward_backward_dump_end(cls):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        if instance.task in PrecisionDebugger.task_not_need_service:
+            return
+        if instance.service:
+            instance.service.forward_backward_dump_end()
+
     @classmethod
     def stop(cls):
         instance = cls._instance
         if not instance:
-            raise Exception("PrecisionDebugger instance is not created.")
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
         if instance.task == Const.GRAD_PROBE:
             instance.gm.stop()
         if instance.task in PrecisionDebugger.task_not_need_service:
@@ -99,18 +126,21 @@ class PrecisionDebugger:
     def step(cls):
         instance = cls._instance
         if not instance:
-            raise Exception("PrecisionDebugger instance is not created.")
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
         if instance.task in PrecisionDebugger.task_not_need_service:
             return
         if instance.service:
             instance.service.step()
+        HOOKCell.cell_count = defaultdict(int)
+        CellProcessor.reset_cell_stats()
+
         Runtime.step_count += 1
 
     @classmethod
     def monitor(cls, opt):
         instance = cls._instance
         if not instance:
-            raise Exception("PrecisionDebugger instance is not created.")
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
         if instance.task != Const.GRAD_PROBE:
             return
         instance.gm.monitor(opt)
@@ -119,7 +149,7 @@ class PrecisionDebugger:
     def _need_service(cls):
         instance = cls._instance
         if not instance:
-            raise Exception("No instance of PrecisionDebugger found.")
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
         if instance.config.execution_mode != MsConst.PYNATIVE_MODE:
             return False
         else:

msprobe/mindspore/dump/dump_tool_factory.py

@@ -1,7 +1,22 @@
+# 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.mindspore.common.const import Const
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
-from msprobe.mindspore.dump.kernel_kbyk_dump import KernelKbykDump
 from msprobe.mindspore.dump.kernel_graph_dump import KernelGraphDump
+from msprobe.mindspore.dump.kernel_kbyk_dump import KernelKbykDump
 
 
 class DumpToolFactory:
@@ -25,6 +40,8 @@ class DumpToolFactory:
 
     @staticmethod
     def create(config: DebuggerConfig):
+        if len(config.data_mode) != 1 or config.data_mode[0] not in Const.GRAPH_DATA_MODE_LIST:
+            raise Exception("data_mode must be one of all, input, output.")
         tool = DumpToolFactory.tools.get(config.level)
         if not tool:
             raise Exception("Valid level is needed.")

msprobe/mindspore/dump/hook_cell/api_registry.py

@@ -16,13 +16,20 @@
 from mindspore import Tensor, ops, mint
 from mindspore.mint.nn import functional
 from mindspore.common._stub_tensor import StubTensor
+from mindspore.communication import comm_func
 
 from msprobe.mindspore.dump.hook_cell.wrap_api import (HOOKTensor, HOOKStubTensor, HOOKFunctionalOP,
-                                                       HOOKMintOP, HOOKMintNNFunctionalOP,
+                                                       HOOKMintOP, HOOKMintNNFunctionalOP, HOOKDistributedOP,
                                                        get_wrap_api_list, setup_hooks)
 from msprobe.core.common.utils import Const
 
 
+def stub_method(method):
+    def wrapped_method(*args, **kwargs):
+        return method(*args, **kwargs)
+    return wrapped_method
+
+
 class ApiRegistry:
     def __init__(self):
         self.tensor_ori_attr = {}
@@ -30,6 +37,7 @@ class ApiRegistry:
         self.functional_ori_attr = {}
         self.mint_ops_ori_attr = {}
         self.mint_func_ops_ori_attr = {}
+        self.distributed_ori_attr = {}
         self.norm_inner_ops_ori_attr = {}
 
         self.tensor_hook_attr = {}
@@ -37,6 +45,7 @@ class ApiRegistry:
         self.functional_hook_attr = {}
         self.mint_ops_hook_attr = {}
         self.mint_func_ops_hook_attr = {}
+        self.distibuted_hook_attr = {}
         self.norm_inner_ops_hook_attr = {}
 
         self.norm_inner_ops = ["norm", "square", "sqrt", "is_complex"]
@@ -47,9 +56,13 @@ class ApiRegistry:
             if Const.SEP in api:
                 sub_module_name, sub_op = api.rsplit(Const.SEP, 1)
                 sub_module = getattr(ori_api_group, sub_module_name)
-                api_ori_attr[api] = getattr(sub_module, sub_op)
+                ori_api_func = getattr(sub_module, sub_op)
             else:
-                api_ori_attr[api] = getattr(ori_api_group, api)
+                ori_api_func = getattr(ori_api_group, api)
+            if ori_api_group == StubTensor:
+                api_ori_attr[api] = stub_method(ori_api_func)
+                continue
+            api_ori_attr[api] = ori_api_func
 
     @staticmethod
     def set_api_attr(api_group, attr_dict):
@@ -74,6 +87,7 @@ class ApiRegistry:
         self.set_api_attr(ops, self.functional_hook_attr)
         self.set_api_attr(mint, self.mint_ops_hook_attr)
         self.set_api_attr(functional, self.mint_func_ops_hook_attr)
+        self.set_api_attr(comm_func, self.distibuted_hook_attr)
 
     def api_set_ori_func(self):
         self.set_api_attr(Tensor, self.tensor_ori_attr)
@@ -81,6 +95,7 @@ class ApiRegistry:
         self.set_api_attr(ops, self.functional_ori_attr)
         self.set_api_attr(mint, self.mint_ops_ori_attr)
         self.set_api_attr(functional, self.mint_func_ops_ori_attr)
+        self.set_api_attr(comm_func, self.distributed_ori_attr)
 
     def initialize_hook(self, hook):
         wrap_api_name = get_wrap_api_list()
@@ -89,6 +104,7 @@ class ApiRegistry:
         self.store_ori_attr(ops, wrap_api_name.ops_api_names, self.functional_ori_attr)
         self.store_ori_attr(mint, wrap_api_name.mint_api_names, self.mint_ops_ori_attr)
         self.store_ori_attr(functional, wrap_api_name.mint_nn_func_api_names, self.mint_func_ops_ori_attr)
+        self.store_ori_attr(comm_func, wrap_api_name.distributed_api_names, self.distributed_ori_attr)
         self.store_ori_attr(ops, self.norm_inner_ops, self.norm_inner_ops_ori_attr)
         setup_hooks(hook)
         for attr_name in dir(HOOKTensor):
@@ -113,6 +129,10 @@ class ApiRegistry:
             if attr_name.startswith(Const.ATTR_NAME_PREFIX):
                 api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
                 self.mint_func_ops_hook_attr[api_name] = getattr(HOOKMintNNFunctionalOP, attr_name)
+        for attr_name in dir(HOOKDistributedOP):
+            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                self.distibuted_hook_attr[api_name] = getattr(HOOKDistributedOP, attr_name)
 
 
 api_register = ApiRegistry()