arXiv CS open access

LoCoJAX: Host-Local Coordination for Small-Model RL on TPUs

Loki Chen et al. · arxiv_cs

arXiv CS · Papers

cloud distributed-computing parallel-computing

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

LoCoJAX: Host-Local Coordination for Small-Model RL on TPUs

arXiv:2611.16301v1 [cs.LG] 12 Nov 2026

Loki Chen1,∗ , Gregory B. Rehm1,∗ , Yang Song1 , Hongye Xie1 , Ming Lei1 , Chongyang Bai1 , Bob Kamma1 , Mingfu Liang1 , Kavosh Asadi1 , Yufei Li1 , Xi Liu1 , Jay Xu1 , Rui Liu1 , Maziar Sanjabi1 , Luke Simon1 , Hamed Firooz1 1

Meta

∗

Equal contribution

Reinforcement learning (RL) increasingly targets small language models, and executing these workloads efficiently across multiple accelerators requires effective coordination. On TPUs, Pathways-based RL uses a single controller to coordinate computation across hosts. This execution model is effective for large models, but its coordination computational costs are harder to amortize when small models repeatedly invoke short computations. To characterize these computational costs, we sweep dense models from 0.6B to 405B parameters on TPU v7x under a decode-heavy workload and measure how host and network computational overhead vary with model size. At 0.6B, host orchestration occupies 36% of step time and network communication 2%, whereas for large models, estimated network computational costs dominate the non-compute time. Motivated by these findings, we present Local Coordination JAX (LoCoJAX), which implements local rollout and reward processing on each TPU host with synchronous global training. In controlled comparisons against a tuned and optimized Pathways baseline on Qwen3-0.6B across three workloads and 8–64 TPU v7x chips, its colocated mode (LoCoJAX-CO) is 1.3–3.6× faster per step; on GSM8K, it completes a step faster on 8 chips than Pathways does on 64. Finally, a Ray-managed disaggregated mode (LoCoJAX-DA) overlaps rollout and training on separate hosts, and placement experiments establish when this overlap compensates for dividing devices and transferring weights: only when rollout dominates the step and receives most of the hosts. On NuminaMath-CoT at 256 chips, giving rollout 75% of hosts cuts step time by about 15% compared with LoCoJAX-CO. Correspondence: Loki Chen at dengcchi@meta.com

Introduction

Reinforcement learning (RL) is now a standard stage of large language model (LLM) post-training (Christiano et al., 2017; Ouyang et al., 2022), particularly for reasoning tasks with verifiable rewards (Shao et al., 2024; Lambert et al., 2024; DeepSeek-AI et al., 2025). As a result, extensive effort has been devoted to characterizing and optimizing RL at different stages of its lifecycle, including rollouts (Fu et al., 2025; Zhou et al., 2025) and trainer-side improvements (Zhong et al., 2025b; Wan et al., 2026). Most of this research has focused on models over 7B in size on GPU clusters, where generating responses, training, and redistributing weights are all computationally expensive (Mei et al., 2025; Sheng et al., 2025; Zhong et al., 2025b; Fu et al., 2025; Zhong et al., 2025a; Yu et al., 2026). However, to the best of our knowledge, comparatively little work has addressed small-model RL, where optimization challenges can differ from those of larger models. Today, RL increasingly targets smaller models: sub-billion and few-billion-parameter reasoners (Dang and Ngo, 2025; Wang et al., 2026; Zhuang et al., 2026; Zeng et al., 2025). These models often serve as fast reasoning agents in latency-sensitive applications (Belcak et al., 2025) or on embedded devices (Liu et al., 2024). Shrinking a model shrinks its compute but not the incidentals around it. Prior work finds framework dispatch (Fernandez et al., 2023), kernel boundaries (Spector et al., 2025), launch latency (Shen et al., 2024), request scheduling (vLLM Team, 2024; Zhou et al., 2025), and single-controller data paths (Zhou et al., 2025) can all become bottlenecks on small models. Taken together, these works show that small-model RL steps are often bound by incidentals outside of LLM computations.

Host orchestration

Share of step time

40%

Network communication

Host comparable at 8B

30% 20% 10%

Network

0% B

0.6

1.7

8B 14B 32B 70B 00B* 00B* 05B 4 ~1 ~2 Model size (128 TPU v7x chips)

Figure 1 Where the Group Relative Policy Optimization (GRPO) step goes as model size grows. For each dense model, we

report the host-orchestration and network shares of step time at its fastest tensor-parallel configuration on 128 TPU v7x chips under a decode-heavy workload with 8,192 forced output tokens (Table 2). Host orchestration is the largest non-compute component for small models and network communication for large ones.

Prior work on small-model RL bottlenecks was all run on GPUs, and to our knowledge, TPUs remain unstudied in this context. TPUs differ from GPUs in hardware, software and control. Chips connect over a dedicated inter-chip interconnect (ICI) in 3D-torus slices (Jouppi et al., 2023). The software stack is JAX and XLA (Bradbury et al., 2018), with kernels in Pallas (JAX Authors, 2026b; Jiang et al., 2026), training in MaxText/Tunix (Google, 2023; Bao et al., 2025), and rollout in vLLM’s TPU backend (vLLM Project, 2025). These differences reach the RL loop: a TPU device can be opened by only one user-space process, so trainer and engine cannot be colocated in separate processes out of the box (Basile et al., 2026). Multi-host programs run either under multi-controller JAX, with one process per host issuing the same computations in the same order (JAX Authors, 2026a), or under Pathways, whose single controller drives every host (Barham et al., 2022); Google’s documented multi-host TPU RL recipes use Pathways (Google Cloud, 2026d,a,b). Prior work still finds large improvements in TPU LLM serving, from vLLM’s TPU backend to its Pallas attention kernel (Lee et al., 2025; vLLM Team, 2025; Jiang et al., 2026), and reports single-controller dispatch and communication computational overheads that grow with scale (Kumar et al., 2021; Barham et al., 2022; Google Cloud, 2026c). Together, these differences make TPUs an optimization landscape distinct from GPUs, and one that small-model RL has not yet explored. Motivated by this gap, we first measured where small-model RL spends its time on TPUs. We profile decode-heavy GRPO under a Pathways recipe on 128 TPU v7x chips, across dense models from 0.6B to 405B parameters, and split each step into compute, network communication and host orchestration (Figure 1, Table 2). The breakdown of step time varies with model size. For Qwen3-0.6B, host orchestration takes 36% of the GRPO step (rollout plus training) and network communication 2%. The two shares are equal at Qwen3-8B (18.3% and 18.5%), and beyond that network computational overhead becomes the dominant factor in latency. Thus, in the evaluated Pathways stack, host orchestration is the largest non-compute component for small models and our design therefore focuses on reducing this computational overhead. Using these findings, we design Local Coordination JAX (LoCoJAX), which moves rollout control from a single controller to each TPU host while training remains globally sharded (Figure 2). Each prompt and its full GRPO response group go to one host, whose engine generates the responses and whose process computes rewards and group advantages without cross-host communication. All hosts then train jointly and load the updated weights into their local engines before the next rollout. This colocated mode (LoCoJAX-CO) changes execution, not the algorithm: training stays synchronous and on-policy, as in the Pathways recipe. Against our Pathways baseline built from public MaxText, Tunix and vLLM TPU components, optimized LoCoJAX-CO is 1.3–3.6× faster per step on Qwen3-0.6B across three workloads and 8–64 TPU v7x chips; on GSM8K, it completes a step faster on 8 chips than Pathways does on 64. 2

(a) Pathways recipe

(b) LoCoJAX-CO (colocated, synchronous)

Global Controller (scheduling, coordination, request admission)

Host 0

Host 1

Host N

JAX controller (admission, GRPO grouping, decode loop)

Host 0

Host 1

Host N

Rollout Engine(s)

Local Rollout Engine

Model (Shard)

Trainer (JAX)

Cross-host collectives (gradient synchronization, parameter update)

Rollout Host 1

Rollout Host Nr

JAX controller

Rollout Engine(s)

Ray (weight distribution, result collection)

Trainer Host 0

Trainer Host 1

Trainer Host Nt

Model (Shard)

Trainer (JAX)

Cross-host collectives (gradient synchronization, parameter update)

Host (TPU node)

Model (sharded across hosts)

Controller / Runtime

Rollout engine (generation, reward computation)

Trainer (optimizer, parameter update)

Cross-host communication (e.g., all-reduce)

Control or data flow

Figure 2 System architectures of Pathways, LoCoJAX-CO, and LoCoJAX-DA. (a) Our evaluated baseline follows Pathways’

documented single-controller SPMD execution model: one CPU-head controller drives a globally distributed rollout engine and sharded training across TPU hosts. (b) LoCoJAX-CO adopts synchronous, colocated execution with independent host-local JAX controllers and rollout engines. (c) LoCoJAX-DA uses Ray-managed disaggregation to separate rollout and training hosts, enabling asynchronous execution and overlap between generation and training.

Profiling both systems on GSM8K at 64 chips localizes most of the measured step-time gap to rollout (Section 4.3). Pathways’ rollout runtime incurs two kinds of recurring computational overhead: its head admits every request serially each step, a computational cost that does not shrink with more chips; and each decode iteration takes more than twice as long as on LoCoJAX’s host-local engines (Figure 7). We also contribute a disaggregated mode (LoCoJAX-DA) that uses Ray (Moritz et al., 2018) to run rollout and training on separate hosts (Section 4.4). LoCoJAX-DA is modestly superior when rollout dominates the step and receives most of the hosts, but on our balanced and prefill-heavy datasets colocation usually wins. As a result, we recommend LoCoJAX-CO as a strong default for small-model RL on TPUs. In summary, this paper makes the following contributions: 1. We characterize decode-heavy GRPO on TPU v7x across dense models from 0.6B to 405B parameters. Host orchestration is the largest component for the smallest models, whereas network communication dominates the non-compute time of the largest. 2. We design and implement LoCoJAX, an end-to-end TPU RL system designed for small models. Compared with a tuned and optimized Pathways baseline on Qwen3-0.6B, LoCoJAX reduces step time by a factor of 1.3–3.6. It outperforms Google’s public recipes on Qwen3-1.7B (5.3×) and Gemma4-26B MoE (2.1×). 3. We identify two rollout-side computational costs—serial request admission and higher per-iteration decode latency—that account for most of the measured performance gap between Pathways and LoCoJAX on GSM8K at 64 chips. 4. We compare colocated and disaggregated LoCoJAX execution at 8–256 chips and three trainer/rollout splits. Colocation is the better default, but disaggregation can improve step time for rollout-heavy workloads at sufficient scale.

LoCoJAX Design

In this section, we introduce LoCoJAX, a system designed for RL training of small models on TPUs. It builds on JAX’s standard multi-controller mode, typically one process per host (JAX Authors, 2026a). LoCoJAX combines local request admission and GRPO grouping with globally sharded training and weight publication to local rollout engines. Figure 2 compares this execution model with the Pathways baseline. We support two 3

Algorithm 1 One colocated LoCoJAX step for Qwen3-0.6B, run on each host h among H hosts. Require: global prompt batch B of P prompts with G responses each; global policy θ with shard θh on host h

(FSDP); local engine Eh (DP8 × TP1)

1: Bh ← LocalShard(B, h) ▷ P/H prompts; balanced by prompt length when optimized 2: Rh ← Eh .Generate(Bh ) ▷ host admits only its P G/H requests 3: rh ← Reward(Rh ); Ah ← GroupAdvantage(rh ) ▷ all G samples of a prompt are local 4: Xh ← Pack(Rh , Ah ) ▷ fixed-length rows; hosts agree on padded rows per device K 5: θ ← TrainStep(θ, Xh ) ▷ one forward/backward pass and update; FSDP collectives 6: Eh .LoadWeights(ReshardToEngine(θ, h))

next rollout

▷ gather FSDP shards into local engine replicas before

Table 1 Stack-comparison workloads (Qwen3-0.6B): prompts × generations per step, prompt and response caps (tokens),

and Qwen3 thinking mode.

Dataset

Character

GSM8K NuminaMath-CoT SWE-agent trajectories

balanced prefill and decode decode-heavy prefill-heavy

Prompts × generations 128 × 16 256 × 8 256 × 8

Prompt cap

Response cap

512 4,096 15,360

2,048 8,192 1,024

Thinking off on off

