maxtext: 一個用於可擴展 LLM 預訓練與後訓練的高性能 JAX 參考實現

它解決了什麼問題

MaxText 旨在為訓練大型語言模型 (LLMs) 提供高性能、可擴展的參考實現。它解決了在從單一主機到數萬個晶片的硬體集群中，無需複雜的手動優化即可實現高模型 FLOPs 利用率 (MFU) 和吞吐量 (tokens/second) 的挑戰。

運作原理

MaxText 使用純 Python 和 JAX 編寫，目標對象為 Google Cloud TPUs 和 GPUs。它利用了 JAX AI 技術棧，包括用於神經網絡的 Flax、用於後訓練的 Tunix、用於檢查點存儲的 Orbax、用於優化的 Optax 以及用於數據加載的 Grain。該庫提供了一系列高性能的模型架構（例如 Gemma、Llama、DeepSeek 和 Qwen），並支持從頭開始的預訓練以及使用監督式微調 (SFT) 和強化學習 (GRPO 和 GSPO) 等技術進行的可擴展後訓練。

對象是誰

它適用於正在構建宏大 LLM 項目的 AI 研究人員和生產工程師，這些人員需要一個可擴展且「無需優化」的框架，以便進行模型設計實驗、預訓練大規模模型，或對現有的開源或專有模型進行後訓練。

重點亮點

• 廣泛的模型支持：包括 Gemma 4、Llama 4、DeepSeek V3、Qwen 3.5 和 Kimi K2 等模型的參考實現。
• 高可擴展性：支持跨數萬個晶片的預訓練。
• 後訓練框架：透過 Tunix 提供 SFT 和 RL (GRPO/GSPO) 的集成支持，並在 RL 工作流中使用 vLLM 進行採樣。
• 硬體優化：專門針對 Google Cloud TPUs 和 GPUs 使用 XLA 編譯器進行了優化。
• 多模態能力：支持 Gemma 3、Gemma 4 和 Llama 4 VLMs 的多模態訓練。