Saved in:
| Main Authors: | , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2512.06392 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866909956112908288 |
|---|---|
| author | Zhou, Runlong Zhang, Lefan Wu, Shang-Chen Zou, Kelvin Zhou, Hanzhi Ye, Ke Feng, Yihao Yin, Dong Garcia, Alex Guillen Babych, Dmytro Chatterjee, Rohit Hopkins, Matthew Kong, Xiang Lan, Chang Li, Lezhi Ma, Yiping Molinari, Daniele Tong, Senyu Sun, Yanchao Voice, Thomas Wang, Jianyu Wang, Chong Wang, Simon Weers, Floris Xu, Yechen Yin, Guolin Yu, Muyang Zhang, Yi Zhou, Zheng Zhuo, Danyang Pang, Ruoming Leong, Cheng |
| author_facet | Zhou, Runlong Zhang, Lefan Wu, Shang-Chen Zou, Kelvin Zhou, Hanzhi Ye, Ke Feng, Yihao Yin, Dong Garcia, Alex Guillen Babych, Dmytro Chatterjee, Rohit Hopkins, Matthew Kong, Xiang Lan, Chang Li, Lezhi Ma, Yiping Molinari, Daniele Tong, Senyu Sun, Yanchao Voice, Thomas Wang, Jianyu Wang, Chong Wang, Simon Weers, Floris Xu, Yechen Yin, Guolin Yu, Muyang Zhang, Yi Zhou, Zheng Zhuo, Danyang Pang, Ruoming Leong, Cheng |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2512_06392 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs Zhou, Runlong Zhang, Lefan Wu, Shang-Chen Zou, Kelvin Zhou, Hanzhi Ye, Ke Feng, Yihao Yin, Dong Garcia, Alex Guillen Babych, Dmytro Chatterjee, Rohit Hopkins, Matthew Kong, Xiang Lan, Chang Li, Lezhi Ma, Yiping Molinari, Daniele Tong, Senyu Sun, Yanchao Voice, Thomas Wang, Jianyu Wang, Chong Wang, Simon Weers, Floris Xu, Yechen Yin, Guolin Yu, Muyang Zhang, Yi Zhou, Zheng Zhuo, Danyang Pang, Ruoming Leong, Cheng Machine Learning Artificial Intelligence 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. |
| title | RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2512.06392 |