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

S CALE RLT: E XACT L OG -P ROBABILITIES AND FAST D EFERRED ROLLOUTS FOR R EINFORCEMENT L EARNING ON L OOPED T RANSFORMERS

arXiv:2610.34427v1 [cs.LG] 25 Oct 2026

Dezhen Lu 1 Zixiang Ni 2 Runguo Li 3

A BSTRACT 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 ), whose value is a no-grad replay through the rollout’s own schedule and whose 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–1.56×, and by 1.18–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. 1

I NTRODUCTION

Looped (recurrent-depth) transformers apply the same block of layers several times per token, trading parameters for depth (Dehghani et al., 2018; Geiping et al., 2025; Zhu et al., 2025). Reinforcement learning with verifiable rewards is a natural way to train their reasoning, and its cost is dominated by rollouts. A looped model makes rollouts unusually expensive: every loop keeps its own KV history, so both compute and KV memory grow with the number of loops. Ouro runs all four of its loops for every token. We study a mixed-depth readout policy for it, in which most tokens are read out after two loops and every fourth token after all four. The deeper loops of the shallow tokens are still needed by later deep tokens, but they can be computed later and in batches (Figure 1a). This deferred schedule saves rollout time, although, as we find, the mixed readout it relies on can cost quality. It also changes the grouping of every matrix multiplication and attention reduction, while the learner of an on-policy algorithm such as PPO or GRPO 1 University of Tübingen, Tübingen, Germany. 2 Xi’an Jiaotong University, Xi’an, China. 3 University of Illinois Urbana-Champaign, Urbana, IL, USA. Correspondence to: Dezhen Lu <dezhen.lu@student.uni-tuebingen.de>, Zixiang Ni <nzx@stu.xjtu.edu.cn>, Runguo Li <runguol2@illinois.edu>. Preprint.

(Schulman et al., 2017; Shao et al., 2024) recomputes the same log-probabilities with a dense pass. This paper makes the two sides agree exactly, then makes the deferred system fast on one node and plans it for models that do not fit on one. Contributions. • Dense replay is not exact. In FP32 about 18% of the sampled tokens miss a 10−5 same-weight gate at 4096 tokens, independently of batching, and nearly all do in bf16 (Section 3). • An exact learner log-probability. A no-grad replay through the rollout’s own schedule supplies the value and the dense learner the gradient (Section 4). The replay is bitwise in every cell on three platforms, in FP32 and bf16, and the gradient is as accurate as any dense gradient, also for the full training loss (Section 5, Section 6). • A fast deferred rollout and training loop. Deferral, one host synchronization per step and per-step CUDA graphs speed up rollouts bit for bit; in a two-node GRPO loop one-wave admission and graphs shorten an update by 2.71×, the loop scales to sixteen GPUs, and we measure what the mixed readout costs in quality (Section 7).

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers (b) Learner log-probability

(a) Deferred rollout, k = 4 readout depth loop 4 loop 3

2

2

2

4

loops 3–4 of t=1..4 one batch at t=4

2

2

2

4

=ℓ bit for bit, else no update

rollout: KV waves, graphs

gate: replay per wave

publish new weights

train: one optimizer step

s

ℓ

ℓ = ℓdense + sg(ℓsched − ℓdense ) value ℓsched , gradient ∇ℓdense

loop 1 t=1 t=2 t=3 t=4 t=5 t=6 t=7 t=8

roll

ed ch

scheduled replay, no grad: rollout’s waves and kernel

loop 2

gate: ℓ

rollout log-probs ℓroll

loops 3–4 of t=5..8 one batch at t=8

(c) One GRPO update sched

dense learner pass (differentiable)

loops 1–2 run every step (blue); the deeper loops wait and run in batches (orange). The function is unchanged, the grouping of every product and reduction is not.

e dens

ℓ

one optimizer step per rollout batch keeps the replay on the rollout’s weights

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.

• A plan for models beyond one node. A ring mapping of layers to GPUs keeps loop-shared weights and perloop KV in place; point-to-point measurements set its cost (Section 8).

computed from the prompt lengths alone. Each wave runs to completion before the next starts; requests never join a running cohort and leave it when they sample a stop token or reach the token limit.

Every system and learning run had its predictions written down before it was submitted; supplementary Appendix E lists every outcome, including the predictions that failed.

Learner. The learner is a Megatron (Shoeybi et al., 2019) module that wraps the same native layers, replays the exact per-loop KV history of each sequence with PyTorch scaled-dot-product attention, and recomputes each layer in the backward pass. It accumulates gradients in FP32 and feeds a fused log-softmax with a hand-written backward; no Transformer Engine kernel is involved. The training loop applies one optimizer step per rollout batch.

2

S ETTING

Model and readout policy. We use Ouro-1.4B (Zhu et al., 2025): 24 decoder layers applied in four loops, hidden size 2048, 16 attention heads. Every loop keeps its own exact KV history. The first response token is read out after all four loops of the prompt prefill. Each later response input runs loops 1–2 immediately and is read out at depth 2, except every k-th input (k = 4 unless noted), which is read out at depth 4. The deferred engine holds the loop-2 states of pending inputs and runs loops 3–4 for all of them at once when a deep readout is due. An eager variant runs loops 3–4 at every step; it produces the same readouts with a different grouping. k = 1 reads out all four loops at every token (“K4”). The mixed-depth readout is our policy, not Ouro’s native decoding, which runs all four loops for every token. What deferral changes. At loop L, a token attends only to the loop-L keys and values of earlier positions, and the input to loop 3 is the token’s own loop-2 output. Deferring loops 3–4 therefore does not change the function, only how the work is grouped: prompt prefill, per-step shallow passes and deferred deep batches put different rows into each matrix multiplication and split each attention reduction differently from a dense pass. Every causal ancestor of a deep readout must itself run all four loops, so almost every token runs them eventually; deferral saves time by batching, not by skipping work. Admission waves. The production adapter splits a rollout batch into waves under a KV byte budget and a request cap,

