Rounding-location alignment stabilizes BF16 reinforcement learning across inference and training stacks
Ke, Cai, and Wang reveal that training instability in decoupled LLM RL pipelines stems from subtle rounding-location discrepancies across inference and training engines, demonstrating that targeted FP32 retention in RoPE, SwiGLU, and RMSNorm prevents policy collapse.
Primary source: arXiv:2610.32951: https://arxiv.org/abs/2610.32951
What’s new: Modern post-training pipelines for large language models decouple inference (generation and rollouts) from training (gradient computation and parameter updates) across different runtimes, such as using specialized inference servers like vLLM alongside training frameworks like Megatron-LM or DeepSpeed. Even when identical bfloat16 (BF16) weights and tensors are shared, subtle floating-point implementation differences cause generated rollouts to diverge from the training graph. In Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning (arXiv:2610.32951, 23 Oct 2026), researchers Yijing Ke, Tianle Cai, and Ruosong Wang analyze how micro-level numerical discrepancies turn intended on-policy reinforcement learning into off-policy divergence, causing catastrophic policy collapse.
Rather than overhauling global precision to FP32 or enforcing rigid full-stack engine symmetry across rollouts and training, the authors pinpoint three critical operation boundaries where intermediate rounding creates large accumulated errors: Rotary Position Embeddings (RoPE), SwiGLU activation functions, and residual RMSNorm layers. By retaining intermediate variables in FP32 precision until their final BF16 outputs across these three regions, the training graph mirrors the inference engine’s numerical trajectory. Across four models in the Qwen family, this targeted rounding-location alignment completely prevents training collapse under BF16 RL. In fact, aligning the RoPE pathway alone proved sufficient to prevent collapse in their evaluation suite, even while other matrix-multiplication and attention kernels maintained numerical discrepancies.
Why it matters: Reinforcement learning with verifiable rewards (RLVR) and PPO-based post-training are notoriously sensitive to numerical drift. Discrepancies between rollout log-probabilities and backward graph evaluations inject spurious policy-ratio penalties or rewards, frequently triggering sudden training blowups. Ke, Cai, and Wang demonstrate that RL practitioners do not need to incur the 2x memory penalty of full FP32 training or sacrifice the serving speed of optimized inference engines. Pinpointing and aligning rounding boundaries in a handful of architectural components restores mathematical stability without changing attention kernels, serving runtimes, or optimization algorithms.
Caveats: While RoPE-path alignment prevented collapse across the evaluated Qwen architectures and benchmark environments, the authors note that deeper networks or more complex multi-step reasoning rollouts may still suffer from accumulated drift in unaligned GEMM operations. Furthermore, retaining FP32 registers across RoPE, SwiGLU, and RMSNorm requires custom kernel definitions or compiler modifications in standard distributed training harnesses.