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

msprobe/mindspore/dump/hook_cell/api_registry.py

@@ -1,4 +1,5 @@
-# Copyright 2024 Huawei Technologies Co., Ltd
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
+# All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,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.
-# ============================================================================
 
 from mindspore import Tensor, ops, mint
 from mindspore.mint.nn import functional
@@ -20,8 +20,21 @@ from mindspore.communication import comm_func
 
 from msprobe.mindspore.dump.hook_cell.wrap_api import (HOOKTensor, HOOKStubTensor, HOOKFunctionalOP,
                                                        HOOKMintOP, HOOKMintNNFunctionalOP, HOOKDistributedOP,
-                                                       get_wrap_api_list, setup_hooks)
+                                                       HOOKTorchOP, HOOKTorchTensor, HOOKTorchFunctionalOP,
+                                                       HOOKTorchDistributedOP, HOOKTorchNpuOP,
+                                                       get_wrap_api_list, get_wrap_torch_api_list, setup_hooks)
 from msprobe.core.common.utils import Const
+from msprobe.mindspore.common.utils import is_mindtorch
+
+if is_mindtorch():
+    import torch
+    import torch_npu
+
+
+def stub_method(method):
+    def wrapped_method(*args, **kwargs):
+        return method(*args, **kwargs)
+    return wrapped_method
 
 
 class ApiRegistry:
@@ -34,6 +47,12 @@ class ApiRegistry:
         self.distributed_ori_attr = {}
         self.norm_inner_ops_ori_attr = {}
 
+        self.torch_ori_attr = {}
+        self.torch_tensor_ori_attr = {}
+        self.torch_functional_ori_attr = {}
+        self.torch_distributed_ori_attr = {}
+        self.torch_npu_ori_attr = {}
+
         self.tensor_hook_attr = {}
         self.stub_tensor_hook_attr = {}
         self.functional_hook_attr = {}
@@ -42,6 +61,12 @@ class ApiRegistry:
         self.distibuted_hook_attr = {}
         self.norm_inner_ops_hook_attr = {}
 
+        self.torch_hook_attr = {}
+        self.torch_tensor_hook_attr = {}
+        self.torch_functional_hook_attr = {}
+        self.torch_distributed_hook_attr = {}
+        self.torch_npu_hook_attr = {}
+
         self.norm_inner_ops = ["norm", "square", "sqrt", "is_complex"]
 
     @staticmethod
@@ -50,9 +75,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):
@@ -72,22 +101,71 @@ class ApiRegistry:
         self.set_api_attr(ops, self.norm_inner_ops_ori_attr)
 
     def api_set_hook_func(self):
