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

msprobe/pytorch/monitor/module_hook.py

@@ -0,0 +1,870 @@
+# Copyright (c) 2024-2024, 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 time
+import json
+import os
+import uuid
+from collections import defaultdict
+from datetime import datetime, timezone
+from functools import partial
+
+import pytz
+import torch
+import torch.distributed as dist
+from msprobe.core.common.const import MonitorConst
+from msprobe.core.common.file_utils import load_json
+from msprobe.core.common.log import logger
+from msprobe.pytorch.monitor.anomaly_analyse import AnomalyDataWriter
+from msprobe.pytorch.monitor.anomaly_detect import AnomalyScanner, SummaryWriterWithAD, AnomalyDataFactory, \
+    CSVWriterWithAD, BaseWriterWithAD, WriterInput
+from msprobe.pytorch.monitor.distributed.wrap_distributed import api_register, create_hooks, op_aggregate, \
+    get_process_group
+from msprobe.pytorch.monitor.features import get_sign_matches
+from msprobe.pytorch.monitor.module_metric import get_metrics, write_metrics_base, get_summary_writer_tag_name, \
+    TensorMetrics, write_metrics_csv, 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
+from msprobe.pytorch.monitor.visualizer import HeatmapVisualizer
+from torch.optim.optimizer import register_optimizer_step_pre_hook, register_optimizer_step_post_hook
+from torch.utils.hooks import BackwardHook
+
+try:
+    import torch_npu
+except ImportError:
+    pass
+
+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")
+
+output_base_dir = os.getenv(MonitorConst.MONITOR_OUTPUT_DIR, MonitorConst.DEFAULT_MONITOR_OUTPUT_DIR)
+
+FORMAT_MAPPING = {
+    MonitorConst.TENSORBOARD: (SummaryWriterWithAD, write_metrics_base),
+    MonitorConst.CSV: (CSVWriterWithAD, write_metrics_csv),
+    MonitorConst.API: (BaseWriterWithAD, write_metrics_base)
+}
+
+
+def param_is_not_tensor_parallel_duplicate(param, tp_group):
+    return (hasattr(param, 'tensor_model_parallel') and param.tensor_model_parallel) or (
+            torch.distributed.get_rank(group=tp_group) == 0
+    )
+
+
+def param_is_data_parallel_duplicate(dp_group):
+    return torch.distributed.get_rank(group=dp_group) != 0
+
+
+class ModuleHookContext:
+    def __init__(self, module_name) -> None:
+        self.step = 0
+        self.micro_step = 0
+        self.actv = defaultdict(dict)
+        self.actvgrad = []
+        self.module_name = module_name
+        self.struct = {}
+        self.format_by_arg = {}
+        self.verified = False
+        self.focused_in_col = 0
+        self.focused_out_col = 0
+        self.ignore_in = False  # no need to care when no key 'input' or 'input_grad' found
+
+    def set_format_by_arg(self, key_name: str, target_config: dict):
+        cared = target_config.get(self.module_name, self.struct)
+        if key_name in cared:
+            if isinstance(cared[key_name], dict):
+                # current cared is self.struct
+                config = cared[key_name].get('config')
+                self.format_by_arg[key_name] = config
+            else:
+                # current cared is target_config[self.module_name]
+                self.format_by_arg[key_name] = cared[key_name]
+        elif key_name in ['input', 'input_grad']:
+            self.ignore_in = True
+
+
+class OptimizerContext:
+    def __init__(self) -> None:
+        self.step = 0
+        self.param_effective_rank = defaultdict(float)
+        self.param_mg_direction = defaultdict(float)
+        self.param_adam_update = defaultdict()
+        self.param_adam_ratio = defaultdict()
+        self.param_weight_grad = defaultdict()
+        self.param_exp_avg = defaultdict()
+        self.exp_avg_metric = {}
+        self.param_exp_avg_sq = defaultdict()
+        self.exp_avg_sq_metric = {}
+        self.metric_dict = {}
+        self.param_metric = {}
+
+
+class CommunicationContext:
+    def __init__(self) -> None:
+        self.data = {}
+
+    @staticmethod
+    def _agg(data):
+        aggregated_data = {}
+        for tag, op2tensorlist in data.items():
+            aggregated_data[tag] = {}
+            for op, tensorlist in op2tensorlist.items():
+                aggregated_data[tag][op] = op_aggregate(op, tensorlist)
+        return aggregated_data
+
+    def reset(self):
+        self.data = {}
+
+    def aggregate(self):
+        self.data = self._agg(self.data)
+
+
+class GradContext:
+    def __init__(self) -> None:
+        self.pre = {}
+        self.post = {}
+        self.acc_metric = {}
+        self.acc = {}
+        self.actv = {}
+
+    def reset(self):
+        self.pre.clear()
+        self.post.clear()
+        self.acc_metric.clear()
+        self.acc.clear()
+        self.actv.clear()
+
+
+class TrainerMon:
+    tensor_metrics = TensorMetrics()
+
+    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"
+        """
+        self.module_fwd_hook_context_by_module = defaultdict(ModuleHookContext)
+        self.module_bwd_hook_context_by_module = defaultdict(ModuleHookContext)
+        self.optimizer_context = defaultdict(OptimizerContext)
+        self.cc_context = defaultdict(CommunicationContext)
+        self.grad_context = GradContext()
+        self.process_group = get_process_group(process_group)
+        self.params_have_main_grad = params_have_main_grad
+        self.opt_ty = opt_ty
+        self.config = load_json(config_file_path)
+        validate_config(self.config)
+
+        self.module_rank_list = self.config.get("module_ranks", [])
+        self.format = self.config.get('format', 'tensorboard')
+        self.eps = self.config.get('eps', 1e-8)
+        self.ops = self.config.get('ops', [])
+        self.ndigits = self.config.get('ndigits', 6)
+        self.all_xy = self.config.get('all_xy', False)
+        self.xy_distribution = self.config.get('xy_distribution', False)
+        self.forward_only = self.config.get('forward_only', False)
+        self.backward_only = self.config.get('backward_only', False)
+        self.ur_distribution = self.config.get('ur_distribution', False)
+        self.mv_distribution = self.config.get("mv_distribution", False)
+        self.wg_distribution = self.config.get("wg_distribution", False)
+        self.param_distribution = self.config.get("param_distribution", False)
+        self.mg_direction = self.config.get('mg_direction', False)
+        self.cc_distribution = self.config.get("cc_distribution", {})
+        if not self.cc_distribution.get('enable', False):
+            self.cc_log_only = False
+        else:
+            self.cc_codeline = self.cc_distribution.get('cc_codeline', [])
+            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)
+            api_register.initialize_hook(*create_hooks(context=self.cc_context, monitor=self))
+            api_register.redirect_api()
+
+        self.common_info()
+
+        alert_setting = self.config.get('alert', {"rules": []})
+        self.alert_rules = AnomalyScanner.load_rules(alert_setting["rules"])
+
+        # 设置时区，使用 'UTC' 作为示例
+        local_tz = pytz.timezone("Asia/Shanghai")  # 根据需要调整为目标时区
+
+        cur_time = datetime.now(local_tz).strftime('%b%d_%H-%M-%S')
+        unique_id = str(uuid.uuid4())[:8]
+
+        if dist.is_initialized():
+            rank = dist.get_rank()
+            tensorboard_dir = os.path.join(output_base_dir, f"{cur_time}-rank{rank}-{unique_id}")
+            pp_stage = dist.get_group_rank(self.process_group, rank)
+            group_mates = dist.get_process_group_ranks(self.process_group)
+        else:
+            rank = 0
+            tensorboard_dir = os.path.join(output_base_dir, f"{cur_time}-{unique_id}")
+            pp_stage = 0
+            group_mates = [0]
+        self.rank = rank
+
+        # 初始化AnomalyData工厂
+        self.anomaly_data_factory = None
+        if alert_setting.get('dump', False):
+            self.anomaly_data_factory = AnomalyDataFactory(rank, pp_stage, group_mates)
+
+        if self.format not in FORMAT_MAPPING:
+            raise ValueError(f"Unsupported format: {self.format}")
+        writer, self.write_metrics = FORMAT_MAPPING[self.format]
+        self.step_count_per_record = self.config.get('step_count_per_record', 1)
+
+        if (rank in self.module_rank_list) or len(self.module_rank_list) == 0:
+            self.summary_writer = writer(
+                WriterInput(
+                    tensorboard_dir,
+                    self.alert_rules,
+                    unique_id,
+                    None,
+                    self.anomaly_data_factory,
+                    self.ndigits,
+                    self.step_count_per_record
+                )
+            )
+            # 初始化anomaly detected文件目录
+            if self.anomaly_data_factory:
+                self.anomaly_data_writer = AnomalyDataWriter(os.path.join(output_base_dir, "anomaly_detected"), rank)
+                self.anomaly_data_writer.init_detected_json()
+
+        # A HeatmapVisualizer instance is associated with an image
+        self.update_heatmap_visualizer = defaultdict(HeatmapVisualizer)
+        self.ratio_heatmap_visualizer = defaultdict(HeatmapVisualizer)
+        self.micro_batch_number = 1
+
+        self.model = None
+        self.weight_hooked = False
+        self.optimizer_hooked = False
+        self.param_registered = False
+        self.vpp = False
+        self.dp_group = None
+        self.tp_group = None
+        self.enable_megatron = False
+
+        self.param2name = defaultdict(str)
+        self.name2index = defaultdict()
+        self.name2indices = defaultdict()
+        self.name2param = {}
+        self.param_name_call_id = {}
+        self.duplicate_param = {}
+        self.name2tag = {}
+        self.call_id = 0
+        self.grad_accs = []
+        self.handles = defaultdict(list)
+
+        self.mix_precision_optimizer_mon = OptimizerMonFactory.create_optimizer_mon(opt_ty)
+        self.print_struct = self.config.get("print_struct", False)
+        self.struct_printed = False
+        self.module_struct = {}
+
+    def __del__(self):
+        if hasattr(self, "summary_writer"):
+            self.summary_writer.close()
+
+    @property
+    def ops(self):
+        return self._ops
+
+    @ops.setter
+    def ops(self, value):
+        self._ops = validate_ops(value)
+
+    @staticmethod
+    def set_wrapped_optimizer(_wrapped_optimizer):
+        OptimizerMon.set_wrapped_optimizer(_wrapped_optimizer)
+
+    @staticmethod
+    def adhoc_check(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
+        TrainerMon.tensor_metrics.stat_insert(target_tensor, ops_list, module_name, tensor_name, rank)
+
+    @staticmethod
+    def build_tbtag_tensor_map(module_name, tag, tensor):
+        metrics = {}
+        rank = dist.get_rank() if dist.is_initialized() else None
+        key = get_summary_writer_tag_name(module_name, tag, rank)
+        if torch.is_tensor(tensor):
+            metrics[key] = tensor
+        return metrics
+
+    @staticmethod
+    def generate_cc_metrics(cc_name, cc_tensor):
+        metrics = defaultdict(dict)
+        rank = dist.get_rank() if dist.is_initialized() else None
+        for op, tag2tensor in cc_tensor.data.items():
+            for tag, tensor in tag2tensor.items():
+                key = get_summary_writer_tag_name(cc_name, tag, rank)
+                metrics[op].update({key: tensor})
+        cc_tensor.reset()
+        return metrics
+
+    def common_info(self):
+        if not self.xy_distribution:
+            logger.info_on_rank_0("> module input/output input_grad/output_grad is not monitored. ")
+        if self.forward_only:
+            logger.info_on_rank_0("> only module forward is monitored. ")
+        if not self.ur_distribution:
+            logger.info_on_rank_0("> update vector and ratio vector of adam is not monitored. ")
+        if not self.mv_distribution:
+            logger.info_on_rank_0("> momentum and variance of adam is not monitored. ")
+        if not self.wg_distribution:
+            logger.info_on_rank_0("> weight grad of specified module is not monitored. ")
+        if not self.mg_direction:
+            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, model: torch.nn.Module, grad_acc_steps):
+        if self.module_rank_list and (self.rank not in self.module_rank_list):
+            return
+
+        if not isinstance(model, list):
+            model = [model]
+        self.model = model
+        self._register_param_name(model)
+
+        self.micro_batch_number = grad_acc_steps
+
+        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(model))})
+
+        hooked_count = 0
+        for vpp_stage, model_chunk in enumerate(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} out of {len(self.config['targets'])} 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)
+
+        if not self.optimizer_hooked:
+            self.hook_optimizer()
+        return
+
+    def generate_param_metrics(self, opt_context):
+        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)
+            grad_dict[tag] = grad
+
+        get_metrics(self.ops, grad_dict, self.eps, self.grad_context.post)
+        return self.grad_context.post, self.grad_context.pre
+
+    def monitor_gnorm_with_ad(self, model, grad_acc_steps=1, optimizer=None, tp_group=None, dp_group=None):
+        """External interface"""
+        logger.info(f'grad acc steps {grad_acc_steps}')
+        self.hook_optimizer(optimizer)
+        self.micro_batch_number = grad_acc_steps
+
+        self.dp_group = dp_group
+        self.tp_group = tp_group
+
+        self._register_param_name(model)
+        self._patch_grad_sync()
+        self.hook_modules(model, grad_acc_steps)
+
+    def generate_param_map(self, tag, param_tensor):
+        metrics = {}
+        rank = dist.get_rank() if dist.is_initialized() else None
+        for name in self.param2name.values():
+            key = get_summary_writer_tag_name(name, tag, rank)
+            if name not in param_tensor or param_tensor[name] is None:
+                continue
+            metrics[key] = param_tensor[name]
+        return metrics
+
+    def generate_xy_metrics(self):
+        actv = {}
+        for fwd_context in self.module_fwd_hook_context_by_module.values():
+            actv.update(fwd_context.actv)
+
+        actv_grad = self.grad_context.actv
+
+        return actv, actv_grad
+
+    def reload_xy(self, xy_distribution=False):
+        self.xy_distribution = xy_distribution
+
+        for handle in self.handles['xy']:
+            handle.remove()
+        self.handles['xy'].clear()
+        self.hook_modules(self.model, self.micro_batch_number)
+        for _, fwd_context in self.module_fwd_hook_context_by_module.items():
+            fwd_context.actv.clear()
+
+    def write_adhoc_check(self, step):
+        TrainerMon.tensor_metrics.flush(self.summary_writer)
+
+    def write_xy_tb(self, step):
+        if not self.xy_distribution:
+            return
+        for _, fwd_context in self.module_fwd_hook_context_by_module.items():
+            if len(fwd_context.actv) == 0:
+                continue
+            self.write_metrics(self.ops, self.summary_writer, fwd_context.actv, step, 'actv')
+            fwd_context.actv.clear()
+        if self.grad_context.actv:
+            self.write_metrics(self.ops, self.summary_writer, self.grad_context.actv, step, 'actv_grad')
+
+    def write_param_tb(self, opt_context):
+        if not self.param_distribution:
+            return
+        self.write_metrics(self.ops, self.summary_writer, opt_context.param_metric, opt_context.step, 'param')
+
+    def write_mv_tb(self, opt_context):
+        if not self.mv_distribution:
+            return
+        self.write_metrics(self.ops, self.summary_writer, opt_context.exp_avg_metric,
+                           opt_context.step, 'exp_avg')
+        self.write_metrics(self.ops, self.summary_writer, opt_context.exp_avg_sq_metric,
+                           opt_context.step, 'exp_avg_sq')
+
+    def write_grad_tb(self, step):
+        if not self.wg_distribution:
+            return
+
+        if self.enable_megatron:
+            self.write_metrics(self.ops, self.summary_writer, self.grad_context.pre, step, 'grad_unreduced')
+        else:
+            self.write_metrics(self.ops, self.summary_writer, self.grad_context.acc_metric, step, 'grad_unreduced')
+        self.write_metrics(self.ops, self.summary_writer, self.grad_context.post, step, 'grad_reduced')
+
+    def hook_optimizer(self, optimizer=None):
+        # in DDP by default use params_have_main_grad
+        def optimizer_pre_step_hook(optimizer, args, kwargs):
+            context = self.optimizer_context[optimizer]
+            if self.opt_ty in MonitorConst.DEEPSPEED_OPT_TY:
+                if context.step == 0:
+                    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.param2name = mv_result.grad
+            else:
+                mv_result = self.mix_precision_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
+            context.param_adam_ratio = mv_result.ratio
+
+            if (self.print_struct and not all(value == {} for value in self.module_struct.values())
+                    and not self.struct_printed):
+                self._smallest_rank_print("> module struct:")
+                self._smallest_rank_print(json.dumps(self.module_struct))
+                self.struct_printed = True
+                if not self.cc_log_only:
+                    raise Exception("exit after first step when print model struct")
+            if self.cc_log_only and context.step > 0:
+                self._smallest_rank_print("> Used communication ops and corresponding stack")
+                self._smallest_rank_print(
+                    json.dumps({k: [i.split(';') for i in v] for k, v in self.cc_logged_stack.items()}))
+                raise Exception("exit after first step when print cc stack")
+
+            self.generate_wgrad_metrics()
+            self.generate_mv_metrics(context)
+            self.generate_param_metrics(context)
+
+            tbtag_tensor_map = {}
+            if self.mg_direction:
+                for param, name in self.param2name.items():
+                    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
+                    if context.step == 0:
+                        same_direction_ratio = torch.tensor(1.)
+                    else:
+                        same_direction_ratio = get_sign_matches(grad, context.param_exp_avg[name])
+                    context.param_mg_direction[name] = same_direction_ratio
+                tbtag_tensor_map.update(self.generate_param_map('mg_direction', context.param_mg_direction))
+
+            metric_dict = {}
+            get_metrics(self.ops, tbtag_tensor_map, self.eps, metric_dict)
+            for cc in self.cc_context.values():
+                cc.aggregate()
+                metric_dict.update(cc.data)
+                cc.reset()
+
+            if not metric_dict:
+                return
+            context.metric_dict = metric_dict
+            return
+
+        def optimizer_post_step_hook(optimizer, args, kwargs):
+            context = self.optimizer_context[optimizer]
+            rank = dist.get_rank() if dist.is_initialized() else None
+
+            if self.anomaly_data_factory:
+                self.anomaly_data_factory.set_call_id(self.param_name_call_id)
+            self.write_xy_tb(context.step)
+            self.write_grad_tb(context.step)
+            self.write_mv_tb(context)
+            self.write_param_tb(context)
+            self.write_adhoc_check(context.step)
+
+            if self.ur_distribution:
+                for param_name, _ in context.param_adam_update.items():
+                    self.update_heatmap_visualizer[param_name].visualize(
+                        get_summary_writer_tag_name(param_name, 'adam_update', rank), context.step, self.summary_writer)
+                for param_name, _ in context.param_adam_ratio.items():
+                    self.ratio_heatmap_visualizer[param_name].visualize(
+                        get_summary_writer_tag_name(param_name, 'adam_ratio', rank), context.step, self.summary_writer)
+
+            if context.metric_dict:
+                self.write_metrics(self.ops, self.summary_writer, context.metric_dict, context.step, 'other')
+            context.metric_dict.clear()
+            context.step += 1
+            if self.anomaly_data_factory:
+                self.anomaly_data_writer.write_detected_json(self.summary_writer.get_anomalies())
+            self.summary_writer.clear_anomalies()
+            self.call_id = 0
+            return
+
+        def patch_step(func, optimizer):
+            def wrapper(*args, **kwargs):
+                optimizer_pre_step_hook(optimizer, args, kwargs)
+                out = func(*args, **kwargs)
+                optimizer_post_step_hook(optimizer, args, kwargs)
+                return out
+
+            return wrapper
+
+        if self.optimizer_hooked:
+            return
+
+        if optimizer:
+            optimizer.__class__.step = patch_step(optimizer.__class__.step, optimizer)
+
+        else:
+            if not self.module_rank_list or (dist.is_initialized() and dist.get_rank() in self.module_rank_list):
+                register_optimizer_step_pre_hook(optimizer_pre_step_hook)
+                register_optimizer_step_post_hook(optimizer_post_step_hook)
+        self.optimizer_hooked = True
+        return
+
+    def _smallest_rank_print(self, msg):
+        if dist.is_initialized():
+            if self.module_rank_list:
+                if dist.get_rank() == min(self.module_rank_list):
+                    logger.info(msg)
+            else:
+                if dist.get_rank() == 0:
+                    logger.info(msg)
+        else:
+            logger.info(msg)
+
+    def _is_target_param(self, param_name, param, prefix):
+        squash_name = prefix + squash_param_name(param_name)
+        name = prefix + param_name
+        for target in self.config['targets'].keys():
+            if param_name.startswith(target) or squash_name.startswith(target) or name.startswith(target):
+                setattr(param, "zero_out_wgrad", True)
+                return True
+
+        return False
+
+    def _register_chunk(self, model_chunk, prefix):
+        for index, (param_name, param) in enumerate(model_chunk.named_parameters()):
+            if not param.requires_grad:
+                continue
+            if self._is_target_param(param_name, param, prefix):
+                name = prefix + squash_param_name(param_name)
+                if name in self.param2name.values():
+                    logger.error(f'same name {name} for different param. Current param is {param_name}. \
+                                    May be error of squash_param_name')
+                    raise Exception("param with same name will be overwritten.")
+                self.param2name[param] = name
+                self.name2param[name] = param
+                self.name2index[name] = index
+
+                if self.tp_group and not param_is_not_tensor_parallel_duplicate(param, self.tp_group):
+                    self.duplicate_param[name] = True
+                if self.dp_group and param_is_data_parallel_duplicate(self.dp_group):
+                    self.duplicate_param[name] = True
+                self.name2tag[name] = {}
+                self.name2tag[name][MonitorConst.PRE_GRAD] = get_summary_writer_tag_name(name, MonitorConst.PRE_GRAD,
+                                                                                         self.rank)
+                self.name2tag[name][MonitorConst.POST_GRAD] = get_summary_writer_tag_name(name, MonitorConst.POST_GRAD,
+                                                                                          self.rank)
+
+    def _register_param_name(self, model):
+        if self.param_registered:
+            return
+
+        if not isinstance(model, list):
+            model = [model]
+
+        if len(model) > 1:
+            self.vpp = True
+            self._smallest_rank_print('vpp enabled')
+
+        for vpp_stage, model_chunk in enumerate(model):
+            prefix = f'{vpp_stage}{MonitorConst.VPP_SEP}'
+            self._register_chunk(model_chunk, prefix)
+
+        self.param_registered = True
+
+    def _is_target_module(self, module_name, targets, vpp_stage):
+        if self.all_xy or self.print_struct:
+            return vpp_stage + squash_param_name(module_name)
+        for pattern in [
+            vpp_stage + squash_param_name(module_name),
+            vpp_stage + module_name,
+        ]:
+            if pattern in targets:
+                return pattern
+        return ""
+
+    def _hook_module(self, target_names, module: torch.nn.Module, vpp_stage=''):
+        if '_modules' not in module.__dict__:
+            # nothing to hook
+            return 0
+
+        def fwd_hook_fun(module, module_input, module_output, name):
+            if is_recomputation():
+                return
+            if module not in self.module_fwd_hook_context_by_module:
+                self.module_fwd_hook_context_by_module[module] = ModuleHookContext(name)
+            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)}
+            if self.print_struct:
+                if context.module_name not in self.module_struct:
+                    self.module_struct[context.module_name] = {}
+                self.module_struct[context.module_name].update(context.struct)
+                return
+            if not module.training:
+                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'])
+            if not context.format_by_arg:
+                return
+            if not context.verified:
+                if not context.ignore_in:
+                    context.focused_in_col = validate_config_spec(context.format_by_arg[MonitorConst.ACTV_IN],
+                                                                  module_input, context.module_name,
+                                                                  MonitorConst.ACTV_IN)
+                context.focused_out_col = validate_config_spec(context.format_by_arg[MonitorConst.ACTV_OUT],
+                                                               module_output, context.module_name,
+                                                               MonitorConst.ACTV_OUT)
+                context.verified = True
+            # expect output be tensor type
+            tbtag_tensor_map = {}
+            if not context.ignore_in:
+                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))
+            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))
+
+            get_metrics(self.ops, tbtag_tensor_map, self.eps, context.actv)
+
+            context.micro_step += 1
+            if context.micro_step == self.micro_batch_number:
+                context.micro_step = 0
+                context.step += 1
+            return
+
+        def bwd_hook_fun(module, input_grad, output_grad):
+            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)}
+            if self.print_struct:
+                if context.module_name not in self.module_struct:
+                    self.module_struct[context.module_name] = {}
+                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'])
+            if not context.format_by_arg:
+                return
+            if not context.verified:
+                if not context.ignore_in:
+                    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.verified = True
+
+            tbtag_tensor_map = {}
+            if not context.ignore_in:
+                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))
+            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))
+
+            if context.micro_step == 0 and context.actvgrad:
+                logger.warning(f"actvgrad context of {context.module_name} is not empty when first micro_step, "
+                               f"maybe something wrong happened. Now clear it.")
+                context.actvgrad.clear()
+
+            get_metrics(self.ops, tbtag_tensor_map, self.eps, self.grad_context.actv)
+
+            context.micro_step += 1
+            if context.micro_step == self.micro_batch_number:
+                context.micro_step = 0
+                context.step += 1
+            return
+
+        if self.backward_only and self.forward_only:
+            logger.warning('not enable backward_only and forward_only simultaneously')
+
+        hooked_count = 0
+        if self.xy_distribution or self.print_struct:
+            for module_name, submodule in module.named_modules():
+                name = self._is_target_module(module_name, target_names, vpp_stage)
+                if not name:
+                    continue
+                if not self.backward_only:
+                    handle = submodule.register_forward_hook(partial(fwd_hook_fun, name=name))
+                    self.handles['xy'].append(handle)
+                if not self.forward_only:
+                    handle = submodule.register_full_backward_hook(bwd_hook_fun)
+                    self.handles['xy'].append(handle)
+                    self.module_bwd_hook_context_by_module[submodule] = ModuleHookContext(name)
+                logger.info_on_rank_0(f"> {name} is monitored successfully")
+                hooked_count += 1
+        return hooked_count
+
+    def _patch_grad_sync(self):
+        def patch_sync(sync_grad_func):
+            def wrapper(bucket):
+                grad_dict = {}
+                for param, name in self.param2name.items():
+                    if param not in bucket.params_list:
+                        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.PRE_GRAD)
+                    if tag is None:
+                        continue
+                    grad_dict[tag] = grad
+                get_metrics(self.ops, grad_dict, self.eps, self.grad_context.pre)
+                out = sync_grad_func(bucket)
+                return out
+
+            return wrapper
+
+        try:
+            from megatron.core.distributed.param_and_grad_buffer import Bucket
+            self.enable_megatron = True
+        except ImportError:
+            self.enable_megatron = False
+
+        if self.enable_megatron:
+            Bucket.start_grad_sync = patch_sync(Bucket.start_grad_sync)  # differ in different megatron version
+        else:
+            self._hook_weights()
+
+    def _hook_weights(self):
+        context = self.grad_context
+
+        @torch.no_grad
+        def param_hook(*args, context_dict, param, key, name):
+            param.micro_step += 1
+            self.param_name_call_id[name] = self.call_id
+            self.call_id += 1
+            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()
+                else:
+                    context_dict[key] = param.grad.clone()
+
+        for param, name in self.param2name.items():
+            key = get_summary_writer_tag_name(name, 'acc_grad', self.rank)
+            setattr(param, 'micro_step', 0)
+            param_tmp = param.expand_as(param)
+            grad_acc = param_tmp.grad_fn.next_functions[0][0]
+            handle = grad_acc.register_hook(
+                partial(param_hook, context_dict=context.acc, param=param, key=key, name=name))
+            self.grad_accs.append(grad_acc)
+            self.handles['wgrads'].append(handle)
+
+        self.weight_hooked = True