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

msprobe/pytorch/monitor/csv2tb.py

@@ -22,13 +22,15 @@ from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 
 from msprobe.core.common.const import MonitorConst
-from msprobe.core.common.file_utils import read_csv, create_directory, remove_path
+from msprobe.core.common.file_utils import read_csv, create_directory, remove_path, recursive_chmod
 from msprobe.core.common.utils import is_int
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.monitor.utils import get_target_output_dir
 
 all_data_type_list = ["actv", "actv_grad", "exp_avg", "exp_avg_sq", "grad_unreduced", "grad_reduced", "param"]
 CSV_FILE_SUFFIX = r"_\d+-\d+\.csv"
+MAX_PROCESS_NUM = 128
 
 
 def parse_step_line(line, ops):
@@ -76,6 +78,7 @@ def write_step(output_dirpath, parse_step_result, rank, data_type):
                 writer.add_scalar(tag, value, step)
 
 
+@recursion_depth_decorator("update_dict", max_depth=50)
 def update_dict(dict1, dict2):
     for key, value in dict2.items():
         if key in dict1:
@@ -115,11 +118,13 @@ def csv2tb_by_step_work(target_output_dirs, output_dirpath, data_type_list):
 def check_process_num(process_num):
     if not is_int(process_num) or process_num <= 0:
         raise ValueError(f"process_num({process_num}) is not a positive integer")
+    if process_num > MAX_PROCESS_NUM:
+        raise ValueError(f"The maximum supported process_num is {MAX_PROCESS_NUM}, current value: {process_num}.")
 
 
 def check_data_type_list(data_type_list):
     if data_type_list is None:
-        logger.info(f"data_type_list is None, use defualt all_data_type_list: {all_data_type_list}")
+        logger.info(f"data_type_list is None, use default all_data_type_list: {all_data_type_list}")
         return
     if not isinstance(data_type_list, list):
         raise ValueError(f"data_type_list({data_type_list}) is not a list")
