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

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:
@@ -51,16 +66,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,31 @@
+# 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
 
 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.common.const import Const as MsConst
-from msprobe.mindspore.runtime import Runtime
-
+from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
 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 +79,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 +96,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,7 +123,7 @@ 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:
@@ -110,7 +134,7 @@ class PrecisionDebugger:
     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 +143,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:

msprobe/mindspore/dump/hook_cell/api_registry.py

@@ -16,9 +16,10 @@
 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
 
@@ -30,6 +31,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 +39,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"]
@@ -74,6 +77,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 +85,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 +94,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 +119,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()

msprobe/mindspore/dump/hook_cell/primitive_hooks.py

@@ -0,0 +1,206 @@
+# Copyright 2024 Huawei Technologies Co., Ltd
+#
+# 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
+
+import mindspore as ms
+from mindspore.common.tensor import Tensor
+from mindspore import ops
+
+from msprobe.mindspore.common.log import logger
+from msprobe.core.common.utils import Const, DumpException
+from msprobe.core.data_dump.data_processor.base import ModuleBackwardInputsOutputs, ModuleForwardInputsOutputs, \
+    ModuleBackwardInputs, ModuleBackwardOutputs
+
+
+class PrimitiveHookService:
+    def __init__(self, service_instance):
+        self.primitive_counters = {}
+        self.service_instance = service_instance
+
+    def wrap_primitive(self, origin_func, primitive_name):
+        """
+        包装原始的 primitive 函数，添加输入和输出的 hook 以捕获前向和反向数据。
+
+        Args:
+            origin_func (callable): 原始 的 primitive 函数。
+            primitive_name (str): 原始的 primitive 名称。
+
+        Returns:
+            callable: 包装后的 primitive 函数。
+        """
+        def create_backward_hook(captured_grads, num_tensors, updated_primitive_name, hook_type):
+            """
+            创建反向 hook 函数，用于捕获梯度。
+
+            Args:
+                captured_grads (list): 用于保存捕获的梯度。
+                num_tensors (int): 张量数量。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+                hook_type (str): hook 类型 (输入/输出)。
+
+            Returns:
+                callable: 反向 hook 函数。
+            """
+            def backward_hook(grad):
+
+                captured_grads.append(grad)
+                backward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.BACKWARD}"
+
+                try:
+                    if len(captured_grads) == num_tensors and hook_type == Const.INPUT:
+                        self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
+                        new_module_input_output = ModuleBackwardOutputs(grad_output=tuple(captured_grads))
+                        self.service_instance.data_collector.backward_output_data_collect(
+                            backward_primitive_name, self, os.getpid(), new_module_input_output
+                        )
+                        captured_grads.clear()
+                    elif len(captured_grads) == num_tensors and hook_type == Const.OUTPUT:
+                        self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
+                        new_module_input_output = ModuleBackwardInputs(grad_input=tuple(captured_grads))
+                        self.service_instance.data_collector.backward_input_data_collect(
+                            backward_primitive_name, self, os.getpid(), new_module_input_output
+                        )
+                        captured_grads.clear()
+
+                except Exception as exception:
+                    logger.error(f"This is a primitive op {hook_type}_backward dump error: {exception}, "
+                                 f"updated_primitive_name: {updated_primitive_name}")
+                    raise DumpException(DumpException.BACKWARD_DATA_COLLECTION_ERROR) from exception
+
+            return backward_hook
+
+        def hook_primitive_inputs(args, captured_grads_input, updated_primitive_name):
+            """
+            针对前向输入添加 hook。
+
+            Args:
+                args (tuple): primitive 输入参数。
+                captured_grads_input (list): 捕获的输入梯度。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+
+            Returns:
+                list: 添加了 hook 的输入。
+            """
+            hooked_inputs = []
+            num_tensors = sum(isinstance(arg, Tensor) for arg in args)
+            input_backward_hook = create_backward_hook(captured_grads_input, num_tensors, updated_primitive_name,
+                                                       Const.INPUT)
+            for arg in args:
+                if isinstance(arg, Tensor):
+                    arg_hooked = ops.HookBackward(input_backward_hook)(arg)
+                    hooked_inputs.append(arg_hooked)
+                else:
+                    hooked_inputs.append(arg)
+            return hooked_inputs
+
+        def hook_primitive_outputs(out, captured_grads_output, updated_primitive_name):
+            """
+            针对前向输出添加 hook。
+
+            Args:
+                out (Tensor/tuple): primitive 输出。
+                captured_grads_output (list): 捕获的输出梯度。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+
+            Returns:
+                Tensor/tuple: 添加了 hook 的输出。
+            """
+            if isinstance(out, tuple):
+                num_output_tensors = sum(isinstance(tensor, Tensor) for tensor in out)
+            else:
+                num_output_tensors = 1
+            output_backward_hook = create_backward_hook(captured_grads_output, num_output_tensors,
+                                                        updated_primitive_name, Const.OUTPUT)
+
+            if isinstance(out, Tensor):
+                return ops.HookBackward(output_backward_hook)(out)
+            elif isinstance(out, tuple):
+                hooked_outputs = []
+                for tensor in out:
+                    if isinstance(tensor, Tensor):
+                        hooked_outputs.append(ops.HookBackward(output_backward_hook)(tensor))
+                    else:
+                        hooked_outputs.append(tensor)
+                return tuple(hooked_outputs)
+            return out
+
+        def wrapped_primitive_call(instance_self, *args, **kwargs):
+            """
+            包装后的 primitive 调用函数，添加输入和输出的 hook。
+
+            Args:
+                instance_self (object): primitive 的实例。
+                *args: primitive 输入参数。
+                **kwargs: primitive 关键字参数。
+
+            Returns:
+                Tensor/tuple: primitive 的返回值。
+            """
+            self.update_primitive_counters(primitive_name)
+            current_count = self.primitive_counters.get(primitive_name, 0)
+            updated_primitive_name = f"{Const.PRIMITIVE_PREFIX}{Const.SEP}{primitive_name}{Const.SEP}{current_count}"
+
+            if not self.service_instance.primitive_switch:
+                return origin_func(*args, **kwargs)
+
+            captured_grads_input, captured_grads_output = [], []
+
+            try:
+                hooked_inputs = hook_primitive_inputs(args, captured_grads_input, updated_primitive_name)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during input hooking: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.INPUT_HOOK_ERROR) from exception
+
+            try:
+                out = origin_func(*hooked_inputs, **kwargs)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during function call: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.FUNCTION_CALL_ERROR) from exception
+
+            forward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.FORWARD}"
+            self.service_instance.data_collector.update_api_or_module_name(forward_primitive_name)
+            if self.service_instance.data_collector:
+                module_input_output = ModuleForwardInputsOutputs(args=hooked_inputs, kwargs=kwargs, output=out)
+                try:
+                    self.service_instance.data_collector.forward_data_collect(forward_primitive_name, instance_self,
+                                                             os.getpid(), module_input_output)
+                except Exception as exception:
+                    logger.error(f"This is a primitive op dump error during forward data collection: {exception}, "
+                                 f"primitive_name: {primitive_name}")
+                    raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+
+                if self.service_instance.data_collector.if_return_forward_new_output():
+                    out = self.service_instance.data_collector.get_forward_new_output()
+
+            try:
+                out = hook_primitive_outputs(out, captured_grads_output, updated_primitive_name)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during output hooking: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.OUTPUT_HOOK_ERROR) from exception
+
+            return out
+
+        return wrapped_primitive_call
+
+    def update_primitive_counters(self, primitive_name):
+        if primitive_name not in self.primitive_counters:
+            self.primitive_counters[primitive_name] = 0
+        else:
+            self.primitive_counters[primitive_name] += 1
+

