diff --git a/debug/accuracy_tools/msprobe/docs/19.monitor.md b/debug/accuracy_tools/msprobe/docs/19.monitor.md
index f57fadf9afda3b4718da21cbce8714e0632e006e..f7c9028b26f557ad9ad40418d0593c0a158905b4 100644
--- a/debug/accuracy_tools/msprobe/docs/19.monitor.md
+++ b/debug/accuracy_tools/msprobe/docs/19.monitor.md
@@ -228,7 +228,7 @@ monitor.monitor_gnorm_with_ad(
module_name可以通过nn.Module的接口named_modules()获取。
#### 打印模型结构
-工具提供可选项`print_struct`打印模型结构，帮助配置targets。工具会在在第一个step后打印结构并停止训练进程，模型结构默认打印在`$MONITOR_OUTPUT_DIR/module_struct.json`。
+工具提供可选项`print_struct`打印模型结构，帮助配置targets。工具会在在第一个step后打印结构并停止训练进程，每张卡上的模型结构默认保存在`$MONITOR_OUTPUT_DIR/module_struct/rank{rank}/module_struct.json`, 其中{rank}为对应的卡号。
```json
{
"print_struct": true
@@ -701,7 +701,7 @@ TrainerMon.monitor_gnorm_with_ad(model, grad_acc_steps, optimizer, dp_group, tp_
| "collect_times" | 可选 | 设置采集次数，达到该次数后停止监控，默认值为100000000，目的是一直采集。 |
| "start_step" | 可选 | 设置开始采集step，模型训练达到start_step后开始监控采集，默认值为0，表示从step0开始监控采集。注：在动态启停模式下该设置不生效，只会从下一步开始监控采集。 |
| "step_interval" | 可选 | 设置采集step间隔，默认值为1，表示每个step均采集监控数据。 |
-| "print_struct" | 可选 | 设置为true后监控工具会打印模型中torch module的名字和详细结构，并在第1个step后退出。不填默认为false。**仅PyTorch场景支持此参数**。 |
+| "print_struct" | 可选 | 设置为true后监控工具会打印每张卡模型中module的名字和详细结构，并在第1个step后退出。不填默认为false。 |
| "module_ranks" | 可选 | 用于在分布式训练场景中希望控制在哪些rank开启module监控。如果不填，则默认在所有rank开启。 列表内rank要求为int类型。 |
| "ur_distribution" | 可选 | 若为true则会统计adam优化器指定模块（targets中指定）参数的update和ratio向量的数值分布，并展示在heatmap里，默认为false，同时format字段必须设置为tensorboard。
依赖histc算子， 需要CANN8.0.rc2以上版本， 否则会有严重的性能问题。**仅PyTorch场景支持此参数**。 |
| "xy_distribution" | 可选 | 若为true则会监控指定module（targets中指定）的输入输出张量。 默认为false。 |
@@ -710,7 +710,7 @@ TrainerMon.monitor_gnorm_with_ad(model, grad_acc_steps, optimizer, dp_group, tp_
| "backward_only" | 可选 | 开启xy_distribution后生效，若为true，仅监控指定module的反向，targets中的input、output不生效。默认为false。 |
| "mv_distribution" | 可选 | 若为true则会监控指定模块中的参数的优化器状态， 默认为false。版本 save module struct to {module_struct_file}")
+ output_dir = os.path.join(get_output_base_dir(), 'module_struct', f'rank{self.rank}')
+ make_dir(output_dir)
+ module_struct_file = os.path.realpath(os.path.join(output_dir, 'module_struct.json'))
+ save_json(module_struct_file, self.module_struct, indent=2)
+ logger.info(f"> save module struct to {module_struct_file}")
self.struct_printed = True
def _hook_module(self, target_names, module, vpp_stage=''):
diff --git a/debug/accuracy_tools/msprobe/pytorch/monitor/module_hook.py b/debug/accuracy_tools/msprobe/pytorch/monitor/module_hook.py
index edda3d4b59f26b0ad837f13ef74447a2f1eab3a2..61e9ed4b00d2e285849a68d200db6bfa137b850a 100644
--- a/debug/accuracy_tools/msprobe/pytorch/monitor/module_hook.py
+++ b/debug/accuracy_tools/msprobe/pytorch/monitor/module_hook.py
@@ -26,7 +26,7 @@ import pandas as pd
from torch.utils.hooks import BackwardHook
from msprobe.core.common.const import MonitorConst, Const
-from msprobe.core.common.file_utils import load_json, save_json
+from msprobe.core.common.file_utils import load_json, save_json, make_dir
from msprobe.core.common.decorator import recursion_depth_decorator
from msprobe.core.monitor.anomaly_processor import AnomalyScanner, AnomalyDataFactory, AnomalyDataWriter
from msprobe.core.common.file_utils import write_df_to_csv
@@ -46,12 +46,6 @@ from msprobe.pytorch.monitor.module_metric import get_metrics, get_summary_write
from msprobe.pytorch.monitor.optimizer_collect import OptimizerMonFactory
from msprobe.pytorch.monitor.visualizer import HeatmapVisualizer
-try:
- from megatron.core import mpu
-except ImportError:
- MPU_IMPORT = False
-else:
- MPU_IMPORT = True
torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
if not torch_version_above_or_equal_2:
@@ -1006,19 +1000,11 @@ class TrainerMon:
logger.info(msg)
def _save_module_struct(self):
- if MPU_IMPORT:
- pp_group = mpu.get_pipeline_model_parallel_group()
- pp_group_list = torch.distributed.get_process_group_ranks(pp_group)
- save_module_struct = not dist.is_initialized() or dist.get_rank() in pp_group_list
- else:
- save_module_struct = (not dist.is_initialized()
- or (self.module_rank_list and dist.get_rank() == min(self.module_rank_list))
- or (not self.module_rank_list and dist.get_rank() == 0))
- if save_module_struct:
- module_struct_file = os.path.realpath(
- os.path.join(get_output_base_dir(), f'{dist.get_rank()}_module_struct.json'))
- save_json(module_struct_file, self.module_struct, indent=2)
- logger.info(f"> save module struct to {module_struct_file}")
+ output_dir = os.path.join(get_output_base_dir(), 'module_struct', f'rank{self.rank}')
+ make_dir(output_dir)
+ module_struct_file = os.path.realpath(os.path.join(output_dir, 'module_struct.json'))
+ save_json(module_struct_file, self.module_struct, indent=2)
+ logger.info(f"> save module struct to {module_struct_file}")
self.struct_printed = True
def _is_target_param(self, param_name, param, prefix):
diff --git a/debug/accuracy_tools/msprobe/pytorch/monitor/module_metric.py b/debug/accuracy_tools/msprobe/pytorch/monitor/module_metric.py
index d0ff16a97d2410362a883429becdd49559674ff3..ad10580ad55b47f8604ee4dfc27a3290ed693bf9 100644
--- a/debug/accuracy_tools/msprobe/pytorch/monitor/module_metric.py
+++ b/debug/accuracy_tools/msprobe/pytorch/monitor/module_metric.py
@@ -34,7 +34,8 @@ def squash_param_name(param_name, enable=True):
if not enable:
return param_name
name = ''
- for pattern in ['layers?\.(.*)', 'embeddings?\.(.*)', 'final.*', 'output.*', 'norm.*']:
+ for pattern in ['^.*\.(layers?\..*)', '^.*\.(embeddings?\..*)', '^.*\.(final.*)', '^.*\.(output.*)',
+ '^.*\.(norm.*)']:
match = re.findall(pattern, param_name)
if match:
name += match[0]