@@ -161,4 +166,5 @@ def csv2tensorboard_by_step(
         p.start()
     for p in processes:
         p.join()
+    recursive_chmod(output_dirpath)
     logger.info(f"output has been saved to: {output_dirpath}")

msprobe/pytorch/monitor/distributed/wrap_distributed.py

@@ -24,6 +24,7 @@ import torch.nn as nn
 from msprobe.core.common.const import MonitorConst
 from msprobe.core.common.file_utils import load_yaml
 from msprobe.pytorch.monitor.module_metric import get_metrics, get_summary_writer_tag_name
+from msprobe.pytorch.common.log import logger
 
 try:
     import torch_npu
@@ -37,6 +38,7 @@ WrapDistributedOps = load_yaml(OpsPath).get("distributed", [])
 
 StackBlackListPath = os.path.join(os.path.dirname(__file__), "stack_blacklist.yaml")
 StackBlackList = load_yaml(StackBlackListPath).get("stack", [])
+MAX_STRING_LENGTH = 1000
 
 distributed_func = {}
 for f in dir(dist):
@@ -139,6 +141,8 @@ def get_process_group(process_group):
 
 
 def stack_filter(stack):
+    if len(stack) > MAX_STRING_LENGTH:
+        logger.warning(f'The character string contains more than {MAX_STRING_LENGTH}. re match is skipped.')
     for pattern in StackBlackList:
         if re.search(pattern, stack):
             return False
@@ -188,10 +192,12 @@ def update_data(old, new):
 
 
 def is_target_line(codeline):
-    stack = get_callstack()
-    whole_stack = ';'.join(stack)
     if codeline == []:
         return True
+    stack = get_callstack()
+    whole_stack = ';'.join(stack)
+    if len(whole_stack) > MAX_STRING_LENGTH:
+        logger.warning(f'The character string contains more than {MAX_STRING_LENGTH}. re match is skipped.')
     for pattern in codeline:
         if re.search(pattern, whole_stack):
             return True

msprobe/pytorch/monitor/module_hook.py

@@ -26,8 +26,9 @@ from torch.utils.hooks import BackwardHook
 
 from msprobe.core.common.const import MonitorConst, Const
 from msprobe.core.common.file_utils import load_json, save_json
+from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.common.utils import is_recomputation
+from msprobe.pytorch.common.utils import is_recomputation, is_float8_tensor
 from msprobe.pytorch.monitor.anomaly_analyse import AnomalyDataWriter
 from msprobe.pytorch.monitor.anomaly_detect import AnomalyScanner, SummaryWriterWithAD, AnomalyDataFactory, \
     CSVWriterWithAD, BaseWriterWithAD, WriterInput
@@ -39,7 +40,7 @@ from msprobe.pytorch.monitor.module_metric import get_metrics, get_summary_write
 from msprobe.pytorch.monitor.module_spec_verifier import validate_config_spec
 from msprobe.pytorch.monitor.optimizer_collect import OptimizerMonFactory
 from msprobe.pytorch.monitor.utils import get_param_struct, validate_config, validate_ops, \
-    get_output_base_dir, get_target_output_dir
+    get_output_base_dir, get_target_output_dir, chmod_tensorboard_dir, validate_set_monitor
 from msprobe.pytorch.monitor.visualizer import HeatmapVisualizer
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
@@ -176,7 +177,8 @@ class GradContext:
 class TrainerMon:
     tensor_metrics = TensorMetrics()
 
-    def __init__(self, config_file_path, process_group=None, params_have_main_grad=True) -> None:
+    # 保留原opt_ty参数, 兼容msprobe1.2.2前旧版本
+    def __init__(self, config_file_path, process_group=None, params_have_main_grad=True, opt_ty=None) -> None:
         # TYPE1: 只在这里初始化的变量, 不会随着训练中途config配置改变而重置
         self.config_file_path = config_file_path
         self.process_group = get_process_group(process_group)
@@ -222,6 +224,7 @@ class TrainerMon:
         self.micro_batch_number = 1
         self.optimizer_class = None
         self.optimizer_mon = None
+        self.optimizer_trans = None
 
         # TYPE3: 会随着训练中途config配置更新或监控状态改变而重置的变量
         self.module_fwd_hook_context_by_module = defaultdict(ModuleHookContext)
@@ -322,8 +325,6 @@ class TrainerMon:
             self.cc_log_only = self.cc_distribution.get('cc_log_only', False)
             self.cc_logged_stack = defaultdict(set)
             self.cc_pre_hook = self.cc_distribution.get('cc_pre_hook', False)
-            self.handles['cc'] = api_register.initialize_hook(*create_hooks(context=self.cc_context, monitor=self))
-            api_register.redirect_api()
 
         self.common_info()
 
@@ -336,11 +337,11 @@ class TrainerMon:
 
         # 初始化writer, 创建输出目录
         if self.format not in FORMAT_MAPPING:
-            logger.error(f"Unsupported format: {self.format}, use default format: {MonitorConst.CSV}")
+            logger.warning(f"Unsupported format: {self.format}, use default format: {MonitorConst.CSV}")
             self.format = MonitorConst.CSV
 
         if self.ur_distribution and self.format != 'tensorboard':
-            logger.error("can only set ur_distribution when format is 'tensorboard', cancel ur_distribution")
+            logger.warning("can only set ur_distribution when format is 'tensorboard', cancel ur_distribution")
             self.ur_distribution = False
 
         writer = FORMAT_MAPPING[self.format]
@@ -363,19 +364,6 @@ class TrainerMon:
                                                              self.rank)
                 self.anomaly_data_writer.init_detected_json()
 
