mindstudio-probe 1.2.2__py3-none-any.whl → 8.1.0__py3-none-any.whl

msprobe/mindspore/monitor/optimizer_collect.py

@@ -0,0 +1,331 @@
+# 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.
+from abc import abstractmethod
+
+from mindspore import mint, ops
+
+from msprobe.mindspore.common.log import logger
+from msprobe.core.common.const import MonitorConst
+
+
+class OptimizerMon(object):
+    def __init__(self, optim) -> None:
+        self.fp16_to_fp32_param = {}
+        self.optim = optim
+        self.state = {}
+
+    def narrow_from_flatten(self, param, flatten_state):
+        return flatten_state
+    
+    def get_state(self, optim):
+        if hasattr(optim, 'chained_optimizers'):
+            for opt in optim.chained_optimizers:
+                self._get_single_state(opt)
+        else:
+            self._get_single_state(optim)
+
+    def fetch_grad(self, monitor, params2name):
+        if not self.fp16_to_fp32_param:
+            self.map_fp16_to_fp32_param(self.optim)
+
+        grad_dict = {}
+        first_param = True
+        for param, name in params2name.items():
+            if monitor.duplicate_param.get(name, False):
+                continue
+            if self.fp16_to_fp32_param and param not in self.fp16_to_fp32_param:
+                continue
+            grad = param.main_grad if monitor.params_have_main_grad else param.grad
+            element_in_cur_partition = self.fp16_to_fp32_param.get(param, param).numel()
+            if param.numel() != element_in_cur_partition:
+                if first_param:
+                    grad = grad.flatten()[-element_in_cur_partition:]
+                else: # supposed to be the last one
+                    grad = grad.flatten()[:element_in_cur_partition]
+            first_param = False
+            if grad is None:
+                continue
+            tag = monitor.name2tag.get(name, {}).get(MonitorConst.POST_GRAD)
+            monitor.register_param_call_id("hook_optimizer", tag)
+            grad_dict[tag] = grad
+        return grad_dict
+    
+    def map_fp16_to_fp32_param(self, optim):
+        pass
+
+    def fetch_mv(self, monitor, params2name):
+        if not self.fp16_to_fp32_param:
+            self.map_fp16_to_fp32_param(self.optim)
+        if not self.state:
+            self.get_state(self.optim)
+
+        exp_avg_dict = {}
+        exp_avg_sq_dict = {}
+        update_dict = {}
+        ratio_dict = {}
+
+        if not self.state:
+            logger.warning('optimizer state can not accessed')
+            return exp_avg_dict, exp_avg_sq_dict, update_dict, ratio_dict
+
+        for lp_param, name in params2name.items():
+            if lp_param in self.fp16_to_fp32_param:
+                hp_param = self.fp16_to_fp32_param[lp_param]
+            else:
+                hp_param = lp_param
+
+            if hp_param in self.state:
+                state_param = self.state.get(hp_param, {})
+                exp_avg = self.narrow_from_flatten(lp_param, state_param.get("exp_avg", None))
+                exp_avg_sq = self.narrow_from_flatten(lp_param, state_param.get("exp_avg_sq", None))
+                if monitor.mv_distribution:
+                    exp_avg_dict[name] = exp_avg
+                    exp_avg_sq_dict[name] = exp_avg_sq
+                if monitor.mg_direction:
+                    exp_avg_dict[name] = exp_avg
+                if monitor.ur_distribution:
+                    if len(self.optim.param_groups) > 1:
+                        logger.info(f"the length of optim.param_groups is {len(self.optim.param_groups)}.")
+                    if 'step' in state_param:
+                        step = state_param['step']  # Optimizer from pytorch or FusedAdam from apex(used by megatron)
+                    elif 'step' in self.optim.param_groups[0]:
+                        step = self.optim.param_groups[0]['step']  # AdamW from mindspeed
+                    else:
+                        logger.warning(f"step of {name} is None, maybe something wrong happened.")
+                        continue
+                    exp_avg_hat = exp_avg / (1 - self.optim.defaults['betas'][0] ** step)
+                    exp_avg_sq_hat = exp_avg_sq / (1 - self.optim.defaults['betas'][1] ** step)
+                    update_dict[name] = exp_avg_hat / (mint.sqrt(exp_avg_sq_hat) + self.optim.defaults['eps'])
+                    ratio_dict[name] = exp_avg_hat / mint.sqrt(exp_avg_sq_hat)
+                    monitor.update_heatmap_visualizer[name].pre_cal(update_dict[name])
+                    monitor.ratio_heatmap_visualizer[name].pre_cal(ratio_dict[name])
+        return exp_avg_dict, exp_avg_sq_dict, update_dict, ratio_dict
+    
+    def _get_single_state(self, optim):
+        state = {}
+        if hasattr(optim, 'param_to_cpu_states_map'):
+            state = optim.param_to_cpu_states_map
+        elif hasattr(optim, 'state'):
+            state = optim.state
+        elif hasattr(optim, 'optimizer') and hasattr(optim.optimizer, 'state'):
+            state = optim.optimizer.state
+        self.state.update(state)
+
+
+class MixPrecisionOptimizerMon(OptimizerMon):
+    """
+    混合精度优化器监控类。在混合精度训练中监控和管理优化器。
+    混合精度训练通过适当降低某些计算的精度来加速训练过程并减少内存消耗。
+    """
+    def map_fp16_to_fp32_param(self, optim):
+        for fp16_group, fp32_group in zip(optim.float16_groups, optim.fp32_from_float16_groups):
+            for fp16_param, fp32_param in zip(fp16_group, fp32_group):
+                self.fp16_to_fp32_param[fp16_param] = fp32_param
+
+
+class MegatronDistributedOptimizerMon(OptimizerMon):
+    def map_fp16_to_fp32_param(self, optim):
+        if not (hasattr(optim, "model_float16_groups") and
+                hasattr(optim, "shard_fp32_from_float16_groups")):
+            raise Exception(
+                "megatron distributed optimizer should have model_float16_groups and shard_fp32_from_float16_groups, "
+                "if not, please check megatron-lm version")
+        for fp16_group, shard_fp32_group in zip(optim.model_float16_groups,
+                                                optim.shard_fp32_from_float16_groups):
+            for fp16_param, shard_fp32_param in zip(fp16_group, shard_fp32_group):
+                self.fp16_to_fp32_param[fp16_param] = shard_fp32_param
+
+
+class MegatronChainedDistributedOptimizerMon(MegatronDistributedOptimizerMon):
+    def map_fp16_to_fp32_param(self, optim):
+        for opt in optim.chained_optimizers:
+            super().map_fp16_to_fp32_param(opt)
+
+
+class MegatronChainedMixPrecisionOptimizerMon(MixPrecisionOptimizerMon):
+    def map_fp16_to_fp32_param(self, optim):
+        for opt in optim.chained_optimizers:
+            super().map_fp16_to_fp32_param(opt)
+
+
+class DeepSpeedZeroOptimizerMon(OptimizerMon):
+    """
+    Base monitor class for DeepSpeed ZeRO optimizer.
+    ZeRO stage 0 no partition
+    ZeRO stage 1 partitions optimizer states across data parallel processes.
+    ZeRO stage 2 additionally partitions gradients.
+    ZeRO stage 3 additionally partitions parameters.
+
+    This class provides monitoring capabilities for ZeRO optimizers by:
+    - Handling gradient collection for different ZeRO stages
+    - Managing optimizer state access for monitoring
+    """
+    def __init__(self, optim):
+        super().__init__(optim)
+        self.stage = ''
+        self.bit16_groups = []
+        self.fp32_flat_groups = []
+        self.param2group = ()
+        self.param2index = []
+        self.group_offset = {}
+
+    @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.optim.state_dict()['param_slice_mappings'][group_idx]
+        hp_address = param_slice_mapping.get(self.optim.param_names.get(lp_param))
+        return hp_address is None
+    
+    def get_position(self, lp_param, group_idx):
+        param_slice_mapping = self.optim.state_dict()['param_slice_mappings'][group_idx]
+        hp_address = param_slice_mapping.get(self.optim.param_names.get(lp_param))
+        return hp_address.start, hp_address.numel
+
+    def get_group_index(self):
+        param2group = {}
+        for group_idx, bit16_group in enumerate(self.bit16_groups):
+            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:
+                param2index = {}
+                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
+        group_idx = self.param2group[param]
+        if self.param_not_in_partition(param, group_idx):
+            return None
+        start, numel = self.get_position(param, group_idx)
+        return flatten_state.narrow(0, start, numel)
+        
+    def map_fp16_to_fp32_param(self, optim):
+        for group_idx, group in enumerate(self.bit16_groups):
+            for param in group:
+                self.fp16_to_fp32_param[param] = self.fp32_flat_groups[group_idx]
+
+    def fetch_grad(self, monitor, params2name):
+        grad_dict = {}
+        for lp_param, name in params2name.items():
+            group_idx = self.param2group[lp_param]
+            param_id = self.get_param_index(lp_param, group_idx)
+            if self.param_not_in_partition(lp_param, group_idx):
+                continue
+            if self.stage == '1or2':
+                param_id = param_id - self.group_offset[group_idx] - 1
+            grad = self.get_grad_for_param(lp_param, group_idx, param_id)
+            tag = monitor.name2tag.get(name, {}).get(MonitorConst.POST_GRAD)
+            monitor.register_param_call_id("hook_optimizer", tag)
+            grad_dict[tag] = grad
+
+        return grad_dict
+
+
+class DeepSpeedZeroOptimizerStage0Mon(DeepSpeedZeroOptimizerMon):
+    def __init__(self, optim):
+        super().__init__(optim)
+        self.stage = '0'
+        self.bit16_groups = optim.bf16_groups
+        self.fp32_flat_groups = optim.fp32_groups_flat_partition
+        self.param2group = self.get_group_index()
+            
+    def get_grad_for_param(self, lp_param, group_idx, param_id):
+        return self.optim.fp32_groups_gradient_dict[group_idx][param_id]
+
+
+class DeepSpeedZeroOptimizerStage1or2Mon(DeepSpeedZeroOptimizerMon):
+    def __init__(self, optim):
+        super().__init__(optim)
+        self.stage = '1or2'
+        self.bit16_groups = optim.bit16_groups
+        self.fp32_flat_groups = optim.single_partition_of_fp32_groups
+        self.param2group = self.get_group_index()
+        self.group_offset = {}
+        self.get_group_offset()
+
+    def get_grad_for_param(self, lp_param, group_idx, param_id):
+        if getattr(self.optim, "cpu_offload", False):
+            grads = self.optim.single_partition_of_fp32_groups[group_idx].grad
+            start, numel = self.get_position(lp_param, group_idx)
+            grad = grads.narrow(0, start, numel)
+        else:
+            grad = self.optim.averaged_gradients[group_idx][param_id]
+        return grad
+
+    def get_group_offset(self):
+        for group_idx, group in enumerate(self.bit16_groups):
+            self.group_offset[group_idx] = -1
+            for lp_param in group:
+                if self.param_not_in_partition(lp_param, group_idx):
+                    self.group_offset[group_idx] = self.get_param_index(lp_param, group_idx)
+                else:
+                    break
+
+
+class DeepSpeedZeroOptimizerStage3Mon(DeepSpeedZeroOptimizerMon):
+    def __init__(self, optim):
+        super().__init__(optim)
+        self.stage = '3'
+        self.bit16_groups = optim.fp16_groups
+        self.fp32_flat_groups = optim.fp32_partitioned_groups_flat
+        self.param2group = self.get_group_index()
+
+    def param_not_in_partition(self, param, group_index):
+        """Each param partioned across all zero ranks"""
+        return False
+    
+    def get_position(self, lp_param, group_idx):
+        param_id = self.optim.get_param_id(lp_param)
+        return self.optim.grad_position[param_id][1:]
+    
+    def get_grad_for_param(self, lp_param, group_idx, param_id):
+        return self.optim.averaged_gradients[group_idx][param_id]
+
+
+class OptimizerMonFactory:
+    _optimizer_mon_map = {
+        "FP32Optimizer": OptimizerMon,
+        "Float16OptimizerWithFloat16Params": MixPrecisionOptimizerMon,
+        "DistributedOptimizer": MegatronDistributedOptimizerMon,
+        "SwapDistributedOptimizer": MegatronDistributedOptimizerMon,
+        "ChainedDistributedOptimizer": MegatronChainedDistributedOptimizerMon,
+        "ChainedSwapDistributedOptimizer": MegatronChainedDistributedOptimizerMon,
+        "ChainedFloat16OptimizerWithFloat16Params": MegatronChainedMixPrecisionOptimizerMon,
+        "BF16_Optimizer": DeepSpeedZeroOptimizerStage0Mon,
+        "DeepSpeedZeroOptimizer": DeepSpeedZeroOptimizerStage1or2Mon,
+        "DeepSpeedZeroOptimizer_Stage3": DeepSpeedZeroOptimizerStage3Mon,
+        "Adam": OptimizerMon
+    }
+
+    @staticmethod
+    def create_optimizer_mon(optimizer):
+        # auto replace opt_ty
+        optimizer_class = optimizer.__class__.__name__
+        if optimizer_class == "ChainedOptimizer":
+            optimizer_class = "Chained" + optimizer.chained_optimizers[0].__class__.__name__
+        logger.info(f'The optimizer type is {optimizer_class}')
+
+        optimizer_mon_class = OptimizerMonFactory._optimizer_mon_map.get(optimizer_class, OptimizerMon)
+        return optimizer_mon_class(optimizer)

