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

arXiv: https://arxiv.org/abs/2512.06392 · HTML: https://mlsanity.com/paper/2512.06392/

## Key facts

- Cited by 1 papers 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.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)

## Builds on (most cited first, 18 of 18)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attention](https://mlsanity.com/paper/2506.13585/) (2025)
- [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://mlsanity.com/paper/2505.24298/) (2025)
- [Spider 2.0: Evaluating Language Models on Real-World Enterprise Text-to-SQL Workflows](https://mlsanity.com/paper/2411.07763/) (2024)
- [Sample More to Think Less: Group Filtered Policy Optimization for Concise Reasoning](https://mlsanity.com/paper/2508.09726/) (2025)
- [Defeating the Training-Inference Mismatch via FP16](https://mlsanity.com/paper/2510.26788/) (2025)
- [Bamboo: Making Preemptible Instances Resilient for Affordable Training of Large DNNs](https://mlsanity.com/paper/2204.12013/) (2022)
- [AXLearn: Modular, Hardware-Agnostic Large Model Training](https://mlsanity.com/paper/2507.05411/) (2025)

## Built on by (most cited first, 1 of 1)

- [ECHO-2: A Large-Scale Distributed Rollout Framework for Cost-Efficient Reinforcement Learning](https://mlsanity.com/paper/2602.02192/) (2026)

## 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.