mindstudio-probe 8.2.0__py3-none-any.whl → 8.3.0__py3-none-any.whl

msprobe/pytorch/common/utils.py

@@ -39,7 +39,6 @@ except ImportError:
 else:
     is_gpu = False
 
-
 torch_without_guard_version = torch.__version__ >= '2.1'
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
@@ -416,7 +415,8 @@ def is_recomputation():
 
     # Identify indices in the call stack where the specific function is being executed
     for idx, frame_info in enumerate(call_stack):
-        if frame_info.function == Const.BACKWARD or frame_info.function == 'checkpoint_function_backward':
+        if (frame_info.function == Const.BACKWARD or frame_info.function == 'checkpoint_function_backward' and
+                "megatron" in frame_info.filename):
             backward_function_indices.append(idx)
 
     # Check if the execution is within 'torch/autograd/function.py' file
@@ -471,3 +471,23 @@ def register_forward_hook(module, forward_hook):
         module.register_forward_hook(forward_hook, with_kwargs=True)
     else:
         module.register_forward_hook(forward_hook)
+
+
+def save_api_data(api_data):
+    """Save data to io stream"""
+    try:
+        io_buff = io.BytesIO()
+        torch.save(api_data, io_buff)
+    except Exception as e:
+        raise RuntimeError(f"save api_data to io_buff failed") from e
+    return io_buff
+
+
+def load_api_data(api_data_bytes):
+    """Load data from bytes stream"""
+    try:
+        buffer = io.BytesIO(api_data_bytes)
+        buffer = torch.load(buffer, map_location="cpu")
+    except Exception as e:
+        raise RuntimeError(f"load api_data from bytes failed") from e
+    return buffer

msprobe/pytorch/compare/utils.py

@@ -27,15 +27,15 @@ def read_pt_data(dir_path, file_name):
         return None
 
     data_path = os.path.join(dir_path, file_name)
-    path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE,
-                               FileCheckConst.PT_SUFFIX, False)
+    path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE, FileCheckConst.PT_SUFFIX)
     data_path = path_checker.common_check()
     try:
         # detach because numpy can not process gradient information
         data_value = load_pt(data_path, to_cpu=True).detach()
     except RuntimeError as e:
         # 这里捕获 load_pt 中抛出的异常
-        logger.error(f"Failed to load the .pt file at {data_path}.")
+        data_path_file_name = os.path.basename(data_path)
+        logger.error(f"Failed to load the .pt file at {data_path_file_name}.")
         raise CompareException(CompareException.INVALID_FILE_ERROR) from e
     except AttributeError as e:
         # 这里捕获 detach 方法抛出的异常

msprobe/pytorch/debugger/debugger_config.py

@@ -48,6 +48,16 @@ class DebuggerConfig:
                 "max_sample": task_config.max_sample
             }
 
+        self.online_run_ut = False
+        if self.task == Const.TENSOR:
+            # dump api tensor and collaborate with online run_ut
+            self.online_run_ut = task_config.online_run_ut if task_config.online_run_ut else False
+            self.nfs_path = task_config.nfs_path if task_config.nfs_path else ""
+            self.tls_path = task_config.tls_path if task_config.tls_path else ""
+            self.host = task_config.host if task_config.host else ""
+            self.port = task_config.port if task_config.port else -1
+            self.online_run_ut_recompute = task_config.online_run_ut_recompute \
+                if isinstance(task_config.online_run_ut_recompute, bool) else False
 
         self.check()
         self._check_statistics_config(task_config)

msprobe/pytorch/dump/module_dump/hook_wrapper.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 from functools import wraps
+from typing import Any, Callable
 
 import torch
 from torch.utils.hooks import BackwardHook
@@ -21,11 +22,17 @@ from torch.utils.hooks import BackwardHook
 from msprobe.core.common.const import Const
 from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.hook_module.api_register import get_api_register
+
+torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
 
 def wrap_setup_backward_hook(func):