msprobe/mindspore/monitor/utils.py

@@ -12,27 +12,36 @@
 # 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 os
+import re
+from datetime import datetime
 from mindspore import dtype as mstype, Tensor
 
 from msprobe.mindspore.monitor.features import FUNC_MAP
 from msprobe.core.common.const import MonitorConst
 from msprobe.core.common.utils import is_int
 from msprobe.core.common.log import logger
+from msprobe.core.common.file_utils import check_file_or_directory_path
 
 
-def get_single_metrics(op_list, tag, tensor, output=None):
+def get_single_metrics(op_list, tag, tensor, eps=1e-8, output=None):
     if output is None:
         output = {}
     if tag not in output:
         output[tag] = {}
     for op in op_list:
         func = FUNC_MAP.get(op)
-        statistic = func(tensor)
+        if op == "zeros":
+            statistic = func(tensor, eps)
+        else:
+            statistic = func(tensor)
         if hasattr(statistic, "dtype") and statistic.dtype == mstype.bfloat16:
             statistic = float(statistic)
             statistic = Tensor(statistic)
-        output[tag][op] = statistic.astype(mstype.float32)
+        if isinstance(statistic, Tensor):
+            output[tag][op] = statistic.astype(mstype.float32)
+        else:
+            output[tag][op] = statistic
 
 
 def get_metrics(op_list, tag2tensor, eps, output=None):