-    def adhoc_check(self, target_tensor: torch.tensor, module_name: str, tensor_name: str, rank_list, ops_list):
-        rank = None
-        if dist.is_initialized():
-            rank = dist.get_rank()
-            if (rank not in rank_list) and len(rank_list) != 0:
-                return
-        self.tensor_metrics.stat_insert(target_tensor, ops_list, module_name, tensor_name, rank)
-
-    def build_tbtag_tensor_map(self, module_name, tag, tensor):
-        key = get_summary_writer_tag_name(module_name, tag, self.rank)
-        self._register_param_call_id("_hook_module", key)
-        return {key: tensor}
-
     def common_info(self):
         if not self.xy_distribution:
             logger.info_on_rank_0("> module input/output input_grad/output_grad is not monitored. ")
@@ -392,94 +380,31 @@ class TrainerMon:
         if not self.cc_distribution.get('enable', False):
             logger.info_on_rank_0("> cc operator is not monitored.")
 
-    def hook_modules(self):
-        if self.module_rank_list and (self.rank not in self.module_rank_list):
-            return
-
-        targets = self.config['targets']
-        module_in_all_stage = [key for key in targets.keys() if MonitorConst.NAME_SEP not in key]
-        for key in module_in_all_stage:
-            struct = targets.pop(key)
-            targets.update({f'{vpp_stage}{MonitorConst.NAME_SEP}{key}': struct for vpp_stage in range(len(self.model))})
-
-        hooked_count = 0
-        for vpp_stage, model_chunk in enumerate(self.model):
-            vpp_stage = f'{vpp_stage}{MonitorConst.NAME_SEP}'
-            targets = [x for x, _ in model_chunk.named_modules()] if self.print_struct else self.config[
-                'targets'].keys()
-            hooked_count += self._hook_module(targets, model_chunk, vpp_stage)
-
-        logger.info_on_rank_0(f"> {hooked_count} modules are monitored.")
-
-        def clone_if_tensor(args):
-            if isinstance(args, tuple):
-                return tuple([clone_if_tensor(arg) for arg in args])
-            elif isinstance(args, torch.Tensor):
-                return args.clone()
-            else:
-                return args
-
-        @torch.no_grad
-        def wrap_hook_setup(setup):
-            def wrapped_setup(*args, **kwargs):
-                args = setup(*args, **kwargs)
-                args = clone_if_tensor(args)
-                return args
-
-            return wrapped_setup
-
-        BackwardHook.setup_output_hook = wrap_hook_setup(BackwardHook.setup_output_hook)
-
-        return
-
-    def generate_param_metrics(self, opt_context):
-        if not self.param_distribution:
-            return
-        get_metrics(self.ops, self.name2param, self.eps, opt_context.param_metric)
-
-    def generate_mv_metrics(self, opt_context):
-        if not self.mv_distribution:
-            return
-        opt_context.exp_avg_metric = {}
-        opt_context.exp_avg_sq_metric = {}
-        m_tag_tensor_map = self.generate_param_map(MonitorConst.EXP_AVG, opt_context.param_exp_avg)
-        v_tag_tensor_map = self.generate_param_map(MonitorConst.EXP_AVG_SQ, opt_context.param_exp_avg_sq)
-        get_metrics(self.ops, m_tag_tensor_map, self.eps, opt_context.exp_avg_metric)
-        get_metrics(self.ops, v_tag_tensor_map, self.eps, opt_context.exp_avg_sq_metric)
-
-    def generate_wgrad_metrics(self):
-        if not self.wg_distribution:
-            return {}, {}
-
-        if self.weight_hooked:
-            get_metrics(self.ops, self.grad_context.acc, self.eps, self.grad_context.acc_metric)
-
-        grad_dict = {}
-        for param, name in self.param2name.items():
-            if self.duplicate_param.get(name, False):
-                continue
-            grad = param.main_grad if self.params_have_main_grad else param.grad
-            if grad is None:
-                logger.warning(f"grad is None: {name}, maybe something wrong happened.")
-                continue
-            tag = self.name2tag.get(name, {}).get(MonitorConst.POST_GRAD)
-            self._register_param_call_id("hook_optimizer", tag)
-            grad_dict[tag] = grad
+    # 保留原接口, 兼容msprobe1.2.2前旧版本
+    def monitor_gnorm_with_ad(self, model, optimizer=None, grad_acc_steps=1, tp_group=None, dp_group=None,
+                              start_iteration=0):
+        if optimizer is None:
+            optimizer = getattr(self, "optimizer_trans", None)  # 兼容老版本可传None的情况, 从set_wrapped_optimizer获取
+            if optimizer is None:
+                logger.error("monitor_gnorm_with_ad: please set_wrapped_optimizer before it or input optimizer!=None")
+                return
+        self.set_monitor(model, optimizer, grad_acc_steps, tp_group, dp_group, start_iteration)
 
