LoCoJAX cuts step time for small-model TPU training
LoCoJAX: Host-Local Coordination for Small-Model RL on TPUs · arXiv:2611.16301
In plain words
Small language models can spend a lot of TPU training time on coordination around the model computations. This paper introduces LoCoJAX, a system that lets each TPU host handle its own requests and generated answers, then brings the hosts together to train.
In tests with Qwen3-0.6B, LoCoJAX's colocated setup completed steps faster than the Pathways baseline across three workloads and the tested chip counts. The detailed comparison covers that model, so it does not show the same result for every model or setup.
The research
For small-model TPU training, time spent coordinating work can matter as much as the model computation itself.
The results change the expectation that adding chips alone addresses the bottleneck: in the tested setup, moving rollout control to hosts reduced step times, while a separate placement test shows that separating training and generation can help for a particular rollout-heavy workload.
LoCoJAX moves rollout control to TPU hosts and keeps training synchronous, reducing measured step times versus Pathways for small-model reinforcement learning.
The system comparison measured median steady-state step times for Qwen3-0.6B on three workloads, matching the work processed by Pathways and LoCoJAX. Across 8–64 TPU chips, optimized LoCoJAX was faster on the tested configurations, though it did not fit the 8-chip SWE-agent setup.
This supports a step-time advantage for LoCoJAX in these comparisons. Because several system implementation choices changed together, the measurements do not identify one change as the cause of the full difference.
The detailed comparison covers Qwen3-0.6B, and the profiling study used a decode-heavy workload. Pathways could not be tuned as extensively as LoCoJAX, trainer precision differed across datasets, and some LoCoJAX configurations did not fit in memory.
What the evidence shows
- The bottleneck shifts
Local Coordination JAX (LoCoJAX) is a TPU reinforcement-learning system that gives each host control of its own rollout work. It addresses a profiling result: in the evaluated Pathways setup, host orchestration was the largest non-compute part of a step for the smallest models, while network communication became dominant for the largest.
The comparison therefore starts with a bottleneck that changes with model size, rather than a single cost shared by all models.
- A size-dependent contrast
Across these two models in this decode-heavy Pathways profile, the dominant non-compute cost shifts from host coordination for Qwen3-0.6B to network communication for Llama-3.1-405B.
That contrast supports focusing LoCoJAX on host coordination for small-model training, while limiting the inference to this workload and evaluated stack.
- Two costs in rollout
For Qwen3-0.6B on GSM8K at 64 chips, the comparison separates two recurring costs: centralized request admission before generation and how each system runs decode iterations.
Their different execution paths help explain the measured step-time gap, but the timings do not isolate synchronization as its sole cause.
- Faster steps in scope
For Qwen3-0.6B, optimized LoCoJAX had shorter median steady-state steps than tuned Pathways across the tested workloads and chip counts from 8 to 64, where it fit.
Because several system choices and optional optimizations differed together, the comparison supports a bounded performance advantage but does not attribute it to one change.
- Placement changes the answer
On NuminaMath-CoT at 256 chips, one tested split assigned 12.5% of hosts to training and the rest to rollout; disaggregated LoCoJAX had a shorter median step than colocated LoCoJAX.
Because placement, overlap, and one-step staleness change together, this supports separation as an option for this rollout-heavy setting, not a general rule that one mode wins.
Deep
- Original paper authors
- Loki Chen, Gregory B. Rehm, Yang Song, Hongye Xie, Ming Lei
- Submitted to arXiv
- License
- arXiv non-exclusive