mindstudio-probe 1.2.2__py3-none-any.whl → 8.1.0__py3-none-any.whl

msprobe/pytorch/attl_manager.py

@@ -0,0 +1,65 @@
+# Copyright (c) 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
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from msprobe.core.common.runtime import Runtime
+from msprobe.core.common.utils import Const
+from msprobe.pytorch.api_accuracy_checker.common.utils import ApiData
+from msprobe.pytorch.common.log import logger
+
+
+class ATTLManager:
+    def __init__(self, config):
+        self.config = config
+        self.attl = None
+
+    def attl_init(self):
+        if self.config.online_run_ut:
+            from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import ATTLConfig, ATTL
+            attl_config = ATTLConfig(is_benchmark_device=False,
+                                     connect_ip=self.config.host,
+                                     connect_port=self.config.port,
+                                     nfs_path=self.config.nfs_path,
+                                     tls_path=self.config.tls_path)
+            need_dump = len(self.config.rank) == 0 or Runtime.current_rank in self.config.rank
+            self.attl = ATTL('npu', attl_config, need_dump=need_dump)
+            if self.config.nfs_path:
+                self.attl.upload("start")
+
+    def attl_send(self, name, args, kwargs, output):
+        api_data = ApiData(
+                    name[:-len(Const.FORWARD_NAME_SUFFIX)],
+                    args,
+                    kwargs,
+                    output,
+                    Runtime.current_iter,
+                    Runtime.current_rank
+                )
+        logger.info(f"tools is dumping api: {api_data.name}, rank: {Runtime.current_rank}")
+        api_type, _, _ = api_data.name.split(Const.SEP)
+        if api_type in [Const.DISTRIBUTED]:
+            logger.info(f"api {api_data.name} is not supported, skip")
+            return
+        if self.config.nfs_path:
+            self.attl.upload(api_data)
+        else:
+            self.attl.send(api_data)
+
+    def attl_stop(self):
+        if self.config.nfs_path:
+            self.attl.upload("end")
+        elif self.attl.socket_manager is not None:
+            logger.info(f"pid: {os.getpid()} finished, start sends STOP signal.")
+            self.attl.socket_manager.send_stop_signal()

msprobe/pytorch/bench_functions/moe_gating_top_k_softmax.py

@@ -29,6 +29,8 @@ def softmax_func(x, axis=None):
 
 def npu_moe_gating_top_k_softmax(x, finished_optional, k):
     input_dtype = x.dtype
+    if x.dim() < 1:
+        raise ValueError("Input x must have at least 1 dimensions.")
     num_expert = x.shape[-1]
     softmax = softmax_func(x, -1)
     softmax = softmax.to(input_dtype)
@@ -36,9 +38,13 @@ def npu_moe_gating_top_k_softmax(x, finished_optional, k):
     expert_idx = expert_idx[:, :k]
     y = torch.gather(softmax, index=expert_idx, dim=-1)
     if finished_optional is not None:
+        if finished_optional.dim() < 1:
+            raise ValueError("Finished_optional must have at least 1 dimensions.")
         finished_optional = finished_optional.view(finished_optional.shape[0], 1)
         finished_optional = finished_optional.expand(-1, k)
         expert_idx = torch.where(finished_optional, num_expert, expert_idx)
+    if y.dim() < 2:
+        raise ValueError("Variable y must have at least 2 dimensions.")
     row_idx = torch.arange(y.shape[0] * y.shape[1]).reshape(y.shape[1], y.shape[0]).t()
 
     return y, expert_idx, row_idx

msprobe/pytorch/bench_functions/npu_fusion_attention.py

@@ -117,6 +117,12 @@ def fusion_attention_forward(forward_params):
     pse = forward_params.pse
     scale = forward_params.scale
     keep_prob = forward_params.keep_prob
+
+    # 除零风险拦截：keep_prob 为 0 时会导致除零错误
+    if keep_prob == 0:
+        raise ValueError("fusion_attention_forward: keep_prob cannot be zero to avoid division by zero.")
+
+
     qk = calculate_qk(q, k, atten_mask, pse, scale)
     softmax_res, softmax_max, softmax_sum = softmax_forward(qk)
     if drop_mask is None or len(drop_mask.shape) == 0:
@@ -137,6 +143,11 @@ def fusion_attention_backward(backward_params):
     pse = backward_params.pse
     scale = backward_params.scale
     keep_prob = backward_params.keep_prob
