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

msprobe/pytorch/monitor/module_hook.py

@@ -22,12 +22,13 @@ from functools import partial
 import pytz
 import torch
 import torch.distributed as dist
-from torch.optim.optimizer import register_optimizer_step_pre_hook, register_optimizer_step_post_hook
 from torch.utils.hooks import BackwardHook
 
-from msprobe.core.common.const import MonitorConst
+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, is_float8_tensor
 from msprobe.pytorch.monitor.anomaly_analyse import AnomalyDataWriter
 from msprobe.pytorch.monitor.anomaly_detect import AnomalyScanner, SummaryWriterWithAD, AnomalyDataFactory, \
     CSVWriterWithAD, BaseWriterWithAD, WriterInput
@@ -37,15 +38,16 @@ from msprobe.pytorch.monitor.features import get_sign_matches
 from msprobe.pytorch.monitor.module_metric import get_metrics, get_summary_writer_tag_name, \
     TensorMetrics, squash_param_name
 from msprobe.pytorch.monitor.module_spec_verifier import validate_config_spec
-from msprobe.pytorch.monitor.optimizer_collect import OptimizerMonFactory, OptimizerMon
-from msprobe.pytorch.monitor.utils import get_param_struct, validate_config, validate_ops, is_recomputation, \
-    get_output_base_dir, get_target_output_dir
+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, 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'
 if not torch_version_above_or_equal_2:
     raise ValueError("monitor require torch>=2.0")
 
