google/tunix

解决的问题

Tunix 是一个基于 JAX 的库，旨在简化并加速大语言模型（LLMs）的后训练过程。它提供了一个可扩展的框架，通过多种对齐和优化技术来微调模型，特别针对 TPU 上的高性能进行了优化。

工作原理

Tunix 作为 JAX 生态系统中的中间层，集成了 Flax、Optax 和 Orbax 等基础工具。它将这些工具与高性能模型实现（如 MaxText）以及推理引擎（如 vLLM 和 SGLang-JAX）连接，以高效处理模型推演和执行。

适用人群

适用于需要使用 JAX 和 TPU 基础设施进行 LLM 可扩展后训练的研究人员和开发者，包括从事监督微调和强化学习的人员。

主要亮点

• 多样化的训练算法：支持监督微调（SFT）、直接偏好优化（DPO）以及多种强化学习（RL）方法，包括 PPO、GRPO 和 GSPO-Token。
• 智能体强化学习：支持多轮工具使用、异步推演以实现高吞吐量轨迹收集，以及轨迹批处理。
• TPU 优化：与 vLLM、SGLang-JAX 和 MaxText 原生集成，实现 TPU 硬件上的顶尖性能。
• 可扩展性：通过 Pathways 实现无缝多主机分布式训练，可扩展至数千个设备。
• 模型支持：兼容 Gemma、Llama 和 Qwen 等主流模型家族。