6.2 The Generation–Training Loop & Rollout Engines¶

In The Anatomy of an RL-for-LLM System we drew the system as four boxes — a generator that samples responses, a reward that scores them, a learner that updates weights, and a controller that orchestrates the cycle. This chapter zooms into the beating heart of that system: the loop that alternates between generating experience and training on it. Everything else in RL infrastructure — Ray placement groups, weight-sync protocols, disaggregation — exists to make this one loop go faster without breaking its math.

The loop is deceptively simple to state. Sample a batch of prompts. Generate completions. Score them. Compute advantages. Take a few gradient steps. Repeat. But each of those verbs hides a different system: generation is an autoregressive, memory-bound inference workload best served by a purpose-built engine like vLLM or SGLang; training is a compute-bound, gradient-accumulating workload best served by FSDP or Megatron. RL for LLMs is the awkward marriage of these two opposite regimes inside a single training step, and almost every interesting engineering decision in the field — colocate or disaggregate, sync or async, on-policy or off — is about managing the seam between them.

The single most important empirical fact to anchor on: generation usually dominates the wall-clock. In a typical GRPO or PPO run, sampling the rollouts can eat 60–80% of each step’s time, with the actual backward pass a relative afterthought. If you do not understand why, you will optimize the wrong thing. So we start there.

On-policy, off-policy, and the “staleness” that RL infra actually fights¶

The definition that matters for systems¶

A learning algorithm is on-policy if the data it trains on was generated by the same policy it is updating, and off-policy if the data came from some other (older or different) policy. In the LLM RL setting this distinction is not abstract — it is literally a question of which version of the weights produced these tokens.

Let \(\pi_\theta\) be the policy we are training and \(\pi_{\theta_{\text{behavior}}}\) be the policy that actually sampled the rollouts. Strictly on-policy means \(\theta_{\text{behavior}} = \theta\) at the moment of every gradient computation. The instant you take one optimizer step, your weights move, and any rollouts you have not yet consumed become off-policy — they were drawn from the now-stale \(\theta_{\text{behavior}}\).

This is why PPO and GRPO carry an importance sampling ratio. The policy-gradient objective is an expectation under \(\pi_\theta\), but we have samples from \(\pi_{\theta_{\text{behavior}}}\). The correction is the per-token ratio

and the clipped surrogate \(\min(r_{i,t}\hat A_i,\ \operatorname{clip}(r_{i,t},1-\epsilon,1+\epsilon)\hat A_i)\) is what keeps the off-policy correction from exploding when \(\theta\) has drifted far from \(\theta_{\text{behavior}}\). The full derivation lives in Policy Gradients & PPO for Language Models; here we care about the systems consequence: the clip exists precisely so that we are allowed to be a little off-policy, which in turn is what lets us reuse a batch of rollouts for several gradient steps instead of regenerating after every single one.

The staleness knob¶

In practice “on-policy vs off-policy” is not binary; it is a dial we call staleness \(s\) — how many optimizer steps old the weights were when they generated the data we are now training on.

Three regimes appear in real systems:

• Synchronous on-policy-ish (\(s\) small, bounded by PPO epochs). Generate a big batch, then take ppo_epochs × num_minibatches gradient steps on it before regenerating. Within those steps \(\theta\) drifts, so the later minibatches are mildly off-policy — the clip absorbs it. This is what TRL’s PPOTrainer and most GRPO recipes do by default. Staleness is small and bounded.
• Asynchronous (\(s\) moderate, e.g. 1–4). The generator runs continuously on its own copy of the weights while the trainer updates; weights are pushed to the generator every few steps. The generator is therefore always a bounded number of steps behind. This trades a little policy-gradient bias for a large throughput win (we overlap generation and training). Covered in depth in Prime-RL, Async RL & Decentralized Training.
• Off-policy / replay (\(s\) large). You keep a buffer of old rollouts and replay them. This is rare in modern LLM RL because the importance weights get huge and the clip throws away most of the gradient — but it shows up in some sample-efficiency-focused research.

Why we tolerate any staleness at all¶

If on-policy is statistically cleaner, why not always run \(s=0\)? Because \(s=0\) is a throughput disaster. Strictly on-policy means: generate a batch, do one gradient step, throw the batch away, generate again. You pay the full (dominant) generation cost to extract a single update. Allowing a handful of minibatch updates per rollout batch — and allowing the sampler to lag slightly — is how RL infra reclaims an order of magnitude of throughput. The art is keeping \(s\) small enough that the policy gradient stays approximately valid (clip fraction reasonable, KL controlled) while large enough that the expensive generators are never idle.

Anatomy of one RL step: the five phases¶

Let us make the loop concrete. One outer step (sometimes called an iteration or a “rollout step”) consists of five phases. We will time each one in the worked example later; first, what they are.

Phase 1 — Rollout / generation¶

