mindstudio-probe 1.1.1__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/wrap_distributed.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -13,16 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import os
 import re
-import inspect
 
 import torch
-import torch.nn as nn
 import torch.distributed as dist
+import torch.nn as nn
 
-from msprobe.core.common.file_utils import load_yaml
 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:
@@ -57,10 +57,13 @@ class DistributedOPTemplate(nn.Module):
         super(DistributedOPTemplate, self).__init__()
         self.op_name_ = str(op_name)
         self.__name__ = self.op_name_
+        self.cc_hooks = []
         for pre_hook in pre_hooks:
-            self.register_forward_pre_hook(pre_hook, with_kwargs=True)
+            handle = self.register_forward_pre_hook(pre_hook, with_kwargs=True)
+            self.cc_hooks.append(handle)
         for hook in post_hooks:
-            self.register_forward_hook(hook, with_kwargs=True)
+            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)
@@ -120,8 +123,11 @@ class ApiRegistry:
 
     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):
@@ -237,15 +243,20 @@ def create_hooks(context, monitor):
         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
+                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)
+                    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

msprobe/pytorch/monitor/features.py

@@ -15,7 +15,7 @@
 
 import torch
 from torch.autograd.functional import jacobian
-from msprobe.core.common.log import logger
+from msprobe.pytorch.common.log import logger
 
 
 @torch.no_grad()
@@ -30,12 +30,12 @@ def get_min(x: torch.tensor):
 
 @torch.no_grad()
 def get_mean(x: torch.tensor):
-    return torch.mean(x)
+    return torch.mean(x.to(torch.float64))
 
 
 @torch.no_grad()
 def get_norm(x: torch.tensor):
-    return torch.norm(x, p=2)
+    return torch.norm(x.to(torch.float64), p=2)
 
 
 @torch.no_grad()