只需对齐 RoPE 等三处舍入位置,就能稳定 LLM 强化学习训练
Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning
做 RL 训练的朋友看看这篇:不用重写训练框架,只在 RoPE 那条路保留 FP32 中间量就能防崩溃,四个 Qwen 模型都验证过了。
RL 训练通常用分离的推理与训练系统,即使权重同为 BF16,实现差异也会让本应 on-policy 的更新变成 off-policy,导致训练崩溃。论文提出在 RoPE、SwiGLU、residual RMSNorm 三个区域把训练侧舍入位置与推理对齐,让选定的中间量保持 FP32。在四个 Qwen 系列模型上,三处全部对齐可避免默认 BF16 基线下的训练崩溃,仅在 RoPE 路径对齐也能奏效。该方法不改动推理引擎、矩阵乘与注意力内核,也不改学习目标。
Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning
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.