-    def requires_clone(tensor):
-        return isinstance(tensor, torch.Tensor) and tensor.requires_grad and torch.is_grad_enabled()
+    def requires_clone(tensor, need_check_leaf=False):
+        need_clone = isinstance(tensor, torch.Tensor) and tensor.requires_grad and torch.is_grad_enabled()
+        if need_check_leaf:
+            need_clone &= tensor.grad_fn is not None
+        return need_clone
 
     @recursion_depth_decorator("Dump: wrap_setup_backward_hook.parse_tensor", max_depth=Const.DUMP_MAX_DEPTH)
     def parse_tensor(item, tensor_list):
@@ -39,20 +46,20 @@ def wrap_setup_backward_hook(func):
                 parse_tensor(value, tensor_list)
 
     @recursion_depth_decorator("Dump: wrap_setup_backward_hook.rebuild_args", max_depth=Const.DUMP_MAX_DEPTH)
-    def rebuild_args(item, tensor_iter):
-        if requires_clone(item):
+    def rebuild_args(item, tensor_iter, need_check_leaf=False):
+        if requires_clone(item, need_check_leaf):
             result = next(tensor_iter)
             if hasattr(result, "_base") and result._base is not None:
                 if torch._C._autograd._get_creation_meta(result) != torch._C._autograd.CreationMeta(0):
                     torch._C._autograd._set_creation_meta(result, torch._C._autograd.CreationMeta(0))
-            return result    
+            return result
         if isinstance(item, list):
             for index, value in enumerate(item):
-                item[index] = rebuild_args(value, tensor_iter)
+                item[index] = rebuild_args(value, tensor_iter, need_check_leaf=True)
             return item
         if isinstance(item, dict):
             for key, value in item.items():
-                item[key] = rebuild_args(value, tensor_iter)
+                item[key] = rebuild_args(value, tensor_iter, need_check_leaf=True)
             return item
         if isinstance(item, tuple):
             if hasattr(item, '_fields'):
@@ -89,3 +96,23 @@ def wrap_setup_backward_hook(func):
 def wrap_setup_input_output_hook():
     BackwardHook.setup_input_hook = wrap_setup_backward_hook(BackwardHook.setup_input_hook)
     BackwardHook.setup_output_hook = wrap_setup_backward_hook(BackwardHook.setup_output_hook)
+
+
+def get_apply_func_wrapper(original_func: Callable) -> Callable:
+    @wraps(original_func)
+    def wrapped_apply(*args, **kwargs) -> Any:
+        api_register = get_api_register()
+        if api_register:
+            api_register.restore_inner_used_api()
+        result = original_func(*args, **kwargs)
+        if api_register:
+            api_register.register_inner_used_api()
+        return result
+
+    return wrapped_apply
+
+
+def wrap_backward_hook_function_apply():
+    if torch_version_above_or_equal_2:
+        original_apply = torch.nn.modules._functions.BackwardHookFunction.apply
+        torch.nn.modules._functions.BackwardHookFunction.apply = get_apply_func_wrapper(original_apply)

msprobe/pytorch/dump/module_dump/module_processer.py

@@ -13,23 +13,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import threading
 import sys
+import threading
 from collections import OrderedDict
 
 import torch
 from torch.utils.hooks import BackwardHook, RemovableHandle
 
 from msprobe.core.common.const import Const
+from msprobe.core.common.megatron_utils import wrap_megatron_step, get_micro_step, is_megatron
 from msprobe.core.common.runtime import Runtime
 from msprobe.core.common.utils import ModuleQueue, ThreadSafe
 from msprobe.core.data_dump.scope import BaseScope, ModuleRangeScope, MixRangeScope
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.common.utils import is_torch_nn_module, register_forward_pre_hook
-from msprobe.pytorch.dump.module_dump.hook_wrapper import wrap_setup_input_output_hook
+from msprobe.pytorch.dump.module_dump.hook_wrapper import (
+    wrap_setup_input_output_hook,
+    wrap_backward_hook_function_apply
+)
+
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
-if torch_version_above_or_equal_2:
+torch_version_above_or_equal_21 = torch.__version__.split('+')[0] >= '2.1'
+if torch_version_above_or_equal_21:
     from torch.utils.checkpoint import _StopRecomputationError
 
 
