RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs

TL;DR

RLAX is a scalable, distributed reinforcement learning framework on TPUs that enhances large language models' reasoning abilities through system innovations and dataset techniques, achieving significant accuracy improvements efficiently.

Contribution

Developed RLAX, a novel scalable RL framework on TPUs with system optimizations and data curation methods for large language models.

Findings

RLAX improves QwQ-32B's pass@8 accuracy by 12.8% in under 13 hours.

Achieves scalable RL training with robustness to preemptions.

Demonstrates effective large-scale RL on TPU infrastructure.

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.