IsoExec: Unified Execution to Eliminate Trainer-Inference Mismatch in SkyRL

11 min read
Alexander Jiang and the SkyRL Team

TL;DR

In theory, on-policy RL assumes that rollout and training evaluate the same policy. In practice, RL training systems often use two separate engines for rollout and training, with different model definitions, kernels, batch shapes, and parallelism layouts. Because floating-point arithmetic is non-associative, those differences can change token probabilities even when executing the same policy. This can make new RL algorithms, changes in harnesses and environments, and improvements to RL infrastructure and hardware kernels hard to debug.

To resolve this problem, we introduce IsoExec, a cross-framework unified execution abstraction for eliminating mismatch between the training and inference engines for RL workloads. IsoExec has two components: an execution contract that specifies and enforces the execution details affecting floating-point rounding across engines, and a unified model with aligned, batch-invariant kernels that are bitwise consistent across training and rollout. We implement IsoExec in SkyRL with vLLM and Megatron. On a single 8×H100 node, with synchronous Qwen3.5-35B-A3B DAPO training, we reduced the average end-to-end rollout-versus-training logprob difference below , with 25% overhead compared with the current SkyRL baseline over 50 steps.

Our main contributions are:

  • Unified execution contract: One numerical execution contract across training and inference, enabling zero contract-covered mismatch with minimal overhead and low debugging cost for new RL algorithms, changes to RL environments and harnesses, and kernel improvements.
  • Parallelism-invariant kernels: Preserve numerics across tensor, expert, and sequence parallelism.
  • Chunkwise-parallel recurrent (CPR) Gated DeltaNet: Align training, prefill, and recurrent decode without serializing long-sequence forward passes.

Introduction

RL workloads require executing the same policy twice to generate rollouts (rollout engine) and train the policy (trainer). The rollout engine samples a token under policy ; the trainer later recomputes its log probability under policy using the same model parameters. Under synchronous RL, typical on-policy training assumes (no train–inference mismatch). From a systems perspective, true on-policy training is hard due to floating-point non-associativity:

RL systems often use existing inference systems such as vLLM and SGLang alongside training systems such as Megatron and FSDP. These systems are optimized for different workloads and use different kernels, batch shapes, execution modes (training, prefill, and decode), and distributed layouts. As a result, they can use different reduction orders when executing the same mathematical model, leading to different output token probability distributions.

ByteDance's VeXact study shows that mismatch alone can destabilize REINFORCE and GRPO runs, distort advantage-weighted loss contributions before a KL estimator reacts, and make importance-sampling or rejection-based fixes sensitive to calibration. Fireworks reports a GLM-5.2 run with train–inference KL around 0.013 where clipping discarded roughly 45% of tokens and reward collapsed around step 20, while a bitwise-aligned run had zero clipped tokens and remained stable.

Previous work on system determinism has focused on parts of the problem. Thinking Machines formalized batch invariance: neither the other elements in a batch nor the batch size should affect the computation for a specific element. The vLLM × TorchTitan bitwise-consistency work showed that importing matched kernels into both engines can reach parity, but still requires two aligned model copies. More recent work on Zero Train–Inference Mismatch for Linear Attention and Async RL shared one model definition between the TorchTitan trainer and vLLM generator and extended parity to Gated DeltaNet, using the recurrent form for all forward computations while retaining the chunked kernel for backward. Tree-Based Invariant Kernels (TBIK) focused on achieving bitwise consistency under different parallelism configurations.

IsoExec eliminates train–inference mismatch with an execution contract and unified model. The execution contract captures rounding-sensitive execution details (e.g., kernel implementation, accumulation dtype, reduction order) in a framework-independent form and enforces them in each runtime. The unified model uses kernels validated for bitwise consistency across training, prefill, and decode. It integrates with vLLM's engine features (e.g., its scheduler, KV cache manager, and CUDA graph capture) and Megatron's training stack.

Unified execution contract

IsoExec's core is the execution contract, which declares every bit-relevant execution choice that both runtimes must specify identically.

"ExecutionContract": {
  "cases": [ ... ],        // logprob computations: trainer_fwd, engine_decode, etc.
  "composition": [ ... ],  // (region, case) -> implementation + pinned constants
  "claims": { ... },       // topology invariance, state invalidation, tolerances
  "identities": { ... }    // semantic / numerical_policy / deployment digests
}