We draw \(G\) completions per prompt (the group, for GRPO/RLOO — see GRPO, RLOO & Critic-Free RL) for a batch of \(P\) prompts, so \(B = P \cdot G\) sequences. This is pure autoregressive decoding: prefill the prompt, then emit one token at a time until EOS or a length cap. It is memory-bandwidth-bound (each decode step reads the whole model and KV cache to produce one token) and the dominant cost. We discuss the engines that do this well in the next section.

Critically, the generator returns not just token ids but the sampling log-probabilities \(\log \pi_{\theta_{\text{behavior}}}(o_{i,t}\mid\cdot)\) — these are the behavior-policy log-probs that go in the denominator of the importance ratio. Getting these to agree with the trainer’s recomputed log-probs is the single subtlest correctness issue in the whole loop (Section on numerics below).

Phase 2 — Reward / scoring¶

Each completion is scored. For RL with verifiable rewards this is a math checker, a code sandbox, or an exact-match rule (Reward Engineering, Verifiers & Sandboxes); for RLHF it is a reward-model forward pass (The RLHF Pipeline & Reward Modeling). Reward is usually cheap relative to generation — a sandbox call is milliseconds, a 7B RM forward is one prefill — but it can become a tail-latency problem if the sandbox is slow or rate-limited, which is why it is often run asynchronously, per-sample, as completions finish rather than as a blocking batch step.

Phase 3 — Experience preparation¶

Now we move to the trainer’s GPUs. We run forward passes to obtain, for every token of every completion:

• \(\log \pi_\theta(o_t\mid\cdot)\) — the current policy’s log-probs (the “old” log-probs used as the ratio denominator if you trust them, or to re-derive a consistent behavior baseline);
• \(\log \pi_{\text{ref}}(o_t\mid\cdot)\) — the frozen reference model’s log-probs, for the KL penalty (skip if \(\beta=0\));
• optionally \(V(s_t)\) — the critic’s value estimates, for PPO/GAE (skip for GRPO/RLOO).

Then we compute advantages. For GRPO this is the group z-score; for PPO it is GAE over the critic’s values. We assemble everything into a flat training batch: (input_ids, response_mask, old_logprobs, ref_logprobs, advantages).

Phase 4 — Learn / minibatch updates¶

We iterate ppo_epochs times over the rollout batch, each time shuffling and splitting it into minibatches, computing the clipped surrogate loss, and stepping the optimizer with gradient accumulation. This is the only phase that consumes gradients, and (perhaps surprisingly) it is often the fastest of the five — a handful of forward+backward passes versus thousands of autoregressive decode steps.

Phase 5 — Weight synchronization¶

After updating \(\theta\), the inference engine still holds the old weights. We must copy the new \(\theta\) into the generator before (or during) the next rollout. In a colocated system this is an in-process tensor copy; in a disaggregated system it is a cross-process or cross-node transfer (NCCL broadcast, RDMA, or a serialized checkpoint). For a 7B model in bf16 that is ~14 GB to move every step — non-trivial, and the subject of Colocated vs Disaggregated RL & Weight Synchronization.

Rollout engines: why you don’t use model.generate in production¶

The toy GRPO code in GRPO, RLOO & Critic-Free RL called model.generate(). That is pedagogically honest and catastrophically slow for real RL. The reason is the same reason production inference uses vLLM or SGLang: HuggingFace’s generate does static batching with a dense KV cache, so the whole batch runs at the speed of its longest sequence and most of the KV memory is wasted on padding.

What a real rollout engine buys you¶

A purpose-built engine (vLLM: Architecture, PagedAttention & Internals, SGLang: RadixAttention & Structured Programs) gives the rollout phase three things that matter enormously for RL:

1. Continuous (in-flight) batching. As soon as one sequence finishes, its slot is filled by a waiting sequence — no waiting for the longest member of a static batch (Continuous Batching & Request Scheduling). In RL, completion lengths vary wildly (some prompts get solved in 50 tokens, some ramble for 4,000), so this is a 2–5× win on its own.
2. PagedAttention. The KV cache is stored in non-contiguous “pages” so memory is allocated on demand instead of pre-padded to max length (PagedAttention & KV-Cache Memory Management). This lets you fit a far larger effective batch and keep the GPU saturated.
3. Prefix sharing. All \(G\) completions in a group share the same prompt prefix. RadixAttention (SGLang) and vLLM’s prefix caching (Prefix Caching & KV-Cache Reuse) compute that shared prefix’s KV once and reuse it across the group — a direct, large saving that is almost unique to the RL access pattern, where you deliberately sample many continuations of one prompt.

Calling vLLM from inside a trainer¶

In production RL, the trainer and a vLLM engine live in the same process (or adjacent processes) and the trainer drives generation through vLLM’s LLM API or its async engine. The crucial extra requirements over plain serving are: (a) you need the sampling log-probs back for the importance ratio, and (b) you need to swap the weights every step without restarting the engine.

