Tunix: A JAX Post-Training Library for SFT, RL and Agentic RL

What Tunix Is For, and Who It Is Not For

Tunix solves a narrow problem: post-training an LLM after the base model exists, using JAX rather than PyTorch. The README describes it as a JAX based library for supervised fine-tuning, reinforcement learning and agentic RL, positioned inside the JAX training stack above Flax, Optax and Orbax, and below optimized models like MaxText and MaxDiffusion. That positioning is the whole pitch. If you are already writing Flax NNX models and running on TPUs, Tunix gives you algorithms and rollout plumbing without leaving JAX. If you are not, the library has little to offer you.

The intended audience is research and infrastructure engineers who treat TPU capacity as a given and want SFT, DPO, PPO or GRPO without assembling the loop themselves. The pyproject classifiers list Development Status :: 3 - Alpha, which is worth taking literally. The README also says the contribution process is still being formalized, so external contributors should expect friction. This is not a library for someone who wants a stable training API with long-term compatibility guarantees.

The JAX Stack Layering and the Rollout Split

The architecture has two distinct halves that matter for how you deploy it. The training half sits between foundational JAX utilities (Flax, Optax, Orbax) and optimized model implementations such as MaxText. Tunix does not ship its own kernels; it composes the ones below it. The README points to a Design Overview page for the full picture, and the repository layout reflects this: the tunix/ package holds algorithms and models, while examples/ holds notebooks and scripts that wire them together.

The second half is rollout, and this is where the design gets opinionated. Tunix integrates with vLLM and SGLang-JAX for generation during RL. That means the training process and the inference process are separate components that must be coordinated, which is why the README documents two install paths rather than one. For agentic RL, the library adds multi-turn tool use, asynchronous rollout for trajectory collection, and trajectory batching and grouping. The async rollout is the interesting part: collecting trajectories from an agent environment is latency-bound, so decoupling it from the training step is a reasonable design. It also means your environment code becomes part of the training loop's critical path, and the README does not describe how environment failures are surfaced.

Installing Tunix and Running a First SFT Job

The README does not inline installation commands. It points to a Quickstart page on the documentation site and states that TPU users integrating vllm and tpu-inference have two supported setup paths: a Docker image build using the repository Dockerfile, and a local TPU VM or developer-machine install using scripts/install_tunix_vllm_requirement.sh. The README is explicit that these are separate entry points and that you do not run the install script inside the Docker build.

The Dockerfile shows what the image does. It starts from python:3.12-slim, creates a virtual environment at /opt/venv, installs uv, copies scripts/install_tunix_vllm_requirement.sh and the requirements/ directory, runs that script, then installs the package itself with uv pip install .

FROM python:3.12-slim
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC
RUN python3.12 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
RUN pip install uv
COPY scripts/install_tunix_vllm_requirement.sh scripts/
COPY requirements/ requirements/
RUN bash scripts/install_tunix_vllm_requirement.sh

After the base install, the Dockerfile unconditionally adds SFT and MaxText dependencies. Note that MaxText is installed with --no-deps from a git URL, so it does not pull its own dependency tree.

RUN uv pip install --upgrade flax && \
    uv pip install torchax aqtp tokamax math_verify drjax && \
    uv pip install --no-deps git+https://github.com/google/maxtext.git

For a local install instead of Docker, the project provides the shell script. The README gives the path but not the invocation flags, so check the script before running it. The pyproject requires Python >=3.11 and pins jax[tpu]>=0.6.0,!=0.7.2,<0.11.1, with an inline comment explaining that 0.7.2 has a performance regression on OSS and that 0.11.1 removed a symbol Flax still subclasses. That pin is the single most likely source of install failure if your environment already has a different JAX.

pip install uv
bash scripts/install_tunix_vllm_requirement.sh
uv pip install .

Once installed, the README directs you to the Quickstart page for starting examples and to the Examples and Guides page for the full list. The examples directory contains notebooks for the main algorithms, including examples/sft/, examples/dpo_gemma.ipynb, examples/grpo_gemma.ipynb, examples/qlora_gemma.ipynb and examples/qlora_llama3_gpu.ipynb. Reading one of those notebooks is the fastest way to see the actual training loop, because the README itself does not show one.

Where Tunix Gets in the Way

