mindstudio-probe 8.3.2__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/pytorch/free_benchmark/result_handlers/base_handler.py

@@ -1,252 +0,0 @@
-# Copyright (c) 2024-2024, 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 math
-from abc import ABC, abstractmethod
-from typing import Any, Optional, Tuple
-
-import numpy as np
-import torch
-from msprobe.core.common.const import Const
-from msprobe.core.common.exceptions import FreeBenchmarkException
-from msprobe.pytorch.free_benchmark import logger
-from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
-from msprobe.pytorch.free_benchmark.common.enums import (
-    FuzzThreshold,
-    NormType,
-    PerturbationMode,
-)
-from msprobe.pytorch.free_benchmark.common.params import (
-    DataParams,
-    HandlerParams,
-    make_unequal_row,
-)
-from msprobe.pytorch.free_benchmark.common.utils import Tools, TorchC
-
-
-class FuzzHandler(ABC):
-    def __init__(self, params: HandlerParams) -> None:
-        self.params = params
-        self.unequal_rows = []
-
-    @staticmethod
-    def pre_process(origin_ouput, perturbed_output):
-        if (
-            isinstance(origin_ouput, tuple)
-            and hasattr(origin_ouput, "values")
-            and hasattr(origin_ouput, "indices")
-        ):
-            origin_ouput = origin_ouput.values
-            perturbed_output = perturbed_output.values
-        if hasattr(perturbed_output, "dtype"):
-            abs_tol = ThresholdConfig.ABS_TOL_VALUE_DICT.get(
-                perturbed_output.dtype, FuzzThreshold.F32_THD
-            )
-        else:
-            abs_tol = FuzzThreshold.F32_THD
-        return (
-            origin_ouput.to(perturbed_output.dtype).to(perturbed_output.device),
-            perturbed_output,
-            abs_tol,
-        )
-
-    @staticmethod
-    def tensor_split_for_error_calculate(origin_output, perturbed_output):
-        """
-        对将投入误差值计算的扰动前后输出张量进行分块
-        :param origin_output: 原始输出
-        :param perturbed_output: 扰动后输出
-        :return origin_output_chunks: 切块后原始输出列表
-        :return perturbed_output_chunks: 切块后扰动后输出列表
-        """
-        single_output_mem = (
-            origin_output.element_size() * origin_output.nelement() / Const.ONE_MB
-        )
-        if single_output_mem == 0 or origin_output.ndim == 0:
-            return [origin_output], [perturbed_output]
-        # 张量大小和批数之间的关系：chunks_exp=math.log(M,2)-4, chunks=2**chunks_exp (M为对比张量数据大小[Mb])
-        chunks_exp = int(math.log(single_output_mem, 2)) - 4
-        chunks = 2**chunks_exp
-        chunks = max(chunks, 1)
-        chunks = min(chunks, ThresholdConfig.TENSOR_SPLIT_MAX_CHUNK)
-        origin_output_chunks = TorchC.tensor_split(
-            TorchC.reshape(origin_output, (-1,)), chunks
-        )
-        perturbed_output_chunks = TorchC.tensor_split(
-            TorchC.reshape(perturbed_output, (-1,)), chunks
-        )
-        return origin_output_chunks, perturbed_output_chunks
-
-    @abstractmethod
-    def get_threshold(self, dtype):
-        pass
-
-    @abstractmethod
-    def handle(self, data_params: DataParams) -> Any:
-        pass
-
-    def get_ratio_from_specific_norm(
-        self, origin_output, perturbed_output, norm_type, abs_tol
-    ):
-        if norm_type == NormType.ENDLESS_NORM:
-            return self.calculate_max_ratio(origin_output, perturbed_output, abs_tol)
-        return ThresholdConfig.COMP_CONSISTENT
-
-    def calculate_max_ratio(self, origin_output, perturbed_output, abs_tol):
-        origin_output_chunks, perturbed_output_chunks = (
-            self.tensor_split_for_error_calculate(origin_output, perturbed_output)
-        )
-        if len(origin_output_chunks) != len(perturbed_output_chunks):
-            err_msg = (
-                f"For {self.params.api_name}, the number of compare tensor chunks is different: "
-                f"{len(origin_output_chunks)} != {len(perturbed_output_chunks)}. please check!"
-            )
-            raise FreeBenchmarkException(
-                FreeBenchmarkException.OutputIndexError, err_msg
-            )
-
-        max_ratio = ThresholdConfig.COMP_CONSISTENT
-        for i, chunk_origin in enumerate(origin_output_chunks):
-            if chunk_origin.nelement() == 0:
-                break
-            chunk_perturbed = perturbed_output_chunks[i]
-            # 如果乘积最小值 < 极小值乘积的负值，认为存在非极小值符号相反的情况
-            if TorchC.lt(
-                TorchC.min(TorchC.mul(chunk_origin, chunk_perturbed)), -(abs_tol**2)
-            ):
-                return ThresholdConfig.SYMBOL_FLIPPING
-            # 求A/B B/A的比值前，将值限制在大于极小值范围内
-            clamp_origin = TorchC.clamp(TorchC.abs(chunk_origin), min=abs_tol)
-            clamp_perturbed = TorchC.clamp(TorchC.abs(chunk_perturbed), min=abs_tol)
-            # 对于计算结果为nan的情况，认为两者没有差异
-            ratio_tensor = TorchC.nan_to_num(
-                TorchC.div(clamp_origin, clamp_perturbed),
-                nan=ThresholdConfig.COMP_CONSISTENT,
-            )
-            # 求A/B 和 B/A比值最大值，其中 B/A的最大值为 A/B的最小值的倒数
-            min_ratio, max_ratio = TorchC.stack([*TorchC.aminmax(ratio_tensor)]).tolist()
-            min_ratio_reciprocal = np.inf if min_ratio == 0 else 1 / min_ratio
-            max_ratio = max(max_ratio, min_ratio_reciprocal)
-        return max_ratio
-
-    def ratio_calculate(self, origin_output, perturbed_output, norm_type) -> float:
-        try:
-            origin_output, perturbed_output, abs_tol = self.pre_process(
-                origin_output, perturbed_output
-            )
-        except Exception as e:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when computing ratio,"
-                f" y1 or y2 dtype is not supported {e}"
-            )
-            return ThresholdConfig.COMP_NAN
-        if self.params.fuzz_stage == Const.BACKWARD:
-            abs_tol = ThresholdConfig.BACKWARD_OUTPUT_LOWER_BOUND
-        else:
-            abs_tol = abs_tol**0.5
-        return self.get_ratio_from_specific_norm(
-            origin_output, perturbed_output, norm_type, abs_tol
-        )
-
-    def npu_compare(
-        self, origin_output, perturbed_output
-    ) -> Tuple[bool, Optional[float]]:
-
-        if isinstance(perturbed_output, int):
-            return origin_output == perturbed_output, None
-        elif isinstance(perturbed_output, float):
-            if perturbed_output == 0:
-                origin_output += FuzzThreshold.F32_THD
-                perturbed_output += FuzzThreshold.F32_THD
-            return (
-                math.isclose(origin_output, perturbed_output),
-                origin_output / perturbed_output,
-            )
-        elif not isinstance(perturbed_output, torch.Tensor):
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name} "
-                f"The compare for output type {type(perturbed_output)} is not supported"
-            )
-            return True, 1
-
-        threshold = self.get_threshold(Tools.get_first_tensor_dtype(origin_output))
-        ratio = self.ratio_calculate(
-            origin_output, perturbed_output, norm_type=NormType.ENDLESS_NORM
-        )
-        if threshold == 0:
-            raise ValueError("Threshold cannot be zero. Check `get_threshold` implementation.")
-        if ratio == ThresholdConfig.SYMBOL_FLIPPING:
-            is_consistent = False
-        else:
-            is_consistent = threshold >= ratio >= 1 / threshold
-        return is_consistent, ratio
-
-    def cmp_output_npu(self, data_params: DataParams):
-        npu_consistent = True
-        max_fuzz_ratio = 0
-        try:
-            if isinstance(data_params.original_result, torch.Tensor):
-                is_consistent, ratio = self.npu_compare(
-                    data_params.original_result, data_params.perturbed_result
-                )
-                npu_consistent = is_consistent
-                max_fuzz_ratio = (
-                    max_fuzz_ratio
-                    if not isinstance(ratio, (int, float))
-                    else max(max_fuzz_ratio, ratio)
-                )
-                data_params.is_consistent = is_consistent
-                if not is_consistent:
-                    self.unequal_rows.append(
-                        make_unequal_row(data_params, self.params, ratio=ratio)
-                    )
-
-            elif isinstance(data_params.original_result, (list, tuple)):
-                for index_, origin_item in enumerate(data_params.original_result):
-                    is_consistent, ratio = self.npu_compare(
-                        origin_item, data_params.perturbed_result[index_]
-                    )
-                    npu_consistent = npu_consistent and is_consistent
-                    max_fuzz_ratio = (
-                        max_fuzz_ratio
-                        if not isinstance(ratio, (int, float))
-                        else max(max_fuzz_ratio, ratio)
-                    )
-                    data_params.is_consistent = is_consistent
-                    if not is_consistent:
-                        self.unequal_rows.append(
-                            make_unequal_row(
-                                data_params, self.params, ratio=ratio, index=index_
-                            )
-                        )
-        except Exception as e:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when campare the result exception raise {e}"
-            )
-        return npu_consistent, max_fuzz_ratio
-
-    def get_unequal_rows(self):
-        return self.unequal_rows
-
-    def _get_default_threshold(self, dtype):
-        if self.params.pert_mode == PerturbationMode.NO_CHANGE:
-            threshold = ThresholdConfig.COMP_CONSISTENT
-        else:
-            threshold = ThresholdConfig.DTYPE_PER_THD.get(
-                dtype, ThresholdConfig.DTYPE_PER_THD.get(torch.float32)
-            )
-        return threshold

