PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yijing Ke, Tianle Cai, Ruosong Wang

arXiv:2610.32951v1cs.LGcs.AI

TL;DR

Separate inference and training implementations can disagree at identical weights, making intended on-policy LLM RL updates off-policy and destabilizing training. This paper aligns selected training-side rounding locations with inference while retaining FP32 intermediates until BF16 outputs. Across four Qwen-family models, all-three-region alignment prevents collapse under the default BF16 baseline, and RoPE-path alignment alone also suffices in the studied runs.

  • Problem

    Separate inference and training implementations can disagree at identical weights, creating training–inference mismatch that can destabilize LLM RL.

    Source: §1 · p. 1

  • Method

    The paper aligns training-side rounding locations with inference in selected regions by retaining intermediates in FP32 until BF16 outputs.

    Source: §1 · p. 2

  • Results

    Across four Qwen-family models, aligning all three regions prevents collapse under the default BF16 baseline, while RoPE-path alignment alone also suffices in every studied run.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support inspecting inference and training graphs to identify and align rounding locations moved by fusion as a targeted intervention for stabilizing BF16 RL.

    Source: §9 · p. 9

    Experiments cover two math datasets and one architectural family, so broader experiments are needed to establish when alignment suffices across architectures and training settings.

    Source: §9 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Separate inference and training implementations can disagree even at identical weights, turning intended on-policy updates off-policy and destabilizing LLM RL. The paper investigates whether targeted rounding-location alignment can stabilize BF16 RL while other execution differences remain.

1 Introduction

  • Training–inference mismatch can make recomputed probabilities differ from those of the policy that generated responses, destabilizing an intended on-policy update.LLM RL commonly uses a separate inference engine and training implementation.Source: §1 · p. 1
  • Prior remedies change global precision or correct the learning algorithm, while exact-agreement approaches align arithmetic and reduction order across engines.The paper asks whether selected execution regions can instead be aligned while other implementation differences remain.Source: §1 · p. 1
  • A rounding-location mismatch arises when fused and separate kernels round intermediates at different points, even if both ultimately return BF16 outputs.Separate kernels may store an intermediate in BF16, while a fused kernel can retain it in FP32.Source: §1 · p. 2
  • The proposed alignment retains selected training intermediates in FP32 before rounding region outputs to BF16 in RoPE, SwiGLU, and residual RMSNorm.The selected regions were identified by inspecting execution graphs; model weights and the learning objective remain BF16 and unchanged, respectively.Source: §1 · p. 2
  • Across four Qwen-family models, aligning all three regions prevents the collapse observed under the default BF16 baseline, and RoPE-path alignment alone also suffices in every studied run.The interventions preserve the inference engine, matrix-multiplication and attention kernels, and learning objective.Source: §1 · p. 2

2 Related Work

Related work addresses training–inference mismatch through algorithmic corrections, execution alignment, and precision or quantization changes. The paper situates selective rounding alignment among approaches that seek numerical agreement between inference and training.

Corrections for training–inference mismatch

  • Mismatch remedies include truncated importance sampling, trust-region masking, learning-rate reductions, and candidate-update checks using an inference-side discrepancy proxy.Zheng et al. analyze numerical disagreement and policy staleness in token-level surrogate objectives.Source: §2 · p. 2

Execution alignment

  • Execution-alignment approaches include invariant reductions, shared implementations, expert-selection replay, module-by-module alignment, and matching RoPE and SwiGLU rounding locations.Zhong et al. report that numerical disagreement alone can cause collapse; Panda matches rounding locations as part of achieving bitwise-identical token log probabilities.Source: §2 · p. 2

Precision and quantization

  • Precision and quantization approaches range from BF16-to-FP16 changes to low-bit or FP8 inference and training schemes, sometimes combined with corrections or modified objectives.Related work also reports BF16 errors disrupting RoPE's relative-position property in long-context training [31].Source: §2 · p. 2

3 Training-Inference Mismatch in LLM RL

LLM RL commonly separates response generation from optimization, so implementation-level numerical differences can create training–inference mismatch before learning updates begin. The section describes sources of that mismatch and the sampling, objective, measurement, and evaluation setup used in the experiments.

3.1 Where the mismatch comes from

  • Mismatch can come from precision and cast locations, different kernels or decompositions, and reduction order; fusion can retain an intermediate in FP32 where separate kernels store it in BF16.Both paths can return BF16 outputs even though the next operation receives different intermediate values.Source: §3.1 · p. 3

