RLAX：面向 TPU 的大规模分布式强化学习框架用于大语言模型
机器学习
2025-12-12 v2 人工智能
摘要
强化学习 (RL) 已成为提高大型语言模型 (LLM) 推理能力事实标准范式。我们开发了 RLAX，一个在 TPU 上实现可扩展 RL 的框架。RLAX 采用参数服务器架构。主训练器定期将更新后的模型权重推送到参数服务器，而推理工作器从参数服务器获取最新权重并生成新的 rollout。我们引入了一套系统技术，使其能够支持多样化的先进 RL 算法的可扩展和可抢占式 RL。为了加快收敛速度并提高模型质量，我们设计了新的数据策展和对齐技术。大规模评估表明，RLAX 在 1024 个 v5p TPU 上仅用 12 小时 48 分钟即可将 QwQ-32B 的 pass@8 正确率提高了 12.8%，在训练期间保持对抢占的鲁棒性。
引用
@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}
}
备注
The submission is being withdrawn because internal stakeholders determined that it is not appropriate to publish work on this topic at this time