Skip to main content

core.trainers.async_grpo.traces

core.trainers.async_grpo.traces

The on_rollouts_scored plugin hook for async GRPO.

Rank 0 collates every micro-batch, so it records the samples of each one and, once an optimizer step has fetched its micro-batches (and scored them, with trainer-side scorers), hands that step’s prompts, completions, rewards and advantages to plugins.

Classes

Name Description
TraceCollator TRL’s collator that records each micro-batch’s samples for the trace hook.
TraceRecorder Rank 0’s samples of each collated micro-batch, until its step pops them.

TraceCollator

core.trainers.async_grpo.traces.TraceCollator(progress=None, recorder=None)

TRL’s collator that records each micro-batch’s samples for the trace hook.

TraceRecorder

core.trainers.async_grpo.traces.TraceRecorder()

Rank 0’s samples of each collated micro-batch, until its step pops them.

Functions

Name Description
rollout_traces The hook’s (prompts, completions, rewards, advantages) for samples.

rollout_traces

core.trainers.async_grpo.traces.rollout_traces(tokenizer, samples, scored=None)

The hook’s (prompts, completions, rewards, advantages) for samples.

Prompt and completion are the decoded tokens before and from the first trained one. scored holds the reward names and, per sample uid, the reward columns and advantage computed by trainer-side scoring; without it they come from what the rollout process stamped on each sample.