diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py index a8cc4f6d1..acde3320d 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,8 @@ 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: + 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} %)" )