@@ -41,7 +50,7 @@ def get_metrics(op_list, tag2tensor, eps, output=None):
     for tag, tensor in tag2tensor.items():
         if tag not in output:
             output[tag] = {}
-        get_single_metrics(op_list, tag, tensor, output)
+        get_single_metrics(op_list, tag, tensor, eps, output)
     return output
 
 
@@ -88,6 +97,11 @@ def validate_ops(ops):
         default_op = MonitorConst.OP_LIST[0]
         valid_ops.append(default_op)
         logger.info(f"There is no valid ops, default op {default_op} is used")
+    # 增加默认shape和dtype参数
+    if "shape" not in valid_ops:
+        valid_ops.append("shape")
+    if "dtype" not in valid_ops:
+        valid_ops.append("dtype")
     return valid_ops
 
 
@@ -95,8 +109,8 @@ def validate_ranks(ranks):
     if not isinstance(ranks, list):
         raise TypeError("module_ranks should be a list")
     for rank in ranks:
-        if not isinstance(rank, str):
-            raise TypeError(f"element in module_ranks should be a str, get {type(rank)}")
+        if not isinstance(rank, int):
+            raise TypeError(f"element in module_ranks should be a int, get {type(rank)}")
 
 
 def validate_targets(targets):
@@ -168,7 +182,7 @@ def validate_alert(alert):
             args = rule.get("args")
             if args and isinstance(args, dict):
                 threshold = args.get("threshold")
