诊断大语言模型强化学习中的训练-推理不匹配

# 诊断大语言模型强化学习中的训练-推理不匹配
来源: https://arxiv.org/html/2605.14220
commentstyle=, keywordstyle=, stringstyle=, basicstyle=, identifierstyle=, backgroundcolor=, breakatwhitespace=false, breakindent=0pt, breaklines=true, captionpos=b, keepspaces=true, numbers=left, numberstyle=, numbersep=3pt, showspaces=false, showstringspaces=false, showtabs=false, framexleftmargin=0pt, frame=lines, rulecolor=, rulesepcolor=, xleftmargin=0pt, xrightmargin=0pt,

Tianle Zhong¹˒²∗Neiwen Ling¹∗Yifan Pi¹Zijun Wei¹ Tianshu Yu¹Geoffrey Fox²Peng Wu¹†Xiao Yu¹† ¹字节跳动²弗吉尼亚大学 ∗同等贡献†通讯作者

###### 摘要

现代大语言模型（LLM）强化学习（RL）系统将生成轨迹与策略优化相分离。理论上，这两个阶段应为相同序列和相同模型权重生成完全匹配的token概率。然而，实现上的差异可能导致其对同一序列赋予不同值，从而引发训练-推理不匹配（Training–Inference Mismatch, TIM）。由于TIM与离策略漂移和常见的稳定机制相互交织，因此难以检查。本文在零不匹配的诊断设置（VeXact）中隔离了TIM，并表明微小的token级数值差异可以独立导致训练崩溃。我们进一步表明TIM改变了有效优化问题，并确定了一组可以缓解TIM的补救措施。我们的结果表明，TIM并非无害的数值噪声，而是一种系统级扰动，在分析LLM RL稳定性时应将其视为首要因素。

## 1 引言