-        self.set_api_attr(Tensor, self.tensor_hook_attr)
-        self.set_api_attr(StubTensor, self.stub_tensor_hook_attr)
-        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)
+        if is_mindtorch():
+            self.set_api_attr(torch, self.torch_hook_attr)
+            self.set_api_attr(torch.Tensor, self.torch_tensor_hook_attr)
+            self.set_api_attr(torch.nn.functional, self.torch_functional_hook_attr)
+            self.set_api_attr(torch.distributed, self.torch_distributed_hook_attr)
+            self.set_api_attr(torch_npu, self.torch_npu_hook_attr)
+        else:
+            self.set_api_attr(Tensor, self.tensor_hook_attr)
+            self.set_api_attr(StubTensor, self.stub_tensor_hook_attr)
+            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)
-        self.set_api_attr(StubTensor, self.stub_tensor_ori_attr)
-        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)
+        if is_mindtorch():
+            self.set_api_attr(torch, self.torch_ori_attr)
+            self.set_api_attr(torch.Tensor, self.torch_tensor_ori_attr)
+            self.set_api_attr(torch.nn.functional, self.torch_functional_ori_attr)
+            self.set_api_attr(torch.distributed, self.torch_distributed_ori_attr)
+            self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
+        else:
+            self.set_api_attr(Tensor, self.tensor_ori_attr)
+            self.set_api_attr(StubTensor, self.stub_tensor_ori_attr)
+            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):
+        setup_hooks(hook)
+        if is_mindtorch():
+            wrap_torch_api_name = get_wrap_torch_api_list()
+            self.store_ori_attr(torch,
+                                wrap_torch_api_name.torch_api_names, self.torch_ori_attr)
+            self.store_ori_attr(torch.Tensor,
+                                wrap_torch_api_name.tensor_api_names, self.torch_tensor_ori_attr)
+            self.store_ori_attr(torch.nn.functional,
+                                wrap_torch_api_name.functional_api_names, self.torch_functional_ori_attr)
+            self.store_ori_attr(torch.distributed,
+                                wrap_torch_api_name.distributed_api_names, self.torch_distributed_ori_attr)
+            self.store_ori_attr(torch_npu,
+                                wrap_torch_api_name.npu_api_names, self.torch_npu_ori_attr)
+            for attr_name in dir(HOOKTorchOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_hook_attr[api_name] = getattr(HOOKTorchOP, attr_name)
+            for attr_name in dir(HOOKTorchTensor):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_tensor_hook_attr[api_name] = getattr(HOOKTorchTensor, attr_name)
+            for attr_name in dir(HOOKTorchFunctionalOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_functional_hook_attr[api_name] = getattr(HOOKTorchFunctionalOP, attr_name)
+            for attr_name in dir(HOOKTorchDistributedOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_distributed_hook_attr[api_name] = getattr(HOOKTorchDistributedOP, attr_name)
+            for attr_name in dir(HOOKTorchNpuOP):
+                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
+                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
+                    self.torch_npu_hook_attr[api_name] = getattr(HOOKTorchNpuOP, attr_name)
+            return
+
         wrap_api_name = get_wrap_api_list()
         self.store_ori_attr(Tensor, wrap_api_name.tensor_api_names, self.tensor_ori_attr)
         self.store_ori_attr(StubTensor, wrap_api_name.stub_tensor_api_names, self.stub_tensor_ori_attr)
@@ -96,7 +174,6 @@ class ApiRegistry:
         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):
             if attr_name.startswith(Const.ATTR_NAME_PREFIX):
                 api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]

msprobe/mindspore/dump/hook_cell/hook_cell.py

@@ -1,4 +1,5 @@
-# Copyright 2024 Huawei Technologies Co., Ltd
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
+# All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,45 +12,66 @@
 # 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 collections import defaultdict
 
 from mindspore import nn
 
-from msprobe.core.common.const import Const
-
-
-class HOOKCell(nn.Cell):
-    cell_count = defaultdict(int)
-    g_stop_hook = False
-
-    def __init__(self, build_hook) -> None:
-        super(HOOKCell, self).__init__()
-        self.changed_status = False
-        self.input_kwargs = {}
-        self.prefix = ""
-        if not HOOKCell.g_stop_hook:
-            HOOKCell.g_stop_hook = True
-            self.changed_status = True
-            if hasattr(self, "prefix_api_name"):
-                self.prefix = self.prefix_api_name
-
-            HOOKCell.cell_count[self.prefix] += 1
-            self.prefix = self.prefix + str(HOOKCell.cell_count[self.prefix] - 1) + Const.SEP
-            forward_hook, backward_hook = build_hook(self.prefix)
-            self.register_forward_hook(forward_hook)
-            self.register_backward_hook(backward_hook)
-
-    # 重载call，加全局标志。
-    def __call__(self, *args, **kwargs):
-        try:
-            self.input_kwargs = kwargs
-            out = super(HOOKCell, self).__call__(*args, **kwargs)
-        except Exception as e:
-            raise e
-        finally:
-            if self.changed_status:
-                self.changed_status = False
-                HOOKCell.g_stop_hook = False
-        return out
+from msprobe.mindspore.common.utils import is_mindtorch, register_backward_hook_functions
+
+
+def add_cell_count(name):
+    HOOKCell.cell_count[name] += 1
+
+
+def get_cell_count(name):
+    return HOOKCell.cell_count[name]
+
+
+def __init__(self, build_hook) -> None:
+    super(HOOKCell, self).__init__()
+    self.changed_status = False
+    self.input_kwargs = {}
+    self.prefix = ""
+    if not HOOKCell.g_stop_hook:
+        HOOKCell.g_stop_hook = True
+        self.changed_status = True
+        if hasattr(self, "prefix_api_name"):
+            self.prefix = self.prefix_api_name
+
+        self.forward_data_collected = False
+        forward_pre_hook, forward_hook, backward_hook, backward_pre_hook = build_hook(self.prefix)
+        self.register_forward_pre_hook(forward_pre_hook)
+        self.register_forward_hook(forward_hook)
+        register_backward_hook_functions["full"](self, backward_hook)
+        register_backward_hook_functions["pre"](self, backward_pre_hook)
+
+
+# 重载call，加全局标志。
+def __call__(self, *args, **kwargs):
+    try:
+        self.input_kwargs = kwargs
+        out = super(HOOKCell, self).__call__(*args, **kwargs)
+    except Exception as e:
+        raise e
+    finally:
+        if self.changed_status:
+            self.changed_status = False
+            HOOKCell.g_stop_hook = False
+    return out
+
+
+hook_cell_dict = {
+    "cell_count": defaultdict(int),
+    "g_stop_hook": False,
+    "add_cell_count": staticmethod(add_cell_count),
+    "get_cell_count": staticmethod(get_cell_count),
+    "__init__": __init__,
+    "__call__": __call__
+}
+
+if is_mindtorch():
+    import torch
+    HOOKCell = type("HOOKCell", (torch.nn.Module,), hook_cell_dict)
+else:
+    HOOKCell = type("HOOKCell", (nn.Cell,), hook_cell_dict)

msprobe/mindspore/dump/hook_cell/primitive_hooks.py

@@ -1,4 +1,5 @@
-# Copyright 2024 Huawei Technologies Co., Ltd
+# 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.
@@ -11,18 +12,16 @@
 # 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 mindspore.common.tensor import Tensor
 
-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
+from msprobe.core.data_dump.data_processor.base import (ModuleBackwardInputs, ModuleBackwardOutputs,
+                                                        ModuleForwardInputsOutputs)
+from msprobe.mindspore.common.log import logger
 
 
 class PrimitiveHookService:
@@ -41,6 +40,7 @@ class PrimitiveHookService:
         Returns:
             callable: 包装后的 primitive 函数。
         """
+
         def create_backward_hook(captured_grads, num_tensors, updated_primitive_name, hook_type):
             """
             创建反向 hook 函数，用于捕获梯度。
@@ -54,26 +54,24 @@ class PrimitiveHookService:
             Returns:
                 callable: 反向 hook 函数。
             """
-            def backward_hook(grad):
 
-                captured_grads.append(grad)
+            def backward_hook(grad):
+                captured_grads.extend(grad)
                 backward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.BACKWARD}"
 
                 try:
-                    if len(captured_grads) == num_tensors and hook_type == Const.INPUT:
+                    if 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:
+                    elif 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}, "
@@ -104,7 +102,7 @@ class PrimitiveHookService:
                     hooked_inputs.append(arg_hooked)
                 else:
                     hooked_inputs.append(arg)
-            return hooked_inputs
+            return tuple(hooked_inputs)
 
         def hook_primitive_outputs(out, captured_grads_output, updated_primitive_name):
             """
@@ -137,6 +135,34 @@ class PrimitiveHookService:
                 return tuple(hooked_outputs)
             return out
 
+        def pre_forward_hook(primitive_name, primitive_instance, args, kwargs):
+            module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
+            try:
+                self.service_instance.data_collector.forward_input_data_collect(
+                    primitive_name,
+                    primitive_instance,
+                    os.getpid(),
+                    module_input_output
+                )
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during forward input data collection: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+
+        def post_forward_hook(primitive_name, primitive_instance, args, kwargs, output):
+            module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+            try:
+                self.service_instance.data_collector.forward_output_data_collect(
+                    primitive_name,
+                    primitive_instance,
+                    os.getpid(),
+                    module_input_output
+                )
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during forward output data collection: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+
         def wrapped_primitive_call(instance_self, *args, **kwargs):
             """
             包装后的 primitive 调用函数，添加输入和输出的 hook。
@@ -165,27 +191,17 @@ class PrimitiveHookService:
                              f"primitive_name: {primitive_name}")
                 raise DumpException(DumpException.INPUT_HOOK_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)
+
+            pre_forward_hook(forward_primitive_name, instance_self, hooked_inputs, kwargs)
             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()
+            post_forward_hook(forward_primitive_name, instance_self, hooked_inputs, kwargs, out)
 
             try:
                 out = hook_primitive_outputs(out, captured_grads_output, updated_primitive_name)
@@ -203,4 +219,3 @@ class PrimitiveHookService:
             self.primitive_counters[primitive_name] = 0
         else:
             self.primitive_counters[primitive_name] += 1
-

msprobe/mindspore/dump/hook_cell/support_wrap_ops.yaml

@@ -15,7 +15,7 @@
 
 # List of ops that register hooks
 
-  
+
 ops:
   - adaptive_avg_pool1d
   - adaptive_avg_pool2d
@@ -85,6 +85,7 @@ ops:
   - relu6
   - celu
   - rrelu
+  - rms_norm
   - selu
   - sigmoid
   - silu
@@ -490,6 +491,31 @@ ops:
   - scatter_update
   - derivative
   - jet
+  - row_stack
+  - gather
+  - arange
+  - cond
+  - slice_scatter
+  - clip_by_norm
+  - eps
+  - layer_norm
+  - cast
+  - numel
+  - permute
+  - select_scatter
+  - group_norm
+  - eq
+  - embedding
+  - ones_like
+  - zeros
+  - nanmean
+  - shape
+  - zeros_like
+  - ones
+  - diagonal_scatter
+  - vander
+  - is_nonzero
+  - rotary_position_embedding
 
 tensor:
   - __abs__
@@ -528,6 +554,7 @@ tensor:
   - acos
   - acosh
   - add
+  - add_
   - addbmm
   - addcdiv
   - addcmul
@@ -582,6 +609,7 @@ tensor:
   - diff
   - digamma
   - div
+  - div_
   - divide
   - equal
   - erf
@@ -714,6 +742,8 @@ tensor:
   - square
   - squeeze
   - std
+  - sub
+  - sub_
   - subtract
   - subtract
   - svd
@@ -958,6 +988,7 @@ mint.nn.functional:
   - one_hot_ext
   - pad
   - relu
+  - relu_
   - sigmoid
   - silu
   - softmax
@@ -992,3 +1023,7 @@ communication.comm_func:
   - broadcast
   - gather_into_tensor
   - scatter_tensor
+  - send
+  - recv
+  - isend
+  - irecv

msprobe/mindspore/dump/hook_cell/wrap_api.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -23,10 +23,16 @@ from mindspore.mint.nn import functional
 from msprobe.core.common.const import Const
 from msprobe.core.common.file_utils import load_yaml
 from msprobe.mindspore.common.const import Const as MsConst
+from msprobe.mindspore.common.utils import is_mindtorch
 from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 
+if is_mindtorch():
+    import torch
+    import torch_npu
+
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, MsConst.SUPPORTED_API_LIST_FILE)
+torch_yaml_path = os.path.join(cur_path, "../../../pytorch/hook_module", MsConst.SUPPORTED_API_LIST_FILE)
 
 
 class HOOKTensor(object):
@@ -53,6 +59,26 @@ class HOOKDistributedOP(object):
     pass
 
 
+class HOOKTorchOP(object):
+    pass
+
+
+class HOOKTorchTensor(object):
+    pass
+
+
+class HOOKTorchFunctionalOP(object):
+    pass
+
+
+class HOOKTorchDistributedOP(object):
+    pass
+
+
+class HOOKTorchNpuOP(object):
+    pass
+
+
 class ApiTemplate(HOOKCell):
     def __init__(self, api_name, api_dict, prefix, hook):
         self.api_name = api_name
@@ -60,7 +86,30 @@ class ApiTemplate(HOOKCell):
         self.prefix_api_name = prefix + str(api_name.split(Const.SEP)[-1]) + Const.SEP
         super().__init__(hook)
 
+    @staticmethod
+    def async_to_sync(output):
+        # Fake handle, used to return after the CommHandle executes the wait method
+        fake_handle = type("FakeHandle", (), {"wait": lambda self: None})()
+        if isinstance(output, tuple) and len(output) == 2 and hasattr(output[1], "wait"):
+            output[1].wait()
+            output = (output[0], fake_handle)
+        elif hasattr(output, "wait"):
+            output.wait()
+            output = fake_handle
+        return output
+
     def construct(self, *args, **kwargs):
+        if self.api_name.startswith(MsConst.DROPOUT_API_NAME_PREFIX):
+            return args[0] if args else kwargs.get(Const.INPUT)
+        
+        output = self.api_func(*args, **kwargs)
+
+        if self.prefix_api_name.startswith(MsConst.DISTRIBUTED_DATA_PREFIX):
+            if kwargs.get("async_op") or self.api_name in ["isend", "irecv"]:
+                output = self.async_to_sync(output)
+        return output
+
+    def forward(self, *args, **kwargs):
         if self.api_name.startswith(MsConst.DROPOUT_API_NAME_PREFIX):
             return args[0] if args else kwargs.get(Const.INPUT)
         return self.api_func(*args, **kwargs)
@@ -77,6 +126,15 @@ class WrapApiName:
         self.distributed_api_names = distributed_api_names
 
 
+class WrapTorchApiName:
+    def __init__(self, torch_api_names, tensor_api_names, functional_api_names, distributed_api_names, npu_api_names):
+        self.torch_api_names = torch_api_names
+        self.tensor_api_names = tensor_api_names
+        self.functional_api_names = functional_api_names
+        self.distributed_api_names = distributed_api_names
+        self.npu_api_names = npu_api_names
+
+
 def get_wrap_api_list():
     api_list = load_yaml(yaml_path)
     tensor_api = api_list.get(MsConst.SUPPORTED_TENSOR_LIST_KEY)
@@ -93,6 +151,21 @@ def get_wrap_api_list():
     return wrap_api_name
 
 
+def get_wrap_torch_api_list():
+    api_list = load_yaml(torch_yaml_path)
+    torch_api = api_list.get("torch")
+    tensor_api = api_list.get("tensor")
+    functional_api = api_list.get("functional")
+    distributed_api = api_list.get("distributed")
+    npu_api = api_list.get("torch_npu")
+    wrap_api_name = WrapTorchApiName(set(torch_api) & set(dir(torch)),
+                                     set(tensor_api) & set(dir(torch.Tensor)),
+                                     set(functional_api) & set(dir(torch.nn.functional)),
+                                     set(distributed_api) & set(dir(torch.distributed)),
+                                     set(npu_api) & set(dir(torch_npu)))
+    return wrap_api_name
+
+
 def wrap_api_func(api_name, api_dict, prefix, hook):
     def api_function(*args, **kwargs):
         return ApiTemplate(api_name, api_dict, prefix, hook)(*args, **kwargs)
@@ -106,6 +179,24 @@ def wrap_api_func_and_bind(api_list, api_dict, prefix, hook, hook_class):
 
 
 def setup_hooks(hook):
+    if is_mindtorch():
+        torch_wrap_api_name = get_wrap_torch_api_list()
+        wrap_api_func_and_bind(torch_wrap_api_name.torch_api_names,
+                               {f: getattr(torch, f) for f in dir(torch)},
+                               MsConst.TORCH_DATA_PREFIX, hook, HOOKTorchOP)
+        wrap_api_func_and_bind(torch_wrap_api_name.tensor_api_names,
+                               {f: getattr(torch.Tensor, f) for f in dir(torch.Tensor)},
+                               MsConst.TENSOR_DATA_PREFIX, hook, HOOKTorchTensor)
+        wrap_api_func_and_bind(torch_wrap_api_name.functional_api_names,
+                               {f: getattr(torch.nn.functional, f) for f in dir(torch.nn.functional)},
+                               MsConst.OPS_DATA_PREFIX, hook, HOOKTorchFunctionalOP)
+        wrap_api_func_and_bind(torch_wrap_api_name.distributed_api_names,
+                               {f: getattr(torch.distributed, f) for f in dir(torch.distributed)},
+                               MsConst.DISTRIBUTED_DATA_PREFIX, hook, HOOKTorchDistributedOP)
+        wrap_api_func_and_bind(torch_wrap_api_name.npu_api_names, {f: getattr(torch_npu, f) for f in dir(torch_npu)},
+                               MsConst.TORCH_NPU_DATA_PREFIX, hook, HOOKTorchNpuOP)
+        return
+
     wrap_api_name = get_wrap_api_list()
     wrap_api_func_and_bind(wrap_api_name.tensor_api_names, {f: getattr(Tensor, f) for f in dir(Tensor)},
                            MsConst.TENSOR_DATA_PREFIX, hook, HOOKTensor)