Repository navigation
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
Open
copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
…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
copybara-service
Bot
requested review from
abheesht17,
hgao327,
jiangyangmu,
lc5211,
s-noghabi,
sizhit2,
tianshub and
wang2yn84
as code owners
October 12, 2026 06:24
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
A slow
StandardRLProgramstep shows up only asperf/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_timeandweight_sync_timematch the MLPerf branch's names and definitions, so dashboards compare across branches. Each microbatch'strain_stepis 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 answerstrain_steponce its programs are dispatched, so device time lands in the first later call that blocks on the trainer's results (normallyget_metrics, insidepolicy_training_time), and an async checkpoint's background write is not counted.rollout/status/<status>_fracfor everyTrajectoryStatus, so a drop insuccess_rateshows which status took its place.rollout/num_turns_{mean,max}androllout/turn_length/{mean,p50,p99,max}(tokens sampled per assistant turn).num_turns_meanused to counttraj["steps"], which is always empty for Token-mode rollouts, so it was never logged; turns are now counted inconversation_text, andcritique_stageno longer stamps the emptystepslist nothing reads. The rollout collector records each turn's sampled token count intraj["assistant_turn_lengths"], next topolicy_versions. Unlikeconversation_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_secandrollout/batch_wall_sec, from theengine.dispatchandworker.rolloutlineage events every rollout already carries.[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_telemetrymodule (StepTimes,RolloutFacts,rollout_health,slowest_rollouts_summary) with its own tests, sorl_program.pyonly wires them in. Each rollout is read once intoRolloutFacts, which both the metrics and the slowest-rollouts line use, androllout_healthreturns the per-status counts alongside the metrics.The status breakdown requires the
statusthatcritique_stagestamps 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.