3 Training-Inference Mismatch in LLM RL

  • Inference samples responses using μθ0, while training recomputes their probabilities with πθ and updates from the same checkpoint; the implementations can differ numerically even at identical weights.Each response receives a correctness reward, and its advantage is centered within the group without standard-deviation normalization.Source: §3 · p. 2, §3.2 · p. 3

3.2 Algorithm and experimental setup

  • The experiments use PPO-clipped GRPO without additional mismatch corrections and hold that objective fixed across conditions.At the start of an update, the training-side ratio is one, so clipping leaves the initial inference–training discrepancy uncorrected; discrepancy is measured at fixed parameters using an estimated per-token forward KL.Source: §3.2 · pp. 3–4
  • Main-text RL experiments train Sanity-Test-R1D-1.5B on a selected MATH subset and evaluate on AIME 2025 problems.vLLM generates responses, while Hugging Face Transformers recomputes probabilities and updates the model using FSDP.Source: §3.2 · p. 4

4 Motivation: Small Execution Changes Can Destabilize RL

A BF16 experiment compares compiled and eager inference and training modes to test whether small execution changes affect RL stability. Aligned execution remains stable, whereas crossed execution collapses; changing inference to match fixed eager training restores stability.

4 Motivation: Small Execution Changes Can Destabilize RL

  • The motivating comparison uses BF16 DeepSeek-R1-Distill-Qwen-1.5B and tests compiled–compiled, compiled–eager, eager–compiled, and eager–eager execution.The first mode in each pair denotes inference; the conditions vary whether inference and training use the same execution mode.Source: §4 · p. 4
  • Both aligned execution-mode conditions maintain high accuracy, while both crossed conditions initially improve and then lose accuracy.Figure 1 reports training accuracy, AIME 2025 accuracy, and inference–training discrepancy bDKL for the four conditions.Source: Figure 1 · p. 4, §4 · p. 4
  • With eager training held fixed, switching inference from compiled to eager restores stability without increasing training-side precision.The four conditions support execution-mode agreement as a stabilizing factor and motivate inspecting compilation-induced numerical changes.Source: §4 · p. 4

5 Qwen Analysis: Where Fusion Moves Rounding Locations

Graph inspection identifies three Qwen regions where fusion changes when intermediates are rounded: the RoPE path, SwiGLU, and residual RMSNorm. Fusion retains selected intermediates in FP32 but stores region outputs in BF16, unlike separate operations that round intermediates earlier.

5 Qwen Analysis: Where Fusion Moves Rounding Locations

  • Graph inspection identifies changed rounding locations in the RoPE path, SwiGLU, and residual RMSNorm.Fused execution retains marked intermediates in FP32, while separate operations round them to BF16 before the next operation; outputs are BF16 in both cases.Source: §5 · p. 5, Figure 2 · p. 5

5.1 The three affected regions

  • In the RoPE path, fusion retains Q/K-normalization intermediates when present and RoPE products and sums in FP32; only rotated Q/K outputs are stored in BF16.Variants without Q/K normalization still fuse RoPE products and sums.Source: §5.1 · p. 5, Figure 2 · p. 5
  • SwiGLU fusion keeps sigmoid and SiLU results in FP32 through gate multiplication, whereas eager execution stores the SiLU output in BF16 before multiplication.The fused product is stored in BF16.Source: §5.1 · p. 5
  • Residual RMSNorm fusion feeds the unrounded FP32 residual sum into RMSNorm and retains its intermediate computations in FP32; eager execution rounds the sum first.The RMSNorm output and a separate copy of the sum for the continuing residual stream are stored in BF16; this region follows both attention and MLP in each decoder block.Source: §5.1 · p. 5
  • The training intervention retains selected intermediates in FP32 and casts region outputs to BF16, preserving the inference engine and compute-intensive matrix-multiplication and attention kernels.Aligning all three regions adds 5.33–10.92% to forward-and-backward latency and 0.41–0.89 GiB of peak allocated memory per GPU across three benchmarked models.Source: §5.1 · p. 5

6 Static Attribution of Policy Discrepancy

The analysis measures how FP32 recomputation in RoPE, SiLU/gate multiplication, and residual addition/RMSNorm changes fixed-weight policy discrepancy. All-three-region FP32 yields the smallest estimated KL for each model, with the RoPE path accounting for the largest share of the joint reduction.

