ScaleRLT: Exact Log-Probabilities and Fast Deferred Rollouts for Reinforcement Learning on Looped Transformers

Dezhen Lu; Zixiang Ni; Runguo Li

cs.LGOct 25, 2026

ScaleRLT: Exact Log-Probabilities and Fast Deferred Rollouts for Reinforcement Learning on Looped Transformers

Authors: Dezhen Lu, Zixiang Ni, Runguo Li

Organizations: University of Tübingen, Tübingen, Germany. · Xi’an Jiaotong University, Xi’an, China. · University of Illinois Urbana-Champaign, Urbana, IL, USA.

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.

Figures & tables

Panel of Figure 1 — Figure 1: ScaleRLT at a glance (schematic). (a) The deferred rollout runs loops 1–2 of every response input at once and batches loops 3–4 until a deep readout (every $k$ -th input) needs them. (b) The learner’s log-probability takes its value from a no-grad replay through the rollout’s own schedule, which the gate requires to equal the rollout bit for bit, and its gradient from the dense learner ( Eq. 1 ). (c) One GRPO update; one optimizer step per rollout batch keeps the replay on the rollout’s weights.

Platform	Length	Cells	Max $\|\Delta\log p\|$	Tokens over $10^{-5}$
L40S	96	54	$5.1\times 10^{-5}$	161/14,496 (1.1%)
L40S	1024	3	$1.4\times 10^{-4}$	404/10,016 (4.0%)
L40S	4096	1	$2.6\times 10^{-4}$	760/4,096 (18.6%)
GH200	96	36	$2.6\times 10^{-5}$	65/9,664 (0.7%)
GH200	1024	4	$8.9\times 10^{-5}$	599/17,984 (3.3%)
GH200	4096	5	$8.0\times 10^{-4}$	13,545/73,408 (18.5%)

Table 1: Dense per-sequence replay (A0) against the rollout’s log-probabilities, FP32 unless marked bf16. “Tokens over

10^{-5}

” counts sampled tokens whose absolute log-probability difference exceeds the gate.

Platform	Cells	Tokens	Lengths	$B$	A2 bitwise	A2x max	A2-solo max
L40S (CPU arms)	10	2,928	96/512	1/2/4/8	10/10	–	$2.2\times 10^{-5}$
L40S	58	28,608	96/1024/4096	1/2/4/8	58/58	$2.1\times 10^{-5}$	$2.1\times 10^{-5}$
GH200	45	101,056	96/1024/4096	1/2/4/8	45/45	$9.5\times 10^{-5}$	$1.0\times 10^{-4}$
GH200, production rollout	54	55,678	$\leq$ 1024	1/2/4/8	54/54	$4.0\times 10^{-5}$	–
GH200, bf16	23	19,008	96/1024/4096	1/2/4/8	23/23	$1.2$	$4.4\times 10^{-1}$
GH200, bf16, production rollout	5	15,724	$\leq$ 1024	8	5/5	$4.4\times 10^{-1}$	–

Table 2: Scheduled replay (A2) against the rollout, FP32 above the line and bf16 below it. “A2 bitwise” counts cells in which every sampled token’s log-probability is reproduced exactly. A2x and A2-solo give the largest absolute difference over all cells. The CPU arms ran PyTorch attention only, so they have no A2x. Lengths are token limits; production-rollout responses end earlier at stop tokens.

Platform	Length	Comparisons	$e(\mathrm{M})$ median	$e(\mathrm{M})$ max	Other dense, median	$e(\mathrm{M})/\max e(Y)$ , max
GH200	96	72	$9.9\times 10^{-6}$	$6.7\times 10^{-5}$	$8.7\times 10^{-6}$	1.62
B300	512	2	$7.2\times 10^{-5}$	$7.2\times 10^{-5}$	$5.6\times 10^{-5}$	1.09
B300	1024	20	$1.2\times 10^{-4}$	$1.3\times 10^{-3}$	$2.0\times 10^{-4}$	1.19
B300	4096	8	$1.1\times 10^{-4}$	$1.5\times 10^{-4}$	$9.5\times 10^{-5}$	0.96

