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

msprobe/pytorch/monitor/optimizer_collect.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.
+# -------------------------------------------------------------------------
+
 from abc import abstractmethod
 
 import torch
@@ -32,7 +34,7 @@ class OptimizerMon(object):
 
     def narrow_from_flatten(self, param, flatten_state):
         return flatten_state
-    
+
     def get_state(self, torch_opt):
         if hasattr(torch_opt, 'chained_optimizers'):
             for opt in torch_opt.chained_optimizers:
@@ -70,7 +72,7 @@ class OptimizerMon(object):
             monitor.register_param_call_id("hook_optimizer", tag)
             grad_dict[tag] = grad
         return grad_dict
-    
+
     def map_fp16_to_fp32_param(self, torch_opt):
         pass
 
@@ -124,7 +126,7 @@ class OptimizerMon(object):
                     monitor.update_heatmap_visualizer[name].pre_cal(update_dict[name])
                     monitor.ratio_heatmap_visualizer[name].pre_cal(ratio_dict[name])
         return MVResult(exp_avg=exp_avg_dict, exp_avg_sq=exp_avg_sq_dict, update=update_dict, ratio=ratio_dict)
-    
+
     def patch_grad_sync(self, monitor):
         def patch_sync(sync_grad_func):
             def wrapper(bucket):
@@ -150,7 +152,7 @@ class OptimizerMon(object):
                 return out
 
             return wrapper
-    
+
         try:
             from megatron.core.distributed.param_and_grad_buffer import Bucket
             self.origin_funcs.append(Bucket.start_grad_sync)
@@ -174,9 +176,8 @@ class OptimizerMon(object):
     def restore_grad_sync(self, monitor):
         if not monitor.enable_megatron:
             return
-        
+
         self.bucket_class.start_grad_sync = self.origin_funcs[0]
-        
 
     def _get_single_state(self, torch_opt):
         state = {}
@@ -191,9 +192,10 @@ class OptimizerMon(object):
 
 class MegatronMixPrecisionOptimizerMon(OptimizerMon):
     """
-    混合精度优化器监控类。在混合精度训练中监控和管理优化器。
+    混合精度优化器监测类。在混合精度训练中监测和管理优化器。
     混合精度训练通过适当降低某些计算的精度来加速训练过程并减少内存消耗。
     """
+
     def map_fp16_to_fp32_param(self, torch_opt):
         for fp16_group, fp32_group in zip(torch_opt.float16_groups, torch_opt.fp32_from_float16_groups):
             for fp16_param, fp32_param in zip(fp16_group, fp32_group):
