Tunix
A JAX-based library for post-training large language models, providing efficient supervised fine-tuning, reinforcement learning and agentic reinforcement learning on TPUs.
Tunix (Tune-in-JAX) is a JAX-based library for post-training large language models, aimed at efficient, scalable training on TPUs. It sits as a post-training layer within the JAX training stack, built on foundational tools such as Flax, Optax and Orbax, and between them and optimized model implementations such as MaxText and MaxDiffusion, and it integrates with inference engines including vLLM and SGLang-JAX for rollout.
It supports supervised fine-tuning, including full-weight and parameter-efficient fine-tuning, DPO and ORPO; reinforcement learning methods including PPO, GRPO, token-level GSPO, DAPO and Dr.GRPO; and agentic reinforcement learning with multi-turn tool use, asynchronous rollout for high-throughput trajectory collection, and trajectory batching and grouping.
Its components are designed to be modular, reusable and composable, and the project is under active development, with recent updates adding support for Gemma models and more efficient kernel execution.