用TPU实现大规模分布式强化学习，12小时提升大模型推理能力12.8%
RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs
- 基于参数服务器架构，支持多工作节点并行训练
- 在1024个v5p TPU上12.8小时提升QwQ-32B pass@8准确率
- 支持中断恢复，适合大规模、高可用的RL训练场景
强化学习（RL）已成为提升大语言模型（LLM）推理能力的标准范式。我们开发了基于TPU的可扩展强化学习框架RLAX，采用参数服务器架构：主训练器定期向参数服务器推送更新权重，大量推理工作节点拉取最新权重生成新轨迹。我们引入一系列系统技术，支持多种先进RL算法的大规模、可抢占式训练。为加速收敛并提升模型质量，设计了新的数据集构建与对齐方法。大规模评估显示，在1024个v5p TPU上仅用12小时48分钟，RLAX使QwQ-32B的pass@8准确率提升12.8%，且训练过程中对中断具有鲁棒性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。