Split data-loader wait from preprocessing in schedule metrics

fast_llm/engine/schedule/config.py

@@ -39,12 +39,6 @@ class ScheduleConfig(Config):
     data_overlap: bool = Field(
         default=True, desc="Overlap the data-parallel network communication.", hint=FieldHint.testing
     )
-    data_batch_warn_time_ms: float = Field(
-        default=1000.0,
-        desc="Warn if a batch takes too long to load.",
-        hint=FieldHint.optional,
-        valid=check_field(Assert.gt, 0),
-    )
     # Enable cpu throttling to avoid lag spikes, see https://arxiv.org/pdf/2211.05953.pdf, appendix D.2.
     throttle_cpu: bool = Field(
         default=True,

fast_llm/engine/schedule/runner.py

@@ -1,7 +1,6 @@
 import collections
 import contextlib
 import dataclasses
-import logging
 import time
 import typing
 
@@ -22,8 +21,6 @@
 from fast_llm.logging import log_memory_usage
 from fast_llm.utils import Assert
 
-logger = logging.getLogger(__name__)
-
 
 @dataclasses.dataclass()
 class BatchContext:
@@ -162,6 +159,7 @@ def run_step(
         if metrics is not None:
             # Always present on logging steps so "no wait" shows as 0 rather than a gap.
             metrics["data_wait_time_ms"] = 0.0
+            metrics["data_preprocessing_time_ms"] = 0.0
         # Set the context.
         context = BatchContext(
             iteration=iteration,
@@ -326,7 +324,16 @@ def _preprocess_data(
             if context.schedule.phase.is_training
             else None
         )
+        measure_time = context.metrics is not None
+        # Time blocked on the data loader (input starvation), kept separate from the CPU
+        # preprocessing below. Preprocessing runs interleaved with the schedule's compute, so its
+        # clock is paused across each yield to exclude the compute happening between micro-batches.
+        if measure_time:
+            wait_start = time.perf_counter()
         model_inputs = [next(data_iterator) for _ in range(self._config.sequential_micro_batches)]
+        if measure_time:
+            context.metrics["data_wait_time_ms"] += (time.perf_counter() - wait_start) * 1000
+            preprocess_start = time.perf_counter()
         model_inputs[0][0].share_batch_data(
             [model_input for model_inputs_ in model_inputs for model_input in model_inputs_], self._distributed
         )
@@ -360,7 +367,11 @@ def _preprocess_data(
                         device=self._distributed.device if self._stages_owned[-1] else "meta",
                     )
                 context.batch[data_index] = kwargs
+                if measure_time:
+                    context.metrics["data_preprocessing_time_ms"] += (time.perf_counter() - preprocess_start) * 1000
                 yield
+                if measure_time:
+                    preprocess_start = time.perf_counter()
 
     def _restore(self, context: BatchContext, step: Step) -> None:
         if step.restore_launch:
@@ -429,17 +440,8 @@ def _backward(self, context: BatchContext, step: Step) -> torch.Tensor:
         return input_grad
 
     def _get_forward_input(self, context: BatchContext, step: Step) -> torch.Tensor:
-        if step.index not in context.batch:
-            start_time = time.perf_counter()
-
-            while step.index not in context.batch:
-                next(context.data_iterator)
-
-            data_time = (time.perf_counter() - start_time) * 1000
-            if context.metrics is not None:
-                context.metrics["data_wait_time_ms"] += data_time
-            if data_time > self._config.data_batch_warn_time_ms:
-                logger.warning(f"Data loading took {data_time:,.2f} ms")
+        while step.index not in context.batch:
+            next(context.data_iterator)
         return context.inputs.pop(step.global_index).detach().requires_grad_(step.stage != 0)
 
     def _send(self, context: BatchContext, step: Step, output: torch.Tensor) -> None:

fast_llm/logging.py

@@ -66,6 +66,7 @@
     "nan_iters",
     "step_time_average_ms",
     "data_wait_time_ms",
+    "data_preprocessing_time_ms",
     "remaining_time",
     "completion_time",
     "percent_done",
@@ -79,6 +80,7 @@
     " | nan iterations: {nan_iters:3.0f}"
     " | average step time {step_time_average_ms:.2f} ms"
     " | data wait: {data_wait_time_ms:.2f} ms"
+    " | data preprocessing: {data_preprocessing_time_ms:.2f} ms"
     " | remaining {remaining_time} "
     " | completion {completion_time} ({percent_done:.2f} %)"
 )