vLLM V1 Migration: Ensuring Backend Correctness in Reinforcement Learning

vLLM V1 Migration: Ensuring Backend Correctness in Reinforcement Learning

ServiceNow AI successfully migrated its RL inference engine from vLLM V0 to V1 by prioritizing backend correctness over objective-side corrections. The team achieved parity between vLLM 0.8.5 (V0) and vLLM 0.18.1 (V1) by resolving four critical discrepancies: processed rollout logprobs, V1-specific runtime defaults, the inflight weight-update path, and the precision of the lm_head projection.

The Migration Objective

When using vLLM as an inference engine for rollout generation in systems like PipelineRL, the engine samples tokens and returns logprobs that the trainer uses to compute essential RL metrics, including policy ratios, KL divergence, clip rate, and entropy. Any discrepancy in how these logprobs are computed creates a train-inference mismatch that can fundamentally alter training dynamics.

To migrate to vLLM V1, ServiceNow AI established a narrow objective: verify that V1 returned rollout logprobs in the form the trainer expected, rerun the workload against the V0 reference, and evaluate objective-level changes only after backend parity was restored. Initial attempts with V1 showed significant deviations in reward, entropy, and clip rate compared to the V0 reference.

Backend Failure Modes

The team categorized potential causes of the train-inference mismatch into three layers:

  1. Semantic mismatch: The backend returns logprobs with a different meaning than what the trainer expects.
  2. Inference-path mismatch: Differences in runtime defaults for caching, scheduling, or request handling cause prompts to follow different execution paths.
  3. Objective mismatch: The RL objective requires correction for remaining staleness or backend mismatch.

By treating the first two as backend behavior problems and ruling them out first, the team avoided the mistake of prematurely applying objective-side corrections to mask backend errors.

vLLM V1 Backend Fixes

Logprob Semantics

By default, vLLM V1 returns logprobs from raw model outputs before logits post-processing (such as temperature scaling and top-k/top-p filtering). PipelineRL requires logprobs from the processed distribution used by the sampler. This was resolved by setting logprobs-mode=processed_logprobs, which eliminated the mean offset in rollout logprobs and centered the mean policy ratio close to 1.0.

Runtime Defaults

To ensure parity, the team explicitly disabled V1-specific defaults that differed from the V0 reference path:

  • Prefix Caching: Disabled (enable-prefix-caching: false) to prevent the reuse of state computed before a weight update, which would introduce a V1-only difference in cache lifetime.
  • Async Scheduling: Disabled (async-scheduling: false) to remove another V1-only degree of freedom.

Inflight Weight Updates

To match the V0 behavior of loading new weights without explicit cached-state invalidation, the team implemented the following V1 update path:

await engine.pause_generation(mode="keep", clear_cache=False)
await engine_client.collective_rpc_async(
    "receive_weight_update",
    args=(request.model_dump_json(),),
)
await engine.resume_generation()

Using mode="keep" and clear_cache=False ensured that the weight synchronization model matched the V0 reference, reducing persistent lag later in training.

Numerical Parity: The fp32 lm_head

Even after backend fixes, final parity required matching the numerical path used to compute logits. The trainer used an fp32 lm_head for the final projection, and the rollout backend had to be updated to match this precision.

This precision is critical because RL updates consume token logprobs directly; small changes in logits can significantly impact policy ratios and KL divergence. This finding aligns with the MiniMax-M1 technical report and the ScaleRL paper, both of which identify fp32 head computation as a necessary design choice for large-scale RL to prevent training/inference token-probability mismatches.

Why Backend Correctness Precedes Objective Correction

While tools like truncated importance sampling or importance-ratio reweighting can correct for stale or asynchronous rollouts, applying them before fixing the backend would have confounded the results. If an objective-side correction is used to compensate for broken inference-backend behavior, the training curve becomes harder to interpret.

The team's approach ensures that any remaining mismatch is a result of the RL objective's design (e.g., async/off-policy issues) rather than a bug in the inference engine. Future improvements will focus on async/off-policy cleanup, such as keeping explicit behavior-policy logprobs from rollout time and recomputing trainer-side old-policy logprobs at optimization time.

Sources