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

msprobe/pytorch/monitor/module_metric.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");
@@ -12,16 +12,12 @@
 # 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 itertools
-import math
 import re
-import statistics
 
 import torch
 
-from msprobe.core.common.const import MonitorConst
-from msprobe.pytorch.monitor.features import square_sum, get_max, get_min, get_zeros, get_nans, get_norm, get_mean
-from msprobe.core.common.log import logger
+from msprobe.pytorch.monitor.features import get_max, get_min, get_zeros, get_nans, get_norm, get_mean
+from msprobe.pytorch.monitor.utils import get_nan_tensor
 
 
 def get_summary_writer_tag_name(module_or_param_name: str, tag: str, rank):
@@ -31,7 +27,9 @@ def get_summary_writer_tag_name(module_or_param_name: str, tag: str, rank):
         return f"{module_or_param_name}/rank{rank}/{tag}"
 
 
-def squash_param_name(param_name):
+def squash_param_name(param_name, enable=True):
+    if not enable:
+        return param_name
     name = ''
     for pattern in ['layers?\.(.*)', 'embeddings?\.(.*)', 'final.*', 'output.*', 'norm.*']:
         match = re.findall(pattern, param_name)
@@ -63,7 +61,7 @@ class TensorMetrics:
         self.metrics = {}  # tensor_tag --> []
         self.cur_idx = {}
 
-    def stat_insert(self, tensor, stat_ops, module_name, tensor_name, rank, eps=1e-8):
+    def stat_insert(self, tensor, stat_ops, module_name, tensor_name, rank):
         """get stats and insert into metrics dictionary"""
         prefix = get_summary_writer_tag_name(module_name, tensor_name, rank)
         for stat_op in stat_ops:
@@ -120,14 +118,14 @@ class NormMetric(Metric):
     @staticmethod
     def get_metric_value(tensor, eps):
         return get_norm(tensor)
-    
+
 
 @register_config_metric("zeros")
 class ZerosMetric(Metric):
     @staticmethod
     def get_metric_value(tensor, eps):
         return get_zeros(tensor, eps)
-    
+
 
 @register_config_metric("nans")
 class NaNsMetric(Metric):
@@ -146,48 +144,29 @@ class IdentMetric(Metric):
 
 
 def get_metrics(ops, tag2tensor, eps, out_dict=None):
+    """
+    :param ops: ["op1", "op2"]
+    :param tag2tensor: {
+    '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.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
+    """
     if out_dict is None:
         out_dict = {}
     for tag, tensor in tag2tensor.items():
         if tag not in out_dict:
             out_dict[tag] = {}
-        for metric_name in ops:    
+        if not torch.is_tensor(tensor):
+            # Non-tensor in/output filled with nan.
+            out_dict[tag].update({metric_name: get_nan_tensor() for metric_name in ops})
+            continue
+        for metric_name in ops:
             fun_metric = config_metric_registry.get(metric_name)
             out_dict[tag][metric_name] = fun_metric.get_metric(tensor, eps)
     return out_dict
-
-
-def write_metrics_base(ops, summary_writer, metric_value, step, prefix=''):
-    if not metric_value:
-        return
-    tensors = []
-    tags = list(itertools.product(metric_value.keys(), ops))
-    for op2tensor in metric_value.values():
-        tensors.extend(op2tensor.values())
-    with torch.no_grad():
-        metric_list = torch.stack(tensors).cpu()
-    for tag, metric in zip(tags, metric_list):
-        summary_writer.add_scalar(tag, metric, step)
-
-
-def write_metrics_csv(ops, summary_writer, metric_value, step, prefix=''):
-    write_metrics_base(ops, summary_writer, metric_value, step, prefix='')
-
-    if not summary_writer.header:
-        # 前向的norm用input.ops_和output.ops_，反向的用input_grad.ops_和output_grad.ops_
-        if prefix in {"actv", "actv_grad"}:
-            if prefix == "actv":
-                input_and_output = [MonitorConst.ACTV_IN, MonitorConst.ACTV_OUT]
-            else:
-                input_and_output = [MonitorConst.ACTVGRAD_IN, MonitorConst.ACTVGRAD_OUT]
-            ops_ = [MonitorConst.DOT.join(i[::-1]) for i in itertools.product(ops, input_and_output)]
-            summary_writer.header = ["module_name", "step", *ops_]
-        else:
-            summary_writer.header = ["param_name", "step", *ops]
-
-        for key in metric_value.keys():
-            if MonitorConst.VPP_SEP in key:
-                summary_writer.header.insert(0, 'vpp_stage')
-            break
-    summary_writer.write_csv(prefix, step)
-    summary_writer.header = []

