maxtext: スケーラブルなLLMの事前学習および事後学習のための高性能JAXリファレンス実装

何を解決するか

MaxTextは、大規模言語モデル（LLM）をトレーニングするための、高性能でスケーラブルなリファレンス実装を提供するために設計されています。単一のホストから数万個のチップに及ぶハードウェアクラスター全体で、複雑な手動の最適化を必要とせずに、高いモデルFLOPs利用率（MFU）とスループット（tokens/second）を達成するという課題に対処します。

仕組み

純粋なPythonとJAXで記述されたMaxTextは、Google CloudのTPUおよびGPUをターゲットとしています。Flax（ニューラルネットワーク用）、Tunix（事後学習用）、Orbax（チェックポインティング用）、Optax（最適化用）、Grain（データロード用）を含むJAX AIスタックを活用しています。このライブラリは、一連の高性能なモデルアーキテクチャ（Gemma、Llama、DeepSeek、Qwenなど）を提供し、スクラッチからの事前学習と、Supervised Fine-Tuning（SFT）やReinforcement Learning（GRPOおよびGSPO）などの手法を用いたスケーラブルな事後学習の両方をサポートしています。

対象者

モデル設計の実験、大規模モデルの事前学習、または既存のオープンソースやプロプライエタリなモデルの事後学習を行うために、スケーラブルで「最適化不要」なフレームワークを必要とする、野心的なLLMプロジェクトを構築しているAI研究者やプロダクションエンジニアを対象としています。

ハイライト

• 幅広いモデルサポート: Gemma 4、Llama 4、DeepSeek V3、Qwen 3.5、Kimi K2などを含むリファレンス実装を提供します。
• 高スケーラビリティ: 数万個のチップにわたる事前学習をサポートします。
• 事後学習フレームワーク: Tunixを介したSFTおよびRL（GRPO/GSPO）の統合サポートを提供し、RLワークフローでのサンプリングにはvLLMが使用されます。
• ハードウェア最適化: XLAコンパイラを使用して、Google CloudのTPUおよびGPU向けに特別に最適化されています。
• マルチモーダル機能: Gemma 3、Gemma 4、およびLlama 4 VLMsのマルチモーダル学習をサポートします。