Async RL in Pure JAX

Async RL in JAX

Scaling RL has become the core approach in unlocking the capabilities of many frontier models today. Thinking beyond current RLVR recipes potentially requires changes across different actors within the stack. Existing frameworks such as Tunix are great for production level RLVR; however, for experimentation, they are quite heavy and require Pathways to scale RL on TPUs. To make this experimentation easier, we introduce a lightweight, performant asynchronous RL stack in pure JAX. We share a work log alongside a few engineering insights we've seen along the way. Code:…

saved by

• Ishaan Panigrahi

related reading

• How to Parallelize a Transformer for Training — an explorable explanationezyang.github.io
• RL Scaling Laws for LLMs - by Cameron R. Wolfe, Ph.D.cameronrwolfe.substack.com
• Is Frontier Asynchronous RL Solved? — Luke J. Huangluk-huang.github.io
• Keep the Tokens Flowing: Lessons from 16 Open-Source RL Librarieshuggingface.co
• State of RL for reasoning LLMs | A. Weersaweers.de
• RL at 1T Scale: prime-rl Performance Deep Diveprimeintellect.ai
• Defeating Nondeterminism in LLM Inference - Thinking Machines Labthinkingmachines.ai
• How To Scale Your Modeljax-ml.github.io
• irhum.github.io - Tensor Parallelism with jax.pjitirhum.github.io
• Rethinking RL Infra for Agents | B'Logbillxbf.github.io
• Composer2.pdfcursor.com
• [2602.17616] Stable Asynchrony: Variance-Controlled Off-Policy RL for LLMsarxiv.org