From 807cc6002a09e2905649d3951a505f3e4c57843b Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 7 Jul 2026 14:21:16 -0400 Subject: [PATCH 1/2] Log data-loader wait time as an input-starvation metric The schedule runner already timed how long it blocked on the data loader but only warned past a threshold and discarded the value. Accumulate it into the step metrics as `data_wait_time_ms` (0 when never starved), so input starvation is visible in the training logs and W&B. Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/schedule/runner.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py index a8cc4f6d1..019a52fee 100644 --- a/fast_llm/engine/schedule/runner.py +++ b/fast_llm/engine/schedule/runner.py @@ -159,6 +159,9 @@ def run_step( assert self._support_training metrics = {} if return_metrics else None + 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 # Set the context. context = BatchContext( iteration=iteration, @@ -433,6 +436,9 @@ def _get_forward_input(self, context: BatchContext, step: Step) -> torch.Tensor: next(context.data_iterator) data_time = (time.perf_counter() - start_time) * 1000 + if context.metrics is not None: + # Time the trainer spent blocked waiting for the data loader — how input-starved it is. + context.metrics["data_wait_time_ms"] = context.metrics.get("data_wait_time_ms", 0.0) + data_time if data_time > self._config.data_batch_warn_time_ms: logger.warning(f"Data loading took {data_time:,.2f} ms") return context.inputs.pop(step.global_index).detach().requires_grad_(step.stage != 0) From 9e4981a54bca0b3fe3c9b9d840a2c2d5f55f9830 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 12:17:33 -0400 Subject: [PATCH 2/2] Review fixes: drop dead default, format data_wait_time_ms like sibling metrics The metric is pre-seeded to 0.0 on logging steps, so the accumulation no longer needs a `.get` fallback. Format it as `.2f ms` in the training log line to match step_time_ms / step_time_average_ms. Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/schedule/runner.py | 3 +-- fast_llm/logging.py | 2 ++ 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py index 019a52fee..acde3320d 100644 --- a/fast_llm/engine/schedule/runner.py +++ b/fast_llm/engine/schedule/runner.py @@ -437,8 +437,7 @@ def _get_forward_input(self, context: BatchContext, step: Step) -> torch.Tensor: data_time = (time.perf_counter() - start_time) * 1000 if context.metrics is not None: - # Time the trainer spent blocked waiting for the data loader — how input-starved it is. - context.metrics["data_wait_time_ms"] = context.metrics.get("data_wait_time_ms", 0.0) + data_time + 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") return context.inputs.pop(step.global_index).detach().requires_grad_(step.stage != 0) diff --git a/fast_llm/logging.py b/fast_llm/logging.py index 2619883d6..20f2c4ecc 100644 --- a/fast_llm/logging.py +++ b/fast_llm/logging.py @@ -65,6 +65,7 @@ "skipped_iters", "nan_iters", "step_time_average_ms", + "data_wait_time_ms", "remaining_time", "completion_time", "percent_done", @@ -77,6 +78,7 @@ " | skipped iterations: {skipped_iters:3.0f}" " | nan iterations: {nan_iters:3.0f}" " | average step time {step_time_average_ms:.2f} ms" + " | data wait: {data_wait_time_ms:.2f} ms" " | remaining {remaining_time} " " | completion {completion_time} ({percent_done:.2f} %)" )