vLLM RL-Kernel v0.1.0：在 ROCm 上实现训练与推理的逐位一致性

vLLM 和 RL-Kernel 已实现一个系统，确保在 AMD Instinct MI300X GPU 上训练和推理的 logprobs 之间实现逐位数值一致性。此集成消除了不同执行引擎处理训练和推理时通常出现的数值策略偏移，确保强化学习（RL）后训练中的重要性比率和 KL 诊断基于完全相同的浮点计算。

训练与推理数值不一致的问题

在强化学习（RL）工作流中，训练和推理通常使用不同的执行引擎和算子。即使使用相同的模型权重和输入，内核实现、并行策略和归约顺序的差异也可能导致 logprobs 分歧。

这种不一致引入了数值策略偏移，进而传播到重要性比率、推理-训练 KL 诊断和裁剪机制中。尽管之前的努力主要集中在 NVIDIA 平台上，但本工作将这些保证扩展到了 ROCm 生态系统。

技术架构：vime 和 RL-Kernel

实现一致性需要对数据流的时间和内核的实际数值执行进行对齐。工作负载被分为两个组件：

  • vime： 对齐训练时间线。它管理哪个 token 批次属于哪个步骤，跟踪哪个权重版本生成了特定 token，并将新权重同步到 vLLM。
  • RL-Kernel： 对齐数值执行。它确保相同的值参与计算，并且在两个引擎中以相同的方式进行分区、合并和舍入。

数值执行契约

为实现逐位一致性，RL-Kernel 在前向路径上强制执行严格的数值执行契约。该契约定义了六个关键可观察行为：

  1. 输入依赖性： 输出所依赖的特定元素和归约域。
  2. 分区： 归约域如何划分为部分结果。
  3. 合并顺序： 部分结果合并的顺序。
  4. 精度： 输入、累加器、中间状态和输出的精度。
  5. 舍入： 舍入或向下转换发生的精确位置。
  6. 原语： 使用的具体数值原语（例如 exp、log、rsqrt、SiLU 和 FMA）。

在 ROCm 上的实现

将这些保证适配到 AMD Instinct MI300X 需要修改多个原生 ROCm 执行路径，以确保其行为与训练端一致。关键的技术修复包括：

  • GEMM 内核： 实现基于确定性 MFMA 的 GEMM 内核。
  • 注意力机制： 修复 AITER/CK 注意力的执行调度，并解决分页 KV 布局状态问题。
  • 通信： 实现 HIP IPC 通信并修复分布式集体操作的秩归约顺序。
  • 数学与融合： 解决由编译器融合和特定数学函数实现引起的最后一位差异。
  • logprobs： 实现词汇表归约并修复跨 TP 词汇表合并问题。

RL-Kernel 严格路径中的边界修复

计算边界 严格路径修复
RMSNorm 归约域、epsilon、残差加法和输出边界
Attention 位置、掩码、逻辑分页 KV、拆分策略、LSE 精度和最终类型转换
GEMM 和 SwiGLU K 维度归约、累加精度、后处理、激活和实例化边界
Linear logprob 真实词汇表范围、目标所有权、本地归约和跨秩 LSE 合并
分布式集体操作 负载所有权、数据类型和固定的秩归约顺序

在 AMD Instinct MI300X 上的验证结果

使用单节点上的 8x AMD Instinct MI300X GPU 进行了端到端实验，模型为 Qwen3-8B（BF16）。训练采用 Megatron（TP4/CP2/PP1），推理使用两个 vLLM 引擎（每个 TP4）。

关键发现

  • 零不一致： 在连续 200 步的推理和训练中，mismatch_count 和 max_abs_diff 始终为零。
  • 与原生 vime 的对比： 虽然原生 vime 在每一步都显示出不一致，但 vime + RL-Kernel 配置在整个轨迹中保持了逐位一致。
  • 稳定性： 在整个 200 步运行期间，平均绝对训练-推理 logprob 差值始终保持为零。

成就总结

在 ROCm 上集成 vime 和 RL-Kernel 为 RL 后训练提供了几个关键能力：

  • 逐位一致性： 在 ROCm 硬件上实现训练与推理 logprobs 的精确匹配。
  • 可验证的可复现性： 稳定的一致性保证使 RL 训练轨迹更易于验证。
  • 故障定位： 能够通过算子消融识别数值发散开始的确切系统边界。
  • 执行证据： 对运行时使用的内核、HIP Graph 执行和分页 KV 路径进行了全面验证。

Sources