@@ -61,7 +67,8 @@ def wrap_forward_with_hook_safety(module):
                 hook_fn = list(module._forward_hooks.values())[0]
                 hook_fn(module, args, kwargs, exception_output)
             raise e
-    if torch_version_above_or_equal_2:
+
+    if torch_version_above_or_equal_21:
         module.forward = wrapped_forward
 
 
@@ -78,10 +85,13 @@ class ModuleProcesser:
     def __init__(self, scope):
         self.scope = scope if isinstance(scope, (ModuleRangeScope, MixRangeScope)) else None
         wrap_setup_input_output_hook()
+        wrap_backward_hook_function_apply()
         try:
             from megatron.core.pipeline_parallel import schedules
             origin_func_id = id(schedules.deallocate_output_tensor)
             schedules.deallocate_output_tensor = wrap_megatron_deallocate(schedules.deallocate_output_tensor)
+            schedules.forward_step = wrap_megatron_step(schedules.forward_step)
+            schedules.backward_step = wrap_megatron_step(schedules.backward_step, is_forward=False)
             for module in list(sys.modules.values()):
                 if module.__name__ == 'schedules':
                     continue
@@ -258,14 +268,16 @@ class ModuleProcesser:
             ModuleProcesser.module_stack[tid] = []
 
         if self.module_stack[tid]:
-            ModuleProcesser.module_node[full_name] = self.module_stack[tid][-1]
+            ModuleProcesser.module_node[full_name] = self.module_stack[tid][-1] if not is_megatron() \
+                else [self.module_stack[tid][-1], get_micro_step()]
         else:
             parent_name = ModuleProcesser.module_queue.find_last(full_name)
-            ModuleProcesser.module_node[full_name] = parent_name
+            ModuleProcesser.module_node[full_name] = parent_name if not is_megatron() \
+                else [parent_name, get_micro_step()]
 
         ModuleProcesser.module_queue.add_name(full_name)
         ModuleProcesser.module_stack[tid].append(full_name)
-        ModuleProcesser.api_parent_node[tid] = full_name
+        ModuleProcesser.api_parent_node[tid] = full_name if not is_megatron() else [full_name, get_micro_step()]
         if self.scope:
             self.scope.begin_module(full_name)
 
@@ -273,14 +285,15 @@ class ModuleProcesser:
         tid = threading.get_ident()
         if torch_version_above_or_equal_2 or is_forward:
             ModuleProcesser.module_queue.remove_name(full_name)
-            ModuleProcesser.api_parent_node[tid] = None
+            ModuleProcesser.api_parent_node[tid] = None if not is_megatron() else [None, get_micro_step()]
             if self.module_stack.get(tid):
                 ModuleProcesser.module_stack[tid].pop()
             if self.module_stack.get(tid):
-                ModuleProcesser.api_parent_node[tid] = ModuleProcesser.module_stack[tid][-1]
+                ModuleProcesser.api_parent_node[tid] = ModuleProcesser.module_stack[tid][-1] if not is_megatron() \
+                    else [ModuleProcesser.module_stack[tid][-1], get_micro_step()]
             if self.scope:
                 self.scope.end_module(full_name)
         else:
             if self.scope:
                 self.scope.begin_module(full_name)
-            ModuleProcesser.api_parent_node[tid] = full_name
+            ModuleProcesser.api_parent_node[tid] = full_name if not is_megatron() else [full_name, get_micro_step()]

msprobe/pytorch/hook_module/api_register.py

@@ -43,7 +43,6 @@ else:
 
 torch_version_above_2 = torch.__version__.split('+')[0] > '2.0'
 
