google/tunix

何を解決するか

Tunixは、大規模言語モデル（LLMs）の後学習を簡素化し、高速化するためのJAXベースのライブラリです。さまざまなアライメントおよび最適化技術を通じてモデルを微調整するスケーラブルなフレームワークを提供し、TPU上で高いパフォーマンスを発揮するように特に最適化されています。

動作方法

TunixはJAXエコシステムにおける中間層として機能し、Flax、Optax、Orbaxといった基盤ツールと統合します。これらのツールを、MaxTextのような高性能なモデル実装やvLLM、SGLang-JAXのような推論エンジンと接続することで、効率的なロールアウトとモデル実行を処理します。

対象ユーザー

JAXとTPUインフラを用いてLLMのスケーラブルな後学習を実施する必要がある研究者や開発者を対象としています。特に、教師あり微調整（SFT）や強化学習（RL）に取り組んでいる方々に適しています。

主な特徴

• 多様な学習アルゴリズム: 教師あり微調整（SFT）、直接的好み最適化（DPO）、およびPPO、GRPO、GSPO-Tokenを含むさまざまな強化学習（RL）手法をサポート。
• エージェント型RL: 複数ターンのツール利用、高スループットの非同期ロールアウト、トラジェクトリーバッチ処理をサポート。
• TPU最適化: vLLM、SGLang-JAX、MaxTextとのネイティブ統合により、TPUハードウェア上で最先端のパフォーマンスを実現。
• スケーラビリティ: Pathwaysを介したシームレスなマルチホスト分散学習を可能にし、数千デバイスまでスケーリング可能。
• モデル互換性: Gemma、Llama、Qwenなど、人気のあるモデルファミリーと互換性があります。