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

发表时间: 2026-08 · Blog post by Alexander Jiang and the SkyRL Team (blog.vllm.ai)

原文: https://blog.vllm.ai/blog/2026-08-21-isoexec

作者/机构:Alexander Jiang and the SkyRL Team

速读

一句话结论 IsoExec 提出了一个跨框架的统一执行抽象,通过执行契约和统一模型消除了强化学习中训练和推理引擎之间的浮点数计算差异,在 Qwen3.5-35B-A3B 的同步强化学习中将对数概率误差降低了四个数量级。

要解决什么问题 在同策略强化学习中,理论上要求生成经验和更新模型评估的是完全相同的策略。但在实际的系统工程中,为了追求极致性能,通常会使用两套不同的引擎,例如用 vLLM 做推理生成,用 Megatron 做训练。这两套引擎针对不同的负载进行了优化,使用了不同的算子实现、批次形状、执行模式以及分布式并行策略。由于计算机中的浮点数加法不具备结合律,计算顺序的改变会导致最终输出的 Token 概率分布产生微小差异。这种训练与推理之间的数值不匹配不仅会增加新算法和算子的调试成本,更会直接破坏训练的稳定性,例如导致 REINFORCE 和 GRPO 运行失稳、优势加权损失失真、重要性采样或拒绝采样机制对校准变得极其敏感,甚至在某些极端情况下会导致大量 Token 被裁剪、奖励值在训练早期直接崩溃。

怎么做的 IsoExec 的核心思路是通过一套严格的跨框架抽象,强制要求训练和推理引擎在所有影响浮点数舍入的执行细节上保持按位一致。该方法主要由统一执行契约和统一模型两个关键部件构成。统一执行契约是一个独立于框架的声明文件,它将模型的前向传播划分为多个算子区域,并为每种执行情况指定确切的底层实现和常量参数(如累加数据类型、归约分解参数)。契约通过声明拓扑约束来验证并行规模,并计算哈希摘要来比对两端的语义和数值策略,最后由各引擎内的适配器在运行时拦截并强制执行这些设定。统一模型则提供了一套在不同并行策略下保持数值不变的算子库。针对张量并行,它没有使用常规的通信归约,而是将矩阵乘法的 K 维度划分为 $N$ 个连续的叶子节点,强制使用确定的 FP32 累加和固定的二叉树归约调度;对于专家并行和序列并行,也同样采用了固定路由顺序和保留输出切片的固定归约树。此外,针对 Gated DeltaNet 这类线性注意力架构在分块并行(用于训练)和循环(用于解码)模式下存在的严重舍入差异,作者设计了分块并行循环算法。该算法在训练和预填充阶段,先计算每个分块边界的循环状态,然后在分块内进行并行的循环扫描;在解码阶段,则保持循环形式,但强制每隔 $C$ 个 Token($C$ 为分块大小)重新同步一次隐藏状态,从而在不将长序列前向传播完全串行化的前提下,对齐了所有阶段的舍入调度。

效果如何 实验在一台 8×H100 节点上进行,使用 DAPO-Math-17k 数据集对 Qwen3.5-35B-A3B 混合专家模型进行同步强化学习训练。对比的基线方法是原生的 SkyRL 协议栈,它代表了当前主流的 vLLM 推理结合 Megatron 训练的异构路线。量化结果显示,在 50 步的训练窗口内,IsoExec 将生成与训练之间的平均绝对对数概率差异从 $1.3 \times 10^{-3}$ 极大地降低到了 $1.1 \times 10^{-7}$,每步的最大差异也从 $0.021$ 降至 $1.3 \times 10^{-6}$。这种严格的对齐是有代价的:在最高吞吐量的配置下,IsoExec 带来了 25.3% 的端到端耗时开销,其中生成阶段变慢了 31.3%,策略训练阶段变慢了 18.6%。具体到分块并行循环算子层面,其训练前向加反向、推理预填充和推理解码的单层延迟分别是原生混合实现的 1.43 倍、1.67 倍和 1.38 倍。作者也坦诚了该方法的局限性:在这个仅包含 50 步的短期实验中,虽然消除了契约覆盖范围内的训练与推理不匹配,但并没有观察到模型最终奖励值有实质性的提升。

主要贡献

在理论上,同策略(on-policy)强化学习(RL)假设模型在 rollout(数据采样)和 training(训练)阶段评估的是完全相同的策略。但在实际的 RL 训练系统中,往往使用两个独立的引擎分别进行 rollout 和训练,这两个引擎在模型定义、内核(kernels)、批次形状(batch shapes)以及并行布局上存在差异。由于浮点运算不具备结合律,这些差异会导致即使执行相同策略,输出的 token 概率也会发生改变。这种训练与推理间的不匹配(train-inference mismatch)使得调试新的 RL 算法、环境调整以及底层基础设施的优化变得极为困难。

为了解决这一问题,本文提出了 IsoExec,这是一个跨框架的统一执行抽象,旨在消除 RL 工作负载中训练引擎和推理引擎之间的不匹配。IsoExec 的实现基于 SkyRL,并结合了 vLLM 和 Megatron。在单节点 8×H100 上进行 Qwen3.5-35B-A3B DAPO 的同步训练时,IsoExec 将端到端的 rollout 与 training 之间的平均对数概率(logprob)差异降低到了 $0.0$ 以下,且在 50 个 step 内相比 SkyRL 基线仅带来 25% 的性能开销。

本文的主要贡献包括:

  • 统一执行契约(Unified execution contract):在训练和推理之间建立单一的数值执行契约,以极低的开销实现契约覆盖范围内的零不匹配,并显著降低了新算法、新环境及内核优化的调试成本。
  • 并行不变内核(Parallelism-invariant kernels):在张量并行(TP)、专家并行(EP)和序列并行(SP)下保持数值计算的绝对一致性。
  • 分块并行循环(CPR)Gated DeltaNet:对齐训练、预填充(prefill)和循环解码(recurrent decode)阶段的计算,且无需将长序列前向传播串行化。

背景知识与设计原则

强化学习中的同策略困境:RL 工作负载需要对同一策略执行两次:一次用于生成 rollouts(rollout 引擎),另一次用于训练策略(trainer)。Rollout 引擎在策略 $\pi_{\theta_{old}}$ 下采样 token;trainer 随后使用相同的模型参数重新计算其在策略 $\pi_{\theta}$ 下的对数概率。在同步 RL 中,典型的同策略训练假设 $\pi_{\theta_{old}} = \pi_{\theta}$(即不存在训练-推理不匹配)。然而从系统角度来看,由于浮点数非结合律的存在,实现真正的同策略训练非常困难。RL 系统通常将现有的推理系统(如 vLLM 和 SGLang)与训练系统(如 Megatron 和 FSDP)结合使用。这些系统针对不同的工作负载进行了优化,使用了不同的内核、批次形状、执行模式(训练、预填充和解码)以及分布式布局。因此,在执行相同的数学模型时,它们可能会使用不同的归约(reduction)顺序,从而导致输出的 token 概率分布不同。

不匹配对模型训练的破坏性影响:字节跳动的研究【1】表明,仅仅是不匹配本身就足以破坏 REINFORCE 和 GRPO 的运行稳定性,在 KL 估计器做出反应之前扭曲优势加权的损失贡献,并使得基于重要性采样或拒绝采样的修复方法对校准变得极其敏感。Fireworks 报告【2】指出,在一次 GLM-5.2 的运行中,当训练-推理 KL 散度约为 0.013 时,梯度裁剪丢弃了约 45% 的 token,并且奖励在第 20 步左右崩溃;而实现了比特级对齐的运行则没有丢弃任何 token,并保持了稳定。

解决不匹配问题的先前探索:先前关于系统确定性的工作主要集中在解决部分问题上。Thinking Machines【3】形式化了批次不变性(batch invariance):批次中的其他元素或批次大小都不应影响特定元素的计算。vLLM × TorchTitan 的比特级一致性工作【4】表明,将匹配的内核导入两个引擎可以实现对齐,但这仍然需要两份对齐的模型副本。最近关于线性注意力和异步 RL 的工作【5】在 TorchTitan 训练器和 vLLM 生成器之间共享了一个模型定义,并将对齐扩展到了 Gated DeltaNet【7】,其方法是对所有前向计算使用循环形式,同时保留用于反向计算的分块内核。基于树的不变内核(TBIK)【6】则专注于在不同的并行配置下实现比特级一致性。

IsoExec的解决思路:IsoExec 通过执行契约和统一模型消除了训练-推理不匹配。执行契约以独立于框架的形式捕获对舍入敏感的执行细节(例如内核实现、累加数据类型、归约顺序),并在每个运行时中强制执行。统一模型使用了经过验证的、在训练、预填充和解码之间具有比特级一致性的内核。它集成了 vLLM 的引擎特性(例如其调度器、KV 缓存管理器和 CUDA 图捕获)以及 Megatron 的训练栈。

方法细节

统一执行契约 (Unified execution contract)

统一执行契约的核心设计:IsoExec 的核心是执行契约(Execution Contract),它声明了两个运行时必须完全一致的、与比特级别相关的每一个执行选择。

"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
}