+
 FORMAT_MAPPING = {
     MonitorConst.TENSORBOARD: SummaryWriterWithAD,
     MonitorConst.CSV: CSVWriterWithAD,
@@ -85,9 +87,6 @@ class ModuleHookContext:
         :param target_config: target obj in config json.
         :return:
         """
-        valid_key = [MonitorConst.ACTV_IN, MonitorConst.ACTV_OUT, MonitorConst.ACTVGRAD_IN, MonitorConst.ACTVGRAD_OUT]
-        if key_name not in valid_key:
-            raise ValueError(f"key({key_name}) error, valid_key: {valid_key}")
         cared = target_config.get(self.module_name, self.struct)
         if key_name in cared:
             target_module_config = cared[key_name]
@@ -178,20 +177,17 @@ class GradContext:
 class TrainerMon:
     tensor_metrics = TensorMetrics()
 
+    # 保留原opt_ty参数, 兼容msprobe1.2.2前旧版本
     def __init__(self, config_file_path, process_group=None, params_have_main_grad=True, opt_ty=None) -> None:
-        """
-        opt_ty: "Megatron_Float16OptimizerWithFloat16Params" or "Megatron_DistributedOptimizer"
-        """
         # TYPE1: 只在这里初始化的变量, 不会随着训练中途config配置改变而重置
         self.config_file_path = config_file_path
         self.process_group = get_process_group(process_group)
         self.params_have_main_grad = params_have_main_grad
-        self.opt_ty = opt_ty
-        self.mix_precision_optimizer_mon = OptimizerMonFactory.create_optimizer_mon(opt_ty)
         self.update_heatmap_visualizer = defaultdict(HeatmapVisualizer)
         self.ratio_heatmap_visualizer = defaultdict(HeatmapVisualizer)
         self.origin_step_func = None
-        self.config_timestamp = 0  # 后面有校验时间戳, 首次监控无需为了更新config文件时间戳而去改, 可通过switch开关直接打开
+        self.origin_start_grad_sync = None
+        self.config_timestamp = 0  # 后面有校验时间戳, 首次监控无需为了更新config文件时间戳而去改, 可通过dynamic_on开关直接打开
         self.config = load_json(config_file_path)
         validate_config(self.config)
 
@@ -219,13 +215,16 @@ class TrainerMon:
             self.pp_stage = 0
             self.group_mates = [0]
 
-        # TYPE2: 只会在monitor_gnorm_with_ad()主调中赋值的变量
+        # TYPE2: 只会在set_monitor()主调中赋值的变量
         self.model = None
         self.vpp = False
         self.dp_group = None
         self.tp_group = None
         self.enable_megatron = False
         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)
@@ -253,7 +252,7 @@ class TrainerMon:
         self.dynamic_enable = os.getenv("DYNAMIC_MONITOR", 'False').lower() == 'true'
         if self.dynamic_enable:
             logger.warning(f"DYNAMIC_MONITOR is set, "
-                           f"please make sure you have 'switch' and 'collect_times' item in {self.config_file_path}")
+                           f"please make sure you have 'dynamic_on' and 'collect_times' in {self.config_file_path}")
             self.monitoring = False
         else:
             self.set_config()
@@ -273,10 +272,6 @@ class TrainerMon:
     def ops(self, value):
         self._ops = validate_ops(value)
 
-    @staticmethod
-    def set_wrapped_optimizer(_wrapped_optimizer):
-        OptimizerMon.set_wrapped_optimizer(_wrapped_optimizer)
-
     @staticmethod
     def has_register_backward_hook(module_name, module):
         if hasattr(module, '_backward_hooks') and \
@@ -308,7 +303,7 @@ class TrainerMon:
         self.has_collect_times = 0  # 重设采集计数器
         self.print_struct = self.config.get("print_struct", False)
         self.module_rank_list = self.config.get("module_ranks", [])
-        self.format = self.config.get('format', 'tensorboard')
+        self.format = self.config.get('format', MonitorConst.CSV)
         self.eps = self.config.get('eps', 1e-8)
         self.ops = self.config.get('ops', [])
         self.ndigits = self.config.get('ndigits', 6)
@@ -330,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()
 
@@ -344,7 +337,13 @@ class TrainerMon:
 
         # 初始化writer, 创建输出目录
         if self.format not in FORMAT_MAPPING:
-            raise ValueError(f"Unsupported format: {self.format}")
+            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.warning("can only set ur_distribution when format is 'tensorboard', cancel ur_distribution")
+            self.ur_distribution = False
+
         writer = FORMAT_MAPPING[self.format]
         self.step_count_per_record = self.config.get('step_count_per_record', 1)
 
@@ -365,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. ")
@@ -393,105 +379,39 @@ class TrainerMon:
             logger.info_on_rank_0('> grad and momentum direction will not be compared.')
         if not self.cc_distribution.get('enable', False):
             logger.info_on_rank_0("> cc operator is not monitored.")
-        if not self.opt_ty:
-            if self.ur_distribution:
-                raise Exception("ur_distribution cannot be enabled with unknown optimizer.")
-            if self.mv_distribution:
-                raise Exception("mv_distribution cannot be enabled with unknown optimizer.")
-
-    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.VPP_SEP not in key]
-        for key in module_in_all_stage:
-            struct = targets.pop(key)
-            targets.update({f'{vpp_stage}{MonitorConst.VPP_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.VPP_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('exp_avg', opt_context.param_exp_avg)
-        v_tag_tensor_map = self.generate_param_map('efxp_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)
-        return self.grad_context.post, self.grad_context.pre
+    # 保留原接口, 兼容msprobe1.2.2前旧版本
+    def set_wrapped_optimizer(self, optimizer):
+        self.optimizer_trans = optimizer
 
-    def monitor_gnorm_with_ad(
+    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}')
         self.micro_batch_number = grad_acc_steps
         self.dp_group = dp_group
         self.tp_group = tp_group
+        self.optimizer_mon, self.optimizer_class = OptimizerMonFactory.create_optimizer_mon(optimizer)
         self.hook_step_final(optimizer)
         if not isinstance(model, list):
             model = [model]
@@ -507,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():
@@ -519,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():
@@ -529,6 +503,8 @@ class TrainerMon:
         return actv, actv_grad
 
     def reload_xy(self, xy_distribution=False):
+        logger.warning("reload_xy() is deprecated and will be removed in a future version. "
+                       "Use DYNAMIC_MONITOR instead.")
         self.xy_distribution = xy_distribution
 
         for handle in self.handles['xy']:
@@ -547,21 +523,23 @@ class TrainerMon:
         for _, fwd_context in self.module_fwd_hook_context_by_module.items():
             if len(fwd_context.actv) == 0:
                 continue
-            self.summary_writer.write_metrics(self.ops, fwd_context.actv, step, 'actv')
+            self.summary_writer.write_metrics(self.ops, fwd_context.actv, step, MonitorConst.ACTV)
             fwd_context.actv.clear()
         if self.grad_context.actv:
-            self.summary_writer.write_metrics(self.ops, self.grad_context.actv, step, 'actv_grad')
+            self.summary_writer.write_metrics(self.ops, self.grad_context.actv, step, MonitorConst.ACTVGRAD)
 
     def write_param_tb(self, opt_context):
         if not self.param_distribution:
             return
-        self.summary_writer.write_metrics(self.ops, opt_context.param_metric, opt_context.step, 'param')
+        self.summary_writer.write_metrics(self.ops, opt_context.param_metric, opt_context.step, MonitorConst.PARAM)
 
     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, opt_context.step, 'exp_avg')
-        self.summary_writer.write_metrics(self.ops, opt_context.exp_avg_sq_metric, opt_context.step, 'exp_avg_sq')
+        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,
+                                          opt_context.step, MonitorConst.EXP_AVG_SQ)
 
     def write_grad_tb(self, step):
         if not self.wg_distribution:
@@ -573,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]
@@ -592,15 +570,13 @@ class TrainerMon:
             # skip generate metrics
             if context.step < self.start_step or (context.step - self.start_step) % self.step_interval != 0:
                 return
-            if self.opt_ty in MonitorConst.DEEPSPEED_OPT_TY:
+            if MonitorConst.DEEPSPEED_ZERO_OPT_FILTER in self.optimizer_class:  # use deepspeed with zero1/2/3
                 if not self.name2indices:
-                    self.name2indices = self.mix_precision_optimizer_mon.get_param_index(self.param2name,
-                                                                                         self.name2index)
-                mv_result = self.mix_precision_optimizer_mon.fetch_mv(self, optimizer, self.param2name,
-                                                                      self.name2indices)
+                    self.name2indices = self.optimizer_mon.get_param_index(self.param2name, self.name2index, optimizer)
+                mv_result = self.optimizer_mon.fetch_mv(self, optimizer, self.param2name, self.name2indices)
                 self.param2name = mv_result.grad
             else:
-                mv_result = self.mix_precision_optimizer_mon.fetch_mv(self, optimizer, self.param2name)
+                mv_result = self.optimizer_mon.fetch_mv(self, optimizer, self.param2name)
             context.param_exp_avg = mv_result.exp_avg
             context.param_exp_avg_sq = mv_result.exp_avg_sq
             context.param_adam_update = mv_result.update
@@ -641,19 +617,13 @@ class TrainerMon:
                 optimizer_pre_step_hook(optimizer, args, kwargs)
                 out = func(*args, **kwargs)
                 return out
-
             return wrapper
 
         if self.optimizer_hooked:
             return
 
-        if optimizer:
-            optimizer.__class__.step = patch_step(optimizer.__class__.step, optimizer)
-            self.handles['optimizer'] = []
-        else:
-            if not self.module_rank_list or (dist.is_initialized() and dist.get_rank() in self.module_rank_list):
-                step_pre_hook = register_optimizer_step_pre_hook(optimizer_pre_step_hook)
-                self.handles['optimizer'] = [step_pre_hook]
+        optimizer.__class__.step = patch_step(optimizer.__class__.step, optimizer)
+
         self.optimizer_hooked = True
         return
 
@@ -677,11 +647,12 @@ class TrainerMon:
             logger.error(f"get config.json wrong because {e}, not updated, please check!!!")
             return
 
-        if config.get("switch", False):
+        if config.get("dynamic_on", False):
             try:
                 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:
@@ -729,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()
 
@@ -745,11 +719,49 @@ class TrainerMon:
                 return out
             return wrapper
 
-        if optimizer:
-            optimizer.__class__.step = patch_step(optimizer.__class__.step, optimizer)
-            self.origin_step_func = optimizer.__class__.step
-        else:
-            register_optimizer_step_post_hook(step_final_hook)
+        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):
@@ -764,17 +776,28 @@ class TrainerMon:
             bwd_context.reset()
         self.grad_context.reset()  # 权重梯度和激活值梯度都在这
 
-        for handle in self.handles['wgrads']:
-            handle.remove()
-        self.handles['wgrads'].clear()
-        self.weight_hooked = False
+        if self.origin_start_grad_sync:  # megatron
+            try:
+                from megatron.core.distributed.param_and_grad_buffer import Bucket
+                Bucket.start_grad_sync = self.origin_start_grad_sync
+                logger.info("remove Bucket start_grad_sync")
+            except ImportError:
+                pass
+            try:
+                from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBucketGroup
+                _ParamAndGradBucketGroup.start_grad_sync = self.origin_start_grad_sync
+                logger.info("remove _ParamAndGradBucketGroup start_grad_sync")
+            except ImportError:
+                pass
+        else:  # not megatron
+            for handle in self.handles['wgrads']:
+                handle.remove()
+            self.handles['wgrads'].clear()
+            self.weight_hooked = False
 
-        if len(self.handles['optimizer']) == 0 and self.optimizer_hooked:
+        if self.optimizer_hooked:
             optimizer.__class__.step = self.origin_step_func
-        else:
-            for handle in self.handles['optimizer']:
-                handle.remove()
-            self.handles['optimizer'].clear()
+
         for _, context in self.optimizer_context.items():
             context.reset()
         self.optimizer_hooked = False
@@ -782,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()
 
@@ -800,17 +824,17 @@ class TrainerMon:
 
     def _remove_all_hooks_final(self, optimizer):
         if self.dynamic_enable:
-            # 结束后自动重置switch为False等待用户手动开启
+            # 结束后自动重置dynamic_on为False等待用户手动开启
             try:
                 config = load_json(self.config_file_path)
-                config['switch'] = False
+                config['dynamic_on'] = False
                 save_json(self.config_file_path, config, indent=2)
                 config_timestamp = os.path.getmtime(self.config_file_path)
                 self.config_timestamp = config_timestamp
                 logger.info(
-                    "Finish monitor, set config'switch=False, will restart by set switch=True and update content")
+                    "Finish monitor, set config'dynamic_on=False, will restart by set it to True and update config")
             except Exception as e:
-                logger.warning(f"Finish monitor, set config'switch=False fail because {e}, please check!!!")
+                logger.warning(f"Finish monitor, set config'dynamic_on=False fail because {e}, please check!!!")
         logger.info("Finish monitor")
         self._remove_all_hooks(optimizer)
 
@@ -871,7 +895,7 @@ class TrainerMon:
 
     def _register_param_name(self):
         for vpp_stage, model_chunk in enumerate(self.model):
-            prefix = f'{vpp_stage}{MonitorConst.VPP_SEP}'
+            prefix = f'{vpp_stage}{MonitorConst.NAME_SEP}'
             self._register_chunk(model_chunk, prefix)
 
     def _is_target_module(self, module_name, targets, vpp_stage):
@@ -900,35 +924,37 @@ class TrainerMon:
             context: ModuleHookContext = self.module_fwd_hook_context_by_module[module]
             if not context.struct:
                 context.struct = {
-                    MonitorConst.ACTV_IN: get_param_struct(module_input),
-                    MonitorConst.ACTV_OUT: get_param_struct(module_output)
+                    Const.INPUT: get_param_struct(module_input),
+                    Const.OUTPUT: get_param_struct(module_output)
                 }
             if self.print_struct:
                 self.module_struct[context.module_name].update(context.struct)
                 return
             if not context.format_by_arg:
-                context.set_format_by_arg(MonitorConst.ACTV_IN, self.config['targets'])
-                context.set_format_by_arg(MonitorConst.ACTV_OUT, self.config['targets'])
+                context.set_format_by_arg(Const.INPUT, self.config['targets'])
+                context.set_format_by_arg(Const.OUTPUT, self.config['targets'])
             if not context.format_by_arg:
                 return
             if not context.verified:
-                context.focused_in_col = validate_config_spec(context.format_by_arg[MonitorConst.ACTV_IN],
+                context.focused_in_col = validate_config_spec(context.format_by_arg[Const.INPUT],
                                                               module_input, context.module_name,
-                                                              MonitorConst.ACTV_IN)
-                context.focused_out_col = validate_config_spec(context.format_by_arg[MonitorConst.ACTV_OUT],
+                                                              Const.INPUT)
+                context.focused_out_col = validate_config_spec(context.format_by_arg[Const.OUTPUT],
                                                                module_output, context.module_name,
-                                                               MonitorConst.ACTV_OUT)
+                                                               Const.OUTPUT)
                 context.verified = True
             # expect output be tensor type
             tbtag_tensor_map = {}
             cared_input = module_input if context.focused_in_col is None else module_input[context.focused_in_col]
             tbtag_tensor_map.update(
-                self.build_tbtag_tensor_map(f'{context.module_name}_{context.micro_step}', MonitorConst.ACTV_IN,
-                                            cared_input))
+                self.build_tbtag_tensor_map(
+                    f'{context.module_name}.{Const.INPUT}{MonitorConst.NAME_SEP}{context.micro_step}',
+                    MonitorConst.ACTV, cared_input))
             cared_output = module_output if context.focused_out_col is None else module_output[context.focused_out_col]
             tbtag_tensor_map.update(
-                self.build_tbtag_tensor_map(f'{context.module_name}_{context.micro_step}', MonitorConst.ACTV_OUT,
-                                            cared_output))
+                self.build_tbtag_tensor_map(
+                    f'{context.module_name}.{Const.OUTPUT}{MonitorConst.NAME_SEP}{context.micro_step}',
+                    MonitorConst.ACTV, cared_output))
 
             get_metrics(self.ops, tbtag_tensor_map, self.eps, context.actv)
             context.micro_step += 1
@@ -940,35 +966,37 @@ class TrainerMon:
             context: ModuleHookContext = self.module_bwd_hook_context_by_module[module]
             if not context.struct:
                 context.struct = {
-                    MonitorConst.ACTVGRAD_IN: get_param_struct(input_grad),
-                    MonitorConst.ACTVGRAD_OUT: get_param_struct(output_grad)
+                    MonitorConst.INPUT_GRAD: get_param_struct(input_grad),
+                    MonitorConst.OUTPUT_GRAD: get_param_struct(output_grad)
                 }
             if self.print_struct:
                 self.module_struct[context.module_name].update(context.struct)
                 return
             if not context.format_by_arg:
-                context.set_format_by_arg(MonitorConst.ACTVGRAD_IN, self.config['targets'])
-                context.set_format_by_arg(MonitorConst.ACTVGRAD_OUT, self.config['targets'])
+                context.set_format_by_arg(MonitorConst.INPUT_GRAD, self.config['targets'])
+                context.set_format_by_arg(MonitorConst.OUTPUT_GRAD, self.config['targets'])
             if not context.format_by_arg:
                 return
             if not context.verified:
-                context.focused_in_col = validate_config_spec(context.format_by_arg[MonitorConst.ACTVGRAD_IN],
-                                                              input_grad, context.module_name,
-                                                              MonitorConst.ACTVGRAD_IN)
-                context.focused_out_col = validate_config_spec(context.format_by_arg[MonitorConst.ACTVGRAD_OUT],
-                                                               output_grad, context.module_name,
-                                                               MonitorConst.ACTVGRAD_OUT)
+                context.focused_in_col = validate_config_spec(
+                    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],
+                    output_grad, context.module_name, MonitorConst.OUTPUT_GRAD)
                 context.verified = True
 
             tbtag_tensor_map = {}
             cared_input_grad = input_grad if context.focused_in_col is None else input_grad[context.focused_in_col]
             tbtag_tensor_map.update(
-                self.build_tbtag_tensor_map(f'{context.module_name}_{context.micro_step}', MonitorConst.ACTVGRAD_IN,
-                                            cared_input_grad))
+                self.build_tbtag_tensor_map(
+                    f'{context.module_name}.{Const.INPUT}{MonitorConst.NAME_SEP}{context.micro_step}',
+                    MonitorConst.ACTV, cared_input_grad))
             cared_output_grad = output_grad if context.focused_out_col is None else output_grad[context.focused_out_col]
             tbtag_tensor_map.update(
-                self.build_tbtag_tensor_map(f'{context.module_name}_{context.micro_step}', MonitorConst.ACTVGRAD_OUT,
-                                            cared_output_grad))
+                self.build_tbtag_tensor_map(
+                    f'{context.module_name}.{Const.OUTPUT}{MonitorConst.NAME_SEP}{context.micro_step}',
+                    MonitorConst.ACTV, cared_output_grad))
 
             if context.micro_step == 0 and context.actvgrad:
                 logger.warning(f"actvgrad context of {context.module_name} is not empty when first micro_step, "
@@ -1006,7 +1034,10 @@ class TrainerMon:
         def patch_sync(sync_grad_func):
             def wrapper(bucket):
                 grad_dict = {}
-                bucket_params_id_list = [id(params) for params in bucket.params_list]
+                # Megatron between core_r0.6.0 and core_r0.8.0, this bucket is Bucket.
+                # When megatron is core_r0.9.0, this bucket is _ParamAndGradBucketGroup.
+                # In megatron version core_r0.9.0, func start_grad_sync from Bucket moved to _ParamAndGradBucketGroup.
+                bucket_params_id_list = [id(params) for params in bucket.params]
                 for param, name in self.param2name.items():
                     if id(param) not in bucket_params_id_list:
                         continue
@@ -1025,18 +1056,28 @@ class TrainerMon:
 
             return wrapper
 
+        if not self.wg_distribution:
+            return
+
         try:
             from megatron.core.distributed.param_and_grad_buffer import Bucket
+            self.origin_start_grad_sync = Bucket.start_grad_sync
+            Bucket.start_grad_sync = patch_sync(Bucket.start_grad_sync)
             self.enable_megatron = True
+            logger.info("megatron version is >= core_r0.6.0 <= core_r0.8.0")
         except ImportError:
             self.enable_megatron = False
 
-        if not self.wg_distribution:
-            return
+        try:
+            from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBucketGroup
+            self.origin_start_grad_sync = _ParamAndGradBucketGroup.start_grad_sync
+            _ParamAndGradBucketGroup.start_grad_sync = patch_sync(_ParamAndGradBucketGroup.start_grad_sync)
+            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
 
-        if self.enable_megatron:
-            Bucket.start_grad_sync = patch_sync(Bucket.start_grad_sync)  # differ in different megatron version
-        else:
+        if not self.enable_megatron:
             self._hook_weights()
 
     def _hook_weights(self):
@@ -1049,10 +1090,14 @@ 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():
             key = get_summary_writer_tag_name(name, 'acc_grad', self.rank)
             setattr(param, 'micro_step', 0)