3

D ENSE R EPLAY D OES N OT R EPRODUCE THE ROLLOUT

Table 1 compares the rollout’s log-probabilities with an independent dense per-sequence pass under the same FP32 weights (arm A0). The gap grows with length: at 96 tokens about 1% of the sampled tokens differ by more than 10−5 , at 4096 tokens about 18% do, on both platforms and even at batch size one. A 10−5 same-weight gate is therefore unpassable at production length, independently of batching. The training loop’s own learner exceeded 10−5 in 33 of 36 cells at 96 tokens (largest 5.9 × 10−5 ) and differed from the rollout by up to 3.5 × 10−4 in the long cells. In FP32 the largest ratio deviation | exp(∆) − 1| is below 10−3 , far inside the PPO clip range, so the problem is the gate, not the size of the update. With bf16 weights the gap is three orders of magnitude larger: 89–98% of the sampled tokens miss and the largest difference is 0.45 at 96 tokens and 1.6 at 4096. A bf16 learner inside a training loop (Section 7.4) missed by a median 0.16 per response (largest 2.7), yet only 0.02% of its tokens left the PPO clip range: at 512 tokens the mismatch fails the gate without changing the update much.

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers Table 1. Dense per-sequence replay (A0) against the rollout’s logprobabilities, FP32 unless marked bf16. “Tokens over 10−5 ” counts sampled tokens whose absolute log-probability difference exceeds the gate. Platform

Length Cells Max |∆ log p|

Tokens over 10−5

−5

L40S L40S L40S GH200 GH200 GH200

96 1024 4096 96 1024 4096

54 3 1 36 4 5

5.1 × 10 161/14,496 (1.1%) 1.4 × 10−4 404/10,016 (4.0%) 2.6 × 10−4 760/4,096 (18.6%) 2.6 × 10−5 65/9,664 (0.7%) 8.9 × 10−5 599/17,984 (3.3%) 8.0 × 10−4 13,545/73,408 (18.5%)

GH200, bf16 GH200, bf16 GH200, bf16

96 1024 4096

20 2 1

4.5 × 10−1 6.7 × 10−1 1.6

4

5,260/5,920 (88.9%) 8,552/8,992 (95.1%) 3,997/4,096 (97.6%)

A H YBRID L OG -P ROBABILITY TAKES I TS VALUE FROM THE ROLLOUT ’ S S CHEDULE

Let ℓsched be the log-probability of sampled token t from t a no-grad, teacher-forced replay through the rollout’s own schedule, and ℓdense the dense learner’s. The learner uses t
ℓt = ℓdense + sg ℓsched − ℓdense , (1) t t t where sg stops the gradient. Its value is ℓsched and its gradit ent is ∇ℓdense (Figure 1b). If the scheduled replay is bitwise t 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 the weights equal those of the rollout, which holds because the loop takes one optimizer step per rollout batch; with several steps per batch, ℓsched would have to be recomputed at each step. The replay needs no gradient, so it can use the rollout’s own attention kernel; a fully differentiable scheduled learner would need a functional KV cache, step-level checkpointing and a differentiable attention backend, and we use one only as a reference arm (A3) at short lengths. Arms. A0: dense per sequence. A1: the batch packed into one dense pass with block-diagonal causal attention. A2: no-grad replay through the rollout’s schedule on the rollout’s backend. A2x: the same schedule on the other attention backend (Triton ↔ PyTorch). A2-solo: the same backend, each request replayed alone. A3: the rollout’s schedule executed differentiably. M: the production learner of Section 2.

5

S CHEDULED R EPLAY I S B ITWISE E XACT

Setup. FP32 weights unless marked bf16, temperature 1, full vocabulary, eight GSM8K-style (Cobbe et al., 2021) prompts, one sampling seed per request. Grid cells vary the batch size (1/2/4/8), the readout schedule (K4, eager, deferred), the length (96, 1024, 4096 tokens), the backend and

the seed; request i of a cell stops after L−8i tokens, so batch membership changes during the cohort. The productionrollout cells instead call the production generation routine with stop tokens and a 1024-token limit. Each hardware platform is a separate cohort with the same pinned set of 137 Python packages (PyTorch 2.12.1 with CUDA 13.0, Triton 3.7.1). Results. Table 2 lists every real-weight cell. A2 is bitwise in all 167 of them. The production-rollout cells include 36 whose batch shrank at stop tokens (166 stopped and 20 length-capped responses), so the replay also covers membership changes that the engine, not the test, decides. Triton and PyTorch rollouts sampled identical tokens in every pair we compared, yet A2x is never bitwise, and neither is A2solo for B ≥ 2. Exactness therefore requires the rollout’s batch composition, step schedule and kernel together; in production this means replaying each admission wave with the rollout’s partition and backend, which the adapter records. bf16. A subset of the grid and the production-rollout cells with bf16 weights, using the same arm code: 28 cells, 34,732 sampled tokens, lengths 96 to 4096, batch 1 to 8, both backends, two seeds, and five production cells whose batches shrank at stop tokens. A2 is bitwise in all 28; A2x is never bitwise, with differences up to 1.2, and A2-solo is never bitwise for B ≥ 2. The exactness of the scheduled replay is a property of the schedule and kernel, not of FP32.