-                if not isinstance(threshold, float) or threshold < 0:
+                if not isinstance(threshold, (float, int)) or threshold < 0:
                     raise TypeError('threshold must be float and not less than 0')
     dump = alert.get('dump')
     if dump and not isinstance(dump, bool):
@@ -209,6 +223,18 @@ def validate_collect_times(collect_times):
         raise ValueError("collect_times must greater than 1")
 
 
+def validate_dynamic_on(dynamic_on):
+    if not isinstance(dynamic_on, bool):
+        raise TypeError('dynamic_on should be a bool')
+
+
+def validate_monitor_mbs_grad(monitor_mbs_grad):
+    if not isinstance(monitor_mbs_grad, bool):
+        logger.warning(f'monitor_mbs_grad should be a bool, actual value is {monitor_mbs_grad}.')
+        return False
+    return monitor_mbs_grad
+
+
 def validate_config(config):
     config['ops'] = validate_ops(config.get('ops', []))
 
@@ -255,9 +281,14 @@ def validate_config(config):
     step_interval = config.get('step_interval', 1)
     validate_step_interval(step_interval)
 
-    collect_times = config.get('collect_times', 1e8)
+    collect_times = config.get('collect_times', int(1e8))
     validate_collect_times(collect_times)
 
+    config["monitor_mbs_grad"] = validate_monitor_mbs_grad(config.get('monitor_mbs_grad', False))
+
+    dynamic_on = config.get('dynamic_on', False)
+    validate_dynamic_on(dynamic_on)
+
     if not targets:
         if xy_distribution:
             config["all_xy"] = True
