core.trainers.async_grpo.scoring

core.trainers.async_grpo.scoring

Scoring async GRPO rollouts on the trainer.

Some rewards need the GPU: a reward model, or EBFT’s feature network. TRL’s rollout process is CUDA-free, so with trainer scorers it scores only its own reward columns and leaves each sample’s advantage pending (see AxolotlRolloutLoop). The trainer then scores whole groups collectively, on every rank, before training on them:

GroupedQueue releases a group’s samples together, once all of them arrived.
StepPlanner replaces TRL’s batchers: it plans one optimizer step at a time (gradient_accumulation_steps micro-batches of one row per rank) so that no group straddles two steps, and no row is empty.
ScoringCollator adds the samples’ ids, their rollout-side rewards and each scorer’s inputs to TRL’s packed micro-batch.
score_step, called from get_batch_samples on every rank, runs each scorer’s forward on the rank’s micro-batches, all-gathers the features, computes group rewards and advantages, and writes them into the micro-batches. Samples that carry no gradient (zero advantage, or truncated under mask_truncated_completions) are taken out of the loss and its normalization, which is what dropping them before packing would have done.

Classes

Name	Description
GroupedQueue	A view of the rollout queue that releases each group’s samples together.
ScoredRow	A `TreeRow` that also carries the group extras assigned to it.
ScoringCollator	TRL’s collator plus what `score_step` needs on every rank.
StepPlanner	Plans one optimizer step at a time, keeping every group inside one step.
TrainerScorer	One reward column computed on the trainer from a group’s samples.

GroupedQueue

core.trainers.async_grpo.scoring.GroupedQueue(inner)

A view of the rollout queue that releases each group’s samples together.

ScoredRow

core.trainers.async_grpo.scoring.ScoredRow()

A TreeRow that also carries the group extras assigned to it.

ScoringCollator

core.trainers.async_grpo.scoring.ScoringCollator(
    progress=None,
    recorder=None,
    scorers=list(),
)

TRL’s collator plus what score_step needs on every rank.

StepPlanner

core.trainers.async_grpo.scoring.StepPlanner(
    dataset,
    *,
    num_processes,
    steps,
    token_budget,
    samples_per_row,
    scorers,
    metrics,
    share_prefixes=True,
    stop_token_ids=None,
    progress=None,
)

Plans one optimizer step at a time, keeping every group inside one step.

A step is steps micro-batches of num_processes rows. Samples are placed like TRL’s TokenBudgetBatcher (rows holding the group first, then the least loaded row they fit), within token_budget unique and trained tokens per row, or samples_per_row samples when token_budget <= 0. When a group does not fit, the step is rebuilt from the groups accepted so far and emitted, and the group opens the next step. Empty rows, which would desync collectives, are filled by spreading the samples, or failing that with a copy of a sample that trains no token and is not scored.

TrainerScorer

core.trainers.async_grpo.scoring.TrainerScorer()

One reward column computed on the trainer from a group’s samples.

group_extras and collate run on rank 0 while steps are planned and packed; features runs on every rank for every micro-batch, so it may run collective forwards (FSDP); group_rewards runs on every rank, on the gathered features.

Methods

Name	Description
collate	Per-row input tensors for `features` and their padding values.
features	`(uids, features)` of the sequences in one rank’s chunk of a micro-batch.
group_extras	Sequences scored with a group but not trained, e.g. a reference answer.
group_rewards	Rewards of a group’s samples, and metrics to average over groups.

collate

core.trainers.async_grpo.scoring.TrainerScorer.collate(rows, index)

Per-row input tensors for features and their padding values.

index is this scorer’s position, which marks its entries in row.extras. Every row must yield input for features, also when it holds only filler samples (uid -1), so that collective forwards line up across ranks. Keys prefixed c{j}_ split a row into chunks, each scored by its own features call; every row of a micro-batch must then have the same number of chunks, for the same reason.

features

core.trainers.async_grpo.scoring.TrainerScorer.features(trainer, inputs)

(uids, features) of the sequences in one rank’s chunk of a micro-batch.

group_extras

core.trainers.async_grpo.scoring.TrainerScorer.group_extras(samples)

Sequences scored with a group but not trained, e.g. a reference answer.

group_rewards

core.trainers.async_grpo.scoring.TrainerScorer.group_rewards(
    sample_features,
    extra_features,
)

Rewards of a group’s samples, and metrics to average over groups.

Functions

Name	Description
check_step_capacity	Fail early when a group can never fit in one optimizer step.
chunk_inputs	One scorer’s micro-batch inputs, split into its `c{j}_` chunks in order.
chunk_sequences	Indices of `sequences` split into chunks that pack to `token_budget` tokens.
pack_sequences	Tree-pack token sequences for a no-grad forward, as `TreeRow.pack` would.
score_step	Score one optimizer step’s micro-batches in place (see module doc).
unpack_sequences	Cut one rank’s padded `pack_sequences` output back to model inputs.

check_step_capacity

core.trainers.async_grpo.scoring.check_step_capacity(
    group_rows,
    steps,
    num_processes,
    token_budget,
    samples_per_row,
)

Fail early when a group can never fit in one optimizer step.

chunk_inputs

core.trainers.async_grpo.scoring.chunk_inputs(inputs)

One scorer’s micro-batch inputs, split into its c{j}_ chunks in order.

chunk_sequences

core.trainers.async_grpo.scoring.chunk_sequences(
    sequences,
    token_budget,
    share_prefixes,
)

Indices of sequences split into chunks that pack to token_budget tokens.

Bounds the memory of a scoring forward, whose sequences (references, whole conversations) the planner’s budget for trained tokens does not cover. A sequence longer than the budget is scored alone.

pack_sequences

core.trainers.async_grpo.scoring.pack_sequences(sequences, share_prefixes=True)

Tree-pack token sequences for a no-grad forward, as TreeRow.pack would.

Besides the model inputs (input_ids, position_ids, tree_*), returns positions: the packed position of each sequence’s tokens, laid end to end, with offsets marking where each sequence starts, and sizes: the lengths of input_ids, tree_kv_index, tree_cu_seqlens_q and offsets, so a padded copy can be cut back.

score_step

core.trainers.async_grpo.scoring.score_step(trainer, batches)

Score one optimizer step’s micro-batches in place (see module doc).

On rank 0, returns the reward names and each sample’s reward columns and advantage, for the on_rollouts_scored hook.

unpack_sequences

core.trainers.async_grpo.scoring.unpack_sequences(inputs)

Cut one rank’s padded pack_sequences output back to model inputs.