-        get_metrics(self.ops, grad_dict, self.eps, self.grad_context.post)
-        unreduced_grad = self.grad_context.acc_metric if self.weight_hooked else self.grad_context.pre
-        return self.grad_context.post, unreduced_grad
+    # 保留原接口, 兼容msprobe1.2.2前旧版本
+    def set_wrapped_optimizer(self, optimizer):
+        self.optimizer_trans = optimizer
 
     def set_monitor(
             self,
             model,
+            optimizer,
             grad_acc_steps=1,
-            optimizer=None,
             tp_group=None,
             dp_group=None,
             start_iteration=0
     ):
         """External interface"""
+        grad_acc_steps, start_iteration = validate_set_monitor(grad_acc_steps, start_iteration)
         global start_step
         start_step = start_iteration
         logger.info(f'grad acc steps {grad_acc_steps}')
@@ -502,8 +427,24 @@ class TrainerMon:
         self.hook_optimizer(optimizer)
         self._patch_grad_sync()
         self.hook_modules()
+        if self.cc_distribution.get('enable', False):
+            self.handles['cc'] = api_register.initialize_hook(*create_hooks(context=self.cc_context, monitor=self))
+            api_register.redirect_api()
         self.monitoring = True
 
+    def adhoc_check(self, target_tensor: torch.tensor, module_name: str, tensor_name: str, rank_list, ops_list):
+        rank = None
+        if dist.is_initialized():
+            rank = dist.get_rank()
+            if (rank not in rank_list) and len(rank_list) != 0:
+                return
+        self.tensor_metrics.stat_insert(target_tensor, ops_list, module_name, tensor_name, rank)
+
+    def build_tbtag_tensor_map(self, module_name, tag, tensor):
+        key = get_summary_writer_tag_name(module_name, tag, self.rank)
+        self._register_param_call_id("_hook_module", key)
+        return {key: tensor}
+
     def generate_param_map(self, tag, param_tensor):
         metrics = {}
         for name in self.param2name.values():
@@ -514,6 +455,44 @@ class TrainerMon:
             metrics[key] = param_tensor[name]
         return metrics
 
+    def generate_param_metrics(self, opt_context):
+        if not self.param_distribution:
+            return
+        get_metrics(self.ops, self.name2param, self.eps, opt_context.param_metric)
+
+    def generate_mv_metrics(self, opt_context):
+        if not self.mv_distribution:
+            return
+        opt_context.exp_avg_metric = {}
+        opt_context.exp_avg_sq_metric = {}
+        m_tag_tensor_map = self.generate_param_map(MonitorConst.EXP_AVG, opt_context.param_exp_avg)
+        v_tag_tensor_map = self.generate_param_map(MonitorConst.EXP_AVG_SQ, opt_context.param_exp_avg_sq)
+        get_metrics(self.ops, m_tag_tensor_map, self.eps, opt_context.exp_avg_metric)
+        get_metrics(self.ops, v_tag_tensor_map, self.eps, opt_context.exp_avg_sq_metric)
+
+    def generate_wgrad_metrics(self):
+        if not self.wg_distribution:
+            return {}, {}
+
+        if self.weight_hooked:
+            get_metrics(self.ops, self.grad_context.acc, self.eps, self.grad_context.acc_metric)
+
+        grad_dict = {}
+        for param, name in self.param2name.items():
+            if self.duplicate_param.get(name, False):
+                continue
+            grad = param.main_grad if self.params_have_main_grad else param.grad
+            if grad is None:
+                logger.warning(f"grad is None: {name}, maybe something wrong happened.")
+                continue
+            tag = self.name2tag.get(name, {}).get(MonitorConst.POST_GRAD)
+            self._register_param_call_id("hook_optimizer", tag)
+            grad_dict[tag] = grad
+
+        get_metrics(self.ops, grad_dict, self.eps, self.grad_context.post)
+        unreduced_grad = self.grad_context.acc_metric if self.weight_hooked else self.grad_context.pre
+        return self.grad_context.post, unreduced_grad
+
     def generate_xy_metrics(self):
         actv = {}
         for fwd_context in self.module_fwd_hook_context_by_module.values():
