[Enhance] drop rms_norm hook impl. to avoid precision problem (fp8)

[Enhance] refactor internal metrics to use TypedDict

Author: nil0x9

xtuner/v1/float8/float8_ops.py

@@ -165,56 +165,4 @@ def float8_desugar_data_and_scale_op(aten_op, args, kwargs=None):
         args[0]._orig_dtype,
         args[0]._scaling_granularity,
         args[0]._group_size,
-    )
-
-
-@implements(
-    [
-        aten._to_copy.default,
-    ]
-)
-def float8_to_copy(aten_op, args, kwargs=None):
-    if kwargs is None:
-        kwargs = {}
-
-    dtype = kwargs.get("dtype", torch.float32)
-
-    tensor = args[0]
-    scaling_granularity = tensor._scaling_granularity
-
-    dequantized: torch.Tensor = None
-
-    if scaling_granularity == ScalingGranularity.TENSORWISE:
-        dequantized = tensor._data.to(torch.float32) * tensor._scale.to(torch.float32)
-
-    elif scaling_granularity == ScalingGranularity.BLOCKWISE:
-        from xtuner.v1.float8.triton_kernels import per_block_dequant_gemm
-
-        if tensor._data.ndim == 2:
-            dequantized = per_block_dequant_gemm(tensor._data, tensor._scale, block_size=tensor._group_size)
-
-        else:
-            raise NotImplementedError(
-                f"{aten_op} with {scaling_granularity} scaling granularity is not implemented. "
-            )
-    elif scaling_granularity == ScalingGranularity.TILEWISE:
-        # For tilewise, scale is per-tile (1x128)
-        original_shape = tensor._data.shape
-        data_flat = tensor._data.view(-1, original_shape[-1])
-        scale_flat = tensor._scale.view(-1, tensor._scale.shape[-1])
-
-        # Expand scale to match data: each scale applies to group_size elements
-        group_size = tensor._group_size
-        num_groups = data_flat.shape[-1] // group_size
-        scale_expanded = scale_flat[:, :num_groups].unsqueeze(-1).expand(-1, -1, group_size)
-        scale_expanded = scale_expanded.contiguous().view(-1, data_flat.shape[-1])
-
-        # Dequantize: data * scale
-        dequantized_flat = data_flat.to(torch.float32) * scale_expanded.to(torch.float32)
-        dequantized = dequantized_flat.view(*original_shape)
-    else:
-        raise NotImplementedError(
-            f"{aten_op} with {scaling_granularity} scaling granularity is not supported. "
-        )
-
-    return dequantized.to(dtype)
+    )

xtuner/v1/train/trainer.py

@@ -47,6 +47,7 @@
     log_format,
     record_git_info,
     InternalMetricsRecorder,
+    InternalMetrics,
 )
 from xtuner.v1.utils.device import get_device, get_torch_device_module
 
@@ -609,7 +610,7 @@ def fit(self):
         self._exp_tracker.close()
         self.logger.info(f"Training finished in {time.time() - train_begin:.2f} seconds")
 
-    def _maybe_check_model_internal_metrics(self, data_batches: list[ModelItem]) -> dict[str, float] | None:
+    def _maybe_check_model_internal_metrics(self, data_batches: list[ModelItem]) -> InternalMetrics | None:
         if self._internal_metrics_interval is None:
             return None
 
