mindstudio-probe 8.3.2__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/pytorch/{pytorch_service.py → dump/pytorch_service.py}

@@ -1,33 +1,34 @@
-# Copyright (c) 2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 from msprobe.core.common.utils import Const
-from msprobe.core.service import BaseService
+from msprobe.core.dump.service import BaseService
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.common.utils import get_rank_if_initialized
-from msprobe.pytorch.dump.module_dump.module_processer import ModuleProcesser
-from msprobe.pytorch.hook_module.api_register import (
+from msprobe.pytorch.dump.module_dump.module_processor import ModuleProcessor
+from msprobe.pytorch.dump.api_dump.api_register import (
     get_api_register,
     ApiTemplate,
     redirect_wait,
     reset_dist_collect_func
 )
-from msprobe.pytorch.hook_module.hook_module import HOOKModule
-from msprobe.pytorch.hook_module.pt_hook_manager import PytorchHookManager
-from msprobe.pytorch.hook_module.register_optimizer_hook import register_optimizer_hook
-from msprobe.pytorch.hook_module.script_wrapper import wrap_script_func, preprocess_func
+from msprobe.pytorch.dump.api_dump.hook_module import HOOKModule
+from msprobe.pytorch.dump.api_dump.pt_hook_manager import PytorchHookManager
+from msprobe.pytorch.dump.api_dump.register_optimizer_hook import register_optimizer_hook
+from msprobe.pytorch.dump.api_dump.script_wrapper import wrap_script_func, preprocess_func
 
 
 class PytorchService(BaseService):
@@ -45,7 +46,7 @@ class PytorchService(BaseService):
     def _init_specific_components(self):
         self.logger = logger
         self.api_register = get_api_register()
-        self.module_processor = ModuleProcesser(self.data_collector.scope)
+        self.module_processor = ModuleProcessor(self.data_collector.scope)
         self.hook_manager = PytorchHookManager(self.data_collector, self.config)
         self.api_template = ApiTemplate
 
@@ -60,12 +61,12 @@ class PytorchService(BaseService):
         redirect_wait()
 
     def _register_module_hook(self):
-        ModuleProcesser.enable_module_dump = True
+        ModuleProcessor.enable_module_dump = True
         self.module_processor.register_module_hook(self.model, self.build_hook)
         self.logger.info(f"The module {self.config.task} hook function is successfully mounted to the model.")
 
     def _reset_status(self):
         super()._reset_status()
-        ModuleProcesser.reset_module_stats()
+        ModuleProcessor.reset_module_stats()
         HOOKModule.reset_module_stats()
         reset_dist_collect_func()

msprobe/pytorch/monitor/csv2tb.py

@@ -1,17 +1,19 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
 import datetime
 import os
 import re

msprobe/pytorch/monitor/data_writers.py

@@ -1,17 +1,19 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
 import itertools
 import os
 from collections import defaultdict

msprobe/pytorch/monitor/distributed/wrap_distributed.py

@@ -1,17 +1,19 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
 
 import inspect
 import os

msprobe/pytorch/monitor/features.py

@@ -1,17 +1,18 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 import torch
 from torch.autograd.functional import jacobian

msprobe/pytorch/monitor/module_hook.py

@@ -1,17 +1,19 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
 import json
 import os
 import uuid
@@ -29,7 +31,7 @@ from torch.utils.hooks import BackwardHook
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
 from msprobe.core.common.const import MonitorConst, Const
-from msprobe.core.common.file_utils import load_json, save_json, make_dir
+from msprobe.core.common.file_utils import load_json, save_json, create_directory
 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
@@ -169,7 +171,7 @@ class TrainerMon:
         self.ratio_heatmap_visualizer = defaultdict(HeatmapVisualizer)
         self.fsdp_post_backward_hook = None
         self.fsdp2_foreach_reduce = None
-        self.config_timestamp = 0  # 后面有校验时间戳, 首次监控无需为了更新config文件时间戳而去改, 可通过dynamic_on开关直接打开
+        self.config_timestamp = 0  # 后面有校验时间戳, 首次监测无需为了更新config文件时间戳而去改, 可通过dynamic_on开关直接打开
         self.config = load_json(config_file_path)
         validate_config(self.config)
 
@@ -210,7 +212,7 @@ class TrainerMon:
         self.optimizer_mon = None
         self.optimizer_trans = None
 
-        # TYPE3: 会随着训练中途config配置更新或监控状态改变而重置的变量
+        # TYPE3: 会随着训练中途config配置更新或监测状态改变而重置的变量
         self.module_fwd_hook_context_by_module = defaultdict(ModuleHookContext)
         self.module_bwd_hook_context_by_module = defaultdict(ModuleHookContext)
         self.feature_hook_context_by_module = defaultdict(FeatureHookContext)
@@ -225,8 +227,7 @@ class TrainerMon:
         self.duplicate_param = {}
         self.name2tag = {}
         self.param_name_call_id = {}
-        self.flat_prefix_names = []
-        self.flat_prefix_reverse_iter = None
+        self.fsdp_param_name_map = {}
         self.call_id = 0
         self.module_struct = defaultdict(dict)
         self.grad_accs = []
@@ -844,6 +845,7 @@ class TrainerMon:
             bwd_context.reset()
         self.grad_context.reset()  # 权重梯度和激活值梯度都在这
 
+        # 还原梯度patch
         self.optimizer_mon.restore_grad_sync(self)
         if self.fsdp_post_backward_hook:  # fsdp
             torch.distributed.fsdp._runtime_utils._post_backward_hook = self.fsdp_post_backward_hook
@@ -914,7 +916,7 @@ class TrainerMon:
 
     def _save_module_struct(self):
         output_dir = os.path.join(get_output_base_dir(), 'module_struct', f'rank{self.rank}')
-        make_dir(output_dir)
+        create_directory(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}")
@@ -939,10 +941,8 @@ class TrainerMon:
                 continue
             if not self.fsdp2_wrapped_module and param.__class__.__name__ == "DTensor":
                 self.fsdp2_wrapped_module = True
-            if self.fsdp_wrapped_module:  # FSDP1需要记录完整的不被target限制的flat权重前缀名，以供后续对flat解包
-                flat_prefix_name, _ = param_name.rsplit(MonitorConst.FSDP_FLAT_SEP, 1)
-                if flat_prefix_name not in self.flat_prefix_names:
-                    self.flat_prefix_names.append(flat_prefix_name)
+            if self.fsdp_wrapped_module:  # FSDP1需要记录完整的不被target限制的flat权重名，以供后续对flat解包
+                self.fsdp_param_name_map[id(param)] = param_name
 
             if self._is_target_param(param_name, param, prefix):
                 name = prefix + squash_param_name(param_name, self.squash_name)
@@ -967,8 +967,6 @@ class TrainerMon:
                     k: get_summary_writer_tag_name(name, k, self.rank)
                     for k in keywords
                 }