The contract handles each computation of a token's logprob by case (e.g., rollout engine_prefill and trainer trainer_fwd). The model's forward operators are partitioned into regions, spans of arithmetic implemented by one kernel that may fuse multiple operations. For every (region, case) pair, the composition selects the implementation and the constants it is pinned to. The constants capture any parameter that can change the bits, including accumulation and boundary dtypes and reduction-decomposition parameters such as split-K and split-KV partition counts. Every operator region is tested for bitwise exactness across cases before its implementations may be registered in the composition.

For example:

"composition": [
  {
    "region": ["gdn.core", "gdn.gating", "norms.l2"],
    "cases": ["trainer_fwd", "trainer_fwd_no_autograd", "engine_prefill", "engine_decode"],
    "impl": {"id": "native_fused_sigmoid", "version": 1, "arch": "sm90"}
  },
  {
    "region": ["moe.combine"],
    "cases": ["engine_prefill", "engine_decode"],       // the trainer side is its own entry
    "impl": {"id": "pik_leaf_tree", "version": 2, "arch": "sm90"},
    "constants": {"leaves": 8, "leaf_dtype": "fp32"},
    "discharge": {"kind": "equivalence_proof", "ref": "gates/ep_invariant_combine"} // proved equivalence
  }
]

Claims state the conditions under which the composition's guarantees hold and are enforced at runtime. For example, a topology claim lists the parallel sizes for which a reduction tree is proven bitwise invariant. When installing a kernel, the adapter compares the runtime's actual parallel size with that list and rejects an unproven size.

The identities are SHA-256 digests of the serialized contract, used to verify contract agreement between trainer and rollout engines. semantic verifies that the runtimes describe the same logical model, numerical_policy covers every execution choice that can affect numerical results (e.g., implementations and versions), and deployment covers settings proven not to affect the bits, such as memory sizing and transport configuration, which the contract does not require to match. Because every kernel admitted to the composition has been pre-validated to preserve the prescribed rounding schedule across cases, matching semantic and numerical_policy digests, together with contract-adapter enforcement, indicate that both sides are executing the same verified numerical policy across covered regions.

IsoExec's unified execution contract across training and inference runtimes.
IsoExec's unified execution contract across training and inference runtimes.

A per-runtime contract adapter installs the contract and implementations in each engine and enforces the contract at runtime. It binds every composition entry to the framework's extension points, such as selecting the specified attention kernel. It then monitors the runtime by checking installed kernels, declared claims, and cross-process identity digests.

Unified model

In SkyRL, we adopt a unified model definition consisting of batch-invariant GEMM, attention, and normalization kernels, along with deterministic MoE routing and combination. The model is also bitwise consistent across tensor-, expert-, and sequence-parallel configurations and uses a chunkwise-parallel recurrent algorithm for GDN hybrid architectures. In our experiments, we applied the abstraction to achieve zero contract-covered mismatch during training on dense (MiMo-7B), MLA MoE (GLM-4.7-Flash), hybrid (Qwen3.5-9B), and hybrid MoE (Qwen3.5-35B-A3B) models. The implementation of IsoExec is available at https://github.com/zanderjiang/SkyRL-IsoExec.

Parallelism-invariant kernels

Training and inference benefit from different distributed strategies. The trainer must fit optimizer state, activations, gradients, and, for MoE models, distributed expert weights. The rollout engine instead needs enough memory capacity for the KV cache without hurting decode latency.

For forward-pass numerics with fixed inputs and weights, the six common parallelism axes affect execution as follows:

  • Data parallelism (DP) changes batch partitioning; batch-invariant kernels preserve each sample's numerics.
  • Pipeline parallelism (PP) moves whole layers across devices without splitting their reductions when boundary dtypes are fixed.
  • Tensor parallelism (TP) splits contraction reductions across ranks.
  • Expert parallelism (EP) distributes expert computation and changes how expert outputs are combined.
  • Sequence parallelism (SP) changes row-parallel reductions from all-reduce to reduce-scatter.
  • Context parallelism (CP) splits the attention reduction across the sequence dimension.

TBIK achieved TP-invariant inference by fixing a global reduction tree across row-parallel GEMMs and the cross-GPU reduction. IsoExec takes the same fixed-reduction idea but applies it along the K dimension. Instead of building the tree over GEMM K-tiles, pik divides the K dimension into contiguous leaves. Each leaf uses deterministic Tensor Core MMA with FP32 accumulation. The contract fixes the rank-to-leaf mapping and binary arithmetic schedule, while NCCL transports partial results instead of requiring custom communication kernels.