+
+    # 除零风险拦截：keep_prob 为 0 时会导致除零错误
+    if keep_prob == 0:
+        raise ValueError("fusion_attention_backward: keep_prob cannot be zero to avoid division by zero.")
+
     dp = torch.matmul(dx, v.permute(0, 1, 3, 2))
     if drop_mask is None or len(drop_mask.shape) == 0:
         drop_res = softmax_res.permute(0, 1, 3, 2)
@@ -164,23 +175,35 @@ def parse_bsnd_args(query, key, head_num, input_layout):
         if input_layout == "BSH":
             b, s1, h1 = query.shape
             _, s2, h2 = key.shape
+            if n1 == 0:
+                raise ValueError("parse_bsnd_args: head_num (n1) cannot be zero to avoid division by zero.")
             d = h1 // n1
+            if d == 0:
+                raise ValueError("parse_bsnd_args: computed head dimension (d) is zero, division by zero risk.")
             n2 = h2 // d
         elif input_layout == "SBH":
             s1, b, h1 = query.shape
             s2, _, h2 = key.shape
+            if n1 == 0:
+                raise ValueError("parse_bsnd_args: head_num (n1) cannot be zero to avoid division by zero.")
             d = h1 // n1
+            if d == 0:
+                raise ValueError("parse_bsnd_args: computed head dimension (d) is zero, division by zero risk.")
             n2 = h2 // d
         elif input_layout == "BSND":
             b, s1, n1, d = query.shape
             _, s2, n2, _ = key.shape
             h1 = n1 * d
             h2 = n2 * d
+            if d == 0:
+                raise ValueError("parse_bsnd_args: head dimension (d) is zero, division by zero risk.")
         elif input_layout == "BNSD":
             b, n1, s1, d = query.shape
             _, n2, s2, _ = key.shape
             h1 = n1 * d
             h2 = n2 * d
+            if d == 0:
+                raise ValueError("parse_bsnd_args: head dimension (d) is zero, division by zero risk.")
     except Exception as e:
         raise ValueError(f"query.shape: {query.shape}, key.shape: {key.shape}, parse_bsnd_args error: {e}") from e
 
@@ -446,6 +469,8 @@ def npu_fusion_attention_forward_patch(*args, **kwargs):
     input_layout = get_input_layout(*args, **kwargs)
 
     b, s1, s2, n1, n2, d, h1, h2, dtype = parse_bsnd_args(args[0], args[1], head_num, input_layout)
