[Enhance] Add type hint and fix typos and code styles

Author: nil0x9

xtuner/v1/utils/internal_metrics.py

@@ -4,6 +4,8 @@
 import torch
 from torch import nn
 import torch.distributed as dist
+from torch.utils.hooks import RemovableHandle
+
 from xtuner.v1.module import (
     RMSNorm,
     MultiHeadAttention,
@@ -35,7 +37,8 @@ class InternalMetricsRecorder:
     def __init__(self, engine: TrainEngine):
         self.model = engine.model
         self.intra_layer_micro_batch = engine.intra_layer_micro_batch
-        self.hooks: list[Any] = []
+        self.hooks: list[RemovableHandle] = []
+        # TODO: refactor with TypeDict
         self.metrics: dict[str, dict[str, Any]] = dict[str, dict[str, Any]](
             weight_rms=dict[str, Any](),
             maxvio=dict[str, Any](),
@@ -46,14 +49,11 @@ def __init__(self, engine: TrainEngine):
             attn_max_logits=dict[str, Any](),
         )
 
-    def register_weight_rms_hook(self, module, layer_name=None):
+    def register_weight_rms_hook(self, module: nn.Module, layer_name: str):
         """
         Register weight RMS hook as a pre-forward hook, as at this point, the parameters are should be 
         all-gathered into current rank.
         """
-        if layer_name is None:
-            layer_name = f"layer_{len(self.weight_rms_dict)}"
-
         def hook(module, args, kwargs=None):
             if layer_name in self.metrics['weight_rms']: # only calculate before the first batch
                 return
@@ -67,10 +67,13 @@ def hook(module, args, kwargs=None):
                 rms = torch.sqrt(l2_norm / total_params)
                 self.metrics['weight_rms'][layer_name] = rms
 
-        hook_handle = module.register_forward_pre_hook(hook)
+        hook_handle: RemovableHandle = module.register_forward_pre_hook(hook)
         self.hooks.append(hook_handle)
 
-    def register_attn_extra_info_hook(self, module, layer_name=None):
+    def register_attn_extra_info_hook(self, module: nn.Module, layer_name: str):
+        """
+        Register attention extra info hook as a forward hook
+        """
         def hook(module, input, output):
             extra_info = output[1]
             if extra_info.get("softmax_lse", None) is not None:
@@ -88,7 +91,7 @@ def hook(module, input, output):
                     prev_logits_max = self.metrics["attn_max_logits"][layer_name]
                     self.metrics["attn_max_logits"][layer_name] = max(prev_logits_max, extra_info["attn_logits"].max())
 
-        hook_handle = module.register_forward_hook(hook)
+        hook_handle: RemovableHandle = module.register_forward_hook(hook)
         self.hooks.append(hook_handle)
 
     @torch.no_grad()
@@ -155,26 +158,26 @@ def get_metrics(self, data_batches: list[ModelItem]):
             self.metrics["maxvio"]["total"] = maxvio
             self.metrics["drop_ratio"]["total"] = drop_ratio
 
-        if len(router_logits_max) > 0:
+        if router_logits_max:
             for layer_name, router_logits_list in router_logits_max.items():
                 # [bsz/intra_layer_micro_batch, ]
                 local_router_logits_max = torch.max(torch.stack(router_logits_list))
                 dist.all_reduce(local_router_logits_max, op=dist.ReduceOp.MAX)
                 self.metrics["router_logits_max"][layer_name] = local_router_logits_max.item()
 
-        if len(router_logits_mean) > 0:
+        if router_logits_mean:
             for layer_name, router_logits_list in router_logits_mean.items():
                 # [bsz/intra_layer_micro_batch, ]
                 local_router_logits_mean = torch.mean(torch.stack(router_logits_list))
                 dist.all_reduce(local_router_logits_mean.div_(dist.get_world_size()), op=dist.ReduceOp.SUM)
                 self.metrics["router_logits_mean"][layer_name] = local_router_logits_mean.item()
 
-        if len(self.metrics["attn_max_lse"]) > 0:
+        if self.metrics["attn_max_lse"]:
             for layer_name, local_attn_max_lse in self.metrics["attn_max_lse"].items():
                 dist.all_reduce(local_attn_max_lse, op=dist.ReduceOp.MAX)
                 self.metrics["attn_max_lse"][layer_name] = local_attn_max_lse.item()
 
-        if len(self.metrics["attn_max_logits"]) > 0:
+        if self.metrics["attn_max_logits"]:
             for layer_name, local_attn_max_logits in self.metrics["attn_max_logits"].items():
                 dist.all_reduce(local_attn_max_logits, op=dist.ReduceOp.MAX)
                 self.metrics["attn_max_logits"][layer_name] = local_attn_max_logits.item()
@@ -187,8 +190,7 @@ def __enter__(self):
                 self.register_attn_extra_info_hook(module, self._clean_module_name(name))
             if isinstance(module, RMS_NORM_MONITOR_MODULES):
                 self.register_weight_rms_hook(module, self._clean_module_name(name))
-            else:
-                pass
+
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):