mindstudio-probe 8.2.1__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,8 +27,7 @@ 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

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,6 +22,9 @@ 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):
@@ -92,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,21 +13,25 @@
 # 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.common.megatron_utils import wrap_megatron_step, get_micro_step, is_megatron
 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'
 torch_version_above_or_equal_21 = torch.__version__.split('+')[0] >= '2.1'
@@ -63,6 +67,7 @@ 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_21:
         module.forward = wrapped_forward
 
@@ -80,6 +85,7 @@ 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)

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/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()

msprobe/visualization/builder/graph_builder.py

@@ -74,6 +74,7 @@ class GraphBuilder:
             config.graph_b.data_source = GraphConst.JSON_BENCH_KEY
             config.graph_b.step = config.step
             config.graph_b.rank = config.rank
+            config.graph_b.compare_mode = config.compare_mode
             node_to_db(config.graph_b, filename)
         config_to_db(config, filename)
 

msprobe/visualization/utils.py

@@ -152,7 +152,8 @@ def load_parallel_param(input_param):
 
 
 def validate_parallel_param(parallel_param, dump_path, log_prefix='[NPU]'):
-    params = [parallel_param.tp, parallel_param.pp, parallel_param.rank_size]
+    pattern = re.compile(r'^[a-z\-]+$')
+    params = [parallel_param.tp, parallel_param.pp, parallel_param.rank_size, parallel_param.vpp]
     ranks = check_and_return_dir_contents(dump_path, Const.RANK)
     if len(ranks) != parallel_param.rank_size:
         logger.error(f'{log_prefix} The parallel param "rank_size" error, '
@@ -161,6 +162,12 @@ def validate_parallel_param(parallel_param, dump_path, log_prefix='[NPU]'):
     if any(x is None for x in params):
         logger.error(f'{log_prefix} The parallel params "tp/pp/rank_size" must not be null!')
         raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if any(isinstance(x, bool) for x in params):
+        logger.error(f'{log_prefix} The parallel params "tp/pp/vpp/rank_size" must not be bool!')
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if any(not isinstance(x, int) for x in params):
+        logger.error(f'{log_prefix} The parallel params "tp/pp/vpp/rank_size" must be int!')
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
     if any(x <= 0 for x in params):
         logger.error(f'{log_prefix} The parallel params "tp/pp/vpp/rank_size" must be greater than 0!')
         raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
@@ -185,6 +192,9 @@ def validate_parallel_param(parallel_param, dump_path, log_prefix='[NPU]'):
     if not isinstance(parallel_param.order, str):
         logger.error(f'{log_prefix} The parallel params "order" must be of string type!')
         raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if not pattern.match(parallel_param.order):
+        logger.error(f'{log_prefix} The parallel params "order" must consist only of lowercase letters and "-"!')
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
 
 
 class ParallelParam:

msprobe/core/compare/diff_analyze/ignore_op_list.yaml

@@ -1,3 +0,0 @@
-npu_fusion_attention:
-  - 4
-  - 5