tunix

Tunix: 一款轻量级 LLM 后训练库 Tunix (Tune-in-JAX) 是一款基于 JAX 的库,旨在简化大型语言模型的后训练工作。它提供了高效且可扩展的支持,包括: - 在 TPU 上的 SOTA 训练性能 - 监督式微调 - 强化学习 (RL) - 代理式 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 (邻近策略优化)