Skip to main content

core.trainers.async_grpo.progress

core.trainers.async_grpo.progress

Progress accounting and termination for async GRPO’s rollout stream.

TRL’s rollout worker cycles the prompts forever, and TRL ends an epoch-bounded run only when an optimizer step ends. Rollouts dropped before training (no reward signal, truncated, too stale, too long) complete no step, so a run whose remaining rollouts are all dropped would generate forever.

SampleFilterQueue drops the rollouts that would train nothing. Every drop and every trained micro-batch reports to RolloutProgress, which keeps TRL’s set of consumed prompt groups (read for the epoch bound and for resume). When the run’s prompt groups are consumed, or a full pass over the prompts trains nothing, the rank-0 pipeline raises RolloutStop. StopAwareDataset turns it into one stop micro-batch that accelerate’s dispatcher hands to every rank, so all ranks end training (or raise the stall error) at the same micro-batch.

Classes

Name Description
ProgressCollator TRL’s collator, reporting trained groups and passing a stream stop on.
RolloutProgress Consumed and trained prompt groups, and the stream’s stop conditions.
RolloutStop Ends the rollout stream; error is set when the run made no progress.
StopAwareDataset A rollout batcher whose stream ends with the RolloutStop it raised.

ProgressCollator

core.trainers.async_grpo.progress.ProgressCollator(progress=None)

TRL’s collator, reporting trained groups and passing a stream stop on.

RolloutProgress

core.trainers.async_grpo.progress.RolloutProgress(
    consumed,
    *,
    epochs_done,
    stall_groups,
)

Consumed and trained prompt groups, and the stream’s stop conditions.

consumed is the trainer’s set TRL counts epochs and resumes from: a group counts once it trained or its rollouts were dropped. stall_groups bounds how many groups may be consumed in a row without one training.

Only rank 0 reports, from the rollout stream and, with trainer-side scoring, from score_step; the stream raises the stop, so that it reaches every rank as a stop micro-batch.

Methods

Name Description
check End the stream when a full pass over the prompts trained no group.
drop Count group_id as consumed without training; may end the stream.
exclude Count a collated group whose every sample scoring took out of the loss.
train Record groups that trained.
check
core.trainers.async_grpo.progress.RolloutProgress.check()

End the stream when a full pass over the prompts trained no group.

drop
core.trainers.async_grpo.progress.RolloutProgress.drop(group_id, reason)

Count group_id as consumed without training; may end the stream.

exclude
core.trainers.async_grpo.progress.RolloutProgress.exclude(group_id, reason)

Count a collated group whose every sample scoring took out of the loss.

Runs outside the stream, so the stall it may complete is raised by the stream’s next check.

train
core.trainers.async_grpo.progress.RolloutProgress.train(group_ids)

Record groups that trained.

RolloutStop

core.trainers.async_grpo.progress.RolloutStop(error=None)

Ends the rollout stream; error is set when the run made no progress.

StopAwareDataset

core.trainers.async_grpo.progress.StopAwareDataset(batcher)

A rollout batcher whose stream ends with the RolloutStop it raised.