IsoExec: Eliminating Trainer-Inference Mismatch in SkyRL

vLLM has introduced IsoExec, a cross-framework unified execution abstraction designed to eliminate the numerical mismatch between training and inference engines in Reinforcement Learning (RL) workloads. By enforcing a strict execution contract and utilizing parallelism-invariant kernels, IsoExec ensures that the rollout engine and the trainer evaluate the same policy with bitwise consistency, reducing the average end-to-end rollout-versus-training logprob difference to below $10^{-7}$ with a 25% performance overhead.

The Problem: Trainer-Inference Mismatch

In on-policy RL, the system must execute the same policy twice: once in a rollout engine to sample tokens and once in a trainer to recompute log probabilities. In practice, these two stages often use different engines (e.g., vLLM for rollout and Megatron for training), which employ different kernels, batch shapes, and parallelism layouts.

Because floating-point arithmetic is non-associative, these systemic differences lead to different reduction orders, changing token probabilities even when the model parameters are identical. This mismatch can destabilize RL algorithms like REINFORCE and GRPO and cause reward collapse. For instance, a GLM-5.2 run reported by Fireworks saw a train-inference KL of approximately 0.013, resulting in 45% of tokens being discarded via clipping and subsequent reward collapse around step 20.

The IsoExec Unified Execution Contract

IsoExec resolves these discrepancies using an execution contract that declares every bit-relevant execution choice both runtimes must follow. This framework-independent contract ensures that rounding-sensitive details are identical across engines.

Contract Structure and Enforcement

  • Cases and Regions: The contract partitions forward operators into "regions" (arithmetic spans implemented by one kernel). It then defines "cases" (e.g., engine_prefill vs. trainer_fwd) to specify how each region is handled.
  • Composition: For every region-case pair, the contract selects a specific implementation and pins it to constants, such as accumulation dtypes and reduction-decomposition parameters (e.g., split-K partition counts).
  • Claims: These are runtime-enforced conditions, such as topology claims that list the specific parallel sizes for which a reduction tree is proven bitwise invariant.
  • Identities: SHA-256 digests (semantic, numerical_policy, and deployment) are used to verify that both the trainer and rollout engines are executing the same verified numerical policy.

A per-runtime contract adapter binds these specifications to the framework's extension points and monitors the runtime to ensure compliance.

Unified Model and Parallelism-Invariant Kernels

IsoExec implements a unified model definition that remains bitwise consistent across various distributed strategies, including tensor, expert, and sequence parallelism.

Achieving Parallelism Invariance

To maintain numerics across different GPU layouts, IsoExec employs a fixed binary reduction tree approach:

  • Tensor Parallelism (TP): Using the pik implementation, IsoExec divides the K dimension into contiguous leaves. Each leaf uses deterministic Tensor Core MMA with FP32 accumulation, and the contract fixes the rank-to-leaf mapping and arithmetic schedule.
  • Expert Parallelism (EP): Expert outputs are combined in a fixed routing order rather than rank order.
  • Sequence Parallelism (SP): IsoExec reuses the same reduction tree as non-SP systems, with each rank retaining its own output slice, ensuring trainer logits remain identical regardless of whether SP is enabled.

Chunkwise-Parallel Recurrent (CPR) Gated DeltaNet

Linear-attention architectures like Gated DeltaNet (GDN) typically use chunkwise-parallel forms for training/prefill and recurrent forms for decode, creating numerical mismatch. Previous attempts to fix this by using recurrent forms everywhere resulted in significant slowdowns (up to 5x on some workloads).

IsoExec introduces Chunkwise-Parallel Recurrent (CPR), which maintains the recurrence as the main function but evaluates it in parallel across chunks. For decode, it resynchronizes the hidden state every $C$ tokens (where $C$ is the chunk size). This achieves bitwise exactness while maintaining high throughput:

Stage Native Mixed Chunkwise Everywhere Recurrent Everywhere CPR
Bitwise Exact No Yes Yes Yes
Trainer Fwd+Bwd 5.177 ms 5.177 ms (1.00x) 22.863 ms (4.42x) 7.386 ms (1.43x)
Rollout Prefill 0.844 ms 0.844 ms (1.00x) 3.639 ms (4.31x) 1.412 ms (1.67x)
Rollout Decode 0.0612 ms 2.2374 ms (36.6x) 0.0612 ms (1.00x) 0.0846 ms (1.38x)

Experimental Results

IsoExec was tested on a single 8xH100 node training Qwen3.5-35B-A3B on DAPO-Math-17k using synchronous RL.

Numerical Accuracy

Across 50 steps, the mean pre-update rollout-versus-training absolute logprob difference was reduced from $10^{-3}$ (native SkyRL) to below $10^{-7}$ (IsoExec). The average per-step maximum difference dropped from 5.073 to $10^{-4}$.

Performance Overhead

Eliminating the mismatch introduced a performance penalty compared to the native SkyRL stack:

Metric Native IsoExec Overhead
Generation 591.3 s 776.6 s 31.3%
Policy Training 498.6 s 591.3 s 18.6%
Full RL Step 1224.6 s 1534.0 s 25.3%

While numerical consistency was achieved, the team did not observe a meaningful reward improvement over the short 50-step test window.

Sources

Related

  • Dispatch
    PipelineRL: Optimizing LLM Reinforcement Learning via Inflight Weight Updates

    Hugging Face and ServiceNow Research have open-sourced PipelineRL, an experimental RL implementation that uses inflight weight updates to eliminate the trade-off between inference throughput and on-policy data collection.

  • Dispatch
    Keep the Tokens Flowing: Lessons from 16 Open-Source RL Libraries

    Hugging Face surveyed 16 open-source RL libraries and found that async RL training separates inference and training onto different GPU pools, uses a rollout buffer, and pushes weights asynchronously, with Ray dominating orchestration and NCCL broadcast the common weight‑sync method.

  • Project
    ServiceNow/PipelineRL

    A scalable asynchronous reinforcement learning framework for LLM agents that uses in-flight weight updates to maximize GPU utilization while maintaining on-policy data freshness.

  • Dispatch
    vLLM V1 Migration: Ensuring Backend Correctness in Reinforcement Learning

    ServiceNow AI achieved parity between vLLM V0 and V1 for RL rollout generation by fixing logprob semantics, runtime defaults, weight update paths, and implementing an fp32 lm_head.

  • Dispatch
    vLLM Elastic Expert Parallelism

    vLLM introduces Elastic Expert Parallelism (Elastic EP), enabling Mixture-of-Experts (MoE) deployments to scale the number of workers up or down at runtime without requiring a server restart.