-        if self.fsdp_wrapped_module:
-            self.flat_prefix_reverse_iter = cycle(reversed(self.flat_prefix_names))  # post_backward_hook调用顺序是反向的
 
     def _register_param_name(self):
         for vpp_stage, model_chunk in enumerate(self.model):
@@ -1189,6 +1187,7 @@ class TrainerMon:
     def _patch_grad_sync(self):
         if not self.wg_distribution:
             return
+
         if self.fsdp_wrapped_module:
             # patch fsdp _runtime_utils._post_backward_hook
             self._patch_fsdp_post_backward_hook()
@@ -1221,24 +1220,32 @@ class TrainerMon:
 
         def patch_post_backward_hook(_post_backward_hook):
             def wrapper(state, handle, *unused):
-                grad_dict = {}
-                local_names = handle.flat_param._fqns
-                offsets = handle._get_flat_param_offsets()
-                shapes = handle.flat_param._shapes
-                flat_prefix = next(self.flat_prefix_reverse_iter)
-                for local_name, (start, end), local_shape in zip(local_names, offsets, shapes):
-                    grad_clip = handle.flat_param.grad[start:end + 1]
-                    grad = grad_clip.reshape(local_shape)
-                    total_name = f"{flat_prefix}{MonitorConst.FSDP_FLAT_SEP}{local_name}"
-                    if total_name not in self.origin2squash:
-                        logger.warning(f"{total_name} not in model.named_parameters(), skip.")
-                        continue
-                    tag = self.name2tag.get(self.origin2squash[total_name], {}).get(MonitorConst.PRE_GRAD)
-                    if tag is None:
-                        continue
-                    grad_dict[tag] = grad
-                    self.register_param_call_id("_post_backward_hook", tag)
-                get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
+                if not hasattr(handle.flat_param, 'micro_step'):
+                    setattr(handle.flat_param, 'micro_step', 0)
+
+                handle_not_in_any_model = all(
+                    model._all_handles is None or handle not in model._all_handles
+                    for model in self.model
+                )
+                if handle_not_in_any_model:
+                    logger.debug(f"Skip handle not along to the monitor model list, handle id: {id(handle)}.")
+                elif self.monitor_mbs_grad or (handle.flat_param.micro_step + 1 == self.micro_batch_number):
+                    grad_dict = {}
+                    offsets = handle._get_flat_param_offsets()
+                    shapes = handle.flat_param._shapes
+                    params = handle.flat_param._params
+                    for param, (start, end), local_shape in zip(params, offsets, shapes):
+                        param_name = self.fsdp_param_name_map.get(id(param), "")
+                        tag = self.name2tag.get(self.origin2squash.get(param_name, ""), {}).get(MonitorConst.PRE_GRAD)
+                        if tag is None:
+                            continue
+                        grad = handle.flat_param.grad[start:end + 1].reshape(local_shape)
+                        if self.monitor_mbs_grad:
+                            tag = tag.replace('/', f'{MonitorConst.NAME_SEP}{handle.flat_param.micro_step}/', 1)
+                        grad_dict[tag] = grad
+                        self.register_param_call_id("_post_backward_hook", tag)
+                    get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
+                handle.flat_param.micro_step = (handle.flat_param.micro_step + 1) % self.micro_batch_number
                 out = _post_backward_hook(state, handle, *unused)
                 return out
 
