Skip to content

Log data-loader wait time as an input-starvation metric#564

Merged
jlamypoirier merged 2 commits into
mainfrom
jlp_data_wait
Jul 10, 2026
Merged

Log data-loader wait time as an input-starvation metric#564
jlamypoirier merged 2 commits into
mainfrom
jlp_data_wait

Conversation

@jlamypoirier

Copy link
Copy Markdown
Collaborator

Claude Opus 4.8 note: extracted from PR #553 (RL diagnostics umbrella).

What

The schedule runner already times how long it blocks on the data loader, but only warned past a threshold and discarded the value. This accumulates 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.

Details

  • Initialized to 0.0 on logging steps so "no wait" reads as 0 rather than a gap in the series.
  • Accumulated across the step's micro-batches in the data-loading path.
  • Flows to the logged metrics via run_step's returned train_metrics, already spread into the trainer's per-step metrics dict.

Single-file, 6-line change to runner.py. This is the last piece of #553; the other parts landed as #555, #559, #560, #561, #562.

🤖 Generated with Claude Code

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 <noreply@anthropic.com>
…g 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 <noreply@anthropic.com>
@jlamypoirier
jlamypoirier merged commit 105cc56 into main Jul 10, 2026
3 checks passed
@jlamypoirier
jlamypoirier deleted the jlp_data_wait branch July 10, 2026 16:31
jlamypoirier added a commit that referenced this pull request Jul 16, 2026
Brings main (#562/#564/#567) into the branch. The one semantic conflict was
policy-gradient metrics: main #562 added reward/staleness data metrics
(`_register_data_metrics`) to the standalone GRPO/GSPO `_forward_backward`,
while this branch reworked registration to compute policy metrics in the fused
kernel and register them through `register_combinable_extras`. Kept the
fused-metrics design and grafted the data metrics onto every registration path:

- `register_combinable_extras` gains `split_index` (base + GRPO/GSPO overrides
  + 3 call sites) so the compiled and monolithic paths can register data metrics.
- GRPO triton-standalone and both `register_combinable_extras` overrides now
  call `_register_data_metrics`; GSPO standalone already did (auto-merged).
- `_policy_metric_definitions` folds in `_data_metric_definitions()`.

`test_lm_head` now supplies reward/model_version/documents_seen and asserts the
data-metric family via a `reference_data_metrics` helper.

CPU validation: test_lm_head + test_lm_losses = 722 passed, 75 skipped
(triton/GPU-only skipped locally).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant