[Enhance] internal_metrics: rm loss_ctx and disable intra-layer bsz for dummy fwd

Author: nil0x9

xtuner/v1/utils/internal_metrics.py

@@ -71,7 +71,6 @@ class InternalMetricsRecorder:
     def __init__(self, internal_metrics_cfg: InternalMetricsConfig, engine: TrainEngine):
         self.internal_metrics_cfg = internal_metrics_cfg
         self.model = engine.model
-        self.intra_layer_micro_batch = engine.intra_layer_micro_batch
         self.hooks: list[RemovableHandle] = []
         self._attn_monitor_type: str | None = None
         self.attn_max_lse: dict[str, torch.Tensor] = {}
@@ -169,25 +168,10 @@ def pop_metrics(self, data_batches: list[ModelItem]):
 
         # do dummy forward to get metrics
         if self.need_dummy_forward:
-            for i in range(0, len(data_batches), self.intra_layer_micro_batch):
-                data_batch = data_batches[i : i + self.intra_layer_micro_batch]
-                seq_ctx_list = []
-                loss_ctx_list = []
-                for data in data_batch:
-                    seq_ctx = data["seq_ctx"]
-                    loss_ctx = data["loss_ctx"]
-                    seq_ctx_list.append(seq_ctx)
-                    loss_ctx_list.append(loss_ctx)
-                if self.intra_layer_micro_batch == 1:
-                    output = self.model(seq_ctx=seq_ctx_list[0], loss_ctx=loss_ctx_list[0], **additional_kwargs)
-                else:
-                    # although we dont need loss at this point, we still need loss_ctx for micro-batch forward
-                    output = self.model(
-                        seq_ctx=seq_ctx_list,
-                        loss_ctx=loss_ctx_list,
-                        **additional_kwargs,
-                    )
-
+            for i in range(0, len(data_batches)):
+                data_batch = data_batches[i]
+                seq_ctx = data_batch["seq_ctx"]
+                output = self.model(seq_ctx=seq_ctx, loss_ctx=None, **additional_kwargs)
                 if (
                     self.internal_metrics_cfg.monitor_moe_load_balance_stats
                     and (cur_tokens_per_expert := output.get("tokens_per_expert_global")) is not None