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

msprobe/pytorch/grad_probe/grad_monitor.py

@@ -0,0 +1,90 @@
+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)

msprobe/pytorch/grad_probe/grad_stat_csv.py

@@ -0,0 +1,129 @@
+from abc import ABC, abstractmethod
+from collections import namedtuple
+import hashlib
+import torch
+from msprobe.core.grad_probe.constant import GradConst
+
+CSV_header_input = namedtuple("CSV_header_input", ["bounds"])
+CSV_content_input = namedtuple("CSV_content_input", ["grad", "bounds"])
+
+
+class GradStatCsv:
+    csv = {}
+
+    @staticmethod
+    def generate_csv_header(level, bounds):
+        header = ["param_name"]
+        for key in level["header"]:
+            csv_header_input = CSV_header_input(bounds=bounds)
+            header.extend(GradStatCsv.csv[key].generate_csv_header(csv_header_input))
+        return header
+
+    @staticmethod
+    def generate_csv_line(param_name, level, grad, bounds):
+        line = [param_name]
+        for key in level["header"]:
+            csv_content_input = CSV_content_input(grad=grad, bounds=bounds)
+            line.extend(GradStatCsv.csv[key].generate_csv_content(csv_content_input))
+        return line
+
+
+def register_csv_item(key, cls=None):
+    if cls is None:
+        # 无参数时，返回装饰器函数
+        return lambda cls: register_csv_item(key, cls)
+    GradStatCsv.csv[key] = cls
+    return cls
+
+
+class CsvItem(ABC):
+    @abstractmethod
+    def generate_csv_header(csv_header_input):
+        pass
+
+    @abstractmethod
+    def generate_csv_content(csv_content_input):
+        pass
+
+
+@register_csv_item(GradConst.MD5)
+class CSV_md5(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["MD5"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        tensor_bytes = grad.cpu().detach().float().numpy().tobytes()
+        md5_hash = hashlib.md5(tensor_bytes)
+        return [md5_hash.hexdigest()]
+
+
+@register_csv_item(GradConst.DISTRIBUTION)
+class CSV_distribution(CsvItem):
+    def generate_csv_header(csv_header_input):
+        bounds = csv_header_input.bounds
+        intervals = []
+        if bounds:
+            intervals.append(f"(-inf, {bounds[0]}]")
+            for i in range(1, len(bounds)):
+                intervals.append(f"({bounds[i-1]}, {bounds[i]}]")
+        if intervals:
+            intervals.append(f"({bounds[-1]}, inf)")
+        intervals.append("=0")
+    
+        return intervals
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        bounds = csv_content_input.bounds
+        grad = grad.cpu().detach()
+        if grad.dtype == torch.bfloat16:
+            grad = grad.to(torch.float32)
+        element_num = grad.numel()
+        grad_equal_0_num = (grad == 0).sum().item()
+        bound = torch.Tensor(bounds)
+        bucketsize_result = torch.bucketize(grad, bound)
+        interval_nums = [(bucketsize_result == i).sum().item() for i in range(len(bound) + 1)]
+        interval_nums.append(grad_equal_0_num)
+        return_list = [x / element_num if element_num != 0 else 0 for x in interval_nums]
+        return return_list
+
+
+@register_csv_item(GradConst.MAX)
+class CSV_max(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["max"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [torch.max(grad).cpu().detach().float().numpy().tolist()]
+
+
+@register_csv_item(GradConst.MIN)
+class CSV_max(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["min"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [torch.min(grad).cpu().detach().float().numpy().tolist()]
+
+
+@register_csv_item(GradConst.NORM)
+class CSV_max(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["norm"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [torch.norm(grad).cpu().detach().float().numpy().tolist()]
+    
+
+@register_csv_item(GradConst.SHAPE)
+class CSV_shape(CsvItem):
+    def generate_csv_header(csv_header_input):
+        return ["shape"]
+
+    def generate_csv_content(csv_content_input):
+        grad = csv_content_input.grad
+        return [list(grad.shape)]

msprobe/pytorch/hook_module/hook_module.py

@@ -17,10 +17,13 @@
 
 import functools
 import threading
+
 import torch
 import torch.nn as nn
 import torch.utils.hooks as full_hooks
+
 from msprobe.core.common.const import Const
+torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 
 
 class HOOKModule(nn.Module):
@@ -46,9 +49,13 @@ class HOOKModule(nn.Module):
             else:
                 HOOKModule.module_count[self.prefix] += 1
                 self.prefix = self.prefix + str(HOOKModule.module_count[self.prefix] - 1) + Const.SEP
-            forward_pre_hook, forward_hook, backward_hook = build_hook(self.prefix)
-            self.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
-            self.register_forward_hook(forward_hook, with_kwargs=True)
+            forward_pre_hook, forward_hook, backward_hook, _ = build_hook(self.prefix)
+            if torch_version_above_or_equal_2:
+                self.register_forward_pre_hook(forward_pre_hook, with_kwargs=True)
+                self.register_forward_hook(forward_hook, with_kwargs=True)
+            else:
+                self.register_forward_pre_hook(forward_pre_hook)
+                self.register_forward_hook(forward_hook)
             self.register_backward_hook(backward_hook)
 
     def __call__(self, *input, **kwargs):
@@ -61,6 +68,10 @@ class HOOKModule(nn.Module):
             HOOKModule.inner_stop_hook[self.current_thread] = False
         return result
 
+    @classmethod
+    def reset_module_stats(cls):
+        cls.module_count = {}
+
     def _call_func(self, *input, **kwargs):
         full_backward_hooks, non_full_backward_hooks = [], []
         if len(self._backward_hooks) > 0:

msprobe/pytorch/hook_module/support_wrap_ops.yaml

@@ -1873,4 +1873,5 @@ distributed:
   - reduce_scatter
   - _reduce_scatter_base
   - _all_gather_base
-  - all_to_all_single
+  - all_to_all_single
+  - all_to_all

msprobe/pytorch/hook_module/utils.py

@@ -16,14 +16,14 @@
 """
 
 import os
-import yaml
+from msprobe.core.common.utils import load_yaml
 
-from msprobe.core.common.file_check import FileOpen
 
-cur_path = os.path.dirname(os.path.realpath(__file__))
-yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    Ops = yaml.safe_load(f)
-    WrapFunctionalOps = Ops.get('functional')
-    WrapTensorOps = Ops.get('tensor')
-    WrapTorchOps = Ops.get('torch')
+def get_ops():
+    cur_path = os.path.dirname(os.path.realpath(__file__))
+    yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
+    ops = load_yaml(yaml_path)
+    wrap_functional = ops.get('functional')
+    wrap_tensor = ops.get('tensor')
+    wrap_torch = ops.get('torch')
+    return set(wrap_functional) | set(wrap_tensor) | set(wrap_torch)

msprobe/pytorch/hook_module/wrap_aten.py

@@ -18,18 +18,17 @@
 import os
 import torch
 
-import yaml
-
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.common.utils import torch_device_guard
 from msprobe.core.common.const import Const
-from msprobe.core.common.file_check import FileOpen
-
+from msprobe.core.common.utils import load_yaml
+from msprobe.pytorch.function_factory import npu_custom_grad_functions
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    WrapAtenOps = yaml.safe_load(f).get('aten')
+ops = load_yaml(yaml_path)
+wrap_aten_ops = ops.get('aten')
+white_aten_ops = ops.get('white_aten_ops', [])
 
 
 aten_func = {}
@@ -38,9 +37,9 @@ for f in dir(torch.ops.aten):
 
 
 def get_aten_ops():
-    global WrapAtenOps
+    global wrap_aten_ops
     _all_aten_ops = dir(torch.ops.aten)
-    return set(WrapAtenOps) & set(_all_aten_ops)
+    return set(wrap_aten_ops) & set(_all_aten_ops)
 
 
 class HOOKAtenOP(object):
@@ -48,7 +47,7 @@ class HOOKAtenOP(object):
 
 
 class AtenOPTemplate(HOOKModule):
-    def __init__(self, op, hook):
+    def __init__(self, op, hook, need_hook=True):
         if isinstance(op, torch._ops.OpOverloadPacket):
             op_name_ = op._qualified_op_name.split("::")[-1]
         else:
@@ -58,10 +57,21 @@ class AtenOPTemplate(HOOKModule):
                 op_name_ = op_name_ + '.' + overload_name
         self.op = op
         self.prefix_op_name_ = "Aten" + Const.SEP + str(op_name_) + Const.SEP
-        super().__init__(hook)
+        self.need_hook = need_hook
+        if self.need_hook:
+            super().__init__(hook)
 
     @torch_device_guard
     def forward(self, *args, **kwargs):
+        if isinstance(self.op, str):
+            if self.op in npu_custom_grad_functions:
+                return npu_custom_grad_functions[self.op](*args, **kwargs)
+            if self.op in white_aten_ops:
+                return eval(f"torch.ops.aten.{self.op}")(*args, **kwargs)
+            if self.op not in aten_func:
+                raise Exception(f"Skip op[{self.op}] accuracy check, because the op is not "
+                                f"in dir(torch.ops.aten) and support yaml.")
+            return aten_func[self.op](*args, **kwargs)
         return self.op(*args, **kwargs)
 
 

msprobe/pytorch/hook_module/wrap_distributed.py

@@ -18,18 +18,15 @@
 import os
 from functools import wraps
 import torch.distributed as dist
-import yaml
 
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.common.utils import torch_device_guard
 from msprobe.core.common.const import Const
-from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.utils import load_yaml
 
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    WrapDistributedOps = yaml.safe_load(f).get('distributed')
 
 
 distributed_func = {}
@@ -38,9 +35,10 @@ for f in dir(dist):
 
 
 def get_distributed_ops():
-    global WrapDistributedOps
     _all_distributed_ops = dir(dist)
-    return set(WrapDistributedOps) & set(_all_distributed_ops)
+    yaml_data = load_yaml(yaml_path)
+    wrap_distributed_ops = yaml_data.get('distributed')
+    return set(wrap_distributed_ops) & set(_all_distributed_ops)
 
 
 class HOOKDistributedOP(object):
@@ -57,7 +55,12 @@ class DistributedOPTemplate(HOOKModule):
 
     @torch_device_guard
     def forward(self, *args, **kwargs):
-        return distributed_func.get(self.op_name_)(*args, **kwargs)
+        if kwargs.get("async_op") or self.op_name_ in ["isend", "irecv"]:
+            handle = distributed_func.get(self.op_name_)(*args, **kwargs)
+            handle.wait()
+            return handle
+        else:
+            return distributed_func.get(self.op_name_)(*args, **kwargs)
 
 
 def wrap_distributed_op(op_name, hook):

msprobe/pytorch/hook_module/wrap_functional.py

@@ -16,15 +16,13 @@
 """
 
 import os
-
 import torch
-import yaml
 
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.common.utils import torch_device_guard
 from msprobe.core.common.const import Const
 from msprobe.pytorch.common.log import logger
-from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.utils import load_yaml
 
 
 def remove_dropout():
@@ -66,14 +64,13 @@ def remove_dropout():
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    WrapFunctionalOps = yaml.safe_load(f).get('functional')
 
 
 def get_functional_ops():
-    global WrapFunctionalOps
+    yaml_data = load_yaml(yaml_path)
+    wrap_functional_ops = yaml_data.get('functional')
     _all_functional_ops = dir(torch.nn.functional)
-    return set(WrapFunctionalOps) & set(_all_functional_ops)
+    return set(wrap_functional_ops) & set(_all_functional_ops)
 
 
 TorchFunctions = {func: getattr(torch.nn.functional, func) for func in get_functional_ops()}

msprobe/pytorch/hook_module/wrap_npu_custom.py

@@ -17,27 +17,33 @@
 
 import os
 import torch
-import torch_npu
-import yaml
 
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.common.utils import torch_device_guard, torch_without_guard_version
 from msprobe.core.common.const import Const
-from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.utils import load_yaml
+from msprobe.pytorch.function_factory import npu_custom_functions
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    WrapNpuOps = yaml.safe_load(f).get('torch_npu')
+
+
+try:
+    import torch_npu
+except ImportError:
+    is_gpu = True
+else:
+    is_gpu = False
 
 
 def get_npu_ops():
-    global WrapNpuOps
     if torch_without_guard_version:
         _npu_ops = dir(torch.ops.npu)
     else:
         _npu_ops = dir(torch_npu._C._VariableFunctionsClass)
-    return set(WrapNpuOps) & set(_npu_ops)
+    yaml_data = load_yaml(yaml_path)
+    wrap_npu_ops = yaml_data.get('torch_npu')
+    return set(wrap_npu_ops) & set(_npu_ops)
 
 
 class HOOKNpuOP(object):
@@ -46,13 +52,19 @@ class HOOKNpuOP(object):
 
 class NpuOPTemplate(HOOKModule):
 
-    def __init__(self, op_name, hook):
+    def __init__(self, op_name, hook, need_hook=True):
         self.op_name_ = op_name
         self.prefix_op_name_ = "NPU" + Const.SEP + str(op_name) + Const.SEP
-        super().__init__(hook)
+        self.need_hook = need_hook
+        if need_hook:
+            super().__init__(hook)
 
     @torch_device_guard
     def forward(self, *args, **kwargs):
+        if not self.need_hook:
+            if self.op_name_ not in npu_custom_functions:
+                raise Exception(f'There is not bench function {self.op_name_}')
+            return npu_custom_functions[self.op_name_](*args, **kwargs)
         if torch_without_guard_version:
             return getattr(torch.ops.npu, str(self.op_name_))(*args, **kwargs)
         else:
@@ -60,7 +72,6 @@ class NpuOPTemplate(HOOKModule):
 
 
 def wrap_npu_op(op_name, hook):
-
     def npu_op_template(*args, **kwargs):
         return NpuOPTemplate(op_name, hook)(*args, **kwargs)
 

msprobe/pytorch/hook_module/wrap_tensor.py

@@ -18,23 +18,22 @@
 import os
 
 import torch
-import yaml
 
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.common.utils import torch_device_guard, parameter_adapter
 from msprobe.core.common.const import Const
-from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.utils import load_yaml
+
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    WrapTensorOps = yaml.safe_load(f).get('tensor')
 
 
 def get_tensor_ops():
-    global WrapTensorOps
     _tensor_ops = dir(torch.Tensor)
-    return set(WrapTensorOps) & set(_tensor_ops)
+    yaml_data = load_yaml(yaml_path)
+    wrap_tensor_ops = yaml_data.get('tensor')
+    return set(wrap_tensor_ops) & set(_tensor_ops)
 
 
 TensorOps = {op: getattr(torch.Tensor, op) for op in get_tensor_ops()}

msprobe/pytorch/hook_module/wrap_torch.py

@@ -16,25 +16,23 @@
 """
 
 import os
-
 import torch
-import yaml
 
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.common.utils import torch_device_guard
 from msprobe.core.common.const import Const
-from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.utils import load_yaml
+
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    WrapTorchOps = yaml.safe_load(f).get('torch')
 
 
 def get_torch_ops():
-    global WrapTorchOps
     _torch_ops = []
-    for operation in WrapTorchOps:
+    yaml_data = load_yaml(yaml_path)
+    wrap_torch_ops = yaml_data.get('torch')
+    for operation in wrap_torch_ops:
         if '.' in operation:
             operation_sub_module_name, operation_sub_op = operation.rsplit('.', 1)
             operation_sub_module = getattr(torch, operation_sub_module_name)

msprobe/pytorch/hook_module/wrap_vf.py

@@ -16,24 +16,22 @@
 """
 
 import os
-
 import torch
-import yaml
 
+from msprobe.core.common.const import Const
+from msprobe.core.common.utils import load_yaml
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
-from msprobe.core.common.file_check import FileOpen
 from msprobe.pytorch.common.utils import torch_device_guard
-from msprobe.core.common.const import Const
+
 
 cur_path = os.path.dirname(os.path.realpath(__file__))
 yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-with FileOpen(yaml_path, 'r') as f:
-    WrapVfOps = yaml.safe_load(f).get('_VF')
 
 
 def get_vf_ops():
-    global WrapVfOps
-    return WrapVfOps
+    yaml_data = load_yaml(yaml_path)
+    wrap_vf_ops = yaml_data.get('_VF')
+    return wrap_vf_ops
 
 
 class HOOKVfOP(object):