mindstudio-probe 1.0.3__py3-none-any.whl → 1.1.0__py3-none-any.whl

msprobe/pytorch/function_factory.py

@@ -1,75 +1,91 @@
-from msprobe.pytorch.common.utils import logger
-from msprobe.pytorch.bench_functions.apply_adam_w import npu_apply_adam_w
-from msprobe.pytorch.bench_functions.confusion_transpose import npu_confusion_transpose, \
-    npu_confusion_transpose_backward
-from msprobe.pytorch.bench_functions.fast_gelu import fast_gelu, npu_fast_gelu_backward
-from msprobe.pytorch.bench_functions.layer_norm_eval import npu_layer_norm_eval
-from msprobe.pytorch.bench_functions.linear import npu_linear, npu_linear_backward
-from msprobe.pytorch.bench_functions.matmul_backward import matmul_backward
-from msprobe.pytorch.bench_functions.npu_fusion_attention import npu_fusion_attention, npu_fusion_attention_grad
-from msprobe.pytorch.bench_functions.rms_norm import npu_rms_norm, npu_rms_norm_backward
-from msprobe.pytorch.bench_functions.rotary_mul import npu_rotary_mul, npu_rotary_mul_backward
-from msprobe.pytorch.bench_functions.scaled_mask_softmax import npu_scaled_masked_softmax, \
-    npu_scaled_masked_softmax_backward
-from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward, swish_grad, swish
-
-
-class Register(dict):
-    def __init__(self, *args, **kwargs):
-        super(Register, self).__init__(*args, **kwargs)
-        self._dict = {}
-
-    def __call__(self, target_func_list):
-        for target in target_func_list:
-            self.register(target)
-        return
-
-    def __setitem__(self, key, value):
-        self._dict[key] = value
-
-    def __getitem__(self, key):
-        return self._dict[key]
-
-    def __contains__(self, key):
-        return key in self._dict
-
-    def __str__(self):
-        return str(self._dict)
-
-    def keys(self):
-        return self._dict.keys()
-
-    def values(self):
-        return self._dict.values()
-
-    def items(self):
-        return self._dict.items()
-
-    def register(self, target):
-
-        def add_register_item(key, value):
-            if key in self._dict:
-                logger.warning(f"{value.__name__} has been registered before, so we will overriden it.")
-            self[key] = value
-            return value
-
-        if callable(target):
-            return add_register_item(target.__name__, target)
-        else:
-            raise Exception(f"The func {target} is not callable.")
-
-
-# register for npu custom bench functions
-npu_custom_functions = Register()
-npu_custom_functions([
-    npu_apply_adam_w, npu_confusion_transpose, fast_gelu, npu_layer_norm_eval, npu_linear, npu_fusion_attention,
-    npu_rms_norm, npu_rotary_mul, npu_scaled_masked_softmax, npu_swiglu
-])
-
-# register for npu custom backward bench functions
-npu_custom_grad_functions = Register()
-npu_custom_grad_functions([
-    npu_confusion_transpose_backward, npu_fast_gelu_backward, npu_linear_backward, matmul_backward,
-    npu_fusion_attention_grad, npu_rms_norm_backward, npu_rotary_mul_backward, npu_scaled_masked_softmax_backward,
-    npu_swiglu_backward
-])
+# 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.bench_functions.apply_adam_w import npu_apply_adam_w
+from msprobe.pytorch.bench_functions.confusion_transpose import npu_confusion_transpose, \
+    npu_confusion_transpose_backward
+from msprobe.pytorch.bench_functions.fast_gelu import npu_fast_gelu, npu_fast_gelu_backward
+from msprobe.pytorch.bench_functions.layer_norm_eval import npu_layer_norm_eval
+from msprobe.pytorch.bench_functions.linear import npu_linear, npu_linear_backward
+from msprobe.pytorch.bench_functions.matmul_backward import matmul_backward
+from msprobe.pytorch.bench_functions.npu_fusion_attention import npu_fusion_attention, npu_fusion_attention_grad, \
+    gpu_fusion_attention
+from msprobe.pytorch.bench_functions.rms_norm import npu_rms_norm, npu_rms_norm_backward
+from msprobe.pytorch.bench_functions.rotary_mul import npu_rotary_mul, npu_rotary_mul_backward
+from msprobe.pytorch.bench_functions.scaled_mask_softmax import npu_scaled_masked_softmax, \
+    npu_scaled_masked_softmax_backward
+from msprobe.pytorch.bench_functions.swiglu import npu_swiglu, npu_swiglu_backward
+from msprobe.pytorch.common.utils import logger
+
+
+class Register(dict):
+    def __init__(self, *args, **kwargs):
+        super(Register, self).__init__(*args, **kwargs)
+        self._dict = {}
+
+    def __call__(self, target_func_list):
+        for target in target_func_list:
+            self.register(target)
+        return
+
+    def __setitem__(self, key, value):
+        self._dict[key] = value
+
+    def __getitem__(self, key):
+        return self._dict[key]
+
+    def __contains__(self, key):
+        return key in self._dict
+
+    def __str__(self):
+        return str(self._dict)
+
+    def keys(self):
+        return self._dict.keys()
+
+    def values(self):
+        return self._dict.values()
+
+    def items(self):
+        return self._dict.items()
+
+    def register(self, target):
+
+        def add_register_item(key, value):
+            if key in self._dict:
+                logger.warning(f"{value.__name__} has been registered before, so we will overriden it.")
+            self[key] = value
+            return value
+
+        if callable(target):
+            return add_register_item(target.__name__, target)
+        else:
+            raise Exception(f"The func {target} is not callable.")
+
+
+# register for npu custom bench functions
+npu_custom_functions = Register()
+npu_custom_functions([
+    npu_apply_adam_w, npu_confusion_transpose, npu_fast_gelu, npu_layer_norm_eval, npu_linear, npu_fusion_attention,
+    npu_rms_norm, npu_rotary_mul, npu_scaled_masked_softmax, npu_swiglu, gpu_fusion_attention
+])
+
+# register for npu custom backward bench functions
+npu_custom_grad_functions = Register()
+npu_custom_grad_functions([
+    npu_confusion_transpose_backward, npu_fast_gelu_backward, npu_linear_backward, matmul_backward,
+    npu_fusion_attention_grad, npu_rms_norm_backward, npu_rotary_mul_backward, npu_scaled_masked_softmax_backward,
+    npu_swiglu_backward
+])