@@ -557,9 +536,9 @@ class TrainerMon:
     def write_mv_tb(self, opt_context):
         if not self.mv_distribution:
             return
-        self.summary_writer.write_metrics(self.ops, opt_context.exp_avg_metric, 
+        self.summary_writer.write_metrics(self.ops, opt_context.exp_avg_metric,
                                           opt_context.step, MonitorConst.EXP_AVG)
-        self.summary_writer.write_metrics(self.ops, opt_context.exp_avg_sq_metric, 
+        self.summary_writer.write_metrics(self.ops, opt_context.exp_avg_sq_metric,
                                           opt_context.step, MonitorConst.EXP_AVG_SQ)
 
     def write_grad_tb(self, step):
@@ -572,7 +551,7 @@ class TrainerMon:
             self.summary_writer.write_metrics(self.ops, self.grad_context.acc_metric, step, 'grad_unreduced')
         self.summary_writer.write_metrics(self.ops, self.grad_context.post, step, 'grad_reduced')
 
-    def hook_optimizer(self, optimizer=None):
+    def hook_optimizer(self, optimizer):
         # in DDP by default use params_have_main_grad
         def optimizer_pre_step_hook(optimizer, args, kwargs):
             context = self.optimizer_context[optimizer]
@@ -638,7 +617,6 @@ class TrainerMon:
                 optimizer_pre_step_hook(optimizer, args, kwargs)
                 out = func(*args, **kwargs)
                 return out
-
             return wrapper
 
         if self.optimizer_hooked:
@@ -674,6 +652,7 @@ class TrainerMon:
                 validate_config(config)
                 self.config = config
                 self.set_config()
+                self.start_step = context.step  # 动态启停时不受原start_step影响，永远从下一步开始
                 logger.warning(f"config is updated at step{context.step - 1}, "
                                f"will start new hook at step{context.step}.")
             except Exception as e:
@@ -721,6 +700,9 @@ class TrainerMon:
                     if self.anomaly_data_factory:
                         self.anomaly_data_writer.write_detected_json(self.summary_writer.get_anomalies())
                     self.summary_writer.clear_anomalies()
+
+                    if self.format == MonitorConst.TENSORBOARD:
+                        chmod_tensorboard_dir(self.tensorboard_dir)
                     self.call_id = 0
                     self.param_name_call_id.clear()
 
@@ -739,7 +721,47 @@ class TrainerMon:
 
         optimizer.__class__.step = patch_step(optimizer.__class__.step, optimizer)
         self.origin_step_func = optimizer.__class__.step
+        return
+
+    def hook_modules(self):
+        if self.module_rank_list and (self.rank not in self.module_rank_list):
+            return
+
+        targets = self.config['targets']
+        module_in_all_stage = [key for key in targets.keys() if MonitorConst.NAME_SEP not in key]
+        for key in module_in_all_stage:
+            struct = targets.pop(key)
+            targets.update({f'{vpp_stage}{MonitorConst.NAME_SEP}{key}': struct for vpp_stage in range(len(self.model))})
+
+        hooked_count = 0
+        for vpp_stage, model_chunk in enumerate(self.model):
+            vpp_stage = f'{vpp_stage}{MonitorConst.NAME_SEP}'
+            targets = [x for x, _ in model_chunk.named_modules()] if self.print_struct else self.config[
+                'targets'].keys()
+            hooked_count += self._hook_module(targets, model_chunk, vpp_stage)
+
+        logger.info_on_rank_0(f"> {hooked_count} modules are monitored.")
+
+        @recursion_depth_decorator('msprobe.pytorch.monitor.clone_if_tensor')
+        def clone_if_tensor(args):
+            if isinstance(args, tuple):
+                return tuple([clone_if_tensor(arg) for arg in args])
+            elif isinstance(args, torch.Tensor) and not is_float8_tensor(args):
+                return args.clone()
+            else:
+                return args
 
+        @torch.no_grad
+        def wrap_hook_setup(setup):
+            def wrapped_setup(*args, **kwargs):
+                args = setup(*args, **kwargs)
+                args = clone_if_tensor(args)
+                return args
+
+            return wrapped_setup
+
+        BackwardHook.setup_input_hook = wrap_hook_setup(BackwardHook.setup_input_hook)
+        BackwardHook.setup_output_hook = wrap_hook_setup(BackwardHook.setup_output_hook)
         return
 
     def _remove_all_hooks(self, optimizer):
@@ -783,6 +805,7 @@ class TrainerMon:
         for handle in self.handles['cc']:
             handle.remove()
         self.handles['cc'].clear()
+        api_register.restore_api()
         for _, context in self.cc_context.items():
             context.reset()
 
@@ -956,7 +979,7 @@ class TrainerMon:
                 return
             if not context.verified:
                 context.focused_in_col = validate_config_spec(
-                    context.format_by_arg[MonitorConst.INPUT_GRAD], 
+                    context.format_by_arg[MonitorConst.INPUT_GRAD],
                     input_grad, context.module_name, MonitorConst.INPUT_GRAD)
                 context.focused_out_col = validate_config_spec(
                     context.format_by_arg[MonitorConst.OUTPUT_GRAD],
@@ -1052,7 +1075,7 @@ class TrainerMon:
             self.enable_megatron = True
             logger.info("megatron version is > core_r0.8.0 <= core_r0.9.0")
         except ImportError:
-            self.enable_megatron = False
+            self.enable_megatron = False | self.enable_megatron
 
         if not self.enable_megatron:
             self._hook_weights()
@@ -1067,9 +1090,12 @@ class TrainerMon:
             if param.micro_step == self.micro_batch_number:
                 param.micro_step = 0
                 if self.params_have_main_grad:
-                    context_dict[key] = param.main_grad.clone()
+                    grad = param.main_grad
                 else:
-                    context_dict[key] = param.grad.clone()
+                    grad = param.grad
+                if is_float8_tensor(grad):
+                    grad = grad.float()
+                context_dict[key] = grad.clone()
 
         logger.info("hooking weights.")
         for param, name in self.param2name.items():

msprobe/pytorch/monitor/module_metric.py

@@ -16,6 +16,7 @@ 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 get_nan_tensor
 
@@ -166,6 +167,8 @@ def get_metrics(ops, tag2tensor, eps, out_dict=None):
             # Non-tensor in/output filled with nan.
             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

@@ -185,7 +185,7 @@ class MegatronChainedDistributedOptimizerMon(MegatronDistributedOptimizerMon):
             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 torch_opt.chained_optimizers:
                 torch_opt.state.update(opt.optimizer.state)
@@ -198,7 +198,7 @@ class MegatronChainedMixPrecisionOptimizerMon(MixPrecisionOptimizerMon):
             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 torch_opt.chained_optimizers:
                 torch_opt.state.update(opt.optimizer.state)
@@ -206,9 +206,60 @@ class MegatronChainedMixPrecisionOptimizerMon(MixPrecisionOptimizerMon):
 
 
 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, torch_opt):

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):