6.0
关注
AI SCORE
开源项目2024-04-24 11:00
Maxtext：JAX高性能大模型训练框架
Hacker News · github.com#JAX#LLM#模型训练
Editor brief · 编辑速览
用JAX构建的可扩展LLM训练框架，强调简洁性和训练效率；适合JAX用户和想探索高性能框架的研究开发者。
注意：我们推荐使用 Python 3.12 运行 MaxText，因为它是我们的主要支持版本。其他 Python 版本可能会遇到兼容性问题。
MaxText 是一个高性能、高度可扩展的开源 LLM 库和参考实现，用纯 Python/JAX 编写，针对 Google Cloud TPUs 和 GPUs 的训练进行了优化。
MaxText 提供了一个高性能模型库供选择，包括 Gemma、Llama、DeepSeek、Qwen 和 Mistral。对于这些模型中的每一个，MaxText 支持预训练（最多支持数万个芯片）和可扩展的后训练，采用流行的技术如监督微调（SFT）和组相对策略优化（GRPO，一种强化学习）以及组序列策略优化（GSPO，一种强化学习）。
MaxText 从单机到大规模集群都能实现高模型浮点操作利用率（MFU）和 tokens/秒，同时保持简洁，无需复杂优化，这得益于 JAX 和 XLA 编译器的强大能力。
MaxText 是研究和生产环境中雄心勃勃的 LLM 项目的起点。我们鼓励你先开箱即用地尝试 MaxText，然后 fork 并修改 MaxText 以满足你的需求。
查看我们的 Read The Docs 站点或开始你的第一次 MaxText 运行。如果你对扩散模型（Wan 2.1、Flux 等）感兴趣，请参阅我们 AI Hypercomputer GitHub 组织中的 MaxDiffusion 仓库。
查看我们的安装指南，通过 pip 从 PyPI 安装 MaxText。
查看我们关于在解耦模式下运行 MaxText 的指南，无需任何 GCP 依赖。详见 Decoupled Mode Guide。
[2026年7月15日] 现已支持 Qwen3 30B 和 GPT-OSS 20B 的 RL。详见 Qwen3 30B RL 教程和 GPT-OSS 20B RL 教程。
[2026年6月9日] 现已支持 Qwen3.5 35B & 397B。
[2026年5月29日] 新的评估框架现已推出，可用于对 MaxText 检查点运行 lm-eval、evalchemy 和自定义基准测试。查看评估指南开始使用。
[2026年5月6日] 现已支持 Kimi-K2-Thinking、Kimi-K2.5（文本）和 Kimi-K2.6（文本）。详见 Run_Kimi.md。
[2026年4月18日] 新增笔记本 maxtext_with_gepa.ipynb，用于使用 MaxText 的 GEPA 框架优化 AIME 提示。
[2026年4月14日] 已移除旧的 MaxText.* 后训练垫片。请参考 src/MaxText/README.md 了解新命令位置和迁移方式。
[2026年4月13日] 现已支持 Kimi-K2，以及 MuonClip 优化器。尝试 kimi-k2-1t 配置并查看用户指南。
[2026年4月10日] 现已支持 DeepSeek-V3.2，具备用于长上下文的 DeepSeek Sparse Attention。使用 deepseek3.2-671b 配置试试吧。详见用户指南。
[2026年4月2日] 现已支持 Gemma 4 多模态模型（26B MoE、31B 密集）！使用我们的 gemma4-26b 和 gemma4-31b 配置试试吧。详见 Run_Gemma4.md。
[2026年3月6日] 现已支持来自 DeepSeek-AI 的新功能：可扩展查询的条件记忆（Engram）和流形约束超连接（mHC）。使用我们的 deepseek-custom 入门配置试试吧。
[2026年3月5日] PyPI 中的新 tpu-post-train 目标。请也将此安装选项用于运行 vllm_decode。详见 MaxText 安装说明。
[2026年3月5日] 现已支持 Qwen3-Next。
[2026年2月27日] 新的 MaxText 结构！MaxText 已根据 RESTRUCTURE.md 进行了重构。请随时分享你的想法和反馈。
[2025年12月22日] 现已支持 Muon 优化器。
[2025年12月10日] 现已支持 DeepSeek V3.1。使用 DeepSeek V3 671B 的现有配置并加载 V3.1 检查点即可使用该模型。
[2025年12月9日] 新的 RL 和 SFT 笔记本教程现已推出。
[2025年12月4日] ReadTheDocs 文档站已重新组织。
[2025年12月3日] 现已推出通过新 RL 教程支持 GSPO 和 GRPO 的多主机功能。
[2025年11月20日] 新指南《MaxText 中的后训练是什么？》现已推出。
[2025年11月6日] 发布 Ironwood TPU 协同设计的 AI 栈。阅读关于其与 MaxText 协同设计的博客文章。
[2025年10月29日] 优化模型分层文档已更新。
[2025年10月12日] 添加版本控制。查看我们的第一组发布说明！
[2025年10月10日] 现已通过 Tunix 推出后训练（SFT、RL）。
[2025年9月26日] MaxText 现已支持词汇表平铺（PR）！调整配置 num_vocab_tiling 以解锁更高效的内存使用。
[2025年9月24日] 现已支持 GPT-OSS 模型族（20B、120B）。
[2025年9月15日] MaxText 现已作为 PyPI 包推出。用户现在可以通过 pip 安装 maxtext。
[2025年9月5日] MaxText 已移至 src 布局作为 RESTRUCTURE.md 的一部分。对于现有环境，请从 MaxText 根目录运行 pip install -e .。
[2025年8月13日] 现已支持 Qwen3 2507 MoE 模型族：MoEs: 235B Thinking & 480B Coder 以及现有密集模型：0.6B、4B、8B、14B 和 32B。
[2025年7月27日] 更新了 TFLOPS/s 计算（PR）以考虑因果注意力，将注意力浮点运算除以二。考虑了滑动窗口和分块注意力减少的注意力浮点运算，详见相关 PR。这些变更影响大序列配置，如本文档所述。
[2025年7月16日] 我们将重新组织 MaxText 仓库以改进组织和清晰度。请查看提议的结构并提供反馈。
[2025年7月11日] 多令牌预测（MTP）训练支持！添加了一个基于预测多个未来令牌的辅助损失，灵感来自 DeepSeek-V3 论文，以提高训练效率。
[2025年6月25日] 现已支持 DeepSeek R1-0528 变体。
[2025年4月24日] 现已支持 Llama 4 Maverick 模型。
MaxText 提供了一个模型库，并展示了如何以高性能和规模执行预训练或后训练。
MaxText 利用 JAX AI 库，通过使用 Flax（神经网络）、Tunix（后训练）、Orbax（检查点）、Optax（优化）和 Grain（数据加载），展示了大规模训练的完整和综合演示。
除了纯文本 LLM 外，我们还支持使用 Gemma 3、Gemma 4 和 Llama 4 VLM 的多模态训练。
如果你从头开始构建模型，MaxText 可以作为实验、构思和灵感的参考实现——只需 fork 并修改 MaxText 来训练你的模型，无论是像 Llama 8B 这样的小密集模型，还是像 DeepSeek-V3 这样的大型 MoE。尝试配置和模型设计，以在 TPU 或 GPU 上构建最高效的模型。
MaxText 为如何在分片、量化和检查点等多个维度上实现最优性能提供了有主见的实现。
如果你正在对模型进行后训练，无论是专有的还是开源的，MaxText 使用 Tunix 提供了一个可扩展的框架。对于 RL（如 GRPO），我们利用 vLLM 进行采样，并利用 Pathways（即将推出）进行多主机。
我们的目标是提供各种模型（维度"a"）和技术（维度"b"），这样你可以轻松探索（a）*（b）的组合，并高效地训练最适合你用例的完美模型。
查看这些入门指南：
MaxText 旨在为你提供最好的开源模型，无论是作为参考实现，还是用于后训练，然后用 vLLM 进行服务。
详见 MaxDiffusion（LTXV、Wan 2.1、Flux、SDXL 等）
请加入我们的 Discord 频道，如果你有反馈，可以在这里提交功能请求、文档请求或错误报告。