6 Static Attribution of Policy Discrepancy

  • All-three-region FP32 reduces estimated per-token bDKL by 25.5%, 9.1%, and 12.5% versus all-BF16 in the three models, respectively.It yields the smallest estimate for every model.Source: §6 · p. 6, Table 1 · p. 6
  • The RoPE path accounts for 81.5%, 91.8%, and 76.9% of the joint KL reduction across the three models.Shapley attribution averages each intervention’s marginal contribution over all application orders.Source: §6 · p. 6
  • Negative SwiGLU shares for Qwen3-1.7B and Qwen3-8B-Base mean that adding this intervention increases discrepancy on average across intervention orders.The nonmonotonic results show that higher intermediate precision does not automatically reduce mismatch.Source: §6 · p. 6

7 Rounding-Location Alignment Stabilizes Training

Across four Qwen-family models, the default BF16 baseline collapses, while the two targeted alignment conditions and the zero-mismatch backend remain stable. RoPE-path alignment alone suffices in every observed run despite small, nonzero inference–training discrepancy.

7 Rounding-Location Alignment Stabilizes Training

  • The BF16 baseline loses accuracy for every model, while FP32 RoPE-path alignment, FP32 at all three regions, and the zero-mismatch backend avoid collapse.Both targeted conditions retain higher AIME 2025 accuracy; they retain small, nonzero bDKL, while the zero-mismatch backend remains at zero.Source: Figure 3 · p. 7, §7 · p. 7
  • RoPE-path alignment alone suffices in all observed runs, showing that eliminating all numerical disagreement is unnecessary for stability in these runs.The zero-mismatch backend uses the same Hugging Face implementation for inference and training with batch-invariant kernels.Source: §7 · p. 7
  • RoPE-path alignment adds 0.90–2.04% to training-computation latency, with unchanged peak allocated memory in the three-model benchmark.This overhead is lower than aligning all three regions.Source: §7 · p. 7
  • The phase-concentration figure reports layer medians, quartile boxes, and full-range whiskers, with changed-rounding intermediates markedly more nonuniform than other signals.The caption identifies the plotted distributions as 16-frequency phase concentrations.Source: Figure 4 · p. 8

8 Structure of BF16 Rounding Phases

The section characterizes how values fall within BF16 rounding cells using a phase-concentration statistic and compares affected intermediates with other model signals. The changed-rounding groups have the largest median concentrations, while the statistic alone does not determine rounding errors or downstream effects.

8 Structure of BF16 Rounding Phases

  • For each signal and layer, phase concentration is the root mean square of the first 16 Fourier magnitudes of its phase distribution; larger values indicate stronger departures from uniformity over those frequencies.The reported summary is the median across layers.Source: §8 · p. 8
  • Across models, median phase concentration is 1.5 × 10−5–8.0 × 10−4 for projection, attention, and RMSNorm outputs, versus 0.038–0.040 for SiLU, 0.054–0.192 for RoPE products, and 0.443–0.472 for residual sums.The three groups whose rounding locations change under fusion have the largest median concentrations; some Q/K/V and attention layers have large maxima but low medians and quartiles, indicating isolated cases.Source: §8 · p. 8
  • Phase measurements describe values within BF16 intervals, but rounding errors and downstream effects also depend on interval widths and subsequent computations.The analysis complements fixed-weight discrepancy and RL-stability measurements by examining affected intermediates directly.Source: §8 · p. 8

9 Discussion

Selective rounding-location alignment can stabilize BF16 RL even when static policy discrepancy remains, while compilation and broader architectural coverage impose practical and evidential boundaries. The appendices describe the experimental setup and analyses of intervention attribution and BF16 rounding patterns.

9 Discussion

  • A 9.1–25.5% reduction in static per-token KL accompanies prevention of collapse, and targeted alignment supports stability without full-stack agreement.The authors suggest that errors at different locations may affect updates differently, so aligning a sensitive region can help while average discrepancy remains.Source: §9 · pp. 8–9
  • Compilation can add initial cost, recompilation for variable response lengths, graph breaks, and integration work; compiling training also does not eliminate engine differences.The discussion also notes that stability can return without increasing training-side precision when inference switches from compiled to eager execution.Source: §9 · p. 9
  • Experiments cover two math datasets and one model family, so broader experiments are needed to establish when alignment suffices across architectures and training settings.Adapting the intervention to other architectures requires inspecting their inference and training graphs to identify relevant rounding locations.Source: §9 · p. 9
  • The takeaway is to inspect both execution graphs and align rounding locations moved by fusion; RoPE-only alignment suffices in the studied runs, without full-stack alignment.Aligning all three regions stabilizes training and gives the lowest per-token KL among the Section 6 conditions, while leaving the inference engine, matrix-multiplication and attention kernels, and objective unchanged.Source: §9 · p. 9

