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

msprobe/pytorch/parse_tool/lib/parse_tool.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# 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.
 # You may obtain a copy of the License at
 #
@@ -13,17 +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 argparse
 import os
 from collections import namedtuple
 
+from msprobe.core.common.file_utils import create_directory
+from msprobe.pytorch.parse_tool.lib.compare import Compare
 from msprobe.pytorch.parse_tool.lib.config import Const
+from msprobe.pytorch.parse_tool.lib.parse_exception import catch_exception, ParseException
 from msprobe.pytorch.parse_tool.lib.utils import Util
-from msprobe.pytorch.parse_tool.lib.compare import Compare
 from msprobe.pytorch.parse_tool.lib.visualization import Visualization
-from msprobe.pytorch.parse_tool.lib.parse_exception import catch_exception, ParseException
-from msprobe.core.common.file_utils import create_directory
+
 
 class ParseTool:
     def __init__(self):
@@ -117,7 +117,8 @@ class ParseTool:
         self.util.check_path_valid(args.golden_dump_path)
         self.util.check_file_path_format(args.my_dump_path, Const.NPY_SUFFIX)
         self.util.check_file_path_format(args.golden_dump_path, Const.NPY_SUFFIX)
-        compare_data_args = namedtuple('compare_data_args', ['my_dump_path', 'golden_dump_path', 'save', 'rtol', 'atol', 'count'])
+        compare_data_args = namedtuple('compare_data_args',
+                                       ['my_dump_path', 'golden_dump_path', 'save', 'rtol', 'atol', 'count'])
         compare_data_args.__new__.__defaults__ = (False, 0.001, 0.001, 20)
         res = compare_data_args(args.my_dump_path, args.golden_dump_path, args.save, args.rtol, args.atol, args.count)
         self.compare.compare_data(res)

msprobe/pytorch/parse_tool/lib/utils.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# 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.
 # You may obtain a copy of the License at
 #
@@ -13,24 +12,24 @@
 # 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 hashlib
 import os
 import re
-import sys
 import subprocess
-import hashlib
+import sys
 import time
-import numpy as np
 from collections import namedtuple
-from msprobe.pytorch.parse_tool.lib.config import Const
-from msprobe.pytorch.parse_tool.lib.file_desc import DumpDecodeFileDesc, FileDesc
-from msprobe.pytorch.parse_tool.lib.parse_exception import ParseException
-from msprobe.core.common.file_utils import change_mode, check_other_user_writable,\
-    check_path_executable, check_path_owner_consistent
+
+import numpy as np
 from msprobe.core.common.const import FileCheckConst
+from msprobe.core.common.file_utils import change_mode, check_other_user_writable, \
+    check_path_executable, check_path_owner_consistent
 from msprobe.core.common.file_utils import check_file_or_directory_path, remove_path, check_file_type, os_walk_for_files
 from msprobe.pytorch.common.log import logger
-
+from msprobe.pytorch.parse_tool.lib.config import Const
+from msprobe.pytorch.parse_tool.lib.file_desc import DumpDecodeFileDesc, FileDesc
+from msprobe.pytorch.parse_tool.lib.parse_exception import ParseException
 
 try:
     from rich.traceback import install
@@ -135,7 +134,7 @@ class Util:
             zero_mask = (data == 0)
             data[zero_mask] += np.finfo(float).eps
         return data
-    
+
     @staticmethod
     def dir_contains_only(path, endfix):
         files = os_walk_for_files(path, Const.MAX_TRAVERSAL_DEPTH)
@@ -143,11 +142,11 @@ class Util:
             if not file['file'].endswith(endfix):
                 return False
         return True
-    
+
     @staticmethod
     def localtime_str():
         return time.strftime("%Y%m%d%H%M%S", time.localtime())
-    
+
     @staticmethod
     def change_filemode_safe(path):
         change_mode(path, FileCheckConst.DATA_FILE_AUTHORITY)