@@ -1183,7 +1184,7 @@ def _log_step(
         if internal_metrics is None:
             internal_metrics = {}
         else:
-            internal_metrics = _flatten_dict(internal_metrics)
+            internal_metrics = _flatten_nested_metrics(internal_metrics)
 
         self.logger.info(
             f"Epoch {self._cur_epoch} Step {self.cur_step}/{self.total_step} data_time: {data_time:.4f} lr: {lr:.6e} time: {step_time:.4f} "
@@ -1437,14 +1438,15 @@ def _setup_env(self):
         logger.info(log_str)
 
 
-def _flatten_dict(d: dict, parent_key: str = '', sep: str = '/') -> dict:
+def _flatten_nested_metrics(metrics: InternalMetrics, sep: str = '/') -> dict:
     items = []
-    for k, v in d.items():
-        new_key = f"{parent_key}{sep}{k}" if parent_key else k
-        if isinstance(v, dict):
-            items.extend(_flatten_dict(v, new_key, sep=sep).items())
-        elif isinstance(v, torch.Tensor):
-            items.append((new_key, v.item()))
+    for name, sub_metrics in metrics.items():
+        if isinstance(sub_metrics, dict):
+            for k, v in sub_metrics.items():
+                if isinstance(v, (float, int)):
+                    items.append((f"{name}{sep}{k}", v))
+                else:
+                    raise ValueError(f"Unsupported metric value type: expected float or int, but got {type(v)}")
         else:
-            items.append((new_key, v))
+            raise ValueError(f"Unsupported metric type for internal metrics: expected dict, but got {type(sub_metrics)}")
     return dict(items)

xtuner/v1/utils/__init__.py

@@ -14,7 +14,7 @@
 from .type_helper import copy_method_signature, copy_signature
 from .update_weights_utils import monkey_unpatch_torch_reductions
 
-from .internal_metrics import InternalMetricsRecorder
+from .internal_metrics import InternalMetricsRecorder, InternalMetrics
 
 IGNORE_INDEX = -100
 
@@ -47,4 +47,5 @@
     "IGNORE_INDEX",
     "monkey_unpatch_torch_reductions",
     "InternalMetricsRecorder",
+    "InternalMetrics",
 ]

xtuner/v1/utils/internal_metrics.py

@@ -17,6 +17,20 @@
 from xtuner.v1.model import MoE
 from xtuner.v1.model.base import ModelItem
 from xtuner.v1.engine.train_engine import TrainEngine
+from xtuner.v1.utils.grad_norm import group_tensors_by_device_mesh_and_placements, cal_total_norm
+
+from typing_extensions import TypedDict
+
+
+class InternalMetrics(TypedDict):
+    weight_rms: dict[str, float]
+    maxvio: dict[str, float]
+    drop_ratio: dict[str, float]
+    router_logits_max: dict[str, float]
+    router_logits_mean: dict[str, float]
+    attn_max_lse: dict[str, float]
+    attn_max_logits: dict[str, float]
+
 
 RMS_NORM_MONITOR_MODULES = (
     nn.Embedding,
@@ -34,36 +48,32 @@ def __init__(self, engine: TrainEngine):
         self.intra_layer_micro_batch = engine.intra_layer_micro_batch
         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](),
-            drop_ratio=dict[str, Any](),
-            router_logits_max=dict[str, Any](),
-            router_logits_mean=dict[str, Any](),
-            attn_max_lse=dict[str, Any](),
-            attn_max_logits=dict[str, Any](),
-        )
-
-    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.
-        """
-        def hook(module, args, kwargs=None):
-            if layer_name in self.metrics['weight_rms']: # only calculate before the first batch
-                return
-            l2_norm = 0.0
-            total_params = 0
-            for param in module.parameters():
-                if param.requires_grad:
-                    l2_norm += torch.norm(param.detach().float(), p=2) ** 2
-                    total_params += param.numel()
-            if total_params > 0:
-                rms = torch.sqrt(l2_norm / total_params)
-                self.metrics['weight_rms'][layer_name] = rms
-
-        hook_handle: RemovableHandle = module.register_forward_pre_hook(hook)
-        self.hooks.append(hook_handle)
+        self.metrics: InternalMetrics = {
+            "weight_rms": {},
+            "maxvio": {},
+            "drop_ratio": {},
+            "router_logits_max": {},
+            "router_logits_mean": {},
+            "attn_max_lse": {},
+            "attn_max_logits": {},
+        }
+        self.attn_max_lse: dict[str, torch.Tensor] = {}
+        self.attn_max_logits: dict[str, torch.Tensor] = {}
+
+    def calculate_module_weight_rms(self, module: nn.Module, layer_name: str, dtype: torch.dtype = torch.float32):
+        all_params = [param for param in module.parameters() if param.requires_grad]
+        if not all_params:
+            return
+        grouped_params = group_tensors_by_device_mesh_and_placements(all_params)
+        total_norms = []
+        total_numel = 0
+        for params in grouped_params.values():
+            total_norm = cal_total_norm(params, norm_type=2.0, foreach=True, dtype=dtype)
+            total_norms.append(total_norm)
+            total_numel += sum(p.numel() for p in params)
+        param_l2_norm = torch.linalg.vector_norm(torch.stack(total_norms), ord=2.0, dtype=dtype)
+        param_rms = param_l2_norm / total_numel**0.5
+        self.metrics['weight_rms'][layer_name] = param_rms.item()
 
     def register_attn_extra_info_hook(self, module: nn.Module, layer_name: str):
         """