@@ -1252,14 +1259,21 @@ class TrainerMon:
     def _patch_fsdp2_foreach_reduce(self):
         def patch_foreach_reduce(foreach_reduce):
             def wrapper(fsdp_params, unsharded_grads, *unused):
-                grad_dict = {}
-                for param, grad in zip(fsdp_params, unsharded_grads):
-                    tag = self.name2tag.get(self.origin2squash[param._param_fqn], {}).get(MonitorConst.PRE_GRAD)
-                    if tag is None:
-                        continue
-                    grad_dict[tag] = grad
-                    self.register_param_call_id("foreach_reduce", tag)
-                get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
+                if not hasattr(fsdp_params[0], 'micro_step'):
+                    setattr(fsdp_params[0], 'micro_step', 0)
+
+                if self.monitor_mbs_grad or (fsdp_params[0].micro_step + 1 == self.micro_batch_number):
+                    grad_dict = {}
+                    for param, grad in zip(fsdp_params, unsharded_grads):
+                        tag = self.name2tag.get(self.origin2squash.get(param._param_fqn, ""), {}).get(MonitorConst.PRE_GRAD)
+                        if tag is None:
+                            continue
+                        if self.monitor_mbs_grad:
+                            tag = tag.replace('/', f'{MonitorConst.NAME_SEP}{fsdp_params[0].micro_step}/', 1)
+                        grad_dict[tag] = grad
+                        self.register_param_call_id("foreach_reduce", tag)
+                    get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
+                fsdp_params[0].micro_step = (fsdp_params[0].micro_step + 1) % self.micro_batch_number
                 out = foreach_reduce(fsdp_params, unsharded_grads, *unused)
                 return out
 
@@ -1285,20 +1299,15 @@ class TrainerMon:
 
             key = get_summary_writer_tag_name(key, 'acc_grad', self.rank)
             self.register_param_call_id("param_hook", key)
-            param.micro_step += 1
-            grad_dict = {}
-            if self.monitor_mbs_grad or (param.micro_step == self.micro_batch_number):
+            if self.monitor_mbs_grad or (param.micro_step + 1 == self.micro_batch_number):
                 if self.params_have_main_grad:
                     grad = param.main_grad
                 else:
                     grad = param.grad
-                grad_dict[key] = grad.clone()
+                get_metrics(self.ops, {key: grad.clone()}, self.eps, self.grad_context.pre)
+            param.micro_step = (param.micro_step + 1) % self.micro_batch_number
 
-            get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
-            if param.micro_step == self.micro_batch_number:
-                param.micro_step = 0
-
-        logger.info("hooking weights.")
+        logger.info("hooking weight grads.")
         for param, name in self.param2name.items():
             setattr(param, 'micro_step', 0)
             param_tmp = param.expand_as(param)
@@ -1307,5 +1316,4 @@ class TrainerMon:
                 partial(param_hook, param=param, name=name))
             self.grad_accs.append(grad_acc)
             self.handles['wgrads'].append(handle)
-
         self.weight_hooked = True

msprobe/pytorch/monitor/module_metric.py

@@ -1,17 +1,19 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
+
 import re
 
 import torch