Repolode
AI-Hypercomputer

AI-Hypercomputer/maxtext

基于 JAX 的开源大模型库与参考实现,用于在 TPU 和 GPU 上预训练与后训练大语言模型。

Star
2,444
Fork
628
未关闭 Issue
477
最近提交
3小时前

这是什么?

MaxText 是基于 JAX 和 XLA 编译器、使用纯 Python 编写的大语言模型库,面向 Google Cloud 的 TPU 与 GPU 训练,可从单机扩展到超大规模集群。项目内置 Gemma、Llama、DeepSeek、Qwen、Mistral、Mixtral、GPT-OSS 等多个模型的实现。

在后训练方面,MaxText 支持监督微调(SFT)以及 GRPO、GSPO 等强化学习方法,训练框架使用 Tunix,RL 采样阶段使用 vLLM。检查点由 Orbax 管理,优化器基于 Optax,数据加载使用 Grain,神经网络构建基于 Flax。另外也支持 Gemma 3、Gemma 4、Llama 4 等多模态视觉语言模型的训练。

项目既可以直接运行现成配置进行训练,也可以作为起点,fork 后修改模型结构,用于研究或生产环境。扩散模型相关工作由同组织下的 MaxDiffusion 项目负责。

核心功能

  • 01覆盖 Gemma、Llama、DeepSeek、Qwen、Mistral 等多个开源模型系列的预训练与后训练。
  • 02提供监督微调(SFT)以及 GRPO、GSPO 等强化学习训练流程。
  • 03面向 Google Cloud TPU 和 GPU,支持多主机与多切片扩展。
  • 04基于 JAX、Flax、Optax、Orbax 和 Grain 构建。
  • 05可通过 PyPI 安装,也提供预构建的容器镜像。
  • 06支持 Gemma 3、Gemma 4、Llama 4 等多模态视觉语言模型训练。

适用场景

  • 在 TPU 或 GPU 集群上从头预训练自定义的稠密或 MoE 语言模型。
  • 在 Google Cloud 上对 Gemma、Llama 等开源模型进行 SFT 或强化学习微调。
  • 将 MaxText 作为参考实现,学习或 fork 大规模大模型训练代码。

快速上手

  1. 1

    按照安装指南从 PyPI 安装 MaxText

  2. 2

    若已有环境,在 MaxText 根目录下执行可编辑安装

    pip install -e .
  3. 3

    按照入门指南完成第一次 MaxText 运行

使用前须知

  • 面向 Google Cloud TPU 和 GPU 设计,需要加速器硬件才能发挥作用。
  • main 分支不保证可用于生产环境,建议使用最新的 PyPI 发行版。
  • Python 3.12 是主要支持版本,其他版本可能存在兼容性问题。

标签

常见问题

maxtext 是什么?
基于 JAX 的开源大模型库与参考实现,用于在 TPU 和 GPU 上预训练与后训练大语言模型。 MaxText 是基于 JAX 和 XLA 编译器、使用纯 Python 编写的大语言模型库,面向 Google Cloud 的 TPU 与 GPU 训练,可从单机扩展到超大规模集群。项目内置 Gemma、Llama、DeepSeek、Qwen、Mistral、Mixtral、GPT-OSS 等多个模型的实现。
maxtext 可以免费使用吗?用的什么许可证?
maxtext 以 Apache-2.0 许可证开源。用于商业产品前,请先确认许可证条款。
maxtext 还在维护吗?
活跃维护中。maxtext 最近一次提交是在 2026年10月21日。
maxtext 适合用来做什么?
在 TPU 或 GPU 集群上从头预训练自定义的稠密或 MoE 语言模型。 在 Google Cloud 上对 Gemma、Llama 等开源模型进行 SFT 或强化学习微调。 将 MaxText 作为参考实现,学习或 fork 大规模大模型训练代码。
maxtext 怎么安装和运行?
按照安装指南从 PyPI 安装 MaxText 若已有环境,在 MaxText 根目录下执行可编辑安装: pip install -e . 按照入门指南完成第一次 MaxText 运行
在哪里获取 maxtext?
源代码和安装说明在 GitHub 上:https://github.com/AI-Hypercomputer/maxtext

概要由 AI 根据项目 README 生成;GitHub 数据每日更新。

相关项目