Asynchronous GRPO

[!IMPORTANT] This trainer requires vllm>=0.22.0 and transformers>=5.2.0. For distributed training, only FSDP2 is supported (DeepSpeed ZeRO is not).

Currently, vllm and transformers have conflicting dependency constraints. To work around this, install vLLM first and then force-install transformers:

bashCopy

pip install 'vllm>=0.22.0'
pip install 'transformers>=5.2.0' --no-deps

Overview

[AsyncGRPOTrainer] implements the same GRPO algorithm but decouples rollout generation from training. A background worker continuously streams completions from a vLLM server while the training loop consumes them, so generation and gradient updates overlap instead of alternating. The API mirrors [GRPOTrainer] — for full details on the GRPO method itself (advantage computation, KL estimation, loss formulation, reward functions, etc.), see the GRPO Trainer documentation. Not all features from [GRPOTrainer] are available; refer to [AsyncGRPOConfig] for the supported parameters.

This trainer was contributed by Quentin Gallouédec and Amine Dirhoussi.

How it differs from [GRPOTrainer]

In the standard [GRPOTrainer], generation and training are sequential: generate a batch, compute the loss, update weights, repeat. Even in vLLM colocate mode, where generation runs on the same GPUs, one phase must finish before the other begins.

[AsyncGRPOTrainer] separates these two concerns:

• Rollout worker (background process) — sends prompts to a vLLM server, scores completions with reward functions, computes advantages, and pushes ready-to-train samples into a queue.
• Training loop (main process) — pulls samples from the queue, computes the clipped surrogate loss, and updates the model weights.

The rollout worker runs in a separate process spawned from the trainer, so reward computation never contends with the training loop for the GIL. This has two consequences for what you can pass as reward_funcs, tools, and environment_factory (for the latter, see the OpenEnv guide, which covers the contract and the available integrations):

[!WARNING] Because we run the rollout worker in a separate process, everything passed to it is pickled. Each reward function, tool, and environment_factory (and anything they close over) must therefore be picklable: use a module-level function, functools.partial, or a callable class instance. Lambdas and closures will raise a TypeError at trainer.train(). This is a difference from [GRPOTrainer], where reward functions are called in-process and closures work.

The rollout process also runs with CUDA_VISIBLE_DEVICES="", so it cannot use the GPU. A GPU-backed reward model (e.g. an AutoModelForSequenceClassification scorer) still loads without error but silently falls back to CPU (note that in [GRPOTrainer], such a reward model shares the trainer's GPUs). Keep reward functions CPU-side and lightweight (verifiers like accuracy_reward, format/length checks).

If you do need a GPU reward model, the recommended approach is to serve it behind its own inference engine (vLLM, TGI, …) on separate GPUs and have a lightweight, picklable reward function call it over HTTP. This keeps the reward model on its own device while the rollout process stays CPU-only, and it scales independently of the trainer.

After every weight_sync_steps training steps, the updated weights are transferred to the vLLM server via NCCL so that subsequent generations reflect the latest policy.

Because generation and training run concurrently, the training samples may have been generated by a slightly older version of the model. The max_staleness parameter controls how many weight updates a sample can lag behind before being discarded.

The number of concurrent requests sent to the vLLM server is controlled by max_inflight_tasks. By default it is set automatically to max_staleness × per_device_train_batch_size × gradient_accumulation_steps × num_processes — the maximum number of samples the trainer can consume before they become stale. Generating more than this is wasteful since the excess samples will be discarded.

Checkpoint and resume: ignore_data_skip defaults to True; the base Trainer's skip-and-replay loop does not apply to a live rollout queue. Instead, the index of the first prompt not yet trained on is saved to rollout_state.json alongside each checkpoint and restored on resume, so the worker fast-forwards to that prompt without replaying samples. It is the trained position, not the generator's: the worker runs ahead of training by the rollout queue depth, and those buffered samples are lost when the run ends, so resuming from the generator's position would skip prompts that were generated but never trained on. Streaming datasets (IterableDataset) cannot be repositioned; their worker restarts from prompt 0 on resume.

Quick start

# train_async_grpo.py
from datasets import load_dataset
from trl.experimental.async_grpo import AsyncGRPOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-4B",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)
trainer.train()

The vLLM server and the trainer must run on separate GPUs. Use CUDA_VISIBLE_DEVICES to partition your GPUs. For example, with 2 GPUs, you can run the vLLM server on GPU 0 and the trainer on GPU 1 as follows:

# Terminal 1: vLLM server on GPU 0 (dev mode + NCCL weight transfer are required)
CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --weight-transfer-config '{"backend":"nccl"}'

[!TIP] Set --max-model-len to the maximum total sequence length (prompt + completion) you expect. A lower value reduces GPU memory usage on the server, freeing more memory for the KV cache and increasing throughput. A good starting point is the prompt length plus max_completion_length from your config.

# Terminal 2: training on GPU 1
CUDA_VISIBLE_DEVICES=1 accelerate launch train_async_grpo.py

Design philosophy

This trainer is intentionally kept minimal and is not meant to grow into a general-purpose solution. If you need a feature that is not supported, we recommend cloning the repository and adapting the trainer to your needs directly. New features will only be considered when there is significant community demand.

Logged metrics

A rollout passes through several stages before it becomes a gradient. Let's first take a look at the different stages, from a prompt to our final batch composition:

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)

Reading it top to bottom:

