Tunix Documentation — Lightweight LLM Post-Training Library in JAX

by Google

IntermediateDocumentationFree~1-2 hours per Colab quick start; several days to adapt a GRPO or agentic RL run to your own model and reward

Google's JAX library for SFT, DPO, GRPO and agentic RL post-training, with Colab notebooks you can run end to end.

Start LearningAdded Oct 12, 2026 · Updated Oct 12, 2026

Overview

Tunix is the official documentation for google/tunix, a JAX-native post-training library built on Flax NNX that Google announced on the Google Developers Blog on September 30, 2025 and now ships as its V2 release. The docs are organised as Quick Start (installation, a GRPO quick start and a multi-node quick start), Design Overview (the high-level architecture plus the SFT, RL and agentic RL designs), an Agentic RL section on core components and async rollout, Performance Considerations (LoRA, batching, collocated versus disaggregated training, weight sync), Reliability (checkpoints, fault tolerance, determinism), Launching Jobs through the Tunix CLI, Rollout backends (vLLM, SGLang, a vanilla sampler and a mock), Algorithms, Models and Metrics. Supported methods are PEFT/LoRA, DPO and ORPO for supervised and preference tuning, and PPO, GRPO, GSPO-Token, DAPO and Dr.GRPO for reinforcement learning, with logit distillation shown in a notebook. The model table covers Gemma 1 through Gemma 4, Llama 3 to 3.2 and Qwen 2.5 through Qwen 3.5. Worked examples include QLoRA on Gemma 270M for translation, GRPO on Gemma 3 1B for GSM8K math, DPO on Gemma 3 1B, Gemma 7B to 2B distillation and a DeepScaleR reproduction. It is TPU-first: multi-host training runs through Pathways and fully optimised models come via MaxText. The launch post reports roughly a 12% relative gain in GSM8K pass@1 from GRPO on Gemma 2 2B-IT.

At a Glance

Topic: Fine-Tuning
Level: Intermediate
Format: Documentation
Cost: Free
Duration: ~1-2 hours per Colab quick start; several days to adapt a GRPO or agentic RL run to your own model and reward
Provider: Google
Hands-on: Yes — code/exercises
Certificate: None

What You’ll Learn

✓Fine-tune Gemma, Llama or Qwen with full weights, LoRA or QLoRA using the PeftTrainer
✓Run preference tuning with DPO and ORPO on paired chosen and rejected responses
✓Train a reasoning model with GRPO on GSM8K, including reward functions and pass@1 evaluation
✓Choose between collocated and disaggregated rollout and tune batching for TPU throughput
✓Plug vLLM or SGLang-JAX in as the rollout engine and set sampling knobs for RL
✓Build multi-turn agentic RL with tool use, async rollout and trajectory grouping
✓Add a new RL algorithm by subclassing AlgorithmConfig and RLLearner and registering a custom loss
✓Launch multi-host jobs with the Tunix CLI and make them recoverable with checkpoints and fault tolerance

Highlights

•Covers the whole post-training ladder in one JAX codebase: SFT, LoRA/QLoRA, DPO, ORPO, PPO, GRPO, GSPO-Token, DAPO, Dr.GRPO and distillation
•White-box design: the docs show exactly which config and learner classes to subclass to change the training loop
•Runnable Colab notebooks for QLoRA, GRPO, DPO and logit distillation on small Gemma models, plus a DeepScaleR reproduction
•Actively maintained by Google (about 2.5k GitHub stars, 2,200+ commits, Gemma 4 support added April 2026), Apache-2.0

Who It’s For

Best For

✓ML engineers post-training open models on Google Cloud TPUs
✓Researchers who want a hackable GRPO/PPO implementation outside the PyTorch ecosystem
✓Teams already on JAX, Flax NNX or MaxText who need SFT, preference tuning and RL in the same stack

Prerequisites

•Working Python and basic JAX/Flax NNX familiarity
•Understanding of supervised fine-tuning and the idea behind policy-gradient RL such as PPO or GRPO
•Access to a TPU (Colab or Google Cloud) for the notebooks; GPU use is not the documented path

FAQ

What is Tunix Documentation — Lightweight LLM Post-Training Library in JAX?

Tunix (Tune-in-JAX) is Google's open-source documentation and library for post-training large language models in JAX, written for ML engineers who fine-tune open models on TPUs. After working through it you can run LoRA/QLoRA, DPO, distillation and GRPO jobs on Gemma, Llama or Qwen checkpoints and add your own RL algorithm.

Is Tunix Documentation — Lightweight LLM Post-Training Library in JAX free?

Tunix Documentation — Lightweight LLM Post-Training Library in JAX is free to access.

What level is Tunix Documentation — Lightweight LLM Post-Training Library in JAX for?

Tunix Documentation — Lightweight LLM Post-Training Library in JAX is aimed at a intermediate audience. Recommended background: Working Python and basic JAX/Flax NNX familiarity, Understanding of supervised fine-tuning and the idea behind policy-gradient RL such as PPO or GRPO, Access to a TPU (Colab or Google Cloud) for the notebooks; GPU use is not the documented path.

How long does Tunix Documentation — Lightweight LLM Post-Training Library in JAX take?

Expect roughly ~1-2 hours per Colab quick start; several days to adapt a GRPO or agentic RL run to your own model and reward. Most learners work through it at their own pace.

What will I learn from Tunix Documentation — Lightweight LLM Post-Training Library in JAX?

You'll learn: Fine-tune Gemma, Llama or Qwen with full weights, LoRA or QLoRA using the PeftTrainer; Run preference tuning with DPO and ORPO on paired chosen and rejected responses; Train a reasoning model with GRPO on GSM8K, including reward functions and pass@1 evaluation; Choose between collocated and disaggregated rollout and tune batching for TPU throughput; Plug vLLM or SGLang-JAX in as the rollout engine and set sampling knobs for RL; Build multi-turn agentic RL with tool use, async rollout and trajectory grouping; Add a new RL algorithm by subclassing AlgorithmConfig and RLLearner and registering a custom loss; Launch multi-host jobs with the Tunix CLI and make them recoverable with checkpoints and fault tolerance.

Topics

tunixjaxgrpopost-trainingtpulora

Sources

This page was written from 6 sources, 2 on domains other than tunix.readthedocs.io.