Rework AsyncGRPO metrics metrics and docs

[AmineDiro]

What does this PR do?

Async GRPO has more places to be slow than sync GRPO: generation, scoring, the queue between the worker and the trainer, weight sync.

The current metric set can't tell you which one is the problem. This PR restructures the metrics around where a number is measured, fixes the reductions, and adds a "Logged metrics" section to the docs explaining what each one means and how to read them together. A rollout passes through several stages before it becomes a gradient:

Dataset row = PROMPT (message list + reward kwargs)
 └─ GROUP                  (1 prompt, `num_generations` rollouts, 1 advantage baseline)
     └─ ROLLOUT            (1 conversation, keyed by `rollout_id`)
         ├─ TURN           (1 vLLM /v1/completions call + tool messages fed back)
         │   └─ TurnRecord (prompt_ids, output_ids, output_log_probs)
         └─ reconcile      (`_chain_to_sequences` classifies drift per turn: CLEAN/REALIGN/FORK)
             └─ SEQUENCE   (≥1 per rollout; new one per fork; dropped if no trained token)
                 └─ SAMPLE (Sequence + group advantage + reward + metrics)

════════════ Trainer process boundary: `rollout_buffer` (mp.Queue) ════════════

                 └─ SAMPLE (Pulled 1 at a time; dropped if staleness > `max_staleness`)
                     └─ ROW            (Planner assigns it to one of `dp` rows, Σ Lᵢ²-balanced)
                         └─ MICRO-BATCH (`dp` rows, one per rank)
                             └─ PACKED ROW (1 concat sequence, `position_ids` reset per sample)
                                 └─ FORWARD (`compute_loss`, bs=1, inter-rank padding stripped)
                     └─ OPTIMIZER STEP (from `grad_accum` micro-batches)

So I added more metrics based on these entities:

namespace	entity it counts
rollout/	represents one full conversation: its turns, its forks, how long it took to generate
completions/	what the model (vLLM) generated for one prompt
tools/	tool calls metrics
sample/	represents one training sample, as it arrives in the queue
batch/	one micro-batch or one optimizer step (built from one or multiple samples)
perf/	measured seconds and FLOPs

Renames

Anyone with a dashboard on this trainer will need to update it:

before	after
training_tok/s	perf/forwarded_tok_s_fwd_bwd, perf/forwarded_tok_s_wall_clock, perf/trained_tok_s_wall_clock
step_time	perf/fwd_bwd_s
iteration_time_s	perf/step_s
forward_time_s	perf/fwd_s
weight_sync_time_s	perf/weight_sync_s (+ _pause_s, _barrier_s, _transfer_s)
queue_wait_time_s	split into perf/rollout_wait_s and sample/time_in_queue_s
generation_tok_per_s	rollout/generated_tok_s (windowed, not cumulative)
scoring_time_ms, wait_scoring_ms	rollout/score_s, rollout/score_wait_s
buffer_qsize	sample/rollout_queue_size
train_seq_len	dropped — superseded by sample/forwarded_tokens_mean and batch/row_tokens_mean

ratio, kl, entropy, clip_ratio/*, aux_loss, completions/mean_length, reward, reward_std, rewards/<func>, tools/call_frequency and tools/failure_frequency keep their names.

queue_wait_time_s is the one worth calling out: it was used for two opposite things — the trainer blocked because the queue was empty, and a sample's own time sitting in it. Those diagnose opposite bottlenecks, so they are now two metrics.

---

Note

Medium Risk
Large change to logging contracts and RolloutWorkerProtocol (custom workers must expose metrics_queue); dashboards need metric renames, but training logic changes are mostly observability and aggregation paths.

Overview
Async GRPO observability is rebuilt so you can tell whether generation, scoring, the rollout queue, or training is the bottleneck. Metrics are grouped by entity (rollout/, completions/, tools/, sample/, batch/, perf/) and documented in a new Logged metrics section in docs/source/async_grpo_trainer.md.

Reduction and naming fixes: _reduce_metric infers sum/mean/max/min and Σnum/Σden for rates (fixing broken throughput like the old training_tok/s). queue_wait_time_s splits into perf/rollout_wait_s (trainer starved) vs sample/time_in_queue_s (per-sample age). Throughput/MFU are reported on _fwd_bwd and _wall_clock bases; weight sync is broken into phase timings under perf/weight_sync_*.

Instrumentation plumbing: The rollout worker pushes metrics through a bounded metrics_queue on RolloutWorkerProtocol; the trainer drains it in log(). The worker logs rollout/scoring/queue backpressure, tool stats, drift/fork tallies from _chain_to_sequences, and windowed rollout/generated_tok_s. The trainer/collator log batch packing (row_fill_frac, pad_frac, token counts), staleness drops, optimizer time via _OptimizerTimeCallback, and per-step aggregates in _log_step_metrics. Per-sample rewards are aggregated on rank 0 in the collator instead of broadcasting metric tensors. Rollout trace logging is throttled per policy version to avoid flooding dashboards.

Tests cover _reduce_metric, worker metric push, protocol metrics_queue, collator metric keys, and drift tally behavior.

^{Reviewed by Cursor Bugbot for commit 4a4a1ae. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Optimizer time logged one step late

Medium Severity

_log_step_metrics runs at the end of the last training_step micro-batch, before optimizer.step(). It flushes and resets perf/optimizer_s while _OptimizerTimeCallback only records the optimizer afterward, so each step logs the previous step's optimizer time and the final step's value is never logged.

Additional Locations (1)

• trl/experimental/async_grpo/async_grpo_trainer.py#L188-L193

 

^{Reviewed by Cursor Bugbot for commit 36bbf31. Configure here.}

[cursor]

Generated throughput counts trained tokens

Medium Severity

rollout/generated_tok_s is fed from _total_completion_tokens, which sums completion_mask over reconciled sequences. After a realign, generated tokens demoted to context are excluded, so the metric under-reports generation throughput relative to the PR's generated-vs-trained definitions and to completions/mean_length.

Additional Locations (1)

• trl/experimental/async_grpo/async_rollout_worker.py#L640-L641

 

^{Reviewed by Cursor Bugbot for commit 36bbf31. Configure here.}

[cursor]

Harness timeout metric never recorded

Low Severity

OpenEnv passes loop_exhausted=timed_out into _push_rollout_metrics, but that rate is only updated when self.tools is truthy. Harness sessions do not populate self.tools, so tools/loop_exhausted_frac stays absent for timed-out OpenEnv rollouts.

Additional Locations (1)

• trl/experimental/async_grpo/openenv_harness.py#L242-L249

 

^{Reviewed by Cursor Bugbot for commit 36bbf31. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

There are 5 total unresolved issues (including 3 from previous reviews).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 094a214. Configure here.}

[cursor]

Unset enqueue time skews queue wait

Medium Severity

sample/time_in_queue_s is always now - sample.enqueued_at. enqueued_at defaults to 0.0 and is only set by the built-in worker’s put_nowait path, so custom workers and _StubRolloutWorker report queue waits of about 1.7e9 seconds.

Additional Locations (1)

• trl/experimental/async_grpo/async_rollout_worker.py#L203-L204

 

^{Reviewed by Cursor Bugbot for commit 094a214. Configure here.}