• a group is one prompt generated num_generations times. The advantage baseline is computed inside it, which is why the group is also the unit that gets scored.
• a rollout is one conversation. Every turn re-tokenizes the whole message list, so drift between turns has to be reconciled at the end.
• reconciling can fork one rollout into several sequences, which is why a single conversation can produce more than one training sample.
• a sample is a sequence plus its advantage, reward and metrics. This is what crosses the process boundary.
• past the boundary the trainer plans samples into rows (one per DP rank), packs each row into one concatenated sequence, and accumulates grad_accum micro-batches into a single optimizer step.

So we split our metrics based on which stage (or entity) they are about:

Let's also define some terms to make token metrics unambiguous:

• generated tokens are what the model produced (completions/*).
• forwarded tokens are every token the forward pass processes: the prompt, any tool results, and the generated tokens.
• trained tokens are the subset the loss is taken over where completion_mask == 1. For example, a fork or a realign can demote already-generated tokens to context, so trained ≠ generated.

How the batch numbers compose

A row is what one DP rank forwards in one micro-batch: several samples concatenated into a single sequence, with position_ids restarting at each sample boundary. The planner decides which samples land in which row (balancing Σ Lᵢ² so no rank straggles), and a data collator does the concatenating.

In all our metrics a "step" always means one FULL optimizer step, never a micro-batch. gradient_accumulation_steps micro-batches make one step, and every metric named _per_step is per optimizer step, matching global_step, logging_steps and the rest of the [Trainer] vocabulary. Where a metric is per micro-batch it says so (batch/microbatches_per_step) or it is a per-row quantity (batch/row_*, batch/samples_per_row).

One step therefore holds a fixed number of row-slots:

row-slots per step = gradient_accumulation_steps x world_size

which makes the batch metrics compose into a ladder you can check against each other:

sample                                     sample/forwarded_tokens_mean
  └─ packed into a ROW                     batch/samples_per_row, batch/row_tokens_mean
      └─ one row per DP rank               = one micro-batch
          └─ gradient_accumulation_steps micro-batches = one STEP
                                           batch/samples_per_step, batch/forwarded_tokens_per_step

So batch/samples_per_step ≈ row-slots x batch/samples_per_row, and likewise for tokens. The two sides agree only up to the variation between micro-batches inside the step — the per-step metrics are sums, the per-row ones are means — so expect a fraction of a percent, not an exact match.

batch/row_fill_frac is the one to watch when samples are long: a 10k-token sample tiles a 32k budget badly (three fit, four never do, so the packer often gets two and the row runs ~77% full), while 1k-token samples tile it almost perfectly. That is quantization, not a bug, and token_budget is the lever — bearing in mind that attention is O(L²) per sequence, so a fuller row of long sequences does not cost linearly more memory.

The rollout queue

The worker pushes scored samples into a queue and the trainer pulls from it. Four metrics describe that one queue, and they answer different questions:

The last two metrics (perf/rollout_wait_s and rollout/backpressure_s) are mirror images and are never both large! Reading them together with the queue size tells you which side is the bottleneck:

• queue near empty, perf/rollout_wait_s high → generation-bound. The trainer is starving; look at rollout/generated_tok_s, rollout/inflight and rollout/score_queue_size.
• queue near full, rollout/backpressure_s high → trainer-bound. Generation is throttled and its output is aging in the queue, so watch sample/staleness_mean climb.
• both near zero → balanced !

Completions

A completion is what vLLM generated for one prompt, counted across every turn of the rollout.

Rollouts

A rollout is one full conversation: a prompt generated to completion, including every tool round-trip it took to get there. num_generations rollouts share a prompt and form a group.

Tools

Logged only when the model executes tools. tools/<name>_* repeats per tool name, the same way rewards/<func> repeats per reward function, because one failing tool is invisible in an average over all of them.

Samples arriving from the queue

Each sample is one RolloutSample: a reconciled sequence plus the group advantage, the reward and its per-sample metrics, assembled by the worker in _score_group once every generation in a group has finished and been scored. The worker pushes them into rollout_buffer; everything below is measured on the trainer side, as it pulls them for training.

Batches

One micro-batch is world_size rows (remember packing flattens into 1 sequence), so one per DP rank, so that rank i forwards row i. gradient_accumulation_steps of those make one optimizer step. A step then covers gradient_accumulation_steps × world_size rows in total. Every _per_step metric below is a sum over the whole step and across every rank; the batch/row_* ones are means over the rows.

Performance

Throughput and MFU are each reported twice, over the same optimizer step, differing only in what they divide by. The suffix names the denominator:

• _fwd_bwd divides by perf/fwd_bwd_s — the compute alone. How efficiently does the trainer run when it has data? If it is low, the trainer is the problem.
• _wall_clock divides by perf/step_s — the whole step, including the time spent waiting for rollouts. What fraction of the allocation actually became training? If this is far below the _fwd_bwd one, generation is probably the bottleneck.

The gap between them is perf/rollout_wait_s plus the optimizer and weight-sync time. But it's useful to look at both: looking only at _fwd_bwd hides the GPU-hours spent generating rollouts, and quoting only _wall_clock could blame the trainer for the generator's latency.

AsyncGRPOConfig

[[autodoc]] trl.experimental.async_grpo.AsyncGRPOConfig

AsyncGRPOTrainer

[[autodoc]] trl.experimental.async_grpo.AsyncGRPOTrainer

RolloutWorkerProtocol

[[autodoc]] trl.experimental.async_grpo.async_grpo_trainer.RolloutWorkerProtocol