大语言模型（LLM）强化学习（RL）已成为基础模型后训练的核心范式，也是近期复杂推理能力进步的关键驱动力（Schulman等, 2017 (https://arxiv.org/html/2605.14220#bib.bib24); Ouyang等, 2022 (https://arxiv.org/html/2605.14220#bib.bib16); Ziegler等, 2019 (https://arxiv.org/html/2605.14220#bib.bib17); Shao等, 2024 (https://arxiv.org/html/2605.14220#bib.bib25); Stiennon等, 2020 (https://arxiv.org/html/2605.14220#bib.bib18); Liu等, 2025b (https://arxiv.org/html/2605.14220#bib.bib54)）。然而，在实践中，RL训练仍然难以稳定：策略可能会迅速退化，导致奖励信号在短训练窗口内下降。

因此，理解导致这些崩溃的原因对于构建可靠的LLM RL训练系统至关重要（Sheng等, 2025b (https://arxiv.org/html/2605.14220#bib.bib46); Hu等, 2024 (https://arxiv.org/html/2605.14220#bib.bib45); Fu等, 2025b (https://arxiv.org/html/2605.14220#bib.bib49); Team, 2025 (https://arxiv.org/html/2605.14220#bib.bib47); Cao等, 2025 (https://arxiv.org/html/2605.14220#bib.bib48); Sheng等, 2025a (https://arxiv.org/html/2605.14220#bib.bib50); MiniMax等, 2025 (https://arxiv.org/html/2605.14220#bib.bib35)）。然而，诊断根本原因很困难，因为许多失败模式是深度纠缠的，并且出现在训练堆栈的不同层次。崩溃可能由多种因素引起，例如超参数调谐不当、奖励指定错误和奖励黑客（Fu等, 2025a (https://arxiv.org/html/2605.14220#bib.bib52); Pan等, 2024 (https://arxiv.org/html/2605.14220#bib.bib53)）。在这些因素中，训练-推理不匹配（TIM）是一个基础设施层面的混淆因素：训练引擎和推理引擎之间的实现差异可能导致即使对于相同的输入和模型权重，token概率也不同。

为了应对这些稳定性挑战，社区已经开发了一系列训练级别的稳定技术，包括重要性采样、拒绝采样和其他形式的保守策略更新（Schulman等, 2017 (https://arxiv.org/html/2605.14220#bib.bib24); Yao等, 2025 (https://arxiv.org/html/2605.14220#bib.bib29); Li等, 2026 (https://arxiv.org/html/2605.14220#bib.bib33); Team等, 2025 (https://arxiv.org/html/2605.14220#bib.bib1); Ring AI Team, 2025 (https://arxiv.org/html/2605.14220#bib.bib21); Zheng等, 2025a (https://arxiv.org/html/2605.14220#bib.bib19); Liu等, 2025a (https://arxiv.org/html/2605.14220#bib.bib20)）。尽管在某些设置下有效，但它们与特定失效机制的联系仍不清楚：同一技术可能纠正PPO小步长离策略漂移、抑制TIM引起的数值异常，或引入额外的优化偏差。如果没有无TIM的诊断基线来隔离这些影响，从业者必须通过试错而非因果诊断来调整干预措施和过滤阈值。

在本文中，我们旨在系统理解TIM对LLM RL稳定性的影响。具体来说，我们旨在回答两个关键问题：首先，TIM是否会导致RL训练不稳定性，如果是，影响程度如何？其次，常见稳定技术如何与TIM相互作用，它们缓解了不匹配的哪些方面，又引入了哪些优化副作用？

为了回答这些问题，我们开发了VeXact¹，一个轻量级生成引擎，能够在VeRL（Sheng等, 2025b (https://arxiv.org/html/2605.14220#bib.bib46)）之上实现与FSDP（Zhao等, 2023 (https://arxiv.org/html/2605.14220#bib.bib38); Rajbhandari等, 2020 (https://arxiv.org/html/2605.14220#bib.bib39)）引擎的零不匹配。VeXact通过统一内核和模型实现与FSDP训练引擎，并采用批量不变内核（He, 2025 (https://arxiv.org/html/2605.14220#bib.bib3)）（§3.1 (https://arxiv.org/html/2605.14220#S3.SS1)）来消除TIM。使用VeXact，我们对LLM RL稳定性进行了细粒度的诊断研究。具体来说，我们的贡献是：

**隔离TIM对LLM RL的影响：** 利用我们的无TIM基线，我们确定TIM本身是触发RL训练崩溃的重要因素（§3.2 (https://arxiv.org/html/2605.14220#S3.SS2)）。

**分析TIM诱导的RL训练崩溃的失败模式。** 然后，我们针对一般设置下TIM在RL训练崩溃中的作用进行消融研究。具体来说，我们分析了为什么在训练侧对数概率重新计算和生成侧对数概率绕过两种情况下RL训练会崩溃。我们发现TIM从根本上改变了优化目标，从而引发了不同的失败（§4.1 (https://arxiv.org/html/2605.14220#S4.SS1)）。

**消融算法性TIM补偿的有效性。** 此外，我们评估了常见稳定技术是否能有效缓解TIM，包括截断重要性采样（Truncated Importance Sampling, TIS）（Yao等, 2025 (https://arxiv.org/html/2605.14220#bib.bib29)）和拒绝采样（Rejection Sampling, RS）（Li等, 2026 (https://arxiv.org/html/2605.14220#bib.bib33)）（§4.2 (https://arxiv.org/html/2605.14220#S4.SS2)）。基于我们的消融研究，我们确定了一种有效的现有算法性TIM补偿组合，可以紧密跟踪我们的无TIM基线。

## 2 LLM RL 中的训练-推理不匹配

由于训练引擎（FSDP（Zhao等, 2023 (https://arxiv.org/html/2605.14220#bib.bib38); Rajbhandari等, 2020 (https://arxiv.org/html/2605.14220#bib.bib39)）、Megatron（Shoeybi等, 2020 (https://arxiv.org/html/2605.14220#bib.bib40)）等）和推理引擎（vLLM（Kwon等, 2023 (https://arxiv.org/html/2605.14220#bib.bib11)）、SGLang（Zheng等, 2024 (https://arxiv.org/html/2605.14220#bib.bib41)）等）之间的实现差异，包括不同的模型/内核实现，即使使用完全相同的模型检查点和输入，词表上下一个token的概率分布也可能不同。这在采样和模型更新之间引入了非预期的离策略偏差。与PPO小步长引入的离策略偏差（Schulman等, 2017 (https://arxiv.org/html/2605.14220#bib.bib24)）不同，TIM离策略偏差是一种基础设施层面的噪声，无法通过简单的PPO裁剪方法解决（在§4 (https://arxiv.org/html/2605.14220#S4)中讨论）。

我们将RL目标中的这个问题公式化如下：给定上下文 \(x\) 和采样响应 \(y=(a_1,\ldots,a_T)\)，令 \(s_t=(x,y_{<t})\)。我们区分三个token级分布：\(\pi_\theta(a_t|s_t)\) 表示当前正在优化的策略；\(\pi_\text{old}^\text{rollout}(a_t|s_t)\) 表示采样token时生成引擎所实现的行为分布；\(\pi_\text{old}^\text{train}(a_t|s_t)\) 表示当算法需要旧策略概率时训练侧使用的参考分布。

在精确的在策略实现中，每个采样token的概率在生成和训练之间应保持一致。当生成执行路径和训练执行路径在相同模型权重和采样序列下为同一token分配不同概率时，就会发生TIM。在token级别，这种差异可以写为

\[\delta_t = \log\pi_\text{old}^\text{train}(a_t|s_t) - \log\pi_\text{old}^\text{rollout}(a_t|s_t). \tag{1}\]

该定义与目标无关：无论更新是使用REINFORCE（Williams, 1992 (https://arxiv.org/html/2605.14220#bib.bib36); Hu等, 2025 (https://arxiv.org/html/2605.14220#bib.bib34)）、PPO（Schulman等, 2017 (https://arxiv.org/html/2605.14220#bib.bib24)）还是GRPO（Shao等, 2024 (https://arxiv.org/html/2605.14220#bib.bib25); Yu等, 2025 (https://arxiv.org/html/2605.14220#bib.bib27)）实现，该不匹配都存在。

表1 (https://arxiv.org/html/2605.14220#S2.T1) 展示了一个示例，来自相同上下文和模型检查点采样的几个token。对于每个token，我们比较生成线程产生的对数概率与训练引擎重新评估的对数概率。理想情况下，对于在策略更新，这两个值应完全相同。然而，我们可以观察到两侧的token对数概率存在明显差异，包括某些位置top-1 token选择被翻转的情况。如图1 (https://arxiv.org/html/2605.14220#S2.F1) 所示，尽管每个训练批次中token级概率的平均差异很小，但对于某些极端token，最大差异甚至可以达到1.0，这在TIM存在时通常可以观察到。

| The problem states that | † | thereexistrealnumbers | \(\log\pi_\text{rollout}\) | -0.279 | -0.063 | -0.314 | -0.694 | -0.000 | -0.030 | -0.000 | -0.000 |
|------------------------|---|---------------------|--------------------------|--------|--------|--------|--------|--------|--------|--------|--------|
| \(\log\pi_\text{train}\) |   |                     | -0.278                   | -0.063 | -0.314 | -0.827 | -0.000 | -0.038 | -0.000 | -0.000 |
| \(\delta_t\)              |   |                     | +0.001                   | +0.000 | +0.000 | -0.133 | +0.000 | -0.008 | +0.000 | +0.000 |

表 1: 在相同 Qwen3-8B (bf16) 权重上，生成栈和训练栈之间的 token 级内核数值漂移，沿着从 AIME-2024 问题上贪心采样响应的一个句子追踪。\(\delta_t = \log\pi_\text{train} - \log\pi_\text{rollout}\)。大多数位置是位接近的，但 † 标记了一个 argmax 翻转：在 *thereexistrealnumbers* 处，\(\pi_\text{train}\) 的 top-1 token 实际上是标点 + 换行符字符串 ":\n\n"（对数概率 -0.577），因此训练侧本会将从句结束为 "The problem states:\n\n"。

(a) \(|\delta_t|\) (最大值)
(b) \(|\delta_t|\) (平均值)

图 1: Qwen3-1.7B GRPO 实验中每个训练批次的统计 \(|\delta_t|\) 最大值和平均值（详细配置见附录 A.1 (https://arxiv.org/html/2605.14220#A1.SS1)）。虽然 \(|\delta_t|\) 的平均值很小，但我们可以观察到一些极端 token 的 \(|\delta_t|\) 接近 1.0。

## 3 使用 VeXact 隔离 TIM

在本节中，我们将隔离 TIM 对 RL 训练稳定性的影响。这需要两个关键要素：(1) 首先，我们引入一个无 TIM 的生成实现 VeXact，作为一个诊断基线，从 RL 循环中移除基础设施引起的不匹配。(2) 其次，我们在 REINFORCE 下评估该基线，这避免了可能掩盖或扭曲 TIM 引起的损失和梯度信号变化的 PPO 比率裁剪。

### 3.1 VeXact：零不匹配生成引擎

对于无 TIM 基线，我们引入了 VeXact，这是一个轻量级生成引擎，其生成 token 对数概率可以与 FSDP 引擎实现位级对齐。

TIM 来自两个来源：(1) 推理引擎和训练引擎之间的模型和内核实现差异。尽管在语义和数学上相同，但它们往往在实现细节上做出不同的决策。例如，推理引擎更喜欢推理优化的内核库（如 FlashInfer），而这在训练引擎中不适用。(2) 内核归约顺序和分块的变化。即使使用相同的内核实现，面向性能的优化（如原子加法）也可能引入非确定性，导致内核为相同输入产生不同输出。此外，即使是确定性内核也可能表现出与批量相关的数值行为：批量大小的变化可以通过自动调优触发不同的发射网格配置，从而改变 GPU 分块策略和归约顺序。由于在有限精度下浮点累加是非结合的，这些执行顺序的变化最终可能导致数值不同的结果。

因此，VeXact 通过以下方式解决这两个不匹配来源：(1) 使用相同的基于 HuggingFace 的模型实现，并在 FSDP 引擎初始化中注册 VeXact 内核实现；(2) 采用确定性和批量不变的内核，固定 GPU 内核实现中的分块和归约顺序。遵循原始的批量不变内核实现（He, 2025 (https://arxiv.org/html/2605.14220#bib.bib3)），VeXact 额外实现了 RMSNorm（Zhang 和 Sennrich, 2019 (https://arxiv.org/html/2605.14220#bib.bib42)）、批量矩阵乘法以及用于高效 MoE 训练/推理的批量不变 Fused MoE 内核。对于注意力实现（Dao 等, 2022 (https://arxiv.org/html/2605.14220#bib.bib6)），我们特别禁用了 KV 分割（Dao 等, 2023 (https://arxiv.org/html/2605.14220#bib.bib37)）以确保确定性。

同时，由于批量不变内核中的固定分块以性能换取数值稳定性，VeXact 通过集成块预填充（Agrawal 等, 2023 (https://arxiv.org/html/2605.14220#bib.bib12)）、CUDAGraph（NVIDIA Corporation, 2025 (https://arxiv.org/html/2605.14220#bib.bib13)）、流水线并行（Huang 等, 2019 (https://arxiv.org/html/2605.14220#bib.bib14)）以及带有抢占回退的乐观 KV 分配，为实际 RL 训练保持了合理的吞吐量。VeXact 保持可修改且非常轻量，其代码行数与 nano-vLLM（GeeeekExplorer, 2025 (https://arxiv.org/html/2605.14220#bib.bib15)）相似。

(a) MoE 训练奖励
(b) MoE 验证奖励
(c) MoE \(\delta_t\)（平均值）
(d) MoE 梯度范数
(e) 密集训练奖励
(f) 密集验证奖励
(g) 密集 \(\delta_t\)（平均值）
(h) 密集梯度范数