Looped transformers improve reinforcement learning speed and accuracy
ScaleRLT: Exact Log-Probabilities and Fast Deferred Rollouts for Reinforcement Learning on Looped Transformers
Machine LearningDistributed, Parallel, and Cluster Computing
Summary
Working with very large language models can slow down training, especially when trying to teach them by reinforcement learning. The authors studied a special kind of model called a looped transformer that reuses the same layers multiple times for each word. They found a way to compute exact scores for the model’s guesses without slowing things down. This lets training happen faster and more accurately on GPUs and multiple machines. It also helped models get better at solving math problems faster.
What this means in practice
- •For machine learning engineers: Speed up reinforcement learning training on looped transformer models by using deferred rollout computations with exact log probabilities.
- •For high performance computing teams: Reduce GPU update times via efficient layer-to-GPU mappings and CUDA graph optimizations for multi-node looped transformer training.
Authors
Dezhen Lu, Zixiang Ni, Runguo Li
Abstract
Looped language models such as Ouro reuse one block of layers several times per token. A rollout engine can batch the deeper loops of several response positions and run them only when a deep readout is due. This deferred schedule is fast, but the learner of an on-policy RL algorithm recomputes the sampled tokens' log-probabilities with a dense pass, and the two computations differ in floating point: in FP32 at 4096 tokens about 18% of the sampled tokens miss a $10^{-5}$ same-weight gate, and with bf16 weights nearly every token does. We use a hybrid learner log-probability, $\ell=\ell_{\mathrm{dense}}+\mathrm{sg}(\ell_{\mathrm{sched}}-\ell_{\mathrm{dense}})$, where $\mathrm{sg}$ denotes stop-gradient. Its value is a no-grad replay through the rollout's own schedule, and its gradient is the dense learner's. The replay reproduces the rollout bit for bit in all 167 FP32 and 28 bf16 real-weight cells we ran on three platforms, including rollouts whose batch shrinks at stop tokens, and the gradient has the same FP32 error as other dense implementations in all 102 pre-registered comparisons from 96 to 4096 tokens; the full GRPO loss of a training step matches a float64 reference to $3.1\times10^{-4}$. On one GPU, deferral raises rollout throughput by $1.47$ to $1.56\times$, and by $1.18$ to $1.54\times$ once every step runs from CUDA graphs, bit for bit. In a two-node GRPO loop, one-wave KV admission and CUDA graphs shorten an update by $2.71\times$ with the gate bitwise at every update, and four nodes keep the update time within 2% at the same per-GPU load. In 100-update runs held-out GSM8K accuracy rose from 23% to 82% with the mixed-depth readout that deferral relies on and from 34% to 90% with the native readout, which reached a higher accuracy in less time. For looped models beyond one node we analyse a ring mapping of layers to GPUs and measure its communication paths.