tunix

What Tunix provides

Tunix, short for Tune-in-JAX, is a JAX based library designed to streamline the post-training of large language models. The README states it provides efficient and scalable support for supervised fine-tuning, reinforcement learning, and agentic reinforcement learning, with an emphasis on training performance on TPUs. It leverages JAX for accelerated computation and integrates with JAX based modeling frameworks such as Flax NNX, and it connects to high performance inference engines like vLLM and SGLang-JAX for rollout. The project is under active development and the README marks the current status as the V2 release. In the high level architecture, Tunix sits in the JAX training stack and positions itself to use foundational tools such as Flax, Optax, and Orbax for efficient model refinement. It acts as an intermediate layer between those core utilities and optimized models such as MaxText and MaxDiffusion, streamlining tuning workflows on top of the XLA and JAX infrastructure. The documentation is hosted separately at tunix.readthedocs.io, and the README links to pages for the design overview, algorithms, performance, rollout, reliability, and examples. The stated scope is post-training rather than pretraining, so the library targets teams that already have a base model and want to adapt it through fine-tuning or reinforcement learning rather than train one from scratch.

Training algorithms supported

The README lists a broad set of post-training algorithms organized into three groups. For supervised fine-tuning, Tunix supports full weights fine-tuning, parameter efficient fine-tuning through LoRA, and preference methods including DPO, which is direct preference optimization, and ORPO, which is odds ratio preference optimization. For reinforcement learning, it supports PPO, or proximal policy optimization, and GRPO, or group relative policy optimization, plus several variants: GSPO-Token for token level group sequence policy optimization, DAPO for direct alignment via preference optimization, and Dr.GRPO for another variant of GRPO. For agentic reinforcement learning, Tunix lists multi-turn tool use, asynchronous rollout for high throughput trajectory collection, and trajectory batching and grouping. The news section records specific milestones. In April 2026 the README says Gemma4 models are supported. In January 2026 it notes efficient kernel execution through splash attention and GMM mixture of experts. In December 2025 it says agentic RL training was released, with efficient support for multi-turn agent environment interaction, tool usage, and async rollout. These entries are presented as project updates rather than benchmark claims. The supported model list is described as growing and includes the Gemma, Llama, and Qwen families, with a documentation page for adding new models. The README links each algorithm to its original paper, so a reader can follow the method behind every variant the library implements.

Infrastructure and collaboration

The framework and infrastructure highlights cover modularity, performance, and stability. On modularity, the README says components are designed to be reusable and composable and easy to customize and extend. On performance, Tunix offers native vLLM and SGLang-JAX integration on TPU for rollout, native MaxText model integration for high performance kernels and model execution, and micro-batching support for component level efficient execution. On stability, it offers seamless multi-host distributed training with Pathways, which the README says can scale up to thousands of devices, and checkpointing with fault tolerance. The getting started section points to a separate installation page and describes two setup paths for TPU users who integrate vLLM and TPU inference: a Docker image build that installs pinned dependencies from requirement files, and a local TPU VM or developer machine install through a script that installs the same requirement files outside Docker. The README also notes a collaboration with GRL, a game reinforcement learning framework from the Hao AI Lab at UCSD, which integrates TPU support so users can run scalable, reproducible RL experiments such as PPO rollouts on Qwen2.5-0.5B-Instruct on TPU v4 meshes with minimal setup. Contributions are welcomed, and the README says the contribution process is still being formalized while pointing to a GitHub discussion forum for feature requests, issues, and questions. A citation entry in BibTeX format is provided for research use.

Editorial conclusion

The repository is published under the Apache 2.0 license and its most recent commit was recorded on 2026-08-26. The source lives at https://github.com/google/tunix.