msprobe/pytorch/free_benchmark/result_handlers/check_handler.py

@@ -1,54 +0,0 @@
-# Copyright (c) 2024-2024, 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 typing import Any
-
-from msprobe.pytorch.free_benchmark import logger
-from msprobe.pytorch.free_benchmark.common.enums import DeviceType
-from msprobe.pytorch.free_benchmark.common.params import DataParams, make_unequal_row
-from msprobe.pytorch.free_benchmark.common.utils import Tools
-from msprobe.pytorch.free_benchmark.compare.single_benchmark import SingleCompare
-from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
-
-
-class CheckerHandler(FuzzHandler):
-    def other_compare(self, data_params: DataParams) -> bool:
-        is_consistent = SingleCompare().compare_seq(
-                    data_params.original_result, data_params.perturbed_result
-                )
-        if not is_consistent:
-            self.unequal_rows.append(
-                make_unequal_row(data_params, self.params)
-            )
-
-    def get_threshold(self, dtype):
-        return self._get_default_threshold(dtype)
-
-    def handle(self, data_params: DataParams) -> Any:
-        if isinstance(data_params.perturbed_result, bool) or not Tools.is_float_tensor(
-            data_params.perturbed_result
-        ):
-            return data_params.original_result
-        try:
-            if self.params.fuzz_device == DeviceType.NPU:
-                self.cmp_output_npu(data_params)
-            else:
-                self.other_compare(data_params)
-        except Exception as e:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when comparing the results, an exception is raised: {e}"
-            )
-        return data_params.original_result