# Sketch of the rollout call inside an RL trainer using vLLM.
# (Real code in TRL/veRL is more involved; this shows the load-bearing parts.)
from vllm import LLM, SamplingParams

# One persistent engine, created ONCE. We will hot-swap weights into it each step.
engine = LLM(
    model="Qwen/Qwen2.5-7B",
    dtype="bfloat16",
    gpu_memory_utilization=0.5,   # leave room for the trainer if COLOCATED
    enable_prefix_caching=True,   # reuse the shared prompt prefix across the group
    enable_sleep_mode=True,       # lets us free the engine's KV cache during training
    max_model_len=4096,
)

sampling = SamplingParams(
    n=8,                 # G: sample 8 completions per prompt in ONE call
    temperature=1.0,
    top_p=1.0,
    max_tokens=2048,
    logprobs=0,          # return the logprob of the SAMPLED token (behavior logprob)
    # ^ logprobs=0 means "0 extra besides the chosen token" -> we still get the
    #   chosen token's logprob, which is exactly the behavior-policy logprob we need.
)

def rollout(prompts):
    """Return, per prompt, G completions + their per-token sampling logprobs."""
    outs = engine.generate(prompts, sampling)         # batched, continuous-batched internally
    batch = []
    for req in outs:                                  # one req per prompt
        for comp in req.outputs:                      # G completions
            token_ids = comp.token_ids
            # vLLM gives a list of {token_id: Logprob} dicts, one per position.
            behavior_logprobs = [
                lp_dict[tid].logprob
                for tid, lp_dict in zip(token_ids, comp.logprobs)
            ]
            batch.append({
                "prompt": req.prompt,
                "response_ids": token_ids,
                "behavior_logprobs": behavior_logprobs,   # log pi_behavior(o_t|.)
            })
    return batch

The weight swap (Phase 5) typically goes through vLLM’s worker API. Conceptually:

def sync_weights_to_engine(engine, state_dict):
    """
    Push freshly-trained weights into the live vLLM engine without a restart.
    `state_dict` maps param name -> bf16 tensor (e.g. policy.state_dict()).

    In current vLLM the engine core runs in its own process, so you do NOT reach
    into internals; the supported entry point is the public `collective_rpc`,
    which invokes a named method on every worker. You make `load_weights`
    callable by passing `worker_extension_cls="mypkg.MyWorkerExt"` when
    constructing the LLM (see vLLM's RLHF example under
    `examples/offline_inference/`). veRL/OpenRLHF wrap this so the named tensors
    map correctly onto vLLM's (possibly tensor-parallel-sharded) layout, and for
    multi-node they broadcast the tensors over NCCL rather than pickling them.
    """
    engine.collective_rpc("load_weights", args=(list(state_dict.items()),))
    # After this, the next engine.generate() samples from the NEW policy.

In a colocated setup (engine and trainer on the same GPUs) there is one more move that matters as much as the transfer itself: the engine’s KV cache is a large, statically reserved pool, and it is dead weight while the trainer runs. vLLM’s sleep mode releases it — engine.sleep(level=1) before Phase ¾ frees the KV blocks (level 2 also offloads the weights) and engine.wake_up() re-allocates before the next rollout — which is what lets a single 80 GB GPU host both an inference engine and an FSDP trainer without either being memory-starved.

The details (NCCL broadcast for multi-GPU engines, parameter name remapping, handling tensor-parallel sharding) are exactly what libraries like veRL (veRL: HybridFlow & The Single-Controller Architecture) and OpenRLHF (OpenRLHF, NeMo-Aligner & Ray-Based Systems) exist to handle robustly. The conceptual point is unchanged: one persistent inference engine, weights hot-swapped each step. Tearing down and rebuilding the engine per step would cost more than the rollout itself.

The same five phases, in a real library¶

You would not hand-write the above for a production run. TRL’s GRPOTrainer is the shortest path to exactly the loop of this chapter, and its config surface is a useful glossary of the knobs we have been naming (TRL: HuggingFace’s RL Library):

from trl import GRPOTrainer, GRPOConfig

def reward_correct(completions, answer, **kwargs):           # Phase 2, verifiable
    return [1.0 if extract_final(c) == a else 0.0 for c, a in zip(completions, answer)]

cfg = GRPOConfig(
    num_generations=8,            # G: group size for the group-relative baseline
    max_completion_length=1024,   # caps the L_g long tail that gates a sync step
    num_iterations=1,             # "ppo_epochs": gradient passes per rollout batch
    beta=0.0,                     # KL-to-reference coefficient; 0 => skip the ref forward
    epsilon=0.2, epsilon_high=0.28,   # asymmetric ("decoupled") clip, DAPO-style
    use_vllm=True,                # Phase 1 goes through vLLM, not model.generate
    vllm_mode="colocate",         # engine shares the trainer's GPUs ("server" = disaggregated)
    per_device_train_batch_size=4, gradient_accumulation_steps=8,
)
GRPOTrainer(model="Qwen/Qwen2.5-1.5B-Instruct", reward_funcs=reward_correct,
            args=cfg, train_dataset=ds).train()

