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

msprobe/pytorch/monitor/csv2tb.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2024-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 datetime
+import os
+import re
+from multiprocessing import Process
+
+import pytz
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from msprobe.core.common.const import MonitorConst
+from msprobe.core.common.file_utils import read_csv, create_directory, remove_path
+from msprobe.core.common.utils import is_int
+from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.monitor.utils import get_target_output_dir
+
+all_data_type_list = ["actv", "actv_grad", "exp_avg", "exp_avg_sq", "grad_unreduced", "grad_reduced", "param"]
+CSV_FILE_SUFFIX = r"_\d+-\d+\.csv"
+
+
+def parse_step_line(data, line_id, name, ops):
+    vp_id = data["vpp_stage"][line_id]
+    module_name = data[name][line_id]
+    step = data["step"][line_id]
+    vpp_name = f"vp{vp_id}:{module_name}"
+    ops_result = {}
+    for op in ops:
+        ops_result[op] = data[op][line_id]
+    return vpp_name, step, ops_result
+
+
+def parse_step_fn(filepath):
+    data = read_csv(filepath)
+
+    header = list(data.keys())
+    name = header[MonitorConst.HEADER_NAME_INDEX]
+    ops = header[MonitorConst.OPS_START_INDEX:]
+
+    parse_step_result = {}
+
+    for line_id in range(len(data)):
+        vpp_name, step, ops_result = parse_step_line(data, line_id, name, ops)
+        if vpp_name not in parse_step_result:
+            parse_step_result[vpp_name] = {}
+        if step in parse_step_result[vpp_name]:
+            raise Exception(f"duplicated step({step})")
+        parse_step_result[vpp_name][step] = ops_result
+    return parse_step_result
+
+
+def write_step(output_dirpath, parse_step_result, rank, data_type):
+    tb_output_path = os.path.join(output_dirpath, f"rank{rank}", data_type)
+    if os.path.exists(tb_output_path):
+        remove_path(tb_output_path)
+        logger.warning(f"existing path {tb_output_path} will be recovered")
+    writer = SummaryWriter(tb_output_path)
+    for vpp_name, step_data_dict in parse_step_result.items():
+        step_data_list = [(step, ops) for step, ops in step_data_dict.items()]
+        step_data_list.sort(key=lambda x: x[0])
+        for step_data in step_data_list:
+            step = step_data[0]
+            ops = step_data[1]
+            for op, value in ops.items():
+                tag = f"{vpp_name}/{op}"
+                writer.add_scalar(tag, value, step)
+
+
+def update_dict(dict1, dict2):
+    for key, value in dict2.items():
+        if key in dict1:
+            if isinstance(dict1[key], dict) and isinstance(value, dict):
+                try:
+                    update_dict(dict1[key], value)
+                except Exception as e:
+                    raise Exception(f"Error updating nested dict failed at key '{key}': {e}") from e
+            else:
+                raise Exception(f"duplicate key: {key}")
+        else:
+            dict1[key] = value
+    return dict1
+
+
+def csv2tb_by_step_work(target_output_dirs, output_dirpath, data_type_list):
+    for directory in tqdm(target_output_dirs):
+        dirpath = directory["path"]
+        rank = directory["rank"]
+        for data_type in data_type_list:
+            all_step_result = {}
+            for filename in os.listdir(dirpath):
+                if not re.match(f"{data_type}{CSV_FILE_SUFFIX}", filename):
+                    continue
+                filepath = os.path.join(dirpath, filename)
+                try:
+                    parse_step_result = parse_step_fn(filepath)
+                except Exception as e:
+                    logger.error(f"csv2tensorboard parse {filepath} failed \n {e}")
+                    break
+
+                all_step_result = update_dict(all_step_result, parse_step_result)
+            if all_step_result:
+                write_step(output_dirpath, all_step_result, rank, data_type)
+
+
+def check_process_num(process_num):
+    if not is_int(process_num) or process_num <= 0:
+        raise ValueError(f"process_num({process_num}) is not a positive integer")
+
+
+def check_data_type_list(data_type_list):
+    if data_type_list is None:
+        logger.info(f"data_type_list is None, use defualt all_data_type_list: {all_data_type_list}")
+        return
+    if not isinstance(data_type_list, list):
+        raise ValueError(f"data_type_list({data_type_list}) is not a list")
+    for data_type in data_type_list:
+        if data_type not in all_data_type_list:
+            raise ValueError(f"data type({data_type}) is not supported, supported data type: {all_data_type_list}")
+
+
+def csv2tensorboard_by_step(
+        monitor_path,
+        time_start=None,
+        time_end=None,
+        process_num=1,
+        data_type_list=None,
+        output_dirpath=None
+):
+    check_process_num(process_num)
+    check_data_type_list(data_type_list)
+    target_output_dirs = get_target_output_dir(monitor_path, time_start, time_end)
+    target_output_dirs = [{"rank": rank, "path": path} for rank, path in target_output_dirs.items()]
+    if output_dirpath is None:
+        local_tz = pytz.timezone("Asia/Shanghai")  # 根据需要调整为目标时区
+        cur_time = datetime.datetime.now(local_tz).strftime("%b%d_%H-%M-%S")
+        output_dirpath = os.path.join(monitor_path, f"{cur_time}-csv2tensorboard_by_step")
+    create_directory(output_dirpath)
+
+    task_num = len(target_output_dirs)
+    task_num_per_pro = task_num // process_num
+    target_data_type = data_type_list if data_type_list else all_data_type_list
+
+    processes = []
+    for pro_id in range(process_num):
+        task_start_id = pro_id * task_num_per_pro
+        task_end_id = (pro_id + 1) * task_num_per_pro if pro_id != process_num - 1 else task_num
+        task_dirs = target_output_dirs[task_start_id: task_end_id]
+
+        p = Process(target=csv2tb_by_step_work, args=(task_dirs, output_dirpath, target_data_type))
+        processes.append(p)
+        p.start()
+    for p in processes:
+        p.join()
+    logger.info(f"output has been saved to: {output_dirpath}")