modes: colocated execution (LoCoJAX-CO) and disaggregated execution managed by Ray (LoCoJAX-DA). Unless marked otherwise, LoCoJAX refers to LoCoJAX-CO. Colocated mode (LoCoJAX-CO). Each host runs one process of a trainer whose state is sharded across all

devices using fully sharded data parallelism (FSDP), alongside a local vLLM rollout engine (Figure 2b). For Qwen3-0.6B, the engine runs eight data-parallel (DP) replicas per host, one per TPU TensorCore, each using tensor parallelism (TP) of degree one (DP8 × TP1). Rollout TP is configurable; larger models can use TP2 or higher. Each step follows five phases: request admission, generation, reward, training, and weight publication (Appendix C). Each host admits its share of the prompts and generates G responses per prompt, where G is the GRPO group size. Because each group stays on one host, that host can evaluate rewards and compute group advantages without cross-host communication. All hosts then train jointly on the batch. The updated weight shards are all-gathered over the TPU ICI, and each host loads the weights into its local engine before the next rollout. This keeps learning synchronous and on-policy. Relative to the Pathways baseline, LoCoJAX changes where requests are admitted and where rollout ranks synchronize: • Per-host request admission. The evaluated Pathways baseline admits all 2,048 requests serially in its head process. LoCoJAX distributes this work: each of H hosts admits 2,048/H requests in parallel with the others (128 requests per host at 64 chips). • Rollout synchronization. The Pathways recipe runs all rollout ranks in lockstep as one SPMD program. During generation, LoCoJAX synchronizes only the devices within each host. Cross-host communication remains necessary for packed-shape agreement, FSDP training collectives, and weight publication. Section 4.3 measures the two recurring controller computational costs: serial request admission once per step and the decode premium once per iteration. Disaggregated mode (LoCoJAX-DA). LoCoJAX-DA uses Ray to launch and coordinate separate trainer and

rollout hosts (Moritz et al., 2018). Ray (Moritz et al., 2018) is a general-purpose distributed actor framework that GPU RL systems such as HybridFlow (Sheng et al., 2025) use to place stages on separate workers; here it provides an alternative to both colocation and Pathways’ single controller for TPU RL. Generation uses the previously published policy while the trainer computes the next update, introducing one-step policy staleness. After each update, trainer hosts publish weights to Ray’s object store, and rollout hosts fetch a full copy over the data-center network (DCN). Rollout data sent to the trainer include the generating policy’s log-probabilities for importance weighting. Section 4.4 gives the tested splits and timing tradeoffs. 4

Share of step time

15 19

Step time vs fastest TP (log)

24 36

outside compute

100%

0.6B TP1

1.7B TP1

4B TP2

8B TP4

14B TP4

32B TP4

70B TP8

100B* 200B* 405B TP16 TP16 TP16

0.6B

2.5×

2.0×

1.5×

70B

network

other computational overhead

405B

14B

~100B

TP8

TP16

1.0× TP1

host orchestration

~200B

TP2

TP4

TP32

Rollout tensor parallelism

compute

(A) Step-time breakdown.

(B) Tensor-parallel scaling.

Figure 3 Scaling study on 128 TPU v7x chips, Pathways, decode-heavy workload (8,192 forced output tokens); medians over runs. (A) Share of step time by component at each model’s fastest TP (Table 2). (B) Step time at each rollout TP

relative to the model’s fastest TP. Filled: each model’s fastest TP (Table 2). 14B has no TP16 (its 40 attention heads do not divide by 16), ∼100B ran at TP16 only, and 0.6B–70B were not run at TP32. Values: Table H1.

LoCoJAX adds four optional optimizations (O1–O4; Table B1): (O1) exporting metrics outside the host barrier, (O2) freezing Python garbage collection, (O3) balancing prompts by length, and (O4) publishing all weights to the local engine in one call. Section 4.2 compares Pathways with LoCoJAX base, which uses the compatible optimizations shared with the baseline, and LoCoJAX optimized, which also enables O1–O4.

Experimental Setup

All experiments run on Google TPU v7x. Each chip has two TensorCores, which JAX exposes as separate devices with about 95 GiB of usable HBM each, and each host has four chips. The stack comparison uses single slices ranging from 8 chips (two hosts, 2×2×2) to 64 chips (16 hosts, 4×4×4). The scale study uses 128-chip slices, and the colocated-vs-disaggregated comparison (Section 4.4) extends to 128 and 256 chips. Pathways jobs additionally use a CPU node to host the single controller. Pathways recipe (baseline). Pathways (Barham et al., 2022) is the runtime used by Google’s documented multi-

