One Additive Term Makes FP8 RL Training Stable

0.715 vs. 0.338. 0.723 vs. 0.095.

These are the training reward comparisons posted by Lewis Tunstall. Under FP8 rollout, Score Centering + masked importance sampling reaches a reward of 0.715, while the ordinary baseline is only 0.338. His caption is just one line: “FP8 rollouts. Stable learning.”

TRL (Hugging Face's reinforcement learning library) has added Score Centering to the async GRPO Trainer. The green line tracks the bf16 reference, while the red line climbs halfway and then turns downward.

RL training for LLMs has a long-standing problem: training-inference mismatch (TIM). The model samples data to generate, and training updates the parameters; ideally the two are consistent. In practice there are always differences—the trainer uses bf16, the sampler uses FP8; KV cache quantization; asynchronous updates mean the trainer is still using old weights. These differences bias the gradient updates, and the bias accumulates: the longer training runs, the more the model drifts.

Existing solutions in the industry fall into two categories. The first is numerical alignment, such as router replay and batch-invariant kernels. It works, but rollout performance takes a hit. The second is importance sampling, which weights updates that deviate from the policy, but after the weights are clipped, new bias is introduced.

Max Ryabinin and Martin Marek re-examined the problem from first principles. They found that the instability caused by TIM has its root in drift: the persistent bias between trainer and sampler that accumulates at every training step. Standard solutions all work around drift. They chose to cancel it directly.

The policy gradient update can be decomposed into two parts: signal and drift. In the on-policy case, the expectation of the trainer score is 0, and drift disappears. In the off-policy case, it is nonzero and can be estimated, then subtracted during the update. They call this Score Centering.

Martin Marek made an animated comparison: on the left, without correction, the optimization trajectory is pulled off course by the sampler and gets farther and farther from the loss minimum; on the right, with Score Centering added, the trajectory heads straight for the minimum.

Implementation is not complicated. During training, obtain the sampler's top-k logprobs from vLLM (the log probabilities of the k tokens with the highest probability, k=128), and model the tail with the trainer's distribution. The loss is computed per token; there are concise implementations in both JAX and PyTorch, and no model architecture changes are needed.

They validated it on models from 0.6B to 30B parameters. Used alone, Score Centering matches or exceeds importance-sampling methods such as TIS and MIS; the more severe the mismatch, the bigger the advantage. Because it is an additive correction term, it can also be combined with importance sampling, and in staleness experiments the combination outperformed the pure IS baseline.

Then came deployment. Dirhousssi Amine set score centering + masked importance sampling as the default configuration for the TRL async GRPO Trainer. Under various staleness and sampling quantization settings, stability is impressive.

There is an engineering cost: the vLLM API frontend must additionally return top-k logprobs, increasing the load. With 8 API servers, the throughput loss for top-8 log-probs can be kept to around 3%. They are also working with vLLM on flat log-probs—packing logprobs into arrays to avoid constructing Python objects token by token. The PR is already open (vllm-project/vllm#62446).

One commenter (@Saksham Jain) said that the numerics gap is the real bug, and importance sampling merely covers it up. Another commenter (@Varun Daiya) observed that centering removes the per-prompt offset that the reward model should never have learned in the first place, so the policy no longer benefits from that offset; he wonders whether length drift will also decrease accordingly.

Paper: arxiv.org/abs/2609.20807. Authors Martin Marek and Max Ryabinin; Martin led the project during a summer internship on Together's Model Shaping team. Code: github.com/martin-marek/score-centering, built on a JAX framework called postax, supporting pretraining, SFT, DPO, and RL for Llama and Qwen. Hugging Face Journal Club has a full derivation video by Kashif Rasul, spanning mathematical details through implementation: YouTube.

The result is the figure at the beginning: the green line under FP8 rollout can finally track the bf16 reference.