msprobe/pytorch/free_benchmark/result_handlers/fix_handler.py

@@ -1,40 +0,0 @@
-# Copyright (c) 2024-2024, 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 typing import Any
-
-from msprobe.core.common.exceptions import FreeBenchmarkException
-from msprobe.pytorch.free_benchmark import logger
-from msprobe.pytorch.free_benchmark.common.params import DataParams
-from msprobe.pytorch.free_benchmark.common.utils import Tools
-from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
-
-
-class FixHandler(FuzzHandler):
-
-    def get_threshold(self, dtype):
-        return self._get_default_threshold(dtype)
-
-    def handle(self, data_params: DataParams) -> Any:
-        try:
-            return Tools.convert_fuzz_output_to_origin(
-                data_params.original_result, data_params.perturbed_result
-            )
-        except FreeBenchmarkException as e:
-            logger.warning(
-                f"[msprobe] Free Benchmark: For {self.params.api_name} "
-                f"Fix output failed because of: \n{e}"
-            )
-            return data_params.original_result

msprobe/pytorch/free_benchmark/result_handlers/handler_factory.py

@@ -1,45 +0,0 @@
-# Copyright (c) 2024-2024, 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 msprobe.pytorch.free_benchmark import FreeBenchmarkException
-from msprobe.pytorch.free_benchmark.common.constant import PreheatConfig
-from msprobe.pytorch.free_benchmark.common.enums import HandlerType
-from msprobe.pytorch.free_benchmark.common.params import HandlerParams
-from msprobe.pytorch.free_benchmark.result_handlers.check_handler import CheckerHandler
-from msprobe.pytorch.free_benchmark.result_handlers.preheat_handler import PreheatHandler
-from msprobe.pytorch.free_benchmark.result_handlers.fix_handler import FixHandler
-
-
-class FuzzHandlerFactory:
-
-    result_handlers = {
-        HandlerType.CHECK: CheckerHandler,
-        HandlerType.FIX: FixHandler,
-        HandlerType.PREHEAT: PreheatHandler,
-    }
-
-    @staticmethod
-    def create(params: HandlerParams):
-        if_preheat = params.preheat_config.get(PreheatConfig.IF_PREHEAT)
-        if not if_preheat:
-            handler = FuzzHandlerFactory.result_handlers.get(params.handler_type)
-        else:
-            handler = FuzzHandlerFactory.result_handlers.get(HandlerType.PREHEAT)
-        if not handler:
-            raise FreeBenchmarkException(
-                FreeBenchmarkException.UnsupportedType,
-                f"无标杆工具支持 [ {HandlerType.CHECK}、{HandlerType.FIX}] 形式",
-            )
-        return handler(params)