@@ -237,6 +239,7 @@ class DeepSpeedZeroOptimizerMon(OptimizerMon):
     - Handling gradient collection for different ZeRO stages
     - Managing optimizer state access for monitoring
     """
+
     def __init__(self, torch_opt):
         super().__init__(torch_opt)
         self.stage = ''
@@ -249,12 +252,12 @@ class DeepSpeedZeroOptimizerMon(OptimizerMon):
     @abstractmethod
     def get_grad_for_param(self, lp_param, group_idx, param_id):
         raise NotImplementedError
-    
+
     def param_not_in_partition(self, lp_param, group_idx):
         param_slice_mapping = self.torch_opt.state_dict()['param_slice_mappings'][group_idx]
         hp_address = param_slice_mapping.get(self.torch_opt.param_names.get(lp_param))
         return hp_address is None
-    
+
     def get_position(self, lp_param, group_idx):
         param_slice_mapping = self.torch_opt.state_dict()['param_slice_mappings'][group_idx]
         hp_address = param_slice_mapping.get(self.torch_opt.param_names.get(lp_param))
@@ -266,7 +269,7 @@ class DeepSpeedZeroOptimizerMon(OptimizerMon):
             for param in bit16_group:
                 param2group[param] = group_idx
         return param2group
-    
+
     def get_param_index(self, lp_param, group_idx):
         if not self.param2index:
             for group in self.bit16_groups:
@@ -274,9 +277,9 @@ class DeepSpeedZeroOptimizerMon(OptimizerMon):
                 for index, param in enumerate(group):
                     param2index[param] = index
                 self.param2index.append(param2index)
-                
+
         return self.param2index[group_idx][lp_param]
-    
+
     def narrow_from_flatten(self, param, flatten_state):
         if flatten_state is None:
             return flatten_state
@@ -285,7 +288,7 @@ class DeepSpeedZeroOptimizerMon(OptimizerMon):
             return None
         start, numel = self.get_position(param, group_idx)
         return flatten_state.narrow(0, start, numel)
-        
+
     def map_fp16_to_fp32_param(self, torch_opt):
         for group_idx, group in enumerate(self.bit16_groups):
             for param in group:
@@ -306,7 +309,7 @@ class DeepSpeedZeroOptimizerMon(OptimizerMon):
             grad_dict[tag] = grad
 
         return grad_dict
-    
+
     def patch_grad_sync(self, monitor):
         pass
 
@@ -321,7 +324,7 @@ class DeepSpeedZeroOptimizerStage0Mon(DeepSpeedZeroOptimizerMon):
         self.bit16_groups = torch_opt.bf16_groups
         self.fp32_flat_groups = torch_opt.fp32_groups_flat_partition
         self.param2group = self.get_group_index()
-            
+
     def get_grad_for_param(self, lp_param, group_idx, param_id):
         return self.torch_opt.fp32_groups_gradient_dict[group_idx][param_id]
 
@@ -354,13 +357,12 @@ class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon):
                 else:
                     break
 
-
     def patch_grad_sync(self, monitor):
         def patch_sync(reduce_func):
             def wrapper(zero_optimizer, *args, **kwargs):
                 grad_dict = {}
                 for i, param, _ in zero_optimizer.params_in_ipg_bucket:
-                    if isinstance(param, int): # for ds >= 0.17.0
+                    if isinstance(param, int):  # for ds >= 0.17.0
                         param = zero_optimizer.bit16_groups[i][param]
                     name = monitor.param2name[param]
                     tag = monitor.name2tag.get(name, {}).get(MonitorConst.PRE_GRAD)
@@ -369,14 +371,15 @@ class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon):
                 get_metrics(monitor.ops, grad_dict, monitor.eps, monitor.grad_context.pre)
                 out = reduce_func(zero_optimizer, *args, **kwargs)
                 return out
-            
+
             return wrapper
+
         try:
             from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
             self.origin_funcs = [
-                DeepSpeedZeroOptimizer.average_tensor, 
+                DeepSpeedZeroOptimizer.average_tensor,
                 DeepSpeedZeroOptimizer.buffered_reduce_fallback
-                ]
+            ]
             DeepSpeedZeroOptimizer.average_tensor = patch_sync(DeepSpeedZeroOptimizer.average_tensor)
             DeepSpeedZeroOptimizer.buffered_reduce_fallback = \
                 patch_sync(DeepSpeedZeroOptimizer.buffered_reduce_fallback)
@@ -389,13 +392,12 @@ class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon):
     def restore_grad_sync(self, monitor):
         if not monitor.enable_deepspeed:
             return
-    
+
         from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
         DeepSpeedZeroOptimizer.average_tensor = self.origin_funcs[0]
         DeepSpeedZeroOptimizer.buffered_reduce_fallback = self.origin_funcs[1]
 
 
-
 class DeepSpeedZeroOptimizerStage3Mon(DeepSpeedZeroOptimizerMon):
     def __init__(self, torch_opt):
         super().__init__(torch_opt)
@@ -407,11 +409,11 @@ class DeepSpeedZeroOptimizerStage3Mon(DeepSpeedZeroOptimizerMon):
     def param_not_in_partition(self, lp_param, group_idx):
         """Each param partioned across all zero ranks"""
         return False
-    
+
     def get_position(self, lp_param, group_idx):
         param_id = self.torch_opt.get_param_id(lp_param)
         return self.torch_opt.grad_position[param_id][1:]
-    
+
     def get_grad_for_param(self, lp_param, group_idx, param_id):
         return self.torch_opt.averaged_gradients[group_idx][param_id]
 

msprobe/pytorch/monitor/utils.py

@@ -1,17 +1,19 @@
-# 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 msprobe.pytorch.common.log import logger

msprobe/pytorch/monitor/visualizer.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
 import numpy as np

msprobe/pytorch/torchair_dump/__init__.py

@@ -0,0 +1,17 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# 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://license.coscl.org.cn/MulanPSL2
+#
+# 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 .torchair_dump import set_fx_dump_config, set_ge_dump_config

msprobe/pytorch/torchair_dump/torchair_dump.py

@@ -0,0 +1,114 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# 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://license.coscl.org.cn/MulanPSL2
+#
+# 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 os
+
+from msprobe.core.common.file_utils import create_directory, check_path_readability, check_path_writability
+from msprobe.core.common.const import Const
+from msprobe.core.common.log import logger
+
+
+
+def try_import_torchair():
+    try:
+        import torch
+        import torch_npu
+        import torchair
+    except ModuleNotFoundError as e:
+        logger.error("torch or torch_npu with torchair not found. Try installing the latest torch_npu.")
+        raise e
+
+
+def set_ge_dump_config(
+        dump_path='',
+        dump_mode='all',
+        fusion_switch_file=None,
+        dump_token=None,
+        dump_layer=None,
+        compiler_config=None
+):
+    try_import_torchair()
+    from torchair.configs.compiler_config import CompilerConfig
+
+    dump_path = os.path.realpath(os.path.join(dump_path, Const.TORCHAIR_GE_DATA_DIRECTORY))
+    create_directory(dump_path)
+    check_path_readability(dump_path)
+    check_path_writability(dump_path)
+    if os.listdir(dump_path):
+        logger.warning('The specified directory for GE-Dump data is not empty, which may result in data mixing.')
+
+    if compiler_config is not None and not isinstance(compiler_config, CompilerConfig):
+        raise TypeError(f'compiler_config must be a CompilerConfig, but got {type(compiler_config)}')
+    config = compiler_config if compiler_config else CompilerConfig()
+    # Generate GE mapping graph
+    config.debug.graph_dump.type = "txt"
+    if hasattr(config.debug.graph_dump, "_path"):  # interface changed since 8.0.RC1.b080
+        setattr(config.debug.graph_dump, "_path", dump_path)
+    else:
+        config.debug.graph_dump.path = dump_path
+
+    if fusion_switch_file is not None:
+        if not os.path.exists(fusion_switch_file):
+            raise FileNotFoundError(f'fusion_switch_file: {fusion_switch_file} not found')
+        config.fusion_config.fusion_switch_file = fusion_switch_file
+
+    # Enable GE dump
+    config.dump_config.enable_dump = True
+    config.dump_config.dump_mode = dump_mode
+    config.dump_config.dump_path = dump_path
+
+    if dump_token is not None:
+        new_token = [str(x) for x in dump_token]
+        config.dump_config.dump_step = "|".join(new_token)
+    if dump_layer is not None:
+        dump_layer = " ".join(dump_layer)
+        config.dump_config.dump_layer = dump_layer
+
+    return config
+
+
+def set_fx_dump_config(dump_path='', op_list=None, compiler_config=None):
+    try_import_torchair()
+    from torchair.configs.compiler_config import CompilerConfig
+
+    dump_path = os.path.realpath(os.path.join(dump_path, Const.TORCHAIR_FX_DATA_DIRECTORY))
+    create_directory(dump_path)
+    check_path_readability(dump_path)
+    check_path_writability(dump_path)
+    if os.listdir(dump_path):
+        logger.warning('The specified directory for FX-Dump data is not empty, which may result in data mixing.')
+
+    if compiler_config is not None and not isinstance(compiler_config, CompilerConfig):
+        raise TypeError(f'compiler_config must be a CompilerConfig, but got {type(compiler_config)}')
+    config = compiler_config if compiler_config else CompilerConfig()
+    if op_list:
+        if not isinstance(op_list, list) or not all(isinstance(item, str) for item in op_list):
+            raise TypeError("op_list must be a list of strings")
+        from msprobe.infer.utils.util import filter_cmd
+        op_list = filter_cmd(op_list)
+        config.debug.data_dump.filter = (
+            lambda x: x
+            if any(keyword.lower() in x.name.lower() for keyword in op_list)
+            else None
+        )
+        
+
+    # Enable FX dump
+    config.debug.data_dump.type = "npy"
+    if hasattr(config.debug.data_dump, "path"):
+        config.debug.data_dump.path = dump_path
+
+    return config

msprobe/scripts/atb/config_example.json

@@ -0,0 +1,10 @@
+{
+    "task": "tensor",
+    "dump_enable": false,
+    "exec_range": "all",
+    "ids": "",
+    "op_name": "",
+    "save_child": false,
+    "device": "",
+    "filter_level": 1
+}

msprobe/scripts/atb/load_atb_probe.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+#
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# 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://license.coscl.org.cn/MulanPSL2
+#
+# 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.
+# -------------------------------------------------------------------------
+
+function get_cxx_abi_option()
+{
+    cxx_abi=""
+    if [[ "$ATB_HOME_PATH" =~ cxx_abi_0$ ]]; then
+        cxx_abi=0
+    else
+        if [[ "$ATB_HOME_PATH" =~ cxx_abi_1$ ]]; then
+            cxx_abi=1
+        else
+            echo "[WARNING] The running environment of ATB model is not ready. ATB probe would not work"
+        fi
+    fi
+}
+
+function get_output_option()
+{
+    output_path="$(pwd)"
+    if [ -n "$1" ]; then
+        if [[ "$1" =~ ^--output=.* ]]; then
+            var="$1"
+            output_path="${var#*=}"
+        fi
+    fi
+    if [ -n "$2" ]; then
+        if [[ "$2" =~ ^--output=.* ]]; then
+            var="$2"
+            output_path="${var#*=}"
+        fi
+    fi
+    left="${output_path%%/*}"
+    right="${output_path#*/}"
+    if [ -n "$left" ]; then
+        if [ "$left" == "$output_path" ]; then
+            output_path=$(realpath $output_path)
+        else
+            output_path="$(realpath $left)"/${right}
+        fi
+    fi
+}
+
+function get_config_option()
+{
+    config_path=$(dirname $(realpath "${BASH_SOURCE[0]}"))/config.json
+    if [ -n "$1" ]; then
+        if [[ "$1" =~ ^--config=.* ]]; then
+            var="$1"
+            config_path="${var#*=}"
+        fi
+    fi
+    if [ -n "$2" ]; then
+        if [[ "$2" =~ ^--config=.* ]]; then
+            var="$2"
+            config_path="${var#*=}"
+        fi
+    fi
+    left="${config_path%%/*}"
+    right="${config_path#*/}"
+    if [ -n "$left" ]; then
+        if [ "$left" == "$config_path" ]; then
+            config_path=$(realpath $config_path)
+        else
+            config_path="$(realpath $left)"/${right}
+        fi
+    fi
+}
+
+if [ $# -gt 2 ]; then
+    echo "[WARNING] The number of parameters should not be greater than 2, but got $#. ATB probe would not work"
+else
+    if [ -z "$ASCEND_HOME_PATH" ]; then
+        echo "[WARNING] ASCEND_HOME_PATH enviroment variable is not set. ATB probe would not work"
+    else
+        get_cxx_abi_option
+        if [ -n "$cxx_abi" ]; then
+            get_output_option "$@"
+            get_config_option "$@"
+            lib_path=$(realpath $(dirname $(realpath "${BASH_SOURCE[0]}"))/../../lib)
+            export LD_PRELOAD=${lib_path}/libatb_probe_abi${cxx_abi}.so:$LD_PRELOAD
+            export ATB_OUTPUT_DIR=$output_path
+            export ATB_DUMP_CONFIG=$config_path
+        fi
+    fi
+fi

msprobe/scripts/atb/unload_atb_probe.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+#
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# 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://license.coscl.org.cn/MulanPSL2
+#
+# 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.
+# -------------------------------------------------------------------------
+
+if [ -n "$LD_PRELOAD" ]; then
+    var=$(echo ":$LD_PRELOAD" | sed 's|:[^:]*libatb_probe_abi0\.so|:|g')
+    var=$(echo ":$var" | sed 's|:[^:]*libatb_probe_abi1\.so|:|g')
+    var=$(echo $var | sed 's|:[:]*:|:|g')
+    var=$(echo $var | sed 's|^:[:]*||')
+    export LD_PRELOAD=${var}
+fi
+unset ATB_OUTPUT_DIR
+unset ATB_DUMP_CONFIG