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

1 INTRODUCTION

• Deferred mixed-depth readout batches loops 3–4 for shallow tokens, saving rollout time while changing multiplication and attention-reduction grouping and potentially costing quality.The deeper loops are delayed, not skipped; later deep readouts still require them.Source: §1 · pp. 1–2
• The hybrid learner takes its log-probability value from a no-grad replay through the rollout schedule and its gradient from the dense learner.The replay must match the rollout bit for bit; the reported gradient accuracy matches other dense implementations.Source: §1 · p. 1, Figure 1 · p. 2
• One-wave admission and per-step CUDA graphs help make the gated training loop fast across multiple nodes.The reported two-node update was shortened by 2.71×, with the gate bitwise at every update.Source: §1 · p. 1, Figure 1 · p. 2
• A ring mapping is proposed for models beyond one node, placing loop-shared weights and per-loop KV on GPUs while measuring communication paths.The paper describes the mapping and its communication paths.Source: §1 · p. 2

2 SETTING

• Ouro uses four loops with separate KV histories; the first response token follows four-loop prompt prefill, while later inputs usually read out after two loops and every fourth after four.The native decoding policy reads out after all four loops for every token.Source: §2 · p. 2
• Deferring loops 3–4 changes work grouping, not the function, and saves time by batching rather than skipping computation.Almost every token eventually runs all four loops because deep-readout ancestors need them.Source: §2 · p. 2
• The production adapter runs admission waves to completion under a KV budget and request cap; requests do not join an active cohort and leave at a stop token or token limit.Wave admission is computed from prompt lengths alone.Source: §2 · p. 2
• The learner replays each sequence’s per-loop KV history, recomputes layers during backpropagation, accumulates FP32 gradients, and takes one optimizer step per rollout batch.Its fused log-softmax uses a hand-written backward; no Transformer Engine kernel is involved.Source: §2 · p. 2

3 DENSE REPLAY DOES NOT REPRODUCE THE ROLLOUT

• At 96 tokens, about 1% of FP32 sampled tokens exceed 10−5, while the training loop’s learner exceeds it in 33 of 36 cells.The training learner differed from the rollout by up to 3.5 × 10−4 in long cells.Source: §3 · p. 2
• In FP32, the largest ratio deviation stays below 10−3, inside the PPO clip range, even though the strict same-weight gate fails.With bf16 weights, 89–98% of sampled tokens miss the gate and the largest difference ranges from 0.45 at 96 tokens to 1.6 at 4096.Source: §3 · p. 2
• At 4096 tokens, about 18% of FP32 sampled tokens exceed the 10−5 gate; with bf16 weights, 89–98% exceed it.The FP32 result holds on both platforms, even at batch size one.Source: §3 · p. 2, Table 1 · p. 3

4 A HYBRID LOG-PROBABILITY TAKES ITS VALUE FROM THE ROLLOUT’S SCHEDULE

• A no-grad, teacher-forced replay through the rollout’s own schedule supplies the sampled token’s log-probability value, while the dense learner supplies its gradient.The replay can use the rollout’s attention kernel because it needs no gradient.Source: §4 · p. 3
• When scheduled replay is bitwise equal to the rollout, the importance ratio at θold is exactly one and the strict gate can be kept.The correction is exact only while replay uses the rollout’s weights; multiple optimizer steps per batch would require recomputing it at each step.Source: §4 · p. 3
• The evaluation distinguishes dense per-sequence and packed passes, scheduled replay on either backend or alone, differentiable scheduled execution, and the production learner.Source: §4 · p. 3

5 SCHEDULED REPLAY IS BITWISE EXACT

• A2 replay was bitwise in all 167 FP32 real-weight cells and all 28 bf16 cells, including production cohorts whose batches shrank at stop tokens.The bf16 cells spanned lengths of 96 to 4096 tokens and batch sizes of 1 to 8.Source: §5 · p. 3
• A2x was never bitwise, and A2-solo was never bitwise for batch sizes of at least 2.Triton and PyTorch rollouts sampled identical tokens in every compared pair, but replaying with a different attention backend did not preserve bitwise equality.Source: §5 · p. 3
• Exact replay requires the rollout’s batch composition, step schedule, and kernel together.In production, the adapter records each admission wave’s partition and backend for replay.Source: §5 · p. 3

6 THE HYBRID GRADIENT IS AS ACCURATE AS ANY DENSE GRADIENT