A Experimental Details

  • Main-text experiments use Sanity-Test-R1D-1.5B training data and evaluate on AIME 2025; Section 6 uses the first 512 prompts and Section 8 the first 256, generating one response per prompt in both.Inference uses vLLM with FlashAttention, while Hugging Face training uses scaled dot-product attention without compilation. Section 6 holds responses and inference log probabilities fixed across all eight Hugging Face recomputations.Source: A Experimental Details · p. 14, §A.1 · p. 14

A.2 RL training and evaluation

  • RL metrics are logged each step and AIME accuracy every 10 steps, averaged across GPU-wise evaluations on the full test set.The supplied configuration describes the training and evaluation setup for Sections 4 and 7.Source: §A.2 · p. 14

A.3 Static attribution

  • The Section 6 Shapley contribution averages each region’s marginal discrepancy reduction over intervention orders; normalized shares sum to one, with negative shares indicating average increases.The analysis uses three interventions and defines discrepancy as bDKL with a subset enabled.Source: §A.3 · p. 14

A.4 BF16 cell phases and Fourier diagnostics

  • Figure 5 shows recurring frequency patterns across models: RoPE products peak at k = 4, 8, and 16, while residual sums are largest at k = 16.Plotted magnitudes are medians across layers; residual-sum magnitudes are averaged within each layer first, and panels use different vertical scales.Source: §A.4 · p. 15, Figure 5 · p. 16

B Additional Experiment on DAPO-Math-17k

On DAPO-Math-17k, targeted rounding alignments and the Zero-mismatch backend sustain training progress, while the BF16 baseline loses accuracy after an initial improvement.

B Additional Experiment on DAPO-Math-17k

  • Both targeted alignments and the Zero-mismatch backend sustain training progress on DAPO-Math-17k, while the BF16 baseline loses accuracy after initially improving.The experiment repeats four conditions with DeepSeek-R1-Distill-Qwen-1.5B and evaluates on AIME 2025.Source: B Additional Experiment on DAPO-Math-17k · p. 16
  • The BF16 baseline’s bDKL rises sharply, while targeted alignments retain small, nonzero bDKL and the Zero-mismatch backend remains at zero.Figure 6 reports training accuracy, AIME 2025 accuracy, and bDKL.Source: B Additional Experiment on DAPO-Math-17k · p. 16
  • RoPE-path rounding-location alignment also stabilizes training on this additional dataset.The PPO-clipped GRPO objective, sampling configuration, optimization hyperparameters, and execution settings match those used for this model in Appendix A and Table 3.Source: B Additional Experiment on DAPO-Math-17k · p. 16

C Training Overhead of Rounding-Location Alignment

The overhead benchmark measures the training-computation cost of rounding-location alignment, while the Zero-mismatch backend uses shared execution arithmetic to preserve rounding behavior.

C Training Overhead of Rounding-Location Alignment

  • The benchmark uses eight GPUs with FSDP, times forward and backward computation including communication, and excludes rollout, optimizer updates, data loading, and gradient clearing.Timing uses fixed pretrained weights, ten warmup steps, and twenty measured steps; step latency is the maximum wall-clock duration across ranks.Source: C Training Overhead of… · pp. 16–17
  • Figure 6 compares training accuracy, AIME 2025 accuracy, and bDKL; the BF16 baseline loses accuracy after initial improvement, while both alignments and Zero-mismatch sustain progress.Observations use causal trailing means over 20 observations for training accuracy and bDKL, and five for AIME 2025 accuracy.Source: Figure 6 · p. 17
  • RoPE-path alignment adds 0.90–2.04% latency with unchanged peak allocated memory; three-region alignment adds 5.33–10.92% latency and 0.41–0.89 GiB per GPU.Table 4 reports training-computation overhead, with latency changes relative to each model’s BF16 baseline.Source: C Training Overhead of… · p. 17, Table 4 · p. 17

D Implementation of the Zero-Mismatch Backend

  • The Zero-mismatch backend shares inference and training implementations and uses batch-invariant Triton kernels for matrix multiplication, reductions, log-softmax, and attention.Its attention computation shares key/value traversal order, accumulates in FP32, and returns BF16 outputs.Source: D Implementation of the Zero-Mismatch… · p. 17

Related papers

Loading 2610.32951v1…