-_inner_used_api = {}
 _supported_api_list_path = (os.path.join(os.path.dirname(os.path.realpath(__file__)), Const.SUPPORT_API_FILE_NAME),)
 _cuda_func_mapping = {"npu_fusion_attention": "gpu_fusion_attention"}
 dist_data_collect_func = {}
@@ -85,6 +84,12 @@ if not is_gpu:
         mindspeed_op_file_list = [op.split(Const.SEP)[0] + Const.PY_SUFFIX for op in mindspeed_op_list]
         dynamic_import_op(mindspeed.ops, mindspeed_op_file_list)
 
+_inner_used_api = {
+    Const.PT_FRAMEWORK + Const.SEP + Const.PT_API_TYPE_TENSOR: (
+        torch.Tensor, "view_as"
+    )
+}
+
 
 @parameter_adapter
 def tensor_module_forward(module, *args, **kwargs):

msprobe/pytorch/monitor/module_hook.py

@@ -19,12 +19,14 @@ import importlib
 from collections import defaultdict
 from datetime import datetime
 from functools import partial
+from itertools import cycle
 
 import pytz
 import torch
 import torch.distributed as dist
 import pandas as pd
 from torch.utils.hooks import BackwardHook
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
 from msprobe.core.common.const import MonitorConst, Const
 from msprobe.core.common.file_utils import load_json, save_json, make_dir
@@ -229,6 +231,8 @@ class TrainerMon:
         self.duplicate_param = {}
         self.name2tag = {}
         self.param_name_call_id = {}
+        self.flat_prefix_names = []
+        self.flat_prefix_reverse_iter = None
         self.call_id = 0
         self.module_struct = defaultdict(dict)
         self.grad_accs = []
@@ -945,13 +949,20 @@ class TrainerMon:
         return False
 
     def _register_chunk(self, model_chunk, prefix):
+        if isinstance(model_chunk, FSDP):
+            if not model_chunk._use_orig_params:
+                raise ValueError("Only Support fsdp1 with use_orig_params=True")
+            self.fsdp_wrapped_module = True
         for (param_name, param) in model_chunk.named_parameters():
             if not param.requires_grad:
                 continue
-            if not self.fsdp_wrapped_module and param_name.startswith("_fsdp_wrapped_module"):
-                self.fsdp_wrapped_module = True
             if not self.fsdp2_wrapped_module and param.__class__.__name__ == "DTensor":
                 self.fsdp2_wrapped_module = True
+            if self.fsdp_wrapped_module:  # FSDP1需要记录完整的不被target限制的flat权重前缀名，以供后续对flat解包
+                flat_prefix_name, _ = param_name.rsplit(MonitorConst.FSDP_FLAT_SEP, 1)
+                if flat_prefix_name not in self.flat_prefix_names:
+                    self.flat_prefix_names.append(flat_prefix_name)
+
             if self._is_target_param(param_name, param, prefix):
                 name = prefix + squash_param_name(param_name, self.squash_name)
                 if name in self.param2name.values():
@@ -975,6 +986,8 @@ class TrainerMon:
                     k: get_summary_writer_tag_name(name, k, self.rank)
                     for k in keywords
                 }
+        if self.fsdp_wrapped_module:
+            self.flat_prefix_reverse_iter = cycle(reversed(self.flat_prefix_names))  # post_backward_hook调用顺序是反向的
 
     def _register_param_name(self):
         for vpp_stage, model_chunk in enumerate(self.model):
