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

msprobe/core/service.py

@@ -0,0 +1,356 @@
+# 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.
+
+
+from abc import ABC, abstractmethod
+import copy
+from collections import defaultdict
+import functools
+import os
+
+from msprobe.core.common.exceptions import DistributedNotInitializedError
+from msprobe.core.common.file_utils import create_directory
+from msprobe.core.common.runtime import Runtime
+from msprobe.core.common.utils import Const, print_tools_ends_info, DumpPathAggregation
+from msprobe.core.data_dump.api_registry import ApiRegistry
+from msprobe.core.data_dump.data_collector import build_data_collector
+from msprobe.core.hook_manager import BaseHookManager
+from msprobe.core.kernel_dump.kernel_config import create_kernel_config_json
+
+
+class BaseService(ABC):
+    def __init__(self, config):
+        self.config = copy.deepcopy(config)
+        self.config.level = getattr(config, 'level_ori', config.level)  # 兼容MindSpore配置
+        self.model = None
+        self.data_collector = build_data_collector(self.config)
+        self.attl_manager = None
+        self.current_iter = 0
+        self.loop = 0
+        self.init_step = 0
+        self.cur_token_id = 0
+        self.first_start = True
+        self.primitive_switch = False
+        self.current_rank = None
+        self.dump_iter_dir = None
+        self.should_stop_service = False
+        self.ori_customer_func = {}
+        self.debug_variable_counter = None
+        self.currrent_step_first_debug_save = True
+        self.logger = None  # 子类中注入
+        self.api_register = None  # 子类中注入
+        self.api_template = None  # 子类中注入
+        self.hook_manager = None  # 子类中注入
+        self._init_specific_components()
+        self._register_api_hook()
+
+    @property
+    def _is_debug_level(self):
+        return self.config.level == Const.LEVEL_DEBUG
+
+    @property
+    def _is_l2_level(self):
+        return self.config.level == Const.LEVEL_L2
+    
+    @property
+    def _is_mix_level(self):
+        return self.config.level == Const.LEVEL_MIX
+
+    @property
+    def _is_need_module_hook(self):
+        return self.config.level in [Const.LEVEL_MIX, Const.LEVEL_L0]
+    
+    @property
+    def _is_need_api_hook(self):
+        return self.config.level in [Const.LEVEL_MIX, Const.LEVEL_L1, Const.LEVEL_L2]
+
+    @property
+    def _is_no_dump_step(self):
+        return (self.config.step and self.current_iter not in self.config.step)
+
+    @property
+    def _is_no_dump_rank(self):
+        return (self.config.rank and self.current_rank not in self.config.rank)
+
+    @property
+    def _need_tensor_data(self):
+        """判断是否需要采集tensor数据"""
+        return bool(
+            self.config.task in self.data_collector.tasks_need_tensor_data or
+            (self.config.task == Const.STATISTICS and self.config.tensor_list)
+        )
+    
+    @property
+    def _is_online_run_ut(self):
+        return getattr(self.config, "online_run_ut", False)
+
+    @property
+    @abstractmethod
+    def _get_framework_type(self):
+        """获取框架类型"""
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def _get_current_rank():
+        """获取当前rank_id"""
+        pass
+
+    @staticmethod
+    def _change_jit_switch(status):
+        """修改JitDump开关，mindspore子类重写"""
+        pass 
+
+    def start(self, model=None, token_range=None):
+        """通用start模板"""
+        self._process_iteration()
+        if self._is_debug_level:
+            return
+        if self._need_stop_service():
+            return 
+        self.model = model
+        self.cur_token_id = 0
+        if self.first_start:
+            try:
+                self.current_rank = self._get_current_rank()
+            except DistributedNotInitializedError:
+                self.current_rank = None
+            Runtime.current_rank = self.current_rank
+            if self._is_no_dump_rank:
+                return
+            self._register_hook()
+            if self._is_need_module_hook:
+                self._register_module_hook()
+            self.first_start = False
+
+            if token_range:
+                self._register_infer_count_hook(self.model, token_range)
+        self.logger.info(f"{Const.TOOL_NAME}: debugger.start() is set successfully")
+        if token_range is None:
+            Runtime.is_running = True
+            self.primitive_switch = True
+            self._change_jit_switch(True)
+            self.logger.info(f"Dump switch is turned on at step {self.current_iter}. ")
+        if self._is_online_run_ut:
+            self._run_ut_dispatch(True)
+        else:
+            self.create_dirs()
+            self.logger.info(f"Dump data will be saved in {self.dump_iter_dir}.")
+
+    def stop(self):
+        """通用stop模板"""
+        if self._is_debug_level or self.should_stop_service:
+            return
+        if self._is_no_dump_step or self._is_no_dump_rank:
+            return
+        self.logger.info(f"{Const.TOOL_NAME}: debugger.stop() is set successfully. "
+                    "Please set debugger.start() to turn on the dump switch again. ")
+        Runtime.is_running = False
+        self.primitive_switch = False
+        self._change_jit_switch(False)
+        if self._is_l2_level:
+            return
+        if self._is_online_run_ut:
+            self._run_ut_dispatch(False)
+        self._process_async_dump()
+        self.data_collector.write_json()
+
+    def step(self):
+        """通用step处理"""
+        if self.should_stop_service:
+            return
+        self._process_async_dump()
+        self.data_collector.write_json()
+        self.currrent_step_first_debug_save = True
+        self.loop += 1
+        self._reset_status()
+
+    def save(self, variable, name, save_backward):
+        '''
+        Args:
+            variable: Union[List[variable], dict{str: variable}, mindspore.tensor, str, float, int]
+            name: str
+            save_backward: boolean
+        Return:
+            void
+        '''
+        if not self._is_debug_level:
+            return
+        self.current_iter = self.loop + self.init_step
+        if self._is_no_dump_step:
+            return
+
+        if self.currrent_step_first_debug_save:
+            try:
+                self.current_rank = self._get_current_rank()
+            except DistributedNotInitializedError:
+                self.current_rank = None
+
+            self.create_dirs()
+            self.debug_variable_counter = defaultdict(int)
+            self.currrent_step_first_debug_save = False
+
+        count = self.debug_variable_counter[name]
+        self.debug_variable_counter[name] += 1
+
+        name_with_count = f"{name}.{count}"
+        grad_name_with_count = f"{name}_grad.{count}"
+
+        # forward save
+        self.data_collector.debug_data_collect_forward(variable, name_with_count)
+
+        # backward save
+        if save_backward:
+            self.data_collector.debug_data_collect_backward(variable, grad_name_with_count)
+
+    def register_custom_api(self, module, api_name, api_prefix):
+        self.ori_customer_func[str(module) + Const.SEP + api_name] = getattr(module, api_name)
+        ApiRegistry.register_custom_api(module, api_name, api_prefix,
+                                          functools.partial(self.build_hook, Const.API), self.api_template)
+
+    def restore_custom_api(self, module, api):
+        ori_func = self.ori_customer_func.get(str(module) + Const.SEP + api)
+        if ori_func:
+            setattr(module, api, ori_func)
+
+
+    def build_hook(self, hook_type, name):
+        return self.hook_manager.build_hook(hook_type, name)
+
+    def create_dirs(self):
+        """统一目录创建逻辑"""
+        create_directory(self.config.dump_path)
+        if Runtime.run_mode == Const.PYNATIVE_GRAPH_MODE:
+            self.dump_iter_dir = os.path.join(self.config.dump_path, Const.PYNATIVE_MODE, f"step{self.current_iter}")
+        else:
+            self.dump_iter_dir = os.path.join(self.config.dump_path, f"step{self.current_iter}")
+
+        cur_rank = self.current_rank if self.current_rank is not None else ''
+        if self._is_l2_level:
+            self._create_l2_dirs(cur_rank)
+        else:
+            self._create_default_dirs(cur_rank)
+
+    @abstractmethod
+    def _init_specific_components(self):
+        """初始化框架特定组件"""
+        pass
+    
+    @abstractmethod
+    def _register_hook(self):
+        """注册hook函数"""
+        pass
+    
+    @abstractmethod
+    def _register_module_hook(self):
+        """注册模块级别的hook函数"""
+    
+    def _need_stop_service(self):
+        if self.should_stop_service:
+            return True
+        end_service = self.config.step and self.current_iter > max(self.config.step) or \
+                      self.data_collector and self.data_collector.data_processor.is_terminated
+        if end_service:
+            if self._is_online_run_ut and self.attl_manager:
+                self.attl_manager.attl_stop()
+            self.primitive_switch = False
+            self._change_jit_switch(False)
+            Runtime.is_running = False
+            self.should_stop_service = True
+            print_tools_ends_info()
+            return True
+        if self._is_no_dump_step:
+            return True
+        return False
+
+    def _register_api_hook(self):
+        if self._is_need_api_hook:
+            self.api_register.initialize_hook(functools.partial(self.build_hook, Const.API))
+            self.api_register.register_all_api()
+            self.logger.info(f"The api {self.config.task} hook function is successfully mounted to the model.")
+
+    def _register_infer_count_hook(self, root_model, token_range):
+        """
+        通过root_model执行的轮次来判断当前在第几个token
+        param root_model: 需要采集的推理模型
+        param token_range: [start, end], 采集infer的token循环范围，左右皆包含在内
+        return: None
+        """
+        def infer_hook(model, args):
+            if self.cur_token_id == token_range[0]:
+                Runtime.is_running = True
+                self.primitive_switch = True
+                self._change_jit_switch(True)
+                self.logger.info(f"Current token id: {self.cur_token_id}, start dump infer token.")
+            elif token_range[0] < self.cur_token_id <= token_range[1]:
+                self.logger.debug(f"Current token id: {self.cur_token_id}.")
+            elif self.cur_token_id == token_range[1] + 1:
+                Runtime.is_running = False
+                self.primitive_switch = False
+                self._change_jit_switch(False)
+                self.logger.info(
+                    f"Current token id: {self.cur_token_id}, exceed token_range, early stop dump infer token.")
+            self.cur_token_id += 1
+        if isinstance(root_model, list):
+            root_model = root_model[0]
+            self.logger.warning("Infer model can only input one to support token_range, choose the first one.")
+        if self._is_online_run_ut:
+            return
+        root_model.register_forward_pre_hook(infer_hook)
+ 
+    def _create_l2_dirs(self, cur_rank):
+        create_directory(self.dump_iter_dir)
+        kernel_config_path = create_kernel_config_json(self.dump_iter_dir, cur_rank)
+        self.config.kernel_config_path = kernel_config_path
+
+    def _create_default_dirs(self, cur_rank):
+        dump_dir = os.path.join(self.dump_iter_dir, f"rank{cur_rank}")
+        create_directory(dump_dir)
+        
+        dump_data_dir = None
+        if self._need_tensor_data:
+            dump_data_dir = os.path.join(dump_dir, "dump_tensor_data")
+            create_directory(dump_data_dir)
+            
+        self._configure_dump_paths(dump_dir, dump_data_dir)
+
+    def _configure_dump_paths(self, dump_dir, dump_data_dir):
+        dump_path_aggregation = DumpPathAggregation()
+        dump_path_aggregation.dump_file_path = os.path.join(dump_dir, "dump.json")
+        dump_path_aggregation.stack_file_path = os.path.join(dump_dir, "stack.json")
+        dump_path_aggregation.construct_file_path = os.path.join(dump_dir, "construct.json")
+        dump_path_aggregation.dump_tensor_data_dir = dump_data_dir
+        dump_path_aggregation.debug_file_path = os.path.join(dump_dir, "debug.json")
+        dump_path_aggregation.free_benchmark_file_path = os.path.join(dump_dir, "free_benchmark.csv")
+        self.data_collector.update_dump_paths(dump_path_aggregation)
+        self.data_collector.initialize_json_file(self._get_framework_type)
+
+    def _process_iteration(self):
+        """处理迭代计数"""
+        self.current_iter = self.loop + self.init_step
+        self.data_collector.update_iter(self.current_iter)
+        Runtime.current_iter = self.current_iter
+
+    def _process_async_dump(self):
+        """处理异步dump逻辑"""
+        if self.config.async_dump and self.config.task in [Const.STATISTICS, Const.TENSOR]:
+            self.data_collector.data_processor.dump_async_data()
+
+    def _reset_status(self):
+        """通用状态重置"""
+        self.data_collector.reset_status()
+        BaseHookManager.params_grad_info.clear()
+        if self._is_l2_level:
+            self.data_collector.data_processor.reset_status()