msprobe/pytorch/monitor/module_spec_verifier.py

@@ -17,7 +17,7 @@ import re
 import abc
 import torch
 
-from msprobe.core.common.log import logger
+from msprobe.pytorch.common.log import logger
 
 # 用于存储所有validator实现类的注册表
 config_validator_registry = {}
@@ -79,6 +79,8 @@ class TupleValidator(ConfigValidator):
 
 def validate_config_spec(config_spec: str, actual_data, module_name: str, data_type: str):
     focused_col = None
+    if not config_spec or not isinstance(config_spec, str):
+        return focused_col
     for _, validator_cls in config_validator_registry.items():
         config_validator = validator_cls()
         pattern_match = config_validator.check_pattern_match(config_spec)

msprobe/pytorch/monitor/optimizer_collect.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,27 +13,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractmethod
 from collections import defaultdict
 
 import torch
 import torch.distributed as dist
 
-from msprobe.core.common.log import logger
+from msprobe.pytorch.common.log import logger
 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
 
@@ -83,11 +76,10 @@ 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):
-            if  fp32_partitioned_groups_flat[group_idx].grad is None:
+            if fp32_partitioned_groups_flat[group_idx].grad is None:
                 if partition_id == dist.get_world_size() - 1 and not self.is_stage3:
                     fp32_partitioned_groups_flat_grad = optimizer.flatten_dense_tensors_aligned(
                         optimizer.averaged_gradients[group_idx],
@@ -102,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]
@@ -111,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:
@@ -151,29 +143,33 @@ class MixPrecisionOptimizerMon(OptimizerMon):
     混合精度训练通过适当降低某些计算的精度来加速训练过程并减少内存消耗。
     """
 
+    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 torch_opt is not None:
+            self.map_fp16_tp_fp32_param(torch_opt)
 
-        if not self.fp16_to_fp32_param and mix_prec_opt is not None:
-            for fp16_group, fp32_group in zip(mix_prec_opt.float16_groups, mix_prec_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
         return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
 
 
 class MegatronDistributedOptimizerMon(OptimizerMon):
-    def fetch_mv(self, monitor, torch_opt, params2name):
-        mix_prec_opt = self.wrapped_optimizer
-        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")
-        if not self.fp16_to_fp32_param and mix_prec_opt is not None:
-            for fp16_group, shard_fp32_group in zip(mix_prec_opt.model_float16_groups,
-                                                    mix_prec_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
+        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):
+        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)
 
@@ -183,15 +179,40 @@ class MegatronFP32OptimizerMon(OptimizerMon):
         return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
 
 
+class MegatronChainedDistributedOptimizerMon(MegatronDistributedOptimizerMon):
+    def fetch_mv(self, monitor, torch_opt, params2name):
+        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):
+            torch_opt.state = {}
+            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):
+        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):
+            torch_opt.state = {}
+            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)
 
 
 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
@@ -210,13 +231,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):
@@ -229,12 +248,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 = []
@@ -242,7 +260,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())
@@ -250,7 +268,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]
@@ -264,32 +282,34 @@ 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)
 
 
 class DummyOptimizerMon(OptimizerMon):
     def fetch_mv(self, monitor, torch_opt, params2name):
-        return MVResult(exp_avg=None, exp_avg_sq=None, update=None, ratio=None)
+        return self._fetch_mv_in_adam(monitor, torch_opt, params2name)
 
 
 class OptimizerMonFactory:
     _optimizer_mon_map = {
-        "Megatron_Float16OptimizerWithFloat16Params": MixPrecisionOptimizerMon,
-        "Megatron_DistributedOptimizer": MegatronDistributedOptimizerMon,
-        "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

@@ -1,11 +1,26 @@
+# 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.
+
+import argparse
 import os
 import re
-import argparse
 from glob import glob
 
 import pandas as pd
 
-from msprobe.core.common.log import logger
+from msprobe.pytorch.common.log import logger
 
 
 def parse_logfile(logfile):
@@ -21,19 +36,19 @@ def parse_logfile(logfile):
 def parse_monitor_output(output_dir):
     reduced = {}
     unreduced = {}
-    for dir in glob(output_dir + '*'):
-        rank = int(re.findall('(?<=rank)[\d]*', dir)[0])
+    for directory in glob(output_dir + '*'):
+        rank = int(re.findall('(?<=rank)[\d]*', directory)[0])
         unreduced[rank] = []
         reduced[rank] = []
-        for file in os.listdir(dir):
-            df = pd.read_csv(os.path.join(dir, file))
+        for file in os.listdir(directory):
+            df = pd.read_csv(os.path.join(directory, file))
             if '_unreduced_' in file:
                 unreduced[rank].append(df)
                 pass
             elif '_reduced_' in file:
                 reduced[rank].append(df)
             else:
-                logger.info(f'unexpected file {file} in {dir}')
+                logger.info(f'unexpected file {file} in {directory}')
     return reduced, unreduced
 
 
@@ -41,7 +56,7 @@ def valid_reduce(reduced, unreduced, tp_size, dp_size, sequence_parallel):
     steps = len(reduced[0])
     world_size = len(reduced)
     errors = []
-    for index, row in unreduced[0][0].iterrows():
+    for _, row in unreduced[0][0].iterrows():
         param = row['param_name']
         is_tp_duplicate = False
         for step in range(2):
@@ -103,7 +118,7 @@ def valid_total_norm(total_norm, reduced, duplicate_embedding):
                 if step == 0:
                     logger.info(f'rank {rank} is duplicated in dp group')
                 continue
-            for index, row in reduced[rank][step].iterrows():
+            for _, row in reduced[rank][step].iterrows():
                 if duplicate_embedding and 'word_embedding' in row['param_name']:
                     continue
                 calculated_norm += row['norm'] ** 2

msprobe/pytorch/monitor/utils.py

@@ -16,13 +16,27 @@ import inspect
 from collections import namedtuple
 from datetime import timezone, timedelta
 from functools import wraps
+from datetime import datetime
+import os
+import re
 
 import torch
 
 from msprobe.core.common.const import MonitorConst, Const
-from msprobe.core.common.log import logger
+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
 
+
+device = "cpu"
+try:
+    import torch_npu
+    device = "npu"
+except ImportError:
+    if torch.cuda.is_available():
+        device = "cuda"
+
+NAN_TENSOR_ON_DEVICE = None
 FILE_MAX_SIZE = 10 * 1024 * 1024 * 1024
 FILE_NAME_MAX_LENGTH = 255
 DIRECTORY_MAX_LENGTH = 4096
@@ -39,6 +53,17 @@ class MsgConst:
     SPECIAL_CHAR = ["\n", "\r", "\u007F", "\b", "\f", "\t", "\u000B", "%08", "%0a", "%0b", "%0c", "%0d", "%7f"]
 
 
+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):
@@ -64,60 +89,19 @@ 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")
-    if not ops:
-        raise TypeError(f"specify ops to calculate metrics. Optional ops: {MonitorConst.OP_LIST}")
-
     valid_ops = []
     for op in ops:
         if op not in MonitorConst.OP_LIST:
             logger.warning(f"op {op} is not supported. Optional ops: {MonitorConst.OP_LIST}")
-        else:
-            valid_ops.append(op)
+            continue
+        valid_ops.append(op)
+    if not valid_ops:
+        default_op = MonitorConst.OP_LIST[0]
+        valid_ops.append(default_op)
+        logger.info_on_rank_0(f"There is no valid ops, default op {default_op} is used")
     return valid_ops
 
 
@@ -164,6 +148,11 @@ def validate_mg_distribution(mg_distribution):
         raise TypeError('mg_distribution should be a bool')
 
 
+def validate_param_distribution(param_distribution):
+    if not isinstance(param_distribution, bool):
+        raise TypeError('param_distribution should be a bool')
+
+
 def validate_cc_distribution(cc_distribution):
     if not isinstance(cc_distribution, dict):
         raise TypeError('cc_distribution should be a dictionary')
@@ -184,6 +173,11 @@ def validate_cc_distribution(cc_distribution):
             raise TypeError(f'{key} of cc_distribution is not supported.')
 
 
+def validate_squash_name(squash_name):
+    if not isinstance(squash_name, bool):
+        raise TypeError('squash_name should be a bool')
+
+
 def validate_alert(alert):
     if not isinstance(alert, dict):
         raise TypeError('alert should be a dictionary')
@@ -240,6 +234,9 @@ def validate_config(config):
     mg_distribution = config.get('mg_distribution', False)
     validate_mg_distribution(mg_distribution)
 
+    param_distribution = config.get('param_distribution', False)
+    validate_param_distribution(param_distribution)
+
     cc_distribution = config.get('cc_distribution', {})
     validate_cc_distribution(cc_distribution)
 
@@ -248,3 +245,42 @@ def validate_config(config):
 
     step_count_per_record = config.get('step_count_per_record', 1)
     validate_step_count_per_record(step_count_per_record)
+
+    squash_name = config.get('squash_name', True)
+    validate_squash_name(squash_name)
+
+    if not targets:
+        if xy_distribution:
+            config["all_xy"] = True
+        config["targets"] = {"": {}}
+
+
+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