@@ -1224,17 +1237,22 @@ class TrainerMon:
         每个forward阶段，fsdp对AccumulateGrad重复注册hook方法，monitor工具内注册hook无法生效，
         因此对_post_backward_hook进行patch，在backward后，reduce_scatter前采集梯度。
         """
+
         def patch_post_backward_hook(_post_backward_hook):
             def wrapper(state, handle, *unused):
                 grad_dict = {}
-                offset = 0
-                for param, name in self.param2name.items():
-                    limit = param.numel()
-                    if not limit:
+                local_names = handle.flat_param._fqns
+                offsets = handle._get_flat_param_offsets()
+                shapes = handle.flat_param._shapes
+                flat_prefix = next(self.flat_prefix_reverse_iter)
+                for local_name, (start, end), local_shape in zip(local_names, offsets, shapes):
+                    grad_clip = handle.flat_param.grad[start:end + 1]
+                    grad = grad_clip.reshape(local_shape)
+                    total_name = f"{flat_prefix}{MonitorConst.FSDP_FLAT_SEP}{local_name}"
+                    if total_name not in self.origin2squash:
+                        logger.warning(f"{total_name} not in model.named_parameters(), skip.")
                         continue
-                    grad = handle.flat_param.grad[offset:offset + limit]
-                    offset += limit
-                    tag = self.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD)
+                    tag = self.name2tag.get(self.origin2squash[total_name], {}).get(MonitorConst.PRE_GRAD)
                     if tag is None:
                         continue
                     grad_dict[tag] = grad
@@ -1242,6 +1260,7 @@ class TrainerMon:
                 get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
                 out = _post_backward_hook(state, handle, *unused)
                 return out
+
             return wrapper
 
         logger.info("Patch fsdp _post_backward_hook, collect pre_grad metrics.")

msprobe/pytorch/online_dispatch/dispatch.py

@@ -17,7 +17,7 @@ import json
 import os
 import time
 import multiprocessing
-from multiprocessing import Pool
+from multiprocessing import Pool, Lock
 
 import torch
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -39,6 +39,7 @@ from msprobe.pytorch.online_dispatch.utils import get_callstack, data_to_cpu, ge
 from msprobe.pytorch.online_dispatch.compare import Comparator
 from msprobe.core.common.utils import check_str_param, safe_get_value
 
+child_global_lock = None
 current_time = time.strftime("%Y%m%d%H%M%S")
 RESULT_FILE_NAME = "accuracy_checking_result_" + current_time + ".csv"
 DETAILS_FILE_NAME = "accuracy_checking_details_" + current_time + ".csv"
@@ -86,14 +87,14 @@ class PtdbgDispatch(TorchDispatchMode):
         yaml_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "torch_ops_config.yaml")
         self.get_ops(yaml_path)
 
-        self.lock = None
+        self.lock = Lock() if process_num > 0 else None
         max_process_num = max(int((multiprocessing.cpu_count() + 1) // Const.CPU_QUARTER), 1)
         if process_num > max_process_num:
             logger.error(f"process_num should be less than or equal to {max_process_num}, but got {process_num}!")
             raise DispatchException(f'process_num should be less than or equal to {max_process_num}, '
                                     f'but got {process_num}!')
         if process_num > 0:
-            self.pool = Pool(process_num)
+            self.pool = Pool(process_num, initializer=self._init_child_process, initargs=(self.lock,))
         if debug:
             logger.info(f'Main pid:{os.getpid()} device:{self.device_id} dump_list:{self.dump_api_list} '
                         f'dump_mode:{self.dump_mode} cpu_path[{self.root_cpu_path}], npu_path[{self.root_npu_path}], '
@@ -114,18 +115,17 @@ class PtdbgDispatch(TorchDispatchMode):
                 logger.error("Please check train log, An exception may have occurred!")
                 return
             check_file_or_directory_path(summary_path, False)
-            fp_handle = FileOpen(summary_path, "r")
-            while True:
-                json_line_data = fp_handle.readline()
-                if json_line_data == '\n':
-                    continue
-                if len(json_line_data) == 0:
-                    break
-                msg = json.loads(json_line_data)
-                if len(msg) < 2:
-                    raise ValueError("JSON data does not contain enough elements. Expected at least 2 elements.")
-                self.all_summary[msg[0]] = msg[1]
-            fp_handle.close()
+            with FileOpen(summary_path, "r") as fp_handle:
+                while True:
+                    json_line_data = fp_handle.readline()
+                    if json_line_data == '\n':
+                        continue
+                    if len(json_line_data) == 0:
+                        break
+                    msg = json.loads(json_line_data)
+                    if len(msg) < 2:
+                        raise ValueError("JSON data does not contain enough elements. Expected at least 2 elements.")
+                    self.all_summary[msg[0]] = msg[1]
 
         if self.debug_flag:
             input_num = 0
@@ -163,11 +163,16 @@ class PtdbgDispatch(TorchDispatchMode):
 
         call_stack = get_callstack()
         self.call_stack_list.append(call_stack)
-        self.api_index += 1
-        if aten_api not in self.single_api_index_dict:
-            self.single_api_index_dict[aten_api] = 1
-        else:
-            self.single_api_index_dict[aten_api] += 1
+
+        self.lock.acquire() if self.process_num > 0 else None
+        try:
+            self.api_index += 1
+            if aten_api not in self.single_api_index_dict:
+                self.single_api_index_dict[aten_api] = 1
+            else:
+                self.single_api_index_dict[aten_api] += 1
+        finally:
+            self.lock.release() if self.process_num > 0 else None
 
         run_param = self.get_run_param(aten_api, func.__name__, aten_api_overload_name)
 
@@ -180,7 +185,7 @@ class PtdbgDispatch(TorchDispatchMode):
         cpu_kwargs = []
         data_to_cpu(args, 0, cpu_args)
         data_to_cpu(kwargs, 0, cpu_kwargs)
-        
+
         cpu_args = safe_get_value(cpu_args, 0, "cpu_args")
         cpu_kwargs = safe_get_value(cpu_kwargs, 0, "cpu_kwargs")
 
@@ -194,7 +199,12 @@ class PtdbgDispatch(TorchDispatchMode):
             try:
                 cpu_out = func(*cpu_args, **cpu_kwargs)
             except RuntimeError as e:
-                self.api_index -= 1
+                self.lock.acquire() if self.process_num > 0 else None
+                try:
+                    self.api_index -= 1
+                    self.single_api_index_dict[aten_api] -= 1
+                finally:
+                    self.lock.release() if self.process_num > 0 else None
                 logger.warning(f"RuntimeError: {e}")
                 logger.warning(f"This aten_api {aten_api} does not support running on cpu, so skip it.")
                 return npu_out
@@ -215,7 +225,7 @@ class PtdbgDispatch(TorchDispatchMode):
             run_param.process_flag = True
             if self.check_fun(func, run_param):
                 data_info = DisPatchDataInfo(cpu_args, cpu_kwargs, self.all_summary, None, npu_out_cpu, cpu_out,
-                                             self.lock)
+                                             child_global_lock)
                 self.pool.apply_async(func=dispatch_multiprocess, args=(run_param, data_info),
                                       error_callback=error_call)
             else:
@@ -233,12 +243,20 @@ class PtdbgDispatch(TorchDispatchMode):
                     return True
         return False
 
+    @staticmethod
+    def _init_child_process(lock):
+        global child_global_lock
+        child_global_lock = lock
+
     def get_dir_name(self, tag):
         # guarantee file uniqueness
         time.sleep(1)
-        time_now = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
+        # 时间格式：年-月-日-时-分-秒-毫秒（精确到千分之一秒）
+        time_now = time.strftime("%Y%m%d%H%M%S%f", time.localtime(time.time()))[:-3]  # 取前3位毫秒
+
         if tag is None or not isinstance(tag, str):
             logger.warning('There is not tag or the type of tag is not string.')
+            # 目录名格式：msprobe_rank{设备ID}_{毫秒时间戳}
             dir_name = f'msprobe_rank{self.device_id}_{time_now}'
         else:
             dir_name = f'msprobe_{tag}_rank{self.device_id}_{time_now}'

msprobe/pytorch/pt_config.py

@@ -35,15 +35,48 @@ from msprobe.pytorch.hook_module.utils import get_ops
 class TensorConfig(BaseConfig):
     def __init__(self, json_config):
         super().__init__(json_config)
+        self.online_run_ut = json_config.get("online_run_ut", False)
+        self.nfs_path = json_config.get("nfs_path", "")
+        self.host = json_config.get("host", "")
+        self.port = json_config.get("port", -1)
+        self.tls_path = json_config.get("tls_path", "./")
+        self.online_run_ut_recompute = json_config.get("online_run_ut_recompute", False)
         self.check_config()
         self._check_summary_mode()
         self._check_file_format()
-
+        if self.online_run_ut:
+            self._check_online_run_ut()
 
     def _check_file_format(self):
         if self.file_format is not None and self.file_format not in ["npy", "bin"]:
             raise Exception("file_format is invalid")
 
+    def _check_online_run_ut(self):
+        if not isinstance(self.online_run_ut, bool):
+            raise Exception(f"online_run_ut: {self.online_run_ut} is invalid.")
+
+        if not isinstance(self.online_run_ut_recompute, bool):
+            raise Exception(f"online_run_ut_recompute: {self.online_run_ut_recompute} is invalid.")
+
+        if self.nfs_path:
+            check_file_or_directory_path(self.nfs_path, isdir=True)
+            return
+
+        if self.tls_path:
+            check_file_or_directory_path(self.tls_path, isdir=True)
+            check_file_or_directory_path(os.path.join(self.tls_path, "client.key"))
+            check_file_or_directory_path(os.path.join(self.tls_path, "client.crt"))
+            check_file_or_directory_path(os.path.join(self.tls_path, "ca.crt"))
+            crl_path = os.path.join(self.tls_path, "crl.pem")
+            if os.path.exists(crl_path):
+                check_file_or_directory_path(crl_path)
+
+        if not isinstance(self.host, str) or not re.match(Const.ipv4_pattern, self.host):
+            raise Exception(f"host: {self.host} is invalid.")
+
+        if not isinstance(self.port, int) or not (0 < self.port <= 65535):
+            raise Exception(f"port: {self.port} is invalid, port range 0-65535.")
+
 
 class StatisticsConfig(BaseConfig):
     def __init__(self, json_config):
@@ -80,6 +113,7 @@ class FreeBenchmarkCheckConfig(BaseConfig):
         self.handler_type = json_config.get("handler_type", PytorchFreeBenchmarkConst.DEFAULT_HANDLER)
         self.fuzz_level = json_config.get("fuzz_level", PytorchFreeBenchmarkConst.DEFAULT_FUZZ_LEVEL)
         self.fuzz_stage = json_config.get("fuzz_stage", PytorchFreeBenchmarkConst.DEFAULT_FUZZ_STAGE)
+        self.list = json_config.get("list")
         self.if_preheat = json_config.get("if_preheat", False)
         self.preheat_step = json_config.get("preheat_step", PytorchFreeBenchmarkConst.DEFAULT_PREHEAT_STEP)
         self.max_sample = json_config.get("max_sample", PytorchFreeBenchmarkConst.DEFAULT_PREHEAT_STEP)
@@ -146,6 +180,11 @@ class FreeBenchmarkCheckConfig(BaseConfig):
             logger.error_log_with_exp(
                 msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
             )
+        if self.fuzz_stage == Const.BACKWARD and not self.list:
+            raise MsprobeException(
+                MsprobeException.INVALID_PARAM_ERROR,
+                f"When fuzz_stage is set to {Const.BACKWARD}, the parameters list must not be empty."
+            )
 
     def _check_fuzz_level(self):
         if self.fuzz_level not in PytorchFreeBenchmarkConst.FUZZ_LEVEL_LIST:
@@ -218,7 +257,12 @@ class RunUTConfig(BaseConfig):
         self.white_list = json_config.get("white_list", Const.DEFAULT_LIST)
         self.black_list = json_config.get("black_list", Const.DEFAULT_LIST)
         self.error_data_path = json_config.get("error_data_path", Const.DEFAULT_PATH)
-
+        self.is_online = json_config.get("is_online", False)
+        self.nfs_path = json_config.get("nfs_path", "")
+        self.host = json_config.get("host", "")
+        self.port = json_config.get("port", -1)
+        self.rank_list = json_config.get("rank_list", Const.DEFAULT_LIST)
+        self.tls_path = json_config.get("tls_path", "./")
         self.check_run_ut_config()
 
     @classmethod
@@ -236,11 +280,22 @@ class RunUTConfig(BaseConfig):
         if not os.path.exists(error_data_path):
             raise Exception("error_data_path: %s does not exist" % error_data_path)
 
+    @classmethod
+    def check_nfs_path_config(cls, nfs_path):
+        if nfs_path:
+            FileChecker(nfs_path, FileCheckConst.DIR, FileCheckConst.READ_ABLE).common_check()
+
+    @classmethod
+    def check_tls_path_config(cls, tls_path):
+        if tls_path:
+            FileChecker(tls_path, FileCheckConst.DIR, FileCheckConst.READ_ABLE).common_check()
 
     def check_run_ut_config(self):
         RunUTConfig.check_filter_list_config(Const.WHITE_LIST, self.white_list)
         RunUTConfig.check_filter_list_config(Const.BLACK_LIST, self.black_list)
         RunUTConfig.check_error_data_path_config(self.error_data_path)
+        RunUTConfig.check_nfs_path_config(self.nfs_path)
+        RunUTConfig.check_tls_path_config(self.tls_path)
 
 
 class GradToolConfig(BaseConfig):

msprobe/pytorch/pytorch_service.py

@@ -15,8 +15,9 @@
 
 from msprobe.core.common.utils import Const
 from msprobe.core.service import BaseService
+from msprobe.pytorch.attl_manager import ATTLManager
 from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.common.utils import get_rank_if_initialized
+from msprobe.pytorch.common.utils import get_rank_if_initialized, torch_version_above_or_equal_2
 from msprobe.pytorch.dump.module_dump.module_processer import ModuleProcesser
 from msprobe.pytorch.hook_module.api_register import get_api_register, ApiTemplate, redirect_wait
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
@@ -24,6 +25,9 @@ from msprobe.pytorch.hook_module.pt_hook_manager import PytorchHookManager
 from msprobe.pytorch.hook_module.register_optimizer_hook import register_optimizer_hook
 from msprobe.pytorch.hook_module.script_wrapper import wrap_script_func, preprocess_func
 
+if torch_version_above_or_equal_2:
+    from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.dump_dispatch import run_ut_dispatch
+
 
 class PytorchService(BaseService):
     @property
@@ -41,10 +45,12 @@ class PytorchService(BaseService):
         self.logger = logger
         self.api_register = get_api_register()
         self.module_processor = ModuleProcesser(self.data_collector.scope)
-        self.hook_manager = PytorchHookManager(self.data_collector, self.config)
+        self.attl_manager = ATTLManager(self.config)
+        self.hook_manager = PytorchHookManager(self.data_collector, self.config, self.attl_manager)
         self.api_template = ApiTemplate
 
     def _register_hook(self):
+        self.attl_manager.attl_init()
         if self._is_mix_level:
             register_optimizer_hook(self.data_collector)
 
@@ -59,6 +65,9 @@ class PytorchService(BaseService):
         self.module_processor.register_module_hook(self.model, self.build_hook)
         self.logger.info(f"The module {self.config.task} hook function is successfully mounted to the model.")
 
+    def _run_ut_dispatch(self, status):
+        if torch_version_above_or_equal_2:
+            run_ut_dispatch(self.attl_manager.attl, status, self.config.online_run_ut_recompute)
 
     def _reset_status(self):
         super()._reset_status()