Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning
作者: Yijing Ke, Tianle Cai, Ruosong Wang
分类: cs.LG
发布日期: 2026-10-23
💡 一句话要点
通过对齐特定部分实现BF16强化学习的稳定性
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 强化学习 大语言模型 数值稳定性 舍入对齐 模型训练 推理优化 Qwen系列模型
📋 核心要点
- 现有的强化学习方法在推理和训练过程中存在实现差异,导致训练不稳定。
- 本文提出通过对齐训练侧的舍入位置与推理,来稳定BF16强化学习,聚焦于RoPE、SwiGLU和残差RMSNorm区域。
- 实验结果表明,对齐所有三个区域可以有效防止模型崩溃,RoPE路径对齐也能在一定程度上保持稳定性。
📝 摘要(中文)
大语言模型的强化学习通常采用分开的推理和训练系统。即使权重和BF16张量相同,实施差异也可能导致预期的在政策更新变为离政策,从而不稳定训练。现有的解决方案通常对齐整个推理和训练栈,改变全局精度或修改学习算法。本文探讨了是否可以通过对齐执行栈的选定部分来稳定BF16 RL,同时保留其他实施差异。我们在RoPE、SwiGLU和残差RMSNorm三个区域对训练侧的舍入位置与推理进行对齐,通过在FP32中保留选定的中间结果,直到其BF16输出。在四个Qwen系列模型中,对齐所有三个区域可以防止在默认BF16基线下观察到的崩溃;仅RoPE路径对齐在研究运行中也足够,尽管仍存在数值差异。这些干预措施保持了推理引擎、矩阵乘法和注意力内核以及学习目标不变。
🔬 方法详解
问题定义:本文旨在解决大语言模型强化学习中推理与训练之间的实现差异导致的训练不稳定性问题。现有方法通常需要对齐整个系统,增加了复杂性和计算成本。
核心思路:论文提出的核心思路是仅对齐执行栈的特定部分(RoPE、SwiGLU和残差RMSNorm),以减少实现差异对训练的影响,从而实现稳定性。
技术框架:整体架构包括三个主要模块:推理引擎、训练过程和舍入位置对齐。推理引擎负责生成输出,训练过程则进行模型更新,而舍入位置对齐确保在特定区域内的数值一致性。
关键创新:最重要的技术创新在于通过选择性对齐特定区域的舍入位置,避免了对整个系统的重构,从而降低了复杂性并提高了训练的稳定性。
关键设计:在设计中,选择在FP32中保留中间结果,直到生成BF16输出,确保了数值的精确性。此外,保持推理引擎和学习目标不变,进一步增强了方法的有效性。
📊 实验亮点
实验结果显示,通过对齐RoPE、SwiGLU和残差RMSNorm三个区域,成功防止了模型崩溃,且RoPE路径对齐在多次实验中也表现出良好的稳定性。与默认的BF16基线相比,本文方法显著提高了训练的可靠性和效果。
🎯 应用场景
该研究的潜在应用领域包括大语言模型的训练和推理优化,尤其是在需要高精度和稳定性的场景中,如自然语言处理、对话系统和文本生成等。未来,该方法可能推动更多高效的强化学习算法的开发,提升模型在复杂任务中的表现。
📄 摘要(原文)
Reinforcement learning for large language models often uses separate systems for inference and training. Even with identical weights and BF16 tensors, implementation differences can make an intended on-policy update off-policy and destabilize training. Existing remedies align the entire inference and training stacks, change global precision, or modify the learning algorithm. We ask whether aligning selected parts of the execution stack can stabilize BF16 RL while other implementation differences remain. We align training-side rounding locations with inference in three regions---RoPE, SwiGLU, and residual RMSNorm---by retaining selected intermediates in FP32 until their BF16 outputs. Across four Qwen-family models, aligning all three regions prevents the collapse observed under the default BF16 baseline; RoPE-path alignment alone also suffices in the studied runs, despite remaining numerical discrepancies. The interventions keep the inference engine, matrix-multiplication and attention kernels, and learning objective unchanged.