msprobe/pytorch/free_benchmark/result_handlers/preheat_handler.py

@@ -1,181 +0,0 @@
-# Copyright (c) 2024-2024, 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 math
-from typing import Any
-
-from msprobe.pytorch.free_benchmark import logger
-from msprobe.pytorch.free_benchmark.common.constant import ThresholdConfig
-from msprobe.pytorch.free_benchmark.common.counter import preheat_counter
-from msprobe.pytorch.free_benchmark.common.enums import DeviceType
-from msprobe.pytorch.free_benchmark.common.params import DataParams, HandlerParams
-from msprobe.pytorch.free_benchmark.common.utils import Tools
-from msprobe.pytorch.free_benchmark.compare.single_benchmark import SingleCompare
-from msprobe.pytorch.free_benchmark.result_handlers.base_handler import FuzzHandler
-
-
-class PreheatHandler(FuzzHandler):
-
-    def __init__(self, params: HandlerParams) -> None:
-        super().__init__(params)
-        self.pure_name = Tools.get_pure_api_name(self.params.api_name)
-
-    def get_threshold(self, dtype):
-        return preheat_counter.get_api_thd(self.pure_name, dtype)
-
-    def compare_npu_and_cpu(self, data_params: DataParams):
-        args = Tools.convert_device_and_dtype(
-            data_params.args, DeviceType.CPU, change_dtype=True
-        )
-        kwargs = Tools.convert_device_and_dtype(
-            data_params.kwargs, DeviceType.CPU, change_dtype=True
-        )
-        cpu_result = data_params.origin_func(*args, **kwargs)
-        return SingleCompare().compare_seq(data_params.original_result, cpu_result)
-
-    def preheat(self, max_fuzz_ratio, cpu_consistent, first_dtype):
-        # 存储当前step所有输出比值和对应npu\cpu比对结果
-        preheat_counter.update_preheat_record(
-            self.pure_name,
-            first_dtype,
-            (max_fuzz_ratio, cpu_consistent),
-        )
-        if self._need_adjust_threshold():
-            self._adjust_threshold()
-
-    def handle(self, data_params: DataParams) -> Any:
-
-        if isinstance(data_params.perturbed_result, bool) or not Tools.is_float_tensor(
-                data_params.perturbed_result
-        ):
-            return data_params.original_result
-
-        if self.params.step == 0:
-            preheat_counter.add_one_step_used_api(self.pure_name)
-            return data_params.original_result
-
-        # 如果当前api,step需要预热
-        npu_consistent, max_fuzz_ratio = self.cmp_output_npu(data_params)
-        data_params.is_consistent = npu_consistent
-
-        preheat_counter.check_step(self.params.step)
-
-        if self.params.preheat_config.get("preheat_step") <= self.params.step:
-            return data_params.original_result
-
-        preheat_counter.add_api_called_time(self.pure_name)
-
-        if not self._is_take_a_sample():
-            return data_params.original_result
-
-        cpu_consistent = True
-        try:
-            cpu_consistent = self.compare_npu_and_cpu(data_params)
-        except Exception as e:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"when campare to cpu exception raise {e}"
-            )
-        try:
-            first_dtype = Tools.get_first_tensor_dtype(data_params.original_result)
-        except RuntimeError:
-            logger.warning_on_rank_0(
-                f"[msprobe] Free Benchmark: For {self.params.api_name}, "
-                f"the output sequence does not contain tensors."
-            )
-        if preheat_counter.get_api_preheat(self.pure_name, str(first_dtype)):
-            self.preheat(max_fuzz_ratio, cpu_consistent, first_dtype)
-
-        return data_params.original_result
-
-    def _is_take_a_sample(self) -> bool:
-        need_sample_set = self._get_need_sample_set()
-        curr_called_seq = preheat_counter.get_api_called_time(self.pure_name)
-        res = curr_called_seq in need_sample_set
-        if res:
-            total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-            logger.info_on_rank_0(
-                f"[msprobe] Free benchmark: preheat sample in step{self.params.step}"
-                f"api_name {self.params.api_name}, "
-                f"curr_called_seq: {curr_called_seq}/{total_count}"
-            )
-            preheat_counter.add_api_sample_time(self.pure_name)
-        return res
-
-    def _get_sample_count_per_step(self) -> set:
-        """
-        每一个step中应该采集的样本数
-        """
-        total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-        preheat_step = self.params.preheat_config.get("preheat_step")
-        max_sample = self.params.preheat_config.get("max_sample")
-        return min(math.ceil(total_count / preheat_step), max_sample)
-
-    def _get_need_sample_set(self):
-        """
-        需要采集的api集合
-        """
-        # 每一步样本数
-        total_count = preheat_counter.get_one_step_used_api(self.pure_name)
-        need_sample_set = set()
-        if total_count == 0:
-            return need_sample_set
-        sample_count_per_step = self._get_sample_count_per_step()
-        prehead_step = self.params.preheat_config.get("preheat_step")
-        for i in range(1, sample_count_per_step + 1):
-            count = (prehead_step * (i - 1) + self.params.step) % total_count
-            if count == 0:
-                count = total_count
-            need_sample_set.add(count)
-        return need_sample_set
-
-    def _need_adjust_threshold(self) -> bool:
-        sample_count_per_step = self._get_sample_count_per_step()
-        sampled_time = preheat_counter.get_api_sample_time(self.pure_name)
-        res = sampled_time >= sample_count_per_step
-        return res
-
-    def _adjust_threshold_for_dtype(self, dtype_str, compare_result):
-        con_ratio = [ratio for ratio, is_consistent in compare_result if is_consistent]
-        incon_ratio = [ratio for ratio, is_consistent in compare_result if not is_consistent]
-        old_thd = preheat_counter.get_api_thd(self.pure_name, dtype_str)
-        new_thd = old_thd
-        # 正例负例都存在
-        if con_ratio and incon_ratio:
-            if min(incon_ratio) > max(con_ratio):
-                new_thd = min(min(incon_ratio), old_thd)
-                preheat_counter.set_api_preheat(self.pure_name, dtype_str, is_preheat=False)
-        elif con_ratio:
-            # 存在漏报
-            if max(con_ratio) > old_thd:
-                new_thd = 1 + ((old_thd - 1) * ThresholdConfig.API_THD_STEP)
-            else:
-                new_thd = 1 + ((old_thd - 1) / ThresholdConfig.API_THD_STEP)
-        else:
-            new_thd = min(min(incon_ratio), old_thd)
-            preheat_counter.set_api_preheat(self.pure_name, dtype_str, is_preheat=False)
-        return new_thd
-
-    def _adjust_threshold(self):
-        for dtype_str, compare_result in preheat_counter.preheat_record[
-            self.pure_name
-        ].items():
-            new_thd = self._adjust_threshold_for_dtype(dtype_str, compare_result)
-            threshold = self._get_default_threshold(
-                preheat_counter.dtype_map.get(dtype_str)
-            )
-            preheat_counter.update_api_thd(
-                self.pure_name, dtype_str, new_thd, threshold
-            )