TRL handles Phases 1–5 internally: it drives vLLM for the rollout, calls your reward functions, recomputes the old/reference log-probs, runs the clipped-surrogate minibatch loop, and syncs weights back into the engine each step. Reading trl/trainer/grpo_trainer.py alongside this chapter is the single best way to see the abstractions land on real code. vllm_mode="server" (backed by the trl vllm-serve CLI) is the disaggregated variant, where the engine lives on separate GPUs and the trainer talks to it over HTTP.

The throughput bottleneck: why generation dominates, with the numbers¶

This is the central quantitative fact of RL infra. Let us derive why generation dominates and then put real numbers on it.

A FLOP and bandwidth accounting¶

Consider one outer step with \(B = P \cdot G\) sequences, prompt length \(L_p\), mean generation length \(L_g\), a model with \(N\) parameters, and ppo_epochs = E with the whole batch reused.

Training compute (Phase 4). A forward+backward pass costs about \(6N\) FLOPs per token (the classic \(2N\) forward, \(4N\) backward; see Scaling Laws: Kaplan, Chinchilla & Beyond). We train on the response tokens of all \(B\) sequences, \(E\) times:

Generation compute (Phase 1). Decoding is \(2N\) FLOPs per generated token (forward only), and we generate \(B \cdot L_g\) tokens:

By raw FLOPs, generation looks cheaper than training (\(2N\) vs \(6NE\)). So why does it dominate the clock? Because the two phases run in completely different efficiency regimes:

• Training is compute-bound and runs at high Model FLOPs Utilization (MFU), often 40–55% of peak on a good FSDP/Megatron setup.
• Decode generation is memory-bandwidth-bound: each decode step must stream all \(N\) parameters (and the growing KV cache) from HBM to produce a single token per sequence. Arithmetic intensity is tiny, so MFU collapses to low single-digit percent. The wall-clock per token is set not by FLOPs but by how fast you can read the weights from memory.

The decode time is better modeled by bandwidth. Per decode step across a batch of \(B\) sequences, you read the weights once (amortized over the batch) plus each sequence’s KV:

and total decode time is roughly \(L_g\) times that. The key term is the weight read: even with a large batch, you pay \(\sim N \cdot b_{\text{param}}\) bytes of weight traffic per decode step, and there are \(L_g\) steps. That is the tax that makes generation slow.

The qualitative ranking is robust across model sizes and clusters:

This is the reason RL-for-LLMs research is obsessed with rollout throughput. Doubling your backward-pass speed barely moves the needle; doubling your generation throughput nearly halves your step time.

Experience collection and advantage computation in detail¶

Between “we have scored completions” and “we can take a gradient step” sits a surprising amount of bookkeeping. Getting it right is mostly about masking and alignment.

Building the experience batch¶

Each rollout produces a variable-length response. We pad to a common length and build a response mask that is 1 on generated tokens and 0 on prompt tokens and padding — every loss term, every advantage, every KL must be masked to the response, never the prompt. Prompt tokens were given, not chosen, so they carry no policy gradient.

import torch
import torch.nn.functional as F

def build_experience_batch(rollouts, pad_id, device):
    """
    rollouts: list of dicts with 'prompt_ids' (Lp,), 'response_ids' (Lg,),
              'behavior_logprobs' (Lg,), and 'reward' (scalar).
    Returns padded tensors + a response mask aligned to the next-token targets.
    """
    seqs, masks, beh_lp, rewards = [], [], [], []
    for r in rollouts:
        ids = torch.cat([r["prompt_ids"], r["response_ids"]])         # (Lp+Lg,)
        m = torch.zeros_like(ids, dtype=torch.float)
        m[len(r["prompt_ids"]):] = 1.0                                # mask = 1 on response
        seqs.append(ids); masks.append(m)
        # behavior logprobs are defined only on response tokens; left-pad with 0s
        beh = torch.zeros_like(ids, dtype=torch.float)
        beh[len(r["prompt_ids"]):] = torch.tensor(r["behavior_logprobs"])
        beh_lp.append(beh); rewards.append(r["reward"])

    maxlen = max(s.numel() for s in seqs)
    def pad(x, val):  # right-pad to maxlen
        return torch.stack([F.pad(t, (0, maxlen - t.numel()), value=val) for t in x])

    input_ids      = pad(seqs,  pad_id).to(device)        # (B, L)
    response_mask  = pad(masks, 0.0).to(device)           # (B, L)
    behavior_lp    = pad(beh_lp, 0.0).to(device)          # (B, L)
    rewards        = torch.tensor(rewards, device=device) # (B,)

    # IMPORTANT: logprobs/loss are computed on NEXT-token prediction, so when we
    # gather log pi(o_t | o_<t) from logits[:, :-1], the aligned mask/targets are
    # shifted by one. Keep the response_mask and shift it where you compute loss.
    return input_ids, response_mask, behavior_lp, rewards

