RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs
RL LLM
Runlong Zhou*, Lefan Zhang*, Shang-Chen Wu*, Kelvin Zou*, Hanzhi Zhou, Ke Ye, Yihao Feng, Dong Yin, Alex Guillen Garcia, Dmytro Babych, Rohit Chatterjee, Matthew Hopkins, Xiang Kong, Chang Lan, Lezhi Li, Yiping Ma, Daniele Molinari, Senyu Tong, Yanchao Sun, Thomas Voice, Jianyu Wang, Chong Wang, Simon Wang, Floris Weers, Yechen Xu, Guolin Yin, Muyang Yu, Yi Zhang, Zheng Zhou, Danyang Zhuo, Ruoming Pang, Cheng Leong
We built a large scale RL training system for LLMs on TPUs, supporting different asynchronous levels while being robust to preemptions.