6

T HE H YBRID G RADIENT I S AS ACCURATE AS A NY D ENSE G RADIENT

Truth and For each cell we compute the graP criterion. P dient of i wi t ℓi,t with two weightings, alternatingsign advantages and all ones. The reference is a float64 evaluation of the same function with float64 normalization, softmax and log-softmax and the engine’s rotary constants, accumulated per sequence. Two independent float64 implementations, the learner itself and a functional dense executor, agree to 4.7 × 10−14 in relative gradient norm on real weights. We report e(X) = ∥∇X − ∇truth ∥/∥∇truth ∥ and fixed the criterion before running: e(M) ≤ 2 maxY e(Y ), where Y ranges over other legitimate dense implementations (A0 and A1 at 96 tokens; A0 and the learner on the math attention backend at longer lengths). Long-sequence rollouts disable stop tokens so that every response reaches its limit, the one deviation from production settings; development seeds are reported apart from fresh ones. Results. Table 3 summarizes the 102 comparisons; the criterion holds in all of them. The production learner’s relative error grows from about 10−5 at 96 tokens to about 10−4 at 1024 tokens and stays there at 4096. The single 1.3 × 10−3 outlier is an ill-conditioned K4 cell whose gradient norm is less than half that of the others; every implementation is

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers 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

Cells

A2 bitwise

A2x max A2-solo max

1/2/4/8 1/2/4/8 1/2/4/8 1/2/4/8

10/10 – 58/58 2.1 × 10−5 45/45 9.5 × 10−5 54/54 4.0 × 10−5

2.2 × 10−5 2.1 × 10−5 1.0 × 10−4 –

19,008 96/1024/4096 1/2/4/8 15,724 ≤1024 8

23/23 1.2 5/5 4.4 × 10−1

4.4 × 10−1 –

L40S (CPU arms) L40S GH200 GH200, production rollout

10 2,928 96/512 58 28,608 96/1024/4096 45 101,056 96/1024/4096 54 55,678 ≤1024

GH200, bf16 GH200, bf16, production rollout

23 5

poor there. The criterion is lenient by design: at 96 tokens the explicit-softmax executor was up to 14× further from truth than the learner in an ill-conditioned seed, while the learner was never more than 1.6× worse than either alternative; at 512–4096 tokens the dense executor was up to 1.9× closer in some cells and the math-backend learner up to 3× further in others. Neither is consistently more accurate. The gradient half of Eq. 1 adds nothing beyond the FP32 noise that any dense learner already has. The full training loss. A further pre-registered check took the gradient that the real training step accumulates just before its optimizer step: GRPO advantages normalized within each group, the T3 value of Eq. 1 inside the PPO objective, the mean over each sample’s response tokens, and accumulation over 32 or 64 micro-batches with the step’s normalization. One update of 64 samples was recomputed by an independently written float64 reference, with float64 normalization and the advantages rebuilt from the raw rewards. The relative error of the FP32 gradient was 3.1 × 10−4 with the native readout and 5.7 × 10−6 with the mixed one (cosine similarities 0.99999996 and 0.99999999998), and the clipped gradient norm the training stack logged matched the reference norm to 1.3 × 10−4 . Three deliberately wrong losses (a token mean over the whole batch, advantages without the standard-deviation division, and fulldepth log-probabilities in place of the recorded readouts, which differs only for the mixed readout) were 2,500 to 4 × 106 times further away, so the check would have caught a defect in the loss path. The native readout’s larger relative error comes with an eight times smaller gradient norm for that batch.

7

B

Tokens Lengths

D EFERRED ROLLOUTS A RE FAST I NSIDE THE T RAINING L OOP

All measurements here are FP32 on GH200 GPUs, with stop tokens disabled for single-GPU throughput and enabled in the training loop.

7.1

One GPU

Deferral. In eight comparisons timed on the same node against the eager schedule, which reads out at the same depths but runs loops 3–4 at every step, deferral raised throughput by 1.47–1.56×, from batch 1 to 64 and from 512 to 4096 tokens, up to 79 GiB, with identical peak memory (Table 10). K4 runs at the eager rate, so the saving comes from batching loops 3–4, not from doing less work. Scheduled replay took 0.86–0.96× the rollout time at the rollout’s peak memory and was bitwise in all three configurations measured. One synchronization per step and CUDA graphs. A profile of the deferred engine showed the GPU busy for only 26–36% of each step: sampling read each row’s token back to the host one row at a time, and a step issued 3,000 to 4,500 CUDA operations. We draw every row on the device and read all tokens and log-probabilities back once per step, keeping each request’s random-number stream and per-row operations unchanged, and replay each per-step pass of 24 layers from a CUDA graph captured per row count and shared across loop depths. The synchronization change alone gives 1.01–1.20×; with graphs the rollout is 1.59–2.49× faster than before, most at small batch (Table 11). Rollouts are bitwise equal to the per-row loop with or without graphs, and the scheduled replay from graphs is 1.44–1.72× faster and bitwise in every pairing of graph and eager rollout and replay, so the replay of Eq. 1 may use graphs whether or not the rollout did. Deferral with CUDA graphs. Graphs remove much of the launch overhead that deferral amortizes, so we measured deferral again with graphs on (Table 4); each run builds a fresh runner that captures its own graphs, as the production engine does. Deferral is still 1.18–1.54× faster than the eager schedule, most at small batch, where the GPU is latency-bound; per token, its forward time is 0.62–0.79× the eager one. Host time is 5–10% of a step at batch 8 and about 30% at batch 64, so the forward pass, not sampling, remains the main cost. With stop tokens the batches keep changing shape, so deferral captures about twice as many graphs (41 against 21), and capture takes 27% of its roll-

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers 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. Platform Length Comparisons e(M) median GH200 B300 B300 B300