Recomputing log-probs (the “old” forward) and the reference¶

As warned above, we recompute the behavior/old log-probs on the trainer so the importance ratio is numerically consistent. We also run the frozen reference for the KL term. Both are no_grad forwards over the entire batch — this is Phase 3, and it costs roughly one (or two, with the reference) generation-length forward pass over all \(B\) sequences. That is why experience-prep is the second-biggest time sink.

@torch.no_grad()
def token_logprobs(model, input_ids):
    """log pi(o_t | o_<t) for the actually-present next tokens. (B, L-1)."""
    logits = model(input_ids).logits[:, :-1, :]          # predict token t+1 from t
    logp   = F.log_softmax(logits.float(), dim=-1)
    targets = input_ids[:, 1:]
    return logp.gather(-1, targets.unsqueeze(-1)).squeeze(-1)

# Phase 3 — experience prep:
old_lp = token_logprobs(policy, input_ids)               # consistent ratio denominator
ref_lp = token_logprobs(ref,    input_ids)               # for KL penalty (if beta>0)
resp_mask = response_mask[:, 1:]                          # shift to align with targets

Advantage computation¶

For GRPO the advantage is the per-group z-score of the reward, broadcast to every token of the response (derived in GRPO, RLOO & Critic-Free RL). For PPO it is GAE over the critic’s value estimates. The infrastructure-level subtleties — whitening, KL-in-reward vs KL-in-loss, clipping the advantage — live in Advantage Estimation, KL Control & Stability Tricks. Here is the GRPO path, which needs no critic:

def grpo_advantages(rewards, group_size, normalize_std=True, eps=1e-4):
    """Group-relative advantage; one scalar per response, then broadcast to tokens."""
    g = rewards.view(-1, group_size)                     # (n_prompts, G)
    adv = g - g.mean(dim=1, keepdim=True)                # baseline = group mean
    if normalize_std:
        adv = adv / (g.std(dim=1, keepdim=True) + eps)   # (the contested std norm)
    return adv.reshape(-1)                               # (B,)

Now notice what happens when every completion in a group earns the same reward — all \(G\) correct on an easy prompt, or all \(G\) wrong on an impossible one. The centered advantage is identically zero, so that group contributes no gradient at all, having burned \(G\) full generations: the most expensive resource in the loop. On a mismatched prompt set this can silently void a large fraction of every batch. DAPO’s dynamic sampling fixes it at the loop level: over-sample prompts, discard the degenerate groups, and keep drawing until the batch is full of groups that actually carry signal — an effective-batch-size guarantee bought with a variable number of generations per step.

def informative_groups(rewards, group_size, min_std=1e-6):
    """Boolean mask over GROUPS: True where the group's rewards have variance."""
    g = rewards.view(-1, group_size)
    return g.std(dim=1) > min_std        # False => advantage is 0 => no gradient

Logging the surviving fraction (the “effective group rate”) is one of the most informative curves in an RLVR run: if it collapses toward zero your prompts have been solved and the curriculum needs harder problems (RL Data, Curriculum & Replay Management).

The result is a flat experience batch: input_ids, resp_mask, old_lp, ref_lp, advantages. Everything downstream is the minibatch loop.

Minibatch updates: the inner loop¶

Now we consume the batch. The outer batch (\(B\) sequences) is split into minibatches, and we do ppo_epochs passes. This is where staleness enters within a step: after the first minibatch updates \(\theta\), the second minibatch’s ratio is no longer exactly 1, and the clip starts doing real work.