Table 3: Gradient error of the production learner (M) against float64 truth. “Other dense, median” pools the alternative dense implementations of the same cells. The last column is the largest ratio of the learner’s error to the worse alternative in the same cell; the pre-registered criterion requires it to be at most 2.

Config	Deferred	Eager	K4	Def./eager	W/o capture	Capture
$B=8$ , 512	316	205	213	1.54	1.56	2%
$B=64$ , 512	1211	1025	976	1.18	1.19	1%
$B=24$ , 2048	591	452	442	1.31	1.31	1%
$B=32$ , 512, stop tokens	485	412	416	1.18	1.37	27%

Table 4: One GH200, FP32, CUDA graphs (tokens per second). “Eager” reads out at the same depths as “Deferred”; K4 reads every token at depth 4. “Def./eager” is the deferred/eager ratio, “W/o capture” the same ratio without graph-capture time, and “Capture” that time’s share of the deferred rollout. One timing per cell.

Configuration	Waves	Rollout s	Gate s	Train s	Step s	Rollout tok/s	Gate bitwise	Peak MiB
Cap 16 (2 waves)	2	80.1	77.6	32.5	195.6	1,336	24/24	67,062
One wave	1	36.5	34.7	29.9	106.2	2,929	24/24	97,208
One wave + graphs	1	18.7	15.9	30.6	70.5	5,728	24/24	97,034
One wave, release	1	37.0	34.2	30.3	106.9	2,892	24/24	63,016
One wave + graphs, release	1	19.2	16.1	31.5	72.2	5,578	24/24	63,246

Table 5: Two nodes

\times

four GH200, data-parallel GRPO, 32 samples per rank, 512-token limit, FP32. Seconds are the median of updates 2–3 at the slowest rank. “Gate bitwise” counts rank-updates whose scheduled replay reproduced the rollout exactly. Peak memory is the largest per-GPU value reported by nvidia-smi (out of 97,871 MiB) before any checkpoint save. “Release” returns cached allocator memory to the device before each rollout and before training.

Arm	Samples	Updates	Update (s)	Rollout	Gate	Train	Peak (MiB)	Gate bitwise	Accuracy
K4, weak	512	20	81.2	23.4	20.9	32.0	49,582	320/320	33.8% $\to$ 84.2%
Mixed, weak	512	20	69.2	18.5	16.2	30.3	50,000	320/320	23.0% $\to$ 76.4%
K4, strong	256	10	60.2	21.6	20.1	15.9	39,262	160/160	33.8% $\to$ 79.6%
K4, no recompute	512	5	68.3	23.7	21.0	19.1	49,518	80/80	–

Table 6: Four nodes

\times

four GH200, data-parallel GRPO, FP32, one job. Seconds are the median of updates 2 onward at the slowest rank; “weak” keeps 32 samples per rank (512 per update), “strong” keeps the two-node total of 256 (16 per rank). Accuracy is held-out GSM8K before training and after the last update.

Updates	8 GPUs, 256/update		1 GPU, 64/update
	Mixed	K4	Mixed	K4
0	23.0%	33.8%	23.0%	33.8%
20	75.8%	84.6%	69.2%	82.6%
40	78.8%	87.8%	74.4%	84.4%
60	81.6%	87.6%	80.8%	86.6%
80	82.0%	89.0%	81.0%	87.6%

Table 7: Held-out accuracy on 500 GSM8K test problems (one sample each, same seeds at every evaluation) during GRPO with the mixed readout and K4, FP32. The two-node runs are pre-registered; the one-GPU runs use a quarter of the samples per update. One training seed per run.

Updates	bf16, T3	bf16, dense ungated	FP32, T3
0	23.6%	23.6%	23.0%
20	50.8%	50.0%	69.2%
40	69.2%	69.0%	74.4%
60	72.4%	70.0%	80.8%
80	74.0%	72.6%	81.0%
100	77.2%	77.0%	78.8%

Table 8: Held-out accuracy during GRPO with the mixed readout on one GH200, bf16 learner and rollout engine, 16 prompts per update: the hybrid learner with its gate (T3) and the dense learner with no gate. The FP32 column is the reduced run of Table 7 . One training seed per run.