96 512 1024 4096

72 2 20 8

−6

9.9 × 10 7.2 × 10−5 1.2 × 10−4 1.1 × 10−4

e(M) max Other dense, median e(M)/ max e(Y ), max 6.7 × 10−5 7.2 × 10−5 1.3 × 10−3 1.5 × 10−4

out; without capture time its advantage would be 1.37×. Keeping the KV pool and its graphs across the admission waves of one update recovers almost none of it: in a preregistered check the second wave’s 6–14 new graphs were nearly all row counts the first wave never saw, sharing saved at most one, and the rollout was bitwise unchanged. The captures that repeat are those of each update’s first wave and the gate’s replay of the rollout’s row counts; avoiding them needs the pool to outlive the rollout, which on one GPU competes with training for memory. 7.2

Two Nodes, Eight GPUs

Setup. Two nodes with four GH200 GPUs each run one data-parallel GRPO job of eight ranks (Figure 1c). Each update samples four responses for each of eight prompts per rank (32 per rank, 256 in total) with stop tokens and a 512-token limit, gates on the scheduled replay, takes one optimizer step and publishes the new weights to the rollout engine. The gradient reduce-scatter and all-gather of 5.6 GB across the two nodes took 0.12 s in a calibration run, against about 15 s of training per update. The same calibration showed that raising the wave cap from 8 to 16 requests cut rollout and gate time to 0.60×, which led to the one-wave policy below. We report the median of updates 2 and 3, each at the slowest rank. Results. Table 5 lists five configurations on the same software. Admitting each rank’s 32 samples in one wave, instead of two waves of 16, cut rollout and gate time to 0.46× and 0.45×; replaying the per-step passes from CUDA graphs in both cut them by a further 0.51× and 0.46×. Together the update went from 195.6 s to 70.5 s (2.77×), or to 72.2 s (2.71×) with the memory release described next. Training is now the largest phase, 43% of the update. The gate was bitwise at every rank and update (120 of 120), and the training stack reported a PPO KL estimate and a learner–rollout log-probability difference of exactly zero at every update. At the first update the graph rollouts equal the eager rollouts bit for bit on all eight ranks, and a separate job on other nodes reproduces them bit for bit. The one-wave and two-wave rollouts sampled the same tokens, but their log-probabilities differ by up to 6.0 × 10−4 : wave composition changes the floating-point result, which is why

8.7 × 10−6 5.6 × 10−5 2.0 × 10−4 9.5 × 10−5

1.62 1.09 1.19 0.96

the gate replays each wave with the rollout’s own partition. Memory. One-wave admission sizes a wave by an estimate of KV bytes, and the estimate is accurate (allocated memory grew by at most 1.012× the estimate), yet the GPUs filled to 97,208 of 97,871 MiB and in one configuration an NCCL buffer allocation failed during the final checkpoint save. Releasing cached allocator memory before each rollout and before training (“release” in Table 5) lowered the peak to 63,016 MiB at the same speed. Allocator snapshots then traced a remaining 14 GiB of pinned memory to one 109 MiB block kept alive by a loop variable; freeing it removed the defect (Appendix B). 7.3

Four Nodes, Sixteen GPUs

A pre-registered job ran the final two-node configuration on four nodes with four arms (Table 6). With the per-rank load unchanged (512 samples per update), the update took 81.2 s with the native readout and 69.2 s with the mixed one, against 79.6 and 69.4 s on eight GPUs: weak scaling within 2%, the gate bitwise at all 640 rank-updates, peak memory unchanged and no pinned memory. Learning matched the two-node runs: 84.2% and 76.4% after 20 updates (84.6% and 75.8% on two nodes), with the native readout again ahead on 68 problems against 29 (exact McNemar p = 9 × 10−5 ). Strong scaling did not meet its prediction: with 16 samples per rank, training halved, from 32.3 to 15.9 s, but the rollout (21.6 against 22.4 s) and the gate (20.1 against 20.1 s) did not move, so the update fell to 0.76× rather than the predicted 0.65×. Both phases are bounded by the longest response, 512 decode steps of one graph replay each, which a smaller wave does not shorten; adding ranks compresses only training, and the update cannot fall below about 42 s in this configuration. Turning activation recomputation off cut training from 32.0 to 19.1 s (0.60×) with the first gradient norm bit for bit identical and the peak unchanged. 7.4

Learning

Setup. We trained each readout policy, the mixed readout of Section 2 and the native K4 readout, for 100 GRPO updates. The pre-registered runs use the two-node setup: four responses for each of 64 GSM8K training prompts per

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers 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. Config

Deferred Eager

B = 8, 512 B = 64, 512 B = 24, 2048 B = 32, 512, stop tokens

316 1211 591 485

205 1025 452 412