def minibatch_update_loop(input_ids, resp_mask, old_lp, ref_lp, advantages,
                          policy, opt, *, ppo_epochs=2, n_minibatches=4,
                          eps_low=0.2, eps_high=0.28, kl_beta=0.0, grad_accum=1):
    """The Phase-4 inner loop: ppo_epochs x n_minibatches clipped-surrogate steps."""
    B = input_ids.size(0)
    A = advantages.unsqueeze(1)                          # (B,1), broadcast over tokens
    micro = 0                                            # counts minibatches for grad accum

    for epoch in range(ppo_epochs):
        perm = torch.randperm(B, device=input_ids.device)    # reshuffle each epoch
        for mb in perm.chunk(n_minibatches):
            # --- forward WITH grad on this minibatch ---
            logits = policy(input_ids[mb]).logits[:, :-1, :]
            logp   = F.log_softmax(logits.float(), dim=-1)
            tgt    = input_ids[mb, 1:]
            new_lp = logp.gather(-1, tgt.unsqueeze(-1)).squeeze(-1)   # (m, L-1)
            mask   = resp_mask[mb]

            # --- clipped surrogate (PPO/GRPO) ---
            ratio     = (new_lp - old_lp[mb]).exp()       # pi_theta / pi_behavior
            unclipped = ratio * A[mb]
            clipped   = torch.clamp(ratio, 1 - eps_low, 1 + eps_high) * A[mb]
            pg_loss   = -torch.min(unclipped, clipped)    # we minimize the negative

            if kl_beta > 0:                               # k3 KL to reference
                log_r = ref_lp[mb] - new_lp               # log(pi_ref/pi_theta)
                kl    = log_r.exp() - log_r - 1.0
                pg_loss = pg_loss + kl_beta * kl

            # token-level reduction (Dr.GRPO/DAPO style): sum over tokens / num tokens
            loss = (pg_loss * mask).sum() / mask.sum().clamp(min=1.0)
            (loss / grad_accum).backward()

            # --- diagnostics you SHOULD log every minibatch ---
            with torch.no_grad():
                clip_frac = ((ratio - 1.0).abs() > eps_low).float().mean().item()
                approx_kl = (old_lp[mb] - new_lp).mul(mask).sum() / mask.sum()
            # if clip_frac is huge or approx_kl spikes -> you're too off-policy.

            # Step every `grad_accum` minibatches. NOTE the placement: the optimizer
            # must step INSIDE the minibatch loop, not once per epoch. That is what
            # makes the *later* minibatches genuinely off-policy w.r.t. `old_lp` --
            # i.e. what gives the clip something to do. Stepping once per epoch would
            # turn `n_minibatches` into plain gradient accumulation over one big
            # on-policy batch, which is a different (and much weaker) algorithm.
            micro += 1
            if micro % grad_accum == 0:
                torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
                opt.step(); opt.zero_grad(set_to_none=True)
    return loss.item(), clip_frac, approx_kl.item()

What to watch in the inner loop¶

Three diagnostics tell you whether your staleness is under control:

• Clip fraction — the share of tokens whose ratio left \([1-\epsilon,1+\epsilon]\). A few percent is healthy; 30%+ means the policy is moving too far per batch (lower the LR, fewer ppo_epochs, smaller batch reuse).
• Approximate KL between old and new policy — how far \(\theta\) drifted within the step. If this blows up across ppo_epochs, your reuse is too aggressive.
• Ratio at the first minibatch of epoch 0. If old_lp was recomputed consistently, this should be ≈ 1.0 with clip_frac ≈ 0. If it isn’t, you have the numerics bug from the warning above.

Overlapping generation and training: from synchronous to async¶

We have established that generation dominates and that the trainer’s GPUs are idle while the engine decodes (and vice versa). The whole frontier of RL infra is about filling those idle windows. There is a ladder of designs.

Rung 0 — Synchronous, colocated, serial¶

The simplest correct system. One pool of GPUs. Generate (GPUs busy with the engine), then train (GPUs busy with FSDP), strictly serial. Each phase idles the other’s machinery. Easiest to reason about, strictly on-policy up to ppo_epochs, but leaves a lot of utilization on the table — when the trainer runs, the inference engine’s memory and SMs sit unused, and vice versa. Colocation details (sharing GPU memory between engine and trainer, gpu_memory_utilization budgeting, sleeping the engine during training) are in Colocated vs Disaggregated RL & Weight Synchronization.

Rung 1 — Disaggregated, synchronous¶

Put generation on one set of GPUs (inference cluster) and training on another (training cluster). Each is now always doing its own job. But if you still gate each training step on a complete fresh rollout, the training GPUs idle during generation and the inference GPUs idle during training — you have moved the bubble, not removed it, unless you pipeline.

Rung 2 — Pipelined / overlapped (one-step-stale)¶

This is the first real win. While the trainer updates on rollout batch \(k\), the generator is already producing rollout batch \(k+1\) using the weights from step \(k-1\) (or \(k\)). Generation of the next batch overlaps with training on the current one. The trainer is never blocked on generation; the generator is never blocked on training. The cost is exactly one step of staleness — the rollouts the trainer consumes were generated by weights one update old — which the PPO/GRPO clip is built to absorb.

Rung 3 — Fully asynchronous¶

Decouple completely. Generators run continuously, dumping completed rollouts into a queue. The trainer pulls from the queue, updates, and periodically (every few steps) pushes new weights to the generators. Staleness is whatever the queue depth and sync interval make it — bounded but variable. This maximizes utilization of both clusters and is how systems like Prime-RL operate (Prime-RL, Async RL & Decentralized Training). The price is the most off-policy of all the rungs; you manage it by bounding the staleness (drop rollouts older than \(s_{\max}\) steps) and by relying on the clip plus, sometimes, an explicit importance-weight correction for the lag.

# Skeleton of an async RL loop (single trainer, N generator workers, a queue).
# This is intentionally minimal; production uses Ray actors + bounded buffers.
import queue, threading, copy

rollout_q = queue.Queue(maxsize=256)     # bounded: backpressure if trainer lags
weight_version = {"v": 0}
shared_weights = {"sd": None, "lock": threading.Lock()}

