NovaSky's IsoExec Fixes the Hidden Math Bug Corrupting AI Training Runs

Reinforcement learning post-training has a dirty secret: the rollout engine and the trainer are supposed to be evaluating the same policy, but they frequently disagree on a token's log-probability. The math is identical, but floating-point arithmetic isn't associative, so any difference in kernels, batch shapes, or parallel layouts between vLLM and Megatron can nudge probabilities apart. The SkyRL team at NovaSky has released IsoExec, an abstraction that forces both engines to execute the same numerical recipe and produce bitwise identical logprobs.

When a rounding bug becomes a training bug

When rollout and training disagree, the policy gradient signal quietly gets corrupted. A GLM-5.2 run with train-inference KL around 0.013 had clipping discard roughly 45% of tokens, causing reward to collapse around step 20, while a bitwise-aligned run had zero clipped tokens and remained stable. That gap separates a run that works from one that silently diverges, and it makes debugging any new RL algorithm brutally hard because you can never tell whether the culprit is your algorithm, your environment, or a reduction order buried inside a fused kernel.

IsoExec has two components: an execution contract that specifies and enforces the details affecting floating-point rounding across engines, and a unified model with aligned, batch-invariant kernels that stay bitwise consistent across training and rollout. In an 8xH100 run training Qwen3.5-35B-A3B with synchronous DAPO, it drove the mean rollout-versus-training logprob difference below 1e-6 while adding about 25% to end-to-end step time.

The execution contract

The core idea is a machine-checkable contract that pins down every choice that can move bits. 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. Those 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 entry gets pre-validated for bitwise exactness before it can be admitted. The contract carries three SHA-256 identity digests: one for semantic equivalence, one for the numerical policy, and one for deployment settings that provably do not affect bits. When the trainer and the rollout engine boot up, they exchange digests and refuse to run if their numerical policies disagree. A per-runtime adapter wires the contract into vLLM's or Megatron's extension points and monitors the kernels that actually get installed.

Making kernels parallelism-invariant

The trainer and the inference engine want completely different parallelism layouts. 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. That mismatch is unavoidable, so the kernels themselves have to produce identical bits regardless of how the work is split.

IsoExec builds on the Tree-Based Invariant Kernels idea but applies it along the K dimension of GEMMs. 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 same fixed-tree trick extends to expert parallelism and sequence parallelism. For expert parallelism, expert outputs are combined in a fixed routing order rather than rank order. For sequence parallelism, the same reduction tree as the non-SP system is reused; each rank keeps its own output slice instead of gathering the full result. The trainer logits come out identical whether SP is on or off.

Solving the linear-attention headache

Gated DeltaNet and similar linear-attention layers have a nastier problem: training uses a chunkwise-parallel algorithm while decode uses a recurrent one. The algorithms are mathematically identical but have different floating-point rounding characteristics. The TorchTitan approach was to just use the recurrent form everywhere, but they report a slowdown of roughly 2-3x on math workloads and about 5x on a terminal-agent workload, making the approach impractical for full training jobs.

IsoExec introduces chunkwise-parallel recurrent (CPR), which keeps the recurrence as the primary algorithm but evaluates it in parallel across chunks. A first pass computes recurrent state at chunk boundaries, then a parallel scan fills in outputs within each chunk. For decode, the recurrent form runs but resynchronizes the hidden state every chunk-size tokens, so the rounding schedule matches prefill and training. The per-layer cost on H100 tells the story:

Stage	Native mixed	Recurrent everywhere	CPR
Trainer fwd+bwd (10240 tok)	5.177 ms	22.863 ms (4.42x)	7.386 ms (1.43x)
Rollout prefill (5x2048 tok)	0.844 ms	3.639 ms (4.31x)	1.412 ms (1.67x)
Rollout decode (256x1 tok)	0.0612 ms	0.0612 ms (1.00x)	0.0846 ms (1.38x)

What it actually costs

Over 50 synchronous DAPO steps on Qwen3.5-35B-A3B, IsoExec collapsed the numerical gap dramatically. The mean pre-update rollout-versus-training absolute logprob difference dropped from 1.6e-2 to 6.7e-7, and the per-step maximum fell from 5.073 down to a small fraction. The cost is real but bounded:

• Generation: 591.3s to 776.6s (31.3% overhead)
• Policy training: 498.6s to 591.3s (18.6% overhead)
• Full RL step: 1224.6s to 1534.0s (25.3% overhead)

Crucially, vLLM's scheduler, paged KV cache, and CUDA graph capture still work. One caveat worth flagging from the authors: over this short 50-step run, they did not observe a meaningful reward improvement from eliminating contract-covered train-inference mismatch. The value shows up in longer runs and harder algorithms where mismatch destabilizes training, not in a quick reward bump on a well-behaved setup.

Where IsoExec fits

IsoExec sits alongside a growing body of work on determinism in LLM systems, including Thinking Machines' batch-invariance work and the earlier vLLM x TorchTitan parity effort. What IsoExec adds is a formal, cross-framework contract that covers dense, MLA MoE, hybrid, and hybrid MoE architectures under multiple parallelism axes at once. The implementation lives at github.com/zanderjiang/SkyRL-IsoExec, and it works with the existing SkyRL, vLLM, and Megatron stacks rather than replacing them.

For teams running production RL training on frontier models, the practical implication is that you can now attribute divergence problems to your algorithm or environment instead of chasing phantoms in kernel reduction orders. That matters for anyone iterating on GRPO, DAPO, or REINFORCE variants where a 1% shift in token probabilities can silently distort your advantage estimates.