K4 Def./eager W/o capture Capture 213 976 442 416

1.54 1.18 1.31 1.18

1.56 1.19 1.31 1.37

2% 1% 1% 27%

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. Configuration Cap 16 (2 waves) One wave One wave + graphs One wave, release One wave + graphs, release

Waves Rollout s Gate s Train s Step s Rollout tok/s Gate bitwise Peak MiB 2 1 1 1 1

80.1 36.5 18.7 37.0 19.2

77.6 34.7 15.9 34.2 16.1

update (256 in total, 32 per rank in one wave), stop tokens, a 512-token limit, temperature 1 and learning rate 10−6 . A reduced version on one GH200 with 16 prompts per update ran first. Data order and seeds are the same across runs. Before training and every 20 updates, each run scores one sample on each of the first 500 GSM8K test problems with the training reward (a correct answer inside \boxed{}), reusing the same per-problem seeds. Results. Table 7 shows held-out accuracy. On two nodes the mixed readout improved by 59.4 points (300 problems gained, 3 lost) and K4 by 56.2 (288 gained, 7 lost). Part of the gain is learning to stop within 512 tokens (the stopped fraction rose from 41% and 49% to 98% and 97%), but not all: among responses that stopped, accuracy rose from 25% to 84% and from 34% to 91%. We did not keep the response texts, so we cannot separate learning to box the answer from better reasoning. The gate was bitwise at every rank and update (1,600 of 1,600 on two nodes, 200 of 200 on one GPU), with a PPO KL estimate of exactly zero; this is also the first evidence of the hybrid learner at batch for K4. Initial accuracies are identical across setups, as the initial weights and evaluation seeds are. The mixed readout costs quality. K4 was 10.8 points ahead before training and still 7.6 ahead after 100 updates on two nodes (9.0 on one GPU). On the last evaluation K4 alone solved 52 problems that the mixed readout missed, against 14 the other way (exact McNemar p = 2.8 × 10−6 ; one training seed per policy). The mixed readout’s update took 69.4 s against 79.6 s (0.87×), with the rollout 1.23× faster per token. K4 reached 84.6% after 20 updates, about 27 minutes, already above the mixed readout’s 82.4% after

32.5 195.6 29.9 106.2 30.6 70.5 30.3 106.9 31.5 72.2

1,336 2,929 5,728 2,892 5,578

24/24 24/24 24/24 24/24 24/24

67,062 97,208 97,034 63,016 63,246

100 updates, about 1.9 hours. In this configuration the mixed readout is not worth its speed. On one GPU the gap was smaller (0.96×): training took 57% of the update, and two waves per update each captured their own graphs, so the rollout was only 1.09× faster per token. Single-GPU training was host-bound, and is no longer. With the same software, the single GPU was busy a median 33–34% of the sampled time against 79–92% on the cluster nodes, bf16 training took as long as FP32, and the autograd thread held one core at its maximum clock. Two pre-registered checks removed the host cost. Keeping the activations instead of recomputing them cut the training phase from 116.6 to 62.9 s per update (0.54×) with the rollout bitwise identical, the first gradient norm within 8 × 10−8 and the peak unchanged, because the peak is set by the rollout’s KV budget. Running every sequence of a micro-batch through the learner at once, right-padded to the longest and one slab per loop over the prefixes the exact-KV schedule requires (float64 agreement with the per-sequence path to 10−10 in a CPU test), cut it to 35.4 s with two sequences per micro-batch (0.30× the original), 0.55 s per sample against 1.01 s on the two-node run, with the GPU busy 92– 98% of the time and the rollout bitwise identical (Table 12). The batched gradient differs from the per-sequence one by 4 × 10−4 in norm; against float64 both FP32 paths were 1–7 × 10−4 from the truth, the batched one closer on the K4 schedule and farther on a mixed one, so they are of the same accuracy class rather than equivalent. bf16, and training without the gate. A pre-registered follow-up repeated the reduced mixed-readout run with a bf16 learner and rollout engine (FP32 master weights and

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers 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. Arm K4, weak Mixed, weak K4, strong K4, no recompute

Samples Updates Update (s) Rollout Gate Train Peak (MiB) Gate bitwise 512 512 256 512

20 20 10 5

81.2 69.2 60.2 68.3

23.4 18.5 21.6 23.7

20.9 16.2 20.1 21.0

32.0 30.3 15.9 19.1

49,582 50,000 39,262 49,518

Accuracy

320/320 33.8% → 84.2% 320/320 23.0% → 76.4% 160/160 33.8% → 79.6% 80/80 –

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 preregistered; the one-GPU runs use a quarter of the samples per update. One training seed per run.

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.

Updates 8 GPUs, 256/update 1 GPU, 64/update Mixed K4 Mixed K4

Updates bf16, T3 bf16, dense ungated FP32, T3

0 23.0% 20 75.8% 40 78.8% 60 81.6% 80 82.0% 100 82.4%

33.8% 23.0% 84.6% 69.2% 87.8% 74.4% 87.6% 80.8% 89.0% 81.0% 90.0% 78.8%

33.8% 82.6% 84.4% 86.6% 87.6% 87.8%

optimizer state) in two arms: the hybrid learner with its gate, and training on the dense log-probabilities with no gate. The gate stayed bitwise at all 100 updates in bf16 with a PPO KL estimate of exactly zero. The ungated arm saw the mismatch at every update, yet reached the same accuracy, 77.0% against 77.2% (Table 8; 42 problems solved by the ungated arm alone and 43 by the gated one, exact McNemar p = 1, and no evaluation during training separated them, p ≥ 0.21). Only 0.02% of its tokens left the PPO clip range, so at 512 tokens the deviation is neither caught by the clip nor large enough to hurt; at this scale the hybrid learner buys certainty, not accuracy. Both bf16 arms trailed the FP32 run for the first 40 updates (51% against 69% after 20) and ended 1.6–1.8 points behind; a consistent but untested explanation is the checkpoint’s bf16 weight grid, which lets at most 26% of the weights leave their initial value within 20 Adam updates at a learning rate of 10−6 . The bf16 update was 0.89× the FP32 one (201.8 against 227.3 s), not the predicted 0.8×; training itself took as long, as expected of a host-bound step.