host RL recipes on Cloud TPU: Tunix and MaxText run multi-host training on Pathways on Cloud (Google Cloud, 2026d,a,b). Our baseline is a GRPO recipe for Pathways that we assembled from public components: a MaxText and Tunix trainer (Google, 2023; Bao et al., 2025) and a vLLM TPU sampler (Kwon et al., 2023; vLLM Project, 2025). Both are built from pinned public main-branch commits (MaxText 3258182f with PR #5289 (Appendix D), Tunix d6893d5e, tpu-inference 518d1aaa, vLLM 51da0ca6; JAX 0.11.0 (Bradbury et al., 2018)). The trainer and sampler are colocated on the same devices, use bf16 trainer weights as in MaxText’s default RL configuration, and are driven by one head process as a single SPMD program. Rollout runs at TP1 with data parallelism over all devices (DP = 2 × chips), and training uses FSDP. We froze the recipe before our experiments. To match work, we use two passes over the same batches, matched sampling, and extra sequence-length logging. The Pathways arm of Table 3 also runs every optimization compatible with both stacks.

3.1

Workloads

Scale study. To measure where time goes as model size changes, every cell runs GRPO on OpenMathInstruct2 (Toshniwal et al., 2024) prompts with 256 prompts × 8 generations = 2,048 completions per step, each forced to 8,192 tokens, using the Pathways scheduler. We chose a very decode-heavy workload because RL for reasoning generates long chains of thought, so most rollout time goes to decode; forcing every completion to 8,192 tokens emulates those long traces, and with prompts of about 150 tokens on average, prefill is negligible. Fixing the output length makes the work per step identical across models. Models are Qwen3-0.6B

Table 2 Step time and its breakdown at each model’s fastest TP, 128 TPU v7x chips, decode-heavy workload (8,192

forced output tokens). Computational overhead is the share of the step outside compute; host and network are its two largest parts. ≈ marks a tie within run-to-run spread; values are for the first-named TP. Medians over 3 runs (∼100B: 2; 405B: 1, a workaround extension); network shares are upper-end estimates. Model

Fastest TP

Qwen3-0.6B TP1 ≈ TP2 Qwen3-1.7B TP1 ≈ TP2 Qwen3-4B TP2 Qwen3-8B TP4 ≈ TP2 Qwen3-14B TP4 Qwen3-32B TP4 ≈ TP8 Llama-3.1-70B TP8 Synthetic ∼100B TP16 only Synthetic ∼200B TP16 Llama-3.1-405B TP16

Step (s) Computational Overhead 68.6 79.6 108.5 134.1 168.9 301.3 425.6 382.9 729.3 3,334.0

Host Network

40.6% 36.4% 29.8% 24.3% 32.2% 17.6% 39.6% 18.3% 34.2% 14.8% 29.2% 8.3% 28.3% 7.0% 33.6% 5.5% 31.8% 4.4% 40.3% 1.4%

2.2% 3.5% 12.2% 18.5% 17.4% 18.6% 20.0% 26.4% 25.9% 38.5%

to 32B (Yang et al., 2025), Llama-3.1-70B, two synthetic dense models with Llama-3.1-405B width (∼100B with 31 layers and ∼200B with 62), and Llama-3.1-405B (Grattafiori et al., 2024); the Llama-3.1 models contribute architecture configurations only. Each model runs on Pathways at 128 chips at every TP from the smallest that fits up to TP8, plus TP16 (not possible for 14B, whose attention heads do not divide by 16); ∼100B runs only at TP16, and ∼200B and 405B run at TP16 and TP32. Weights are randomly initialized. Pathways and LoCoJAX are compared on Qwen3-0.6B over three public datasets, GSM8K (Cobbe et al., 2021), NuminaMath-CoT (Li et al., 2024) and SWE-agent trajectories (Nebius, 2024; Yang et al., 2024), chosen to vary where work falls (Table 1). Each run covers the same 20 batches twice (40 steps) and scores the second pass, so both stacks see the same prompts in the same order; we verify this at run time from both logs. On GSM8K we use bf16 trainer weights, which is also MaxText’s default RL configuration. On NuminaMath and SWE-agent we keep our stack’s default fp32 trainer weights, since both Pathways and LoCoJAX already produce matching output lengths there. Stack comparison.

3.2

Measurement

Step time. Step time is the median over steady-state steps, read from each stack’s own clock. Compile time

and end-to-end wall-clock time are not part of the comparison.

Time breakdown. For the scale study we split one profiled step into compute, network and host orchestration. The network includes TP collectives, other collectives, host-offload copies, the DP all-reduce, and time spent waiting at the DP check and for TP partners. Host orchestration is host-side work outside device programs: controller round trips, scheduling and the host side of weight sync. Computational overhead is the share of the step outside compute (1 − compute/step). Device-level TP comparison. Pathways step time at small model sizes moves by several seconds with the

placement of its head node, which is as large as most TP effects there. We therefore compare TP degrees at the device level, in a separate 64-chip TP sweep on the same workload, using decode device time per step from xprof traces of one host, which is insensitive to head placement.

3.3

Comparison Protocol

Appendix A gives the full measurement protocol and verdict rules. Equal work. A pair counts only if both runs process the same prompts (per-step mean prompt length within

0.5 tokens on all 40 steps), their mean and longest output lengths agree within a factor of [0.8, 1.25], and their truncation shares differ by at most 10 points. Output length matters because rollout time is set by the longest response. Statistics. We report the speedup S = TPathways /Tours , and the same ratio per processed token. Confidence

intervals are 95% percentile bootstraps over steady steps (10,000 resamples). We call our stack faster only if the lower bound of both ratios exceeds 1.111, a margin that covers the roughly 10% shift the Pathways head node can cause on its own. GSM8K Pathways cells at 16–64 chips therefore use heads from one reference pool of CPU head hosts (the 8-chip head ran elsewhere). 6

Results

4.1

Scaling Study

We ran GRPO on Pathways for dense models from 0.6B to 405B on 128 TPU v7x chips, at each model’s fastest TP (Table 2, Figure 3A). What changes with model size is where that time goes. Host orchestration takes 36% of the step at 0.6B and falls steadily to 1% at 405B. Network does the opposite, rising from 2% at 0.6B to 26% at ∼200B and 39% at 405B. Larger models need more tensor parallelism, and more TP means more collective traffic, so the fastest TP grows from TP1 at 0.6B to TP16 from ∼100B up (Figure 3B). Small models do not improve from going wider: in a separate 64-chip sweep, moving 0.6B from TP1 to TP2 saves 0.21 ms of weight streaming per decode iteration but adds 0.7–0.9 ms of collectives, while at 4B the saving grows to 1.16 ms and TP2 wins. Small models therefore stay at TP1, where they lose almost nothing for the network, and their computational overhead is almost all host. The two ends of the model range need different fixes. Large models need high TP, so their levers are the choice of TP and computationally cheaper collectives. For small models the computational cost is host orchestration: controller round trips, request admission and scheduling, and the host side of weight sync. LoCoJAX targets that computational cost. It keeps TP1 rollout and takes the single controller out of the loop, giving each host its own controller and rollout engine. Section 4.2 compares it with Pathways on Qwen3-0.6B.

4.2

LoCoJAX vs. Pathways

LoCoJAX takes less time per step than Pathways on every dataset and chip count we ran (Table 3, Figure 4). LoCoJAX base is 1.95–2.97× faster on short-output GSM8K, 1.23–1.86× on decode-heavy NuminaMath and 1.64–2.05× on prefill-heavy SWE. LoCoJAX optimized adds O1–O4 and is 2.13–3.56×, 1.25–1.89× and 1.72–2.32× faster. The gap is largest on GSM8K, where each step does the least work, which fits Section 4.1; there LoCoJAX on 8 chips takes less time per step than Pathways on 64, even though the smaller run generates each batch in two waves. Our trainer does not accumulate gradients, so its training memory grows as chips are removed: it does not fit SWE at 8 chips (it needs 157.6 GiB per device), where Pathways still runs. Model size. LoCoJAX’s advantage broadly shrinks as models grow, consistent with the falling host share in

Figure 1 (Tables E1 and E2). On GSM8K at 64 chips, Pathways takes 3.25× the step time of LoCoJAX at 1.7B, 1.72× at 8B, 1.88× at 14B and 1.58× at 32B; on NuminaMath at 128 chips, it takes 1.88× at 1.7B, 1.51× at 4B and 1.30× at 8B. LoCoJAX optimizations. As part of the LoCoJAX protocol, we found four additional host-side optimizations (Table B1). We call LoCoJAX with these optimizations LoCoJAX optimized and describe their operation in Appendix B.

4.3

What Coordination Explains

Figure 5 isolates what the Pathways controller adds to our jobs. Before any engine runs, the Pathways head tokenizes the batch and adds all 2,048 requests one at a time, about 1.0 s per step (timers off). This is serial work in one process; it stays at 0.88–0.98 s from 8 to 64 chips (Figure 5a) while the rest of the step shrinks, and its share of the Pathways step nearly doubles, from 8% to 14% (Figure 5b). What it scales with is the batch, so scaling out with a larger batch makes it worse, not better. LoCoJAX splits the same work across hosts: at 64 chips, each of the 16 hosts admits its own 128 requests, in parallel, in under 0.06 s (Appendix C, tradeoff 2). Meanwhile, Pathways spends the equivalent of 48% of a whole LoCoJAX step before its first engine starts. The head interval is not all of the controller computational cost. The controller impacts Pathways twice (Figure 7): once per step, when the head admits the batch (0.94 s with its other work before the engine call), and once per decode iteration. Each Pathways decode iteration takes 2.15 ms against 0.90 ms on LoCoJAX, and over the roughly 1,800 iterations of a step that premium adds 2.25 s, the largest term of the 64-chip gap. This matches the Pathways structure in Figure 2: Pathways runs every iteration as one lockstep program across all chips, while each LoCoJAX host steps its own engine. 7

Pathways

GSM8K (short outputs) Step time (s, log scale)

10.82 8.56

6.75 5.56

5.07

100

3.32

200

117.3 154.7 152.2 89.4

100 64.2 70.3 45.5

2.51x

44.6 2.36

2.03

3.32x

16 32 64 TPU v7x chips

34.7

30 28.4

1.86x

65.2

61.1

52.8

28.0 1.89x

20 8

163.1 105.0

84.1 2.69

SWE (prefill-heavy)

190.3

4.08 2.83

LoCoJAX optimized

NuminaMath (decode-heavy) 200

8.41

LoCoJAX base

31.6

39.4

19.2 2.05x 17.0 2.32x

16 32 64 TPU v7x chips

Bold: Pathways / LoCoJAX step time at 64 chips.

Figure 4 Step time vs TPU v7x chips (log–log), Qwen3-0.6B: Pathways, LoCoJAX base, and LoCoJAX optimized, which adds optimizations O1–O4 (Table B1). Values and speedups: Table 3. Bold: ratio of Pathways to each LoCoJAX line at 64 chips.

Table 3 Steady-state step time, Pathways vs colocated LoCoJAX base and optimized, Qwen3-0.6B; speedup over

Pathways in parentheses. LoCoJAX base runs the shared optimizations of the Pathways arm; LoCoJAX optimized adds O1–O4 (Table B1). At NuminaMath 16 chips, optimized is 7% slower than base in same-window pairs too. LoCoJAX does not fit SWE at 8 chips (needs 157.6 of 94.7 GiB per device).

4.4

Dataset

Chips

Pathways (s)

LoCoJAX base (s)

LoCoJAX optimized (s)

GSM8K (short outputs)

8 16 32 64

10.82 8.56 8.41 6.75

5.56 (1.95x) 4.08 (2.10x) 2.83 (2.97x) 2.69 (2.51x)

5.07 (2.13x) 3.32 (2.58x) 2.36 (3.56x) 2.03 (3.32x)

NuminaMath (decode-heavy)

8 16 32 64

190.3 117.3 70.3 52.8

154.7 (1.23x) 84.1 (1.39x) 45.5 (1.55x) 28.4 (1.86x)

152.2 (1.25x) 89.4 (1.31x) 44.6 (1.58x) 28.0 (1.89x)

SWE (prefill-heavy)

8 16 32 64

163.1 105.0 65.2 39.4

N/A (no fit) 64.2 (1.64x) 34.7 (1.88x) 19.2 (2.05x)

N/A (no fit) 61.1 (1.72x) 31.6 (2.06x) 17.0 (2.32x)

Colocated vs. Disaggregated

We compare LoCoJAX-CO with LoCoJAX-DA (Section 2; Figure 6; Table F1). For every cell we run LoCoJAX-DA at three static splits that give the trainer 50%, 25% and 12.5% of the hosts, three launches each, and report all of them; a split is missing only where the host count, our Ray launcher or trainer memory rules it out (Appendix F). LoCoJAX-DA wins only where one phase dominates the step, and the three workloads show three ways this plays out. GSM8K: rollout still scales at small sizes. At 8 and 16 chips LoCoJAX-DA is slower (1.61× and 1.13×) because

rollout still speeds up with every host it gets, so moving hosts to the trainer requires more than overlapping training saves (2.9 vs 1.1 s at 8 chips). At 128 chips this reverses (0.90×): LoCoJAX-CO stops scaling (2.02 to 1.73 s from 64 to 128 chips, no improvement at 256) as the slowest rollout host falls further behind the average (38% to 55% of the step), and the one-step-stale schedule lets LoCoJAX-DA’s rollout hosts skip that wait. NuminaMath: decode-heavy, so rollout needs the hosts. Long outputs make rollout most of the step. LoCoJAX-

Pathways step time

LoCoJAX step time (optimized)

(a) Seconds per step

Seconds

8.56

8.41

5.07 3.32 2.36

2 1

6.75

0.88

2.03

0.98

16 32 TPU v7x chips

0.92

Share of Pathways step

10.82

(b) Head interval / Pathways step

20%

0.5

Pathways head interval

14.4%

15% 10.7%

10%

11.6%

8.1%

5% 0%

16 32 TPU v7x chips

Figure 5 Pathways head interval vs step time, GSM8K, Qwen3-0.6B, 2,048 requests, timers off. (a) Step time of

Pathways and of LoCoJAX optimized (Table 3), and the Pathways head interval: time in the head process from the start of the global step to the opening of the generation progress bar, which contains the open-source request admission. (b) Head interval as a share of the Pathways step. Head intervals are means of two launch medians over steps 3–18, read from the Pathways launches behind Figure 4. Pathways heads ran in the reference node group at 16–64 chips and in a different node group at 8 chips. LoCoJAX generates in two waves at 8 chips.

DA loses at 32 chips when half the hosts generate (1.26×), and wins at 64 and 256, where three quarters do (0.95×, narrow, and 0.85×). At 128 chips, with an even split, it is unresolved; a rollout-heavier split tried afterwards reads 0.87× (provisional; Appendix F). SWE: prefill-heavy, so the trainer is the bottleneck. Long prompts and short outputs make training on the long

packed sequences the larger phase. LoCoJAX-DA gives the trainer only half the hosts, so that phase takes longer, and there is little rollout to overlap it with. LoCoJAX-CO wins at every size (1.14–1.30×; 32 chips added post hoc), and every split that moved hosts from training to rollout was slower still (Appendix F). LoCoJAX-DA was not run on NuminaMath at 16 chips, where LoCoJAX-CO runs, and a 2-host SWE trainer runs out of memory: with fewer hosts, the trainer holds more model and optimizer state per chip, and LoCoJAX’s trainer does not accumulate gradients, so the whole packed batch must fit at once. These are timing results: on GSM8K, training reward with LoCoJAX-DA stays within our pre-registered 0.005 band of LoCoJAX-CO at 64 and 128 chips, and in an earlier, longer setup both LoCoJAX-DA splits at 32 chips failed a long-run quality gate. Colocation is the better default; disaggregation is beneficial when rollout dominates the step and one step of staleness is acceptable.

Related Work

RL post-training optimization. Most RL systems work targets larger LLMs on GPUs, where each stage is

computationally expensive and the question is how to place and overlap them. Approaches include reallocating parallelism across stages (Mei et al., 2025; Sheng et al., 2025), fusing stages to remove pipeline bubbles (Zhong et al., 2025b), disaggregating generation from training (Zhong et al., 2025a), removing per-layer synchronization barriers under sequence-length imbalance (Wan et al., 2026), and choosing placement per workload (Yu et al., 2026). Closest to ours, RL in the Wild finds that serving-derived rollout engines add CPU scheduling computational overhead and that single-controller data paths bottleneck training (Zhou et al., 2025). All of these study GPUs at 1.5B parameters and above. We study small models on TPUs, where the documented recipe places a single controller on the rollout path, and find that removing it yields substantial improvements. Small-model RL. Small models are a common post-training target, trained with RL as computationally cheap

reasoners (Dang and Ngo, 2025; Wang et al., 2026; Zhuang et al., 2026) or examined at the low end of scaling studies (Zeng et al., 2025). Some use on-policy distillation as students (Zhong et al., 2026). These studies locate the bottlenecks in learning: optimization instability and length limits in long training (Dang and Ngo, 2025), 9

LoCoJAX-CO

LoCoJAX-DA, 50% trainer hosts

GSM8K (short outputs) 10

8.20

LoCoJAX-DA, 25% trainer hosts

NuminaMath (decode-heavy)

Step time (s)

89.4

3.75 2.87

2.56 2.63

3.33

2.90

30 2.22

1.88 2.02 2.38 1.98 1.80 1.77 1.73 1.94 1.62 1.56

56.1

26.5

21.5 21.0

16.8 11.8

19.6 14.8

16.8 12.1

40.6

20.9

27.9 21.4 17.0

40.2

31.1

25.1

40.4

44.5 34.1 42.0

5.08

SWE (prefill-heavy) 50

100

4.32

LoCoJAX-DA, 12.5% trainer hosts

10.0

14.3

11.8

10 7.2 6.3

1 8

16 32 64 128 256 TPU v7x chips

16 32 64 128 256 TPU v7x chips hollow: provisional cell

Figure 6 Step time vs TPU v7x chips (log–log), LoCoJAX-CO vs LoCoJAX-DA at three trainer/rollout host splits,

Qwen3-0.6B; medians over launches. Hollow markers: provisional cells. Missing LoCoJAX-DA splits are ruled out by host count, our Ray launcher or trainer memory (Appendix F). The LoCoJAX-CO point at NuminaMath 16 chips comes from Figure 4. Median step times per split: Table F1.

limited capacity at 0.5B (Zhuang et al., 2026), rigid format rewards and mismatched data difficulty (Zeng et al., 2025), among others. Systems work on small computations finds execution computational overheads instead (framework dispatch (Fernandez et al., 2023), CPU-side serving and scheduling (vLLM Team, 2024), kernel boundaries (Spector et al., 2025)), but in inference, not RL. For small models, RAPID is the closest work: on a single 4-GPU server with 0.5–1.5B models, it cuts RL running time by 11–34% by sampling large inference batches and taking several importance-weighted off-policy updates on each, reducing switches between inference and training (Huang et al., 2025). LoCoJAX-CO instead keeps GRPO synchronous and on-policy and removes multi-host control computational overheads. The two optimizations target different computational costs and could be combined. TPU and accelerator-native RL. In classical deep RL, SEED RL centralizes inference on accelerators (Espeholt

et al., 2020), and Podracer either compiles the whole loop into one XLA program or splits TPU cores between actors and learners (Hessel et al., 2021). Both treat host-device coordination as the main computational cost, but for small networks without autoregressive decoding, KV caches or weight publication. TPU LLM RL is recent. Google’s public recipes run GRPO through MaxText or Tunix on Pathways (Google Cloud, 2026d; Google, 2026c,a; Barham et al., 2022); Tunix’s agentic stack adds asynchronous rollouts to keep TPUs busy during environment calls (Gao et al., 2026). Concurrently, Tunix added a distributed RL runtime with separate trainer and rollout workers (Google, 2026b), which we did not evaluate. Our work was inspired by Tunix RFC #706, which proposed multi-controller JAX for post-training but was not implemented (Wang, 2025). LoCoJAX applies it to the full RL loop by keeping variable-length generation local. To our knowledge, no published work studies sub-billion LLM RL on TPU or compares single- and multi-controller execution for it.

Discussion

Small-model RL on TPU is bound less by compute than by what surrounds it. Our characterization shows that at TP1, 41% of the Qwen3-0.6B Pathways step is outside compute: 36% host orchestration and 2% network (Section 4.1). LoCoJAX optimizes the host side; LoCoJAX optimized is 2.1–3.6× faster than Pathways on GSM8K, 1.3–1.9× on NuminaMath and 1.7–2.3× on SWE (Section 4.2). These results show the advantage of host-local rollouts over a centralized Pathways controller for small-model RL. Our work draws inspiration from an unimplemented Tunix RFC (Wang, 2025) calling for multi-controller JAX for Tunix. This paper implements the RFC for RL and adds evidence about when it matters and how it 10

speeds up small models. We show that the Pathways scheduler spends excessive time on host orchestration with small models. A single controller adds time to a small-model RL step in two places: once per step when it admits the batch and once per decode iteration (Figure 7). The Pathways head spends about 1 s per step admitting the batch, a computational cost that stays flat from 8 to 64 chips while the rest of the step time shrinks, so its share of the Pathways step nearly doubles, from 8% to 14% (Figure 5). Every result rests on public models, datasets and frameworks, so others can replicate our results. We note that our Pathways baseline is a tuned one. This was because we found no peer-reviewed results for small-model RL on TPUs. The only other TPU baselines are Google’s public reference recipes, which we run with only the deviations listed in Appendix G, without these optimizations. Against them, on identical workloads, colocated LoCoJAX is 5.3× faster per step on Qwen3-1.7B with Tunix (Google, 2026c) and 2.1× on Gemma4-26B MoE with MaxText (Google, 2026a) (Appendix G). Our contributions also highlight optimization choices for future RL workflows. Three choices mattered. (i) Keep per-request work off a single controller. Work that one process does for every request grows with the batch and does not shrink with chips; give each host its own controller and rollout engine (Section 4.3). (ii) Choose TP by model size, not by the smallest TP that fits. At 0.6B, TP2’s weight-bandwidth saving cannot cover its collective penalty, so TP1 wins; but from 4B up, higher TP wins (Section 4.1, Figure 3B). (iii) Colocate by default. Disaggregate only when one phase dominates the step and the split gives that phase enough hosts (Section 4.4). Our disaggregation recommendation gives further color to existing literature on how RL work should be placed. GPU systems for large models increasingly disaggregate generation from training and run them asynchronously (Zhong et al., 2025a; Fu et al., 2025) or place stages across heterogeneous devices (He et al., 2026); colocated designs instead time-share the same devices across phases (Sheng et al., 2025; Mei et al., 2025; Zhong et al., 2025b), and RLinf finds that the best placement depends on the workload (Yu et al., 2026). To our knowledge, we give the first such comparison for small-model RL on TPU, and it yields concrete evidence that disaggregation is beneficial only when one phase dominates the step and the split gives that phase enough hosts. At small scale rollout still speeds up with every host it gets, so colocation wins (LoCoJAX-DA takes 1.61× the colocated step on GSM8K at 8 chips). At large scale the workload decides, and so does the split. On decode-heavy NuminaMath, rollout dominates. With three quarters of the hosts on rollout, LoCoJAX-DA beats colocation at 64 and 256 chips and is faster without a verdict at 32 and, post hoc, 128, while an even split never does. On prefill-heavy SWE, training dominates and colocation wins at every size: an even split already leaves the trainer half the hosts, and moving more of them to rollout makes LoCoJAX-DA 1.9–4.1× slower. Thus, for small models, colocation is the right default; disaggregation is beneficial on decode-heavy work, but only with a split that gives rollout most of the hosts. LoCoJAX shows that scheduler defaults built for large SPMD programs can be a poor fit for small-model RL. Our scaling study shows that for small-model RL on TPUs, reducing host-orchestration computational cost is important. In fact, Pathways’ per-step excess over LoCoJAX on GSM8K stays at about 4.1–5.6 s from 8 to 64 chips (Table 3; the head’s serial admission is about 1 s of it) and does not fall systematically as chips are added, so LoCoJAX’s speedup over Pathways on GSM8K grows from 8 to 32 chips as its own step shrinks. Thus, by removing Pathways, we were able to scale our RL jobs more effectively. In summary, LoCoJAX shows why scheduler-side bottlenecks deserve direct attention when training small models.

6.1

Limitations

The Pathways runtime is closed source and we were unable to tune it the way we tuned LoCoJAX. The characterization study covers a decode-heavy workload only; behavior may differ on prefill-heavy datasets. The characterization and the stack comparison use different setups and are not directly comparable. On GSM8K we use bf16 weights, and on NuminaMath and SWE-agent we keep our stack’s default fp32 trainer weights, since both Pathways and LoCoJAX already produce matching output lengths there. However, some discrepancy in speedup factor may be attributable to the precision difference. LoCoJAX was not run on NuminaMath or SWE-agent configurations where we estimated it would not fit, because it does not accumulate gradients. Finally, LoCoJAX is rigorously evaluated on only one model (Qwen3-0.6B). We did run others (Table E1), but tuning them as heavily as Qwen3-0.6B was infeasible.

References T. Bao, J. Carpenter, L. Chai, H. Gao, Y. Jiang, S. Noghabi, et al. Tunix (Tune-in-JAX). Open-source software, 2025. https://github.com/google/tunix. Accessed 2026-10. P. Barham, A. Chowdhery, J. Dean, S. Ghemawat, S. Hand, D. Hurt, et al. Pathways: Asynchronous Distributed Dataflow for ML. In MLSys, 2022. arXiv:2203.12533. C. Basile, M. Voznesensky, G. Park, B. Wilson, P. Holloway, Z. Wan, et al. TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU. arXiv:2610.24161, 2026. P. Belcak, G. Heinrich, S. Diao, Y. Fu, X. Dong, S. Muralidharan, Y. C. Lin, and P. Molchanov. Small Language Models are the Future of Agentic AI. arXiv:2506.02153, 2025. James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs. https://github.com/jax-ml/jax, 2018. P. F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei. Deep Reinforcement Learning from Human Preferences. In NeurIPS, 2017. arXiv:1706.03741. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training Verifiers to Solve Math Word Problems. arXiv:2110.14168, 2021. Q.-A. Dang and C. Ngo. Reinforcement Learning for Reasoning in Small LLMs: What Works and What Doesn’t. arXiv:2503.16219, 2025. DeepSeek-AI, D. Guo, D. Yang, H. Zhang, J. Song, P. Wang, et al. DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948, 2025. L. Espeholt, R. Marinier, P. Stanczyk, K. Wang, and M. Michalski. SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference. In ICLR, 2020. arXiv:1910.06591. Jared Fernandez, Jacob Kahn, Clara Na, Yonatan Bisk, and Emma Strubell. The Framework Tax: Disparities Between Inference Efficiency in NLP Research and Deployment. In EMNLP, 2023. W. Fu, J. Gao, X. Shen, C. Zhu, Z. Mei, C. He, et al. AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning. In NeurIPS, 2025. arXiv:2505.24298. Haoyu Gao, Lance Wang, Shadi Noghabi, Tianshu Bao, and Weiren Yu. Scaling Agentic RL: High-Throughput Agentic Training with Tunix. Google Developers Blog, July 2026. https://developers.googleblog.com/scaling-agentic-rlhigh-throughput-agentic-training-with-tunix/. Accessed 2026-10. Gemma Team. Gemma 4 Technical Report. arXiv:2607.02770, 2026. Google. MaxText: A Simple, Performant, and Scalable JAX LLM. Open-source software, 2023. https://github.com/AIHypercomputer/maxtext. Accessed 2026-10. Google. MaxText GRPO recipe for Gemma 4 26B (run_gemma4_26b_rl.sh). MaxText repository, 2026a. https: //github.com/AI-Hypercomputer/maxtext/blob/3258182fb3570822bb118823621918c677ba9f50/src/maxtext/tr ainers/post_train/rl/scripts/run_gemma4_26b_rl.sh. Accessed 2026-10. Google. Tunix distributed RL runtime (dist_rl). Tunix repository, commits ec932ea1d and ba17165f0, 2026b. https://github.com/google/tunix/tree/ba17165f00f de2b25aaf304850afa4d80e161fe2/tunix/dist_rl. Added October 22–23, 2026; accessed 2026-10. Google. Tunix GRPO recipe for Qwen3 (run_qwen3.sh). Tunix repository, 2026c. https://github.com/google/tunix/ blob/d6893d5ee97a91be6e0027efa048b9c0baf2844f/examples/rl/grpo/gsm8k/run_qwen3.sh. Accessed 2026-10. Google Cloud. Building production AI on Cloud TPUs with JAX. Cloud TPU documentation, 2026a. https: //docs.cloud.google.com/tpu/docs/jax-ai-stack. Accessed 2026-10. Google Cloud. Introduction to Pathways on Cloud. AI Hypercomputer documentation, 2026b. https://docs.cloud.goo gle.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro. Accessed 2026-10. Google Cloud. Port JAX workloads to Pathways. AI Hypercomputer documentation, 2026c. https://docs.cloud.googl e.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/porting-jax-workloads.

Google Cloud. Run multi-host RL training for Qwen3-30b-a3b on TPU v6e. AI Hypercomputer documentation, 2026d. https://docs.cloud.google.com/ai-hypercomputer/docs/tutorials/tpu/run-qwen3-30b-a3b-rl-maxtext. Accessed 2026-10. Aaron Grattafiori et al. The Llama 3 Herd of Models. arXiv:2407.21783, 2024. Y. He, S. Zhang, J. Gai, X. Zhang, B. Han, B. Wang, H. Rangwala, and G. Karypis. HetRL: Efficient Reinforcement Learning for LLMs in Heterogeneous Environments. In MLSys, 2026. arXiv:2512.12476. M. Hessel, M. Kroiss, A. Clark, I. Kemaev, J. Quan, T. Keck, et al. Podracer Architectures for Scalable Reinforcement Learning. arXiv:2104.06272, 2021. L. Huang, S. Anupam, I. Lee, S. Li, and O. Bastani. RAPID: An Efficient Reinforcement Learning Algorithm for Small Language Models. arXiv:2510.03515, 2025. JAX Authors. Introduction to Multi-Controller JAX (aka Multi-Process/Multi-Host JAX). JAX documentation, 2026a. https://docs.jax.dev/en/latest/multi_process.html. Accessed 2026-10. JAX Authors. Pallas: a JAX kernel language. JAX documentation, 2026b. https://docs.jax.dev/en/latest/pallas/in dex.html. J. Jiang, Y. Chen, B. A. Hechtman, F. Zhang, and Y. Mu. Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU. arXiv:2604.15464, 2026. Norman P. Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, Cliff Young, Xiang Zhou, Zongwei Zhou, and David Patterson. TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings. In ISCA, 2023. Sameer Kumar, James Bradbury, Cliff Young, Yu Emma Wang, Anselm Levskaya, Blake Hechtman, Dehao Chen, HyoukJoong Lee, Mehmet Deveci, Naveen Kumar, Pankaj Kanwar, Shibo Wang, Skye Wanderman-Milne, Steve Lacy, Tao Wang, Tayo Oguntebi, Yazhou Zu, Yuanzhong Xu, and Andy Swing. Exploring the Limits of Concurrency in ML Training on Google TPUs. In MLSys, 2021. W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. In SOSP, 2023. arXiv:2309.06180. N. Lambert, J. Morrison, V. Pyatkin, S. Huang, H. Ivison, F. Brahman, et al. Tülu 3: Pushing Frontiers in Open Language Model Post-Training. arXiv:2411.15124, 2024. Mark Lee et al. AXLearn: Modular Large Model Training on Heterogeneous Infrastructure. arXiv:2507.05411, 2025. Jia Li, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu. NuminaMath. Hugging Face dataset AI-MO/NuminaMath-CoT, 2024. https://huggingface.co/datasets/AIMO/NuminaMath-CoT. Z. Liu, C. Zhao, F. Iandola, C. Lai, Y. Tian, I. Fedorov, Y. Xiong, E. Chang, Y. Shi, R. Krishnamoorthi, L. Lai, and V. Chandra. MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases. In ICML, 2024. arXiv:2402.14905. Z. Mei, W. Fu, K. Li, G. Wang, H. Zhang, and Y. Wu. ReaL: Efficient RLHF Training of Large Language Models with Parameter Reallocation. In MLSys, 2025. arXiv:2406.14088. P. Moritz, R. Nishihara, S. Wang, A. Tumanov, R. Liaw, E. Liang, et al. Ray: A Distributed Framework for Emerging AI Applications. In OSDI, 2018. arXiv:1712.05889. Nebius. SWE-agent-trajectories. Hugging Face dataset nebius/SWE-agent-trajectories, 2024. https://huggingface.co /datasets/nebius/SWE-agent-trajectories. L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, et al. Training Language Models to Follow Instructions with Human Feedback. In NeurIPS, 2022. arXiv:2203.02155. Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, X. Bi, et al. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300, 2024. Gerald Shen, Zhilin Wang, Olivier Delalleau, Jiaqi Zeng, Yi Dong, Daniel Egert, et al. NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment. In COLM, 2024.

G. Sheng, C. Zhang, Z. Ye, X. Wu, W. Zhang, R. Zhang, et al. HybridFlow: A Flexible and Efficient RLHF Framework. In EuroSys, 2025. arXiv:2409.19256. Benjamin Spector, Jordan Juravsky, Stuart Sul, Owen Dugan, Dylan Lim, Dan Fu, Simran Arora, and Christopher Ré. Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B. Hazy Research Blog, 2025. https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles. Shubham Toshniwal, Wei Du, Ivan Moshkov, Branislav Kisacanin, Alexan Ayrapetyan, and Igor Gitman. OpenMathInstruct-2: Accelerating AI for Math with Massive Open-Source Instruction Data. arXiv:2410.01560, 2024. vLLM Project. tpu-inference: TPU Backend for vLLM. Open-source software, 2025. https://github.com/vllmproject/tpu-inference. Accessed 2026-10. vLLM Team. vLLM v0.6.0: 2.7x Throughput Improvement and 5x Latency Reduction. vLLM Blog, 2024. https: //vllm.ai/blog/2024-09-05-perf-update. vLLM Team. vLLM TPU: A New Unified Backend Supporting PyTorch and JAX on TPU. vLLM Blog, 2025. https://vllm.ai/blog/2025-10-16-vllm-tpu. X. Wan, P. Qi, G. Huang, C. Ruan, M. Lin, and J. Li. Revisiting Parameter Server in LLM Post-Training. In ICLR, 2026. arXiv:2601.19362. L. Wang. RFC: Multi-Controller JAX (McJAX) Support for Tunix. GitHub discussion #706, google/tunix, 2025. https://github.com/google/tunix/discussions/706. Accessed 2026-10. S. Wang, J. Asilis, Ö. F. Akgül, E. B. Bilgin, O. Liu, and W. Neiswanger. Tina: Tiny Reasoning Models via LoRA. In ICLR, 2026. arXiv:2504.15777. An Yang et al. Qwen3 Technical Report. arXiv:2505.09388, 2025. John Yang, Carlos E. Jimenez, Alexander Wettig, Kilian Lieret, Shunyu Yao, Karthik Narasimhan, and Ofir Press. SWE-agent: Agent-computer interfaces enable automated software engineering. In Advances in Neural Information Processing Systems (NeurIPS), 2024. C. Yu, Y. Wang, Z. Guo, H. Lin, S. Xu, H. Zang, et al. RLinf: Flexible and Efficient Large-Scale Reinforcement Learning via Macro-to-Micro Flow Transformation. In OSDI, pages 829–846, 2026. arXiv:2509.15965. W. Zeng, Y. Huang, Q. Liu, W. Liu, K. He, Z. Ma, et al. SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models in the Wild. In COLM, 2025. arXiv:2503.18892. Q. Zhong, M. Zheng, M. Song, X. Lin, J. Sun, H. Jiang, et al. SOD: Step-wise On-policy Distillation for Small Language Model Agents. arXiv:2605.07725, 2026. Y. Zhong, Z. Zhang, X. Song, H. Hu, C. Jin, B. Wu, et al. StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream Generation. arXiv:2504.15930, 2025a. Y. Zhong, Z. Zhang, B. Wu, S. Liu, Y. Chen, C. Wan, et al. Optimizing RLHF Training for Large Language Models with Stage Fusion. In NSDI, 2025b. arXiv:2409.13221. J. Zhou, Q. Hu, Y. Jin, Z. Wang, P. Sun, Y. Gu, W. Zhang, M. Zhai, X. Zhang, and W. Zhang. RL in the Wild: Characterizing RLVR Training in LLM Deployment. arXiv:2509.25279, 2025. X. Zhuang, P. Ma, Z. Jia, Z. Cao, and S. Liu. Effective Learning for Small Reasoning Models: An Empirical Study on 0.5B Reasoning LLMs. In CPAL, 2026. arXiv:2506.13404.

AI Use Statement The investigators designed the technical approach, the research methodology, and the experimental plan. Claude Opus 5.5 assisted with writing the paper text, implementing code, and writing scripts to collect results. All results reported in this paper were validated by the investigators independently of AI assistance.

Appendix A

Measurement Protocol and Statistics

Step time. A launch’s step time is its median over steady steps: the second pass (steps 23–38) of a 40-step run

that covers the same 20 batches twice. Start-up and compilation are excluded. Where scored steps recompile, we also report the ratio with compile time charged back. A cell averages its launches; ratios are geometric means over launches, or medians of the per-launch comparisons where a table says so. Launches. Each arm has three launches unless a trial is counted as provisional, in which case less than 3 were

run.

Pathways vs LoCoJAX. For each LoCoJAX launch we bootstrap the step-time ratio against the Pathways

arm over steady steps (95% percentile interval, 10,000 resamples), per step and per token. “LoCoJAX faster” requires both lower limits above 1.111 in every comparison, equal work (Section 3.3) and three launches per arm. Within LoCoJAX. LoCoJAX-DA against LoCoJAX-CO uses the 95% Welch interval on log launch values, and

on NuminaMath and SWE also the paired-step bootstrap, which must fall on the same side of 1. Ablation rungs instead require disjoint launch ranges and a ratio that differs from 1 by more than either arm’s relative spread.

LoCoJAX Optimizations

LoCoJAX optimized as evaluated in Table 3 adds 4 optional optimizations to LoCoJAX base, which already runs the shared optimizations (Table B1). Three act on LoCoJAX’s own per-host learner and the fourth on weight publication.

Critical-Path Accounting

This appendix gives our critical-path accounting: the equations and the four tradeoffs it explains. Coordination timings follow Section 4.3. Section 4 reports measured wall times; nothing here is used to predict them. We model Section 4’s steady-state step time as its longest dependent chain: X Tstep (x) = max dv (x) π

(1)

v∈π

where π is a path through configuration x’s acyclic one-step graph, with edges for data dependencies and shared resources (colocated devices, the Pathways head). Durations dv exclude waiting (slack, not computational cost). Decompositions use means; medians do not add. Decode iteration (per TensorCore, TP degree t):

diter (t) = κ(t) + (Qw /t + QKV )/βeff + CTP (t)

(2)

Admission (n requests, one process of stack s):

dadm = α0s + α1s n, or max α0s + α1s nh

when each host h admits its own nh

(3)

Here κ is the fixed per-kernel time, Qw the TP1 weight bytes, QKV the per-core KV bytes (size unmeasured; inferred unchanged as t doubles), βeff the measured marginal HBM rate, CTP the Section 4.1 TP penalty, and α0s , α1s stack s’s fixed and per-request computational costs. Other operations use measured durations.

Compute is omitted (0.02–0.24 ms per iteration at peak FLOP rate, against 0.70–5.27 ms of TP1 weight time). Disaggregation (one-step-stale rollouts; Hr rollout, Ht trainer hosts) and colocation (all H hosts) have steady-state periods (step times) λDA = max{DR (Hr ) + dL , DT (Ht ) + dpub } (4)

− A + τ, λco = DR (H) + DT (H) + dpub

where DR , DT are rollout and training time, dpub weight publication, dL rollout weight load, A ≥ 0 the straggler wait that staleness absorbs, and τ the residual Ray computational overhead (including the unmeasured batch transfer). With rollout always the longer branch, disaggregation is beneficial only if O + A > F + dL + τ , with fragmentation F = DR (Hr ) − DR (H) and hidden training O = DT (H) + dpub . The four tradeoffs:

1. Weight streaming vs TP computational overhead (Eq. 2, Section 4.1, first decode window). TP1 → TP2 saves 0.21 ms per iteration at 0.6B and 1.16 ms at 4B (bytes term at βeff = 3.32 TB/s: 0.18, 1.20 ms), against penalties of 0.68 and 0.77–0.85 ms: TP2 wins at 4B, not 0.6B. 2. Centralized vs host-local admission (Eq. 3, Section 4.3). At 64 chips the Pathways head admits all 2,048 requests serially in about 1.0 s. Each of the 16 LoCoJAX hosts admits its own 128 requests in parallel with the others, in under 0.06 s. 3. Global vs local rollout control (Eq. 1, Section 4.3). Pathways runs every decode iteration as one lockstep program across all chips, while each LoCoJAX host steps its own engine. On GSM8K at 64 chips, Pathways takes 2.15 ms per decode iteration against LoCoJAX’s 0.90 ms; the 1.25 ms premium adds about 2.25 s over roughly 1,800 iterations. 4. Colocated vs disaggregated (Eq. 4, Section 4.4). On GSM8K at 8 chips F = 2.9 s exceeds O = 1.1 s and Ray takes 1.6× the colocated step time (one launch per arm, provisional); at 128 chips O = 0.32 s and A = 0.58 s outweigh F = 0.23 s, dL = 0.49 s and τ = 0.03 s, and Ray is about 10% faster. Scope. This is an accounting of measured step times, not a validated predictor: F and O are derived, A simulated, dL a prior and τ a residual. Unification. One accounting covers both kinds of change. LoCoJAX versus Pathways changes the graph of

Eq. 1: which process serializes per-request work (the Pathways head for all n requests, or each host for its own nh , Eq. 3) and where hosts join (every engine step or only phase boundaries); colocation versus disaggregation changes a resource edge (Eq. 4). The upstream changes mostly keep who serializes and where hosts join. They change node terms instead: the per-iteration terms of Eq. 2, the work a node processes, how many serial calls it issues, local edges and lifecycle nodes within a phase, and which configurations count. Admission and weight-publication slicing (#2232) share the form α1s n on one serializing process, each with its own α. Each change is booked once as its change in Tstep ; changes that move the same term are not added. Like Eqs. 1–4, this explains measured steps and predicts none. Table B1 LoCoJAX optimization levers. Outside this list, both stacks also used LoCoJAX’s sampler, which does not

compute the top-k and top-p masks (unused here); LoCoJAX keeps it with its levers off. Lever

What it changes

Code it acts on

Applies to

Exports metrics off the per-host barrier Freezes Python garbage collection after warmup Balances prompts across hosts by length Publishes all weights to the local engine in one call

LoCoJAX per-host learner LoCoJAX per-host learner LoCoJAX per-host learner LoCoJAX weight publish

LoCoJAX only

O2 O3 O4

LoCoJAX only LoCoJAX only LoCoJAX only

(a) Request admission, once per step about 1.0 s

Pathways under 0.06 s per host

LoCoJAX 0.0

0.5

1.0

1.5

Seconds per step

(b) Decode, once per iteration 2.15 ms / iteration

Pathways 0.90 ms / iteration

LoCoJAX 0

Milliseconds (first 10 of about 1,800 iterations)

Figure 7 The two recurring controller computational costs of Section 4.3, GSM8K, Qwen3-0.6B, 64 chips. (a) Request

admission once per step: the Pathways head admits all 2,048 requests in about 1.0 s, while each LoCoJAX host admits its own 128 requests in under 0.06 s. (b) Ten of about 1,800 decode iterations: each Pathways iteration takes 2.15 ms against 0.90 ms for LoCoJAX, adding about 2.25 s per step.

Upstream Contributions

Method. We profiled an end-to-end RL step, found the critical-path node (Appendix C) whose exposed time was large or growing, traced it to a code path and made a targeted change.

The work falls into three kinds: • Measured or reported benefits: families A–H. • Integration that makes a benefit reachable from MaxText: #3395, #5236, #5255 and #5289. Their numbers are their family’s benefit, not an extra one. • Correctness and observability fixes: family I and the packing crash fix #2309. These are not speedups. Effects are per mechanism, in native units, and not additive. Pairs that move the same term are reported once and never summed: #2231 with #5289, #2086 with #5236, #3102 with #3395, and #2233 with its detokenization tail. What the evaluated runs contain. For each image we checked whether the merge is present, whether its setting

is reachable and set, and whether a log line shows the path ran. The frozen recipe has log evidence for families C, E, F and G, and only configuration or code evidence for A, B, D and H.

Acknowledgments note. We thank the tpu-inference, Tunix and MaxText maintainers for reviewing and merging these changes and for the upstream implementations they build on or that covered ours; Loki Chen authored tpu-inference #3101, #3102 and #3174, and Tianyu Gu authored the other 18 merged PRs in Table D2.

Supporting Details for Section 4.2

Model size. Tables E1 and E2 compare the frozen Pathways recipe with LoCoJAX optimized, each in its

own configuration (on NuminaMath and SWE, also a bf16 against an fp32 trainer), so they are descriptive. LoCoJAX’s advantage shrinks as models grow on every workload, from 1.8–3.3× at 1.7B to 1.0–1.6× at 32B: Pathways adds a roughly fixed time per step, which weighs less as the step grows (Appendix E.1). 8 chips, one wave. At 8 chips LoCoJAX generates in two waves (Table 3). Rerun post hoc in one wave, with

two engine settings changed, LoCoJAX optimized is 1.93× and base 1.84× faster than Table 3’s Pathways arm, which ran in a different window (descriptive). 17

E.1

A Simple Model of the Pathways/LoCoJAX Ratio

Split each step into work that both stacks do, C, and stack-specific time outside it, Os : TLoCo = C + OLoCo and TPW = aC + OPW . Eliminating C gives, with ∆ = OPW − a OLoCo , TPW = a TLoCo + ∆,

R=a+

∆ . TLoCo

(5)

With a = 1, Pathways takes a fixed ∆ longer per step, so the ratio falls as the step grows and rises as LoCoJAX’s step shrinks. We fit ∆ per block of cells and check each cell against the ∆ fitted on the others (Table E3, Figure 8). This accounts for measured step times; it does not predict them. Where it holds. On GSM8K (Qwen3-0.6B, Pathways vs LoCoJAX base), ∆ stays at 4.0–4.8 s from 4 to 128

chips while LoCoJAX’s step falls from 10.06 to 2.54 s, so the ratio rises from 1.44 to 2.58 and flattens above 32 chips only because LoCoJAX’s step stops shrinking. A free fit gives a = 1.03 [0.89, 1.17]. On NuminaMath at 128 chips, ∆ is 20.5–24.4 s from 1.7B to 14B while LoCoJAX’s step grows 2.9×; ∆ falls slightly with size, so a = 1 is approximate (free fit a = 0.92 [0.83, 1.00]).

Where it fails. At 32B the measured ratio is 1.03 against 1.15 predicted, but this cell runs rollout at TP8 and

only LoCoJAX offloads activations, so it does not show the gap closing with model size alone. On SWE, ∆ is 13.3 s at 1.7B but 7.5 s at 4B and 8B, so no single ∆ fits; with outputs of about 140 tokens, SWE’s ∆ cannot come mainly from per-iteration decode computational cost. On the GSM8K size ladder, the model is within 0.15 of the measured ratios from 4B to 32B but over-predicts 1.7B (3.53 vs 3.25).

Colocated vs. Disaggregated: Details

This appendix gives the median step time of every LoCoJAX-CO cell and LoCoJAX-DA split (Table F1), and some detail behind that section. NuminaMath and SWE. Colocated LoCoJAX (LoCoJAX-CO) runs a trainer shard and a rollout engine on

every host and updates synchronously; Ray (LoCoJAX-DA) splits the slice into trainer and rollout hosts that run concurrently. No Ray point exists for NuminaMath at 8–16 chips or SWE at 8–16 chips: colocated LoCoJAX does not fit SWE at 8 chips (Table 3), and Ray was not run on NuminaMath at 16 chips (projected trainer memory about 118 GB). At SWE 32 chips the 4/4 split fits (a memory probe peaked at 91.2 of 94.7 GiB per device), while a 2-host trainer (2/6) runs out of memory (157.6 vs 94.7 GiB). That SWE 32-chip cell was added after the main results, in a later window, and is never pooled across windows: with three launches per arm, colocated takes 31.1 s and the 4/4 split 40.4 s per step, a Ray/colocated ratio of 1.299x [1.296, 1.302; 1.268, 1.318] (median basis; Welch; bootstrap), colocated faster (post hoc).

Comparison with Published TPU Recipes

Baselines and matching. We found no peer-reviewed results for RL post-training of small LLMs on TPUs, so we

compare against the closest public baselines: Google’s reference GRPO recipes for Qwen3-1.7B (Yang et al., 2025) with Tunix (Google, 2026c) and for Gemma4-26B MoE (Gemma Team, 2026) with MaxText (Google, 2026a). We changed none of the recipes’ workload settings. Each LoCoJAX arm reproduces its recipe’s workload: the same checkpoint, tokenizer and prompt template (rendered prompts are byte-identical), the same dataset and per-step prompt order, rewards, sampling (temperature, top-k, top-p, stop tokens), generation cap, prompts × generations per step, and the same GRPO loss with its KL term and estimator. The adapter also matches (LoRA rank 64 for Qwen3, full fine-tuning for Gemma), as do the optimizer and learning-rate schedule. Only the system differs: the controller, placement, rollout engine, weight synchronization, gradientaccumulation layout (one optimizer update per step on both sides), and LoCoJAX’s optimizations. Both sides of each pair run the same JAX, libtpu, vLLM and tpu-inference versions. For Gemma4-26B they also share one container image and MaxText build. For Qwen3-1.7B the recipe keeps its own pinned Tunix commit. Reference Recipe Results On both of Google’s published TPU RL recipes we ran, LoCoJAX takes less time

per training step than the recipe on the same chips. On one host (4 chips), against Tunix’s Qwen3-1.7B recipe 18

Table D1 Upstream changes by mechanism. Effects come from the PR on its author’s workload [PR], our logs [Ours] or

earlier measurements not re-verified [Hist.], and are not additive. Active in: Table 3 images in which the change runs. Per-PR detail: Table D2. Family

Bottleneck

Effect

Active in

A. Attention block #3102; #3395 sizes

Block-size settings ignored on the MaxText path

#3395: step 267.9 → 218.6 s [PR]

B. EOS-check control

#3174, #3509

#3509: 216 µs of a 5.4 ms iteration [PR]

C. Reward parallelism

#1904; #5255, #5235

Cross-DP EOS all-reduce in every decode iteration, even with the check off Serial host scoring of 2,048 completions

Pathways; LoCoJAX uses own tiles Both

D. GC lifecycle

#2086; #5236

E. Sequence packing F. KV-cache lifecycle G. Weight-sync dispatches H. Post-processing overlap I. Correctness, observability

PRs

Full gc.collect() after each weight sync #5100, #2309; covered 76% of training slots were padding by #2076 [Hist.]

#1904: reward 29.6 → 3.5 s [PR]; no improvement on our workload [Ours] ∼1.2 s per step [PR] #5100: training 15.2 → 6.7 s [PR]

#2231: sync 6.7 → 4.7 s [PR]; our A/B: sync 3.53 → 0.55 s, n=1 [Ours] #2232 One dispatch per scanned layer (308 sync 4.6 → 0.4 s [PR] per sync), ∼15 ms each through the Pathways proxy #2233 Detokenization and log-probs ran ∼2 s per step [PR]; ours after the last request inconclusive #3101, #3374, #3466, State leak, compile failure, wrong None claimed #3510, #2088 logits from over-token-budget padding, profiler crashes #2231; #5289

KV pool freed and reallocated at every weight sync

Pathways

Both Pathways; LoCoJAX uses own packer Pathways; LoCoJAX uses own lever Pathways

Both #3466 in both

Table D2 Per-PR inventory. All merged except #3511, closed in favor of #3490, and #2087, superseded by upstream

#2076. #5270 (CI repair) is excluded from all claims. PR

Short title

tpu-inference #3101 tpu-inference #3102 tpu-inference #3174 tpu-inference #3374 tpu-inference #3395 tpu-inference #3466 tpu-inference #3509 tpu-inference #3510 tpu-inference #3511 Tunix #1904 Tunix #2086 Tunix #2087 Tunix #2088 Tunix #2231 Tunix #2232 Tunix #2233 Tunix #2309 MaxText #5100 MaxText #5235 MaxText #5236 MaxText #5255 MaxText #5270 MaxText #5289

Fix leaking mutable default metadata Env overrides for attention block sizes EOS-check interval from env Continue decode on flax_nnx; keep compile options Forward block-size overrides on MaxText path Never pad past the token budget Skip EOS reduction when the check is off Cap request IDs in trace annotation Opt-in vocab-sharded sampling Optional process-pool reward evaluation Configurable per-step gc.collect Build packed chunks in numpy Disable profiler on empty window Optionally keep KV cache across weight sync One jitted unroll per scanned parameter Overlap detokenization and log-probs with decode Fix packed-row conversion on TPU RL packing and vLLM block-size options Linear-time answer extraction Forward Tunix GC option Picklable reward wrapper and worker count Restore SparseCore MoE layout flag Forward KV-free flag to Tunix

Table E1 Model size at 128 chips: frozen Pathways recipe vs LoCoJAX optimized (descriptive). Times: geometric

means of launch median step times; TP: rollout TP of both stacks. Ratio: Pathways/LoCoJAX with Welch 95% interval; at 32B, the per-launch interval, which includes 1. a Value of record; times from a confirmatory rerun (1.91). † Post-result. ‡ LoCoJAX adopts Pathways’ activation offload, which slows LoCoJAX only. b TP2 and TP4 not measured. c Pathways times on a basis that favours Pathways. Qwen3

Pathways (s)

LoCoJAX (s)

NuminaMath 1.7B 1 4B 1 8B 1 14B 2 32B 8b

51.1 70.3 89.4 99.2 154.1

26.7 46.5 68.9 78.5 149.7

SWE c 1.7B 4B 8B 14B

1 1 1 2–8

Ratio [95% CI] 1.88 [1.86, 1.90]a 1.51 [1.49, 1.53] 1.30 [1.28, 1.32] 1.26 [1.22, 1.30]† 1.03 [0.95, 1.10]†‡

26.1 12.8 2.04 [1.68, 2.46]† 35.3 27.8 1.27 [1.26, 1.28]† 43.0 35.5 1.21 [1.07, 1.37]†‡ LoCoJAX does not fit (4 settings)

Table E2 GSM8K size ladder at 64 chips: frozen Pathways recipe vs LoCoJAX optimized (descriptive). Times: means

of launch mean step times; TP: rollout TP of both stacks. Ratio: geometric mean over launches with Welch 95% interval. † Post-result, pooled across windows. a On launch medians: 2.47 [2.18, 2.80]. Qwen3

Pathways (s)

LoCoJAX (s)

1.7B 4B 8B 14B 32B

1 1 1 2 4

11.73 14.79 18.24 21.34 31.10

3.60 6.48 10.57 11.32 19.67

Ratio [95% CI] 3.25 [3.16, 3.36] 2.28 [2.17, 2.40]†a 1.72 [1.71, 1.73] 1.88 [1.66, 2.14] 1.58 [1.49, 1.67]

with its default sampler, where no multi-host controller is involved, colocated LoCoJAX takes 5.34× less time per step and disaggregated LoCoJAX 3.48× less. On 64 chips, against MaxText’s Gemma4-26B MoE recipe on Pathways over steady steps inside learning-rate warmup, the ratios are 2.14× (colocated) and 1.10× (disaggregated) (Table G1). Where the improvement comes from. On Qwen3-1.7B, nearly all of the improvement is in generation. LoCoJAX

decodes a step’s 32 rows in one continuous batch instead of two 16-row loops (about 2.5 s of the improvement), and its engine’s per-iteration computational cost is lower (1.8 vs 8.1 ms; about 4.8 s); together they cut generation from 8.5 to about 1.0 s per step. The trainer improves by 0.15 s, and weight synchronization adds 0.25 s colocated and 0.83 s disaggregated. On Gemma4-26B, colocation gives rollout all 64 chips at the recipe’s per-chip rollout efficiency, and the in-place weight publish takes 0.66 s instead of the recipe’s 12.2 s reshard. Disaggregated LoCoJAX keeps the recipe’s 32 + 32 split. Its rollout is faster than the recipe’s on the same 32 chips (179 vs 238 s), but moving the weights through host memory to the rollout hosts requires 55 s per step, so the net improvement is small. Limitations. (i) Google publishes the Qwen3 recipe for TPU v6e-8. We ran it on 4 v7x chips (8 JAX devices,

the same mesh). (ii) The Pathways server version pinned by the Gemma recipe is incompatible with the client in our environment, so Google’s side uses a newer Pathways release. We also made 15 attention block-size settings explicit at their resolved values, because under the RL configuration they otherwise resolve to empty values and the run fails at start-up. (iii) We report steady-state step time. It excludes start-up and compilation. It also excludes the 13 of the 16 baseline Gemma launches that did not reach the measured window; we did not establish why they stopped, and the cause may lie in our environment or configuration. Google’s three Gemma runs take 293.7, 298.1 and 350.3 s per step at equal work. We quote the median; the mean, 314.0 s, would raise every Gemma speedup by about 5%. (iv) The three runs of each LoCoJAX arm share one sampling seed, and Google’s three Qwen3 runs are bit-identical, so for these arms the intervals cover timing noise, not seed variance.

NuminaMath, 128 chips (Δ = 22 s)

GSM8K, 64 chips, descriptive (Δ = 9.1 s)

SWE, 128 chips (Δ = 7.5−13.3 s)

GSM8K 0.6B, 4−128 chips, vs LoCoJAX base (Δ = 4.2 s), panel (b) only

Pathways / LoCoJAX step time

(a) Ratio vs model size

(b) Same ratios vs LoCoJAX step time

3.0

3.0 32-128c

2.5

2.0

1.5

16c

8B: LoCoJAX offload

1.0

32B: TP8, LoCoJAX offload

1.7B

14B

32B

Qwen3 model size

100 150

LoCoJAX step time TLoCo (s, log scale)

Figure 8 Pathways/LoCoJAX step-time ratio and the model R = 1 + ∆/TLoCo . (a) Against model size: NuminaMath

and SWE at 128 chips and the GSM8K size ladder at 64 chips (frozen Pathways recipe vs LoCoJAX optimized). Lines: the model at each block’s mean ∆; shaded: range of ∆; hollow: outside the fit (32B, SWE 8B). (b) The same ratios against LoCoJAX step time, plus GSM8K Qwen3-0.6B at 4–128 chips (Pathways vs LoCoJAX base). Dotted blue: NuminaMath fit with a free slope (a = 0.92).

Table E3 Inputs and held-out checks of the model (Appendix E.1). Times: geometric means over three launches per

arm (s). Measured: TPW /TLoCo with Welch 95% interval (not computed for GSM8K). Model: 1 + ∆/TLoCo with ∆ fitted on the block’s other cells; ∗ outside the fit. NuminaMath and SWE: frozen Pathways recipe vs LoCoJAX optimized; NuminaMath 1.7B from the confirmatory rerun. GSM8K: an earlier reading; the reading of record differs most at 16 chips (Pathways 8.25 s). ‡ Two generation waves. § Additional cell, same head group. Cell

TPW TLoCo

NuminaMath, 128 chips 1.7B 51.12 26.71 4B 70.34 46.53 8B 89.43 68.95 14B (TP2) 99.20 78.54 32B (TP8)∗ 154.08 149.67 SWE, 128 chips 1.7B 26.10 4B 35.26 8B∗ 43.02

∆ Measured 24.42 23.81 20.48 20.66 4.42

Model

1.91 [1.83, 2.01] 1.51 [1.49, 1.53] 1.30 [1.28, 1.32] 1.26 [1.22, 1.30] 1.03 [1.02, 1.04]

1.81 1.47 1.33 1.29 1.15

12.82 13.28 2.04 [1.68, 2.46] 27.75 7.51 1.27 [1.26, 1.28] 35.55 7.47 1.21 [1.07, 1.37]

1.59 1.48 1.29

GSM8K 0.6B, Pathways vs LoCoJAX base 4 chips 14.53 10.06 4.47 1.44 8 chips‡ 9.64 5.54 4.10 1.74 16 chips 8.08 3.98 4.10 2.03 32 chips 7.65 2.88 4.77 2.66 64 chips 6.54 2.55 3.99 2.56 128 chips§ 6.55 2.54 4.02 2.58

1.42 1.77 2.07 2.44 2.68 2.69

Table F1 Median step time (s), LoCoJAX-CO vs LoCoJAX-DA at three trainer-host shares, Qwen3-0.6B. Fastest per

row in bold; –: not run; † provisional; ‡ added post hoc, in a later window; ∗ split chosen after the cell’s main result.

LoCoJAX-DA Dataset

Chips LoCoJAX-CO

50%

25% 12.5%

GSM8K

8 16 32 64 128 256

5.08 3.33 2.38

8.20 4.32 2.87 1.98

– 3.75 2.63 1.94

1.56 1.62

1.88 1.80

NuminaMath 16 32 64 128 256

89.4

– – 44.5 56.1 42.0 27.9 34.1 26.5 19.6 21.4 17.0†∗ 14.3 14.8∗ 12.1

25.1 16.8 11.8∗

32‡ 64 128 256

31.1 16.8 10.0 6.3†

– – 40.6 21.5∗

SWE

2.02 1.73 1.77

40.4 20.9 11.8 7.2

– 40.2 21.0∗ 11.8∗

– – – 2.56 2.90 2.22 – –

Full Scaling Study Results

Table H1 gives every cell of the scaling study of Section 4.1, of which Table 2 shows only each model’s fastest TP. All cells use the setup of Sections 3.1 and 3.2: GRPO on Pathways on 128 TPU v7x chips, 256 prompts × 8 generations = 2,048 completions per step, each forced to 8,192 output tokens (decode-heavy), with randomly initialized weights. The sampler holds 1,024 sequences at a time, so each of its DP replicas holds 1,024/DP sequences and rollout takes two decode waves; at 405B, a replica holds 16 sequences at TP16 (what fits in HBM) and 32 at TP32, so rollout takes 8 waves.

Computational Cost to Learn

The step-time comparison of Section 4.2 measures time per step, not training outcome. This appendix asks whether the faster step also reaches a quality target sooner. Separate GRPO runs train Qwen3-0.6B on GSM8K (three seeds per cell, 4–128 chips) to a frozen held-out target, and we charge each stack the TPU v7x chip-hours it needs. The main pairing is LoCoJAX base against Pathways, both with the same shared optimizations; “the frozen Pathways recipe” runs without them, and LoCoJAX optimized adds O1–O4 (Table B1). Both stacks train an identical recipe, and the target was fixed before any run.

I.1

Protocol

Table I3 gives the recipe, which differs from the step-time study’s (Table 1). Both stacks keep fp32 trainer weights and a bf16 sampler, and use the Dr. GRPO advantage, no KL term and one update per batch, so the importance ratio is 1. The target is 85% GSM8K test accuracy at temperature 0.7. A cell is charged at the first of two consecutive scored checkpoints whose three-seed score clears the target, and its computational cost is chips × wall time from job start to that checkpoint, summed over seeds. Start-up is included; queueing, evaluation and the Pathways CPU head are not, which favours Pathways. Intervals are 95% seed bootstraps with the charged checkpoints held fixed.

Table G1 Google’s published TPU RL recipes vs LoCoJAX (GRPO, TPU v7x, same chips; steady-state seconds per

step, three runs per arm). Qwen3-1.7B: steps 4–150, mean of run medians. Gemma4-26B: median of run medians over the steady steps both systems share, within steps 4–19. Speedups per generated token (Qwen3): 5.06× and 3.49×.

LoCoJAX Model, chips Qwen3-1.7B, 4 (Tunix)

Google

s/step speedup

9.14 1.71 2.62 1.00× 5.34× 3.48×

Gemma4-26B, 64 s/step (MaxText) speedup

298.1 139.1 269.2a 1.00× 2.14× 1.10×

a Compared on the eight steady steps DA shares with Google’s runs, where Google takes 295.8 s.

I.2

Results

On GSM8K, Pathways needs 1.27–2.77× the chip-hours of LoCoJAX base to reach the target at 4–128 chips (Table I1). The ratio holds when both stacks are charged at the same checkpoint and over the steady segment from update 25 to 275 (1.23–3.01×; Table I2); the 4-chip cell sits on the bar. The ratios do not rise monotonically with chips, and the one exception traces to start-up. GSM8K at 64 chips (2.33×) sits below 32 chips (2.49×) with overlapping intervals, and the steady segment reverses the order (2.58 against 2.39×). In one cell the rule charges the stacks at different checkpoints, in Pathways’ favour: at 16 chips LoCoJAX base misses update 300 by 3 of 3,957 answers (same-checkpoint ratio in Table I1). Two more seeds per cell give 1.89× [1.76, 2.03] and 2.21× [2.05, 2.38] at 16 and 64 chips (n=5). On GSM8K, start-up (to update 25) holds 24–36% of Pathways’ extra chip-hours at 4, 8, 32 and 128 chips, but with Pathways’ resource-manager CPU limit raised fourfold, a separate pairing finds no start-up difference (0.97–1.41× from job creation to the first update at 4, 16 and 64 chips; every interval includes 1). The steady segment carries the robust part of the gap. O1, O2 and O4 contribute little: LoCoJAX base needs 1.01–1.11× the chip-hours of LoCoJAX optimized on GSM8K. Held-out quality matches within evaluation noise: Pathways is 0.65 pp higher on GSM8K, below the noise floor of a single checkpoint (Table I6).

I.3

Limitations

The comparison has two limits. (i) It covers one model and one dataset, Qwen3-0.6B on GSM8K, with three seeds per cell and at most 348 updates. (ii) Pathways kept the template resource-manager CPU limit, which slowed its first weight transfer. Replacing that transfer still gives 1.26–2.54×, but at 4 chips the improvement sits on the 1.25 bar and two sensitivities move a cell below it (Table I4).

Table H1 Full accounting of the scaling study (Section 4.1, Table 2): every model of the study at every TP measured with the study recipe on 128 TPU v7x chips (controls and recipe deviations are listed separately), decode-heavy workload (8,192 forced output tokens), random initialization. Each block starts at the smallest TP that fits (∼100B: TP16 only). Bold: the model’s fastest TP (or its only TP), the row shown in Table 2; ≈: tied with it within run-to-run spread by Table 2’s rule (the runner-up’s and the fastest cell’s run ranges overlap, or their median gap is below the larger range). TP, DP: the sampler’s tensor and data parallelism (TP × DP = 256 TensorCores); m: sequences per replica (m × DP = 1,024, so the 2,048 completions take 2 decode waves; 405B: see its block). n: valid runs. Step: median over runs of each run’s steady-step median (2 runs: mean); range: fastest and slowest run, given when n ≥ 2. Shares of the profiled step, each the median over runs (2 runs: mean): computational overhead = share outside compute; host = host orchestration; network = upper-end estimate, including waits for TP partners and at the DP check; TP coll. = TP collectives, part of network; other = computational overhead − host − network (unassigned decode residual and launch-skew idle, not host orchestration); compute = 100 − computational overhead. Rounded parts may differ from their sum by 0.1. Share of the step (%) TP

Qwen3-0.6B 1 256

2 4 8 16

128 64 32 16

Qwen3-1.7B 1 256

2 4 8 16

128 64 32 16

Waves

Step (s)

Range (s)

Computational Overhead

4 8 16 32 64

2 2 2 2 2

3 3 3 1 3

68.6

67.8–69.3 67.4–72.2 73.9–77.0 – 172.5–178.9

40.6

36.4

2.2

0.0

71.4 74.9 100.4 174.1

45.8 44.0 45.6 50.8

28.7 22.6 18.6 8.3

14.8 17.7 24.8 38.7

5.3 12.1 21.6 33.8

4 8 16 32 64

2 2 2 2 2

3 3 1 1 3

79.4–84.6 75.3–80.1 – – 178.6–182.8

29.8

24.3

3.5

0.0

80.1 87.3 103.6 182.6

43.4 48.0 44.5 52.7

26.5 21.8 15.0 10.2

14.6 23.2 27.2 39.2

4.8 15.5 23.8 31.7

4 8 16 32 64

2 2 2 2 2

3 3 3 1 3

123.6

118.7–124.2 106.5–109.2 115.2–118.8 – 232.7–240.4

25.4

21.9

2.2

0.0

1.2

74.6

32.2

17.6

12.2

5.4

2.4

67.8

41.6 40.0 50.7

17.7 12.0 8.8

20.9 25.9 38.8

15.0 23.3 32.3

3.1 2.2 3.2

58.4 60.0 49.3

4 8 16 32 64

2 2 2 2 2

1 3 3 3 3

167.2 137.3

– 134.7–140.2 134.0–134.4 146.6–152.4 240.3–248.2

21.8 28.5

18.3 16.6

2.7 10.0

0.0 4.3

0.8 2.0

78.2 71.5

≈ (tie)

39.6

18.3

18.5

13.5

2.7

60.4

fastest

40.8 50.2

15.3 9.4

23.8 37.9

21.6 32.9

1.7 2.9

59.2 49.8

79.6

Host

Network

TP coll.

Other

Compute

Fastest

2.0

59.4

fastest

2.4 3.7 2.2 3.7

54.2 56.0 54.4 49.2

≈ (tie)

2.0

70.2

fastest

2.4 3.1 2.3 3.2

56.6 52.0 55.5 47.3

≈ (tie)

Qwen3-4B

1 2

4 8 16

256 128 64 32 16

108.5

118.5 135.3 236.3

fastest

Qwen3-8B

1 2 4

8 16

256 128 64 32 16

134.1

151.2 247.2

Qwen3-14B — no TP16 (its 40 attention heads do not divide by 16); TP1 only with non-default attention tiles, listed separately

2 4

128 64 32

8 16 32

2 2 2

1 3 3

190.0

8 16 32 64

2 2 2 2

1 3 3 3

350.6

2 2 2

3 3 3

457.3

168.9

178.6

– 163.7–169.6 178.0–179.4

24.1

13.7

9.0

3.5

1.5

34.2

14.8

17.4

12.2

2.0

75.9 65.8

37.6

13.2

22.9

20.4

1.5

62.4

– 298.4–302.5 303.3–305.9 463.2–471.3

17.9

7.9

8.7

3.0

1.3

82.1

29.2

8.3

18.6

10.6

2.3

70.8

fastest

32.8 46.0

6.3 5.2

24.5 37.7

18.5 29.5

2.0 3.1

67.2 54.0

≈ (tie)

451.3–459.6 421.6–430.1 604.7–611.4

22.2

7.2

13.3

1.7

77.8

fastest

Qwen3-32B

2 4

8 16

128 64 32 16

301.3

303.8 468.4

Llama-3.1-70B

4 8

64 32 16

16 32 64

425.6

605.2

9.0

28.3

7.0

20.0

17.2

1.3

71.7

43.9

5.0

36.3

31.7

2.6

56.1

fastest

Synthetic ∼100B — TP16 only (one model added between 70B and ∼200B; smaller TPs not sized); 405B width, 31 layers; random init, not a released model 16 16 64 2 2 382.9 381.1–384.8 33.6 5.5 26.4 20.2 1.7 66.4 only TP Synthetic ∼200B — 405B width, 62 layers; random init, not a released model 16 16 64 2 3 729.3 729.0–733.2

31.8

4.4

25.9

20.0

1.5

68.2

41.4

3.0

35.1

31.9

3.2

58.6

128

965.9

963.2–968.6

fastest

Llama-3.1-405B — 16 sequences per replica at TP16 (what fits in HBM) and 32 at TP32, so both run 8 decode waves; one TP16 run

Memory workarounds: TP16: a compiler flag, freeing the sampler’s weight copy and the KV cache, and q/k/v remat; TP32: q/k/v remat only 16 16 8 1 3,334.3 – 40.3 1.4 38.5 31.2 0.4 59.7 fastest 32 8 32 8 2 4,102.7 4,020.5–4,184.9 56.0 0.9 53.8 49.5 1.2 44.0 16

Table H2 Controls and diagnostics of the scaling study, kept out of Table H1 and never pooled with it; columns and

statistics as there. (a) Wave control: Synthetic ∼200B at TP16 with 16 sequences per replica, so 8 decode waves as at 405B TP16 (grid row: m = 64, 2 waves). Workaround controls: the 405B memory workarounds applied to Synthetic ∼200B (TP16: the compiler flag and freeing the sampler’s weight copy and the KV cache, without remat; TP32: q/k/v remat), compared with steps 1–3 of the grid runs of the same TP (728.4 s at TP16, 963.8 s at TP32). ∗ Four-step control run, no steady readout by design: median over steps 1–3 (after the warm-up step 0) of the step plus the following weight sync (the readout’s basis); not profiled, so no breakdown. Weight pinning (with freeing): both attempts failed in the first weight sync, refused on Pathways. (b) Llama-3.1-8B architecture, randomly initialized, compared with Qwen3-8B at TP16 in Table H1. (c) Qwen3-14B at TP1 with a different attention-tile setting (a recipe deviation); with the default tiles 14B runs from TP2. (d) 16 replicas of TP8 on 64 chips; chip count and per-replica work differ from the 128-chip grid, so they are not comparable with Table H1. GB200 runs (another platform, qualitative only) are not shown. Share of the step (%) Run set

TP DP

m Waves

n Step (s)

(a) Synthetic ∼200B controls for the 405B rows (128 chips) Wave control (m = 16) 16 16 16 8 1 1,668.8 TP16 package without remat 16 16 64 2 1∗ 742.6∗ TP32 q/k/v remat 32 8 128 2 1∗ 962.3∗ TP16 weight pinning (PIN) 16 16 64 2 0/2 –

Range (s)

Computational Overhead Host Network TP coll. Other Compute Note 40.9 – – –

1.7 – – –

39.1 – – –

32.8 – – –

0.1 – – –

59.1 8 waves, as 405B TP16 – +1.9% vs grid, same steps – −0.2% vs grid, same steps – refused on Pathways

(b) Model-family control at TP16, 128 chips (never pooled with Qwen3-8B) Llama-3.1-8B shape 16 16 64 2 3 216.0 213.3–224.7

50.8

9.4

38.3

32.8

3.1

49.2 Qwen3-8B: 247.2 s

14.1 12.1

1.2

0.0

0.8

85.9 default tiles: TP2 up

(d) Legacy TP8 × DP16 rows on 64 chips (chip count differs from the 128-chip grid: confounded) Qwen3-0.6B 8 16 64 2 1 146.8 – 37.4 14.3 Qwen3-1.7B 8 16 64 2 1 150.2 – 36.7 12.9 Qwen3-4B 8 16 64 2 1 194.1 – 31.2 7.8 Qwen3-8B 8 16 64 2 1 208.8 – 30.3 8.9 Qwen3-14B 8 16 64 2 1 250.4 – 28.2 7.3 Qwen3-32B 8 16 64 2 1 420.5 – 23.2 4.3 Llama-3.1-70B 8 16 64 2 1 673.8 – 29.6 3.2

19.2 19.8 19.5 17.8 17.6 15.1 23.5

17.1 17.7 17.1 15.9 15.6 14.2 20.5

4.0 4.0 3.8 3.6 3.4 3.7 2.8

62.6 63.3 68.8 69.7 71.8 76.8 70.4

246.7

– – – –

–

Table I1 TPU v7x chip-hours to the frozen held-out target, summed over 3 seeds, LoCoJAX base and Pathways with

the same shared optimizations. Under the dataset: target and ∆, held-out quality of Pathways minus LoCoJAX (Table I6). Ratio: Pathways over LoCoJAX base [95% seed bootstrap]; bar 1.25 (pre-declared). m On the bar (interval 1.245–1.289). a LoCoJAX charged at update 325, Pathways at 275 (same checkpoint: 1.84). Dataset GSM8K 85% (T0.7) ∆ +0.65 pp

Chips LoCoJAX Pathways Ratio [interval] 4 8 16 32 64 128

23.3 28.3 48.0a 64.7 121.0 228.4

29.6 43.0 77.6 160.7 282.1 633.0

1.27m 1.52 [1.42, 1.63] 1.62 [1.49, 1.75] 2.49 [2.30, 2.64] 2.33 [2.14, 2.48] 2.77 [2.49, 3.00]

Table I2 Full grid, LoCoJAX base against Pathways with the same shared optimizations. Chip-hours to the target (3

seeds summed) with the charged checkpoint pair; ratio Pathways over LoCoJAX base [95% seed bootstrap]. Same checkpoint: both stacks charged at each update of the earlier stack’s pair. Steady: chip-hours from update 25 to 275. i Interval 1.201–1.263. Status: against the pre-declared 1.25 bar.

Dataset Chips LoCoJAX base (pair) Pathways (pair)

Ratio [interval]

Same checkpoint Steady Status

GSM8K

1.267 [1.245, 1.289] 1.518 [1.424, 1.629] 1.616 [1.492, 1.754] 2.485 [2.301, 2.640] 2.332 [2.138, 2.477] 2.772 [2.495, 2.999]

1.267 / 1.264 1.518 / 1.509 1.838 / 1.841 2.485 / 2.499 2.332 / 2.368 2.772 / 2.811

4 8 16 32 64 128

23.326 (275/300) 28.300 (275/300) 48.013 (325/348) 64.667 (275/300) 120.996 (275/300) 228.373 (275/300)

29.550 (275/300) 42.967 (275/300) 77.591 (275/300) 160.711 (275/300) 282.116 (275/300) 632.996 (275/300)

1.232i 1.493 1.757 2.394 2.578 3.014

on the bar supported supported supported supported supported

Table I3 Computational cost-to-learn recipe (Qwen3-0.6B, GRPO, 3 seeds per cell), identical on both stacks. P × G:

prompts × generations per update; caps in tokens. Scored checkpoints: updates 275, 300, 325 and 348.

P×G

Dataset Training data (revision) GSM8K openai/gsm8k @740312ad: first 7,424 training rows, 6 epochs

Caps Updates Evaluation

128 × 16 256 / 7,936

348 test split (1,319 rows), T 0.7

Target 3-seed sum ≥ 3,364 of 3,957 (85%) at updates 275–348

Table I4 Sensitivities of the grid, hypothetical on measured inputs. Values are Pathways over LoCoJAX base; chip

counts in parentheses. The template resource-manager CPU limit slowed Pathways’ first weight transfer to about 8–11 min in the 16- and 32-chip GSM8K seeds.

Sensitivity

Effect

Verdict change

Slow first weight transfer replaced GSM8K 1.265 / 1.263 (4), 1.508 (8), 1.516 (16; same checkpoint (raised limit) or removed 1.724), 2.239 (32), 2.260–2.308 (64), 2.541 (128) Limit’s effect on steady steps GSM8K (4) 1.21–1.24; every other cell ≥ 1.39 with all Pathways (1.031–1.062; 1.062 [1.000, 1.155] chip-hours divided by 1.062 at n=3, 1.031 [0.980, 1.090] at n=5, inconclusive) One-checkpoint counterfactuals at GSM8K 1.838 (16), 2.342 (32), 2.214 (64) near-threshold margins Head placement: GSM8K 1.218 (4), 1.399 (8), 2.384 (32), 2.156 (64), 2.566 (128; off-reference-group Pathways seeds slow head nodes 2.530) divided by 1.13

none GSM8K (4) below the bar

none GSM8K (4), whose seeds show no group effect

Table I5 Other pairings of the same runs: chip-hour ratios to the target per cell (3 seeds summed); only the first column

is a fair comparison. Fair: Pathways over LoCoJAX base. Own best: Pathways over LoCoJAX optimized. Frozen / Pathways: what the shared optimizations save Pathways. Base / optimized: what O1, O2 and O4 save LoCoJAX. Frozen / optimized: the product of the Fair, Frozen / Pathways and Base / optimized columns.

Dataset

Chips

Fair

Own best

Frozen / Pathways

Base / optimized

Frozen / optimized (not fair)

GSM8K

4 8 16 32 64 128

1.267 1.518 1.616 2.485 2.332 2.772

1.281 1.554 1.711 2.618 2.590 3.064

1.385 1.325 1.233 1.162 1.133 1.028

1.011 1.023 1.059 1.053 1.111 1.106

1.774 2.058 2.110 3.043 2.934 3.150

Table I6 Held-out quality, Pathways minus LoCoJAX (Welch 95% CI; runs per stack). Pooled over chip counts, which

share seeds and data order, so the interval is optimistic. Floor: evaluation noise of one checkpoint.

Dataset (score)

Pathways − LoCoJAX

Floor

GSM8K (T 0.7, mean of updates 275 and 300)

+0.65 pp [+0.28, +1.02] (18 v 18)

∼1 pp

LoCoJAX: Host-Local Coordination for Small-Model RL on TPUs

Analyze & explore related research in Conceptio

Related documents