msprobe/mindspore/dump/hook_cell/support_wrap_ops.yaml

@@ -185,6 +185,7 @@ ops:
   - float_power
   - fmod
   - frac
+  - flash_attention_score
   - gcd
   - hypot
   - igamma
@@ -876,16 +877,60 @@ mint.ops:
   - zeros
   - zeros_ex
   - zeros_like
-
-mint.nn:
-  - Dropout
-  - Embedding
-  - Fold
-  - LayerNorm
-  - Linear
-  - MaxPool2d
-  - Unfold
-  - Upsample
+  - inverse
+  - select
+  - item
+  - unsqueeze
+  - median
+  - floor
+  - histc
+  - special
+  - arctan2
+  - sign
+  - concat
+  - atanh
+  - greater_equal
+  - eye
+  - fix
+  - argmin
+  - asinh
+  - atan
+  - nan_to_num
+  - tan
+  - round
+  - cosh
+  - norm
+  - roll
+  - log1p
+  - reshape
+  - arccos
+  - outer
+  - arcsin
+  - rand_like
+  - acosh
+  - multinomial
+  - logical_xor
+  - acos
+  - linalg
+  - sinc
+  - arcsinh
+  - asin
+  - narrow
+  - arctanh
+  - trace
+  - erfc
+  - bernoulli
+  - expm1
+  - logaddexp
+  - sinh
+  - arccosh
+  - atan2
+  - rand
+  - arange
+  - trunc
+  - arctan
+  - swapaxes
+  - transpose
 
 mint.nn.functional:
   - absolute_import
