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..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: @@ -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: 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} %)" )