+    if d == 0:
+        raise ValueError("npu_fusion_attention_forward_patch: head dimension (d) is zero, division by zero risk.")
     if n1 == n2 and s1 == s2:
         logger.debug(f"running case : BNSD = {b}_{n1}_{s1}_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
     else:
@@ -478,6 +503,8 @@ def npu_fusion_attention_backward_patch(*args, **kwargs):
         raise ValueError(f"Unsupported npu_fusion_attention_grad args {args}.")
 
     b, s1, s2, n1, n2, d, h1, h2, dtype = parse_bsnd_args(args[0], args[1], args[4], args[5])
+    if d == 0:
+        raise ValueError("npu_fusion_attention_backward_patch: head dimension (d) is zero, division by zero risk.")
     if n1 == n2 and s1 == s2:
         logger.info(f"running case : bnsd = {b}_{n1}_{s1}_{d}, sparse = {kwargs.get('sparse_mode', 0)}")
     else:

msprobe/pytorch/common/utils.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");
@@ -24,11 +24,12 @@ from functools import wraps
 import numpy as np
 import torch
 import torch.distributed as dist
+
 from msprobe.core.common.exceptions import DistributedNotInitializedError
 from msprobe.core.common.file_utils import (FileCheckConst, change_mode,
                                             check_file_or_directory_path, check_path_before_create, FileOpen)
 from msprobe.core.common.log import logger
-from msprobe.core.common.utils import check_seed_all
+from msprobe.core.common.utils import check_seed_all, is_save_variable_valid
 from packaging import version
 
 try:
@@ -38,7 +39,9 @@ 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'
 
 if not is_gpu and not torch_without_guard_version:
     from torch_npu.utils.device_guard import torch_device_guard as torch_npu_device_guard
@@ -57,7 +60,7 @@ def parameter_adapter(func):
 
     @wraps(func)
     def inner(self, *args, **kwargs):
-        if self.op_name_ == "__getitem__" and len(args) > 1 and isinstance(args[1], torch.Tensor):
+        if self.api_name == "__getitem__" and len(args) > 1 and isinstance(args[1], torch.Tensor):
             input_tensor = args[0]
             indices = args[1]
             if indices.dtype == torch.uint8:
@@ -77,7 +80,7 @@ def parameter_adapter(func):
                 else:
                     res = [input_tensor[tensor_index] for tensor_index in indices]
                     return getattr(torch._C._VariableFunctionsClass, "stack")(res, 0)
-        if self.op_name_ == "__eq__" and len(args) > 1 and args[1] is None:
+        if self.api_name == "__eq__" and len(args) > 1 and args[1] is None:
             return False
         return func(self, *args, **kwargs)
 
@@ -261,6 +264,10 @@ class Const:
     NPU = 'NPU'
     DISTRIBUTED = 'Distributed'
 
+    HIFLOAT8_TYPE = "torch_npu.HiFloat8Tensor"
+    FLOAT8_E5M2_TYPE = "torch.float8_e5m2"
+    FLOAT8_E4M3FN_TYPE = "torch.float8_e4m3fn"
+
     RAISE_PRECISION = {
         torch.float16: torch.float32,
         torch.bfloat16: torch.float32,
@@ -309,14 +316,14 @@ def print_rank_0(message):
         logger.info(message)
 
 
-def load_pt(pt_path, to_cpu=False):
+def load_pt(pt_path, to_cpu=False, weights_only=True):
     pt_path = os.path.realpath(pt_path)
     check_file_or_directory_path(pt_path)
     try:
         if to_cpu:
-            pt = torch.load(pt_path, map_location=torch.device("cpu"), weights_only=True)
+            pt = torch.load(pt_path, map_location=torch.device("cpu"), weights_only=weights_only)
         else:
-            pt = torch.load(pt_path, weights_only=True)
+            pt = torch.load(pt_path, weights_only=weights_only)
     except Exception as e:
         raise RuntimeError(f"load pt file {pt_path} failed") from e
     return pt
@@ -391,7 +398,7 @@ def save_api_data(api_data):
         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
+        raise RuntimeError("save api_data to io_buff failed") from e
     return io_buff
 
 
@@ -401,7 +408,7 @@ def load_api_data(api_data_bytes):
         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
+        raise RuntimeError("load api_data from bytes failed") from e
     return buffer
 
 
@@ -419,7 +426,11 @@ def is_recomputation():
         bool: True if in the re-computation phase, False otherwise.
     """
     backward_function_indices = []
-    call_stack = inspect.stack()
+    try:
+        call_stack = inspect.stack()
+    except Exception as e:
+        logger.warning(f"Failed to capture stack trace, recomputation validation may be incorrect, error info: {e}.")
+        return False
 
     # Identify the function 'backward' is being executed within the 'torch/_tensor.py' file.
     for frame_info in call_stack:
@@ -449,9 +460,11 @@ def is_recomputation():
 
 def check_save_param(variable, name, save_backward):
     # try catch this api to skip invalid call
-    if not isinstance(variable, (list, dict, torch.Tensor, int, float, str)):
+    valid_data_types = (torch.Tensor, int, float, str)
+    if not is_save_variable_valid(variable, valid_data_types):
+        valid_data_types_with_nested_types = valid_data_types + (dict, tuple, list)
         logger.warning("PrecisionDebugger.save variable type not valid, "
-                       "should be one of list, dict, torch.Tensor, int, float or string. "
+                       f"should be one of {valid_data_types_with_nested_types}"
                        "Skip current save process.")
         raise ValueError
     if not isinstance(name, str):
@@ -466,10 +479,31 @@ def check_save_param(variable, name, save_backward):
         raise ValueError
 
 
-def replace_last_occurrence(text, old, new):
-    if text is None:
-        return text
-    index = text.rfind(old)
-    if index != -1:
-        return text[:index] + text[index:].replace(old, new, 1)
-    return text
+def is_torch_nn_module(variable):
+    return isinstance(variable, torch.nn.Module) and not isinstance(variable, torch.jit.ScriptModule)
+
+
+def is_hifloat8_tensor(tensor):
+    if not is_gpu and hasattr(torch_npu, "HiFloat8Tensor") and isinstance(tensor, torch_npu.HiFloat8Tensor):
+        return True
+    return False
+
+
+def is_float8_tensor(tensor):
+    if str(tensor.dtype) in [Const.FLOAT8_E5M2_TYPE, Const.FLOAT8_E4M3FN_TYPE]:
+        return True
+    return is_hifloat8_tensor(tensor)
+
+
+def register_forward_pre_hook(module, forward_pre_hook):
+    if torch_version_above_or_equal_2:
+        module.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
+    else:
+        module.register_forward_pre_hook(forward_pre_hook)
+
+
+def register_forward_hook(module, forward_hook):
+    if torch_version_above_or_equal_2:
+        module.register_forward_hook(forward_hook, with_kwargs=True)
+    else:
+        module.register_forward_hook(forward_hook)

msprobe/pytorch/compare/distributed_compare.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-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");
@@ -13,41 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
-from msprobe.core.common.exceptions import FileCheckException
-from msprobe.core.common.file_utils import create_directory
-from msprobe.core.common.utils import CompareException, check_compare_param, check_configuration_param, get_dump_mode, \
-    set_dump_path
-from msprobe.core.compare.acc_compare import ModeConfig
-from msprobe.core.compare.utils import check_and_return_dir_contents, extract_json, set_stack_json_path
-from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.compare.pt_compare import PTComparator, compare
+from msprobe.core.compare.utils import compare_distributed_inner
+from msprobe.pytorch.compare.pt_compare import compare
 
 
 def compare_distributed(npu_dump_dir, bench_dump_dir, output_path, **kwargs):
-    if kwargs.get("suffix"):
-        logger.error("Argument 'suffix' is not supported for compare_distributed.")
-        raise CompareException(CompareException.INVALID_PARAM_ERROR)
-    is_print_compare_log = kwargs.get("is_print_compare_log", True)
-    # get the ranks and match by order
-    npu_ranks = sorted(check_and_return_dir_contents(npu_dump_dir, 'rank'))
-    bench_ranks = sorted(check_and_return_dir_contents(bench_dump_dir, 'rank'))
-    if len(npu_ranks) != len(bench_ranks):
-        logger.error(
-            "The number of ranks in the two runs are different. "
-            "Unable to match the ranks. "
-            "Please use another folder to compare or use compare() api and manually match the ranks.")
-        raise CompareException(CompareException.INVALID_PATH_ERROR)
-    for nr, br in zip(npu_ranks, bench_ranks):
-        npu_data_dir = os.path.join(npu_dump_dir, nr)
-        bench_data_dir = os.path.join(bench_dump_dir, br)
-        npu_path = extract_json(npu_data_dir, stack_json=False)
-        bench_path = extract_json(bench_data_dir, stack_json=False)
-
-        dump_result_param = {
-            "npu_json_path": npu_path,
-            "bench_json_path": bench_path,
-            "is_print_compare_log": is_print_compare_log
-        }
-        compare(input_param=dump_result_param, output_path=output_path, suffix=f'_{nr}-{br}', **kwargs)
+    compare_distributed_inner(npu_dump_dir, bench_dump_dir, output_path, compare, **kwargs)

msprobe/pytorch/compare/pt_compare.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");
@@ -13,92 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os.path
+from msprobe.core.compare.acc_compare import Comparator, ModeConfig, MappingConfig, setup_comparison
+from msprobe.pytorch.compare.utils import read_pt_data
 
-import torch
 
-from msprobe.core.common.const import FileCheckConst
-from msprobe.core.common.exceptions import FileCheckException
-from msprobe.core.common.file_utils import FileChecker, create_directory, load_yaml
-from msprobe.core.common.utils import CompareException, check_compare_param, check_configuration_param, get_dump_mode, \
-    set_dump_path
-from msprobe.core.compare.acc_compare import Comparator, ModeConfig
-from msprobe.core.compare.utils import set_stack_json_path
-from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.common.utils import load_pt
-
-
-class PTComparator(Comparator):
-    def __init__(self, mode_config, data_mapping=None):
-        super().__init__(mode_config)
-
-        self.stack_mode = mode_config.stack_mode
-        self.auto_analyze = mode_config.auto_analyze
-        self.fuzzy_match = mode_config.fuzzy_match
-        self.dump_mode = mode_config.dump_mode
-
-        self.frame_name = PTComparator.__name__
-        self.data_mapping = data_mapping
-        if isinstance(self.data_mapping, str) or self.data_mapping is None:
-            self.data_mapping_dict = self.load_mapping_file(self.data_mapping)
-        elif isinstance(self.data_mapping, dict):
-            self.data_mapping_dict = self.data_mapping
-        else:
-            raise TypeError(f"The type of parameter `data_mapping` must be dict, str or None, but got "
-                            f"{type(self.data_mapping)}")
-
-    @staticmethod
-    def load_mapping_file(mapping_file):
-        if isinstance(mapping_file, str):
-            mapping_dict = load_yaml(mapping_file)
-        else:
-            mapping_dict = {}
-        return mapping_dict
-
-    def read_npy_data(self, dir_path, file_name):
-        if not 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)
-        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}.")
-            raise CompareException(CompareException.INVALID_FILE_ERROR) from e
-        except AttributeError as e:
-            # 这里捕获 detach 方法抛出的异常
-            logger.error(f"Failed to detach the loaded tensor.")
-            raise CompareException(CompareException.DETACH_ERROR) from e
-        if data_value.dtype == torch.bfloat16:
-            data_value = data_value.to(torch.float32)
-        data_value = data_value.numpy()
-        return data_value
+def read_real_data(npu_dir, npu_data_name, bench_dir, bench_data_name, _) -> tuple:
+    n_value = read_pt_data(npu_dir, npu_data_name)
+    b_value = read_pt_data(bench_dir, bench_data_name)
+    return n_value, b_value
 
 
 def compare(input_param, output_path, **kwargs):
-    try:
-        auto_analyze = kwargs.get('auto_analyze', True)
-        fuzzy_match = kwargs.get('fuzzy_match', False)
-        data_mapping = kwargs.get('data_mapping', None)
-        suffix = kwargs.get('suffix', '')
-
-        set_dump_path(input_param)
-        dump_mode = get_dump_mode(input_param)
-        if "stack_json_path" in input_param:
-            stack_mode = kwargs.get('stack_mode', False)
-        else:
-            stack_mode = set_stack_json_path(input_param)  # set stack_mode and set "stack_json_path" in input_param
-        check_configuration_param(stack_mode, auto_analyze, fuzzy_match, input_param.get('is_print_compare_log', True))
-        create_directory(output_path)
-        check_compare_param(input_param, output_path, dump_mode, stack_mode)
-    except (CompareException, FileCheckException) as error:
-        logger.error('Compare failed. Please check the arguments and do it again!')
-        raise CompareException(error.code) from error
+    config = setup_comparison(input_param, output_path, **kwargs)
 
-    mode_config = ModeConfig(stack_mode, auto_analyze, fuzzy_match, dump_mode)
-    pt_comparator = PTComparator(mode_config, data_mapping)
-    pt_comparator.compare_core(input_param, output_path, suffix=suffix)
+    mode_config = ModeConfig(config.stack_mode, config.auto_analyze, config.fuzzy_match,
+                             config.dump_mode, config.compared_file_type)
+    mapping_config = MappingConfig(data_mapping=config.data_mapping)
+    pt_comparator = Comparator(read_real_data, mode_config, mapping_config)
+    pt_comparator.compare_core(input_param, output_path, suffix=config.suffix)

msprobe/pytorch/compare/utils.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2025-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
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import torch
+
+from msprobe.core.common.utils import logger, CompareException
+from msprobe.core.common.file_utils import FileChecker, FileCheckConst
+from msprobe.pytorch.common.utils import load_pt
+
+
+def read_pt_data(dir_path, file_name):
+    if not 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)
+    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}.")
+        raise CompareException(CompareException.INVALID_FILE_ERROR) from e
+    except AttributeError as e:
+        # 这里捕获 detach 方法抛出的异常
+        logger.error(f"Failed to detach the loaded tensor.")
+        raise CompareException(CompareException.DETACH_ERROR) from e
+    if data_value.dtype == torch.bfloat16:
+        data_value = data_value.to(torch.float32)
+    data_value = data_value.numpy()
+    return data_value

msprobe/pytorch/debugger/debugger_config.py

@@ -13,11 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
-
 from msprobe.core.common.const import Const
 from msprobe.core.common.exceptions import MsprobeException
 from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.common.utils import is_torch_nn_module
 
 
 class DebuggerConfig:
@@ -60,6 +59,7 @@ class DebuggerConfig:
                 if isinstance(task_config.online_run_ut_recompute, bool) else False
 
         self.check()
+        self._check_statistics_config(task_config)
 
         if self.level == Const.LEVEL_L2:
             self.is_backward_kernel_dump = False
@@ -78,10 +78,13 @@ class DebuggerConfig:
         if not isinstance(self.async_dump, bool):
             raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
                                    f"The parameters async_dump should be bool.")
-        if self.async_dump and self.task == Const.TENSOR and not self.list:
-            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
-                                   f"The parameters async_dump is true in tensor task, the parameters list cannot be "
-                                   f"empty.")
+        if self.async_dump and self.task == Const.TENSOR:
+            if self.level == Const.LEVEL_DEBUG:
+                self.list = [] # async_dump + debug level case ignore list
+            if not self.list and self.level != Const.LEVEL_DEBUG:
+                raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                                    f"The parameters async_dump is true in tensor task, the parameters list cannot be "
+                                    f"empty.")
         if self.task == Const.STRUCTURE and self.level not in [Const.LEVEL_L0, Const.LEVEL_MIX]:
             logger.warning_on_rank_0(
                 f"When the task is set to structure, the level should be one of {[Const.LEVEL_L0, Const.LEVEL_MIX]}. "
@@ -93,25 +96,24 @@ class DebuggerConfig:
         self.check_kwargs()
         return True
 
-    def check_model(self, instance, start_model):
-        if self.level not in [Const.LEVEL_L0, Const.LEVEL_MIX]:
-            if instance.model is not None or start_model is not None:
-                logger.info_on_rank_0(
-                    f"The current level is not L0 or mix level, so the model parameters will not be used.")
+    def check_model(self, instance, start_model, token_range=None):
+        instance.model = start_model if start_model is not None else instance.model
+        if self.level not in [Const.LEVEL_L0, Const.LEVEL_MIX] and token_range is None:
             return
-        if start_model is None and instance.model is None:
+
+        if instance.model is None:
             logger.error_on_rank_0(
-                f"For level {self.level}, PrecisionDebugger or start interface must receive a 'model' parameter.")
+                f"For level {self.level} or non-empty token_range, "
+                f"PrecisionDebugger or start interface must receive a 'model' parameter.")
             raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, f"missing the parameter 'model'")
 
-        instance.model = start_model if start_model is not None else instance.model
-        if isinstance(instance.model, torch.nn.Module):
+        if is_torch_nn_module(instance.model):
             return
 
         error_model = None
         if isinstance(instance.model, (list, tuple)):
             for model in instance.model:
-                if not isinstance(model, torch.nn.Module):
+                if not is_torch_nn_module(model):
                     error_model = model
                     break
         else:
@@ -119,7 +121,7 @@ class DebuggerConfig:
 
         if error_model is not None:
             error_info = (f"The 'model' parameter must be a torch.nn.Module or list[torch.nn.Module] "
-                          f"type, currently there is a {type(error_model)} type.")
+                          f"type, currently there is an unsupported {type(error_model)} type.")
             raise MsprobeException(
                 MsprobeException.INVALID_PARAM_ERROR, error_info)
 
@@ -130,8 +132,23 @@ class DebuggerConfig:
         if not self.list or len(self.list) != 1:
             raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
                                    f"When level is set to L2, the list must be configured as a list with one api name.")
+        if self.task != Const.TENSOR:
+            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                                   f"When level is set to L2, the task must be set to tensor.")
+
         api_name = self.list[0]
         if api_name.endswith(Const.BACKWARD):
             self.is_backward_kernel_dump = True
             api_forward_name = api_name[:-len(Const.BACKWARD)] + Const.FORWARD
             self.list.append(api_forward_name)
+
+    def _check_statistics_config(self, task_config):
+        if self.task != Const.STATISTICS:
+            return
+        self.tensor_list = []
+        if not hasattr(task_config, "tensor_list"):
+            return
+        if self.level == Const.LEVEL_DEBUG and task_config.tensor_list:
+            logger.warning_on_rank_0("When level is set to debug, the tensor_list will be invalid.")
+            return
+        self.tensor_list = task_config.tensor_list