Skip to content

Async Reinforcement Learning

Overview

In a standard RL training loop, generation and training happen sequentially: the policy generates rollouts, then training runs on those rollouts, and the cycle repeats. During generation the training accelerators sit idle, and vice versa.

The one-off pipelining approach separates the generation and training phases into two parallel coroutines, allowing the model to generate new samples while simultaneously training on previously generated data. This can lead to better GPU utilization and greater training throughput.

However, this overlap introduces a complication: weights must be updated in the inference engine mid-flight, while requests may still be in progress.

The Pause and Resume API

To safely update weights while the inference engine is running, vLLM provides pause_generation and resume_generation methods. These let the trainer coordinate a clean window for weight synchronization without losing in-flight work.

pause_generation

await engine.pause_generation(mode="keep", clear_cache=True)

The mode parameter controls how in-flight requests are handled, and whether new requests are admitted while paused:

Mode In-flight requests New requests
"abort" Aborted immediately, returning partial results (default) Rejected with HTTP 503 until resume
"wait" Running requests finish before the pause completes; queued ones run after resume Rejected with HTTP 503 until resume
"keep" Frozen in the queue; they resume when resume_generation is called Accepted and queued

"abort" and "wait" treat the pause as a generation boundary: from the moment the pause starts (for "wait", including while in-flight requests drain) until resume_generation returns, the engine rejects new requests from every API server rather than carrying them across the boundary. Clients should retry on 503. A request racing the pause call itself may instead be treated as in-flight, and input to an already-open streaming-input session is still accepted. Other pauses, such as sleep, keep queueing new requests.

The clear_cache parameter controls whether to clear the KV cache and prefix cache after pausing.

resume_generation

await engine.resume_generation()

Resumes the scheduler after a pause. Any requests frozen with mode="keep" will continue generating.

HTTP Endpoints

With VLLM_SERVER_DEV_MODE=1, the vLLM HTTP server exposes the same functionality via:

  • POST /pause?mode=keep - Pause generation
  • POST /resume - Resume generation
  • POST /abort_requests - Abort in-flight requests without pausing the scheduler (send {} to abort all, or {"request_ids": [...]})
  • POST /update_weight_version - Set the weight_version label without changing weights (send {"new_version": "v8"})
  • GET /weight_info - Return the latest committed weight_version

Data Parallelism

When using data parallelism with vLLM's internal load balancer (i.e. data_parallel_backend="ray"), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an external load balancer (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to every engine instance individually before and after the weight update, and set the weight version on every instance too.

Typical Async RL Flow

A typical async RL loop with weight syncing looks like this:

  1. Start generating rollouts from the current policy
  2. Once trainer has new weights to update to, pause generation with mode="keep"
  3. Sync the updated weights from the trainer to the inference engine (see Weight Transfer) and set their weight version
  4. Resume generation -- in-flight requests continue with the new weights
  5. Repeat

The key insight is that requests paused with mode="keep" will produce tokens from the old weights before the pause and tokens from the new weights after resume. The clear_cache parameter controls whether the KV cache is invalidated during the pause. When clear_cache=True, previously cached key-value entries are discarded, so all tokens generated after resume will be computed entirely with the new weights. When clear_cache=False, existing KV cache entries are retained, meaning some tokens in context may still reflect the old weights (stale KV cache).

Weight Versions in Outputs

To tell which tokens came from which weights, set a weight-version label while generation is paused, after the new weights are loaded: pass weight_version to finish_weight_update, or call update_weight_version (POST /update_weight_version over HTTP). Labels are strings you choose, such as the trainer step; vLLM doesn't order or check them.

The final output of every sequence then carries weight_versions, one span per label over its output-token indexes. A request paused with mode="keep" across one update gets two spans:

completion = final_output.outputs[0]
for span in completion.weight_versions:
    # The tokens in [span.start, span.end) were sampled by weights `span.version`.
    print(span.version, completion.token_ids[span.start : span.end])

The offsets cover the whole output, so with streamed delta outputs, collect the tokens before slicing. /inference/v1/generate returns the same list on each choice; see Tokens In <> Tokens Out API.

/v1/completions and /v1/chat/completions return it on each choice when the request sets return_weight_versions: true. The offsets count the choice's output tokens, which return_token_ids returns as token_ids, and a streamed choice gets the list on its last chunk.

An aborted request returns spans over its partial output. If you resubmit that output as part of a new prompt, the new request's spans start again at 0, so shift them by the number of tokens already generated before joining the two lists.

Example

The async RLHF example demonstrates this pattern with vllm.AsyncLLMEngine, NCCL weight transfer, and mid-flight pause/resume with validation.