def generator_worker(engine, prompt_stream):
    while True:
        # 1) refresh local weights if the trainer pushed a newer version
        with shared_weights["lock"]:
            local_v = weight_version["v"]
            if shared_weights["sd"] is not None:
                sync_weights_to_engine(engine, shared_weights["sd"])
        # 2) generate a group, tag it with the weight version that produced it
        prompts = next(prompt_stream)
        for r in rollout(prompts):           # uses the engine (continuous-batched)
            r["gen_version"] = local_v
            rollout_q.put(r)                 # blocks if queue full -> backpressure

def trainer_loop(policy, opt, ref, max_staleness=4, batch_size=512):
    step = 0
    while True:
        # pull a batch of FRESH-ENOUGH rollouts
        batch = []
        while len(batch) < batch_size:
            r = rollout_q.get()
            if step - r["gen_version"] <= max_staleness:   # drop stale rollouts
                batch.append(r)
        # standard experience-prep + minibatch update (Phases 3-4)
        input_ids, resp_mask, beh_lp, rewards = build_experience_batch(batch, ...)
        # ... recompute old_lp/ref_lp, compute advantages, minibatch_update_loop ...
        step += 1
        # push new weights to generators every few steps (Phase 5)
        if step % 2 == 0:
            with shared_weights["lock"]:
                shared_weights["sd"] = copy.deepcopy(policy.state_dict())
                weight_version["v"] = step

The async ladder is the single biggest lever on RL throughput, and choosing a rung is the defining architectural decision of an RL system. Synchronous-colocated is simplest and most on-policy; fully-async is fastest and most off-policy. Most production systems live at rung 2 or a bounded rung 3 — overlapping generation and training with a small, capped staleness, so the clip stays valid while the expensive generators are never idle.

Tying it together: the full synchronous loop¶

Here is the whole synchronous loop assembled from the pieces, so the five phases are visible end to end. This is essentially what a single-controller trainer (veRL-style) executes per step, minus the distributed plumbing.

def rl_outer_step(prompts, golds, *, engine, policy, ref, opt, tok, group_size=8):
    # ---- Phase 1: ROLLOUT (inference engine, continuous-batched) -------------
    rollouts = rollout(prompts)                      # G completions/prompt + beh logprobs

    # ---- Phase 2: REWARD (verifier / RM / sandbox) ---------------------------
    for r, gold in zip(rollouts, expand(golds, group_size)):
        response_text = tok.decode(r["response_ids"], skip_special_tokens=True)
        r["reward"] = reward_fn(r["prompt"], response_text, gold)

    # ---- Phase 3: EXPERIENCE PREP (trainer-side no_grad forwards) -------------
    input_ids, resp_mask, beh_lp, rewards = build_experience_batch(rollouts, pad_id, dev)
    old_lp = token_logprobs(policy, input_ids)       # consistent ratio denominator
    ref_lp = token_logprobs(ref,    input_ids)       # for KL (if used)
    advantages = grpo_advantages(rewards, group_size, normalize_std=False)  # Dr.GRPO

    # ---- Phase 4: LEARN (ppo_epochs x minibatches, the only grad-bearing part) -
    loss, clip_frac, approx_kl = minibatch_update_loop(
        input_ids, resp_mask[:, 1:], old_lp, ref_lp, advantages, policy, opt)

    # ---- Phase 5: WEIGHT SYNC (push new theta into the engine) ----------------
    sync_weights_to_engine(engine, policy.state_dict())

    return {"reward": rewards.mean().item(), "loss": loss,
            "clip_frac": clip_frac, "kl": approx_kl,
            "gen_len_mean": resp_mask.sum(1).mean().item()}

Read top to bottom, this is the entire RL-for-LLM training loop. This is also, almost line for line, the loop that gives Stack-100M its narrow RLVR stage: at 100M parameters the whole thing is colocated on a single GPU — one vLLM engine at gpu_memory_utilization≈0.3 with sleep mode, \(G=8\), short \(L_g\), num_iterations=1, \(\beta=0\), and a synchronous rung-0 schedule, because the model is small enough that generation is cheap and simplicity beats overlap (Post-Training: SFT, DPO, and Narrow RLVR (GRPO) That Works at 100M).

Every other chapter in this Part is an elaboration of one of these five lines: which engine serves Phase 1 (vLLM/SGLang), how Phase 2’s verifiers are built (Reward Engineering, Verifiers & Sandboxes), how Phase 4’s advantages and KL are stabilized (Advantage Estimation, KL Control & Stability Tricks), how Phase 5’s weight sync works across nodes (Colocated vs Disaggregated RL), and how to overlap them all for throughput (Scaling RL: Throughput, Load Balancing & The Latest Tricks).

Further reading¶