@@ -265,3 +296,34 @@ def validate_config(config):
         config["is_select"] = False
     else:
         config["is_select"] = True
+
+
+def time_str2time_digit(time_str):
+    time_format = '%b%d_%H-%M-%S'
+    try:
+        time_digit = datetime.strptime(time_str, time_format)
+    except Exception as e:
+        raise RuntimeError(f"illegal timestamp: {time_str}, timestamp should be prefix \
+                           of existing output dirpath, like 'Dec03_21-34-40'.") from e
+    return time_digit
+
+
+def get_target_output_dir(monitor_path, time_start, time_end):
+    check_file_or_directory_path(monitor_path, isdir=True)
+    time_start = time_str2time_digit(time_start) if time_start is not None else time_start
+    time_end = time_str2time_digit(time_end) if time_end is not None else time_end
+    if time_start and time_end and time_start > time_end:
+        raise ValueError(f"time_start({time_start}) greater than time_end({time_end})")
+    result = {}
+    for dirname in os.listdir(monitor_path):
+        match = re.match(MonitorConst.OUTPUT_DIR_PATTERN, dirname)
+        if not match:
+            continue
+        time_tag = match.group(1)
+        rank = match.group(2)
+        target_time = time_str2time_digit(time_tag)
+        start_ok = time_start is None or target_time >= time_start
+        end_ok = time_end is None or target_time <= time_end
+        if start_ok and end_ok:
+            result[rank] = os.path.join(monitor_path, dirname)
+    return result

msprobe/mindspore/ms_config.py

@@ -29,6 +29,7 @@ class TensorConfig(BaseConfig):
         self.check_mode = None
         self.file_format = json_config.get("file_format")
         self.check_config()
+        self._check_summary_mode()
         self._check_config()
 
     def _check_config(self):
@@ -42,12 +43,23 @@ class StatisticsConfig(BaseConfig):
         self.file_format = None
         self.check_mode = None
         self.check_config()
-        self._check_config()
+        self._check_summary_mode()
 
-    def _check_config(self):
-        single_opt = ["statistics", "md5"]
+        self.tensor_list = json_config.get("tensor_list", [])
+        self._check_str_list_config(self.tensor_list, "tensor_list")
+        self.stat_cal_mode = json_config.get("device", "host")
+        self.device_stat_precision_mode = json_config.get("precision", "high")
+        self._check_stat_params()
+
+    def _check_stat_params(self):
+        if self.stat_cal_mode not in ["device", "host"]:
+            raise Exception("Config param [device] is invalid, expected from [\"device\", \"host\"]")
+        if self.device_stat_precision_mode not in ["high", "low"]:
+            raise Exception("Config param [precision] is invalid, expected from [\"high\", \"low\"]")
+
+    def _check_summary_mode(self):
         muti_opt = ["md5", "max", "min", "mean", "l2norm"]
-        if isinstance(self.summary_mode, str) and self.summary_mode not in single_opt:
+        if isinstance(self.summary_mode, str) and self.summary_mode not in Const.SUMMARY_MODE:
             raise Exception("summary_mode is invalid")
         if isinstance(self.summary_mode, list) and not all(opt in muti_opt for opt in self.summary_mode):
             raise Exception("summary_mode is invalid")
@@ -132,14 +144,3 @@ def parse_task_config(task, json_config):
     if task not in TaskDict:
         raise Exception("task is invalid.")
     return TaskDict.get(task)(task_map)
-
-
-def parse_json_config(json_file_path):
-    if not json_file_path:
-        raise Exception("json file path is None")
-    json_config = load_json(json_file_path)
-    common_config = parse_common_config(json_config)
-    if not common_config.task:
-        common_config.task = Const.STATISTICS
-    task_config = parse_task_config(common_config.task, json_config)
-    return common_config, task_config

msprobe/mindspore/overflow_check/overflow_check_tool_factory.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,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from msprobe.core.common.log import logger
 from msprobe.mindspore.common.const import Const
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
 from msprobe.mindspore.overflow_check.kernel_graph_overflow_check import KernelGraphOverflowCheck
@@ -44,6 +45,7 @@ class OverflowCheckToolFactory:
             raise Exception("Valid level is needed.")
         tool = tool.get(config.execution_mode)
         if not tool:
