tunix

Tunix 是一个基于 JAX 的轻量级大语言模型（LLM）后训练库，专注于高效、可扩展的模型调优。它支持监督微调（SFT）、基于偏好的强化学习（如 DPO、ORPO）和策略优化算法（如 PPO、GRPO），并原生集成 Flax NNX、vLLM 和 SGLang-JAX 等生态工具，在 TPU 上实现高性能训练。项目采用模块化设计，适配 XLA 编译与 Orbax 检查点管理，适用于需要在 JAX 生态中快速迭代 LLM 对齐与能力增强的研究与工程场景。