Path	4 MiB	64 MiB	256 MiB
NIC, one pair	21.8	23.8	23.8
NIC, one pair to a different GPU index	21.5	23.6	23.7
NVLink, 8-rank communicator, defaults	45.6	51.4	51.4
NVLink, node-local communicator	73.3	121.4	124.3
NVLink, 8-rank, 8 channels per net peer	68.7	122.1	124.3
NVLink, copy engine via CUDA IPC	90.7	129.3	132.0

Table 9: Point-to-point bandwidth between two nodes of four GH200 GPUs, unidirectional GB/s, NCCL send/recv unless noted, median of three repeats. “8 channels per net peer” sets NCCL’s per-network-peer channel count to 8. The boundary split uses the best of 1, 4, 8 or 16 pipelined chunks at each size.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Length	$B$	Eager tok/s	Deferred tok/s	Ratio	Peak GiB	Note
512	1	10.8	16.5	1.53	6.2
512	8	85.2	128.7	1.51	12.0
512	32	269.1	402.8	1.50	31.9
512	64	479.7	705.2	1.47	58.4
2048	8	83.0	127.3	1.53	30.0
2048	16	139.8	236.1	–	54.6	different nodes

Appendix

Table 10: Single GH200, FP32, stop tokens disabled: eager and deferred schedules with identical readouts (tokens per second, median of two timings). Peak memory is the same for both. Two cells give no clean ratio: eager and deferred ran on different nodes in one, and in the other the two deferred timings differ by 10% and the ratio is far below its neighbours for a reason we did not find.

$B$	Length	Per-row sync	One sync/step	+ CUDA graphs	Speedup	Bitwise
8	512	138	165	342	2.49 $\times$	yes
64	512	888	894	1412	1.59 $\times$	yes
24	2048	336	396	654	1.95 $\times$	yes
12	4096	194	213	308	1.59 $\times$	yes

Appendix

Table 11: Single GH200, FP32, deferred schedule, stop tokens disabled (tokens per second). “Per-row sync” is the engine before the change. “Bitwise” means the one-sync and the graph rollouts both reproduce the per-row rollout exactly. Peak memory is unchanged.

$N$	Path	Recompute	Train (s)	per sample (s)	GPU busy	Train peak (GiB)	Gradient norm
1	per sequence	off	62.8	0.98	84%	44.8	3.44061
1	batched	off	54.1	0.85	73%	40.4	3.44061
2	batched	off	35.4	0.55	92%	51.2	3.43930
3	batched	off	36.2	0.57	96%	64.4	3.43887
4	batched	off	out of memory
8	batched	off	out of memory

Appendix

Table 12: Batched learner forward on one GH200 (K4, FP32, 64 samples per update, 3 updates per arm; mean of updates 2–3):

N

sequences per learner micro-batch, the per-sequence or batched path, and activation recomputation off or on. Train peak is the allocator’s peak during the training phase; the first update’s gradient norm shows the rounding difference between the paths. The rollout was bitwise identical in every arm. Without recomputation the activations of a padded slab grow by about 12 GiB per sequence, so three sequences are the most that fit next to the rollout’s 80 GB peak.

Shape	Params	States/replica	Rollout weights	FSDP regather	Ring boundary	Ratio
$h=4096$ , 48 layers	10.9B	76 GiB	20 GiB	174 GB	0.54 GB	324
$h=6144$ , 60 layers	29.7B	207 GiB	55 GiB	475 GB	0.81 GB	590
$h=8192$ , 80 layers	69.3B	484 GiB	129 GiB	1108 GB	1.07 GB	1032

Appendix

Table 13: Three looped shapes with

K=4

loops, GQA with 8 KV heads, SwiGLU width

3.5h

. Learner state per replica assumes a distributed optimizer over

D=8

replicas. The last three columns are per 8192-token micro-batch, forward and backward, in bf16: the weight bytes FSDP gathers per GPU, the activation bytes one ring boundary carries (all

K

loops), and their ratio. Model arithmetic, not measurements.

