Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions fast_llm/engine/schedule/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions fast_llm/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@
"skipped_iters",
"nan_iters",
"step_time_average_ms",
"data_wait_time_ms",
"remaining_time",
"completion_time",
"percent_done",
Expand All @@ -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} %)"
)
Expand Down
Loading