IsoExec: Eliminating Trainer-Inference Mismatch in SkyRL

vLLM 推出了 IsoExec，這是一種跨框架的統一執行抽象，旨在消除強化學習 (RL) 工作負載中訓練與推理引擎之間的數值不匹配問題。透過強制執行嚴格的執行合約 (execution contract) 並利用與並行性無關的內核 (parallelism-invariant kernels)，IsoExec 確保了 rollout engine 與 trainer 能夠以位元一致性 (bitwise consistency) 評估相同的策略，將平均端到端 rollout-versus-training logprob 差異降低至 $10^{-7}$ 以下，且僅產生 25% 的效能開銷。

The Problem: Trainer-Inference Mismatch

在 on-policy RL 中，系統必須執行兩次相同的策略：一次在 rollout engine 中進行 token 採樣，另一次在 trainer 中重新計算對數機率 (log probabilities)。在實務中，這兩個階段通常使用不同的引擎（例如，用於 rollout 的 vLLM 與用於訓練的 Megatron），它們採用不同的內核、batch shapes 與並行佈局 (parallelism layouts)。

由於浮點數運算不具備結合律 (non-associative)，這些系統性差異會導致不同的歸約順序 (reduction orders)，即使模型參數完全相同，也會改變 token 的機率。這種不匹配可能會使 REINFORCE 與 GRPO 等 RL 演算法變得不穩定，並導致獎勵崩潰 (reward collapse)。例如，Fireworks 報告的一項 GLM-5.2 運行結果顯示，其訓練與推理的 KL 散度約為 0.013，導致 45% 的 token 被透過 clipping 捨棄，並在第 20 步左右發生了隨後的獎勵崩潰。

The IsoExec Unified Execution Contract

IsoExec 使用 執行合約 (execution contract) 來解決這些差異，該合約聲明了兩個運行時 (runtimes) 必須遵循的所有與位元相關的執行選擇。

Contract Structure and Enforcement

Cases and Regions: 合約將前向運算符 (forward operators) 劃分為「區域 (regions)」（由單個內核實現的算術跨度）。接著定義「案例 (cases)」（例如，engine_prefill 與 trainer_fwd）來指定如何處理每個區域。
Composition: 對於每個 region-case 對，合約會選擇一個特定的實現，並將其固定為常數，例如累加資料型別 (accumulation dtypes) 與歸約分解參數 (reduction-decomposition parameters，例如 split-K partition counts)。
Claims: 這些是運行時強制執行的條件，例如 topology claims，列出了歸約樹 (reduction tree) 被證明在位元層級上與並行性無關的特定並行尺寸。
Identities: 使用 SHA-256 摘要 (semantic, numerical_policy, 與 deployment) 來驗證 trainer 與 rollout engine 是否正在執行相同的已驗證數值策略。

每個運行時的 合約適配器 (contract adapter) 將這些規範綁定到框架的擴展點，並監控運行時以確保合規。

Unified Model and Parallelism-Invariant Kernels

IsoExec 實現了一種統一的模型定義，使其在各種分佈式策略（包括 tensor, expert, 與 sequence parallelism）中保持位元一致性。

Achieving Parallelism Invariance

為了在不同的 GPU 佈局中維持數值一致性，IsoExec 採用了固定的二元歸約樹 (binary reduction tree) 方法：

Tensor Parallelism (TP): 使用 pik 實現，IsoExec 將 K 維度劃分為連續的葉子節點 (leaves)。每個葉子節點使用確定性的 Tensor Core MMA 與 FP32 累加，且合約固定了 rank-to-leaf 映射與算術排程 (arithmetic schedule)。
Expert Parallelism (EP): Expert 的輸出是以固定的路由順序 (routing order) 而非 rank 順序進行組合。
Sequence Parallelism (SP): IsoExec 重用與非 SP 系統相同的歸約樹，每個 rank 保留其自身的輸出切片 (output slice)，確保無論是否啟用 SP，trainer logits 都能保持一致。

Chunkwise-Parallel Recurrent (CPR) Gated DeltaNet

線性注意力架構（如 Gated DeltaNet (GDN)）通常在訓練/prefill 階段使用 chunkwise-parallel 形式，而在 decode 階段使用 recurrent 形式，這會造成數值不匹配。先前嘗試透過在所有地方都使用 recurrent 形式來修復此問題的方法，導致了顯著的減速（在某些工作負載下高達 5 倍）。

IsoExec 引入了 Chunkwise-Parallel Recurrent (CPR)，它將遞歸 (recurrence) 作為主要功能，但會在各個 chunkwise 之間以並行方式進行評估。對於 decode，它每 $C$ 個 token 會重新同步隱藏狀態 (hidden state)，（其中 $C$ 是 chunk size）。這在保持高吞吐量的同時實現了位元精確度：

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 (2.66x)	0.0612 ms (1.00x)	0.0846 ms (1.38x)

Experimental Results

IsoExec 在單個 8xH100 節點上，使用同步 RL 在 DAPO-Math-17k 上訓練 Qwen3.5-35B-A3B 進行了測試。

Numerical Accuracy

在 50 個步驟中，更新前的平均 rollout-versus-training 絕對 logprob 差異從 $10^{-3}$ (native SkyRL) 降低至 $10^{-7}$ 以下 (IsoExec)。每個步驟的最大平均差異從 5.073 降至 $10^{-4}$。

Performance Overhead

消除不匹配引入了與原生 SkyRL 堆疊相比的效能懲罰：

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%

雖然達成了數值一致性，但團隊在短暫的 50 步測試窗口內並未觀察到顯著的獎勵提升。