The dependency pins are the first real constraint. jax[tpu]>=0.6.0,!=0.7.2,<0.11.1 excludes a specific release for performance reasons and caps the upper bound for a Flax compatibility reason. transformers!=4.57.2 is excluded because of a linked issue, and safetensors is capped below 0.8 because higher versions break with jax >=0.9. A library that pins its transitive dependencies this tightly is telling you it has not yet absorbed the churn in its own stack. If you need a JAX version outside that range for another workload, you cannot share the environment.

The second constraint is hardware. The package name in pyproject is google-tunix and the JAX extra is jax[tpu]. The README's performance claims are framed around TPUs and Pathways scaling to thousands of devices. There is a GPU example notebook (examples/qlora_llama3_gpu.ipynb), so GPU work is not absent, but the documented setup paths and the Dockerfile are built for the TPU plus vLLM plus MaxText combination. Running this on a single GPU means departing from the documented path.

The third issue is maturity. The pyproject classifier says Alpha. The README says the contribution process is still being formalized. Model support is described as a growing list rather than a fixed one, and the news section shows Gemma4 support arriving in 2026/04 with training recipes still to come. If you need a model that is not on the supported list, the README points to a Models page covering how to add new ones, which means writing model code yourself. That is a meaningful amount of work to take on for a training run.

Tunix Against a PyTorch Post-Training Stack

The obvious alternative for most teams is a PyTorch-based post-training framework, and the difference is not the algorithm list. PPO, GRPO, DPO and SFT exist in both worlds, and the papers Tunix cites (PPO, GRPO, DAPO, ORPO, GSPO-Token, Dr.GRPO) are framework-agnostic. The difference is the execution model. Tunix composes JAX transformations and hands the compiled graph to XLA, which is what makes the TPU story coherent: MaxText provides the kernels, Pathways handles multi-host distribution, and vLLM or SGLang-JAX handles rollout. A PyTorch stack gets its performance from CUDA kernels and a different distributed strategy, and its ecosystem of trainers is larger because the user base is larger.

So the trade is concrete. Choose Tunix when your models are already Flax NNX and your compute is TPU, because rewriting a model into a different framework costs more than the library saves. Choose a PyTorch trainer when your team's expertise, your checkpoints and your serving path are all PyTorch, because Tunix will not meet you halfway. The README's collaboration note with GRL, a game RL framework from UCSD's Hao AI Lab, shows the intended integration pattern: another project brings the environment and Tunix brings the TPU training loop. That is a reasonable division of labor, and it also tells you Tunix expects to be one component in a larger system rather than the whole pipeline.

Maintenance Cadence, Licensing and Upgrade Cost

The last push to the repository was on 2026-06-11, which is the same date as the v0.1.7 release. The release history shows v0.1.6 in 2026-03-13 and v0.1.5 in 2025-11-21, so the cadence over that window is roughly one tagged release per quarter, with the project describing itself as under active development. The pyproject version is 0.1.8, ahead of the most recent tagged release, which is normal for a main branch but means installing from the repository gives you code that has not been tagged.

Upgrade cost is dominated by the pinned dependencies rather than by API churn. Because jax[tpu] is capped below 0.11.1 and transformers excludes 4.57.2, moving Tunix forward may require moving other parts of your stack, and the inline comments in pyproject show the maintainers tracking those conflicts explicitly. The v0.1.5 release title, Critical Issue Fix for v0.1.4, is a reminder that early versions can carry defects serious enough to warrant a dedicated patch release. Budget for reading release notes before upgrading rather than assuming compatibility.

The licence is Apache-2.0, declared both in the LICENSE file and in pyproject. That is a permissive licence with an explicit patent grant, which is generally straightforward for commercial use, but the repository also pulls in dependencies under their own terms, including MaxText installed from a git URL. Whether those terms suit your distribution model is a question for your own review, not something the README answers.

Editorial conclusion

Adopt Tunix if your training stack is already JAX and Flax NNX and you are targeting TPUs, because the library assumes that environment rather than adapting to yours. Do not adopt it if you are running PyTorch on GPUs and want a drop-in trainer; the dependency set pins jax[tpu] and the Dockerfile installs vLLM and MaxText for a TPU workflow. Before committing, verify that the model you need appears in the supported model list, check that your JAX version satisfies the jax[tpu]>=0.6.0,!=0.7.2,<0.11.1 pin, and read the reliability page to confirm what checkpointing and fault tolerance actually cover.

Frequently asked questions

Official sources

1. Official README
2. Project repository
3. Release notes