8

L OOPED M ODELS M AP ONTO A GPU R ING

This section is a design with measurements of its communication paths; the mapped engine is not yet implemented. When one node is not enough. With a distributed optimizer over D data-parallel replicas, each replica holds about 6 +

0 20 40 60 80 100

23.6% 50.8% 69.2% 72.4% 74.0% 77.2%

23.6% 50.0% 69.0% 70.0% 72.6% 77.0%

23.0% 69.2% 74.4% 80.8% 81.0% 78.8%

12/D bytes per parameter. For three looped shapes at D = 8 the rollout weights alone fit one node of four 95 GiB GPUs, but not the combination that RL on one set of GPUs requires: learner state, a second copy of the weights for the rollout engine, and KV for every loop. For a 69B shape the learner state and rollout weights are already 613 GiB (Table 13). The mapping. Place the layers on S GPUs of N nodes as contiguous arcs, one arc per node, connected in a ring. Each loop sends the hidden state once around the ring: S boundary crossings, of which only the N node boundaries use the network; N is a lower bound for any placement on N nodes. On nodes of four GPUs three quarters of the crossings stay on NVLink. A layer’s K loops run on the same GPU, so the gradient of the shared weights accumulates locally and data-parallel traffic stays at one copy of the parameters; the KV of each layer and loop never moves; and a GPU’s layers are visited K times per forward pass, like the K virtual stages of an interleaved pipeline schedule (Narayanan et al., 2021), whose bubble shrinks by about K (existing schedules assume distinct layers per virtual stage). Deferral does not reduce the bytes that cross the ring, but it groups loops 3–4 into fewer and larger messages. Against sharded data parallelism. FSDP and ZeRO-3 (Zhao et al., 2023; Rajbhandari et al., 2019) gather each layer’s weights when the layer runs. A looped model runs each layer K times forward and K times backward per micro-batch, so it regathers every weight 2K times unless it keeps the gathered weights and gives up the memory saving. The ring moves no weights: for an 8192-token micro-batch,

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers 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. Path

4 MiB 64 MiB 256 MiB

NIC, one pair NIC, one pair to a different GPU index NVLink, 8-rank communicator, defaults NVLink, node-local communicator NVLink, 8-rank, 8 channels per net peer NVLink, copy engine via CUDA IPC Node boundary split over 4 NICs, 8 channels per net peer

21.8 21.5 45.6 73.3 68.7 90.7 13.4

23.8 23.6 51.4 121.4 122.1 129.3 90.4

23.8 23.7 51.4 124.3 124.3 132.0 92.6

FSDP gathers 324–1032 times more bytes per GPU than one ring boundary carries. Measured communication paths. Table 9 reports point-topoint bandwidth between two nodes of four GH200 GPUs, connected by NVLink inside a node and by one 200 Gb/s NIC per GPU with GPUDirect RDMA between nodes. A NIC carries 23.8 GB/s, 95% of its line rate, four pairs run concurrently at nearly that rate each, and the remote GPU index does not matter on this switched fabric. NCCL send/recv over NVLink reached only 51 GB/s inside an eight-rank communicator, because NCCL assigns point-topoint channels per peer from the minimum over all peers and network peers get two by default; a node-local communicator or eight channels per network peer give 124 GB/s without slowing the NIC, and CUDA IPC with the copy engine gives 132 GB/s. Splitting a node boundary over four NICs with an NVLink scatter and gather delivers 92 GB/s at 256 MiB in eight pipelined chunks; at 64 MiB and below chunking does not pay. A host-staged transfer reached 4.9 GB/s, so activations should never leave the GPU, while a 1 KiB host message took 38 µs one way. Small messages are bound by the host cost of each call, about 60 µs in torch.distributed; an eight-hop ring of 64 KiB messages took 159 µs per loop, and 203 µs when captured in a CUDA graph. With these bandwidths a per-hop model, which matched the measured 256 MiB rings within 1%, puts the forward activations of one loop around a two-node ring for the 69.3B shape at 17.8 ms per micro-batch with tuned NVLink channels and 9.4 ms with the four-NIC split, against 26.9 ms with NCCL defaults, while the forward and backward compute per stage and loop is about 1050 ms under an assumed 400 TFLOP/s, so the ring’s communication can be hidden. The open costs are the pipeline bubble, memory balance across arcs, and the per-call host overhead of the many small messages of decoding.

9

R ELATED W ORK

Looped and recurrent-depth transformers (Dehghani et al., 2018; Geiping et al., 2025; Zhu et al., 2025) reuse layers to add depth at fixed parameter count; we apply a mixed-

