RLAX: Large-Scale, Distributed Reinforcement Learning for Large Language Models on TPUs
On the map: LLM reasoning → Reinforcement learning for reasoning
Abstract (arXiv)
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.
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning. Builds on 18 of its 38 references that are in the map. First posted on arXiv 6 Dec 2025.
Builds on
18 of 38 references.
Closest ancestors · 6
Same direction, newest first
- 2025Sample More to Think Less: Group Filtered Policy Optimization for Concise Reasoning13 Aug
- 2025AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning30 May
- 2025Understanding R1-Zero-Like Training: A Critical Perspective26 Mar
- 2025DAPO: An Open-Source LLM Reinforcement Learning System at Scale18 Mar
- 2025DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning22 Jan
- 2017Proximal Policy Optimization Algorithms20 Jul
Other references · 12
- 2022Training language models to follow instructions with human feedbackcited by 15,971
- 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 9,379
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
- 2024HybridFlow: A Flexible and Efficient RLHF Frameworkcited by 2,057
- 2015High-Dimensional Continuous Control Using Generalized Advantage Estimationcited by 2,026
- 2025MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attentioncited by 263
- 2024Spider 2.0: Evaluating Language Models on Real-World Enterprise Text-to-SQL Workflowscited by 232
- 2025Defeating the Training-Inference Mismatch via FP16cited by 66
- 2022Bamboo: Making Preemptible Instances Resilient for Affordable Training of Large DNNscited by 49
- 2025AXLearn: Modular, Hardware-Agnostic Large Model Trainingcited by 4
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2025A Survey of Reinforcement Learning for Large Reasoning Models
- 2025RLVE: Scaling Up Reinforcement Learning for Language Models with Adaptive Verifiable Environments
- 2025Reinforced MLLM: A Survey on RL-Based Reasoning in Multimodal Large Language Models
- 2025RLAE: Reinforcement Learning-Assisted Ensemble for LLMs
- 2025Laminar: A Scalable Asynchronous RL Post-Training Framework
- 2025LlamaRL: A Distributed Asynchronous Reinforcement Learning Framework for Efficient Large-scale LLM Training
- 2025Reinforcement Learning Meets Large Language Models: A Survey of Advancements and Applications Across the LLM Lifecycle
- 2025A Technical Survey of Reinforcement Learning Techniques for Large Language Models