@@ -920,3 +965,30 @@ mint.nn.functional:
   - softplus
   - tanh
   - unfold
+  - mse_loss
+  - adaptive_avg_pool1d
+  - binary_cross_entropy
+  - adaptive_avg_pool2d
+  - hardsigmoid
+  - selu
+  - softshrink
+  - prelu
+  - logsigmoid
+  - hardswish
+  - mish
+  - log_softmax
+  - hardshrink
+  - l1_loss
+  - elu
+
+communication.comm_func:
+  - all_reduce
+  - all_gather_into_tensor
+  - reduce
+  - reduce_scatter_tensor
+  - all_to_all_single_with_output_shape
+  - all_to_all_with_output_shape
+  - batch_isend_irecv
+  - broadcast
+  - gather_into_tensor
+  - scatter_tensor

msprobe/mindspore/dump/hook_cell/wrap_api.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,19 +12,18 @@
 # 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 mindspore import Tensor, ops, mint
-from mindspore.mint.nn import functional
+from mindspore import Tensor, mint, ops
 from mindspore.common._stub_tensor import StubTensor
+from mindspore.communication import comm_func
+from mindspore.mint.nn import functional
 
-from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 from msprobe.core.common.const import Const
-from msprobe.mindspore.common.const import Const as MsConst
 from msprobe.core.common.file_utils import load_yaml
-
+from msprobe.mindspore.common.const import Const as MsConst
+from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, MsConst.SUPPORTED_API_LIST_FILE)
@@ -51,6 +49,10 @@ class HOOKMintNNFunctionalOP(object):
     pass
 
 
+class HOOKDistributedOP(object):
+    pass
+
+
 class ApiTemplate(HOOKCell):
     def __init__(self, api_name, api_dict, prefix, hook):
         self.api_name = api_name
@@ -65,12 +67,14 @@ class ApiTemplate(HOOKCell):
 
 
 class WrapApiName:
-    def __init__(self, tensor_api_names, stub_tensor_api_names, ops_api_names, mint_api_names, mint_nn_func_api_names):
+    def __init__(self, tensor_api_names, stub_tensor_api_names, ops_api_names, mint_api_names, mint_nn_func_api_names,
+                 distributed_api_names):
         self.tensor_api_names = tensor_api_names
         self.stub_tensor_api_names = stub_tensor_api_names
         self.ops_api_names = ops_api_names
         self.mint_api_names = mint_api_names
         self.mint_nn_func_api_names = mint_nn_func_api_names
+        self.distributed_api_names = distributed_api_names
 
 
 def get_wrap_api_list():
@@ -79,11 +83,13 @@ def get_wrap_api_list():
     ops_api = api_list.get(MsConst.SUPPORTED_OPS_LIST_KEY)
     mint_api = api_list.get(MsConst.SUPPORTED_MINT_LIST_KEY)
     mint_nn_func_api = api_list.get(MsConst.SUPPORTED__MINT_NN_FUNC_LIST_KEY)
+    distributed_api = api_list.get(MsConst.SUPPORTED_COMM_LIST_KEY)
     wrap_api_name = WrapApiName(set(tensor_api) & set(dir(Tensor)),
                                 set(tensor_api) & set(dir(StubTensor)),
                                 set(ops_api) & set(dir(ops)),
                                 set(mint_api) & set(dir(mint)),
-                                set(mint_nn_func_api) & set(dir(functional)))
+                                set(mint_nn_func_api) & set(dir(functional)),
+                                set(distributed_api) & set(dir(comm_func)))
     return wrap_api_name
 
 
@@ -111,3 +117,5 @@ def setup_hooks(hook):
                            MsConst.MINT_DATA_PREFIX, hook, HOOKMintOP)
     wrap_api_func_and_bind(wrap_api_name.mint_nn_func_api_names, {f: getattr(functional, f) for f in dir(functional)},
                            MsConst.MINT_NN_FUNC_DATA_PREFIX, hook, HOOKMintNNFunctionalOP)
+    wrap_api_func_and_bind(wrap_api_name.distributed_api_names, {f: getattr(comm_func, f) for f in dir(comm_func)},
+                           MsConst.DISTRIBUTED_DATA_PREFIX, hook, HOOKDistributedOP)