msprobe/core/single_save/single_comparator.py

@@ -0,0 +1,243 @@
+# 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 multiprocessing
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+from msprobe.core.common.file_utils import check_file_or_directory_path, create_directory, save_excel
+from msprobe.core.common.log import logger
+
+
+@dataclass
+class CompareResult:
+    max_abs_error: float
+    max_relative_error: float
+    same_percentage: float
+    first_mismatch_index: int
+    percentage_within_thousandth: float
+    percentage_within_hundredth: float
+
+
+class SingleComparator:
+    result_header = [
+        'step', 
+        'rank', 
+        'micro_step', 
+        'id', 
+        'shape1', 
+        'shape2', 
+        '相同元素百分比(%)', 
+        '首个不匹配元素索引', 
+        '最大绝对误差', 
+        '最大相对误差', 
+        '误差在千分之一内元素占比(%)', 
+        '误差在百分之一内元素占比(%)'
+    ]
+
+    @classmethod
+    def compare(cls, dir1, dir2, output_path="./msprobe_compare_output", num_processes=8):
+        data_dir1 = os.path.join(dir1, "data")
+        data_dir2 = os.path.join(dir2, "data")
+        check_file_or_directory_path(data_dir1, isdir=True)
+        check_file_or_directory_path(data_dir2, isdir=True)
+        # 确保输出目录存在，如果不存在则创建
+        if not os.path.exists(output_path):
+            create_directory(output_path)
+        cls.compare_data(data_dir1, data_dir2, output_path, num_processes)
+
+    @classmethod
+    def compare_arrays(cls, array1, array2) -> CompareResult:
+        """
+        比较两个NumPy数组，计算最大绝对误差、最大相对误差和相同元素的百分比
+        """
+        # 计算每个维度上的最小尺寸
+        min_shape = [min(s1, s2) for s1, s2 in zip(array1.shape, array2.shape)]
+        # 截取数组到相同的形状
+        sliced_array1 = array1[tuple(slice(0, s) for s in min_shape)]
+        sliced_array2 = array2[tuple(slice(0, s) for s in min_shape)]
+
+        abs_error = np.abs(sliced_array1 - sliced_array2)
+        max_abs_error = np.max(abs_error)
+        
+        # 计算相对误差，处理分母为零的情况
+        with np.errstate(divide='ignore', invalid='ignore'):
+            relative_error = np.abs(sliced_array1 - sliced_array2) / \
+                np.maximum(np.abs(sliced_array1), np.abs(sliced_array2))
+        relative_error = np.nan_to_num(relative_error)
+        max_relative_error = np.max(relative_error)
+
+        same_elements = np.sum(sliced_array1 == sliced_array2)
+        total_elements = sliced_array1.size
+        same_percentage = (same_elements / total_elements) * 100
+
+        # 展平数组
+        flat_array1 = sliced_array1.flatten()
+        flat_array2 = sliced_array2.flatten()
+
+        # 计算从第几个元素开始对不上
+        mismatch_indices = np.nonzero(flat_array1 != flat_array2)[0]
+        first_mismatch_index = mismatch_indices[0] if mismatch_indices.size > 0 else None
+
+        # 计算误差在千分之一内的元素占比
+        threshold = 0.001 * np.maximum(np.abs(sliced_array1), np.abs(sliced_array2))
+        error_within_thousandth = np.sum(abs_error <= threshold)
+        percentage_within_thousandth = (error_within_thousandth / total_elements) * 100
+
+        # 计算误差在百分之一内的元素占比
+        threshold = 0.01 * np.maximum(np.abs(sliced_array1), np.abs(sliced_array2))
+        error_within_hundredth = np.sum(abs_error <= threshold)
+        percentage_within_hundredth = (error_within_hundredth / total_elements) * 100
+
+        return CompareResult(
+            max_abs_error,
+            max_relative_error,
+            same_percentage,
+            first_mismatch_index,
+            percentage_within_thousandth,
+            percentage_within_hundredth
+        )
+
+    @classmethod
+    def get_steps(cls, tag_path):
+        for step_folder in os.listdir(tag_path):
+            if step_folder.startswith('step'):
+                try:
+                    step = int(step_folder[4:])
+                except Exception as e:
+                    raise RuntimeError(f"parse step number error") from e
+                yield step, os.path.join(tag_path, step_folder)
+
+    @classmethod
+    def get_ranks(cls, step_path):
+        for rank_folder in os.listdir(step_path):
+            if rank_folder.startswith('rank'):
+                try:
+                    rank = int(rank_folder[4:])
+                except Exception as e:
+                    raise RuntimeError(f"parse rank number error") from e
+                yield rank, os.path.join(step_path, rank_folder)
+
+    @classmethod
+    def get_micro_steps(cls, rank_path):
+        for micro_step_folder in os.listdir(rank_path):
+            if micro_step_folder.startswith('micro_step'):
+                try:
+                    micro_step = int(micro_step_folder[10:])
+                except Exception as e:
+                    raise RuntimeError(f"parse nicro_step number error") from e
+                yield micro_step, os.path.join(rank_path, micro_step_folder)
+            else:
+                yield 0, rank_path
+
+    @classmethod
+    def get_arrays(cls, micro_step_path):
+        for file in os.listdir(micro_step_path):
+            if file.endswith('.npy'):
+                try:
+                    parts = file.rsplit('.', 2)
+                    if len(parts) > 1 and parts[-2].isdigit():
+                        array_id = int(parts[-2])
+                    else:
+                        array_id = 0
+                except ValueError:
+                    array_id = 0
+                yield array_id, os.path.join(micro_step_path, file)
+
+    @classmethod
+    def get_array_paths(cls, dir_path):
+        """
+        获取目录中所有符合结构的NumPy数组文件路径
+        """
+        array_paths = {}
+        if not os.path.exists(dir_path):
+            return array_paths
+        for tag in os.listdir(dir_path):
+            tag_path = os.path.join(dir_path, tag)
+            if not os.path.isdir(tag_path):
+                continue
+            for step, step_path in cls.get_steps(tag_path):
+                for rank, rank_path in cls.get_ranks(step_path):
+                    for micro_step, micro_step_path in cls.get_micro_steps(rank_path):
+                        for array_id, array_path in cls.get_arrays(micro_step_path):
+                            array_paths.setdefault(tag, []).append((step, rank, micro_step, array_id, array_path))
+        return array_paths
+
+    @classmethod
+    def compare_single_tag(cls, tag, array_paths1, array_paths2, output_dir):
+        try:
+            data = []
+            paths1 = array_paths1.get(tag, [])
+            paths2 = array_paths2.get(tag, [])
+            path_dict1 = {(step, rank, micro_step, array_id): path for step, rank, micro_step, array_id, path in paths1}
+            path_dict2 = {(step, rank, micro_step, array_id): path for step, rank, micro_step, array_id, path in paths2}
+            common_keys = set(path_dict1.keys()) & set(path_dict2.keys())
+            for key in common_keys:
+                try:
+                    array1 = np.load(path_dict1[key])
+                    array2 = np.load(path_dict2[key])
+                    result = cls.compare_arrays(array1, array2)
+                    step, rank, micro_step, array_id = key
+                    data.append([
+                        step, rank, micro_step, array_id, 
+                        list(array1.shape), list(array2.shape), 
+                        result.same_percentage, 
+                        result.first_mismatch_index, 
+                        result.max_abs_error, 
+                        result.max_relative_error, 
+                        result.percentage_within_thousandth, 
+                        result.percentage_within_hundredth
+                    ])
+                except Exception as e:
+                    logger.error(f"Error comparing {path_dict1[key]} and {path_dict2[key]}: {e}")
+
+            df = pd.DataFrame(data, columns=SingleComparator.result_header)
+            df = df.sort_values(by=['step', 'rank', 'micro_step', 'id'])
+            # 构建输出文件的完整路径
+            output_file_path = os.path.join(output_dir, f'{tag}.xlsx')
+            save_excel(output_file_path, df)
+        except Exception as e:
+            logger.error(f"Error processing tag {tag}: {e}")
+
+    @classmethod
+    def compare_data(cls, dir1, dir2, output_dir, num_processes=8):
+        """
+        比较两个目录中的NumPy数组文件，并将结果保存到指定目录的Excel文件中
+        """
+
+        array_paths1 = cls.get_array_paths(dir1)
+        array_paths2 = cls.get_array_paths(dir2)
+
+        all_tags = set(array_paths1.keys()) | set(array_paths2.keys())
+
+        with multiprocessing.Pool(processes=num_processes) as pool:
+            args = [(tag, array_paths1, array_paths2, output_dir) for tag in all_tags]
+            try:
+                results = pool.starmap_async(cls.compare_single_tag, args)
+                with tqdm(total=len(all_tags), desc="Processing data") as pbar:
+                    while not results.ready():
+                        pbar.n = len(all_tags) - results._number_left
+                        pbar.refresh()
+                    results.wait()
+                    results.get()
+            except Exception as e:
+                logger.error(f"Multiprocessing error: {e}")
+            finally:
+                pool.close()
+                pool.join()