@@ -208,7 +207,7 @@ class Util:
 
     def list_numpy_files(self, path, extern_pattern=''):
         return self.list_file_with_pattern(path, Const.NUMPY_PATTERN, extern_pattern,
-                                            self._gen_numpy_file_info)
+                                           self._gen_numpy_file_info)
 
     def create_columns(self, content):
         if not Columns:

msprobe/pytorch/parse_tool/lib/visualization.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# 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.
 # You may obtain a copy of the License at
 #
@@ -13,14 +12,14 @@
 # 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 json
-import numpy as np
 
+import numpy as np
+from msprobe.core.common.file_utils import FileOpen, load_npy, save_npy_to_txt
 from msprobe.pytorch.parse_tool.lib.config import Const
-from msprobe.pytorch.parse_tool.lib.utils import Util
 from msprobe.pytorch.parse_tool.lib.parse_exception import ParseException
-from msprobe.core.common.file_utils import FileOpen, load_npy, save_npy_to_txt
+from msprobe.pytorch.parse_tool.lib.utils import Util
 
 
 class Visualization:
@@ -77,7 +76,7 @@ class Visualization:
                             self.util.log.info("  File \"{}\", line {}, in {}".format(item[0], item[1], item[2]))
                             self.util.log.info("    {}".format(item[3]))
                         continue
-                if len(msg) > 5 and len(msg[5]) >=  3:
+                if len(msg) > 5 and len(msg[5]) >= 3:
                     summery_info = "  [{}][dtype: {}][shape: {}][max: {}][min: {}][mean: {}]" \
                         .format(msg[0], msg[3], msg[4], msg[5][0], msg[5][1], msg[5][2])
                     if not title_printed:

msprobe/pytorch/service.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");
@@ -19,7 +19,7 @@ from collections import namedtuple
 
 import torch
 from msprobe.core.common.const import Const
-from msprobe.core.common.exceptions import DistributedNotInitializedError, MsprobeException
+from msprobe.core.common.exceptions import DistributedNotInitializedError
 from msprobe.core.common.file_utils import create_directory
 from msprobe.core.common.utils import print_tools_ends_info
 from msprobe.core.data_dump.data_collector import build_data_collector
@@ -29,10 +29,10 @@ from msprobe.pytorch.api_accuracy_checker.common.utils import ApiData
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.common.utils import get_rank_if_initialized
 from msprobe.pytorch.dump.kernel_dump.kernel_config import create_kernel_config_json
-from msprobe.pytorch.hook_module import remove_dropout
+from msprobe.pytorch.dump.module_dump.module_processer import ModuleProcesser
 from msprobe.pytorch.hook_module.api_registry import api_register
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
-from msprobe.pytorch.module_processer import ModuleProcesser
+from msprobe.pytorch.hook_module.register_optimizer_hook import register_optimizer_hook
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 if torch_version_above_or_equal_2:
@@ -48,100 +48,175 @@ class Service:
         self.data_collector = build_data_collector(config)
         self.module_processor = ModuleProcesser(self.data_collector.scope)
         self.switch = False
+        self.inner_switch = False
         self.current_iter = 0
         self.first_start = True
         self.current_rank = None
         self.dump_iter_dir = None
         self.should_stop_service = False
         self.attl = None
-
-    @staticmethod
-    def forward_backward_dump_end():
-        logger.info_on_rank_0("Data needed ends here.")
-        api_register.api_originality()
-
-    @staticmethod
-    def is_registered_backward_hook(module):
-        if hasattr(module, '_backward_hooks') and \
-                len(module._backward_hooks) > 0 and \
-                module._is_full_backward_hook is False:
-            return True
-        return False
-
-    def check_register_full_backward_hook(self, module):
-        if self.is_registered_backward_hook(module):
-            module._backward_hooks.clear()
-            module._is_full_backward_hook = None
-            logger.warning("Found deprecated backward hooks. Removing them and switching to full backward hooks.")
+        self.params_grad_info = {}
+        # 提前注册，确保注册尽可能多的API hook
+        self.register_api_hook()
 
     def build_hook(self, module_type, name):
         def pre_hook(api_or_module_name, module, args, kwargs):
-            if not self.should_execute_hook():
+            if not self.should_execute_hook(module_type, module, True):
                 return args, kwargs
 
+            self.inner_switch = True
             if module_type == BaseScope.Module_Type_Module:
-                api_or_module_name = module.mindstudio_reserved_name
+                api_or_module_name = module.mindstudio_reserved_name[-1]
+            else:
+                module.forward_data_collected = True
+                HOOKModule.add_module_count(name)
             self.data_collector.update_api_or_module_name(api_or_module_name)
 
             if self.config.online_run_ut:
+                self.inner_switch = False
                 return None, None
             if self.data_collector:
                 module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
-                self.data_collector.pre_forward_data_collect(api_or_module_name, module, pid, module_input_output)
+                self.data_collector.forward_input_data_collect(api_or_module_name, module, pid, module_input_output)
+
+            self.inner_switch = False
             return args, kwargs
 
+        def grad_hook(module, ori_name, param_name):
+            def hook_fn(grad):
+                if not self.should_execute_hook(module_type, module, False):
+                    return grad
+                self.inner_switch = True
+                self.data_collector.params_data_collect(ori_name, param_name, pid, grad)
+                self.inner_switch = False
+                return grad
+
+            return hook_fn
+
+        def register_param_hook(ori_name, module, params_dict):
+            '''
+            注册参数hook
+            '''
+            # data_mode为forward时，不注册参数hook
+            if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
+                for param_name, param in params_dict.items():
+                    if param.requires_grad:
+                        param.register_hook(grad_hook(module, ori_name, param_name))
+
+        def init_params_grad_info(module, params_dict):
+            '''
+            初始化参数梯度信息, 在前向hook结束后, 将参数梯度信息写入cache_data中用于占位
+            '''
+            if not params_dict:
+                return
+            if not (Const.FORWARD in self.config.data_mode and Const.BACKWARD not in self.config.data_mode):
+                grad_name = module.params_grad_name if hasattr(module, 'params_grad_name') else None
+                # 判断是否已经在cache_data中进行了占位, 若没有则先写入cache_data中
+                if not self.params_grad_info.get(grad_name):
+                    data_info = {grad_name: {key: [None] for key, value in params_dict.items() if value.requires_grad}}
+                    # 当模块中的参数有requires_grad属性为True时，才会进行梯度计算，此时才需要占位
+                    if data_info.get(grad_name):
+                        # 将grad_name的data_info先写入cache_data中, 梯度计算后再更新
+                        self.data_collector.handle_data(grad_name, data_info,
+                                                        flush=self.data_collector.data_processor.is_terminated)
+                    # 记录当前模块的参数梯度信息已占位
+                    self.params_grad_info[grad_name] = True
+
         def forward_hook(api_or_module_name, module, args, kwargs, output):
-            if not self.should_execute_hook():
+            if not self.should_execute_hook(module_type, module, True):
                 return None
 
-            if module_type == BaseScope.Module_Type_Module:
-                api_or_module_name = module.mindstudio_reserved_name
-            self.data_collector.update_api_or_module_name(api_or_module_name)
-
+            self.inner_switch = True
             if self.config.online_run_ut:
+                self.data_collector.update_api_or_module_name(api_or_module_name)
                 if self.data_collector.scope and not self.data_collector.scope.check(api_or_module_name):
                     return None
-                api_data = ApiData(name[:-1], args, kwargs, output, self.current_iter, self.current_rank)
+                api_data = ApiData(
+                    api_or_module_name[:-len(Const.FORWARD_NAME_SUFFIX)],
+                    args,
+                    kwargs,
+                    output,
+                    self.current_iter,
+                    self.current_rank
+                )
                 self.attl_send(api_data)
+                self.inner_switch = False
                 return None
 
-            if self.data_collector:
-                module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
-                self.data_collector.forward_data_collect(api_or_module_name, module, pid, module_input_output)
-                if self.data_collector.if_return_forward_new_output():
-                    return self.data_collector.get_forward_new_output()
+            module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+            if module_type == BaseScope.Module_Type_Module:
+                api_or_module_name = module.mindstudio_reserved_name[-1]
+                self.data_collector.update_api_or_module_name(api_or_module_name)
+                params_dict = {key.split(Const.SEP)[-1]: value for key, value in module.named_parameters(recurse=False)}
+                setattr(module_input_output, Const.PARAMS, params_dict)
+                # 判断是否需要注册参数hook
+                if not hasattr(module, 'params_grad_name') and params_dict:
+                    ori_name = api_or_module_name.rsplit(Const.SEP, 2)[0]
+                    grad_name = ori_name + Const.SEP + Const.PARAMS_GRAD
+                    # 首次执行前向hook时，添加params_grad_name属性，并注册参数hook
+                    setattr(module, 'params_grad_name', grad_name)
+                    register_param_hook(ori_name, module, params_dict)
+                self.data_collector.forward_data_collect(
+                    api_or_module_name,
+                    module,
+                    pid,
+                    module_input_output
+                )
+                init_params_grad_info(module, params_dict)
+            else:
+                self.data_collector.update_api_or_module_name(api_or_module_name)
+                self.data_collector.forward_output_data_collect(
+                    api_or_module_name,
+                    module,
+                    pid,
+                    module_input_output
+                )
+
+            if self.data_collector.if_return_forward_new_output():
+                forward_new_output = self.data_collector.get_forward_new_output()
+                self.inner_switch = False
+                return forward_new_output
+            self.inner_switch = False
             return output
 
         def forward_hook_torch_version_below_2(api_or_module_name, module, args, output):
             return forward_hook(api_or_module_name, module, args, {}, output)
 
         def backward_hook(api_or_module_name, module, grad_input, grad_output):
-            if not self.should_execute_hook():
+            if not self.should_execute_hook(module_type, module, False):
                 return
 
+            self.inner_switch = True
             if module_type == BaseScope.Module_Type_Module:
-                api_or_module_name = module.mindstudio_reserved_name
+                api_or_module_name = module.mindstudio_reserved_name[-1]
             self.data_collector.update_api_or_module_name(api_or_module_name)
 
             if self.config.online_run_ut:
+                self.inner_switch = False
                 return
 
             if self.data_collector:
                 # 此处获取到的grad_input实际为反向过程的输出数据，grad_output为反向过程的输入数据，因此传入时调换顺序
                 module_input_output = ModuleBackwardInputsOutputs(grad_input=grad_output, grad_output=grad_input)
                 self.data_collector.backward_data_collect(api_or_module_name, module, pid, module_input_output)
+            self.inner_switch = False
 
         pid = os.getpid()
-        forward_name_template = name + Const.FORWARD
-        backward_name_template = name + Const.BACKWARD
-        pre_forward_hook_fn = functools.partial(pre_hook, forward_name_template)
-        forward_hook_fn = functools.partial(forward_hook, forward_name_template)
-        backward_hook_fn = functools.partial(backward_hook, backward_name_template)
-        forward_hook_torch_version_below_2_fn = functools.partial(forward_hook_torch_version_below_2,
-                                                                  forward_name_template)
+        full_forward_name = None
+        full_backward_name = None
+        if module_type == BaseScope.Module_Type_API:
+            full_forward_name = name + str(HOOKModule.get_module_count(name)) + Const.SEP + Const.FORWARD
+            full_backward_name = name + str(HOOKModule.get_module_count(name)) + Const.SEP + Const.BACKWARD
+        pre_forward_hook_fn = functools.partial(pre_hook, full_forward_name)
+        forward_hook_fn = functools.partial(forward_hook, full_forward_name)
+        backward_hook_fn = functools.partial(backward_hook, full_backward_name)
+        forward_hook_torch_version_below_2_fn = functools.partial(
+            forward_hook_torch_version_below_2,
+            full_forward_name
+        )
         return HookFn(pre_forward_hook_fn, forward_hook_fn, backward_hook_fn, forward_hook_torch_version_below_2_fn)
 
