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.
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, ℓ=ℓdense+sg(ℓsched−ℓdense), where 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×10−4. On one GPU, deferral raises rollout throughput by 1.47 to 1.56×, and by 1.18 to 1.54× 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× 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
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 ∣Δlogp∣
Tokens over 10−5
L40S
96
54
5.1×10−5
161/14,496 (1.1%)
L40S
1024
3
1.4×10−4
404/10,016 (4.0%)
L40S
4096
1
2.6×10−4
760/4,096 (18.6%)
GH200
96
36
2.6×10−5
65/9,664 (0.7%)
GH200
1024
4
8.9×10−5
599/17,984 (3.3%)
GH200
4096
5
8.0×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×10−5
L40S
58
28,608
96/1024/4096
1/2/4/8
58/58
2.1×10−5
2.1×10−5
GH200
45
101,056
96/1024/4096
1/2/4/8
45/45
9.5×10−5
1.0×10−4
GH200, production rollout
54
55,678
≤ 1024
1/2/4/8
54/54
4.0×10−5
–
GH200, bf16
23
19,008
96/1024/4096
1/2/4/8
23/23
1.2
4.4×10−1
GH200, bf16, production rollout
5
15,724
≤ 1024
8
5/5
4.4×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(M) median
e(M) max
Other dense, median
e(M)/maxe(Y) , max
GH200
96
72
9.9×10−6
6.7×10−5
8.7×10−6
1.62
B300
512
2
7.2×10−5
7.2×10−5
5.6×10−5
1.09
B300
1024
20
1.2×10−4
1.3×10−3
2.0×10−4
1.19
B300
4096
8
1.1×10−4
1.5×10−4
9.5×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 × 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% → 84.2%
Mixed, weak
512
20
69.2
18.5
16.2
30.3
50,000
320/320
23.0% → 76.4%
K4, strong
256
10
60.2
21.6
20.1
15.9
39,262
160/160
33.8% → 79.6%
K4, no recompute
512
5
68.3
23.7
21.0
19.1
49,518
80/80
–
Table 6: Four nodes × 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 ×
yes
64
512
888
894
1412
1.59 ×
yes
24
2048
336
396
654
1.95 ×
yes
12
4096
194
213
308
1.59 ×
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× that on 4; 1.066×
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× ; 2.2× . Four-NIC split (unpipelined) at most 0.5× the direct time; 0.54× . 4 KiB half round trip at most 20 μ s (NVLink) and 40 μ s (NIC); 101 and 89 μ 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 μ s; 61. CUDA-graph ring at most 15 μ s per hop; 25
Appendix
Table 14: Outcomes of the pre-registered system, gradient, communication and learning runs.
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.
Andrei Marchenko, Viacheslav Bezrukov, Oleg Kashurin +5
Looped language models improve reasoning and knowledge manipulation by applying shared computation repeatedly. Existing systems usually repeat an entire layer stack, although a mixer and a dense feed-forward network (FFN) perform different operations and have different costs. We ask a narrower question: what should loop? We view recurrence as repeated composition of a state update and argue that an application is valuable when it exposes a new cross-position influence direction that remains observable at the task readout. Iterative Transport Rank (ITR) describes the cumulative influence trajectory; marginal ITR describes the nonredundant influence contributed by successive applications. This view motivates MixerLoop, which repeats each Gated DeltaNet mixer while applying its dense FFN once. We compare MixerLoop with no recurrence and full-block recurrence at 15M and 110M parameters under the same data, initialization, and architecture. A finite context-off intervention tests whether later mixer applications produce distinct, non-negligible, and beneficial changes at the final language-model readout. MixerLoop surpasses FullLoop on aggregate CORE at 15M and retains 41.5% of its CORE improvement at 110M while reducing recurrent-backbone projection FLOPs by 45.9%. These results show that the benefits of recurrent depth can be retained without repeatedly executing the dense FFN.
Looped Transformers increase computation per token by repeatedly applying a shared backbone, offering a way to improve language modeling and downstream task performance at a fixed backbone parameter count. However, each loop depends on the preceding loop, so increasing recurrent depth also increases decoding latency. Existing approaches include Parallel Loop Transformers, which reduce latency by overlapping loop computation across token positions, but yield rapidly diminishing marginal accuracy gains at larger depths. To address this challenge, we introduce LoopFold, which folds computation at multiple recurrent depths into parallel batches for the same token prefix. A lightweight predictor estimates inputs to later loops so they can start before earlier loops finish. A corrector then refines the outputs of these later loops using newly available results from earlier loops. All components are trained jointly using only the final next-token prediction loss. In experiments with 1.4B models, deeper LoopFold configurations achieve progressively higher average zero-shot accuracy across eight tasks, rising from 53.41% to 56.38%. Vanilla looped models and PLT show limited gains over the same depth range. Compared with a nine-loop vanilla model, LoopFold uses five sequential stages and additional parallel computation to reduce decoding latency by 46%.
Feichen Song, Boyi Zeng, Yiqin Hao +7
Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute