Google Introduces Tunix: An Open-Source JAX Library for LLM Post-Training
Tunix

Image Courtesy: Google

Google has released Tunix, a new open-source, JAX-native library designed to simplify post-training workflows for large language models (LLMs). Tunix provides developers and researchers with a unified toolkit for aligning pre-trained models into production-ready systems.

Why Tunix Matters

Model pre-training has seen rapid progress, but alignment and post-training often require fragmented or custom solutions. Tunix addresses this challenge by delivering a comprehensive and developer-friendly library, built for performance on TPUs and optimized when combined with Google’s MaxText framework.

Key Features

  • Complete Post-Training Algorithms: Includes trainers for Supervised Fine-Tuning (SFT), preference tuning, knowledge distillation, and advanced Reinforcement Learning (RL) methods like PPO, GRPO, and GSPO.

  • White-Box Design: Offers full control, allowing developers to customize training loops and post-training code directly.

  • JAX-Native: Integrates seamlessly with the JAX ecosystem and open-source models.

Initial Release Highlights

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

  • Reinforcement Learning (RL):

    • PPOLearner for Proximal Policy Optimization (PPO).

    • GRPOLearner for Group Relative Policy Optimization (GRPO).

    • GSPO-token for token-level flexibility in multi-turn RL training.

  • Knowledge Distillation: The DistillationTrainer supports both logit-based distillation and attention transfer for model compression.

  • Agentic AI Support: Enables training of LLM-based agents that can reason and interact with external environments.

Performance Gains

On the GSM8K math reasoning benchmark, Tunix’s GRPO implementation improved the Gemma 2 2B-IT model’s pass@1 accuracy by ~12%, with consistent gains across other evaluation metrics, including partial and format accuracy.

Getting Started

Tunix is available now on PyPI:

pip install google-tunix

The official repository includes examples, canonical implementations, and onboarding notebooks to help developers get started quickly.