Inference rollouts

A rollout is the policy acting: sampling completions for a prompt, or taking an agent’s turns in an environment. In LLM RL, generating rollouts usually takes more time than the policy update, so the rollout tier decides how fast the loop runs.

Rollouts need three things from the platform:

• An inference engine on a GPU, such as vLLM, that batches many sequences per call.
• The current policy weights, which change every iteration.
• Parallelism, so many prompts generate at once.

Generate a group per call

GRPO compares several completions of the same prompt, so ask the engine for the whole group in one call. vLLM’s SamplingParams(n=GROUP_SIZE) samples them as one batch, and the task returns them together for scoring.

import flyte

rollout_env = flyte.TaskEnvironment(
    name="rollout",
    image=image,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu="L4:1"),
)


@rollout_env.task(retries=2)
async def generate_batch(
    base: flyte.io.Dir, prompts: list[str], adapter: flyte.io.Dir, version: int, seed: int
) -> list[list[str]]:
    from vllm import LLM, SamplingParams
    from vllm.lora.request import LoRARequest

    engine = LLM(model=await base.download(), enable_lora=True)
    lora = LoRARequest(f"policy-v{version}", version + 1, await adapter.download())
    sampling = SamplingParams(n=GROUP_SIZE, temperature=1.0, seed=seed)
    outputs = engine.generate(prompts, sampling, lora_request=lora)
    return [[o.text for o in out.outputs] for out in outputs]

Move weights from the trainer to the generators

Pass the weights as a directory

The trainer returns the new policy as a flyte.io.Dir, and the next round of rollout tasks receives it as an input:

1. The train task saves the weights, or only the LoRA adapter, and returns await flyte.io.Dir.from_local(path).
2. The driver passes that Dir to the next round of rollout calls.
3. Each rollout task downloads it, or reads it by its path, before generating.

The weights travel through your object store, so the trainer and the generators never have to run at the same time or share a network connection. A LoRA adapter is a few megabytes, so the transfer is fast even when the base model is large. With full-weight updates, write the checkpoint once from rank 0 and pass the same Dir to every worker. To load full weights onto the GPU without writing them to local disk first, pass the Dir’s path to the model streamer. See Stream model weights.

Keep rollouts on-policy

Training on rollouts from an older policy biases the update. Record which policy version produced each rollout, so the trainer can detect stale samples or correct for them. Pass the version number into the rollout task and return it with each rollout. To compute importance ratios, also store the token IDs and sampling log-probabilities the engine returns: set SamplingParams(logprobs=0), which returns only the sampled token’s log-probability, and keep it with the completion. Sample at temperature=1.0 with top_p=1.0: vLLM reports log-probabilities after its sampling adjustments, so they equal the policy’s own only at those settings. See Overlap generation with training for the loss that uses them.

Related

• Durable agentic RL with SkyRL and Harbor: rollouts that are multi-turn agent trajectories.
• Training environment: where the policy’s actions run.
• Training loop: how the driver fans out rollouts and overlaps them with training.
• Batch inference: the same engine patterns for scoring a dataset.