Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning

Paper: 2610.32951
Authors: Yijing Ke, Tianle Cai, Ruosong Wang
Categories: cs.LG

The Gap

Reinforcement Learning (RL) post-training—via PPO, GRPO, or REINFORCE—has become the engine behind frontier reasoning models. To achieve industrial throughput, RL systems are architected as decoupled dual stacks:

1. Rollout Inference Stack: Optimized inference engines (such as vLLM, SGLang, or TensorRT-LLM) employ continuous batching, fused CUDA/Triton kernels, and aggressive KV cache optimizations.
2. Training Stack: Distributed training frameworks (such as Megatron-LM, DeepSpeed, or PyTorch FSDP) handle distributed forward-backward propagation and optimizer updates in BF16 mixed precision.

The Silent Catastrophic Collapse

Even when both engines run with bit-for-bit identical weights and BF16 tensors, RL training frequently encounters catastrophic divergence, gradient explosion, or entropy collapse.

Why? Because subtle low-level execution differences between inference fused kernels and training operations cause the rollout policy probabilities $\pi_{\text{old}}(y \mid x)$ to numerically diverge from the training forward pass probabilities $\pi_\theta(y \mid x)$.

At iteration step 0, the policy importance sampling ratio: $r_t(\theta) = \frac{\pi_\theta(y_t \mid x, y_{<t})}{\pi_{\text{old}}(y_t \mid x, y_{<t})}$ is mathematically assumed to equal $1.0000$. But due to minute numerical drift across deep layers, $r_t(\theta)$ drifts significantly away from 1. An intended on-policy algorithm unintentionally operates as an uncontrolled off-policy update with high-variance gradients!

Prior remedies have proven prohibitively painful:

• Full Stack Alignment: Forcing training frameworks to use the exact same kernels as inference ruins distributed training scaling.
• Full FP32 Migration: Doubles VRAM usage, cuts max context length in half, and halves GPU training throughput.
• Heuristic Clipping / Penalties: Mutes policy ratios, crippling exploration on hard reasoning problems.

   THE DIVERGENCE BETWEEN INFERENCE ROLLOUT AND TRAINING FORWARD

   Rollout Phase (vLLM / TensorRT in BF16):
     Prompt x ---> Fast Fused Kernels ---> Generates y ~ pi_old(y|x)
                   Log-probabilities: log pi_old(y|x) recorded.
                                  |
                                  v
   Training Phase (Megatron / PyTorch FSDP in BF16):
     Feed (x, y) into identical weights ---> Forward Pass computes log pi_theta(y|x)
                                  |
                                  v
   Numerical Discrepancy Check at Step 0:
     Theoretical Expectation: pi_theta / pi_old = 1.0000
     Reality: Accumulating low-level rounding discrepancies cause ratio to drift!
     Result: Off-policy gradient shocks ===> Policy Collapse!

The Increment

One sentence: The authors discover that full-stack kernel synchronization is completely unnecessary; identifying and aligning where intermediate FP32 values are rounded to BF16 in just three critical regions—RoPE, SwiGLU, and residual RMSNorm—is sufficient to stabilize BF16 LLM reinforcement learning across diverse model sizes.

Core Mechanism

Rather than treating the dual-system discrepancy as an intractable black box, the authors perform systematic numerical surgery to trace where numerical divergence originates:

1. Diagnosis of Three Divergence Hotspots:
   • The authors pinpoint three specific operators responsible for over 90% of logit variance between inference and training forward passes:
     • Rotary Position Embedding (RoPE): Trigonometric sinusoidal rotations $\mathbf{R}_{\Theta, m}^d$. Inference engines often compute rotation frequencies and coordinate multiplications in continuous FP32 registers before casting. In contrast, standard training implementations cast intermediate rotated halves into BF16 prematurely, accumulating significant angular drift over long context lengths.
     • SwiGLU Activation: $\text{SwiGLU}(x) = (\text{SiLU}(x W_{\text{gate}}) \otimes (x W_{\text{up}})) W_{\text{down}}$. Fused inference kernels evaluate the elementwise product in FP32 registers. Standard training implementations cast the gate and up projections to BF16 before computing the Hadamard product.
     • Residual Connection & RMSNorm: $\text{RMSNorm}(x + f(x))$. Inference kernels accumulate residual additions in FP32 and compute variance in full precision before final quantization, whereas naive training pipelines introduce intermediate BF16 truncations.
