2605.14220-training-inference-mismatch-llm-rl
Diagnosing Training Inference Mismatch in LLM Reinforcement Learning
这篇论文把 LLM RL 中训练侧与推理侧对同一 token 序列给出的 logprob 不一致定义为 Training Inference Mismatch (TIM),并用 VeXact 构造 FSDP trainer 与 rollout engine bitwise 对齐的 zero mismatch 基线;实验证明 TIM 这种看似微小的 token level 数值差异可以单独触发 RL training collapse,...
Source
- Title: Diagnosing Training Inference Mismatch in LLM Reinforcement Learning
- arXiv: https://arxiv.org/abs/2605.14220
- HTML: https://arxiv.org/html/2605.14220
- PDF: https://arxiv.org/pdf/2605.14220
- Code/Project: https://github.com/verl-project/vexact
- Authors: Tianle Zhong, Neiwen Ling, Yifan Pi, Zijun Wei, Tianshu Yu, Geoffrey Fox, Peng Wu, Xiao Yu
- Submitted: 2026-05-14
- Current version read: v1, submitted 2026-05-14
- Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL)
作者与关系
- Tianle Zhong: ByteDance; The University of Virginia。Equal contribution。
- Neiwen Ling: ByteDance。Equal contribution;arXiv 提交者。
- Yifan Pi: ByteDance。
- Zijun Wei: ByteDance。
- Tianshu Yu: ByteDance。
- Geoffrey Fox: The University of Virginia。
- Peng Wu: ByteDance。Corresponding author。
- Xiao Yu: ByteDance。Corresponding author。
关系判断:
- 同机构作者群:ByteDance 作者群为核心,覆盖 7 位作者;The University of Virginia 作者群包含 Tianle Zhong 与 Geoffrey Fox。
- 跨机构桥接:Tianle Zhong 同时标注 ByteDance 与 UVA,是 ByteDance-UVA 合作的直接桥接作者。
- 项目组织:代码仓库位于
verl-project/vexact,README 描述 VeXact 是面向 VeRL/VeOmni FSDP 的 zero-mismatch rollout engine,并与 verl 集成。 - 与已存档作者重叠:未发现与已归档论文作者直接重叠。Acknowledgments 提到 Xibin Wu;Xibin Wu 是
2409.19256HybridFlow 的 ByteDance 作者,因此存在 ByteDance/VERL 系统线的协作线索。 - 与已存档论文的主题或方法关系:本文强连接
2409.19256、2025-09-10、2606.00135。它把 HybridFlow/VERL 的 rollout-training 分离进一步推进到 logprob bitwise alignment 问题,并直接引用 Thinking Machines 的 batch-invariant kernel 思路。 - 需要后续确认:VeXact 仓库、VeOmni、verl 版本仍在演化,后续分析使用 VeXact 的论文时需记录 commit、attention backend、kernel set、FSDP/TP 支持范围。
一句话结论
这篇论文把 LLM RL 中训练侧与推理侧对同一 token 序列给出的 logprob 不一致定义为 Training-Inference Mismatch (TIM),并用 VeXact 构造 FSDP trainer 与 rollout engine bitwise 对齐的 zero-mismatch 基线;实验证明 TIM 这种看似微小的 token-level 数值差异可以单独触发 RL training collapse,且 recomputation、bypass、TIS、sequence rejection 等补丁分别缓解不同症状,系统级 zero-mismatch 仍是诊断和稳定训练的基准。
阅读目标与判断边界
本笔记关注:
- TIM 的定义、来源和它与一般 PPO off-policy drift 的区别。
- VeXact 如何构造 zero-mismatch rollout baseline。
- TIM 如何改变 REINFORCE、PPO/GRPO 的有效优化目标。
- recomputation、bypass、truncated importance sampling、sequence-level rejection 的作用边界。
- 这篇论文和 HybridFlow/VERL、Thinking Machines batch invariance、tool-calling RL 的关系。
判断边界:
- 论文是 2026 年 5 月 arXiv v1,结论需要后续版本和更多系统配置复验。
- 实验主要覆盖 Qwen3-1.7B dense、Qwen3-30B-A3B MoE、FSDP2、vLLM/VeXact、数学推理数据集和 AIME 评测。
- 本文主要诊断 logprob mismatch 对 RL 稳定性的影响,没有声称所有 RL collapse 都由 TIM 造成。
- 公式在 arXiv HTML 中部分渲染缺失,本笔记以正文语义、图表说明和 appendix 配置为依据。
论文脉络
1. 问题背景
现代 LLM RL 系统通常把 rollout generation 和 policy optimization 拆成两个阶段。rollout 阶段常由 vLLM、SGLang 等推理引擎生成 token 并记录 logprob;训练阶段由 FSDP、Megatron 等训练引擎重新计算当前 policy、old policy 或 reference policy 的 logprob,然后执行 REINFORCE、PPO、GRPO 等更新。
理论上,如果模型权重、输入 token 序列和采样轨迹相同,rollout engine 和 trainer engine 对每个 token 的概率应保持一致。实际系统中,两侧可能使用不同模型实现、kernel、attention backend、KV cache 组织、tiling、reduction order、MoE routing 或 batch-dependent kernel configuration,导致同一 token 在两侧得到不同 logprob。作者把这类系统级差异称为 Training-Inference Mismatch。
TIM 会在 RL 中制造一类额外 off-policy bias。它和 PPO mini-step 引入的 policy drift 不同:PPO drift 来自多次更新后 current policy 离采样时 policy 变远;TIM 来自同一 checkpoint 在 rollout path 与 training path 上的数值概率景观不一致。
2. 核心假设或切入点
作者的核心判断是:TIM 是 RL stability 的一阶变量。平均 token-level mismatch 可能很小,但最大 token mismatch 会出现极端值,甚至导致 top-1 token 翻转。这种 heavy-tailed mismatch 进入 PPO/GRPO loss 后,会被 advantage 符号、ratio clipping 和 token/sequence 聚合放大,最终改变优化器真正看到的目标。
为了隔离 TIM,作者需要一个 zero-mismatch baseline。普通 vLLM rollout 与 FSDP trainer 同时包含 TIM、PPO mini-step drift、ratio clipping、rejection sampling 等因素,很难做因果判断。VeXact 的作用就是把 rollout engine 和 FSDP trainer 的 logprob 对齐到 bitwise level,让研究者把 TIM 打开或关掉。
3. 方法 / 系统 / 理论框架
论文提出 VeXact:一个轻量 zero-mismatch rollout engine。
VeXact 处理两类 mismatch 来源:
- 模型和 kernel 实现差异:推理引擎和训练引擎常用不同 HuggingFace/model runner、attention kernel、MoE kernel、sampler 和 fused ops。
- kernel reduction order 与 tiling 差异:即使 kernel 语义相同,batch size、auto-tuning、launch-grid、tiling 策略变化也会改变浮点 accumulation order。
VeXact 的主要手段:
- 使用与 FSDP trainer 一致的 HuggingFace-based model implementation。
- 在 FSDP engine 初始化时注册 VeXact kernel implementation,让 trainer 与 rollout 使用同一套可对齐算子。
- 使用 deterministic and batch-invariant kernels,固定 tiling 和 reduction order。
- 参考 Thinking Machines 的 batch-invariant kernel 实现,并额外实现 RMSNorm、batched matrix multiplication、batch-invariant Fused MoE kernels。
- attention 部分禁用 KV splitting 以保证 determinism。
- 为了保留可用吞吐,集成 chunked prefill、CUDAGraph、pipeline parallelism、optimistic KV allocation 和 preemption fallback。
论文还构造了两个诊断层次:
- REINFORCE:单次 on-policy update,避免 PPO mini-step 和 ratio clipping 掩盖 TIM,适合作为因果隔离实验。
- GRPO/PPO-style setting:研究实际训练中 recomputation、bypass、TIS、sequence rejection 等机制如何与 TIM 交互。
4. 结论链条
论文的证据链是:
- rollout engine 与 trainer engine 在同一 checkpoint 和同一 sampled sequence 上会产生 token logprob mismatch。
- TIM 的 batch mean 可能小,但 token max 可以很大,并可能发生 top-1 flip。
- VeXact 提供 zero-mismatch rollout baseline,使 TIM 的因果影响可隔离。
- 在 REINFORCE 中,vLLM non-exact rollout 发生 reward/gradient instability,VeXact 更稳定;TIM 单独足以触发 collapse。
- 在 GRPO 中,recomputation 和 bypass 都不能从根源消除 TIM:前者用 trainer-side denominator 扭曲 sampled distribution,后者仍在 trainer-side probability landscape 上优化。
- KL estimator 对早期 TIM failure 不够敏感,因为问题先体现在 advantage-weighted loss contribution 的符号不平衡和 skew 上。
- TIS 与 sequence-level rejection 可以接近 VeXact baseline,但属于 post-hoc sample filtering,需要 VeXact 这类 zero-mismatch reference 来校准阈值。
关键实验/定理
结果 1:TIM 的 token-level mismatch 可出现极端值
- 设置:Qwen3-8B bf16,在同一 checkpoint 和同一 AIME-2024 greedy sampled response 上比较 rollout-side logprob 与 trainer-side re-evaluated logprob。
- 指标:每个 batch 中 token-level mismatch 的 mean 和 max。
- 结果:mean 通常较小,但 max 在部分 token 上接近 1.0;表格示例中出现 top-1 token choice flip。
- 解读:平均误差掩盖 heavy-tailed token outlier;少数关键 token 会改变 loss ratio、argmax 或后续轨迹。
结果 2:REINFORCE 中 TIM 可以单独触发训练不稳定
- 设置:dense Qwen3-1.7B 和 MoE Qwen3-30B-A3B;使用 REINFORCE with batch-whitened advantages;比较 vLLM non-exact rollout 和 VeXact。dense 使用 Sanity-Test-R1D-1.5B 训练并在 AIME 2024 评估;MoE 使用 DAPO dataset 训练并在 AIME 2024 评估。
- 指标:training reward、AIME 2024 validation reward、TIM mean、gradient norm。
- 结果:vLLM non-exact rollout 在 reward 和 gradient signals 上出现不稳定;VeXact 显著更稳定并继续改善。
- 解读:REINFORCE 没有 PPO mini-step ratio clipping 的复杂干扰,因此该实验支持 TIM 本身是 destabilizing factor。
结果 3:GRPO 中 recomputation 和 bypass 都会失败
- 设置:Qwen3-1.7B、FSDP、数学推理 workload,GRPO;比较 VeXact、vLLM recomputation、vLLM bypass。
- 指标:training reward、AIME24 validation reward、loss、gradient norm、KL estimators、zero-centered loss contribution。
- 结果:VeXact 保持稳定;recomputation 先早期 reward degrade,随后出现 gradient-norm spike 并 collapse;bypass 出现 reward degradation,但 loss spike 更不明显。
- 解读:recomputation 的 denominator 来自 trainer-side reference,偏离真正采样分布;bypass 使用 rollout-side old logprob,但 numerator 和 backward gradient 仍在 trainer-side probability landscape 上,优化器会学习 trainer path 的数值伪影。
结果 4:KL estimator 难以发现早期 TIM failure
- 设置:在 recomputation 和 bypass 下跟踪常见 PPO ratio KL estimators。
- 指标:基于 ratio 的 KL probes 与 reward degradation 时间关系。
- 结果:bypass 下 KL probes 上升明显;recomputation 下,前 700 步 reward 已进入 degradation,但 KL estimators 仍接近 VeXact baseline。
- 解读:TIM 早期问题可能先出现在 objective-space 的 advantage-weighted contribution skew,全局 probability-space divergence 会更晚暴露。
结果 5:TIS + sequence rejection 可接近 VeXact
- 设置:比较四类 algorithmic patches:基于 correction ratio 或 PPO ratio 的 sequence-level rejection,以及 TIS 结合 correction-ratio sequence rejection 的不同阈值配置。
- 指标:training reward、validation reward、loss、gradient norm。
- 结果:使用 correction ratio 作为 filtering signal 比 PPO ratio 更有效;加入 TIS 后,correction-ratio sequence rejection 可以更接近 VeXact。
- 解读:TIM 同时以局部 token outlier 和序列级 accumulated mismatch 表现出来,需要 token-level 和 sequence-level 共同过滤。但这种方式会丢弃样本,需要 zero-mismatch baseline 校准。
结果 6:实验规模和配置
- 设置:appendix 中给出三组配置。Dense GRPO:Qwen3-1.7B,global batch 64、mini-batch 16、rollout group 8、prompt 1024、response 8192、1 node 8 H100。Dense REINFORCE:Qwen3-1.7B,global batch 64、2 nodes 16 H100。MoE REINFORCE:Qwen3-30B-A3B,global batch 512、prompt 2048、response 20480、8 nodes 64 H100。三者 engine 均为 FSDP2 + vLLM/VeXact。
- 指标:训练稳定性和验证 reward。
- 解读:实验覆盖 dense 与 MoE、GRPO 与 REINFORCE,但仍集中在数学推理和有限系统组合上。
证据链强度评估
强证据
- 论文把 TIM 从一般 off-policy drift 和 reward/hyperparameter 问题中分离出来,问题定义清晰。
- VeXact 作为 zero-mismatch baseline 很关键,使 vLLM non-exact 与 VeXact 的对比更接近因果诊断。
- REINFORCE 实验降低了 PPO/GRPO clipping 与 mini-step 干扰,对证明 TIM 本身影响稳定性有较强支撑。
- recomputation/bypass 的机制分析具体,说明“把 rollout logprob 传过来”或“trainer 重算 logprob”各自为什么仍会出问题。
- 与 Thinking Machines batch-invariant inference、HybridFlow/VERL 基础设施形成明确技术承接关系。
中等强度证据
- GRPO correction ablation 展示了可行 patch 组合,但阈值和配置仍依赖 VeXact reference 校准。
- MoE 实验很有价值,但 MoE routing、Fused MoE kernel 和 batch-invariant 实现的细节会显著影响可迁移性。
- AIME/math workloads 对 reasoning RL 很代表性,但不覆盖 tool-use agent、code sandbox、multi-modal 或开放式 reward model 场景。
需要谨慎的推论
- TIM 足以触发 collapse,不代表所有 collapse 都可归因于 TIM;reward misspecification、数据分布、learning rate、advantage normalization 仍可能独立导致失败。
- VeXact 禁用 KV splitting 等选择有稳定性价值,也可能牺牲部分吞吐;大规模生产训练需要重新平衡 determinism 和 throughput。
- algorithmic patches 接近 VeXact 的结果来自有限配置,跨模型、跨任务、跨异步 RL 系统可能需要重新调参。
本地讨论补充
1. 固定 rollout 序列在训练时不会重新分叉
对同一条已经 rollout 出来的序列,trainer 侧通常使用 teacher forcing / prefill 式 forward 来重算 logprob。训练时不会重新逐 token 采样,也不会让这条固定 trajectory 生成出另一条 token 序列。
因此 TIM 的核心影响发生在概率解释和梯度上:
trainer update uses:
这里
一轮更新可以理解为:
- sample:
- update:increase/decrease
- deploy:next rollout samples from
如果
2. PPO 能缓解更新幅度,不能自动消除分布错位
PPO/GRPO 的 clipping 控制 current policy 相对 old policy 的 ratio 幅度。它默认样本来自的 behavior policy 与 loss 中使用的 old policy 属于同一个概率空间。
recomputation 下可以写成:
理想的 behavior denominator 应该对应
那么 token-level ratio 会带入一个误差因子:
当
delta_logprob有 heavy tail,p99/max 明显偏大。- mismatch 与 advantage 相关,导致 advantage-weighted loss contribution 出现系统性 skew。
- mismatch 落在 clipping boundary 附近,改变 token 是否被 clipped。
- 长序列或 tool-use 轨迹让 token-level mismatch 在 sequence level 累积。
- MoE routing、attention split、precision 或 backend 差异触发计算路径变化。
bypass 可以把 rollout old logprob 直接传给 trainer,使 denominator 更接近 behavior policy:
但 numerator 和 backward gradient 仍来自
3. Dense 与 MoE 的放大机制不同
dense model 中,rollout/trainer logprob mismatch 多数时候表现为连续值偏移:某个 token 的 logprob、ratio、loss contribution 发生小幅变化。PPO/GRPO 可以通过 clipping、KL、TIS 或 rejection sampling 吸收一部分误差。
MoE 增加了离散 expert routing。router 根据 routing logits 选择 top-k experts;当两个 expert score 接近时,微小数值扰动可能翻转 top-1/top-2 expert。expert 改变后,token 进入不同 FFN 路径,差异会从 logprob drift 升级为 computation path change。
因此 MoE 下的风险链条更容易是:
small numerical perturbation
-> router top-k flip
-> different expert path
-> hidden-state/logit/logprob drift
-> ratio and loss-contribution skew
-> policy update bias
-> next-rollout behavior degradation
这解释了为什么本文专门在 Qwen3-30B-A3B MoE 上做实验,并且 VeXact 需要 batch-invariant Fused MoE kernels。
4. 判断 TIM 是否重要的实用指标
在具体训练栈中,TIM 应先作为可测量风险处理,再决定是否值得牺牲吞吐换 zero-mismatch。优先记录:
abs(delta_logprob)的 p50 / p95 / p99 / max。- top-1 token flip rate。
- PPO ratio error:
。 delta_logprob与 advantage 的相关性。- clipped tokens 中有多少由 TIM 触发。
- sequence-level accumulated mismatch。
- MoE router top-k flip rate。
- rollout backend 与 trainer backend 的 attention、precision、KV splitting、MoE kernel、batch-invariant kernel 配置。
更稳妥的结论是:TIM 的平均误差小,不说明训练影响一定小;TIM 存在,也不说明 PPO/GRPO 一定会失稳。关键取决于 token-level tail、advantage correlation、sequence accumulation、MoE routing flip 和系统实现差异。
主要启发
- RLHF/RLVR 框架需要把 rollout engine 与 trainer engine 的 logprob consistency 作为一等指标,同时报告 reward curve 和 throughput。
- 在 on-policy RL 中,sampler 的 behavior policy 和 trainer 的 optimization policy 必须明确定义;实现路径不同会把系统细节变成隐式 optimization perturbation。
- 对 VERL/HybridFlow 这类系统,编排 rollout 和 policy update 只是第一层;第二层是保证 rollout-side probability landscape 与 trainer-side landscape 可校验、可对齐。
- 对 tool-calling RL,TIM 影响可能更大:训练中 fixed trajectory 不会重新分叉,但下一轮 rollout 中的工具调用参数、工具返回、轨迹长度和 reward 会受到更新后 rollout policy 的影响。
- 对评测与复现,必须记录 trainer backend、rollout backend、precision、attention backend、KV splitting、batch-invariant kernel 开关、recomputation/bypass 策略和 correction thresholds。
- VeXact 的更大价值可能是 calibration tool:先在 zero-mismatch 环境中确认算法补丁的真实作用,再迁移到高吞吐 rollout engine。
局限
- 论文是 arXiv v1,实验覆盖面有限,更多模型、任务、backend、硬件组合仍需复验。
- VeXact 为了 zero-mismatch 采用固定 tiling、batch-invariant kernels 和禁用 KV splitting 等策略,吞吐成本需要和 vLLM/SGLang 高性能路径系统比较。
- 本文没有给出大规模异步 agent RL、tool-use、code execution reward 或 multi-modal RL 的完整实证。
- correction methods 的阈值依赖实验配置;缺少 VeXact baseline 时仍容易走向经验调参。
- 跨 tensor parallel size、MoE routing、quantization、speculative decoding、distributed reduction 的 zero-mismatch 仍是后续系统难题。
- 论文重点分析 stability,对最终模型泛化、样本效率和长期 capability tradeoff 的结论仍有限。
跨论文关系
- 与
2025-09-10的作者关系:未发现作者重叠。主题关系最直接。Thinking Machines 文章解释 batch-invariant inference 如何消除服务端 batch/slicing 非确定性;本文引用该思路并把它放入 LLM RL,证明 trainer-rollout mismatch 会改变优化目标并触发 collapse。 - 与
2409.19256的作者关系:未发现直接作者重叠,但本文和 HybridFlow 同属 ByteDance/verl 系统谱系,并在 acknowledgments 中感谢 HybridFlow 作者 Xibin Wu。方法关系很强:HybridFlow/VERL 解决 RLHF dataflow 编排和训练/推理系统衔接;本文进一步解决衔接后的 logprob bitwise alignment。 - 与
2606.00135的作者关系:未发现作者重叠。主题关系很强。2606.00135关注 tool-calling RL 的 harness、rollout 和 policy update 效率;本文说明 rollout backend 与 trainer backend 的数值差异会让 policy update 偏离真实采样分布,尤其会影响长轨迹 tool-use。 - 与
2606.06453的作者关系:未发现作者重叠。系统关系中等。Vortex 关注 sparse attention serving 的吞吐与可编程性;本文关注 deterministic/batch-invariant inference 的稳定性。两者都说明 attention/kernel 细节会影响上层 agent/RL 行为。 - 与
2605.30290的作者关系:未发现作者重叠。主题关系中等。Self-trained verification 依赖 verifier/reasoner 闭环和 RLVR;本文提示 verifier-guided rollout 与 trainer 之间也可能存在 logprob/trajectory consistency 风险。 - 与
2606.04075的作者关系:未发现作者重叠。主题关系中等。SocioHack 关注 reward 目标诱导模型搜索制度漏洞;本文关注系统数值 mismatch 如何诱导优化器学习数值伪影。两者都属于 RL 闭环中的隐藏偏差。 - 与
2510.19315的作者关系:未发现作者重叠。关系较弱;一个是 RL 系统数值稳定性,一个是 Transformer 理论复杂性。 - 与
2605.31514的作者关系:未发现作者重叠。方法论关系中等:两者都提醒研究者谨慎归因,把模型能力、接口、系统路径和测量条件分开记录。 - 新增后应更新的索引 cluster:新增 “Training-Inference Mismatch 与 Zero-Mismatch RL” cluster,并将其连接到 HybridFlow/VERL、TML batch-invariant inference 和 tool-calling RL。
Reference Intake Brief
Target
- Intended target system:
paper archive root论文存档。 - Existing related assets:
papers-index.md、2409.19256-hybridflow-rlhf-framework.md、2025-09-10-defeating-nondeterminism-llm-inference.md、2606.00135-agentic-tool-calling-rl-training.md。 - Proposed form: 新建独立 Markdown 文档,并更新总索引。
Reusable Elements
- TIM definition:同一 checkpoint、同一输入、同一 sampled sequence 下 rollout-side logprob 与 trainer-side logprob 不一致。
- VeXact baseline:FSDP trainer 与 rollout engine bitwise alignment。
- Failure analysis:recomputation denominator skew、bypass probability landscape mismatch、KL estimator blind spot、zero-centered loss contribution skew。
- Mitigation map:TIS、correction-ratio sequence rejection、token+sequence filtering、VeXact calibration。
Risks
- Copyright/over-copying: 本笔记采用转述,避免复制长段原文。
- Unsourced or unverifiable claims: 元数据来自 arXiv abs/html;项目现状来自 GitHub README;跨论文关系为本地分析判断。
- Tone/brand mismatch: 保持本目录技术笔记风格。
- Safety/compliance issues: 论文是 RL 系统稳定性与可复现主题,无直接安全滥用流程。
- Overlap with existing assets: 与
2025-09-10、2409.19256、2606.00135关系强,但本文作为 TIM/zero-mismatch RL 诊断论文单独存档。
Skipped
| Material | Reason |
|---|---|
| arXiv HTML 中缺失的完整公式符号 | HTML 渲染丢失部分数学变量,本笔记保留语义和机制。 |
| VeXact 仓库全部实现细节 | 当前目标是论文分析;代码可后续按 commit 做系统阅读。 |
| 所有 correction objective 展开 | appendix 已列出变体,本笔记保留作用机制和实验结论。 |
Recommendation
Decision: merge
Why: 该论文直接连接 HybridFlow/VERL、Thinking Machines batch-invariant inference 和 tool-calling RL,补齐了本目录中 RL rollout 与 trainer 数值一致性的关键节点。