Tunix：轻量级大语言模型后训练库

Tunix（Tune-in-JAX） 是一个基于 JAX 的库，旨在简化大型语言模型的后训练流程。它为以下任务提供了高效且可扩展的支持：

• 在 TPU 上实现 SOTA 训练性能
• 监督微调
• 强化学习 (RL)
• 智能体强化学习

Tunix 利用 JAX 的强大计算能力，实现了加速运算，并与基于 JAX 的建模框架（如 Flax NNX）无缝集成，同时与高性能推理引擎（如 vLLM 和 SGLang-JAX）结合用于回放。有关详细文档，请参阅 Tunix 官网。

当前状态：V2 版本发布

Tunix 目前仍在积极开发中。我们的团队正致力于扩展其功能、易用性和性能。敬请期待未来的更新和新特性！请参阅 演讲与公告 以获取最新动态、演讲和博客文章。

高层次架构

Tunix 是 JAX 训练栈中的先进后训练库，能够充分利用 Flax、Optax、Orbax 等基础工具，实现高效的模型优化。它位于这些核心工具与 MaxText、MaxDiffusion 等优化模型之间，作为中间层，在 XLA 和 JAX 基础设施之上简化调优工作流。更多架构细节请参阅 设计概述。

核心特性

• 监督微调 (SFT)：
  • 全参数微调
  • PEFT（参数高效微调）
  • DPO（直接偏好优化）
  • ORPO（几率比偏好优化）
• 强化学习 (RL)：
  • PPO（近端策略优化）
  • GRPO（群体相对策略优化）
  • GSPO-Token（基于 token 的群体序列策略优化）
  • DAPO（通过偏好优化进行直接对齐）
  • Dr.GRPO（分布鲁棒 GRPO）
• 智能体强化学习：
  • 多轮工具使用
  • 异步回放以实现高吞吐量轨迹收集
  • 轨迹批处理与分组

最新消息

• [2026年1月] Tunix 模型现已支持高效的内核执行（splash attn, GMM MoE）。
• [2025年12月] 智能体强化学习训练 已发布，高效支持多轮智能体与环境交互、工具使用、异步回放等功能。

框架与基础设施亮点

• 模块化：
  • 组件设计为可重用和可组合
  • 易于自定义和扩展
• 性能与效率：
  • 原生集成 vLLM 和 SGLang-JAX 以实现高效的 TPU 回放
  • 原生集成 MaxText 模型，提供高性能内核和模型执行
  • 支持 微批次，实现组件级别的高效执行
• 稳定性：
  • 通过 Pathways 实现无缝的多主机分布式训练，可扩展至数千台设备
  • 检查点与容错

开始使用

安装： 请跳转至 安装 页面，安装 Tunix 并运行您的第一个训练任务。

示例： 为了帮助您快速上手，我们准备了大量详细的示例和教程。您可以查看 快速入门 获取一组优秀的入门示例，以及 示例与指南 查看我们提供的所有笔记本和示例的完整列表。

支持的模型

Tunix 支持不断增长的模型列表，包括 Gemma、Llama 和 Qwen 系列。请参阅 模型 页面，了解完整列表及如何添加新模型的说明。

贡献与反馈

我们欢迎各方贡献！由于 Tunix 尚处于早期开发阶段，贡献流程仍在逐步完善中。详细的贡献流程已在 这里 说明。在此期间，您可以通过我们的 Tunix GitHub 讨论区 提出功能请求、报告问题或提问。

合作与伙伴关系

由 UCSD 的 Hao AI Lab 开发的 GRL（游戏强化学习），是一个开源框架，用于通过多轮强化学习在具有挑战性的游戏中对大型语言模型进行后训练。借助与 Tunix 的合作，GRL 集成了无缝的 TPU 支持，使用户能够快速运行可扩展、可重复的强化学习实验（例如在 Qwen2.5-0.5B-Instruct 上进行 PPO 回放），只需在 TPU v4 网格上进行 最小化设置。这一合作使社区能够进一步提升 LLM 的能力，将 Tunix 的优化 TPU 运行时与 GRL 的灵活游戏强化学习流水线相结合，从而推动前沿研究并确保易于复现。

引用 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}},
}

致谢

感谢所有优秀的贡献者！

Tunix 快速上手指南

环境准备

• 系统要求：Linux (推荐)，Python 3.8+
• 前置依赖：确保已安装 JAX 和 JAXlib（pip install jax jaxlib）。TPU 支持需 Google Cloud 环境；本地开发建议使用 CPU/GPU。

安装步骤

使用国内镜像加速安装（推荐）：

pip install git+https://github.com/google/tunix.git -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下为最简 SFT（监督微调）示例：

from tunix import sft

# 加载模型（支持 Gemma、Llama 等）
model = sft.load_model("gemma-2b")
# 初始化训练器
trainer = sft.Trainer(model)
# 开始训练（替换 data 为实际数据路径）
trainer.train(data="path/to/your/data")