2. Rounding-Location Alignment (RLA):
   • The proposed fix is surgically non-invasive: preserve intermediate accumulations in FP32 until the final operator output within these three specific regions on the training side.
   • It requires zero changes to:
     • The high-throughput inference engine (e.g., vLLM runs unmodified).
     • Heavy GEMM matrix multiplication or FlashAttention kernels.
     • Model weights, precision formats, or the underlying RL loss objective.
3. The RoPE-Dominance Phenomenon:
   • Across four Qwen-family architectures (1.5B, 7B, 14B), aligning all three regions completely eradicates the training collapse observed under standard BF16 baselines.
   • Remarkably, in the vast majority of runs, aligning the RoPE path alone is sufficient to fully stabilize RL training, even while other low-level numerical discrepancies persist in downstream layers!

   ROUNDING-LOCATION ALIGNMENT (RLA) ARCHITECTURE

   [Training Forward Pipeline]
   
   1. RoPE Region:
      Standard BF16: [Sin/Cos FP32] -> Cast BF16 -> Rotate -> Cast BF16 (Phase Error!)
      RLA Fix:       [Sin/Cos FP32] ---------> FP32 Rotation ---------> Cast BF16 Output
   
   2. SwiGLU Region:
      Standard BF16: Gate(BF16) * Up(BF16) -> BF16 Product (Logit Drift!)
      RLA Fix:       Keep Gate & Up intermediates in FP32 -> Multiply -> Cast BF16 Output
   
   3. Residual RMSNorm:
      Standard BF16: Residual sum rounded to BF16 before variance computation
      RLA Fix:       FP32 residual accumulation -> FP32 variance -> Cast BF16 Output

   RESULT:
   - Inference Engine (vLLM): 100% Unchanged, maximum rollout speed.
   - Matrix Multiplications: 100% Standard BF16 Tensor Core execution.
   - RL Stability: Perfect numerical convergence with ZERO policy collapse!

Key Concepts

• Rounding-Location Alignment (RLA): The practice of matching the exact operational stages where intermediate high-precision numbers (FP32) are cast to low-precision formats (BF16) between inference and training.
• On-Policy Assumption Drift: The breakdown of the policy gradient theorem that occurs when sampling probabilities and training forward probabilities diverge due to hardware kernel implementation differences.
• RoPE Phase Shift: Angular numerical distortion that accumulates across token positions when rotary embedding coordinates are truncated to BF16 before trigonometric multiplication.

Framework Shift

Before (The False Trade-Offs of RL Stabilization):
  Option A: Re-implement inference kernels identically in training -> Degrades training throughput.
  Option B: Force entire training loop to FP32 -> Halves context length, doubles VRAM usage.
  Option C: Add aggressive KL/ratio clipping -> Stifles exploration on reasoning benchmarks.

After (Rounding-Location Alignment):
  Targeted Surgery: Retain FP32 registers only across RoPE, SwiGLU, and Residual RMSNorm.
  Zero Compromise: Keep inference engine untouched, keep BF16 Tensor Core speed.
  Flawless Stability: Completely prevents RL training collapse across all model scales.

Expert Review

Topic Value: Monumental for AI systems engineering. Unintended numerical drift between inference engines and training stacks has been one of the most frustrating, time-wasting silent bugs plaguing frontier post-training teams.

Method Maturity: Pinpoint diagnosis and surgical implementation. By identifying that RoPE, SwiGLU, and RMSNorm rounding locations drive the divergence, it solves the problem with minimal lines of code.

Experimental Rigor: Tested on four separate models within the Qwen family under rigorous RL training regimes, demonstrating conclusive stabilization curves and isolating the outsized role of RoPE rounding.

Writing Quality: Clear, direct, and actionable. Strips away the mystique surrounding RL training instabilities and replaces it with concrete systems engineering analysis.

Overall Verdict: strong accept — An essential operational insight that every team training reasoning LLMs with RL should implement immediately.

Takeaways

• Identical model weights and BF16 tensors do not guarantee on-policy reinforcement learning; where intermediate rounding occurs can break training.
• You do not need to overhaul your entire training framework or switch to FP32; retaining FP32 precision strictly within RoPE, SwiGLU, and RMSNorm solves the issue.
• If you can only fix one operator in your training pipeline, fix the RoPE rotary embedding calculation—its phase stability is the single most dominant factor in preventing policy ratio drift.