google-tunix

What it is and what it does

Tunix is a JAX-native post-training framework designed to streamline fine-tuning and reinforcement learning workflows for large language models. It sits as an intermediate layer in the JAX training stack, integrating with Flax, Orbax, and high-performance inference engines like vLLM and SGLang-JAX for rollout generation. The library provides modular, composable components for supervised fine-tuning (full-weight and parameter-efficient variants), preference optimization (DPO, ORPO), reinforcement learning algorithms (PPO, GRPO, GSPO-Token, DAPO, Dr.GRPO), and agentic RL with multi-turn tool use and asynchronous trajectory collection.

The package is optimized for TPU execution with native integration into XLA and Pathways infrastructure, supporting seamless multi-host distributed training at scale. It includes micro-batching for efficient component-level execution, checkpointing and fault tolerance, and support for a growing list of model families (Gemma, Llama, Qwen). Tunix is under active development (v2 release, alpha status) with ongoing expansion of capabilities.

Use it for:

• Fine-tune open-source LLMs on custom instruction or domain-specific data using supervised fine-tuning.
• Run preference optimization workflows (DPO, ORPO) to align model outputs with human preferences.
• Train models with PPO or GRPO on reward signals for complex reasoning or task-specific behavior optimization.
• Build agentic RL systems where models learn to use tools and interact with environments over multiple turns.
• Integrate with JAX-based modeling frameworks and inference engines to streamline end-to-end post-training pipelines.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Tunix is a JAX-based library for post-training large language models with support for supervised fine-tuning, reinforcement learning, and agentic RL workflows, optimized for TPU execution.

Yes, if you are post-training LLMs on TPU infrastructure and want a modular, actively-maintained JAX-native framework. The low install friction and permissive license are favorable. However, the 26 runtime dependencies and alpha status mean you should verify that optional dependencies match your use case and accept that the API may evolve.

Install

google-tunix on PyPI

pip install google-tunix

uv add google-tunix

poetry add google-tunix

Installing google-tunix

Before you install

Low install friction with a pure-Python wheel. The package is actively maintained (last commit 2026-08-14, 64 days since release) and in alpha status. However, it declares 26 runtime dependencies including heavy scientific stacks (JAX via google-metrax, Flax, TensorFlow, transformers), which will pull in substantial transitive dependencies.

License in practice

Apache-2.0 permissive license allows commercial and private use, modification, and redistribution with minimal restrictions—suitable for most production and research contexts.

Quickstart

pip install google-tunix

import tunix
from tunix.algorithms import supervised_finetune
# Configure and run SFT on your model and dataset

Requires Python 3.11 or later. JAX and its dependencies (including google-metrax) may require additional system libraries and are typically used with TPU or GPU hardware.

Verify before relying

• Whether the 26 runtime dependencies are all strictly required or some are optional/conditional based on use case.
• Whether Tunix can be used effectively on GPUs or if TPU is a hard requirement for documented performance.
• Stability and API surface maturity given the alpha development status and active feature expansion.

Package facts

Evidence: google_tunix-0.1.7-py3-none-any.whl

Tags

More Artificial Intelligence packages

Further reading