mindstudio-probe 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

msprobe/pytorch/monitor/module_metric.py

@@ -16,8 +16,9 @@ import re
 
 import torch
 
+from msprobe.pytorch.common.utils import is_float8_tensor
 from msprobe.pytorch.monitor.features import get_max, get_min, get_zeros, get_nans, get_norm, get_mean
-from msprobe.pytorch.monitor.utils import NAN_TENSOR_ON_DEVICE
+from msprobe.pytorch.monitor.utils import get_nan_tensor
 
 
 def get_summary_writer_tag_name(module_or_param_name: str, tag: str, rank):
@@ -147,13 +148,13 @@ def get_metrics(ops, tag2tensor, eps, out_dict=None):
     """
     :param ops: ["op1", "op2"]
     :param tag2tensor: {
-    '0:fc_0/input': torch.randn([3, 4]),
-    '0:fc_0/output': torch.randn([3, 3])
+    '0:fc.input:0/actv': torch.randn([3, 4]),
+    '0:fc.output:0/actv': torch.randn([3, 3])
     }
     :param eps: float 1e-8
     :param out_dict:{
-    '0:fc_0/input': {"op1": op1(torch.randn([3, 4])), "op2": op2(torch.randn([3, 4]))}
-    '0:fc_0/output': {"op1": op1(torch.randn([3, 3])), "op2": op2(torch.randn([3, 3]))}
+    '0:fc.input:0/actv': {"op1": op1(torch.randn([3, 4])), "op2": op2(torch.randn([3, 4]))}
+    '0:fc.output:0/actv': {"op1": op1(torch.randn([3, 3])), "op2": op2(torch.randn([3, 3]))}
     }
     :return: out_dict
     """
@@ -164,8 +165,10 @@ def get_metrics(ops, tag2tensor, eps, out_dict=None):
             out_dict[tag] = {}
         if not torch.is_tensor(tensor):
             # Non-tensor in/output filled with nan.
-            out_dict[tag].update({metric_name: NAN_TENSOR_ON_DEVICE for metric_name in ops})
+            out_dict[tag].update({metric_name: get_nan_tensor() for metric_name in ops})
             continue
+        if is_float8_tensor(tensor):
+            tensor = tensor.float()
         for metric_name in ops:
             fun_metric = config_metric_registry.get(metric_name)
             out_dict[tag][metric_name] = fun_metric.get_metric(tensor, eps)

msprobe/pytorch/monitor/optimizer_collect.py

@@ -23,16 +23,10 @@ from msprobe.pytorch.monitor.utils import MVResult, MVGradResult
 
 
 class OptimizerMon(object):
-    wrapped_optimizer = None
-
     def __init__(self) -> None:
         self.fp16_to_fp32_param = {}
         self.is_stage3 = False
 
-    @classmethod
-    def set_wrapped_optimizer(cls, wrapped_optimizer):
-        cls.wrapped_optimizer = wrapped_optimizer
-
     def fetch_mv(self, monitor, torch_opt, params2name):
         pass
 
@@ -82,7 +76,6 @@ class OptimizerMon(object):
         ratio_dict = defaultdict()
         param2name = defaultdict()
         fp32_partitioned_groups_flat_grad = defaultdict()
-        mix_prec_opt = OptimizerMon.wrapped_optimizer
         partition_id = dist.get_rank()
 
         def get_flatten_grad(self, optimizer, group_idx):
@@ -101,7 +94,7 @@ class OptimizerMon(object):
                 return fp32_partitioned_groups_flat[group_idx].grad
 
         for group_idx in range(len(fp32_partitioned_groups_flat)):
-            fp32_partitioned_groups_flat_grad[group_idx] = get_flatten_grad(self, mix_prec_opt, group_idx)
+            fp32_partitioned_groups_flat_grad[group_idx] = get_flatten_grad(self, torch_opt, group_idx)
 
         for name in params2name.values():
             start_idx, end_idx, group_idx, group_with_rank = name2indices[name]