The fixed binary reduction tree used by pik to preserve numerics across parallelism layouts.
The fixed binary reduction tree used by pik to preserve numerics across parallelism layouts.

Additionally, IsoExec applies the same principle to EP and SP. For expert parallelism, we combine expert outputs in a fixed routing order rather than rank order. For sequence parallelism, we reuse the same reduction tree as the non-SP system; each rank keeps its own output slice instead of gathering the full result. This makes trainer logits bitwise identical whether SP is enabled or disabled.

Chunkwise-parallel recurrent (CPR) GDN

Eliminating mismatch is more complex for linear-attention architectures because training and inference use different algorithms. Existing GDN systems use a chunkwise-parallel form for training and prefill but a recurrent form for decode. Although mathematically identical, the algorithms have different floating-point rounding characteristics. Comparing GDN layer outputs from FLA's chunkwise-parallel kernel with vLLM's fused recurrent kernel, we observed a mean per-element absolute difference of approximately and a maximum difference of 0.25.

Tthe TorchTitan team addressed this issue by using the recurrent form for both rollout prefill and the trainer's forward pass, and the chunkwise form only for backward computation. Although this removes mismatch for GDN, it makes rollout prefill and the trainer's forward pass serial in the sequence length. They report a slowdown of roughly 2–3× on math workloads and about 5× on a terminal-agent workload, making the approach impractical for full training jobs.

To ensure zero contract-covered mismatch while achieving high throughput for prefill, training, and decode, we designed chunkwise-parallel recurrent (CPR). CPR keeps the recurrence as the main function but evaluates it in parallel across chunks. For training and prefill, as in the chunkwise-parallel form, a first pass computes the recurrent state at every chunk boundary; a parallel recurrent scan then computes the outputs within each chunk. For decode, we use the recurrent form but resynchronize the hidden state every decoded tokens, where is the chunk size. This ensures a consistent rounding schedule across prefill, training, and decode.

Per-layer cost:

StageShapeNative mixedChunkwise everywhereRecurrent everywhereCPR
Bitwise exactNoYesYesYes
Trainer forward + backward1 × 10,240 tokens5.177 ms5.177 ms (1.00×)22.863 ms (4.42×)7.386 ms (1.43×)
Rollout-engine prefill5 × 2,048 tokens0.844 ms0.844 ms (1.00×)3.639 ms (4.31×)1.412 ms (1.67×)
Rollout-engine decode256 sequences × 1 token0.0612 ms2.2374 ms (36.6×)0.0612 ms (1.00×)0.0846 ms (1.38×)

Per-layer latency on H100 (). The trainer and rollout engine use their production TP layouts and kernels. Each ratio is relative to the native mixed implementation for that stage; smaller is better.

Results

For our experiments, we compared IsoExec against SkyRL's native stack on a single 8×H100 node by training Qwen3.5-35B-A3B on DAPO-Math-17k with synchronous RL. With identical setups, IsoExec had 25% end-to-end overhead compared with the native SkyRL stack using vLLM and Megatron under the highest-throughput synchronous-RL configuration we evaluated.

Logprob diffs

Rollout-versus-training absolute logprob differences for the native SkyRL stack and IsoExec.
Rollout-versus-training absolute logprob differences for the native SkyRL stack and IsoExec.

Across 50 steps, the mean pre-update rollout-versus-training absolute logprob difference reduced from to , its standard deviation reduced from to , and the average per-step maximum reduced from 5.073 to .

Performance

Average RL step timing for the native SkyRL stack and IsoExec over 50 steps.
Average RL step timing for the native SkyRL stack and IsoExec over 50 steps.

The average step times over the same 50-step window were:

MetricNativeIsoExecOverhead
Generation591.3 s776.6 s31.3%
Policy training498.6 s591.3 s18.6%
Full RL step1224.6 s1534.0 s25.3%

Rewards

Pass@16 and raw reward for the native SkyRL stack and IsoExec over 50 steps.
Pass@16 and raw reward for the native SkyRL stack and IsoExec over 50 steps.

Over this short 50-step run, we did not observe a meaningful reward improvement from eliminating contract-covered train–inference mismatch.

Next steps

  • Blackwell support
  • Context parallelism invariance
  • Sparse attention
  • Block-FP8 MoE

Acknowledgements

This work was done by Alexander Jiang and the SkyRL team. Thanks to Charlie Ruan, Sumanth Hegde, Eric Tang, Philipp Moritz, Yichuan Wang, Mayank Mishra, and Lingxiao Ma for helpful discussions.