depth readout policy to Ouro and make its rollouts exact and fast. RL frameworks for language models such as HybridFlow (Sheng et al., 2024) co-locate or separate rollout and training engines, and RL objectives such as PPO and GRPO (Schulman et al., 2017; Shao et al., 2024) rely on the learner reproducing the rollout’s probabilities. When separate rollout and training stacks disagree, the update is effectively off-policy; truncated importance sampling corrects for the mismatch with capped learner-to-rollout probability ratios (Yao et al., 2025), and training in FP16 instead of bf16 reduces it (Qi et al., 2025). Our hybrid learner removes the mismatch from the value instead of correcting for it, and keeps the dense learner only for the gradient. Paged KV management (Kwon et al., 2023) underlies our admission budget. Pipeline and sharded data parallelism (Shoeybi et al., 2019; Narayanan et al., 2021; Rajbhandari et al., 2019; Zhao et al., 2023) are the baselines for Section 8; the loop structure changes their trade-off because layers are reused within one forward pass.

10

L IMITATIONS

Learning quality. The learning evidence is one seed per readout policy and setup, 100 updates of 256 (two nodes) or 64 (one GPU) samples, FP32, and a 500-problem evaluation with one sample per problem. Precision and optimizer. The scheduled replay is exact in FP32 and bf16, and the gate stayed bitwise through 100 bf16 training updates, but the gradient and throughput results are FP32, and the bf16 training evidence is one seed on one GPU at 512 tokens, where the ungated learner did as well. Whether the gate matters for learning at longer sequences, where the dense mismatch is ten times larger, is open. The runs used Megatron’s PyTorch Adam fallback. Gradient criterion. The criterion of Section 6 is on a surrogate loss; the full GRPO loss was checked against float64 on one batch at 512 tokens. Under T3 the ratio is exactly one, so the PPO clipping branch never runs. Readout policy. Deferral needs the mixeddepth readout, which trailed the native one by 7.6–9 points after 100 updates while an update was only 4–13% shorter. We tried one readout interval (k = 4). Scale. One model (Ouro-1.4B), up to 4096 tokens in the exactness study and up to four nodes of data parallelism, where strong scaling is bounded by the per-token latency of rollout and gate. The mapping of Section 8 is a design backed by point-to-point measurements.

11

C ONCLUSION

A deferred rollout of a looped transformer and a dense learner compute the same function in different floating-point orders, so a strict same-weight gate fails at production length. Replaying the rollout’s own schedule restores bitwise agreement in every cell we ran, in FP32 and bf16, while the dense

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers

learner keeps supplying a gradient as accurate as any dense gradient. With that gate in place, deferral, one host synchronization per step, CUDA graphs and one-wave admission make the training loop fast and keep it bitwise on two and four nodes. The mixed readout that deferral needs, however, learned less than the native readout in our runs, which reached a higher accuracy in less time.

Yao, F., Liu, L., Zhang, D., Dong, C., Shang, J., and Gao, J. On the Rollout-Training Mismatch in Modern RL Systems. OPT 2025: Optimization for Machine Learning, NeurIPS 2025 workshop, 2025.

R EFERENCES

Zhu, R.-J., Wang, Z., Hua, K., Zhang, T., Li, Z., Que, H., et al. Scaling Latent Reasoning via Looped Language Models. arXiv:2510.25741, 2025.

Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., et al. Training Verifiers to Solve Math Word Problems. arXiv:2110.14168, 2021. Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, Ł. Universal Transformers. arXiv:1807.03819, 2018. Geiping, J., McLeish, S., Jain, N., Kirchenbauer, J., Singh, S., Bartoldson, B. R., et al. Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach. arXiv:2502.05171, 2025. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180, 2023. Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. arXiv:2104.04473, 2021. Qi, P., Liu, Z., Zhou, X., Pang, T., Du, C., Lee, W. S., and Lin, M. Defeating the Training-Inference Mismatch via FP16. arXiv:2510.26788, 2025. Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054, 2019. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal Policy Optimization Algorithms. arXiv:1707.06347, 2017. Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., et al. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300, 2024. Sheng, G., Zhang, C., Ye, Z., Wu, X., Zhang, W., et al. HybridFlow: A Flexible and Efficient RLHF Framework. arXiv:2409.19256, 2024. Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053, 2019.

Zhao, Y., Gu, A., Varma, R., Luo, L., Huang, C.-C., et al. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. arXiv:2304.11277, 2023.

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers

A

S INGLE -GPU ROLLOUT TABLES

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. Length 512 512 512 512 2048 2048 2048 4096 4096 4096

B Eager tok/s Deferred tok/s Ratio Peak GiB Note 1 8 32 64 8 16 24 4 8 12

10.8 85.2 269.1 479.7 83.0 139.8 260.3 56.2 105.5 140.2

16.5 128.7 402.8 705.2 127.3 236.1 392.4 85.9 132.9 218.3

1.53 1.51 1.50 1.47 1.53 – 1.51 1.53 1.26 1.56

6.2 12.0 31.9 58.4 30.0 54.6 different nodes 79.2 29.7 54.0 unexplained 78.3

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. B Length Per-row sync One sync/step + CUDA graphs Speedup Bitwise 8 64 24 12

B

512 512 2048 4096

138 888 336 194

165 894 396 213

342 1412 654 308

2.49× 1.59× 1.95× 1.59×

yes yes yes yes

T HE P INNED -M EMORY D EFECT

