Fine-tune and RL-train LLMs without escaping JAX

What it does

Tunix handles the messy middle of LLM development: supervised fine-tuning (full weights and LoRA), preference optimization like DPO and ORPO, and reinforcement learning from PPO through several GRPO variants. It also covers agentic RL, with multi-turn tool use and asynchronous trajectory collection. The library acts as middleware between standard JAX utilities—Flax, Optax, Orbax—and optimized model implementations such as MaxText, with an eye on TPU clusters.

The interesting bit

Rather than treating rollout as an external chore, Tunix bakes in native vLLM and SGLang-JAX support on TPU for generating trajectories, while micro-batching and Pathways-based multi-host distribution aim to keep thousands of devices busy. It reflects Google’s attempt to build a unified post-training layer that stays entirely within the JAX stack rather than bridging out to PyTorch-centric tooling.

Key highlights

• Algorithm buffet: SFT, DPO, ORPO, PPO, GRPO, GSPO-Token, DAPO, and Dr.GRPO.
• Agentic RL with multi-turn tool use, async rollouts, and trajectory batching.
• TPU-native scale via Pathways, with checkpointing and fault tolerance built in.
• Tight integration with MaxText models and native vLLM / SGLang-JAX rollout on TPUs.
• Supports Gemma (including Gemma4), Llama, and Qwen families.

Caveats

• Explicitly labeled as early-stage and under active development; contribution guidelines are still being formalized.
• The entire value proposition assumes a JAX/XLA/TPU stack—PyTorch or GPU-centric shops should look elsewhere unless they plan to migrate.

Verdict

A sensible choice for teams already committed to JAX on TPUs who want one stop for SFT through agentic RL. If your stack is outside that world, this is mostly a curiosity for now.