Run	Predictions	Met	Not met (predicted; measured)
Two-node calibration, 16 samples/rank	6	5	Update time on 8 GPUs at most $1.05\times$ that on 4; $1.066\times$
Two-node admission and CUDA graphs	9	8	Peak memory at most 90 GB; 97,208 MiB
Two-node memory rerun	8	7	Reserved memory before training within 2 GiB of that before the rollout; 13.8–14.2 GiB once per rank
Allocator snapshots, one GPU	4	4	–
Point-to-point, first round	10	6	NCCL NVLink at least 90 GB/s; 52. NVLink over NIC $4$ – $10\times$ ; $2.2\times$ . Four-NIC split (unpipelined) at most $0.5\times$ the direct time; $0.54\times$ . 4 KiB half round trip at most 20 $\mu$ s (NVLink) and 40 $\mu$ s (NIC); 101 and 89 $\mu$ s
Point-to-point, second round	8	5	Single-process copy-engine slowness is host-side; it is not. Host time per call 10–60 $\mu$ s; 61. CUDA-graph ring at most 15 $\mu$ s per hop; 25

Appendix

Table 14: Outcomes of the pre-registered system, gradient, communication and learning runs.

RL for Language Models

Reinforcement learning for language models uses feedback-driven policy updates to improve reasoning, alignment, and interactive behavior.

See all 438 papers

Closing the Loop: Practical Training Recipes for Looped Language Models
Sep 30, 2026Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin +5Efficient Language Model Training LLM Training
RND NLP, DAIMLD, Russian Federation
Allocating Recurrent Compute in Looped Language Models
Aug 18, 2026Ruhai Lin, Yiyang Guo, Rui-Jie Zhu +2Recurrent Language Models Recurrent Transformers
University of California, Santa Cruz
LoopFold: Parallelizing Recurrent Depth in Looped Transformers
Oct 26, 2026Feichen Song, Boyi Zeng, Yiqin Hao +7Recurrent Transformers Efficient Transformer Inference
Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute

Sep 30, 2026J/K move · Enter open · S save

Closing the Loop: Practical Training Recipes for Looped Language Models

Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin, Inessa Fedorova+4

RND NLP, DAIMLD, Russian Federation

Efficient Language Model Training LLM Training Recursive Language Models Recurrent Language Models

Open paper AskPDF arXiv

Looped language models increase effective depth by repeatedly applying a shared block of layers, but existing large-scale recipes require multi-stage training over trillions of tokens, while the benefits of recurrence remain difficult to separate from differences in data and training. In this work, we establish practical training recipes for looped language models, with three main results. (1) We develop a compute-efficient from-scratch pipeline that reduces the training budget from 7.7T tokens in Ouro to 310B tokens while retaining strong reasoning performance. Pretraining followed by high-quality mid-training, together with learning-rate warmup and stronger exit-gate regularization, enables stable recurrent training without prior multi-stage schedules. (2) Under controlled comparisons, our 1.4B LoopLM outperforms a parameter-matched dense model trained on the same data and token budget on all 12 evaluated benchmarks, including +14 points on GSM8K, +10 on MATH, and +22 on DROP. At matched inference compute, it approaches a 3.9B dense model on mathematical reasoning and reading comprehension while using only 36% as many parameters. (3) We introduce a minimal recipe for converting pretrained dense models into looped ones: a single learned input-mixing scalar and a smoothed exit loss, with no step-specific parameters. Applied to Qwen3-1.7B-Base, Looped Qwen improves over an identically continued dense baseline on every evaluated benchmark across two data regimes, with statistically clear gains on GSM8K, MATH, and MMLU-Pro on the curated mixture. Together, these results make looped language models substantially cheaper to train from scratch and practical to introduce into existing pretrained checkpoints, while isolating the gains due to recurrence itself.

Abstract

RL for Language Models

Explore similar work

Closing the Loop: Practical Training Recipes for Looped Language Models

Allocating Recurrent Compute in Looped Language Models

LoopFold: Parallelizing Recurrent Depth in Looped Transformers