One-wave admission sizes a wave by an estimate of KV bytes, and the estimate is accurate: the allocated memory grew during the rollout by at most 1.012× the estimate. The GPUs nevertheless filled up, to 97,208 MiB of 97,871 MiB, and in one configuration an NCCL buffer allocation failed during the final checkpoint save, which then hung until the collective timed out. Releasing cached allocator memory before each rollout and before training lowered the peak to 63,016 MiB at the same speed. Allocator logging then showed one remaining defect: once per rank, about 14 GiB stayed reserved after the release. Allocator snapshots on one GPU traced it to a single 109 MiB block, the FP32 logits of the last sequence of a diagnostic dense replay that runs between the rollout and the gate. A loop variable kept the block alive; it was allocated in a KV segment the rollout had just freed, so neither the gate nor the release could reuse that segment. Freeing the variable removed the defect: in the training runs of Section 7.4, the reserved memory after the release equalled that before the rollout at every one of 200 updates. The KV budget of one-wave admission is now derived from free device memory only, so memory pinned in this way cannot be counted as available.

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers

C

BATCHED L EARNER F ORWARD

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.

D

N

Path

Recompute

Train (s)

1 1 2 3 4 8 4 8 16 32

per sequence batched batched batched batched batched batched batched batched batched

off off off off off off on on on on

62.8 54.1 35.4 36.2 out of memory out of memory 40.2 39.1 36.3 36.6

per sample (s)

GPU busy

Train peak (GiB)

Gradient norm

0.98 0.85 0.55 0.57

84% 73% 92% 96%

44.8 40.4 51.2 64.4

3.44061 3.44061 3.43930 3.43887

0.63 0.61 0.57 0.57

97% 97% 97% 98%

33.5 35.6 39.8 49.0

3.43869 3.43918 3.43918 3.43918

M EMORY A RITHMETIC FOR THE R ING M APPING

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. Shape Params States/replica Rollout weights FSDP regather Ring boundary Ratio h = 4096, 48 layers h = 6144, 60 layers h = 8192, 80 layers

E

10.9B 29.7B 69.3B

76 GiB 207 GiB 484 GiB

20 GiB 55 GiB 129 GiB

174 GB 475 GB 1108 GB

0.54 GB 0.81 GB 1.07 GB

324 590 1032

P RE -R EGISTERED P REDICTIONS AND T HEIR O UTCOMES

Each run below had its predictions written down before it was submitted. We list every prediction that failed. The singleGPU CUDA-graph measurement of Section 7.1 was not pre-registered. Three further submissions failed before measuring anything and were repeated after the cause was fixed: a launcher missing from the path, a network plugin that fell back to sockets because of a missing C++ runtime symbol, and a hang caused by creating a process subgroup that not all ranks joined in the same way.

ScaleRLT: Exact Log-Probabilities for RL on Looped Transformers Table 14. Outcomes of the pre-registered system, gradient, communication and learning runs. Run

Predictions Met Not met (predicted; measured) 6

5 Update time on 8 GPUs at most 1.05× that on 4; 1.066×

9

8 Peak memory at most 90 GB; 97,208 MiB

8

7 Reserved memory before training within 2 GiB of that before the rollout; 13.8–14.2 GiB once per rank 4 – 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 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 3 – 4 Host time at least 40% of a step at batch ≤ 12; 5–10%. (A first attempt with deterministic kernels was voided: it slowed every schedule several-fold, unlike production) 6 Strong-scaling update at most 0.65× the eight-GPU one; 0.76× (rollout and gate are per-token latency) 5 Mixed-readout update at most 0.8× the K4 update; 0.87× 5 The same; 0.96× (training was 57% of an update, and the rollout only 1.09× faster per token) 5 Ungated arm’s clipped fraction at least 1%; 0.02%. bf16 gated update at most 0.8× the FP32 one; 0.89× (training itself took as long) 5 –

Two-node calibration, 16 samples/rank Two-node admission and CUDA graphs Two-node memory rerun Allocator snapshots, one GPU Point-to-point, first round

4 10

Point-to-point, second round

8

bf16 replay check Deferral with CUDA graphs

3 5

Four nodes, sixteen GPUs (Section 7.3) Learning, two nodes (Section 7.4) Learning, one GPU (Section 7.4)

7 6 6

bf16 training, one GPU (Section 7.4)

7

No recompute, one GPU (Section 7.4) Full GRPO-loss gradient (Section 6)

5 5

Cross-wave KV pool (Section 7.1)

5

Batched forward (Section 7.4)

9

3 Batched FP32 gradient within 2× the per-sequence error; 2.04× (native readout). Every wrong loss at least 10× further; the full-depth one equals the native readout’s correct loss (design error) 3 Later-wave captures at most half; 18 of 20 and 10 of 10. Rollout and gate at most 0.9×; 1.08× and 0.98× 3 Four sequences fit without recomputation; out of memory. Gradient norm within 10−4 of the per-sequence path; 4 × 10−4 . Dense metric at most 0.6×; 0.8× (already compute-bound). Batched FP32 gradient within 2× the per-sequence error against float64 on both schedules; one of two

Related documents

• ScaleRLT: Exact Log-Probabilities and Fast Deferred Rollouts for Reinforcement Learning on Looped Transformers#1457581
• Improving Test-Time Scaling with Adaptive Looped Transformers#1130395
• Improving Test-Time Scaling with Adaptive Looped Transformers#1108746
• Real-Time Instruction Internalization for Large Language Models#272196
• Scheduling Recursive Reasoning in Looped Transformers#1131854