契约的组成部分与验证逻辑:契约按“情况”(case)(例如 rollout 的 engine_prefill 和 trainer 的 trainer_fwd)来处理 token 对数概率的每次计算。模型的前向算子被划分为“区域”(regions),即由一个可能融合多个操作的内核实现的算术跨度。对于每个(region, case)对,“组合”(composition)负责选择具体的实现及其绑定的常量。这些常量捕获了任何可能改变比特的参数,包括累加和边界数据类型,以及归约分解参数(如 split-K 和 split-KV 分区数)。每个算子区域在被注册到组合之前,都必须在不同情况之间进行比特级别的精确性测试。例如:

"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)的约束作用:声明(Claims)说明了组合保证成立并能在运行时强制执行的条件。例如,拓扑声明列出了已证明其归约树具有比特不变性的并行规模。在安装内核时,适配器会将运行时的实际并行规模与该列表进行比较,并拒绝执行未经证明的规模。

标识符(Identities)的校验机制:标识符是序列化契约的 SHA-256 摘要,用于验证 trainer 和 rollout 引擎之间的契约一致性。semantic 用于验证运行时描述的是相同的逻辑模型;numerical_policy 涵盖所有可能影响数值结果的执行选择(如实现和版本);deployment 则涵盖被证明不影响比特的设置(如内存大小和传输配置),契约不要求这些设置必须匹配。由于加入组合的每个内核都经过预先验证,能够在不同情况下保持规定的舍入计划,因此匹配的 semantic 和 numerical_policy 摘要,加上契约适配器的强制执行,表明双方在覆盖的区域内执行的是相同的、已验证的数值策略。

IsoExec在训练和推理运行时之间的统一执行契约。
IsoExec在训练和推理运行时之间的统一执行契约。

契约适配器(Contract adapter)的执行:每个运行时都有一个专属的契约适配器,负责在每个引擎中安装契约和实现,并在运行时强制执行契约。它将每个组合条目绑定到框架的扩展点(例如选择指定的注意力内核)。随后,它通过检查已安装的内核、声明的条件以及跨进程的标识符摘要来持续监控运行时。

统一模型 (Unified model)

统一模型定义:在 SkyRL 中,采用了一个统一的模型定义,该定义包含批次不变的 GEMM、注意力和归一化内核,以及确定性的 MoE 路由和组合。该模型在张量、专家和序列并行配置中也保持比特级一致性,并对 GDN 混合架构使用分块并行循环(CPR)算法。在实验中,将该抽象应用于密集模型(MiMo-7B)、MLA MoE(GLM-4.7-Flash)、混合模型(Qwen3.5-9B)和混合 MoE(Qwen3.5-35B-A3B)模型的训练中,实现了契约覆盖范围内的零不匹配。

并行不变内核 (Parallelism-invariant kernels)

不同分布式策略的挑战:训练和推理受益于不同的分布式策略。训练器必须容纳优化器状态、激活值、梯度,以及(对于 MoE 模型)分布式专家权重。而 rollout 引擎则需要足够的内存容量来存储 KV 缓存,同时不能损害解码延迟。

并行维度对执行的影响:对于具有固定输入和权重的前向传播数值,六个常见的并行轴对执行有如下影响:

  • 数据并行(DP):改变批次划分;批次不变内核保留每个样本的数值。
  • 流水线并行(PP):跨设备移动整个层,在边界数据类型固定时不拆分其归约。
  • 张量并行(TP):跨 rank 拆分收缩(contraction)归约。
  • 专家并行(EP):分布专家计算并改变专家输出的组合方式。
  • 序列并行(SP):将行并行归约从 all-reduce 更改为 reduce-scatter。
  • 上下文并行(CP):跨序列维度拆分注意力归约。

基于固定归约树的并行不变性设计:TBIK【6】通过在行并行 GEMM 和跨 GPU 归约之间固定全局归约树,实现了 TP 不变的推理。IsoExec 采用了相同的固定归约思想,但将其应用于 K 维度。pik 没有在 GEMM K-tiles 上构建树,而是将 K 维度划分为 $N$ 个连续的叶子节点。每个叶子节点使用确定性的 Tensor Core MMA 和 FP32 累加。契约固定了 rank 到叶子的映射和二元算术调度,而 NCCL 传输部分结果,无需自定义通信内核。

pik用于在不同并行布局中保持数值不变的固定二叉归约树。
pik用于在不同并行布局中保持数值不变的固定二叉归约树。

EP 与 SP 的并行不变性扩展:IsoExec 将相同的原则应用于 EP 和 SP。对于专家并行,以固定的路由顺序而不是 rank 顺序组合专家输出。对于序列并行,重用与非 SP 系统相同的归约树;每个 rank 保留自己的输出切片,而不是收集完整结果。这使得无论是否启用 SP,训练器的 logits 在比特级别上完全相同。

分块并行循环(CPR)GDN