msprobe/pytorch/functional/module_dump.py

@@ -0,0 +1,84 @@
+# 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 torch
+import torch.nn as nn
+from msprobe.core.common.const import Const
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.data_dump.scope import BaseScope
+from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.debugger.precision_debugger import PrecisionDebugger
+from msprobe.pytorch.hook_module.api_registry import api_register
+from msprobe.pytorch.service import torch_version_above_or_equal_2
+
+hook_handle_list = []
+
+
+def module_dump(module, dump_name):
+    if not isinstance(module, nn.Module):
+        logger.error("The parameter module in module_dump must be a Module subclass.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if not isinstance(dump_name, str):
+        logger.error("The parameter dump_name in module_dump must be a str type.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+
+    api_register.api_originality()
+    register_hook(module, dump_name)
+
+
+def module_dump_end():
+    api_register.api_modularity()
+    remove_hook()
+    hook_handle_list.clear()
+
+
+def register_hook(module, dump_name):
+    prefix = BaseScope.Module_Type_Module + Const.SEP + dump_name + Const.SEP + module.__class__.__name__ + Const.SEP
+
+    pdg = PrecisionDebugger()
+    _, forward_hook, backward_hook, forward_hook_torch_version_below_2 = \
+        pdg.service.build_hook(BaseScope.Module_Type_Module, prefix)
+
+    if torch_version_above_or_equal_2:
+        forward_hook_handle = module.register_forward_hook(forward_hook, with_kwargs=True)
+        hook_handle_list.append(forward_hook_handle)
+    else:
+        pdg.service.check_register_full_backward_hook(module)
+        full_backward_hook_handle = module.register_full_backward_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
+        forward_hook_handle = module.register_forward_hook(forward_hook_torch_version_below_2)
+        hook_handle_list.extend([full_backward_hook_handle, forward_hook_handle])
+    pdg.service.check_register_full_backward_hook(module)
+    full_backward_hook_handle = module.register_full_backward_hook(backward_hook)
+
+    forward_pre_hook_handle = module.register_forward_pre_hook(
+        pdg.service.module_processor.node_hook(prefix + Const.FORWARD, Const.START))
+    forward_hook_handle = module.register_forward_hook(
+        pdg.service.module_processor.node_hook(prefix + Const.FORWARD, Const.STOP))
+    hook_handle_list.extend([full_backward_hook_handle, forward_pre_hook_handle, forward_hook_handle])
+
+    if torch_version_above_or_equal_2:
+        backward_pre_hook_handle = module.register_full_backward_pre_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.START))
+        pdg.service.check_register_full_backward_hook(module)
+        full_backward_hook_handle = module.register_full_backward_hook(
+            pdg.service.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
+        hook_handle_list.extend([backward_pre_hook_handle, full_backward_hook_handle])
+
+
+def remove_hook():
+    for hook_handle in hook_handle_list:
+        if isinstance(hook_handle, torch.utils.hooks.RemovableHandle):
+            hook_handle.remove()

msprobe/pytorch/grad_probe/grad_monitor.py

@@ -1,90 +1,91 @@
-import os
-from collections import defaultdict
-
-import torch
-if int(torch.__version__.split('.')[0]) >= 2:
-    from torch.optim.optimizer import register_optimizer_step_pre_hook
-from msprobe.pytorch.grad_probe.grad_stat_csv import GradStatCsv
-from msprobe.core.grad_probe.utils import check_numeral_list_ascend, data_in_list_target
-from msprobe.core.grad_probe.constant import GradConst, level_adp
-from msprobe.core.common.file_check import create_directory
-from msprobe.core.common.log import logger
-from msprobe.core.common.utils import remove_path, write_csv, save_npy
-from msprobe.pytorch.common.utils import  get_rank_id, print_rank_0, save_pt
-
-
-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 recoverd")
-        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.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
-            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 recoverd")
-                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)
-        if int(torch.__version__.split('.')[0]) >= 2:
-            register_optimizer_step_pre_hook(optimizer_pre_step_hook)
+import os
+from collections import defaultdict
+
+import torch
+if int(torch.__version__.split('.')[0]) >= 2:
+    from torch.optim.optimizer import register_optimizer_step_pre_hook
+from msprobe.pytorch.grad_probe.grad_stat_csv import GradStatCsv
+from msprobe.core.grad_probe.utils import check_numeral_list_ascend, data_in_list_target
+from msprobe.core.grad_probe.constant import level_adp
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.file_utils import remove_path, save_npy, write_csv, create_directory
+from msprobe.pytorch.common.utils import get_rank_id, print_rank_0
+
+
+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 recoverd")
+        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 recoverd")
+                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)