Repolode
Submit
AI-Hypercomputer

AI-Hypercomputer/maxtext

A JAX-based open-source LLM library and reference implementation for pre-training and post-training large models on Google Cloud TPUs and GPUs.

View on GitHubWebsite
Stars
2,444
Forks
628
Open issues
477
Last commit
4 hours ago

What is it?

MaxText is a Python library built on JAX and the XLA compiler for training large language models on Google Cloud TPUs and GPUs. It ships a collection of model implementations, including Gemma, Llama, DeepSeek, Qwen, Mistral, Mixtral and GPT-OSS, and covers pre-training from a single host up to very large clusters.

For post-training, MaxText supports supervised fine-tuning (SFT) and reinforcement learning methods such as GRPO and GSPO, using Tunix for the training framework and vLLM for sampling during RL. Checkpointing relies on Orbax, optimization on Optax, data loading on Grain, and neural networks on Flax. Multi-modal training is available for Gemma 3, Gemma 4 and Llama 4 VLMs.

The code is intended both as a runnable training stack and as a starting point: users are encouraged to run the existing configs and then fork and modify the models for their own research or production work. Diffusion models are handled by the separate MaxDiffusion project.

Key features

  • 01Pre-training and post-training support for many open model families, including Gemma, Llama, DeepSeek, Qwen and Mistral.
  • 02Supervised fine-tuning and reinforcement learning (GRPO, GSPO) pipelines.
  • 03Targets Google Cloud TPUs and GPUs, with multi-host and multi-slice scaling.
  • 04Built on JAX, Flax, Optax, Orbax and Grain.
  • 05Available via PyPI and pre-built container images.
  • 06Multi-modal training for Gemma 3, Gemma 4 and Llama 4 VLMs.

Use cases

  • Pre-training a custom dense or MoE language model on TPU or GPU clusters.
  • Fine-tuning open models such as Gemma or Llama with SFT or RL on Google Cloud.
  • Using MaxText as a reference implementation to study or fork large-scale LLM training code.

Quick start

  1. 1

    Install MaxText from PyPI following the installation guide

  2. 2

    For an existing environment, run an editable install from the MaxText root

    pip install -e .
  3. 3

    Follow the Get Started guide for your first MaxText run

Things to know

  • Built for Google Cloud TPUs and GPUs, so it needs accelerator hardware to be useful.
  • The main branch is not expected to be production-ready; the latest PyPI release is recommended.
  • Python 3.12 is the primary supported version; other versions may have compatibility issues.

Topics

#jax#llm#tpu#pre-training#fine-tuning#reinforcement-learning#distributed-training#deepseek#gemma2#gemma3

FAQ

What is maxtext?
A JAX-based open-source LLM library and reference implementation for pre-training and post-training large models on Google Cloud TPUs and GPUs. MaxText is a Python library built on JAX and the XLA compiler for training large language models on Google Cloud TPUs and GPUs. It ships a collection of model implementations, including Gemma, Llama, DeepSeek, Qwen, Mistral, Mixtral and GPT-OSS, and covers pre-training from a single host up to very large clusters.
Is maxtext free to use? What license does it have?
maxtext is open source under the Apache-2.0 license. Check the license terms before using it in commercial products.
Is maxtext still maintained?
Actively maintained. The last commit to maxtext was on Oct 21, 2026.
What can I use maxtext for?
Pre-training a custom dense or MoE language model on TPU or GPU clusters. Fine-tuning open models such as Gemma or Llama with SFT or RL on Google Cloud. Using MaxText as a reference implementation to study or fork large-scale LLM training code.
How do I install maxtext?
Install MaxText from PyPI following the installation guide For an existing environment, run an editable install from the MaxText root: pip install -e . Follow the Get Started guide for your first MaxText run
Where can I get maxtext?
The source code and installation instructions are on GitHub: https://github.com/AI-Hypercomputer/maxtext

Summary generated by AI from the project's README. GitHub stats refresh daily.

Related projects