@@ -110,9 +103,9 @@ class OptimizerMon(object):
             fp32_param = fp32_partitioned_groups_flat[group_idx][start_idx: end_idx]
             fp32_param.grad = fp32_partitioned_groups_flat_grad[group_idx][start_idx: end_idx]
             param2name[fp32_param] = name
-            if not mix_prec_opt.state:
+            if not torch_opt.state:
                 continue
-            state_param = list(mix_prec_opt.state.values())[group_idx]
+            state_param = list(torch_opt.state.values())[group_idx]
             exp_avg = state_param.get("exp_avg", None)
             exp_avg_sq = state_param.get("exp_avg_sq", None)
             if exp_avg is None or exp_avg_sq is None:
@@ -150,36 +143,33 @@ class MixPrecisionOptimizerMon(OptimizerMon):
     混合精度训练通过适当降低某些计算的精度来加速训练过程并减少内存消耗。
     """
 
-    def map_fp16_tp_fp32_param(self, mix_prec_opt):
-        for fp16_group, fp32_group in zip(mix_prec_opt.float16_groups, mix_prec_opt.fp32_from_float16_groups):
+    def map_fp16_tp_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):
                 self.fp16_to_fp32_param[fp16_param] = fp32_param
 
     def fetch_mv(self, monitor, torch_opt, params2name):
-        mix_prec_opt = self.wrapped_optimizer
-
-        if not self.fp16_to_fp32_param and mix_prec_opt is not None:
-            self.map_fp16_tp_fp32_param(mix_prec_opt)
+        if not self.fp16_to_fp32_param and torch_opt is not None:
+            self.map_fp16_tp_fp32_param(torch_opt)
 
         return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
 
 
 class MegatronDistributedOptimizerMon(OptimizerMon):
-    def map_fp16_tp_fp32_param(self, mix_prec_opt):
-        if not (hasattr(mix_prec_opt, "model_float16_groups") and
-                hasattr(mix_prec_opt, "shard_fp32_from_float16_groups")):
+    def map_fp16_tp_fp32_param(self, torch_opt):
+        if not (hasattr(torch_opt, "model_float16_groups") and
+                hasattr(torch_opt, "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(mix_prec_opt.model_float16_groups,
-                                                mix_prec_opt.shard_fp32_from_float16_groups):
+        for fp16_group, shard_fp32_group in zip(torch_opt.model_float16_groups,
+                                                torch_opt.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
 
     def fetch_mv(self, monitor, torch_opt, params2name):
-        mix_prec_opt = self.wrapped_optimizer
-        if not self.fp16_to_fp32_param and mix_prec_opt is not None:
-            self.map_fp16_tp_fp32_param(mix_prec_opt)
+        if not self.fp16_to_fp32_param and torch_opt is not None:
+            self.map_fp16_tp_fp32_param(torch_opt)
 
         return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
 
@@ -191,43 +181,89 @@ class MegatronFP32OptimizerMon(OptimizerMon):
 
 class MegatronChainedDistributedOptimizerMon(MegatronDistributedOptimizerMon):
     def fetch_mv(self, monitor, torch_opt, params2name):
-        mix_prec_opt = self.wrapped_optimizer
-
-        if not self.fp16_to_fp32_param and mix_prec_opt is not None:
-            for opt in mix_prec_opt.chained_optimizers:
+        if not self.fp16_to_fp32_param and torch_opt is not None:
+            for opt in torch_opt.chained_optimizers:
                 self.map_fp16_tp_fp32_param(opt)
 
-        if not isinstance(torch_opt, torch.optim.Optimizer):
+        if not isinstance(torch_opt, torch.optim.Optimizer) and not hasattr(torch_opt, 'state'):
             torch_opt.state = {}
-            for opt in mix_prec_opt.chained_optimizers:
+            for opt in torch_opt.chained_optimizers:
                 torch_opt.state.update(opt.optimizer.state)
         return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
 
 
 class MegatronChainedMixPrecisionOptimizerMon(MixPrecisionOptimizerMon):
     def fetch_mv(self, monitor, torch_opt, params2name):
-        mix_prec_opt = self.wrapped_optimizer
-
-        if not self.fp16_to_fp32_param and mix_prec_opt is not None:
-            for opt in mix_prec_opt.chained_optimizers:
+        if not self.fp16_to_fp32_param and torch_opt is not None:
+            for opt in torch_opt.chained_optimizers:
                 self.map_fp16_tp_fp32_param(opt)
 
-        if not isinstance(torch_opt, torch.optim.Optimizer):
+        if not isinstance(torch_opt, torch.optim.Optimizer) and not hasattr(torch_opt, 'state'):
             torch_opt.state = {}
-            for opt in mix_prec_opt.chained_optimizers:
+            for opt in torch_opt.chained_optimizers:
                 torch_opt.state.update(opt.optimizer.state)
         return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
 
 
 class DeepSpeedZeroOptimizerStage0Mon(OptimizerMon):
-    def fetch_mv(self, monitor, torch_opt, params2name):
-        return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
+    def get_group_index(self, torch_opt):
+        bit16_groups = torch_opt.bf16_groups
+        param2group = defaultdict()
+        for group_idx, bit16_group in enumerate(bit16_groups):
+            for param in bit16_group:
+                param2group[param] = group_idx
+        return param2group
+
+    def fetch_mv(self, monitor, torch_opt, params2name, name2indices=None):
+        param2group = self.get_group_index(torch_opt)
+        exp_avg_dict = defaultdict(float)
+        exp_avg_sq_dict = defaultdict(float)
+        update_dict = defaultdict()
+        ratio_dict = defaultdict()
+
+        param_slice_mappings = torch_opt.state_dict()['param_slice_mappings']
+        for param, name in params2name.items():
+            group_idx = param2group[param]
+            state = torch_opt.optimizer.state[torch_opt.fp32_groups_flat_partition[group_idx]]
+            if state.get('exp_avg', None) is None:
+                logger.warning(f"optimizer state is None. Something is wrong if this is not the first step")
+                break
+            param_slice_mapping = param_slice_mappings[group_idx]
+            hp_address = param_slice_mapping.get(torch_opt.param_names[param])
+            if hp_address is None:
+                continue
+            start = hp_address.start
+            numel = hp_address.numel
 
+            if monitor.mv_distribution:
+                exp_avg_dict[name] = state['exp_avg'].narrow(0, start, numel)
+                exp_avg_sq_dict[name] = state['exp_avg_sq'].narrow(0, start, numel)
+            if monitor.mg_direction:
+                exp_avg_dict[name] = state['exp'].narrow(0, start, numel)
+            if monitor.ur_distribution:
+                if len(torch_opt.param_groups) > 1:
+                    logger.info(f"the length of torch_opt.param_groups is {len(torch_opt.param_groups)}.")
+                if 'step' in state:
+                    step = state['step']  # Optimizer from pytorch or FusedAdam from apex(used by megatron)
+                elif 'step' in torch_opt.param_groups[0]:
+                    step = torch_opt.param_groups[0]['step']  # AdamW from mindspeed
+                else:
+                    logger.warning(f"step of {name} is None, maybe something wrong happened.")
+                    continue
+                exp_avg = state['exp_avg'].narrow(0, start, numel)
+                exp_avg_sq = state['exp_avg_sq'].narrow(0, start, numel)
+                exp_avg_hat = exp_avg / (1 - torch_opt.defaults['betas'][0] ** step)
+                exp_avg_sq_hat = exp_avg_sq / (1 - torch_opt.defaults['betas'][1] ** step)
+                update_dict[name] = exp_avg_hat / (torch.sqrt(exp_avg_sq_hat) + torch_opt.defaults['eps'])
+                ratio_dict[name] = exp_avg_hat / torch.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 MVResult(exp_avg=exp_avg_dict, exp_avg_sq=exp_avg_sq_dict, update=update_dict, ratio=ratio_dict)
+    
 
 class DeepSpeedZeroOptimizerStage3Mon(OptimizerMon):
-    def get_param_index(self, params2name, name2index):
-        mix_prec_opt = OptimizerMon.wrapped_optimizer
-        fp16_groups = mix_prec_opt.fp16_partitioned_groups
+    def get_param_index(self, params2name, name2index, torch_opt):
+        fp16_groups = torch_opt.fp16_partitioned_groups
         name2indices = defaultdict()
         index_length = defaultdict()
         index = 0
@@ -246,13 +282,11 @@ class DeepSpeedZeroOptimizerStage3Mon(OptimizerMon):
 
     def fetch_mv(self, monitor, torch_opt, params2name, name2indices=None):
         self.is_stage3 = True
-        mix_prec_opt = OptimizerMon.wrapped_optimizer
-        fp32_partitioned_groups_flat = mix_prec_opt.fp32_partitioned_groups_flat
+        fp32_partitioned_groups_flat = torch_opt.fp32_partitioned_groups_flat
         return self._fetch_mv_grad_in_adam(monitor, torch_opt, params2name, name2indices, fp32_partitioned_groups_flat)
 
 
 class DeepSpeedZeroOptimizerStage1or2Mon(OptimizerMon):
-
     @staticmethod
     def get_group_index(fp32_length, world_size, index):
         for i in range(len(fp32_length) - 1):
@@ -265,12 +299,11 @@ class DeepSpeedZeroOptimizerStage1or2Mon(OptimizerMon):
                 return sub_interval_start, min(sub_index, world_size - 1)
         return fp32_length[-1], 0
 
-    def get_param_index(self, params2name, name2index):
-        mix_prec_opt = OptimizerMon.wrapped_optimizer
-        padding = mix_prec_opt.groups_padding
+    def get_param_index(self, params2name, name2index, torch_opt):
+        padding = torch_opt.groups_padding
         world_size = dist.get_world_size()
         fp32_length = [0]
-        for fp32_group_index, single_partition_of_fp32_group in enumerate(mix_prec_opt.single_partition_of_fp32_groups):
+        for fp32_group_index, single_partition_of_fp32_group in enumerate(torch_opt.single_partition_of_fp32_groups):
             fp32_length.append(len(single_partition_of_fp32_group) * world_size + fp32_length[fp32_group_index])
 
         bf16_groups = []
@@ -278,7 +311,7 @@ class DeepSpeedZeroOptimizerStage1or2Mon(OptimizerMon):
         index_length = defaultdict()
         index = 0
         idx = 0
-        for group_idx, bf16_group in enumerate(mix_prec_opt.bit16_groups):
+        for group_idx, bf16_group in enumerate(torch_opt.bit16_groups):
             bf16_groups.extend(bf16_group)
             for param in bf16_group:
                 param_length = len(param.flatten())
@@ -286,7 +319,7 @@ class DeepSpeedZeroOptimizerStage1or2Mon(OptimizerMon):
                 index_length[idx] = (index, index + param_length, group_idx, group_index, group_with_rank)
                 index += param_length
                 idx += 1
-        group_length = len(bf16_groups) / len(mix_prec_opt.bit16_groups)
+        group_length = len(bf16_groups) / len(torch_opt.bit16_groups)
         for _, name in params2name.items():
             name_index = name2index[name]
             start_idx, end_idx, group_idx, group_index, group_with_rank = index_length[name_index]
@@ -300,8 +333,7 @@ class DeepSpeedZeroOptimizerStage1or2Mon(OptimizerMon):
         return name2indices
 
     def fetch_mv(self, monitor, torch_opt, params2name, name2indices=None):
-        mix_prec_opt = OptimizerMon.wrapped_optimizer
-        fp32_partitioned_groups_flat = mix_prec_opt.single_partition_of_fp32_groups
+        fp32_partitioned_groups_flat = torch_opt.single_partition_of_fp32_groups
         return self._fetch_mv_grad_in_adam(monitor, torch_opt, params2name, name2indices, fp32_partitioned_groups_flat)
 
 
@@ -312,22 +344,23 @@ class DummyOptimizerMon(OptimizerMon):
 
 class OptimizerMonFactory:
     _optimizer_mon_map = {
-        "Megatron_Float16OptimizerWithFloat16Params": MixPrecisionOptimizerMon,
-        "Megatron_DistributedOptimizer": MegatronDistributedOptimizerMon,
-        "Megatron_ChainedDistributedOptimizer": MegatronChainedDistributedOptimizerMon,
-        "Megatron_ChainedFloat16OptimizerWithFloat16Params": MegatronChainedMixPrecisionOptimizerMon,
-        "Megatron_FP32Optimizer": MegatronFP32OptimizerMon,
-        "DeepSpeedZeroOptimizer_Stage0": DeepSpeedZeroOptimizerStage0Mon,
-        "DeepSpeedZeroOptimizer_Stage1_or_2": DeepSpeedZeroOptimizerStage1or2Mon,
+        "FP32Optimizer": MegatronFP32OptimizerMon,
+        "Float16OptimizerWithFloat16Params": MixPrecisionOptimizerMon,
+        "DistributedOptimizer": MegatronDistributedOptimizerMon,
+        "ChainedDistributedOptimizer": MegatronChainedDistributedOptimizerMon,
+        "ChainedFloat16OptimizerWithFloat16Params": MegatronChainedMixPrecisionOptimizerMon,
+        "BF16_Optimizer": DeepSpeedZeroOptimizerStage0Mon,
+        "DeepSpeedZeroOptimizer": DeepSpeedZeroOptimizerStage1or2Mon,
         "DeepSpeedZeroOptimizer_Stage3": DeepSpeedZeroOptimizerStage3Mon,
-        "unknown": DummyOptimizerMon
+        "Adam": DummyOptimizerMon
     }
 
     @staticmethod
-    def create_optimizer_mon(opt_ty: str):
-        if not opt_ty:
-            return DummyOptimizerMon()
-        optimizer_mon_class = OptimizerMonFactory._optimizer_mon_map.get(opt_ty)
-        if not optimizer_mon_class:
-            raise Exception("opt_ty should be one of: " + ", ".join(OptimizerMonFactory._optimizer_mon_map.keys()))
-        return optimizer_mon_class()
+    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__
+
+        optimizer_mon_class = OptimizerMonFactory._optimizer_mon_map.get(optimizer_class, DummyOptimizerMon)
+        return optimizer_mon_class(), optimizer_class

msprobe/pytorch/monitor/unittest/test_monitor.py

@@ -92,7 +92,7 @@ def valid_reduce(reduced, unreduced, tp_size, dp_size, sequence_parallel):
     if errors:
         logger.info(errors)
     else:
-        logger.info(f'grad mean is in consist between unreduced grad and reduced grad monitord.')
+        logger.info(f'grad mean is in consist between unreduced grad and reduced grad monitored.')
 
 
 def assert_equal(a, b):

msprobe/pytorch/monitor/utils.py

@@ -25,7 +25,7 @@ import torch
 from msprobe.core.common.const import MonitorConst, Const
 from msprobe.pytorch.common.log import logger
 from msprobe.core.common.utils import is_int
-from msprobe.core.common.file_utils import check_file_or_directory_path
+from msprobe.core.common.file_utils import check_file_or_directory_path, recursive_chmod
 
 
 device = "cpu"
@@ -36,7 +36,7 @@ except ImportError:
     if torch.cuda.is_available():
         device = "cuda"
 
-NAN_TENSOR_ON_DEVICE = torch.tensor(torch.nan, device=device)
+NAN_TENSOR_ON_DEVICE = None
 FILE_MAX_SIZE = 10 * 1024 * 1024 * 1024
 FILE_NAME_MAX_LENGTH = 255
 DIRECTORY_MAX_LENGTH = 4096
@@ -57,6 +57,13 @@ def get_output_base_dir():
     return os.getenv(MonitorConst.MONITOR_OUTPUT_DIR, MonitorConst.DEFAULT_MONITOR_OUTPUT_DIR)
 
 
+def get_nan_tensor():
+    global NAN_TENSOR_ON_DEVICE
+    if not NAN_TENSOR_ON_DEVICE:
+        NAN_TENSOR_ON_DEVICE = torch.tensor(torch.nan, device=device)
+    return NAN_TENSOR_ON_DEVICE
+
+
 def filter_special_chars(func):
     @wraps(func)
     def func_level(msg):
@@ -82,48 +89,6 @@ def get_param_struct(param):
     return res
 
 
-def is_recomputation():
-    """Check if the current operation is in the re-computation phase.
-
-    This function inspects the current call stack to indicate whether the current operation is in the
-    re-computation phase. We use a blacklist mechanism, now supported megatron and mindspeed framework.
-    megatron: The 'backward' function is called by the 'torch/autograd/function.py' file.
-    mindspeed: The 'checkpoint_function_backward' function is called by the 'torch/autograd/function.py'
-    file or the custom module(use CheckpointWithoutOutput) with the 'backward' function is executed within the
-    'torch/_tensor.py' file.
-
-    Returns:
-        bool: True if in the re-computation phase, False otherwise.
-    """
-    backward_function_indices = []
-    call_stack = inspect.stack()
-
-    # Identify the function 'backward' is being executed within the 'torch/_tensor.py' file.
-    for frame_info in call_stack:
-        if frame_info.function == Const.BACKWARD and frame_info.filename.endswith('torch/_tensor.py'):
-            del call_stack
-            return True
-
-    # Identify indices in the call stack where the specific function is being executed
-    for idx, frame_info in enumerate(call_stack):
-        if frame_info.function == Const.BACKWARD or frame_info.function == 'checkpoint_function_backward':
-            backward_function_indices.append(idx)
-
-    # Check if the execution is within 'torch/autograd/function.py' file
-    for idx in backward_function_indices:
-        # The Megatron and MindSpeed L0&L1 scenes
-        if idx + 1 < len(call_stack) and call_stack[idx + 1].filename.endswith('torch/autograd/function.py'):
-            del call_stack
-            return True
-        # The latest MindSpeed L2 and ModelLink scenes
-        if idx + 2 < len(call_stack) and call_stack[idx + 2].filename.endswith('torch/autograd/function.py'):
-            del call_stack
-            return True
-
-    del call_stack
-    return False
-
-
 def validate_ops(ops):
     if not isinstance(ops, list):
         raise TypeError("ops should be a list")
@@ -140,6 +105,15 @@ def validate_ops(ops):
     return valid_ops
 
 
+def validate_ndigits(ndigits):
+    if not ndigits:
+        return
+    if not is_int(ndigits) or ndigits <= 0:
+        raise ValueError(f"ndigits({ndigits}) is not a positive integer, current is: {ndigits}.")
+    if ndigits > MonitorConst.MAX_NDIGITS:
+        raise ValueError(f"The maximum supported ndigits is {MonitorConst.MAX_NDIGITS}, current value: {ndigits}.")
+
+
 def validate_ranks(ranks):
     if not isinstance(ranks, list):
         raise TypeError("module_ranks should be a list")
@@ -241,9 +215,17 @@ def validate_step_count_per_record(step_count_per_record):
         raise ValueError("step_count_per_record must smaller than 1e6")
 
 
+def validate_dynamic_on(dynamic_on):
+    if not isinstance(dynamic_on, bool):
+        raise TypeError('dynamic_on should be a bool')
+
+
 def validate_config(config):
     config['ops'] = validate_ops(config.get('ops', []))
 
+    ndigits = config.get('ndigits')
+    validate_ndigits(ndigits)
+
     eps = config.get('eps', 1e-8)
     if not isinstance(eps, float):
         raise TypeError("eps should be a float")
@@ -281,9 +263,20 @@ def validate_config(config):
     step_count_per_record = config.get('step_count_per_record', 1)
     validate_step_count_per_record(step_count_per_record)
 
+    config["start_step"] = validate_int_arg(config.get("start_step"), "start_step",
+                                            MonitorConst.DEFAULT_START_STEP, MonitorConst.DEFAULT_START_STEP)
+    config["collect_times"] = validate_int_arg(config.get("collect_times"), "collect_times",
+                                               MonitorConst.DEFAULT_MIN_COLLECT_TIMES,
+                                               MonitorConst.DEFAULT_MAX_COLLECT_TIMES)
+    config["step_interval"] = validate_int_arg(config.get("step_interval"), "step_interval",
+                                               MonitorConst.DEFAULT_STEP_INTERVAL, MonitorConst.DEFAULT_STEP_INTERVAL)
+
     squash_name = config.get('squash_name', True)
     validate_squash_name(squash_name)
 
+    dynamic_on = config.get('dynamic_on', False)
+    validate_dynamic_on(dynamic_on)
+
     if not targets:
         if xy_distribution:
             config["all_xy"] = True
@@ -292,6 +285,8 @@ def validate_config(config):
 
 def time_str2time_digit(time_str):
     time_format = '%b%d_%H-%M-%S'
+    if not isinstance(time_str, str):
+        raise TypeError(f"time_str:{time_str} should be a str")
     try:
         time_digit = datetime.strptime(time_str, time_format)
     except Exception as e:
@@ -319,3 +314,40 @@ def get_target_output_dir(monitor_path, time_start, time_end):
         if start_ok and end_ok:
             result[rank] = os.path.join(monitor_path, dirname)
     return result
+
+
+def chmod_tensorboard_dir(path):
+    """
+        format配置为tensorboard时，需要补充文件权限设置
+    """
+    try:
+        recursive_chmod(path)
+    except Exception as e:
+        logger.warning(f"chmod tensorboard dir wrong because {e}, not updated, please check!!!")
+
+
+def validate_set_monitor(grad_acc_steps, start_iteration):
+    """
+    validate parameters of set_monitor.
+    """
+    grad_acc_steps = validate_int_arg(grad_acc_steps, "grad_acc_steps",
+                                      MonitorConst.DEFAULT_GRAD_ACC_STEPS, MonitorConst.DEFAULT_GRAD_ACC_STEPS)
+
+    start_iteration = validate_int_arg(start_iteration, "start_iteration",
+                                       MonitorConst.DEFAULT_START_ITERATION, MonitorConst.DEFAULT_START_ITERATION)
+    return grad_acc_steps, start_iteration
+
+
+def validate_int_arg(value, name, minimum, default_value):
+    """Validate int args, if any exception occurs, use the default value."""
+    if value is None:
+        return default_value
+    try:
+        if not is_int(value):
+            raise TypeError(f"{name} must be int")
+        if value < minimum:
+            raise ValueError(f"{name} must greater than {minimum}")
+    except Exception as e:
+        value = default_value
+        logger.warning(f"Validate {name} failed, {e}, replaced with default value {value}.")
+    return value

msprobe/pytorch/online_dispatch/compare.py

@@ -125,8 +125,6 @@ class Saver:
 
     def write_summary_csv(self, test_result):
         test_rows = []
-        if self.stack_info:
-            test_rows[0].append(self.COLUMN_STACK_INFO)
 
         check_op_str_pattern_valid(test_result.api_name)
         df_row = [test_result.api_name, test_result.is_fwd_success, test_result.is_bwd_success]

msprobe/pytorch/online_dispatch/dispatch.py

@@ -16,6 +16,7 @@
 import json
 import os
 import time
+import multiprocessing
 from multiprocessing import Pool
 
 import torch
@@ -52,6 +53,7 @@ class PtdbgDispatch(TorchDispatchMode):
             return
         if dump_path is None:
             logger.error("Please set dump_path when dump_mode is config!")
+            raise DispatchException("Please set dump_path when dump_mode is config!")
         check_file_or_directory_path(dump_path, True)
 
         self.device_id = torch_npu._C._npu_getDevice()
@@ -85,6 +87,11 @@ class PtdbgDispatch(TorchDispatchMode):
         self.get_ops(yaml_path)
 
         self.lock = None
+        max_process_num = max(int((multiprocessing.cpu_count() + 1) // Const.CPU_QUARTER), 1)
+        if process_num > max_process_num:
+            logger.error(f"process_num should be less than or equal to {max_process_num}, but got {process_num}!")
+            raise DispatchException(f'process_num should be less than or equal to {max_process_num}, '
+                                    f'but got {process_num}!')
         if process_num > 0:
             self.pool = Pool(process_num)
         if debug:
@@ -115,6 +122,8 @@ class PtdbgDispatch(TorchDispatchMode):
                 if len(json_line_data) == 0:
                     break
                 msg = json.loads(json_line_data)
+                if len(msg) < 2:
+                    raise ValueError("JSON data does not contain enough elements. Expected at least 2 elements.")
                 self.all_summary[msg[0]] = msg[1]
             fp_handle.close()
 

msprobe/pytorch/online_dispatch/dump_compare.py

@@ -19,6 +19,8 @@ import os
 from datetime import datetime, timezone
 
 import torch
+from msprobe.core.common.const import Const
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.core.common.file_utils import FileOpen, save_npy, save_json
 from msprobe.pytorch.common.log import logger
 
@@ -91,6 +93,7 @@ def support_basic_type(data):
     return False
 
 
+@recursion_depth_decorator("dump_data")
 def dump_data(data, prefix, dump_path):
     if isinstance(data, (tuple, list)) and data:
         for i, item in enumerate(data):

msprobe/pytorch/online_dispatch/utils.py

@@ -27,8 +27,10 @@ else:
     pta_cpu_device = torch.device("cpu")
 
 from msprobe.core.common.const import CompareConst
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.common.log import logger
 
+
 cpu_device = torch._C.device("cpu")
 COLOR_RED = '\033[31m'
 COLOR_GREEN = '\033[32m'
@@ -85,6 +87,7 @@ def get_callstack():
     return callstack
 
 
+@recursion_depth_decorator("data_to_cpu")
 def data_to_cpu(data, deep, data_cpu):
     global cpu_device
     list_cpu = []

msprobe/pytorch/parse_tool/lib/interactive_cli.py

@@ -45,12 +45,7 @@ class InteractiveCli(cmd.Cmd):
 
     @catch_exception
     def default(self, line=""):
-        self.util.execute_command(line)
-        return False
-
-    @catch_exception
-    def do_run(self, line=""):
-        self.util.execute_command(line)
+        self.stdout.write("Command invalid, Only support command start with cad/vc/dc/pk/cn/pt\n")
 
     @catch_exception
     def do_vc(self, line=""):

msprobe/pytorch/parse_tool/lib/utils.py

@@ -119,6 +119,7 @@ class Util:
 
     @staticmethod
     def deal_with_dir_or_file_inconsistency(output_path):
+        logger.warning(f"Trying to delete {output_path}")
         remove_path(output_path)
         raise ParseException("Inconsistent directory structure or file.")
 
@@ -264,7 +265,7 @@ class Util:
             match = re_pattern.match(name)
             if not match:
                 continue
-            if extern_pattern != '' and re_pattern.match(extern_pattern) and not re.match(extern_pattern, name):
+            if extern_pattern != '' and re_pattern.match(extern_pattern) and not name.startswith(extern_pattern):
                 continue
             file_list[name] = gen_info_func(name, match, file["root"])
         return file_list