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 等主流模型家族。