Google's Tunix Unifies the Messy Open-Source Post-Training Stack in JAX
Google open-sources Tunix, a JAX-native post-training library covering SFT, RL, distillation, and agentic training in one unified toolkit for TPUs
PRO- Google open-sourced Tunix, a JAX-native LLM post-training library covering SFT, RL, distillation, and agentic training in one toolkit.
- Supports PPO, GRPO, GSPO, DPO, DAPO, and Dr.GRPO out of the box, plus LoRA/QLoRA for parameter-efficient fine-tuning.
- Agentic RL module enables multi-turn tool-use training with async rollouts, trajectory batching, and weight-sync locking for distributed setups.
- 12% relative improvement in pass@1 accuracy on GSM8K math reasoning using GRPO on Gemma 2 2B-IT.
- White-box design lets you customize training loops directly without navigating abstraction layers; integrates with vLLM and SGLang-JAX for fast rollouts.
- Free, Apache 2.0; install via
pip install google-tunix; supports Gemma, Llama, and Qwen model families.
Tunix (short for Tune-in-JAX) is Google's answer to a problem that's been quietly frustrating ML teams for years: once you have a pre-trained model, the path to a production-ready, aligned system is a mess of fragmented tools, custom scripts, and framework-specific hacks. Tunix consolidates that entire post-training stack into a single, open-source JAX-native library.
The gap it fills
Model pre-training has seen rapid progress, but alignment and post-training often require fragmented or custom solutions. If you've ever tried to run GRPO on a JAX model while also needing DPO and a distillation pipeline, you know the pain. Most existing tooling is PyTorch-first, and stitching together reward models, RL trainers, and fine-tuning loops across different frameworks is a significant engineering burden.
Tunix closes this critical gap by providing a comprehensive and developer-friendly toolkit for aligning models at scale. Built for performance on TPUs, especially when combined with MaxText, it offers a complete suite of production-ready trainers for SFT, preference tuning, knowledge distillation, and advanced RL methods like PPO, GRPO, and GSPO in a single, unified library.
What's inside the box
The library is organized around four core capability areas, each with dedicated trainer classes:
- Supervised Fine-Tuning (SFT): Supports full fine-tuning and parameter-efficient approaches like LoRA and QLoRA (via qwix).
- Preference Tuning:
Implements Direct Preference Optimization (DPO) with the
DPOLearner, removing the need for a reward model. DPO trains the model directly on pairs of preferred vs. rejected responses, skipping the expensive step of training a separate reward model.
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