-            raise Exception(f"Overflow check is not supported in {config.execution_mode} mode "
-                            f"when level is {config.level}.")
+            logger.error(f"Overflow check is not supported in {config.execution_mode} mode "
+                         f"when level is {config.level}.")
+            raise ValueError
         return tool(config)

msprobe/mindspore/task_handler_factory.py

@@ -29,11 +29,14 @@ class TaskHandlerFactory:
     }
 
     @staticmethod
-    def create(config: DebuggerConfig):
+    def create(config: DebuggerConfig, model=None):
         task = TaskHandlerFactory.tasks.get(config.task)
         if not task:
             raise Exception("Valid task is needed.")
-        handler = task.create(config)
+        if task == DumpToolFactory:
+            handler = task.create(config, model)
+        else:
+            handler = task.create(config)
         if not handler:
             raise Exception("Can not find task handler")
         return handler

msprobe/msprobe.py

@@ -22,6 +22,8 @@ from msprobe.core.common.log import logger
 from msprobe.core.compare.utils import _compare_parser
 from msprobe.core.compare.compare_cli import compare_cli
 from msprobe.core.compare.merge_result.merge_result_cli import _merge_result_parser, merge_result_cli
+from msprobe.core.config_check.config_check_cli import _config_checking_parser, \
+            _run_config_checking_command
 
 
 def is_module_available(module_name):
@@ -51,6 +53,9 @@ def main():
     graph_service_cmd_parser = subparsers.add_parser('graph')
     op_generate_cmd_parser = subparsers.add_parser('op_generate')
     merge_result_parser = subparsers.add_parser('merge_result')
+    config_checking_parser = subparsers.add_parser('config_check')
+    nan_analyze_parser = subparsers.add_parser('nan_analyze')
+    _config_checking_parser(config_checking_parser)
     _compare_parser(compare_cmd_parser)
     _merge_result_parser(merge_result_parser)
 
@@ -71,6 +76,7 @@ def main():
         from msprobe.visualization.graph_service import _pt_graph_service_parser, _pt_graph_service_command
         from msprobe.pytorch.api_accuracy_checker.generate_op_script.op_generator import _op_generator_parser, \
             _run_operator_generate_commond
+        from msprobe.nan_analyze.analyzer import _nan_analyze_parser, _run_nan_analyze
 
         _run_ut_parser(run_ut_cmd_parser)
         _run_ut_parser(multi_run_ut_cmd_parser)
@@ -80,6 +86,7 @@ def main():
         _run_overflow_check_parser(run_overflow_check_cmd_parser)
         _pt_graph_service_parser(graph_service_cmd_parser)
         _op_generator_parser(op_generate_cmd_parser)
+        _nan_analyze_parser(nan_analyze_parser)
     elif framework_args.framework == Const.MS_FRAMEWORK:
         from msprobe.mindspore.api_accuracy_checker.cmd_parser import add_api_accuracy_checker_argument
         from msprobe.visualization.graph_service import _ms_graph_service_parser, _ms_graph_service_command
@@ -91,6 +98,10 @@ def main():
 
         _ms_graph_service_parser(graph_service_cmd_parser)
 
+        from msprobe.mindspore.api_accuracy_checker.generate_op_script.op_generator import _op_generator_parser, \
+            _run_operator_generate_commond
+        _op_generator_parser(op_generate_cmd_parser)
+
     args = parser.parse_args(sys.argv[1:])
     if sys.argv[2] == Const.PT_FRAMEWORK:
         if not is_torch_available:
@@ -118,6 +129,10 @@ def main():
             compare_cli(args)
         elif sys.argv[3] == "merge_result":
             merge_result_cli(args)
+        elif sys.argv[3] == "config_check":
+            _run_config_checking_command(args)
+        elif sys.argv[3] == "nan_analyze":
+            _run_nan_analyze(args)
     else:
         if not is_module_available(Const.MS_FRAMEWORK):
             logger.error("MindSpore does not exist, please install MindSpore library")
@@ -134,9 +149,13 @@ def main():
             mul_api_checker_main(args)
         elif sys.argv[3] == "graph":
             _ms_graph_service_command(args)
+        elif sys.argv[3] == 'op_generate':
+            _run_operator_generate_commond(args)
         elif sys.argv[3] == "code_mapping":
             from msprobe.mindspore.code_mapping.main import code_mapping_main
             code_mapping_main(args)
+        elif sys.argv[3] == "config_check":
+            _run_config_checking_command(args)
 
 
 if __name__ == "__main__":

msprobe/nan_analyze/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 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.