• The gradient criterion compares learner error with twice the error of other legitimate dense implementations against a float64 reference.The criterion was met in all 102 comparisons.Source: §6 · p. 3
• The learner met the pre-registered gradient-error criterion in all 102 comparisons from 96 to 4096 tokens.Its relative error was about 10^-5 at 96 tokens and about 10^-4 at 1024 tokens, remaining there at 4096 tokens.Source: §6 · p. 3
• The gradient half of Eq. 1 adds no more error than the FP32 noise already present in dense learners.Neither alternative dense implementation was consistently more accurate; one ill-conditioned K4 cell was an outlier where every implementation performed poorly.Source: §6 · pp. 3–4
• The full training-step gradient’s relative error was 3.1 × 10^-4 with the native readout and 5.7 × 10^-6 with the mixed readout.The FP32 gradient was compared with an independently written float64 reference.Source: §6 · p. 4

7.1 One GPU

• Deferral raised throughput 1.47–1.56× over eager execution, and remained 1.18–1.54× faster when both schedules used CUDA graphs.The deferral comparison covered batch sizes 1–64 and lengths 512–4096 tokens; the graph comparison found the largest gains at small batch.Source: §7.1 · p. 4, Table 4 · p. 6
• Moving sampling to one device synchronization per step and replaying passes from CUDA graphs made rollouts 1.59–2.49× faster than before, bitwise equal to the per-row loop.The scheduled replay from graphs was also bitwise across graph and eager rollout/replay pairings.Source: §7.1 · p. 4

7.2 Two Nodes, Eight GPUs

• One-wave admission plus CUDA graphs reduced the two-node update from 195.6 s to 72.2 s (2.71×), with a bitwise gate at every rank and update.Releasing cached allocator memory lowered peak use to 63,016 MiB at the same speed.Source: §7.2 · p. 5, Table 5 · p. 6

7.4 Learning

• Keeping learner activations cut single-GPU training from 116.6 to 62.9 s per update; batching sequences cut it further to 35.4 s with two sequences per micro-batch.The rollout remained bitwise identical in both checks; the batched gradient was in the same accuracy class as the per-sequence gradient, not equivalent.Source: §7.4 · p. 6

7.3 Four Nodes, Sixteen GPUs

• At four nodes, weak scaling stayed within 2% of the eight-GPU update time, while strong scaling reached 0.76× rather than its predicted 0.65×.With 16 samples per rank, training halved, but rollout and gate times did not move.Source: §7.3 · p. 5, Table 6 · p. 7

7.4 Learning

• K4 was 7.6 points ahead after 100 updates on two nodes, and reached 84.6% after 20 updates versus 82.4% for the mixed readout after 100.The mixed readout’s update took 69.4 s against 79.6 s for K4.Source: §7.4 · p. 6, Table 7 · p. 7
• Held-out accuracy improved by 59.4 points with the mixed readout and 56.2 points with K4; both gates were bitwise at every rank and update.Among responses that stopped, accuracy rose from 25% to 84% for the mixed readout and from 34% to 91% for K4.Source: §7.4 · p. 6, Table 7 · p. 7
• At 512 tokens, the ungated bf16 arm reached 77.0% accuracy against 77.2% with the gate, while only 0.02% of its tokens left the PPO clip range.The gate was bitwise at all 100 updates; at this scale the hybrid learner buys certainty, not accuracy.Source: §7.4 · p. 7

8 LOOPED MODELS MAP ONTO A GPU RING

• For a 69B shape, learner state and rollout weights already require 613 GiB, before accounting for per-loop KV on one node.The rollout weights alone fit one node of four 95 GiB GPUs for the three shapes at D = 8, but the combined RL components do not.Source: §8 · p. 7
• The design places contiguous layer arcs across GPUs in a ring, keeping shared weights and per-loop KV in place while sending each loop’s hidden state around it.Only node boundaries use the network; deferral groups loops 3–4 into fewer, larger messages but does not reduce bytes crossing the ring.Source: §8 · p. 7
• FSDP gathers each layer’s weights 2K times per micro-batch unless it retains gathered weights, while the ring moves no weights; FSDP carries 324–1032× more bytes per GPU.The comparison is for an 8192-token micro-batch.Source: §8 · pp. 7–8
• A NIC carried 23.8 GB/s, while tuned NVLink communication reached 124 GB/s and CUDA IPC reached 132 GB/s; the measured-path design remains unimplemented.The two-node 69.3B forward-activation estimate was 17.8 ms with tuned NVLink channels and 9.4 ms with a four-NIC split, versus 26.9 ms with NCCL defaults; pipeline bubble, arc balance, and small-message host overhead remain open costs.Source: §8 · pp. 7–8, Table 9 · p. 8

9 RELATED WORK

