Inference Systems Advanced topics Notebook

Inference in RL training

Reasoning models (the o1 and DeepSeek-R1 kind) are trained with large-scale reinforcement learning, and every RL step first has the current model generate a large number of answers (rollouts), then updates the model by their rewards. Rollout usually takes more than half of RL training time, so the inference engine has become a core component of the training system, bringing new problems absent from online serving: long-tail generations slowing the whole batch, training and inference computing inconsistent probabilities, syncing new weights to the inference engine every step, and training and inference competing for the memory of the same GPUs. This chapter explains these problems with simulations and measurements.

Self-test: if you can answer these, skip this chapter

1. Why is GPU utilization during rollout often low? What mitigations are there?
2. For tokens sampled by the inference engine, if the training side recomputes their log probabilities, are the results identical? How different? What does this mean for RL?
3. When training and inference share GPUs, how is memory switched between them? How are weights synced?
4. How do synchronous and asynchronous RL differ? What does asynchrony cost?

Answers (try first, then expand to compare)

1. Answer lengths are long-tailed: each batch waits for its few longest answers, and most of the time only a handful of requests are still running, so the batch is small and GPU utilization low. Mitigations: partial rollout (carry unfinished answers over to the next round), asynchronous RL, over-sampling (send extra requests and stop once there are enough).
2. Not exactly: the inference and training sides differ in kernels, batching and parallelism; in BF16 the same token's probability measurably differs by about 10%, and by more when kernels and parallelism differ. This effectively turns RL off-policy with biased gradients; correct it with importance sampling (TIS, MIS), or use batch-invariant kernels consistent with the training side.
3. When co-located, the inference engine goes to sleep during the training phase: it frees the physical memory held by KV and weights while keeping the virtual addresses (sleep), and wakes up after training (wake). Weights are synced to the inference engine via CUDA IPC (same machine), NCCL broadcast or layer-by-layer pipelining, re-sharding when the training and inference sides partition differently.
4. Synchronous RL generates with the latest weights every step and trains only after generation finishes; asynchronous RL overlaps generation and training (samples come from slightly older weights). The cost of asynchrony is policy staleness: samples are not generated by the current policy, requiring importance correction or a cap on staleness.

One RL training step

Figure: one RL step: rollout, reward, training, weight sync; the longest answer sets the step's time

Take GRPO as an example (the post-training chapter of the LLM handbook covers the algorithm):

1. rollout: for a batch of questions, sample G answers per question (for example 512 questions × 8 answers); this is pure inference, done by vLLM or SGLang;
2. Scoring: score each answer with rules or a reward model, and compute advantages relative to the group;
3. Computing probabilities: the training framework (FSDP, Megatron) runs one forward pass over all answers to get the log probabilities under the current policy (and the reference model);
4. Update: backpropagation and the optimizer update;
5. Weight sync: pass the new weights to the inference engine and move on to the next step.

The training framework and the inference engine can be co-located (taking turns on the same GPUs, verl's default) or disaggregated (each with its own GPUs, common in slime, AReaL and others). For the training side's view (a complete minimal GRPO loop, re-sharding weights between training and inference, the trade-offs of co-location versus disaggregation), see training frameworks and RL training systems in the distributed training handbook.

Problem one: the long tail

Reasoning models' answer lengths vary enormously: most are hundreds to thousands of tokens, a few run past ten thousand. Rollout must wait for the longest one to finish before the next step. Simulate an inference engine (a 7B model on one GPU, at most 512 concurrent) processing 1024 answers with log-normally distributed lengths:

import math
import random

random.seed(0)
params, weights, kv_per_token, bw, peak, prompt_len = 7.6e9, 15.2e9, 57344, 3.35e12, 989e12, 512

def step_time(batch, context_tokens):              # one decode step: roofline estimate
    return max(2 * params * batch / (peak * 0.5), (weights + kv_per_token * context_tokens) / bw) + 0.4e-3

lengths = [min(16384, max(64, int(random.lognormvariate(math.log(1500), 0.9)))) for _ in range(1024)]

def rollout(lengths, max_seqs=512, cap=None):
    queue, running, elapsed, timeline, tokens = list(range(len(lengths))), {}, 0.0, [], 0
    target = [l if cap is None else min(l, cap) for l in lengths]
    while queue or running:
        while queue and len(running) < max_seqs:
            running[queue.pop(0)] = 0
        dt = step_time(len(running), sum(prompt_len + g for g in running.values()))
        elapsed += dt
        timeline.append((dt, len(running)))
        tokens += len(running)
        for i in list(running):
            running[i] += 1
            if running[i] >= target[i]:
                del running[i]
    return elapsed, timeline, tokens

elapsed, timeline, tokens = rollout(lengths)
full_speed = 512 / step_time(512, 512 * (prompt_len + 1500))              # tokens generated per second at a full batch
low = sum(dt for dt, batch in timeline if batch < 128) / elapsed
print(f"回答长度：中位数 {sorted(lengths)[512]}，最长 {max(lengths)}")
print(f"rollout 用时 {elapsed:.0f} s；如果全程保持满批次，只需 {tokens / full_speed:.0f} s")
print(f"有 {low:.0%} 的时间，批次中的请求不到 128 个（容量的四分之一）")
capped, _, capped_tokens = rollout(lengths, cap=4096)
print(f"每轮最多生成 4096 个 token（partial rollout）：本轮用时 {capped:.0f} s，"
      f"{sum(l > 4096 for l in lengths)} 个回答留到下一轮继续")

output

回答长度：中位数 1549，最长 16384
rollout 用时 203 s；如果全程保持满批次，只需 102 s
有 54% 的时间，批次中的请求不到 128 个（容量的四分之一）
每轮最多生成 4096 个 token（partial rollout）：本轮用时 96 s，145 个回答留到下一轮继续

More than half the time is spent in the phase where "only a few long answers remain in the batch": decode is then badly memory-bound, and most of the GPU's capacity is wasted. Common mitigations:

• partial rollout: cap generation in each round, save the state of unfinished answers, and continue them in the next round (with a slightly older policy, introducing a little off-policy);
• Asynchronous RL: the inference engine generates nonstop, and the training side starts updating as soon as it has enough samples, pipelining the two (the asynchronous modes of AReaL and slime); the cost is again policy lag, corrected with importance sampling and similar methods;
• Over-sampling: send more requests than needed, and once there are enough, drop (or carry over to the next round) the remaining long tail;
• Larger inference scale, PD disaggregation, speculative decoding: in the long-tail phase the batch is small, exactly where speculative decoding is most effective.

Problem two: inconsistent probabilities between training and inference

RL's loss uses "the probability at sampling time". In theory, the probability when the inference engine samples should equal exactly the probability the training framework recomputes for the same sequence; but the two compute differently: inference computes token by token with a KV Cache, training does one forward pass over the whole sequence; inference may use different kernels, different batch composition, even quantization. Measure it on the same model: sample 120 tokens step by step with a KV Cache, record each token's log probability, then recompute with one forward pass over the whole sequence:

import torch
from transformers import AutoTokenizer
from mini_llm import KVCache, Transformer

torch.set_num_threads(16)
path = "models/Qwen3-0.6B"
tok = AutoTokenizer.from_pretrained(path)
prompt = tok(tok.apply_chat_template([{"role": "user", "content": "1 到 100 的和是多少？请一步步推理。"}],
                                     tokenize=False, add_generation_prompt=True, enable_thinking=False)).input_ids
for dtype in (torch.float32, torch.bfloat16):
    model = Transformer.from_pretrained(path, dtype=dtype)
    gen, seq, rollout_logp = torch.Generator().manual_seed(0), list(prompt), []
    cache = KVCache(model.cfg.num_hidden_layers)
    with torch.no_grad():
        logits = model(torch.tensor([prompt]), cache)[0, -1].float()
        for _ in range(120):                                 # inference side: sample token by token, recording the log probability at sampling time
            probs = logits.softmax(-1)
            t = torch.multinomial(probs, 1, generator=gen).item()
            rollout_logp.append(probs[t].log().item())
            seq.append(t)
            logits = model(torch.tensor([[t]]), cache)[0, -1].float()
        full = model(torch.tensor([seq]))[0].float().log_softmax(-1)   # training side: one forward pass over the whole sequence
    train_logp = torch.tensor([full[len(prompt) - 1 + i, seq[len(prompt) + i]].item() for i in range(120)])
    diff = train_logp - torch.tensor(rollout_logp)
    print(f"{str(dtype):15s} |Δ log p| 平均 {diff.abs().mean():.1e}、最大 {diff.abs().max():.1e}；"
          f"重要性比 p_train/p_rollout 的范围 [{diff.exp().min():.3f}, {diff.exp().max():.3f}]")

torch.float32   |Δ log p| 平均 1.2e-06、最大 1.2e-05；重要性比 p_train/p_rollout 的范围 [1.000, 1.000]
torch.bfloat16  |Δ log p| 平均 9.7e-03、最大 1.0e-01；重要性比 p_train/p_rollout 的范围 [0.909, 1.107]

In FP32 the two nearly agree; in BF16, the same token's probability can differ by 10% between the two computations. This is not a bug but the normal error of different accumulation orders at BF16 precision (see which optimizations change the output). For RL, though, it means "the policy used to compute the loss" and "the policy that actually sampled" are not the same, making training implicitly off-policy, which can accumulate into instability or even collapse. Responses in real systems:

• Importance sampling correction: weight the loss by the ratio \(p_\text{train}/p_\text{rollout}\), truncating ratios that are too large (TIS), or simply drop samples whose deviation is too large (MIS). This requires the inference engine to return the logprobs at sampling time;
• Make both sides compute the same: use "batch-independent" (batch-invariant) kernels on the inference side, or the same kernels and computation order on both sides; SGLang and vLLM both offer deterministic/batch-invariant modes, at some cost in performance;
• Raise precision: use FP32 for critical parts (such as the lm_head and the logits computation).

Problem three: switching memory and syncing weights

In co-located deployments, training and inference take turns on the same GPUs:

• Switching memory: after rollout, the inference engine must free the memory held by the KV Cache (and the weights, if the training side has its own copy) for training, and restore it after training. vLLM's sleep mode (enable_sleep_mode, llm.sleep() / llm.wake_up()) and SGLang's release_memory_occupation / resume_memory_occupation (based on torch_memory_saver) are designed for this: free the physical memory but keep the virtual addresses, so restoring needs no re-capture of CUDA Graphs;
• Weight sync: after each training step, the new weights must get into the inference engine. Co-located, tensors in GPU memory can be shared directly through CUDA IPC handles; disaggregated, through NCCL broadcast or a dedicated transfer service. A large model's weights are hundreds of GB, so sync time is not negligible: pipeline it layer by layer, send only what changed, and handle the re-sharding caused by training and inference using different parallelism (different TP and EP splits).

Source code

• vLLM: RL examples are in examples/rl/ (including three ways to sync weights: HTTP, NCCL and IPC), and vllm/distributed/weight_transfer/ implements weight transfer; the Rust frontend also supports the weight-sync lifecycle for RL. Logprobs at sampling time are returned via SamplingParams(logprobs=...), and logprobs_mode decides whether raw or processed values are returned.
• SGLang: srt/weight_sync/ and srt/checkpoint_engine/ (weight updates), srt/batch_invariant_ops/ (batch-invariant operators); interfaces such as update_weights_from_tensor and update_weights_from_distributed are called by training frameworks such as verl and slime.

How to explain it

RL infrastructure has been a hot direction for inference roles since 2025. To explain "what special support does an inference engine need in RL training", go through the four problems: the long tail (partial rollout, asynchronous RL, over-sampling; ideally with a quantitative sense like "half the time goes to the long tail"), inconsistent probabilities (differences from computation order in BF16, importance sampling correction, batch-invariant kernels), switching memory (sleep/wake, freeing physical memory while keeping virtual addresses) and weight sync (IPC, NCCL, layer-by-layer pipelining, re-sharding). Being able to describe the architecture of any of verl, slime or AReaL earns big bonus points.

Exercises

1. Why can weight sync be fast when co-located? The training side uses FSDP (parameters sharded across GPUs) and the inference side TP=2, both on the same GPUs. What does syncing the weights involve?

Approach

Under FSDP each GPU holds only a shard of the parameters, so first gather a layer's full parameters (all-gather), then split them by the inference side's TP and write them into the inference engine's corresponding parameter tensors. Since it is the same GPUs, this can go layer by layer: gather a layer, hand the tensor handle to the inference process through CUDA IPC (no copy needed), let the inference process take its split, then free it. The whole process moves data only between and within GPUs, without going through the CPU or writing to disk.

2. Policy lag in asynchronous RL. In asynchronous RL, some answer was generated by the policy from k steps ago. What problems does this cause? How do you control it?

Answer

The training samples no longer come from the current policy, so the policy gradient estimate is biased, more so the larger the lag, and at worst training becomes unstable. Controls: cap the maximum lag (for example at most 1–2 steps behind the training side), correct with importance sampling and truncate the ratios, and drop or down-weight samples that lag too much. This is the same toolkit as correcting "inconsistent probabilities between training and inference", which is why it is so critical that the inference engine returns accurate sampling logprobs.

Summary

• Every RL step starts with a rollout, making the inference engine a core component of RL training systems; training and inference can be co-located or disaggregated.
• The long tail of answer lengths keeps rollout in a small-batch, low-utilization state most of the time; partial rollout, asynchronous RL and over-sampling mitigate it.
• In BF16 the inference and training sides' probabilities measurably differ by about 10% (more when kernels and parallelism differ), requiring importance sampling correction or batch-invariant kernels.
• Co-located deployments need memory switching (sleep/wake) and efficient weight sync (IPC, NCCL, layer-by-layer pipelining).

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

• 156. RL rollout 的前缀共享与并发 (easy)

Found a mistake or a better way to say it? Edit this page on GitHub · Report an issue