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.

Keep the engine loaded between calls

Loading a base model into vLLM takes much longer than generating a group, so don’t do it on every call. A reusable container keeps a pool of rollout workers running between calls:

rl_grpo_lora.py
                
rollout_env = flyte.TaskEnvironment(
    name="rl-grpo-rollout",
    image=image,
    resources=flyte.Resources(cpu=4, memory="24Gi", gpu=flyte.GPU("L4", 1), shm="auto"),
    reusable=flyte.ReusePolicy(replicas=(1, 4), concurrency=1, idle_ttl=300, scaledown_ttl=120),
    secrets=[HF_SECRET],
    env_vars={"VLLM_USE_V1": "1"},
)

replicas=(1, 4) scales the pool from one to four GPU workers with the number of groups in flight.
concurrency=1 sends one call at a time to each worker. A single vLLM engine batches internally and isn’t safe to drive from several calls at once.
idle_ttl=300 shuts the whole pool down after 300 seconds with no calls, and scaledown_ttl=120 removes each idle worker after 120 seconds. Set idle_ttl longer than one training step, so the pool survives while the trainer runs. To keep every worker warm across the step, not just one, set scaledown_ttl longer than the step too.

Load the engine once per worker and cache it at module level:

rl_grpo_lora.py
                
                    
                
            
                
            

                
                    
                
            
@alru_cache(maxsize=1)
async def _load_engine(base_uri: str) -> Any:
    """Build the vLLM engine once per warm replica (cached); the frozen base stays resident in GPU."""
    from vllm import LLM

    local_base: str = await flyte.io.Dir.from_existing_remote(base_uri).download()  # plain-HF base
    logger.info("Building warm vLLM engine from %s", local_base)
    return LLM(
        model=local_base,
        enable_lora=True,
        max_lora_rank=LORA_RANK,
        max_loras=1,
        trust_remote_code=True,
        gpu_memory_utilization=0.85,
        max_model_len=2048,
        enforce_eager=True,  # skip CUDA-graph capture → faster cold start for an MVP
    )

@alru_cache(maxsize=None)
async def _adapter_local_path(adapter_uri: str) -> str:
    """Download a LoRA adapter once per warm replica (cached by its remote URI → one download/version)."""
    return await flyte.io.Dir.from_existing_remote(adapter_uri).download()

The first call on a worker builds the engine; every later call reuses it. The image needs the unionai-reuse package for reusable environments.

Generate with the current adapter

Each call attaches the current LoRA adapter to the resident base model, by version:

rl_grpo_lora.py
                
                    
                
            
                
            

                
                    
                
            
@rollout_env.task
async def generate(
    base: flyte.io.Dir,
    question: str,
    answer: str,
    adapter: flyte.io.Dir,
    version: int,
    group_id: int,
) -> list[Rollout]:
    """Generate a GROUP_SIZE group of completions for one prompt, using the current LoRA adapter.

    The frozen base loads exactly once per replica (cached ``_load_engine``); each adapter version is
    downloaded once (cached ``_adapter_local_path``) and attached per request via ``LoRARequest`` — the
    base weights in GPU memory are never touched.
    """
    from vllm import SamplingParams
    from vllm.lora.request import LoRARequest

    engine: Any = await _load_engine(base.path)
    adapter_path: str = await _adapter_local_path(adapter.path)

    # lora_int_id must be >= 1 and unique per adapter; version starts at 0 so shift by 1.
    lora: LoRARequest = LoRARequest(f"policy-v{version}", version + 1, adapter_path)
    sampling: SamplingParams = SamplingParams(
        n=GROUP_SIZE,
        temperature=SAMPLING_TEMPERATURE,
        top_p=1.0,
        max_tokens=MAX_NEW_TOKENS,
    )

    prompt: str = build_prompt(question)
    # vLLM's generate() is a blocking call, so we run it via asyncio.to_thread to keep it off the event
    # loop — that lets the reusable replica's background actor heartbeat stay responsive while the GPU
    # is busy. (We deliberately do not use flyte.extras.DynamicBatcher here: it batches many concurrent
    # producers, whereas this env runs concurrency=1 and each call already submits a full group of
    # GROUP_SIZE sequences as one vLLM batch.)
    outputs: Any = await asyncio.to_thread(engine.generate, [prompt], sampling, lora_request=lora)
    completions: list[str] = [o.text for o in outputs[0].outputs]
    logger.info("group %s: generated %d completions (adapter v%d)", group_id, len(completions), version)
    return [
        Rollout(group_id=group_id, question=question, completion=c, answer=answer) for c in completions
    ]

Each adapter version is downloaded once per worker. The base weights in GPU memory never reload, so moving to a new policy costs one small download.

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:

The train task saves the weights, or only the LoRA adapter, and returns await flyte.io.Dir.from_local(path).
The driver passes that Dir to the next round of rollout calls.
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.

When a volume helps

A volume is a versioned file system that tasks mount. For moving weights, it doesn’t beat a Dir:

Speed. A rollout task reads the whole adapter or checkpoint, which a Dir download or the model streamer already does at object-store speed. Reads through a volume pass through a file-system layer, and loading weights into memory that way is slower than streaming them. See Performance and trade-offs.
Durability. Both are stored in your object store, and both are immutable once a task returns them.

A volume does help in these cases:

Many rollout pods on one node. With allow_volumes(shared_node_cache=True), read-only mounts on a node share one cache, so each node fetches the weights once instead of once per pod. This matters when you pack several single-GPU rollout pods onto each multi-GPU node.
Reading part of a large tree. A task that needs only some files, such as one shard of a sharded checkpoint, mounts the volume and reads those files, instead of downloading the whole Dir.
Data that grows across iterations. A rollout buffer that every iteration adds to can live on a volume, with each version storing only what changed. See Keep rollouts on a volume.
A browsable history of policies. A volume declared as an artifact publishes each committed version with a parent edge to the one before it. See Versions, artifacts and lineage.

A volume has one writer while it is mounted, so in a multi-node trainer only one rank can write to it.

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.

RL for LLMs with GRPO and LoRA: the rollout pool on this page, with the full loop around it.

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.