RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs
Abstract
Reinforcement learning (RL) has emerged as the de-facto paradigm for improving the reasoning capabilities of large language models (LLMs). We have developed RLAX, a scalable RL framework on TPUs. RLAX employs a parameter-server architecture. A master trainer periodically pushes updated model weights to the parameter server while a fleet of inference workers pull the latest weights and generates new rollouts. We introduce a suite of system techniques to enable scalable and preemptible RL for a diverse set of state-of-art RL algorithms. To accelerate convergence and improve model quality, we have devised new dataset curation and alignment techniques. Large-scale evaluations show that RLAX improves QwQ-32B's pass@8 accuracy by 12.8% in just 12 hours 48 minutes on 1024 v5p TPUs, while remaining robust to preemptions during training.
Cite
@article{arxiv.2512.06392,
title = {RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs},
author = {Runlong Zhou and Lefan Zhang and Shang-Chen Wu and Kelvin Zou and Hanzhi Zhou and Ke Ye and Yihao Feng and Dong Yin and Alex Guillen Garcia and Dmytro Babych and Rohit Chatterjee and Matthew Hopkins and Xiang Kong and Chang Lan and Lezhi Li and Yiping Ma and Daniele Molinari and Senyu Tong and Yanchao Sun and Thomas Voice and Jianyu Wang and Chong Wang and Simon Wang and Floris Weers and Yechen Xu and Guolin Yin and Muyang Yu and Yi Zhang and Zheng Zhou and Danyang Zhuo and Ruoming Pang and Cheng Leong},
journal= {arXiv preprint arXiv:2512.06392},
year = {2025}
}
Comments
The submission is being withdrawn because internal stakeholders determined that it is not appropriate to publish work on this topic at this time