• Looped and recurrent-depth transformers reuse layers to add depth at fixed parameter count; this paper applies mixed-depth readout to Ouro.Prior examples are attributed to Dehghani et al. (2018), Geiping et al. (2025), and Zhu et al. (2025).Source: §9 · p. 8
• RL frameworks such as HybridFlow (Sheng et al., 2024) co-locate or separate rollout and training engines, while PPO and GRPO rely on learners reproducing rollout probabilities.The cited objectives are PPO (Schulman et al., 2017) and GRPO (Shao et al., 2024).Source: §9 · p. 8
• Prior mismatch responses use capped importance sampling or FP16 instead of bf16 (Qi et al., 2025); the hybrid learner instead removes mismatch from the value.The dense learner supplies the gradient; cited system baselines include pipeline and sharded data parallelism (Shoeybi et al., 2019; Narayanan et al., 2021; Rajbhandari et al., 2019; Zhao et al., 2023).Source: §9 · p. 8

10 LIMITATIONS

• Learning-quality evidence uses one seed per readout policy and setup, 100 updates, and a 500-problem evaluation with one sample per problem.Runs use 256 samples per update on two nodes or 64 on one GPU.Source: §10 · p. 8
• The gate stayed bitwise through 100 bf16 updates, but gradient and throughput results are FP32, and whether the gate matters for learning at longer sequences remains open.At 512 tokens, the one-seed bf16 ungated learner did as well as the gated learner.Source: §10 · p. 8
• The full GRPO loss was checked against float64 on one 512-token batch, while the gradient criterion applies to a surrogate loss and T3’s ratio is exactly one.With a ratio of one, the PPO clipping branch does not run.Source: §10 · p. 8
• With k = 4, mixed readout trailed native readout by 7.6–9 points after 100 updates, while updates were only 4–13% shorter.The paper tried one readout interval.Source: §10 · p. 8
• The study uses Ouro-1.4B, tests exactness up to 4096 tokens and data parallelism up to four nodes, and presents the ring mapping as a design backed by point-to-point measurements.Strong scaling is bounded by per-token rollout and gate latency.Source: §10 · p. 8

11 CONCLUSION

• Scheduled replay restores tested FP32/bf16 bitwise agreement with dense-gradient accuracy; multi-node training stays bitwise, while mixed readout learns less than native readout.The conclusion attributes the fast loop to deferral, one host synchronization per step, CUDA graphs, and one-wave admission.Source: §11 · pp. 8–9

A SINGLE-GPU ROLLOUT TABLES

• Single-GPU rollout tables report unchanged peak memory, and one-sync and graph rollouts reproduce per-row execution bitwise; two Table 10 ratios have no clean interpretation.One anomalous deferred-timing ratio is unexplained, and another compares runs on different nodes.Source: Table 11 · p. 10, Table 10 · p. 10

B THE PINNED-MEMORY DEFECT

• A loop variable retained a 109 MiB diagnostic-logits block in freed KV memory, pinning about 14 GiB per rank; freeing it restored pre-rollout reserved memory across 200 updates.The allocator issue was traced to a dense replay between rollout and gate.Source: B THE PINNED-MEMORY DEFECT · p. 10

C BATCHED LEARNER FORWARD

• Without recomputation, padded-slab activations grow about 12 GiB per sequence, limiting learner micro-batches to three beside the rollout’s 80 GB peak.The table also notes a gradient-norm rounding difference between per-sequence and batched paths.Source: Table 12 · p. 11

D MEMORY ARITHMETIC FOR THE RING MAPPING

• The ring-mapping arithmetic assumes K = 4 loops and D = 8 optimizer replicas, comparing per-GPU FSDP weight-gather bytes with ring-boundary activation bytes for 8192-token micro-batches.The table’s model arithmetic covers three looped shapes and reports forward/backward quantities in bf16.Source: Table 13 · p. 11

• Each run's predictions were recorded before submission, and every failed prediction is listed.Source: E PRE-REGISTERED PREDICTIONS AND THEIR… · p. 11
• The single-GPU CUDA-graph measurement in Section 7.1 was not pre-registered.Source: E PRE-REGISTERED PREDICTIONS AND THEIR… · p. 11
• Three submissions failed before measurement and were repeated after fixes to a missing launcher, a network-plugin runtime symbol, and inconsistent process-subgroup creation.Source: E PRE-REGISTERED PREDICTIONS AND THEIR… · p. 11
• The outcomes cover pre-registered system, gradient, communication, and learning runs.Table 14 is identified as the summary of these runs.Source: Table 14 · p. 12