Skip to content

A slow StandardRLProgram step shows up only as perf/global_step_time. Telling whether rollouts, training, or weight sync were slow, and which rollouts held up the batch, meant scanning trajectory logs. Log that directly: - #2826

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_997541325

Conversation

@copybara-service

Copy link
Copy Markdown

A slow StandardRLProgram step shows up only as perf/global_step_time. Telling whether rollouts, training, or weight sync were slow, and which rollouts held up the batch, meant scanning trajectory logs. Log that directly:

  • orchestrator/{exposed_generation_time,policy_training_time,reference_logps_time,sampler_agreement_time,checkpoint_time,weight_sync_time}: seconds the train stage waited on each phase of the step. exposed_generation_time, policy_training_time and weight_sync_time match the MLPerf branch's names and definitions, so dashboards compare across branches. Each microbatch's train_step is timed inside its own task, so packing or the next queue wait does not inflate it. These are orchestrator wait times, not device times: the trainer answers train_step once its programs are dispatched, so device time lands in the first later call that blocks on the trainer's results (normally get_metrics, inside policy_training_time), and an async checkpoint's background write is not counted.
  • rollout/status/<status>_frac for every TrajectoryStatus, so a drop in success_rate shows which status took its place.
  • rollout/num_turns_{mean,max} and rollout/turn_length/{mean,p50,p99,max} (tokens sampled per assistant turn). num_turns_mean used to count traj["steps"], which is always empty for Token-mode rollouts, so it was never logged; turns are now counted in conversation_text, and critique_stage no longer stamps the empty steps list nothing reads. The rollout collector records each turn's sampled token count in traj["assistant_turn_lengths"], next to policy_versions. Unlike conversation_masks, which the overlong filter zeroes, these lengths survive for rollouts that were masked out, which are exactly the ones a runaway turn produces.
  • rollout/duration_sec/{p50,p90,p99,max}, rollout/tail_sec and rollout/batch_wall_sec, from the engine.dispatch and worker.rollout lineage events every rollout already carries.
  • One INFO line per step naming the slowest rollouts, for example [rollout] step=36 slowest 3 of 256: traj_deepswe_590_g13 dur=963.2s status=MAX_CONTEXT_LIMIT_REACHED turns=14 max_turn_length=57344 worker=rollout-2 | ....

The helpers live in a new step_telemetry module (StepTimes, RolloutFacts, rollout_health, slowest_rollouts_summary) with its own tests, so rl_program.py only wires them in. Each rollout is read once into RolloutFacts, which both the metrics and the slowest-rollouts line use, and rollout_health returns the per-status counts alongside the metrics.

The status breakdown requires the status that critique_stage stamps on every scored rollout. The turn and duration metrics skip rollouts that lack their inputs, since engine error placeholders carry no conversation and lineage is optional, and are omitted when no rollout has them, so a missing input is never counted as zero.

…me`. Telling whether rollouts, training, or weight sync were slow, and which rollouts held up the batch, meant scanning trajectory logs. Log that directly:

* `orchestrator/{exposed_generation_time,policy_training_time,reference_logps_time,sampler_agreement_time,checkpoint_time,weight_sync_time}`: seconds the train stage waited on each phase of the step. `exposed_generation_time`, `policy_training_time` and `weight_sync_time` match the MLPerf branch's names and definitions, so dashboards compare across branches. Each microbatch's `train_step` is timed inside its own task, so packing or the next queue wait does not inflate it. These are orchestrator wait times, not device times: the trainer answers `train_step` once its programs are dispatched, so device time lands in the first later call that blocks on the trainer's results (normally `get_metrics`, inside `policy_training_time`), and an async checkpoint's background write is not counted.
* `rollout/status/<status>_frac` for every `TrajectoryStatus`, so a drop in `success_rate` shows which status took its place.
* `rollout/num_turns_{mean,max}` and `rollout/turn_length/{mean,p50,p99,max}` (tokens sampled per assistant turn). `num_turns_mean` used to count `traj["steps"]`, which is always empty for Token-mode rollouts, so it was never logged; turns are now counted in `conversation_text`, and `critique_stage` no longer stamps the empty `steps` list nothing reads. The rollout collector records each turn's sampled token count in `traj["assistant_turn_lengths"]`, next to `policy_versions`. Unlike `conversation_masks`, which the overlong filter zeroes, these lengths survive for rollouts that were masked out, which are exactly the ones a runaway turn produces.
* `rollout/duration_sec/{p50,p90,p99,max}`, `rollout/tail_sec` and `rollout/batch_wall_sec`, from the `engine.dispatch` and `worker.rollout` lineage events every rollout already carries.
* One INFO line per step naming the slowest rollouts, for example `[rollout] step=36 slowest 3 of 256: traj_deepswe_590_g13 dur=963.2s status=MAX_CONTEXT_LIMIT_REACHED turns=14 max_turn_length=57344 worker=rollout-2 | ...`.

The helpers live in a new `step_telemetry` module (`StepTimes`, `RolloutFacts`, `rollout_health`, `slowest_rollouts_summary`) with its own tests, so `rl_program.py` only wires them in. Each rollout is read once into `RolloutFacts`, which both the metrics and the slowest-rollouts line use, and `rollout_health` returns the per-status counts alongside the metrics.

The status breakdown requires the `status` that `critique_stage` stamps on every scored rollout. The turn and duration metrics skip rollouts that lack their inputs, since engine error placeholders carry no conversation and lineage is optional, and are omitted when no rollout has them, so a missing input is never counted as zero.

PiperOrigin-RevId: 997541325

This branch was successfully deployed

1 active deployment
testing — 03ff845a Deployed Oct 12, 2026 by copybara-service[bot] via tunix_tpu_unit_tests / run_core #11280
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.

2 participants