跳到正文
vLLM Blog· Alexander Jiang and the SkyRL Team·· 2026-08-21精选AI 评分62

SkyRL 提出 IsoExec:统一执行消除训练与推理引擎不匹配

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

AI 导读

SkyRL 团队提出 IsoExec,用跨框架的统一执行抽象消除 RL 训练与推理引擎之间的数值不匹配,包含一份执行契约和一套在训练、prefill、decode 间逐位一致的统一模型,已在 vLLM 与 Megatron 上实现。

推荐理由

SkyRL 团队公开了消除训练与推理引擎数值不一致的执行契约设计,并给出 8×H100 上的开销与 logprob 差异数据。

正文 · AI 翻译

TL;DR

理论上,on-policy RL 假设 rollout 和训练评估的是同一个策略。但在实践中,RL 训练系统通常使用两个独立的引擎分别进行 rollout 和训练,它们的模型定义、kernel、batch 形状和并行布局各不相同。由于浮点运算不满足结合律,即使执行同一策略,这些差异也会改变 token 概率。这会让新的 RL 算法、harness 和环境变更,以及 RL 基础设施和硬件 kernel 的改进都难以调试。

为了解决这个问题,我们引入了 IsoExec,一种跨框架的统一执行抽象,用于消除 RL 工作负载中训练引擎与推理引擎之间的不一致。IsoExec 包含两个组件:一个执行契约,用于指定并强制跨引擎影响浮点舍入的执行细节;以及一个统一模型,具有对齐的、batch 不变的 kernel,在训练和 rollout 之间实现逐位一致。我们在 SkyRL 中结合 vLLM 和 Megatron 实现了 IsoExec。在单台 8×H100 节点上,使用同步 Qwen3.5-35B-A3B DAPO 训练,我们将端到端 rollout 与训练之间的平均 logprob 差异降至 以下,在 50 步内相比当前 SkyRL 基线仅有 25% 的开销。

我们的主要贡献是:

  • 统一执行契约:训练和推理共用一套数值执行契约,以最小开销实现契约覆盖范围内的零不一致,并降低新 RL 算法、RL 环境和 harness 变更以及 kernel 改进的调试成本。
  • 并行不变 kernel:在张量并行、专家并行和序列并行下保持数值一致。
  • 分块并行循环(CPR)Gated DeltaNet:对齐训练、prefill 和循环解码,而无需串行化长序列前向传播。

引言

RL 工作负载需要执行同一策略两次,以生成 rollout(rollout 引擎)并训练策略(训练器)。rollout 引擎在策略 下采样一个 token;训练器随后使用相同的模型参数在策略 下重新计算其对数概率。在同步 RL 中,典型的 on-policy 训练假设 (无训练–推理不一致)。从系统角度看,由于浮点运算不满足结合律,真正的 on-policy 训练很难实现:

RL 系统通常将 vLLM 和 SGLang 等现有推理系统与 Megatron 和 FSDP 等训练系统一起使用。这些系统针对不同的工作负载进行了优化,并使用不同的 kernel、batch 形状、执行模式(训练、prefill 和解码)以及分布式布局。因此,它们在执行相同的数学模型时可能使用不同的归约顺序,从而导致输出 token 概率分布不同。

字节跳动的 VeXact 研究表明,仅这种不一致本身就可能使 REINFORCE 和 GRPO 训练不稳定,在 KL 估计器做出反应之前扭曲优势加权损失贡献,并使基于重要性采样或拒绝的修复方法对校准敏感。Fireworks 报告了一次 GLM-5.2 运行,其训练–推理 KL 约为 0.013,裁剪丢弃了大约 45% 的 token,奖励在第 20 步左右崩溃;而逐位对齐的运行则没有裁剪任何 token,并保持稳定。

此前关于系统确定性的工作只关注了问题的部分环节。Thinking Machines 将批次不变性形式化:批次中的其他元素以及批次大小都不应影响特定元素的计算。vLLM × TorchTitan 的位级一致性工作表明,将匹配的内核导入两个引擎可以达到一致,但仍需要两份对齐的模型副本。近期关于 Zero Train–Inference Mismatch for Linear Attention and Async RL 的工作在 TorchTitan 训练器和 vLLM 生成器之间共享同一份模型定义,并将一致性扩展到 Gated DeltaNet,在所有前向计算中使用循环形式,同时在反向中保留分块内核。Tree-Based Invariant Kernels (TBIK) 则专注于在不同并行配置下实现位级一致性。

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

统一执行契约

IsoExec 的核心是执行契约,它声明了两个运行时必须完全相同指定的每一个与位相关的执行选择。

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

该契约按用例(例如 rollout engine_prefill 和 trainer trainer_fwd)处理 token logprob 的每一次计算。模型的前向算子被划分为区域,即由单个内核实现的算术跨度,该内核可能融合多个操作。对于每一个(区域,用例)对,组合选择实现及其所固定的常量。这些常量捕获任何可能改变位的参数,包括累加和边界数据类型,以及归约分解参数,例如 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
  }
]

声明陈述组合的保证成立的条件,并在运行时强制执行。例如,拓扑声明列出了归约树被证明具有位级不变性的并行规模。在安装内核时,适配器会将运行时的实际并行规模与该列表进行比较,并拒绝未经证明的规模。

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

IsoExec's unified execution contract across training and inference runtimes.
IsoExec 在训练与推理运行时之间统一的执行合约。