-    def start(self, model, api_origin=False):
+    def start(self, model):
         if self.need_stop_service():
             return
 
@@ -155,10 +230,8 @@ class Service:
 
             if self.config.rank and self.current_rank not in self.config.rank:
                 return
-            self.register_hook_new()
+            self.register_module_hook()
             self.first_start = False
-        if api_origin:
-            api_register.api_modularity()
         if self.config.online_run_ut and torch_version_above_or_equal_2:
             run_ut_dispatch(self.attl, True, self.config.online_run_ut_recompute)
         self.switch = True
@@ -170,30 +243,31 @@ class Service:
     def stop(self):
         if self.should_stop_service:
             return
-        if self.config.level == "L2":
-            return
         if self.config.step and self.current_iter not in self.config.step:
             return
         if self.config.rank and self.current_rank not in self.config.rank:
             return
         self.switch = False
+        if self.config.level == Const.LEVEL_L2:
+            return
         if self.config.online_run_ut and torch_version_above_or_equal_2:
             run_ut_dispatch(self.attl, False, self.config.online_run_ut_recompute)
             return
+        if self.config.async_dump:
+            self.data_collector.fill_stack_tensor_data()
+            self.data_collector.data_processor.dump_async_data()
         self.data_collector.write_json()
 
     def step(self):
         if self.should_stop_service:
             return
+        if self.config.async_dump:
+            self.data_collector.fill_stack_tensor_data()
+            self.data_collector.data_processor.dump_async_data()
+        self.data_collector.write_json()
         self.current_iter += 1
         self.data_collector.update_iter(self.current_iter)
-
-        ModuleProcesser.reset_module_stats()
-        HOOKModule.reset_module_stats()
-        self.data_collector.data_writer.reset_cache()
-
-        if self.config.level == Const.LEVEL_L2:
-            self.data_collector.data_processor.reset_status()
+        self.reset_status()
 
     def need_stop_service(self):
         if self.should_stop_service:
@@ -204,8 +278,6 @@ class Service:
             if self.config.online_run_ut:
                 # send stop signal if online_run_ut
                 self.attl_stop()
-            if self.config.level in [Const.LEVEL_L1, Const.LEVEL_L2, Const.LEVEL_MIX]:
-                api_register.api_originality()
             self.switch = False
             self.should_stop_service = True
             print_tools_ends_info()
@@ -214,10 +286,18 @@ class Service:
             return True
         return False
 
-    def should_execute_hook(self):
-        if not self.switch:
+    def should_execute_hook(self, hook_type, module, is_forward):
+        is_module_hook = hook_type == BaseScope.Module_Type_Module
+        if is_module_hook and not self.switch:
+            return False
+        elif not is_module_hook and is_forward and not self.switch:
             return False
-        if self.data_collector and self.data_collector.data_processor.is_terminated:
+        elif not is_module_hook and not is_forward and not module.forward_data_collected:
+            return False
+
+        if self.inner_switch:
+            return False
+        if not self.data_collector or self.data_collector.data_processor.is_terminated:
             return False
         return True
 