msprobe/pytorch/monitor/distributed/distributed_ops.yaml

@@ -0,0 +1,19 @@
+distributed:
+  - send
+  - recv
+  - isend
+  - irecv
+  - batch_isend_irecv
+  - broadcast
+  - all_reduce
+  - reduce
+  - scatter
+  - reduce_scatter
+  - _reduce_scatter_base
+  - reduce_scatter_tensor
+  - all_gather
+  - gather
+  - _all_gather_base
+  - all_gather_into_tensor
+
+  

msprobe/pytorch/monitor/distributed/stack_blacklist.yaml

@@ -0,0 +1,5 @@
+stack:
+- msprobe/pytorch/monitor/distributed
+- site-packages/torch/nn/modules/module.py
+- multiprocessing
+- debugpy

msprobe/pytorch/monitor/distributed/wrap_distributed.py

@@ -0,0 +1,283 @@
+# Copyright (c) 2024-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 inspect
+import os
+import re
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+from msprobe.core.common.const import MonitorConst
+from msprobe.core.common.file_utils import load_yaml
+from msprobe.pytorch.monitor.module_metric import get_metrics, get_summary_writer_tag_name
+
+try:
+    import torch_npu
+except ImportError:
+    pass
+
+RANK = None
+
+OpsPath = os.path.join(os.path.dirname(__file__), "distributed_ops.yaml")
+WrapDistributedOps = load_yaml(OpsPath).get("distributed", [])
+
+StackBlackListPath = os.path.join(os.path.dirname(__file__), "stack_blacklist.yaml")
+StackBlackList = load_yaml(StackBlackListPath).get("stack", [])
+
+distributed_func = {}
+for f in dir(dist):
+    distributed_func[f] = getattr(dist, f)
+
+ORIGIN_WAIT = getattr(dist.Work, 'wait')
+PENDING_ASYNC_CC_BY_HANDLE = {}
+
+
+def get_distributed_ops():
+    global WrapDistributedOps
+    _all_distributed_ops = dir(dist)
+    return set(WrapDistributedOps) & set(_all_distributed_ops)
+
+
+class DistributedOPTemplate(nn.Module):
+    def __init__(self, op_name, pre_hooks, post_hooks):
+        super(DistributedOPTemplate, self).__init__()
+        self.op_name_ = str(op_name)
+        self.__name__ = self.op_name_
+        self.cc_hooks = []
+        for pre_hook in pre_hooks:
+            handle = self.register_forward_pre_hook(pre_hook, with_kwargs=True)
+            self.cc_hooks.append(handle)
+        for hook in post_hooks:
+            handle = self.register_forward_hook(hook, with_kwargs=True)
+            self.cc_hooks.append(handle)
+
+    def forward(self, *args, **kwargs):
+        return distributed_func.get(self.op_name_)(*args, **kwargs)
+
+
+class ApiRegistry:
+    def __init__(self):
+        self.distributed_attr_origin = {}
+        self.distributed_attr_hooked = {}
+
+    @staticmethod
+    def store_ori_attr(ori_api_group, api_list, api_ori_attr):
+        for api in api_list:
+            if '.' in api:
+                sub_module_name, sub_op = api.rsplit('.', 1)
+                sub_module = getattr(ori_api_group, sub_module_name)
+                api_ori_attr[api] = getattr(sub_module, sub_op)
+            else:
+                api_ori_attr[api] = getattr(ori_api_group, api)
+
+    @staticmethod
+    def set_api_attr(api_group, attr_dict):
+        for cc_api_name, cc_api_entry_func in attr_dict.items():
+            if '.' in cc_api_name:
+                sub_module_name, sub_op = cc_api_name.rsplit('.', 1)
+                sub_module = getattr(api_group, sub_module_name, None)
+                if sub_module is not None:
+                    setattr(sub_module, sub_op, cc_api_entry_func)
+            else:
+                setattr(api_group, cc_api_name, cc_api_entry_func)
+
+    @staticmethod
+    def redirect_wait():
+        global ORIGIN_WAIT
+        global PENDING_ASYNC_CC_BY_HANDLE
+
+        def wrapped_wait(work):
+            def wrapped_wait(*args, **kwargs):
+                ORIGIN_WAIT(*args, **kwargs)
+                if args[0] in PENDING_ASYNC_CC_BY_HANDLE:
+                    store_func = PENDING_ASYNC_CC_BY_HANDLE.pop(args[0])
+                    store_func()
+
+            return wrapped_wait
+
+        dist.Work.wait = wrapped_wait(dist.Work)
+
+    def redirect_api(self):
+        self.set_api_attr(dist, self.distributed_attr_hooked)
+        self.set_api_attr(dist.distributed_c10d, self.distributed_attr_hooked)
+        self.redirect_wait()
+
+    def restore_api(self):
+        self.set_api_attr(dist, self.distributed_attr_origin)
+        self.set_api_attr(dist.distributed_c10d, self.distributed_attr_origin)
+        setattr(dist.Work, 'wait', ORIGIN_WAIT)
+
+    def initialize_hook(self, pre_hooks, post_hooks):
+        self.store_ori_attr(dist, get_distributed_ops(), self.distributed_attr_origin)
+        cc_hooks = []
+        for op_name in get_distributed_ops():
+            self.distributed_attr_hooked[op_name] = DistributedOPTemplate(op_name, pre_hooks, post_hooks)
+            cc_hooks.extend(self.distributed_attr_hooked[op_name].cc_hooks)
+        return cc_hooks
+
+
+def get_process_group(process_group):
+    return (
+        process_group
+        if isinstance(process_group, dist.ProcessGroup)
+        else dist.GroupMember.WORLD
+    )
+
+
+def stack_filter(stack):
+    for pattern in StackBlackList:
+        if re.search(pattern, stack):
+            return False
+    return True
+
+
+def get_callstack():
+    callstack = []
+    for (_, path, line, func, _, _) in inspect.stack():
+        stack_line = f'{path}[{line}]'
+        if stack_filter(stack_line):
+            callstack.append(stack_line + '   ' + func)
+    return callstack
+
+
+@torch.no_grad()
+def op_aggregate(op, tensorlist):
+    if isinstance(tensorlist, torch.Tensor):
+        return tensorlist
+    if not tensorlist:
+        return torch.tensor(torch.nan)
+    if op == 'min':
+        return min(tensorlist)
+    if op == 'max':
+        return max(tensorlist)
+    if op == 'norm':
+        return sum(tensorlist)
+    if op == 'zeros':
+        return sum(tensorlist) / len(tensorlist)
+    if op == 'nans':
+        return sum(tensorlist)
+    if op == 'mean':
+        return sum(tensorlist) / len(tensorlist)
+    return torch.tensor(torch.nan)
+
+
+def update_data(old, new):
+    for tag, op2tensor in new.items():
+        if tag not in old:
+            old[tag] = {}
+        for op, tensor in op2tensor.items():
+            if op not in old[tag]:
+                old[tag][op] = [tensor]
+            else:
+                old[tag][op].append(tensor)
+    return old
+
+
+def is_target_line(codeline):
+    stack = get_callstack()
+    whole_stack = ';'.join(stack)
+    if codeline == []:
+        return True
+    for pattern in codeline:
+        if re.search(pattern, whole_stack):
+            return True
+    return False
+
+
+@torch.no_grad()
+def catch_data(cc_context, cc_name, ops, args, prefix):
+    tensor_args = {}
+    for arg in args:
+        if isinstance(arg, torch.Tensor):
+            key = get_summary_writer_tag_name(cc_name, f'{prefix}_{len(tensor_args)}', RANK)
+            tensor_args[key] = arg
+        elif isinstance(arg, list):
+            if isinstance(arg[0], torch.Tensor):
+                stacked_arg = torch.stack(arg)
+            elif isinstance(arg[0], dist.P2POp):
+                stacked_arg = torch.stack([op.tensor for op in arg])
+            key = get_summary_writer_tag_name(cc_name, f'{prefix}_{len(tensor_args)}', RANK)
+            tensor_args[key] = stacked_arg
+
+    new_data = get_metrics(ops, tensor_args, 1e-8)
+    cc_context.data = update_data(cc_context.data, new_data)
+
+
+def create_async_callback_func(context, cc_name, ops, args, prefix):
+    def store_data():
+        catch_data(context, cc_name, ops, args, prefix)
+
+    return store_data
+
+
+def create_hooks(context, monitor):
+    def cc_log_hook(module, args, kwargs):
+        stack = ';'.join(get_callstack())
+        monitor.cc_logged_stack[module.op_name_].add(stack)
+        return
+
+    def cc_pre_hook(module, args, kwargs):
+        if not is_target_line(monitor.cc_codeline):
+            return
+        args = args + tuple(kwargs.values())
+        catch_data(context[module.op_name_], module.op_name_, monitor.ops, args, MonitorConst.PREFIX_PRE)
+        return
+
+    def cc_hook(module, args, kwargs, out=None):
+        if not is_target_line(monitor.cc_codeline):
+            return out
+        args = args + tuple(kwargs.values())
+        if out:  # async
+            if isinstance(out, dist.Work):
+                PENDING_ASYNC_CC_BY_HANDLE[out] = create_async_callback_func(
+                    context[module.op_name_],
+                    module.op_name_,
+                    monitor.ops, args,
+                    MonitorConst.PREFIX_POST
+                )
+            elif isinstance(out, list):  # batch_isend_irecv
+                for out_element in out:
+                    PENDING_ASYNC_CC_BY_HANDLE[out_element] = create_async_callback_func(
+                        context[module.op_name_],
+                        module.op_name_,
+                        monitor.ops, args,
+                        MonitorConst.PREFIX_POST
+                    )
+            return out
+        catch_data(context[module.op_name_], module.op_name_, monitor.ops, args, MonitorConst.PREFIX_POST)
+        return out
+
+    global RANK
+    pre_hooks = []
+    hooks = []
+    RANK = dist.get_rank()
+    if dist.is_initialized() and RANK not in monitor.module_rank_list and monitor.module_rank_list != []:
+        return [pre_hooks, hooks]
+
+    if monitor.cc_log_only:
+        pre_hooks.append(cc_log_hook)
+        return [pre_hooks, hooks]
+
+    if monitor.cc_pre_hook:
+        pre_hooks.append(cc_pre_hook)
+
+    hooks.append(cc_hook)
+
+    return [pre_hooks, hooks]
+
+
+api_register = ApiRegistry()

