MaxText

注意: 我们建议使用 Python 3.12 运行 MaxText，因为它是我们主要支持的版本。其他 Python 版本可能会遇到兼容性问题。

MaxText 是一个高性能、高度可扩展的开源大模型库及参考实现，完全用 Python 和 JAX 编写，专为在 Google Cloud 的 TPU 和 GPU 上进行训练而设计。

MaxText 提供了一系列高性能模型供选择，包括 Gemma、Llama、DeepSeek、Qwen 和 Mistral 等。对于这些模型，MaxText 支持预训练（最多可扩展至数万台芯片）以及可扩展的后训练，涵盖监督微调（SFT）、组相对策略优化（GRPO，一种强化学习方法）和组序列策略优化（GSPO，另一种强化学习方法）等常用技术。

借助 JAX 和 XLA 编译器的强大功能，MaxText 在从单机到超大规模集群的场景下，都能实现高模型浮点运算利用率（MFU）和每秒处理的 token 数，同时保持代码简洁且基本无需手动优化。

MaxText 是科研和生产领域中雄心勃勃的大模型项目的起点。我们鼓励您先从开箱即用的 MaxText 开始实验，然后根据自身需求对 MaxText 进行 fork 和修改。

请访问我们的 Read The Docs 网站 或直接按照 快速入门指南 进行您的首次 MaxText 运行。如果您对扩散模型（如 Wan 2.1、Flux 等）感兴趣，请查看我们 AI Hypercomputer GitHub 组织下的 MaxDiffusion 仓库。

安装

请参阅我们的安装指南，了解如何通过 PyPI 使用 pip 安装 MaxText。

解耦模式

请参阅我们的解耦模式运行指南，在不依赖任何 GCP 服务的情况下运行 MaxText：解耦模式指南。

🔥 最新消息 🔥

• [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笔记本教程已在MaxText示例目录中提供。
• [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系列模型：MoE包括235B Thinking和480B Coder，以及现有的密集模型0.6B、4B、8B、14B和32B。
• [2025年7月27日] 更新了TFLOPS/s计算（PR），以考虑因果注意力机制，将注意力浮点运算量减半。在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 提供一个可扩展的框架。对于强化学习（如 GRPO），我们利用 vLLM 进行采样，并将很快借助 Pathways 实现多主机支持。

我们的目标是提供多种模型（维度“a”）和多种技术（维度“b”），以便您可以轻松探索 (a) * (b) 的组合，并高效地训练出最适合您应用场景的完美模型。

请查看以下入门指南：

• 监督微调（SFT）
  • 单主机 TPU 上的 SFT
  • 多主机 TPU 上的 SFT
• 强化学习（RL）
  • 单主机 TPU 上的 RL
  • 多主机 TPU 上的 RL

模型库

MaxText 致力于为您提供最佳的开源模型，无论是作为参考实现，还是用于后训练后再通过 vLLM 提供服务。

MaxText 支持的 JAX 模型

• Google
  • Gemma 4（26B MoE，31B 密集）
  • Gemma 3（4B、12B、27B）
  • Gemma 2（2B、9B、27B）
  • Gemma 1（2B、7B）
• Alibaba
  • Qwen 2.5（1.5B、7B、14B）
  • Qwen 3 MoE 2507（235B、480B）
  • Qwen 3 MoE（30B、235B）
  • Qwen 3 密集（0.6B、1.7B、4B、8B、14B、32B）
• DeepSeek
  • DeepSeek V3.1（671B）
  • DeepSeek V3 0324（671B）及 DeepSeek R1 0528（671B）
  • DeepSeek V2（16B、236B）
• Meta
  • Llama 4 Scout（109B）及 Maverick（400B）
  • Llama 3.3 70B、3.1（8B、70B、405B）、3.0（8B、70B、405B）
  • Llama 2（7B、13B、70B）
• OpenAI
  • GPT-OSS（20B、120B）
  • GPT3（52K、6B、22B、175B）
• Mistral
  • Mixtral（8x7B、8x22B）
  • Mistral（7B）
• 扩散模型
  • 详见 MaxDiffusion（LTXV、Wan 2.1、Flux、SDXL 等）

参与进来

欢迎加入我们的 Discord 频道，如果您有任何反馈，可以在此处提交功能请求、文档请求或错误报告：这里。

许可证

Apache 2.0 许可证

MaxText 快速上手指南

MaxText 是一个基于纯 Python 和 JAX 编写的高性能、高可扩展性开源大语言模型（LLM）库，专为 Google Cloud TPU 和 GPU 训练设计。它支持 Gemma、Llama、DeepSeek、Qwen、Mistral 等主流模型的预训练及后训练（SFT、RL）。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu)
• Python 版本: Python 3.12 (官方主要支持版本，其他版本可能存在兼容性问题)
• 硬件加速:
  • Google Cloud TPUs (推荐)
  • NVIDIA GPUs (需安装对应的 CUDA/cuDNN 环境)

前置依赖

确保已安装以下基础工具：

• pip (包管理工具)
• git (代码克隆)
• JAX 相关依赖将通过安装命令自动处理

注意：MaxText 强依赖 XLA 编译器以发挥最佳性能，请确保您的硬件驱动支持最新的 XLA 版本。

安装步骤

推荐使用 PyPI 进行安装，这是最简便的方式。

方式一：通过 PyPI 安装（推荐）

pip install maxtext

如果您需要针对 TPU 后训练或运行 vllm_decode，请使用以下专用目标安装：

pip install maxtext[tpu-post-train]

方式二：从源码安装（开发模式）

如果您需要修改源码或体验最新功能：

git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
pip install -e .

国内加速提示：如果下载速度较慢，可临时使用国内镜像源（如清华源）：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple maxtext

基本使用

MaxText 的核心操作通过配置文件驱动。以下是最简单的运行示例流程。

1. 准备配置文件

MaxText 内置了多种模型配置。例如，要运行一个小型的 Gemma 模型，您可以直接使用仓库中的预设配置（假设已克隆源码或配置位于当前路径）：

# 查看可用的模型配置示例
ls src/maxtext/configs/models/

2. 运行训练任务

使用 train.py 脚本启动训练。以下是一个在单主机上运行 Gemma 2B 模型进行简单训练的示例命令：

python3 src/maxtext/train.py \
  src/maxtext/configs/base.yml \
  src/maxtext/configs/models/gemma2-2b.yml \
  run_name=my_first_run \
  dataset_path=your_dataset_path \
  steps=100

参数说明：

• base.yml: 基础配置文件。
• models/gemma2-2b.yml: 指定模型架构。
• run_name: 本次运行的名称（用于保存检查点）。
• dataset_path: 您的训练数据路径（支持 Grain 格式）。
• steps: 训练步数（示例设为 100 步用于测试）。

3. 运行后训练 (SFT/RL)

MaxText 同样支持监督微调 (SFT) 和强化学习 (RL)。使用 Tunix 框架进行后训练的典型命令如下：

# SFT 示例
python3 src/maxtext/post_training/sft.py \
  src/maxtext/configs/base.yml \
  src/maxtext/configs/models/llama3-8b.yml \
  checkpoint_path=path/to/pretrained_checkpoint \
  dataset_path=path/to/sft_dataset

下一步

• 详细文档: 访问 MaxText ReadTheDocs 获取完整的教程和 API 参考。
• 社区交流: 加入 Discord 频道 获取实时支持。
• 多模态支持: 如需使用 Gemma 4 或 Llama 4 等多模态模型，请参考官方最新的配置示例。