msprobe/pytorch/grad_probe/grad_monitor.py

@@ -1,108 +0,0 @@
-# Copyright (c) 2024-2024, 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 collections import defaultdict
-
-import torch
-from msprobe.core.common.file_utils import remove_path, save_npy, write_csv, create_directory
-from msprobe.core.grad_probe.constant import level_adp
-from msprobe.core.grad_probe.utils import check_numeral_list_ascend, data_in_list_target
-from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.common.utils import get_rank_id, print_rank_0
-from msprobe.pytorch.grad_probe.grad_stat_csv import GradStatCsv
-
-if int(torch.__version__.split('.')[0]) >= 2:
-    from torch.optim.optimizer import register_optimizer_step_pre_hook
-
-
-class GradientMonitor:
-
-    def __init__(self, common_config, task_config):
-        level = task_config.grad_level
-        if level not in level_adp:
-            raise Exception(f"level is valid, not in {level_adp.keys()}")
-        self._level_adp = level_adp[level]
-        self._param_list = task_config.param_list
-        self._target_ranks = common_config.rank
-        logger.info(f"target rank {self._target_ranks}")
-        self._target_step = common_config.step
-        logger.info(f"target step {self._target_step}")
-        self._bounds = task_config.bounds
-        check_numeral_list_ascend(self._bounds)
-        self._output_path = common_config.dump_path
-        if not os.path.exists(self._output_path):
-            create_directory(self._output_path)
-        else:
-            logger.warning(f"the file in {self._output_path} will be deleted")
-        self._step = -1
-        self._param2name = defaultdict(str)
-
-    @property
-    def output_path(self):
-        return self._output_path
-
-    @staticmethod
-    def save_grad_direction(param_name, grad, save_path):
-        if not os.path.exists(save_path):
-            create_directory(save_path)
-        param_grad = grad.clone().detach()
-        is_positive = param_grad > 0
-        save_filepath = os.path.join(save_path, f"{param_name}.npy")
-        save_npy(is_positive.cpu().numpy(), save_filepath)
-
-    def monitor(self, model):
-        print_rank_0("> parameter names:")
-        for name, param in model.named_parameters():
-            self._param2name[param] = name
-            print_rank_0(f"\t{name}")
-        setattr(self, "_rank", get_rank_id())
-        if torch.distributed.is_initialized() and not data_in_list_target(getattr(self, "_rank"), self._target_ranks):
-            return
-        self._hook_optimizer()
-
-    def _hook_optimizer(self):
-        def optimizer_pre_step_hook(optimizer, args, kargs):
-            self._step += 1
-            logger.info(f"grad_probe: optimizer step {self._step}")
-            if not data_in_list_target(self._step, self._target_step):
-                return
-            output_lines = []
-            for param, param_name in self._param2name.items():
-                if not data_in_list_target(param_name, self._param_list):
-                    continue
-                grad = param.main_grad if hasattr(param, "main_grad") else param.grad
-                if grad is None:
-                    logger.info(f"grad is None: {param_name}")
-                    continue
-                grad_info = GradStatCsv.generate_csv_line(param_name, self._level_adp, grad, self._bounds)
-                output_lines.append(grad_info)
-                if self._level_adp["have_grad_direction"]:
-                    GradientMonitor.save_grad_direction(param_name, grad,
-                                                        f'{self._output_path}/rank{self._rank}/step{self._step}')
-            output_dirpath = os.path.join(self._output_path, f"rank{getattr(self, '_rank')}")
-            if not os.path.isdir(output_dirpath):
-                create_directory(output_dirpath)
-            output_path = os.path.join(output_dirpath, f"grad_summary_{self._step}.csv")
-            if os.path.exists(output_path):
-                logger.warning(f"{output_path} will be deleted")
-                remove_path(output_path)
-            header_result = GradStatCsv.generate_csv_header(self._level_adp, self._bounds)
-            output_lines.insert(0, header_result)
-            write_csv(output_lines, output_path)
-            logger.info(f"write grad data to {output_path}")
-
-        if int(torch.__version__.split('.')[0]) >= 2:
-            register_optimizer_step_pre_hook(optimizer_pre_step_hook)