tunix
用 JAX 搭积木做 LLM 后训练，算法随便换
Tunix 是 Google 开源的一个轻量级大模型后训练（post-training）库，基于 JAX 生态构建，专注于在模型预训练完成后做进一步优化。它解决的核心问题是：现有后训练流程往往依赖庞大且耦合的框架，难以灵活替换算法组件。Tunix 把监督微调（SFT）、奖励建模、强化学习（如 PPO、GRPO 等）拆成模块化组件，让研究者可以像搭积木一样组合不同训练阶段。库内提供可复用的训练循环、数据管道和常见 RL 算法实现，并与 JAX/Flax 的自动微分和并行能力打通，方便在 TPU/GPU 上做分布式训练。整体定位偏研究友好，代码结构清晰，适合快速验证新的后训练策略，而不是开箱即用的产品级方案。
项目数据
使用教程
⚠ 该项目 README 未提供完整安装说明，以下步骤为通用流程，请以官方文档为准。
环境要求
- Python 环境（具体版本以官方文档为准）
- JAX/Flax NNX 生态依赖，含 Optax、Orbax 等
- 可选 TPU 或 GPU 计算资源，用于分布式训练
- 若需 rollout 推理，可配合 vLLM 或 SGLang-JAX
- Docker（仅镜像构建路径需要）
安装与启动步骤
-
1克隆仓库
从 GitHub 获取 Tunix 源码，后续安装与示例都基于该目录。注意保持网络可访问 GitHub。
git clone https://github.com/google/tunix.git && cd tunix -
2查看安装文档
README 未给出具体安装命令，需先访问官方文档 quickstart 的 Installation 章节，按最新说明操作。
-
3本地安装 TPU 依赖
在 TPU VM 或开发机上运行脚本，按 requirements 文件安装 vLLM 相关依赖，与 Docker 构建路径二选一。
bash scripts/install_tunix_vllm_requirement.sh -
4Docker 构建镜像
使用仓库根目录 Dockerfile 构建镜像，内部会从 requirements.txt 和 special_requirements.txt 安装固定依赖。
docker build -t tunix .
如何确认成功
按官方文档 Quick Start 运行一个训练示例，无导入或依赖报错即表示安装成功。
常见问题
Q：Tunix 如何安装？
A：README 未给出具体 pip 命令，需访问 https://tunix.readthedocs.io/en/latest/quickstart.html#installation 查看官方安装步骤。
Q：Docker 和安装脚本可以同时用吗？
A：不可以，README 明确说明这是两个独立入口；构建 Docker 镜像时无需在容器内再运行 install_tunix_vllm_requirement.sh。
Q：支持哪些硬件？
A：README 声明在 TPU 上有 SOTA 训练性能，也支持 GPU；具体依赖与并行配置以官方文档为准。
Q：可以做哪些训练任务？
A：支持监督微调（SFT）、强化学习（如 PPO、GRPO 等）、Agentic RL 以及奖励建模等后训练环节。
注意事项
- README 未给出 pip 安装命令，实际安装请以官方文档 quickstart 的 Installation 章节为准。
- Docker 构建与本地脚本是两种独立安装路径，避免重复安装依赖。
- 项目处于活跃开发中（V2 Release），接口与依赖可能变化。
核心亮点
- 模块化设计，SFT、奖励模型、RL 各阶段可独立替换组合
- 基于 JAX/Flax，天然支持 TPU/GPU 分布式训练与自动微分
- 由 Google 团队维护，代码结构清晰，适合研究新后训练算法
不足之处
- 偏研究定位，缺少端到端产品化封装，上手门槛较高
- 生态和文档相比 PyTorch 系后训练库仍较新，社区案例有限
适用场景
- 研究团队快速验证新的强化学习后训练算法
- 在 TPU 集群上对开源大模型做 SFT 与 RLHF 实验
- 需要灵活替换奖励模型或 RL 组件的定制化训练流程
替代项目
trl、verl、OpenRLHF
项目介绍
上一篇：slime
下一篇：create-llm
同类项目推荐
Long-RL
开源
让强化学习轻松驾驭超长序列，训练更稳更快
Long-RL: Scaling RL to Long Sequences (NeurIPS 2025)
TPA
开源
把注意力复杂度从平方降到线性，长序列不再卡顿
[NeurIPS 2025 Spotlight] TPA: Tensor ProducT ATTenTion Transformer (https://arxiv.or···
minimind
开源
两小时造出你的小模型，练手入门不求人
Train a 64M-parameter LLM from scratch in just 2h!
pruna
开源
一键优化 AI 模型，推理提速又省资源。
Pruna is a model optimization framework built for developers, enabling you to delive···