@@ -244,50 +324,26 @@ class Service:
         construct_file_path = os.path.join(dump_dir, "construct.json")
         free_benchmark_file_path = os.path.join(self.config.dump_path, "free_benchmark.csv")
         self.data_collector.update_dump_paths(
-            dump_file_path, stack_file_path, construct_file_path, dump_data_dir, free_benchmark_file_path)
-
-    def register_hook_new(self):
-        logger.info_on_rank_0("The {} hook function is successfully mounted to the model.".format(self.config.task))
-        if self.config.level in ["L0", "mix"]:
-            if self.model is None:
-                logger.error_log_with_exp("The model is None.", MsprobeException.INVALID_PARAM_ERROR)
-            logger.info_on_rank_0("The init dump mode is enabled, and the module dump function will not be available")
-            for name, module in self.model.named_modules():
-                if module == self.model:
-                    continue
-                prefix = BaseScope.Module_Type_Module + Const.SEP + name + Const.SEP + \
-                         module.__class__.__name__ + Const.SEP
-
-                pre_forward_hook, forward_hook, backward_hook, forward_hook_torch_version_below_2 = self.build_hook(
-                    BaseScope.Module_Type_Module, prefix)
-                if torch_version_above_or_equal_2:
-                    module.register_forward_hook(forward_hook, with_kwargs=True)
-                else:
-                    self.check_register_full_backward_hook(module)
-                    module.register_full_backward_hook(
-                        self.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
-                    module.register_forward_hook(forward_hook_torch_version_below_2)
-                self.check_register_full_backward_hook(module)
-                module.register_full_backward_hook(backward_hook)
-
-                module.register_forward_pre_hook(
-                    self.module_processor.node_hook(prefix + Const.FORWARD, Const.START))
-                module.register_forward_hook(
-                    self.module_processor.node_hook(prefix + Const.FORWARD, Const.STOP))
-                if torch_version_above_or_equal_2:
-                    module.register_full_backward_pre_hook(
-                        self.module_processor.node_hook(prefix + Const.BACKWARD, Const.START))
-                    self.check_register_full_backward_hook(module)
-                    module.register_full_backward_hook(
-                        self.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
-
-        if self.config.level in ["mix", "L1", "L2"]:
-            api_register.initialize_hook(functools.partial(self.build_hook, BaseScope.Module_Type_API),
-                                         self.config.online_run_ut)
+            dump_file_path, stack_file_path, construct_file_path, dump_data_dir, free_benchmark_file_path
+        )
+        self.data_collector.initialize_json_file(framework=Const.PT_FRAMEWORK)
+
+    def register_api_hook(self):
+        if self.config.level in [Const.LEVEL_MIX, Const.LEVEL_L1, Const.LEVEL_L2]:
+            logger.info_on_rank_0(f"The api {self.config.task} hook function is successfully mounted to the model.")
+            api_register.initialize_hook(
+                functools.partial(self.build_hook, BaseScope.Module_Type_API),
+                self.config.online_run_ut
+            )
             api_register.api_modularity()
 
-        if Const.STATISTICS == self.config.task or Const.TENSOR == self.config.task:
-            remove_dropout()
+        if self.config.level == Const.LEVEL_MIX:
+            register_optimizer_hook(self.data_collector)
+
+    def register_module_hook(self):
+        if self.config.level in [Const.LEVEL_L0, Const.LEVEL_MIX]:
+            logger.info_on_rank_0(f"The module {self.config.task} hook function is successfully mounted to the model.")
+            self.module_processor.register_module_hook(self.model, self.build_hook)
 
     def attl_init(self):
         if self.config.online_run_ut:
@@ -319,3 +375,17 @@ class Service:
         elif self.attl.socket_manager is not None:
             logger.info(f"pid: {os.getpid()} finished, start send STOP signal.")
             self.attl.socket_manager.send_stop_signal()
+
+    def reset_status(self):
+        ModuleProcesser.reset_module_stats()
+        HOOKModule.reset_module_stats()
+        self.data_collector.data_writer.reset_cache()
+        self.params_grad_info.clear()
+
+        if self.config.level == Const.LEVEL_L2:
+            self.data_collector.data_processor.reset_status()
+            return
+        if self.config.step and self.current_iter not in self.config.step:
+            return
+        if self.config.rank and self.current_rank not in self.config.rank:
+            return