每个运行时都有一个 合约适配器,它在各引擎中安装合约与实现,并在运行时强制执行合约。它将每个组合条目绑定到框架的扩展点,例如选择指定的注意力内核。随后它通过检查已安装的内核、声明的断言以及跨进程身份摘要来监控运行时。

统一模型

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

并行不变内核

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

对于输入和权重固定的前向传播数值计算,六种常见的并行轴对执行的影响如下:

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

TBIK 通过在行并行 GEMM 和跨 GPU 归约中固定全局归约树,实现了 TP 不变推理。IsoExec 采用相同的固定归约思路,但将其应用于 K 维度。它不是基于 GEMM K-tile 构建树,而是 pik 将 K 维度划分为连续的叶子节点。每个叶子节点使用确定性 Tensor Core MMA 和 FP32 累加。该契约固定了 rank 到叶子的映射和二进制算术调度,而 NCCL 负责传输部分结果,无需自定义通信内核。

The fixed binary reduction tree used by pik to preserve numerics across parallelism layouts.
pik 用于在不同并行布局间保持数值一致性的固定二叉归约树。

此外,IsoExec 将同样的原则应用于 EP 和 SP。对于专家并行,我们按固定的路由顺序而非 rank 顺序组合专家输出。对于序列并行,我们复用与非 SP 系统相同的归约树;每个 rank 保留自己的输出切片,而不是收集完整结果。这使得无论 SP 是否启用,训练器 logits 都逐位相同。

分块并行循环(CPR)GDN

消除不匹配对于线性注意力架构更为复杂,因为训练和推理使用不同的算法。现有的 GDN 系统在训练和预填充时使用分块并行形式,但在解码时使用循环形式。尽管数学上等价,这些算法具有不同的浮点舍入特性。将 FLA 的 分块并行内核与 vLLM 的融合循环内核的 GDN 层输出进行比较,我们观察到平均每元素绝对差约为 且最大差为 0.25。

TorchTitan 团队 通过将循环形式用于 rollout 预填充和训练器的前向传播,而将分块形式仅用于反向计算,解决了这个问题。尽管这消除了 GDN 的不匹配,但它使 rollout 预填充和训练器的前向传播在序列长度上变为串行。他们报告在数学工作负载上大约慢 2–3 倍,在终端代理工作负载上大约慢 5 倍,使得该方法对于完整训练任务不切实际。

为了确保零契约覆盖的不匹配,同时为预填充、训练和解码实现高吞吐量,我们设计了 分块并行循环(CPR)。CPR 将循环作为主函数,但在块间并行评估它。对于训练和预填充,如同分块并行形式,第一遍计算每个块边界处的循环状态;然后并行循环扫描计算每个块内的输出。对于解码,我们使用循环形式,但每 个解码 token 重新同步隐藏状态,其中 是块大小。这确保了预填充、训练和解码之间一致的舍入调度。

每层成本:

阶段形状原生混合处处分块处处循环CPR
逐位精确—否是是是
训练器前向 + 反向1 × 10,240 tokens5.177 ms5.177 ms (1.00×)22.863 ms (4.42×)7.386 ms (1.43×)
Rollout 引擎预填充5 × 2,048 tokens0.844 ms0.844 ms (1.00×)3.639 ms (4.31×)1.412 ms (1.67×)
Rollout 引擎解码256 序列 × 1 token0.0612 ms2.2374 ms (36.6×)0.0612 ms (1.00×)0.0846 ms (1.38×)

H100 上的每层延迟()。训练器和 rollout 引擎使用其生产 TP 布局和内核。每个比率相对于该阶段的原生混合实现;越小越好。

结果

在我们的实验中,我们在单个 8×H100 节点上,通过在 DAPO-Math-17k 上以同步 RL 训练 Qwen3.5-35B-A3B,将 IsoExec 与 SkyRL 的原生技术栈进行了对比。在相同设置下,在我们评估的最高吞吐同步 RL 配置中,与使用 vLLM 和 Megatron 的原生 SkyRL 技术栈相比,IsoExec 的端到端开销为 25%。

对数概率差异

Rollout-versus-training absolute logprob differences for the native SkyRL stack and IsoExec.
原生 SkyRL 技术栈与 IsoExec 的 rollout 与训练之间的绝对对数概率差异。

在 50 步中,更新前 rollout 与训练之间的平均绝对对数概率差异从 降至 ,其标准差从 降至 ,每步平均最大值从 降至 。

性能

Average RL step timing for the native SkyRL stack and IsoExec over 50 steps.
原生 SkyRL 技术栈与 IsoExec 在 50 步内的平均 RL 步耗时。

在同一 50 步窗口内的平均步耗时为:

指标原生IsoExec开销
生成591.3 s776.6 s31.3%
策略训练498.6 s591.3 s18.6%
完整 RL 步1224.6 s1534.0 s25.3%

奖励

Pass@16 and raw reward for the native SkyRL stack and IsoExec over 50 steps.
原生 SkyRL 技术栈与 IsoExec 在 50 步内的 Pass@16 和原始奖励。

在这段短短的 50 步运行中,我们并未观察到消除契约覆盖的训练–推理不匹配所带来的有意义的奖励提升。

后续工作

  • Blackwell 支持
  • 上下文并行不变性
  • 稀疏注意力
  • Block-FP8 MoE

致谢

这项工作由 Alexander Jiang 和 SkyRL 团队完成。感谢 Charlie Ruan、Sumanth Hegde、Eric Tang、Philipp Moritz、Yichuan Wang、Mayank Mishra 和 Lingxiao Ma 的有益讨论。

来源:vLLM Blog · vllm.ai