Tunix: A Lightweight LLM Post-Training Library

Tunix: A Lightweight LLM Post-Training Library#

Tunix (Tune-in-JAX) is a JAX based library designed to streamline the post-training of Large Language Models. It provides efficient and scalable support for:

  • SOTA Training performance on TPUs

  • Supervised Fine-Tuning

  • Reinforcement Learning (RL)

  • Agentic RL

Tunix leverages the power of JAX for accelerated computation and seamless integration with JAX-based modeling frameworks like Flax NNX, and integrates with high-performance inference engines like vLLM and SGLang-JAX for rollout.

Current Status: V2 Release

Tunix is under active development. Our team is actively working on expanding its capabilities, usability and performance. Stay tuned for upcoming updates and new features! See Talks and Announcements for latest updates, talks, and blog posts.

High Level Architecture#

Tunix serves as a state-of-the-art post-training library within the JAX training stack, positioned to leverage foundational tools like Flax, Optax, Orbax, etc. for efficient model refinement. It sits as an intermediate layer between these core utilities and optimized models like MaxText and MaxDiffusion, streamlining tuning workflows on top of the XLA and JAX infrastructure.

Tunix in JAX ecosystem

See Design Overview for more details on the architecture.

Key Features#

  • Supervised Fine-Tuning (SFT):

    • Full Weights Fine-Tuning

    • PEFT (Parameter-Efficient Fine-Tuning)

    • DPO (Direct Preference Optimization)

      • ORPO (Odds Ratio Preference Optimization)

  • Reinforcement Learning (RL):

    • PPO (Proximal Policy Optimization)

    • GRPO (Group Relative Policy Optimization)

      • GSPO-Token (Token-level Group Sequence Policy Optimization)

      • DAPO (Direct Alignment via Preference Optimization)

      • Dr.GRPO (Distributionally Robust GRPO)

  • Agentic RL:

    • Multi-turn tool use

    • Asynchronous rollout for high-throughput trajectory collection

    • Trajectory batching and grouping

Framework & Infra Highlights#

  • Modularity:

    • Components are designed to be reusable and composable

    • Easy to customize and extend

  • Performance & Efficiency:

    • Native vLLM and SGLang-JAX on TPU integration for performant rollout

    • Native MaxText model integration for high performance kernels and model execution

    • Micro-batching support for component level efficient execution

  • Stability

Get Started#

Jump to Quick Start to install Tunix and run your first training job.

Supported Models#

Tunix supports a growing list of models including Gemma, Llama, and Qwen families. See Models for a full list and details on how to add new ones.

Citing Tunix#

@misc{tunix2025,
  title={Tunix (Tune-in-JAX)},
  author={Bao, Tianshu and Carpenter, Jeff and Chai, Lin and Gao, Haoyu and Jiang, Yangmu and Noghabi, Shadi and Sharma, Abheesht and Tan, Sizhi and Wang, Lance and Yan, Ann and Yu, Weiren and others},
  year={2025},
  howpublished={\url{https://github.com/google/tunix}},
}

Acknowledgements#

Thank you to all our wonderful contributors!

Contributors