From c9a13beb481fc41d0354a3d62e4ee9afa3e90ca2 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 15 Jul 2026 16:56:27 -0400 Subject: [PATCH 1/2] Split data-loader wait from preprocessing in schedule metrics #564 timed `next(context.data_iterator)` in `_get_forward_input`, but that iterator is the `_preprocess_data` generator, so the metric conflated the real data-loader block with CPU preprocessing (share_batch_data / preprocess_batch). In the RL path this reads near zero because the loader is drained into a buffer before the schedule runs. Move the timing into `_preprocess_data`: time only the raw `next(data_iterator)` pulls as `data_wait_time_ms` (input starvation), and the generator's own active runtime -- paused across each yield to exclude schedule compute between micro-batches -- as a new `data_preprocessing_time_ms`. Both are logged to the terminal. Drop the now-redundant `data_batch_warn_time_ms` warning and its config field, since the wait is logged explicitly. Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/schedule/config.py | 6 ------ fast_llm/engine/schedule/runner.py | 26 +++++++++++++++----------- fast_llm/logging.py | 2 ++ 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/fast_llm/engine/schedule/config.py b/fast_llm/engine/schedule/config.py index 29720b90b..54ae03fcc 100644 --- a/fast_llm/engine/schedule/config.py +++ b/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, diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py index acde3320d..3af49cf2b 100644 --- a/fast_llm/engine/schedule/runner.py +++ b/fast_llm/engine/schedule/runner.py @@ -162,6 +162,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 +327,15 @@ 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. + wait_start = time.perf_counter() if measure_time else 0.0 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 +369,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 +442,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: diff --git a/fast_llm/logging.py b/fast_llm/logging.py index 20f2c4ecc..9e809e803 100644 --- a/fast_llm/logging.py +++ b/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} %)" ) From d0d229fa3cbdd8f1566ea864db8eb1d5bd9a4d26 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 15 Jul 2026 17:19:07 -0400 Subject: [PATCH 2/2] Review fixes: drop now-dead module logger, align wait_start guard Removing the slow-load warning left the module `logger` (and `import logging`) with no consumer -- drop both. Guard `wait_start` with `if measure_time:` to match its sibling `preprocess_start` instead of a ternary fallback. Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/schedule/runner.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py index 3af49cf2b..58c641191 100644 --- a/fast_llm/engine/schedule/runner.py +++ b/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: @@ -331,7 +328,8 @@ def _preprocess_data( # 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. - wait_start = time.perf_counter() if measure_time else 0.0 + 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