google/tunix

What it solves

Tunix is a JAX-based library designed to simplify and accelerate the post-training of Large Language Models (LLMs). It provides a scalable framework for refining models through various alignment and optimization techniques, specifically optimized for high performance on TPUs.

How it works

Tunix acts as an intermediate layer in the JAX ecosystem, integrating foundational tools like Flax, Optax, and Orbax. It connects these utilities with high-performance model implementations (such as MaxText) and inference engines (like vLLM and SGLang-JAX) to handle efficient rollouts and model execution.

Who it’s for

It is intended for researchers and developers who need to perform scalable post-training on LLMs using JAX and TPU infrastructure, including those working on supervised fine-tuning and reinforcement learning.

Highlights

• Diverse Training Algorithms: Supports Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and various Reinforcement Learning (RL) methods including PPO, GRPO, and GSPO-Token.
• Agentic RL: Features support for multi-turn tool use, asynchronous rollout for high-throughput trajectory collection, and trajectory batching.
• TPU Optimization: Native integration with vLLM, SGLang-JAX, and MaxText for state-of-the-art performance on TPU hardware.
• Scalability: Enables seamless multi-host distributed training via Pathways, scaling up to thousands of devices.
• Model Support: Compatible with popular model families such as Gemma, Llama, and Qwen.