AI-Hypercomputer/maxtext

A simple, performant, and scalable Jax LLM!

What it solves

MaxText is designed to provide a high-performance, scalable reference implementation for training Large Language Models (LLMs). It addresses the challenge of achieving high Model FLOPs Utilization (MFU) and throughput (tokens/second) across a wide range of hardware, from single hosts to massive clusters of Google Cloud TPUs and GPUs, while keeping the codebase simple and "optimization-free."

How it works

Written in pure Python and JAX, MaxText leverages the XLA compiler to optimize performance. It integrates with a suite of JAX AI libraries including Flax (neural networks), Tunix (post-training), Orbax (checkpointing), Optax (optimization), and Grain (dataloading). The library supports both pre-training from scratch and scalable post-training using techniques like Supervised Fine-Tuning (SFT) and Reinforcement Learning (GRPO and GSPO).

Who it’s for

It is intended for researchers and developers building ambitious LLM projects in both research and production environments who need a scalable, high-performance framework for training and fine-tuning popular open-source models.

Highlights

  • Broad Model Support: Includes reference implementations for Gemma, Llama, DeepSeek, Qwen, and Mistral, including MoE (Mixture of Experts) and multimodal models.
  • Scale: Supports pre-training on up to tens of thousands of chips.
  • Post-Training Framework: Provides a scalable framework via Tunix for SFT and RL (using vLLM for sampling).
  • Hardware Optimization: Specifically targeted at Google Cloud TPUs and GPUs for maximum efficiency.
  • Multimodal Capabilities: Supports multimodal training for Gemma 3, Gemma 4, and Llama 4 VLMs.

Related

  • Project
    google/tunix

    A JAX-based post-training library for LLMs that provides scalable support for supervised fine-tuning and reinforcement learning on TPUs.

  • Project
    NVIDIA/JAX-Toolbox

    NVIDIA JAX‑Toolbox is a collection of up‑to‑date Docker containers, docs, and news that make it easy to run the JAX ML framework (and related projects like MaxText, AxLearn, Equinox, AlphaFold 3) on NVIDIA GPUs, with GPU‑specific XLA optimizations for large‑scale training and inference.

  • Project
    sgl-project/sglang-jax

    A high-performance JAX-based inference engine optimized for Google TPUs, providing high throughput and low latency for LLMs and multimodal models.

  • Project
    erfanzar/EasyDeL

    A JAX and Flax NNX-based framework for the high-performance training, fine-tuning, and serving of LLMs and multimodal models on TPU/GPU clusters.

  • Project
    google-deepmind/gemma

    A JAX library for deploying and fine-tuning the Gemma family of open-weights large language models, supporting multi-modal conversations and multiple hardware backends.