@@ -72,19 +82,19 @@ def register_attn_extra_info_hook(self, module: nn.Module, layer_name: str):
         def hook(module, input, output):
             extra_info = output[1]
             if extra_info.get("softmax_lse", None) is not None:
-                if layer_name not in self.metrics["attn_max_lse"]:
+                if layer_name not in self.attn_max_lse:
                     # original shape: [n_head, seq]
-                    self.metrics["attn_max_lse"][layer_name] = extra_info["softmax_lse"].max()
+                    self.attn_max_lse[layer_name] = extra_info["softmax_lse"].max()
                 else:
-                    prev_lse_max = self.metrics["attn_max_lse"][layer_name]
-                    self.metrics["attn_max_lse"][layer_name] = max(prev_lse_max, extra_info["softmax_lse"].max())
+                    prev_lse_max = self.attn_max_lse[layer_name]
+                    self.attn_max_lse[layer_name] = max(prev_lse_max, extra_info["softmax_lse"].max())
             if extra_info.get("attn_logits", None) is not None:
-                if layer_name not in self.metrics["attn_max_logits"]:
+                if layer_name not in self.attn_max_logits:
                     # original shape: [b, n_head, seq, seq]
-                    self.metrics["attn_max_logits"][layer_name] = extra_info["attn_logits"].max()
+                    self.attn_max_logits[layer_name] = extra_info["attn_logits"].max()
                 else:
-                    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())
+                    prev_logits_max = self.attn_max_logits[layer_name]
+                    self.attn_max_logits[layer_name] = max(prev_logits_max, extra_info["attn_logits"].max())
 
         hook_handle: RemovableHandle = module.register_forward_hook(hook)
         self.hooks.append(hook_handle)
@@ -150,8 +160,8 @@ def get_metrics(self, data_batches: list[ModelItem]):
                 {f"layer{idx}": maxvio_all_layers[idx].item() for idx in range(max_load_i.shape[0])}
             )
             maxvio = maxvio_all_layers.mean()
-            self.metrics["maxvio"]["total"] = maxvio
-            self.metrics["drop_ratio"]["total"] = drop_ratio
+            self.metrics["maxvio"]["total"] = maxvio.item()
+            self.metrics["drop_ratio"]["total"] = drop_ratio.item()
 
         if router_logits_max:
             for layer_name, router_logits_list in router_logits_max.items():
@@ -168,12 +178,12 @@ def get_metrics(self, data_batches: list[ModelItem]):
                 self.metrics["router_logits_mean"][layer_name] = local_router_logits_mean.item()
 
         if self.metrics["attn_max_lse"]:
-            for layer_name, local_attn_max_lse in self.metrics["attn_max_lse"].items():
+            for layer_name, local_attn_max_lse in self.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 self.metrics["attn_max_logits"]:
-            for layer_name, local_attn_max_logits in self.metrics["attn_max_logits"].items():
+        if self.attn_max_logits:
+            for layer_name, local_attn_max_logits in self.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()
 
@@ -184,7 +194,7 @@ def __enter__(self):
             if isinstance(module, ATTENTION_CLS):
                 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))
+                self.calculate_module_weight_rms(module, self._clean_module_name(name), dtype=torch.float32)
 
         return self