• Sheng, Zhang, Ye, et al., HybridFlow (veRL): A Flexible and Efficient RLHF Framework (2024) — the single-controller architecture and the generation/training overlap that this chapter abstracts.
• Hu, et al., OpenRLHF: An Easy-to-Use, Scalable and High-Performance RLHF Framework (2024) — Ray-based actor placement, vLLM-driven rollouts, weight synchronization in practice.
• Kwon, Li, Zhuang, et al., Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM) (2023) — the rollout engine that makes Phase 1 fast.
• Zheng, Yin, Xie, et al., SGLang: Efficient Execution of Structured Language Model Programs (2024) — RadixAttention and prefix sharing, ideal for group sampling.
• Schulman, Wolski, Dhariwal, Radford, Klimov, Proximal Policy Optimization Algorithms (2017) — the clipped surrogate that licenses bounded off-policy reuse.
• DeepSeek-AI, DeepSeek-R1 (2025) and Shao et al., DeepSeekMath (2024) — the GRPO loop whose five phases this chapter dissects.
• Noukhovitch, et al., Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models (2024) — quantifies the staleness-vs-throughput tradeoff of overlapping generation and training.
• HuggingFace TRL (PPOTrainer, GRPOTrainer) and the veRL / OpenRLHF repositories — production implementations of the loop; see TRL: HuggingFace’s RL Library, veRL: HybridFlow & The Single-Controller Architecture, and OpenRLHF, NeMo-Aligner & Ray-Based Systems.

Exercises¶

1. (Conceptual) The chapter’s warning box insists you should recompute the “old”/behavior log-probs on the trainer rather than feed vLLM’s sampling log-probs straight into the importance ratio’s denominator. Suppose you ignore this advice and use the sampler’s log-probs as the denominator and the trainer’s recomputed log-probs as the numerator. At the very first minibatch of epoch 0 — before any optimizer step has moved \(\theta\) — what value should the ratio \(r_{i,t}\) take, what will you actually observe, and why? What downstream metric will most visibly reveal the bug?

2. (Conceptual) Rung 2 of the async ladder (“pipelined / one-step-stale”) is described as costing “exactly one step of staleness.” Explain concretely which weights generated the rollouts the trainer consumes, why one step of staleness is tolerable here, and what specific piece of machinery in the PPO/GRPO objective makes it safe. Then explain why the fully-async Rung 3 needs an explicit max_staleness cap when Rung 2 does not.

3. (Quantitative) Using the chapter’s FLOP/bandwidth accounting, work an RL step for a 13B model in bf16 on a single H100 (HBM bandwidth \(= 3.35\) TB/s; peak bf16 \(= 990\) TFLOP/s dense; assume 45% MFU for training). Batch: \(P = 32\) prompts, \(G = 16\) (so \(B = 512\)), prompt length \(L_p = 256\), mean generation length \(L_g = 512\), ppo_epochs \(E = 2\). Compute: (a) the weight-read time per decode step and the total weight-read-bound decode time over \(L_g\) steps; (b) the training FLOPs \(C_{\text{train}}\) and the training time on a single H100; © comment on the ratio and on what changes if the whole training pass is sharded across 8 GPUs.

4. (Quantitative) A GRPO group has \(G = 4\) completions for one prompt, with rewards \(\mathbf{r} = [1, 0, 1, 0]\) (verifiable pass/fail). (a) Compute the group-relative advantage for each completion without std normalization (Dr.GRPO style, normalize_std=False). (b) Recompute with std normalization (normalize_std=True, eps = 1e-4), using the population std that torch.std(..., unbiased=False) would give. © A second prompt’s group returns \(\mathbf{r} = [1, 1, 1, 1]\) (all correct). What advantage does each member get, with and without std normalization, and why does this case motivate the eps term and the “contested std norm” comment in the code?

5. (Implementation) The chapter’s minibatch_update_loop computes clip_frac as ((ratio - 1.0).abs() > eps_low).float().mean() — but this is unmasked (it averages over prompt/padding tokens too) and it uses the symmetric eps_low even though the loop supports an asymmetric clip range [1 - eps_low, 1 + eps_high] (the DAPO “decoupled clip”). Write a corrected masked_clip_frac(ratio, mask, eps_low, eps_high) that (a) counts a token as clipped only if ratio < 1 - eps_low or ratio > 1 + eps_high, and (b) averages only over response tokens using mask. Then show the one-line change to the diagnostics block. Keep the chapter’s style.

def masked_clip_frac(ratio, mask, eps_low, eps_high):
    """Fraction of RESPONSE tokens whose ratio left [1-eps_low, 1+eps_high].
    ratio, mask: (m, L-1) aligned to next-token targets. mask is 1 on response.
    """
    clipped = (ratio < 1.0 - eps_low) | (ratio > 1.0 + eps_high)  # (m, L-1) bool
    clipped = clipped.float() * mask                              # zero out non-response
    return clipped.sum() / mask.sum().clamp(min=1.0)              # masked mean

        with torch.no_grad():
            clip_frac = masked_clip_frac(ratio, mask, eps_low, eps_high).item()
            approx_kl = (old_lp[mb] - new_lp).mul(mask).sum() / mask.sum()