线性注意力架构的不匹配问题:对于线性注意力架构,消除不匹配更加复杂,因为训练和推理使用不同的算法。现有的 GDN 系统在训练和预填充时使用分块并行(chunkwise-parallel)形式,而在解码时使用循环(recurrent)形式。尽管在数学上相同,但这些算法具有不同的浮点舍入特性。比较 FLA【8】的分块并行内核与 vLLM 的融合循环内核的 GDN 层输出,观察到平均每个元素的绝对差异约为 $4 \times 10^{-4}$,最大差异为 $0.25$。

先前方案的局限性:TorchTitan 团队【5】通过在 rollout 预填充和训练器前向传播中都使用循环形式,仅在反向传播中使用分块形式来解决此问题。虽然这消除了 GDN 的不匹配,但它使得 rollout 预填充和训练器前向传播在序列长度上变为串行。他们报告在数学工作负载上大约有 2-3 倍的减速,在终端代理工作负载上大约有 5 倍的减速,使得该方法对于完整的训练任务不切实际。

CPR算法的设计与工作原理:为了在确保契约覆盖范围内零不匹配的同时,实现预填充、训练和解码的高吞吐量,设计了分块并行循环(CPR,chunkwise-parallel recurrent)。CPR 将循环作为主要函数,但跨分块并行评估它。对于训练和预填充,与分块并行形式一样,第一遍在每个分块边界计算循环状态;然后并行的循环扫描计算每个分块内的输出。对于解码,使用循环形式,但每解码 $C$ 个 token 重新同步隐藏状态,其中 $C$ 是分块大小。这确保了在预填充、训练和解码之间具有一致的舍入计划。

每一层的性能开销如下表所示(在 H100 上的每层延迟。训练器和 rollout 引擎使用其生产 TP 布局和内核。每个比率是相对于该阶段的原生混合实现的;越小越好):

阶段 形状 原生混合 (Native mixed) 全部分块 (Chunkwise everywhere) 全部循环 (Recurrent everywhere) CPR
比特级精确 — 否 是 是 是
训练器前向 + 反向 1 × 10,240 tokens 5.177 ms 5.177 ms (1.00×) 22.863 ms (4.42×) 7.386 ms (1.43×)
Rollout引擎预填充 5 × 2,048 tokens 0.844 ms 0.844 ms (1.00×) 3.639 ms (4.31×) 1.412 ms (1.67×)
Rollout引擎解码 256 sequences × 1 token 0.0612 ms 2.2374 ms (36.6×) 0.0612 ms (1.00×) 0.0846 ms (1.38×)

实验环境

  • 硬件配置:单个 8×H100 节点。
  • 模型架构:Qwen3.5-35B-A3B(混合 MoE 模型)。
  • 数据集:DAPO-Math-17k。
  • 软件配置:使用 vLLM 和 Megatron 的 SkyRL 栈进行同步 RL 训练配置。

实验结果

在实验中,将 IsoExec 与 SkyRL 的原生栈进行了对比。在相同的设置下,使用最高吞吐量的同步 RL 配置,IsoExec 相比原生 SkyRL 栈(使用 vLLM 和 Megatron)产生了 25% 的端到端开销。

对数概率差异 (Logprob diffs):
在 50 个 step 中,更新前的 rollout 与 training 之间的平均绝对对数概率差异从 $2 \times 10^{-4}$ 降低到了 $0.0$,其标准差从 $3 \times 10^{-4}$ 降低到了 $0.0$,平均每步最大差异从 $0.05$ 降低到了 $0.0$。
原生SkyRL栈和IsoExec的Rollout与训练绝对对数概率差异。

性能开销 (Performance):
在相同的 50 步窗口内,平均 step 时间如下:

  • 原生栈:生成 591.3 秒,策略训练 498.6 秒,完整 RL step 1224.6 秒。
  • IsoExec:生成 776.6 秒(开销 31.3%),策略训练 591.3 秒(开销 18.6%),完整 RL step 1534.0 秒(开销 25.3%)。
    在50个step中,原生SkyRL栈和IsoExec的平均RL step耗时。

奖励表现 (Rewards):
在这短短的 50 步运行中,并未观察到通过消除契约覆盖的训练-推理不匹配带来了有意义的奖励提升。
在50个step中,原生SkyRL栈和IsoExec的Pass@16和原始奖励。

结论

本文提出了 IsoExec 抽象,成功消除了强化学习训练过程中由于底层引擎和分布式策略差异导致的浮点数不匹配问题。通过统一执行契约、并行不变内核以及 CPR GDN 设计,在引入可接受的性能开销(25%)的前提下,实现了训练与推理间对数概率差异的清零。未来的工作计划包括:

  • 支持 Blackwell 架构。
  • 实现上下文并行(Context parallelism)不变性。
  • 支持稀疏注意力(Sparse attention)。
  • 支持 Block-FP8 MoE。

参考文献汇总