A Jax-Native LLM Post-Training Library (github.com)

Tunix (Tune-in-JAX) is a new JAX-native library for post-training LLMs that consolidates supervised fine-tuning, reinforcement learning, preference tuning and knowledge distillation into a single, TPU- and JAX-optimized stack. Announced as an early-development project, Tunix targets researchers and engineers who use Flax/NNX and TPU meshes by providing native support for common sharding strategies (data parallel, FSDP, tensor parallel), distributed multi-host training, and optimized rollouts (vLLM). Its modular design and prebuilt recipes aim to make large-scale, reproducible post-training workflows—especially RLHF-style pipelines—much easier to run and extend on accelerators. Technically, Tunix already implements full-weight and PEFT fine-tuning (LoRA/Q-LoRA), PPO plus advanced group/token-level policy optimizers (GRPO, GSPO-token), Direct Preference Optimization (DPO), and several distillation strategies (logit matching, attention/feature transfer and projection). It also supports agentic RL features—async rollouts, multi-turn/multi-step training and tool use—and integrates with GRL (Hao AI Lab/UCSD) to run scalable game-based RL experiments on TPU v4 meshes. Installable via PyPI or GitHub, Tunix is positioned as a practical, extensible toolkit for scaling post-training research in the JAX ecosystem; features, docs and contribution processes are still maturing.