Skip to main content

core.trainers.async_grpo.rollout

core.trainers.async_grpo.rollout

Rollout-process extensions for TRL’s async GRPO worker.

Everything here runs inside the CUDA-free child process TRL’s AsyncRolloutWorker spawns, so it keeps its imports light.

Classes

Name Description
AxolotlAsyncRolloutWorker AsyncRolloutWorker whose child process runs :class:AxolotlRolloutLoop,
AxolotlRolloutLoop TRL’s rollout loop with GRPOTrainer’s reward options and parallel scoring.
RewardProcessPool Long-lived processes that run synchronous reward functions.

AxolotlAsyncRolloutWorker

core.trainers.async_grpo.rollout.AxolotlAsyncRolloutWorker(
    loop_cls=None,
    **kwargs,
)

AsyncRolloutWorker whose child process runs :class:AxolotlRolloutLoop, or the subclass passed as loop_cls.

AxolotlRolloutLoop

core.trainers.async_grpo.rollout.AxolotlRolloutLoop(
    reward_weights=None,
    multi_objective_aggregation='sum_then_normalize',
    scale_rewards='group',
    reward_num_workers=1,
    trainer_scored=False,
    **kwargs,
)

TRL’s rollout loop with GRPOTrainer’s reward options and parallel scoring.

Subclasses that score each rollout themselves set _provides_rollout_reward and return the reward from _generate_one; it is scored as one more reward column, named _rollout_reward_name, after the reward functions.

With trainer_scored, the trainer adds reward columns that need the GPU, so this process scores its own columns but leaves the advantage pending: samples carry a NaN advantage and a payload under SCORE_PAYLOAD (the group’s row count, the columns scored here and _trainer_score_payload).

RewardProcessPool

core.trainers.async_grpo.rollout.RewardProcessPool(reward_funcs, num_workers)

Long-lived processes that run synchronous reward functions.

Each call runs on a worker’s main thread, so CPU-bound rewards scale past the GIL and signal.alarm timeouts (math_verify) stay enabled; TRL’s default asyncio.to_thread has to switch them off.

Functions

Name Description
compute_group_advantages Return the weighted rewards and the advantages of one prompt group.

compute_group_advantages

core.trainers.async_grpo.rollout.compute_group_advantages(
    rewards_per_func,
    reward_weights,
    aggregation='sum_then_normalize',
    scale_rewards='group',
)

Return the weighted rewards and the advantages of one prompt group.

Parameters

Name Type Description Default
rewards_per_func np.ndarray (num_funcs, group_size) rewards, NaN where a function did not score a completion. required
reward_weights np.ndarray (num_funcs,) weight of each reward function. required
aggregation str sum_then_normalize normalizes the weighted sum; normalize_then_sum (GDPO) normalizes each function within the group before the weighted sum. 'sum_then_normalize'
scale_rewards str group divides by the group’s standard deviation, none only centers. 'group'

A completion no function scored keeps a NaN reward and a zero advantage. A group whose learning signal is constant gets exactly zero advantages, so it can be told apart from a group that merely has small ones.