Reference documentation

Reference documentation#

Deep dive into MaxText architecture, models, and core concepts.

📊 Performance Metrics

Understanding Model Flops Utilization (MFU), calculation methods, and why it matters for performance optimization.

Performance metrics
🤖 Models

Supported models and architectures, including Llama, Qwen, and Mixtral. Details on tiering and new additions.

Models
🏗️ Architecture

High-level overview of MaxText design, JAX/XLA choices, and how components interact.

Architecture
💡 Core Concepts

Key concepts including checkpointing strategies, quantization, tiling, and Mixture of Experts (MoE) configuration.

Core concepts

📚 API Reference#

Find comprehensive API documentation for MaxText modules, classes, and functions in the API Reference page.