msprobe/pytorch/monitor/features.py

@@ -0,0 +1,108 @@
+# 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
+from torch.autograd.functional import jacobian
+from msprobe.pytorch.common.log import logger
+
+
+@torch.no_grad()
+def square_sum(x: torch.tensor):
+    return (x * x).sum()
+
+
+@torch.no_grad()
+def get_min(x: torch.tensor):
+    return torch.min(x)
+
+
+@torch.no_grad()
+def get_mean(x: torch.tensor):
+    return torch.mean(x.to(torch.float64))
+
+
+@torch.no_grad()
+def get_norm(x: torch.tensor):
+    return torch.norm(x.to(torch.float64), p=2)
+
+
+@torch.no_grad()
+def get_max(x: torch.tensor):
+    return torch.max(x)
+
+
+@torch.no_grad()
+def get_zeros(x: torch.tensor, eps: float):
+    return torch.sum(torch.abs(x) < eps) / x.numel()
+
+
+@torch.no_grad()
+def get_sign_matches(x: torch.tensor, y: torch.tensor):
+    xs = x.sign()
+    ys = y.sign()
+    try:
+        same_direction_ratio = ((xs * ys).sum() / ys.numel() + 1) / 2
+    except RuntimeError as e:
+        logger.info(f"RuntimeError: {e}")
+        same_direction_ratio = torch.tensor(0.)
+    return same_direction_ratio
+
+
+@torch.no_grad()
+def eff_rank(param: torch.tensor, threshold=1e-10):
+    U, S, Vh = torch.linalg.svd(param.float())
+    rank = torch.sum(S > threshold)
+    return rank
+
+
+# modular neural tangent kernel
+@torch.no_grad()
+def mNTK(module: torch.nn.Module, x: torch.tensor):
+    J_theta_l = jacobian(module, x)
+    mntk = torch.matmul(J_theta_l, J_theta_l.t())
+    return mntk
+
+
+@torch.no_grad()
+def power_iteration(a, num_iterations):
+    b = torch.randn(a.size(1), 1)
+    for _ in range(num_iterations):
+        b = torch.matmul(a, b)
+        b_norm = torch.norm(b)
+        b = b / b_norm if b_norm != 0 else 0
+    eigval = torch.matmul(torch.matmul(b.t(), a), b)
+    return eigval
+
+
+@torch.no_grad()
+def lambda_max_subsample(module: torch.nn.Module, x: torch.tensor, num_iterations=100, subsample_size=None):
+    mntk = mNTK(module, x)
+    if subsample_size is None:
+        subsample_size = min(mntk.size(0), mntk.size(1))
+    idx = torch.randperm(mntk.size(0))[:subsample_size]
+    subsampled = mntk[idx, :]
+    subsampled = subsampled[:, idx]
+    eigval = power_iteration(subsampled, num_iterations)
+    return eigval
+
+
+@torch.no_grad()
+def cal_histc(tensor_cal, bins_total, min_val, max_val):
+    return torch.histc(tensor_cal, bins=bins_total, min=min_val, max=max_val)
+
+
+@torch.no_grad()
+def get_nans(t):
+    return torch.isnan(t).sum()