From bef18ba2b4e633ddf1ed08947d390455857732d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Melissa=20Weber=20Mendon=C3=A7a?= Date: Thu, 25 Sep 2025 13:06:00 -0300 Subject: [PATCH] Reorganize documentation --- docs/conf.py | 7 +- docs/explanations.md | 6 +- .../alternatives.md | 2 +- docs/{guides => explanations}/checkpoints.md | 14 +- .../jax_ai_libraries_chosen.md | 23 +-- .../llm_calculator.ipynb | 0 .../performance_metrics.md | 15 +- docs/explanations/quantization.md | 2 - docs/explanations/steps_model.md | 33 ---- docs/guides.md | 12 +- docs/guides/checkpointing_solutions.md | 9 + .../emergency_checkpointing.md | 16 +- .../gcs_checkpointing.md | 8 +- .../multi_tier_checkpointing.md | 22 +-- docs/guides/custom_model.md | 12 +- docs/guides/data_input_pipeline.md | 14 +- .../data_input_grain.md | 0 .../data_input_hf.md | 2 +- .../data_input_tfds.md | 0 .../data_pipeline_perf.md | 6 +- docs/guides/features_and_diagnostics.md | 19 +- docs/guides/gcp_workload_observability.md | 8 +- docs/guides/knowledge_distillation.md | 24 +-- docs/guides/monitor_goodput.md | 37 ++-- docs/guides/pallas_kernels_performance.md | 187 ++++++++---------- docs/guides/run_maxtext.md | 10 + .../run_maxtext_localhost.md | 34 ++-- .../run_maxtext_single_host_gpu.md} | 18 +- .../run_maxtext_via_multihost_job.md | 0 .../run_maxtext_via_multihost_runner.md | 0 .../run_maxtext_via_pathways.md | 10 +- .../{ => run_maxtext}/run_maxtext_via_xpk.md | 76 +++---- docs/guides/understand_logs_and_metrics.md | 85 ++++---- docs/guides/use_vertex_ai_tensorboard.md | 2 +- docs/reference.md | 2 +- docs/reference/architecture_overview.md | 16 +- docs/reference/benchmark_and_performance.md | 73 +++---- docs/reference/config_base.md | 6 + .../supported_models_and_architectures.md | 18 +- docs/tutorials.md | 3 +- docs/tutorials/first_run.md | 12 +- docs/tutorials/full_finetuning.md | 14 +- docs/tutorials/grpo.md | 10 +- docs/tutorials/pretraining.md | 8 +- docs/tutorials/sft.md | 4 +- docs/tutorials/sft_on_multi_host.md | 18 +- .../tpu/llama2}/run_llama2.md | 1 - requirements_docs.txt | 1 + 48 files changed, 401 insertions(+), 498 deletions(-) rename docs/{reference => explanations}/alternatives.md (98%) rename docs/{guides => explanations}/checkpoints.md (94%) rename docs/{guides => explanations}/jax_ai_libraries_chosen.md (96%) rename docs/{guides => explanations}/llm_calculator.ipynb (100%) rename docs/{guides => explanations}/performance_metrics.md (86%) delete mode 100644 docs/explanations/steps_model.md create mode 100644 docs/guides/checkpointing_solutions.md rename docs/guides/{ => data_input_pipeline}/data_input_grain.md (100%) rename docs/guides/{ => data_input_pipeline}/data_input_hf.md (98%) rename docs/guides/{ => data_input_pipeline}/data_input_tfds.md (100%) rename docs/{explanations => guides/data_input_pipeline}/data_pipeline_perf.md (90%) create mode 100644 docs/guides/run_maxtext.md rename docs/guides/{ => run_maxtext}/run_maxtext_localhost.md (79%) rename docs/guides/{single_host_gpu.md => run_maxtext/run_maxtext_single_host_gpu.md} (91%) rename docs/guides/{ => run_maxtext}/run_maxtext_via_multihost_job.md (100%) rename docs/guides/{ => run_maxtext}/run_maxtext_via_multihost_runner.md (100%) rename docs/guides/{ => run_maxtext}/run_maxtext_via_pathways.md (97%) rename docs/guides/{ => run_maxtext}/run_maxtext_via_xpk.md (86%) create mode 100644 docs/reference/config_base.md rename {docs/tutorials => end_to_end/tpu/llama2}/run_llama2.md (99%) diff --git a/docs/conf.py b/docs/conf.py index 7b99d4a834..a8f646d46b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -36,6 +36,7 @@ extensions = [ "myst_nb", "sphinx_design", + "sphinx_copybutton", ] templates_path = ["_templates"] @@ -59,7 +60,7 @@ # Remove specific documents from ToC exclude_patterns = [ - "guides/run_maxtext_via_multihost_job.md", - "guides/run_maxtext_via_multihost_runner.md", - "guides/llm_calculator.ipynb", + "guides/run_maxtext/run_maxtext_via_multihost_job.md", + "guides/run_maxtext/run_maxtext_via_multihost_runner.md", + "explanations/llm_calculator.ipynb", ] diff --git a/docs/explanations.md b/docs/explanations.md index d7b0a39f24..41404f9d5b 100644 --- a/docs/explanations.md +++ b/docs/explanations.md @@ -19,9 +19,11 @@ ```{toctree} :maxdepth: 1 -explanations/steps_model.md +explanations/jax_ai_libraries_chosen.md +explanations/alternatives.md +explanations/checkpoints.md explanations/quantization.md explanations/sharding.md -explanations/data_pipeline_perf.md explanations/tiling.md +explanations/performance_metrics.md ``` diff --git a/docs/reference/alternatives.md b/docs/explanations/alternatives.md similarity index 98% rename from docs/reference/alternatives.md rename to docs/explanations/alternatives.md index cb44295f01..f515f52f19 100644 --- a/docs/reference/alternatives.md +++ b/docs/explanations/alternatives.md @@ -14,7 +14,7 @@ limitations under the License. --> -# Comparison to Alternatives +# Comparison to alternatives MaxText is heavily inspired by [MinGPT](https://github.com/karpathy/minGPT)/[NanoGPT](https://github.com/karpathy/nanoGPT), elegant standalone GPT implementations written in PyTorch and targeting Nvidia GPUs. MaxText is more complex, supporting more industry standard models and scaling to tens of thousands of chips. Ultimately MaxText has an MFU more than three times the [17%](https://twitter.com/karpathy/status/1613250489097027584?cxt=HHwWgIDUhbixteMsAAAA) reported most recently with that codebase, is massively scalable and implements a key-value cache for efficient auto-regressive decoding. diff --git a/docs/guides/checkpoints.md b/docs/explanations/checkpoints.md similarity index 94% rename from docs/guides/checkpoints.md rename to docs/explanations/checkpoints.md index 294dab42d4..7817d5f379 100644 --- a/docs/guides/checkpoints.md +++ b/docs/explanations/checkpoints.md @@ -16,7 +16,7 @@ # Checkpoints -## Checkpoint Formats +## Checkpoint formats Checkpoint formats in MaxText can be categorized along two axes: whether they include **training states** (e.g., optimizer properties) and whether the model's parameter weights are **stacked** or **unstacked** (aka scanned/unscanned). This results in the four types summarized below: @@ -27,13 +27,13 @@ Checkpoint formats in MaxText can be categorized along two axes: whether they in We discuss these two axes respectively: -### Training States +### Training states Checkpoints with a **training state** contain more than just the model's parameter weights. They also include the **optimizer state** (e.g., momentum values), which is essential for resuming a training run exactly where it left off. These "training checkpoints" are typically saved as snapshots during training to allow for recovery if the process is interrupted. In contrast, **inference checkpoints** contain only the parameter weights. We also call them parameter only/param-only checkpoints. This is the format most commonly used for sharing models on public platforms like HuggingFace, as they are smaller and ready for immediate use in inference or for fine-tuning. -### Stacked Checkpoints and JAX Scan Function +### Stacked checkpoints and JAX scan function The concept of stacked vs. unstacked checkpoints is specific to JAX-based models that use the `jax.lax.scan` function ([doc](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)). `scan` is a powerful JAX feature that compiles sequential operations (like the layers of a Transformer) into a single, highly optimized kernel, avoiding the overhead of a Python for-loop. @@ -78,11 +78,11 @@ In MaxText, we treat **Stacked Inference Checkpoints** as the default format for --- -## Using Checkpoints in Practice +## Using checkpoints in practice Beyond understanding the formats, it's crucial to know how to use checkpoints in your training workflows. MaxText uses flags in the configuration file or on the command line to manage checkpoints. -### Saving Checkpoints During Training +### Saving checkpoints during training MaxText automatically saves checkpoints periodically during a training run. These are **Stacked Training Checkpoints** that contain the full state needed to resume. @@ -97,4 +97,6 @@ Furthermore, MaxText supports emergency checkpointing, which saves a local copy - `local_checkpoint_directory`: The local path for storing emergency checkpoints. - `local_checkpoint_period`: The interval, in training steps, for saving local checkpoints. -More configs about checkpoints can be found in [here](https://github.com/AI-Hypercomputer/maxtext/blob/518a87037abb2497a2514ff0c8ffc263c69c6f9f/MaxText/configs/base.yml#L23-L65). +More configs about checkpoints can be found in [here](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/configs/base.yml#L23-L65). + +For practical guides on checkpointing, please refer to [](../guides/checkpointing_solutions.md). diff --git a/docs/guides/jax_ai_libraries_chosen.md b/docs/explanations/jax_ai_libraries_chosen.md similarity index 96% rename from docs/guides/jax_ai_libraries_chosen.md rename to docs/explanations/jax_ai_libraries_chosen.md index 8188e41c1d..4ea0cb77f2 100644 --- a/docs/guides/jax_ai_libraries_chosen.md +++ b/docs/explanations/jax_ai_libraries_chosen.md @@ -1,4 +1,4 @@ -# The JAX Ecosystem in MaxText: An Opinionated Guide +# The JAX ecosystem in MaxText: an opinionated guide MaxText is built on a curated stack of JAX libraries, each chosen for a specific purpose. This document provides an opinionated view on *why* MaxText uses the following key components of the JAX ecosystem: @@ -11,11 +11,9 @@ MaxText is built on a curated stack of JAX libraries, each chosen for a specific This stack isn't just a random collection of tools; it represents a design philosophy centered around **explicitness, composability, and performance at scale**. - This document provides an opinionated view on *why* MaxText uses these specific libraries, explaining the design decisions that make them ideal for building and training large-scale models. - -## Flax: For Functional Model Definition +## Flax: For functional model definition **What is it?** Flax is a high-performance neural network library for JAX that is designed to be flexible, explicit, and easy to use. @@ -27,8 +25,7 @@ With its latest generation API, NNX, Flax provides a modern, object-oriented (OO For more information on using Flax, please refer to https://github.com/google/flax - -## Optax: For Composable Optimization +## Optax: For composable optimization **What is it?** Optax is a gradient processing and optimization library for JAX. It reimagines the optimizer as a series of composable functional transformations. @@ -38,8 +35,7 @@ For more information on using Flax, please refer to https://github.com/google/fl For more information on using Optax, please refer to https://github.com/google-deepmind/optax - -## Orbax: For Robust Checkpointing +## Orbax: For robust checkpointing **What is it?** Orbax is a library for checkpointing JAX programs, designed for large-scale, potentially unreliable environments. @@ -54,8 +50,7 @@ For massive models, saving and loading state is a critical part of the training For more information on using Orbax, please refer to https://github.com/google/orbax - -## Grain: For Deterministic, Multi-Host Data Loading +## Grain: For deterministic, multi-host data loading **What is it?** Grain is a high-performance data loading library designed for deterministic, global shuffle and multi-host data loading. @@ -67,8 +62,7 @@ Its APIs are explicitly designed for the multi-host paradigm, simplifying the pr For more information on using Grain, please refer to https://github.com/google/grain and the grain guide in maxtext located at https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_grain.md - -## Qwix: For Native JAX Quantization +## Qwix: For native JAX quantization **What is it?** Qwix is a Jax quantization library supporting Quantization-Aware Training (QAT) and Post-Training Quantization (PTQ) @@ -79,8 +73,7 @@ We chose Qwix because it provides the necessary primitives **natively within the For more information on how to quantize your model using Qwix, please refer to https://github.com/google/qwix - -## Tunix: For Comprehensive Post-Training +## Tunix: For comprehensive post-training **What is it?** Tunix is a JAX-based library designed for a wide range of post-training tasks, including Supervised Fine-Tuning (SFT), Reinforcement Learning (RL), and Parameter-Efficient Fine-Tuning (PEFT) methods like LoRA. @@ -95,4 +88,4 @@ MaxText leverages Tunix as its core library for post-training, offering a unifie We chose Tunix because it provides a **comprehensive, performant, and JAX-native solution for the entire post-training lifecycle**. Its integration with libraries like vLLM and its alignment with the NNX ecosystem make it a powerful tool for both full model adaptation and parameter-efficient tuning. -For more information on using Tunix, please refer to https://github.com/google/tunix \ No newline at end of file +For more information on using Tunix, please refer to https://github.com/google/tunix diff --git a/docs/guides/llm_calculator.ipynb b/docs/explanations/llm_calculator.ipynb similarity index 100% rename from docs/guides/llm_calculator.ipynb rename to docs/explanations/llm_calculator.ipynb diff --git a/docs/guides/performance_metrics.md b/docs/explanations/performance_metrics.md similarity index 86% rename from docs/guides/performance_metrics.md rename to docs/explanations/performance_metrics.md index 23e530e554..b30f02979b 100644 --- a/docs/guides/performance_metrics.md +++ b/docs/explanations/performance_metrics.md @@ -14,7 +14,8 @@ limitations under the License. --> -# Performance Metrics +(performance-metrics)= +# Performance metrics ## MFU @@ -57,11 +58,11 @@ $$ Hence, MFU is the fraction of peak hardware performance actually utilized by the model, and can be expressed in different units — step time, throughput, or raw flops/s. -### MaxText Calculating + Reporting -In MaxText, we sum all of the matmuls performed in one step, see [calculate_tflops_per_device](https://github.com/AI-Hypercomputer/maxtext/blob/e969faabbb571285a51545530f34d8f0a9f237e9/MaxText/maxtext_utils.py#L297) -and divide it by the measured (via python `time.time()`) step time. In each step we print the resulting Model Flops per second [`per_device_tflops_per_sec`](https://github.com/AI-Hypercomputer/maxtext/blob/e969faabbb571285a51545530f34d8f0a9f237e9/MaxText/metric_logger.py#L193-L194). One can calculate the MFU by dividing this number by the peak tflops of the hardware (e.g., $918e^{12}$ FLOPS/s for Trillium). +### MaxText calculating + reporting +In MaxText, we sum all of the matmuls performed in one step, see [calculate_tflops_per_device](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/maxtext_utils.py#L454) +and divide it by the measured (via python `time.time()`) step time. In each step we print the resulting Model Flops per second [`per_device_tflops_per_sec`](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/metric_logger.py#L211-L213). One can calculate the MFU by dividing this number by the peak tflops of the hardware (e.g., $918e^{12}$ FLOPS/s for Trillium). -### Causal Attention +### Causal attention Due to causality only half of the (query, key) pairs need to be computed, those with query_idx >= key_idx. This accounts for the fact only prior tokens can be used to predict future ones. Prior to https://github.com/AI-Hypercomputer/maxtext/pull/1988 MaxText did not account for sparsity for theoretical flops, and used Attention Flops ~= 4 * sequence^2 * batch * heads * head_dim @@ -98,6 +99,6 @@ $$\begin{align*} This shows any of step time, tokens/s or MFU can be used to determine how long training will take and are proportionally (or inversely proportionally) related. MFU is most useful to compare across different models/hardwares and while optimizing performance, whereas step time or tokens/second may be more useful when these are fixed. -## Why not Hardware Flops? +## Why not hardware flops? -Hardware (e.g., XLA reported) FLOPs do not accurately reflect computation efficiency as they depend on the program / implementation, not just on the model and its inherent computations (higher hardware FLOPs does not necessarily mean less room for improvement). For example, they include remat and potentially auxiliary operations (such as reshaping for dropping moe [here](https://github.com/AI-Hypercomputer/maxtext/blob/4b6142950aff5d9ba42d830efc5ce4c4ac9d4135/MaxText/layers/moe.py#L1267)), which are an implementation detail and not part of the model. In addition, XLA reported FLOPs may not be accurate with pallas kernels. Hardware flops utilization is not (inversely) proportional to step time as opposed to MFU, since hardware flops can change with implementation details like remat policies. +Hardware (e.g., XLA reported) FLOPs do not accurately reflect computation efficiency as they depend on the program / implementation, not just on the model and its inherent computations (higher hardware FLOPs does not necessarily mean less room for improvement). For example, they include remat and potentially auxilliary operations (such as reshaping for dropping moe [here](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/layers/moe.py#L1544)), which are an implementation detail and not part of the model. In addition, XLA reported FLOPs may not be accurate with pallas kernels. Hardware flops utilization is not (inversely) proportional to step time as opposed to MFU, since hardware flops can change with implementation details like remat policies. diff --git a/docs/explanations/quantization.md b/docs/explanations/quantization.md index 69930cb7ce..a6e8ede859 100644 --- a/docs/explanations/quantization.md +++ b/docs/explanations/quantization.md @@ -35,8 +35,6 @@ You can turn on the quantization by adding the following flag `--quantization` a - 'intmp' for mixed precision weight only quantization based on config file - 'fp8' for 8-bit floating-point GeMMs on NVIDIA GPUs. - - ## How QAT works with Qwix The core idea behind QAT is to insert "fake quantization" operations into the model's computation graph. During the training forward pass, these operations simulate the effect of quantizing weights and activations to a lower precision. For the backward pass, Qwix uses the Straight-Through Estimator (STE) to approximate the gradients, allowing the model to learn effectively despite the non-differentiable nature of quantization. diff --git a/docs/explanations/steps_model.md b/docs/explanations/steps_model.md deleted file mode 100644 index fc418a9717..0000000000 --- a/docs/explanations/steps_model.md +++ /dev/null @@ -1,33 +0,0 @@ - - -# Steps to build a Model - -```{figure} ../_static/build_model.png -Stages of LLM Model Development from pre-training to fine tuning and finally serving a model. -``` - -Model building starts with Pre-training a base model architecture. Pre-training is the process where you take a model architecture, which starts with random weights and train with a very large corpus in the scale of trillions of tokens. E.g. Google’s Gemma models was pre-trained on 6 Trillion tokens; LLama 3 was trained with 15 Trillion tokens - -Post the pre-training most model producers will publish a checkpoint of the weights of the model. The corpus used for pre-training these models are usually a large public corpus like Common Crawl, public code bases, books etc. - -Though these may be a great way to answer very general questions or prompts, they usually fail on very domain specific questions and answers like Medical and Life Sciences, Engineering, etc. - -Customers and enterprises usually like to continue training a pre-trained model or performing a full fine tuning of the models using their own datasets. These datasets are usually in billions of tokens. This allows better prompt understanding when questions are asked on keywords and terms specific to their model or domain specific question. - -Post a Full Fine Tuning, most models go through a process of Instruction Fine Tuning(PEFT/LoRA), Supervised Fine Tuning and RLHF to improve the model quality and follow prompt answers better. - -PEFT/Lora, Supervised Finetuning are less expensive operations compared to full fine tuning. diff --git a/docs/guides.md b/docs/guides.md index c7d070ae6e..4411edc4df 100644 --- a/docs/guides.md +++ b/docs/guides.md @@ -19,24 +19,16 @@ ```{toctree} :maxdepth: 1 -guides/checkpoints.md +guides/run_maxtext.md guides/custom_model.md -guides/run_maxtext_localhost.md -guides/run_maxtext_via_xpk.md -guides/run_maxtext_via_pathways.md guides/data_input_pipeline.md -guides/single_host_gpu.md guides/knowledge_distillation.md guides/gcp_workload_observability.md guides/monitor_goodput.md guides/use_vertex_ai_tensorboard.md guides/features_and_diagnostics.md guides/pallas_kernels_performance.md -guides/performance_metrics.md guides/understand_logs_and_metrics.md -guides/checkpointing_solutions/gcs_checkpointing.md -guides/checkpointing_solutions/emergency_checkpointing.md -guides/checkpointing_solutions/multi_tier_checkpointing.md -guides/jax_ai_libraries_chosen.md guides/xprof_user_guide.md +guides/checkpointing_solutions.md ``` diff --git a/docs/guides/checkpointing_solutions.md b/docs/guides/checkpointing_solutions.md new file mode 100644 index 0000000000..b73d26b731 --- /dev/null +++ b/docs/guides/checkpointing_solutions.md @@ -0,0 +1,9 @@ +# Checkpointing solutions + +```{toctree} +:maxdepth: 1 + +checkpointing_solutions/gcs_checkpointing.md +checkpointing_solutions/emergency_checkpointing.md +checkpointing_solutions/multi_tier_checkpointing.md +``` diff --git a/docs/guides/checkpointing_solutions/emergency_checkpointing.md b/docs/guides/checkpointing_solutions/emergency_checkpointing.md index 552f90f83a..f162765cbe 100644 --- a/docs/guides/checkpointing_solutions/emergency_checkpointing.md +++ b/docs/guides/checkpointing_solutions/emergency_checkpointing.md @@ -1,4 +1,4 @@ -# Emergency Checkpointing +# Emergency checkpointing Emergency checkpointing is a vital feature for large-scale, multi-slice training. It enables rapid saving and restoration of model state from local, in-memory checkpoints in response to hardware failures, host errors, or preemptions. This feature becomes increasingly critical as the number of hosts and devices grows, which raises the probability of a failure. @@ -11,11 +11,10 @@ Emergency checkpointing is a vital feature for large-scale, multi-slice training * **Ramdisk Mounted via Jobset**: Each workload pod must have a [ramdisk directory mounted by Jobset](https://cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing#update-jobset) using the Multi-Tier Checkpointing CSI driver. This provides a high-speed, in-memory storage location for checkpoints. * **Supported TPU types**: [v4](https://cloud.google.com/tpu/docs/v4), [v5e](https://cloud.google.com/tpu/docs/v5e), [v5p](https://cloud.google.com/tpu/docs/v5p), and [v6e](https://cloud.google.com/tpu/docs/v6e) -## Cluster Creation using XPK +## Cluster creation using XPK To run workloads with Emergency Checkpointing, you need a Google Kubernetes Engine (GKE) cluster with the necessary drivers and features enabled. You can create a properly configured cluster using the **XPK** or by setting it up manually with `gcloud` commands following [Google Cloud Checkpointing Documentation](https://cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing). - The `xpk` script provides a streamlined way to create a GKE cluster with all the required MTC settings. The key flags used are: * `--enable-mtc`: Enables the Multi-Tier Checkpointing feature. @@ -23,8 +22,7 @@ The `xpk` script provides a streamlined way to create a GKE cluster with all the * `--mtc-ramdisk-size`: Allocates an in-memory ramdisk on each node for fast, local checkpoints. * `--mtc-gcs-bucket`: Specifies the GCS bucket. It is not utilized in emergency checkpointing, but is needed to deploy checkpointing configurations. - -### Calculating Ramdisk Size Per Host +### Calculating ramdisk size per host The total size of a full training checkpoint (including model weights and optimizer state) can be estimated based on the number of model parameters. A good rule of thumb: @@ -39,7 +37,7 @@ The formula is: It's a good practice to add a **10-15% buffer** . -### Example Calculation +### Example calculation Let's walk through an example for a large model. @@ -97,7 +95,7 @@ In this scenario, you should configure each pod in that slice with a ramdisk of --gke-version=${GKE_VERSION} ``` -## MaxText Configuration +## MaxText configuration MaxText provides a set of configuration flags to control checkpointing options. This configuration manages a `two-tiered checkpointing` system designed for both durability and rapid recovery. @@ -105,7 +103,6 @@ MaxText provides a set of configuration flags to control checkpointing options. * **Persistent Checkpoints**: These are standard checkpoints saved periodically and much more rarely to durable storage(i.e. GCS bucket). They ensure that you can recover your training state even after a complete cluster failure. This is controlled by `enable_checkpointing`, and `checkpoint_period`. - | Flag | Description | Type | Default | | :--- | :--- | :--- | :--- | | `enable_checkpointing` | A master switch to enable (`True`) or disable (`False`) saving checkpoints during the training run. | `boolean` | `False` | @@ -116,7 +113,7 @@ MaxText provides a set of configuration flags to control checkpointing options. | `checkpoint_period` | The interval, in training steps, for how often a checkpoint is saved to **persistent storage**. | `integer` | `10000` | | `enable_single_replica_ckpt_restoring` | If `True`, one replica reads the checkpoint from storage and then broadcasts it to all other replicas. This can significantly speed up restoration on multi-host systems by reducing redundant reads from storage. | `boolean` | `False` | -## Workload Creation using XPK +## Workload creation using XPK The flags below would give the user access to the ramdisk in their workload: @@ -125,7 +122,6 @@ The flags below would give the user access to the ramdisk in their workload: | `--mtc-enabled` | Enables the Multi-Tier Checkpointing feature, by mounting ramdisk to the workload pods, using csi drivers. | | `--ramdisk-directory` | Specifies the mount path inside each pod where the high-speed ramdisk will be accessible. Your training application should write its local, emergency checkpoints to this path. | - ### Example XPK workload creation command 1. **Set up environment variables:** diff --git a/docs/guides/checkpointing_solutions/gcs_checkpointing.md b/docs/guides/checkpointing_solutions/gcs_checkpointing.md index 14f92397e6..1ba471492b 100644 --- a/docs/guides/checkpointing_solutions/gcs_checkpointing.md +++ b/docs/guides/checkpointing_solutions/gcs_checkpointing.md @@ -7,7 +7,7 @@ the training in sync with rest of the workers, then the model's state is asynchronously serialized and written over the network to the specified GCS bucket. -## Checkpoint Loading Priority +## Checkpoint loading priority The system follows a specific order when deciding which checkpoint to load at startup. The first valid condition met is the one executed: @@ -21,7 +21,7 @@ The system follows a specific order when deciding which checkpoint to load at st * **Note**: These two options are mutually exclusive and will cause an error if both are set. 3. **Initialize from Scratch**: We don't load a checkpoint and initialize state instead. -### MaxText Configuration +### MaxText configuration | Flag | Description | Type | Default | | :--- | :--- | :--- | :--- | @@ -34,7 +34,7 @@ The system follows a specific order when deciding which checkpoint to load at st | `lora_input_adapters_path` | Specifies a parent directory containing LoRA (Low-Rank Adaptation) adapters. | `string` | `""` (disabled) | | `force_unroll` | If `True`, unrolls the loop when generating a parameter-only checkpoint. | `boolean` | `False` | -## Storage and Format Configuration +## Storage and format configuration These settings control the underlying storage mechanism ([Orbax](https://orbax.readthedocs.io)) for performance and compatibility. @@ -46,4 +46,4 @@ These settings control the underlying storage mechanism ([Orbax](https://orbax.r | `checkpoint_storage_concurrent_gb` | Controls the concurrent I/O limit in gigabytes for the checkpointer. Larger models may require increasing this value to avoid I/O bottlenecks. | `integer` | `96` | | `enable_orbax_v1` | A boolean flag to explicitly enable features and behaviors from Orbax version 1. | `boolean` | `False` | | `source_checkpoint_layout` | Specifies the format of the checkpoint being **loaded**. This tells the system how to interpret the files at the source path.
**Options**: `"orbax"`, `"safetensors"` | `string` | `"orbax"` | -| `checkpoint_conversion_fn` | A user-defined function to process a loaded checkpoint dictionary into a format that the model can understand. This is essential for loading checkpoints from different frameworks or formats (e.g., converting keys from a Hugging Face SafeTensors file). | `function` or `None` | `None` | \ No newline at end of file +| `checkpoint_conversion_fn` | A user-defined function to process a loaded checkpoint dictionary into a format that the model can understand. This is essential for loading checkpoints from different frameworks or formats (e.g., converting keys from a Hugging Face SafeTensors file). | `function` or `None` | `None` | diff --git a/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md b/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md index 4d05dcb718..421e53fe09 100644 --- a/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md +++ b/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md @@ -1,15 +1,15 @@ -# Multi-Tier Checkpointing +# Multi-tier checkpointing Multi-tier checkpointing is a solution designed to optimize the storage and management of checkpoints for large-scale machine learning (ML) training jobs, particularly those utilizing thousands of nodes. It aims to increase **"Goodput"** (the time spent making progress) and decrease costs by reducing wasted progress and the **mean-time-to-recovery (MTTR)** from failures. -## Purpose and Benefits +## Purpose and benefits * **Addresses frequent interruptions**: Large-scale ML training jobs are prone to frequent interruptions (potentially hourly), and recovery from these can be slow. * **Improves Goodput**: By saving checkpoints more frequently and efficiently, multi-tier checkpointing reduces the amount of lost progress when a failure occurs, thereby increasing the overall Goodput of the training process. * **Reduces MTTR**: The multi-tiered approach allows for faster restoration of training progress after a disruption. * **Optimized restore**: During the ML training workload's startup, available checkpoint shards are asynchronously copied to the local ramdisk. These shards are pulled from the fastest available source, whether from local peer nodes or the backup on GCS persistent storage. This process ensures the data is ready to be picked up by Orbax from the ramdisk, minimizing startup delays. -## Architecture and Tiers +## Architecture and tiers Multi-tier checkpointing stores checkpoints across multiple tiers of storage: @@ -17,13 +17,13 @@ Multi-tier checkpointing stores checkpoints across multiple tiers of storage: * **In-cluster (peer replication)**: Checkpoints are replicated to other nodes or slices within the cluster. * **GCS (persistent storage)**: Checkpoints are backed up to GCS for long-term durability and global accessibility. This tier is used for less frequent, but more robust, saves. -## Implementation Details +## Implementation details * **GKE Component**: A managed GKE component is involved in handling high-scale checkpointing, including controllers, daemonsets, worker discovery, and rank assignment. * **Local Storage**: For multi-tier checkpointing, local storage (such as ramdisk provided by a CSI ephemeral driver) is used for checkpoints, persisting across workload pod deletions. * **Replication Service**: A replication service in a managed GKE component replicates checkpoints in-cluster and backs up local checkpoint files to GCS at certain intervals and is responsible for fetching latest checkpoint files to nodes without local checkpoints during restoration. -## Comparison with Other Checkpointing Methods +## Comparison with other checkpointing methods * **GCS Checkpointing**: This involves saving the model state directly to durable storage like GCS. However, this can be slow at larger model/cluster scales, blocking training, and leading to redundant data copies. * **Emergency/Ramdisk Checkpointing**: While this method uses a low-latency ramdisk for checkpointing, Orbax manages the GCS save and restore operations at the workload level. As a result, saving to GCS blocks the training process during the device-to-host data transfer. @@ -39,7 +39,7 @@ Multi-tier checkpointing stores checkpoints across multiple tiers of storage: * **Supported TPU types**: [v4](https://cloud.google.com/tpu/docs/v4), [v5e](https://cloud.google.com/tpu/docs/v5e), [v5p](https://cloud.google.com/tpu/docs/v5p), and [v6e](https://cloud.google.com/tpu/docs/v6e) * **Cluster version**: Gke cluster version needs to be later than [1.32.3-gke.1170000](https://cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing#existing-cluster). -## Cluster Creation using XPK +## Cluster creation using XPK To run workloads with Multi-Tier Checkpointing (MTC), you need a Google Kubernetes Engine (GKE) cluster with the necessary drivers and features enabled. You can create a properly configured cluster using the **XPK** or by setting it up manually with `gcloud` commands following [Google Cloud Checkpointing Documentation](https://cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing). @@ -53,7 +53,7 @@ The [xpk script](https://github.com/AI-Hypercomputer/xpk/blob/develop/xpk.py) pr | `--mtc-gcs-bucket` | Specifies the GCS bucket. | -### Calculating Ramdisk Size Per Host +### Calculating ramdisk size per host The total size of a full training checkpoint (including model weights and optimizer state) can be estimated based on the number of model parameters. A good rule of thumb: @@ -68,7 +68,7 @@ The formula is: It's a good practice to add a **10-15% buffer** . -### Example Calculation +### Example calculation Let's walk through an example for a large model. @@ -126,7 +126,7 @@ In this scenario, you should configure each pod in that slice with a ramdisk of --gke-version=${GKE_VERSION} ``` -## MaxText Configuration +## MaxText configuration This configuration manages a `multi-tiered checkpointing` system designed for both durability and rapid recovery. @@ -142,7 +142,7 @@ This configuration manages a `multi-tiered checkpointing` system designed for bo | `local_checkpoint_period` | The interval, in training steps, for how often a **Multi-tier checkpoint** is saved in local ramdisks. | `integer` | `0` | | `multi_tier_checkpointing_backup_interval_minutes`| The interval, in minutes, for how often a **Multi-tier checkpoint** is saved to backup from local ramdisks. | `integer` | `0` | -### Workload Creation using XPK +### Workload creation using XPK The flags below would give the user access to the ramdisk in their workload: @@ -184,4 +184,4 @@ The flags below would give the user access to the ramdisk in their workload: --ramdisk-directory=${RAMDISK_DIRECTORY} \ --mtc-enabled \ --command "python3 src/MaxText/train.py src/MaxText/configs/base.yml base_output_directory=$OUTPUT_PATH dataset_path=$DATA_PATH steps=120 per_device_batch_size=6 enable_checkpoint_cloud_logger=True checkpoint_period=${CHECKPOINT_PEROID} enable_multi_tier_checkpointing=True local_checkpoint_period=${LOCAL_CHECKPOINT_PERIOD} local_checkpoint_directory=/${RAMDISK_DIRECTORY} multi_tier_checkpointing_backup_interval_minutes=${MULTI_TIER_CHECKPOINTING_BACKUP_INT_MIN}" - ``` \ No newline at end of file + ``` diff --git a/docs/guides/custom_model.md b/docs/guides/custom_model.md index 5f04904afb..eb214d92d2 100644 --- a/docs/guides/custom_model.md +++ b/docs/guides/custom_model.md @@ -20,9 +20,9 @@ This document provides a guide to optimize and customize your LLM model configurations for higher performance (i.e. MFU) on Cloud TPU. Note that this document focuses exclusively on performance tuning. The analysis of model quality and convergence behavior is outside of scope. -## Step 1. Identify Initial Configs +## Step 1. Identify initial configs -To begin, identify your model's size, review open-source model configs, and establish the initial configurations for each block. You can use our [reference calculator (on Colab)](https://colab.research.google.com/github/AI-Hypercomputer/maxtext/blob/main/docs/guides/llm_calculator.ipynb) to estimate parameters and FLOPs for dense, Mixtral-like Mixture of Experts (MoE), and DeepSeek-like MoE models to help you estimate the parameter count and FLOPs. +To begin, identify your model's size, review open-source model configs, and establish the initial configurations for each block. You can use our [reference calculator (on Colab)](https://colab.research.google.com/github/AI-Hypercomputer/maxtext/blob/main/docs/explanations/llm_calculator.ipynb) to estimate parameters and FLOPs for dense, Mixtral-like Mixture of Experts (MoE), and DeepSeek-like MoE models to help you estimate the parameter count and FLOPs. Based on resources like [Language Modeling from Scratch](https://github.com/stanford-cs336/spring2025-lectures/blob/e9cb2488fdb53ea37f0e38924ec3a1701925cef3/nonexecutable/2025%20Lecture%203%20-%20architecture.pdf), we observe common architectural ratios for dense models, as shown below: @@ -35,7 +35,7 @@ For MoE models, * sparsity (`num_experts / num_experts_per_tok`): 4-32 * `moe_mlp_dim / emb_dim`: 0.3-3 -## Step 2. Consider TPU Best Practices +## Step 2. Consider TPU best practices ### Model configs @@ -65,7 +65,7 @@ Use these general runtime configurations to improve your model's performance. * **Benchmark**. For consistent speed tests, set `reuse_example_batch=1` to repeatedly use the same data batch, isolating computation speed from data loading. Or use on-the-fly generated data by setting `dataset_type=synthetic`. (roofline-sharding)= -## Step 3. Choose Efficient Sharding Strategies Using Roofline Analysis +## Step 3. Choose efficient sharding strategies using Roofline Analysis To achieve good performance, it's often necessary to co-design the model's dimensions (like the MLP dimension) along with the sharding strategy. We have included examples for Trillium that demonstrate which sharding approaches work well for specific models. We recommend reading [](sharding) and Jax’s [scaling book](https://jax-ml.github.io/scaling-book/sharding/). @@ -139,7 +139,7 @@ TP=4 corresponds to a custom Trillium mesh, an 8x8 ring of 2x2 subrings (the TP To use TP=16, we need M > 80k (ideally larger, 100k+). We have used this in a custom dense model (900B, M=131k), which performs very well even at 1k per device tokens (scaling to 25k+ with a reasonable global batch). -## Step 4. Analyze Experiments +## Step 4. Analyze experiments With your configs, begin experimenting to evaluate the model's performance. We strongly recommend capturing a profile by following these [instructions](https://docs.jax.dev/en/latest/profiling.html#). If you are using MaxText, this can be done by simply setting `profiler=xplane` in your configuration. @@ -209,4 +209,4 @@ Objective was to demonstrate achieving reasonable MFU on a low batch setting (2k | **Total Params** | 1.04E+13 | | **Active Params** | 3.76E+11 | | **MFU (1 pod Trillium)** | 34.5% | -| **MFU(16 pods Trillium)** | 26.2% | \ No newline at end of file +| **MFU(16 pods Trillium)** | 26.2% | diff --git a/docs/guides/data_input_pipeline.md b/docs/guides/data_input_pipeline.md index f84f959b88..af350eaaa5 100644 --- a/docs/guides/data_input_pipeline.md +++ b/docs/guides/data_input_pipeline.md @@ -21,10 +21,9 @@ Currently MaxText has three data input pipelines: | Pipeline | Dataset formats | Features | Limitations | | -------- | --------------- | -------- | ----------- | -| **[Grain](data_input_grain.md)** (recommended)| [ArrayRecord](https://github.com/google/array_record) (random access, available through [Tensorflow Datasets](https://www.tensorflow.org/datasets/catalog/overview), or [conversion](https://github.com/google/array_record/tree/main/beam))
[Parquet](https://arrow.apache.org/docs/python/parquet.html) (sequential access) | With arrayrecord: fully deterministic, resilient to preemption; global shuffle
With parquet: performant; fully deterministic, resilient to preemption; hierarchical shuffle | | -| **[Hugging Face](data_input_hf.md)** | datasets in [Hugging Face Hub](https://huggingface.co/datasets)
local/Cloud Storage datasets in json, parquet, arrow, csv, txt (sequential access) | no download needed, convenience;
multiple formats | limit scalability using the Hugging Face Hub (no limit using Cloud Storage);
non-deterministic with preemption
(deterministic without preemption)
| -| **[TFDS](data_input_tfds.md)** | TFRecord (sequential access), available through [Tensorflow Datasets](https://www.tensorflow.org/datasets/catalog/overview) | performant | only supports TFRecords;
non-deterministic with preemption
(deterministic without preemption) | - +| **[Grain](data_input_pipeline/data_input_grain.md)** (recommended)| [ArrayRecord](https://github.com/google/array_record) (random access, available through [Tensorflow Datasets](https://www.tensorflow.org/datasets/catalog/overview), or [conversion](https://github.com/google/array_record/tree/main/beam))
[Parquet](https://arrow.apache.org/docs/python/parquet.html) (sequential access) | With arrayrecord: fully deterministic, resilient to preemption; global shuffle
With parquet: performant; fully deterministic, resilient to preemption; hierarchical shuffle | | +| **[Hugging Face](data_input_pipeline/data_input_hf.md)** | datasets in [Hugging Face Hub](https://huggingface.co/datasets)
local/Cloud Storage datasets in json, parquet, arrow, csv, txt (sequential access) | no download needed, convenience;
multiple formats | limit scalability using the Hugging Face Hub (no limit using Cloud Storage);
non-deterministic with preemption
(deterministic without preemption)
| +| **[TFDS](data_input_pipeline/data_input_hf.md)** | TFRecord (sequential access), available through [Tensorflow Datasets](https://www.tensorflow.org/datasets/catalog/overview) | performant | only supports TFRecords;
non-deterministic with preemption
(deterministic without preemption) | ## Multihost dataloading best practice Training in a multi-host environment presents unique challenges for data input pipelines. An effective data loading strategy must address three key issues: @@ -46,7 +45,8 @@ In MaxText, this is best supported by the ArrayRecord format using the Grain inp ```{toctree} :hidden: -data_input_grain -data_input_hf -data_input_tfds +data_input_pipeline/data_input_grain +data_input_pipeline/data_input_hf +data_input_pipeline/data_input_tfds +data_input_pipeline/data_pipeline_perf.md ``` diff --git a/docs/guides/data_input_grain.md b/docs/guides/data_input_pipeline/data_input_grain.md similarity index 100% rename from docs/guides/data_input_grain.md rename to docs/guides/data_input_pipeline/data_input_grain.md diff --git a/docs/guides/data_input_hf.md b/docs/guides/data_input_pipeline/data_input_hf.md similarity index 98% rename from docs/guides/data_input_hf.md rename to docs/guides/data_input_pipeline/data_input_hf.md index 2b8571a459..2e72cd332a 100644 --- a/docs/guides/data_input_hf.md +++ b/docs/guides/data_input_pipeline/data_input_hf.md @@ -34,7 +34,7 @@ hf_eval_files: 'gs:////*-validation-*.parquet' # match the val # or a local path containing tokenizer in a format supported by transformers.AutoTokenizer tokenizer_path: 'google-t5/t5-large' # for using https://huggingface.co/google-t5/t5-large ``` -## Limitations and Recommendations +## Limitations and recommendations 1. Streaming data directly from Hugging Face Hub may be impacted by the traffic of the server. During peak hours you may encounter "504 Server Error: Gateway Time-out". It's recommended to download the Hugging Face dataset to a Cloud Storage bucket or disk for the most stable experience. 2. Streaming data directly from Hugging Face Hub works in multi-host settings with a small number of hosts. With a host number larger than 16, you might encounter a "read time out" error. 3. Only supports num_epoch=1 at the moment. diff --git a/docs/guides/data_input_tfds.md b/docs/guides/data_input_pipeline/data_input_tfds.md similarity index 100% rename from docs/guides/data_input_tfds.md rename to docs/guides/data_input_pipeline/data_input_tfds.md diff --git a/docs/explanations/data_pipeline_perf.md b/docs/guides/data_input_pipeline/data_pipeline_perf.md similarity index 90% rename from docs/explanations/data_pipeline_perf.md rename to docs/guides/data_input_pipeline/data_pipeline_perf.md index ea115367bb..7230aac9b9 100644 --- a/docs/explanations/data_pipeline_perf.md +++ b/docs/guides/data_input_pipeline/data_pipeline_perf.md @@ -26,7 +26,7 @@ Ideally, while the accelerator is busy with the current step's computation (e.g. If data loading is faster than computation, the pipeline is not a bottleneck. Optimizing it further will **not** improve training speed. -```{figure} ../_static/data_input_goal.png +```{figure} ../../_static/data_input_goal.png Status of TPU and CPU during training stages. ``` @@ -36,7 +36,7 @@ For this overlap to happen, the data pipeline (on the CPU) and the model computa * **Good (parallel)**: The trace on the right shows the CPU (bottom tracks) is busy fetching data at the same time the TPU (top track) is computing. * **Bad (sequential)**: The trace on the left shows a **gap** in TPU utilization, where the TPU is idle. This gap is often caused by **forcing synchronization** (as explained in the following section), not necessarily a slow pipeline. While speeding up data loading might narrow this gap, only removing the synchronization eliminates the gap and achieves true parallelism. -```{figure} ../_static/data_profile.png +```{figure} ../../_static/data_profile.png Example profiles of sequential (left) vs. parallel (right) data loading with TPU computation. ``` @@ -55,5 +55,5 @@ You can check if your data pipeline meets the performance goal in two ways: ## How to speed up a slow data pipeline If your profile confirms that data loading is parallel but still slower than computation, then data loading is the bottleneck. Here are a few ways to speed it up: -1. **Tune Grain**: If you are using the [Grain data pipeline](../guides/data_input_grain.md), start by tuning the `grain_worker_count`. If adjusting the worker count isn't enough, use the [Grain performance and debugging tool](https://google-grain.readthedocs.io/en/latest/tutorials/dataset_debugging_tutorial.html) to find the specific bottleneck. +1. **Tune Grain**: If you are using the [Grain data pipeline](data_input_grain.md), start by tuning the `grain_worker_count`. If adjusting the worker count isn't enough, use the [Grain performance and debugging tool](https://google-grain.readthedocs.io/en/latest/tutorials/dataset_debugging_tutorial.html) to find the specific bottleneck. 2. **Pre-process offline**: Perform as much data preparation as possible offline. Apply only light-weight preprocessing during training. diff --git a/docs/guides/features_and_diagnostics.md b/docs/guides/features_and_diagnostics.md index fe209b7406..58049503f7 100644 --- a/docs/guides/features_and_diagnostics.md +++ b/docs/guides/features_and_diagnostics.md @@ -14,9 +14,9 @@ limitations under the License. --> -# Features and Diagnostics +# Features and diagnostics -## Collect Stack Traces +## Collect stack traces When running a Single Program, Multiple Data (SPMD) job on accelerators, the overall process can hang if there is any error or any VM hangs/crashes for some reason. In this scenario, capturing stack traces will help to identify and troubleshoot the issues for the jobs running on TPU VMs. The following configurations will help to debug a fault or when a program is stuck or hung somewhere by collecting stack traces. Change the parameter values accordingly in `src/MaxText/configs/base.yml`: @@ -30,10 +30,11 @@ jsonPayload.verb="stacktraceanalyzer" Here is the related PyPI package: https://pypi.org/project/cloud-tpu-diagnostics. -## Ahead of Time Compilation (AOT) +(aot-compilation)= +## Ahead of Time compilation (AOT) To compile your training run ahead of time, we provide a tool `train_compile.py`. This tool allows you to compile the main `train_step` in `train.py` for target hardware (e.g. a large number of v5e devices) without using the full cluster. -### TPU Support +### TPU support You may use only a CPU or a single VM from a different family to pre-compile for a TPU cluster. This compilation helps with two main goals: @@ -66,7 +67,7 @@ python3 -m MaxText.train_compile src/MaxText/configs/base.yml compile_topology=v per_device_batch_size=4 steps=10000 learning_rate=1e-3 ``` -**Step 2: Run train.py and load the compiled function** +**Step 2: Run `train.py` and load the compiled function** To load the compiled train_step, you just need to pass `compiled_trainstep_file=my_compiled_train.pickle` into `train.py`: ```sh @@ -74,13 +75,13 @@ To load the compiled train_step, you just need to pass `compiled_trainstep_file= export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true" python3 -m MaxText.train src/MaxText/configs/base.yml run_name=example_load_compile \ compiled_trainstep_file=my_compiled_train.pickle \ - global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3 \ + global_parameter_scale=16 per_device_batch_size=4 steps=10000 learning_rate=1e-3 \ base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket ``` In the save step of example 2 above we included exporting the compiler flag `LIBTPU_INIT_ARGS` and `learning_rate` because those affect the compiled object `my_compiled_train.pickle.` The sizes of the model (e.g. `global_parameter_scale`, `max_sequence_length` and `per_device_batch`) are fixed when you initially compile via `compile_train.py`, you will see a size error if you try to run the saved compiled object with different sizes than you compiled with. However a subtle note is that the **learning rate schedule** is also fixed when you run `compile_train` - which is determined by both `steps` and `learning_rate`. The optimizer parameters such as `adam_b1` are passed only as shaped objects to the compiler - thus their real values are determined when you run `train.py`, not during the compilation. If you do pass in different shapes (e.g. `per_device_batch`), you will get a clear error message reporting that the compiled signature has different expected shapes than what was input. If you attempt to run on different hardware than the compilation targets requested via `compile_topology`, you will get an error saying there is a failure to map the devices from the compiled to your real devices. Using different XLA flags or a LIBTPU than what was compiled will probably run silently with the environment you compiled in without error. However there is no guaranteed behavior in this case; you should run in the same environment you compiled in. -### GPU Support +### GPU support Ahead-of-time compilation is also supported for GPUs with some differences from TPUs: 1. GPU does not support compilation across hardware: A GPU host is still required to run AoT compilation, but a single GPU host can compile a program for a larger cluster of the same hardware. @@ -102,7 +103,7 @@ python3 -m MaxText.train_compile src/MaxText/configs/base.yml compile_topology=a **Step 2: Run `train.py` and load the compiled function** -To load the compiled train_step, you just need to pass `compiled_trainstep_file=my_compiled_train.pickle` into `train.py`: +To load the compiled `train_step`, you just need to pass `compiled_trainstep_file=my_compiled_train.pickle` into `train.py`: ```sh # Run the below on each of the 4 target A3 hosts. export XLA_FLAGS="--xla_gpu_enable_async_collectives=true" @@ -115,6 +116,6 @@ python3 -m MaxText.train src/MaxText/configs/base.yml run_name=example_load_comp As in the TPU case, note that the compilation environment must match the execution environment, in this case by setting the same `XLA_FLAGS`. -## Automatically Upload Logs to Vertex Tensorboard +## Automatically upload logs to Vertex AI Tensorboard MaxText supports automatic upload of logs collected in a directory to a Tensorboard instance in Vertex AI. Follow [](use_vertex_ai_tensorboard.md) to know more. diff --git a/docs/guides/gcp_workload_observability.md b/docs/guides/gcp_workload_observability.md index 3a04a86b33..5328120ed7 100644 --- a/docs/guides/gcp_workload_observability.md +++ b/docs/guides/gcp_workload_observability.md @@ -14,7 +14,7 @@ limitations under the License. --> -# Enable GCP Workload Observabiltiy +# Enable GCP workload observabiltiy This guide provides an overview on how to enable GCP workload observability for your MaxText workload. ## Overview @@ -27,10 +27,10 @@ Users should work with their Customer Engineer (CE) and the Google team to defin This guide layouts how to enable the feature for your MaxText workload. -## Enabling GCP Workload Observabiltiy +## Enabling GCP workload observabiltiy User can control which metric they want to report via config: -### Heartbeat metric +### Heartbeat metric - This metric will be a boolean flag. - To turn on this metric, set `report_heartbeat_metric_for_gcp_monitoring` to `True` - To control the frequency of heartbeat reporting (default is every 5 seconds), set `heartbeat_reporting_interval_in_seconds` to your desired value. @@ -39,4 +39,4 @@ User can control which metric they want to report via config: - This metric will be a double, capturing the training step time in seconds. - To turn on this metric, set `report_performance_metric_for_gcp_monitoring` to `True` -For an example, please refer to [base.yml](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/configs/base.yml). \ No newline at end of file +For an example, please refer to [base.yml](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/configs/base.yml). diff --git a/docs/guides/knowledge_distillation.md b/docs/guides/knowledge_distillation.md index 39de42ecea..8b515141a2 100644 --- a/docs/guides/knowledge_distillation.md +++ b/docs/guides/knowledge_distillation.md @@ -55,16 +55,16 @@ cd maxtext uv pip install -r requirements.txt ``` -### 1. Obtain and Prepare the Teacher Model +### 1. Obtain and prepare the teacher model -#### a. Download Model from Hugging Face +#### a. Download model from Hugging Face ```bash huggingface-cli login # Provide your Hugging Face token huggingface-cli download deepseek-ai/DeepSeek-V2-Lite-Chat --repo-type model --local-dir ~/deepseek2-16b-chat ``` -#### b. Convert Checkpoint to MaxText Format +#### b. Convert checkpoint to MaxText format MaxText requires checkpoints to be in a specific format. You'll need to convert the downloaded Hugging Face checkpoints to a MaxText-compatible checkpoint. ```bash @@ -76,15 +76,15 @@ python3 -m MaxText.convert_deepseek_family_unscanned_ckpt \ --model_size deepseek2-16b ``` -### 2. Obtain and Prepare the Student Model +### 2. Obtain and prepare the student model -#### a. Download Model from Hugging Face +#### a. Download model from Hugging Face ```bash huggingface-cli download meta-llama/Llama-2-7b-chat-hf --repo-type model --local-dir ~/llama2-7b-chat ``` -#### b. Convert Checkpoint to MaxText Format +#### b. Convert checkpoint to MaxText format MaxText requires checkpoints to be in a specific format. You'll need to convert the downloaded Hugging Face checkpoints to a MaxText-compatible checkpoint. ```bash @@ -96,10 +96,10 @@ python3 -m MaxText.llama_or_mistral_ckpt \ --model-size llama2-7b ``` -### 3. Generate Dataset using the Teacher Model +### 3. Generate dataset using the teacher model Once the teacher model's checkpoint is in the MaxText format, you can run inference to generate the dataset that will be used to fine-tune the student model. -### 3.a. Run the JetStream Server +### 3.a. Run the JetStream server Example command to run JetStream server on `v4-8`: @@ -118,7 +118,7 @@ python3 -m MaxText.maxengine_server src/MaxText/configs/base.yml \ Set `multi_sampling` to `True` to generate multiple independent completions per prompt. -### 3.b. Generate Dataset using JetStream Server +### 3.b. Generate dataset using JetStream server In a new tab in your terminal, run the following command to generate dataset from teacher model. Note that this is an example command to run on `v4-8`: ```bash @@ -141,10 +141,10 @@ It's important to note that some prompts may be filtered out by pre-processing l Additionally, the generated dataset can be uploaded to either Hugging Face or Google Cloud Storage (GCS). To upload to Hugging Face, use the `upload-to-hf --hf-repo-id ` flags. To upload to GCS, use the `upload-to-gcs --gcs-bucket --gcs-data-path ` flags. -### 4. Fine-tune the Student Model using Supervised Fine Tuning +### 4. Fine-tune the student model using Supervised Fine Tuning (SFT) You can now fine-tune your smaller student model using supervised fine-tuning technique in MaxText. -### 4.a. Fine-tune the Student Model using Dataset Generated in Step 3 +### 4.a. Fine-tune the student model using dataset generated in Step 3 Example command to run fine-tuning on v4-8: @@ -162,7 +162,7 @@ python3 -m MaxText.sft_trainer src/MaxText/configs/sft.yml \ hf_access_token=$HF_TOKEN ``` -### 4.b. **[OPTIONAL]** Fine-tune the Student Model using the Original Dataset +### 4.b. **[OPTIONAL]** Fine-tune the student model using the original dataset The checkpoint from the student model's fine-tuning (on the teacher-generated dataset) can be used for a subsequent fine-tuning stage. In this step, the student model is fine-tuned on the original dataset that was initially provided to the teacher model for generating the dataset. diff --git a/docs/guides/monitor_goodput.md b/docs/guides/monitor_goodput.md index e936cb2262..4d08c47bc7 100644 --- a/docs/guides/monitor_goodput.md +++ b/docs/guides/monitor_goodput.md @@ -54,24 +54,23 @@ project, then do the following: 5. To learn more about Google Cloud Logging, visit this [page](https://cloud.google.com/logging/docs). -### Access Scopes +### Access scopes - You will need both read and write access scopes for cloud logging on both the - GPU or TPU and CPU node pools. Full cloud logging access is granted by the - following access scope during node pool creation: +You will need both read and write access scopes for cloud logging on both the +GPU or TPU and CPU node pools. Full cloud logging access is granted by the +following access scope during node pool creation: - - `https://www.googleapis.com/auth/cloud-platform` +- `https://www.googleapis.com/auth/cloud-platform` - XPK adds this access scope to the GPU, TPU and CPU node pools, so XPK is the recommended method to create clusters and node-pools in you intend to run your workloads on GKE. +XPK adds this access scope to the GPU, TPU and CPU node pools, so XPK is the recommended method to create clusters and node-pools in you intend to run your workloads on GKE. - Instructions on how to create clusters using XPK can be - found [here](https://github.com/AI-Hypercomputer/xpk/blob/main/README.md#cluster-create) and how to create workloads using XPK can be found - [here](https://github.com/AI-Hypercomputer/xpk/blob/main/README.md#workload-create). - - ```{note} - Access Scopes are immutable and workloads can only be migrated to new node pools with required access scopes. Access scopes on already created clusters cannot be updated. - ``` +Instructions on how to create clusters using XPK can be +found [here](https://github.com/AI-Hypercomputer/xpk/blob/main/README.md#cluster-create) and how to create workloads using XPK can be found +[here](https://github.com/AI-Hypercomputer/xpk/blob/main/README.md#workload-create). +```{note} +Access Scopes are immutable and workloads can only be migrated to new node pools with required access scopes. Access scopes on already created clusters cannot be updated. +``` ### Monitoring @@ -81,7 +80,7 @@ Ensure unique `run_name` for each new experiment or run Please use a unique workload name, unless you intend to monitor cumulative Goodput/Badput metrics of a previous workload along with your current workload -#### How to Monitor Goodput and Badput +#### How to monitor Goodput and Badput MaxText enables Goodput recording and monitoring by default with `enable_goodput_recording=True` and `monitor_goodput=True`. You can configure the goodput upload frequency by setting `goodput_upload_interval_seconds`. @@ -90,7 +89,7 @@ python3 -m MaxText.train src/MaxText/configs/base.yml base_output_directory=$OUT dataset_path=$DATA_PATH run_name=goodput-test-run steps=200 goodput_upload_interval_seconds=30 ``` -#### How to Monitor Step Time Deviation +#### How to monitor step time deviation MaxText enables step time deviation monitoring by default with `monitor_step_time_deviation=True`. You can configure the upload frequency by setting `step_deviation_interval_seconds`. @@ -112,7 +111,7 @@ python3 -m MaxText.train src/MaxText/configs/base.yml base_output_directory=$OUT run_name=goodput-test-run steps=200 goodput_upload_interval_seconds=30 enable_pathways_goodput=True ``` -#### How to enable Checkpoint Logging +#### How to enable checkpoint logging Checkpoint logging is currently supported through Orbax. The Goodput library reads these logs to compute checkpointing badput. To enable checkpoint logging set the `enable_checkpoint_cloud_logger` MaxText flag to `True`. @@ -126,7 +125,7 @@ If checkpointing is enabled, please enable the `enable_checkpoint_cloud_logger` 1. MaxText installs the required packages on setup: `tensorboard-plugin-profile`, `tensorflow` and `tensorboard`. 2. Follow the Tensorboard URL on MaxText logs to view all metrics in one location. -#### Visualize Goodput, Badput and Step Deviation on Google Cloud Monitoring +#### Visualize Goodput, Badput and step deviation on Google Cloud Monitoring By default, performance data ([goodput](https://cloud.google.com/monitoring/api/metrics_gcp#:~:text=workload/goodput_time), [badput](https://cloud.google.com/monitoring/api/metrics_gcp#:~:text=workload/badput_time), and [step deviation](https://cloud.google.com/monitoring/api/metrics_gcp#:~:text=workload/performance)) is automatically sent to Google Cloud Monitoring, enabling visualization on dashboards. @@ -173,7 +172,7 @@ python3 -m MaxText.train src/MaxText/configs/base.yml base_output_directory=$OUT Setting `monitor_goodput` to `False` will disable both tensorboard and GCM monitoring. -### Monitoring Raw Metrics and Dashboards +### Monitoring raw metrics and dashboards Goodput, Badput and Step Time Deviation metrics can be monitored using GCM Metrics Explorer: @@ -190,4 +189,4 @@ Goodput, Badput and Step Time Deviation metrics can be monitored using GCM Metri Represents the workload's performance metric, specifically step deviation in this context, measured by `compute.googleapis.com/workload/performance`. 3. Navigate to [Dashboards](https://console.cloud.google.com/monitoring/dashboards). -4. Create a custom dashboard if there isn't one and add useful widgets with the above mentioned metrics. \ No newline at end of file +4. Create a custom dashboard if there isn't one and add useful widgets with the above mentioned metrics. diff --git a/docs/guides/pallas_kernels_performance.md b/docs/guides/pallas_kernels_performance.md index 8e6edfbb59..eb8e7ede4b 100644 --- a/docs/guides/pallas_kernels_performance.md +++ b/docs/guides/pallas_kernels_performance.md @@ -14,7 +14,7 @@ limitations under the License. --> -# Performance Optimizations with Pallas Kernels +# Performance optimizations with Pallas kernels New to Pallas? Start with the [official docs](https://docs.jax.dev/en/latest/pallas/index.html). @@ -22,11 +22,11 @@ While JAX and the XLA compiler provide excellent out-of-the-box performance, som This guide explains **when** to consider Pallas, a **workflow** for developing and tuning kernels, and how Pallas is **used in MaxText** today. -## **🧠 The Pallas Mindset: When to Write a Custom Kernel** +## 🧠 The Pallas mindset: when to write a custom kernel Think in **roofline** terms ([All About Rooflines](https://jax-ml.github.io/scaling-book/roofline/)) and in terms of **structure the compiler can’t see**: -* **Roofline framing.** Is your op **compute-limited** (MXU at or near peak) or **bandwidth-limited** (HBM↔on-chip transfers dominate)? Pallas tends to shine when you can reduce bandwidth pressure or avoid wasted work via better tiling and scheduling. +* **Roofline framing.** Is your op **compute-limited** (MXU at or near peak) or **bandwidth-limited** (HBM↔on-chip transfers dominate)? Pallas tends to shine when you can reduce bandwidth pressure or avoid wasted work via better tiling and scheduling. * **Compiler invisibles.** Irregular sparsity, ragged batch shapes, non-contiguous memory access, and domain-specific invariants are all signals that a custom kernel could help. **Know when XLA is enough.** Before writing a custom kernel, always [profile your baseline](#1-high-level-profiling). If a standard operation (like a dense [`jnp.matmul`](https://docs.jax.dev/en/latest/_autosummary/jax.numpy.matmul.html)) is already performing well, the XLA compiler is doing its job. In these cases, a Pallas kernel will increase code complexity and maintenance burden with minimal performance improvement. @@ -36,10 +36,9 @@ Think in **roofline** terms ([All About Rooflines](https://jax-ml.github.io/scal **Autodiff note.** Pallas kernels require writing the autodiff rule manually (e.g., with [`jax.custom_vjp`](https://docs.jax.dev/en/latest/_autosummary/jax.custom_vjp.html)), since unlike other transforms such as [`shard_map`](https://docs.jax.dev/en/latest/_autosummary/jax.shard_map.html), it is very difficult to automatically infer the dual of the memory pipeline. +## 💡 Use Cases -## **💡 Use Cases** - -### **1\. Irregular Compute (MoE, ragged activations)** +### 1. Irregular compute (MoE, ragged activations) For dense, regular GEMMs, XLA’s libraries are hard to beat. The exception is **Mixture-of-Experts (MoE)** MLPs with **ragged token→expert layouts** (some tokens routed to different experts; shapes are irregular). Zero-padding to make dense tensors wastes FLOPs; a custom kernel can operate only on the actually-selected tokens. @@ -47,116 +46,108 @@ For dense, regular GEMMs, XLA’s libraries are hard to beat. The exception is * **Note:** *Megablox* is an efficient, non-capped MoE implementation in JAX. *Megablocks* refers to the equivalent PyTorch implementation. See [arXiv:2211.15841](https://arxiv.org/abs/2211.15841) for more details. -### **2\. Memory-Access-Bound Work (Attention)** +### 2. Memory-Access-Bound work (attention) Attention kernels are classically **bandwidth-limited** if you materialize the full \[L,L\] score matrix. A Pallas kernel can block **Q/K/V** into tiles that fit on-chip and perform **online softmax accumulation**, never storing the massive intermediate. * MaxText uses a Pallas attention kernel for training (Flash/Splash-style) and **paged/ragged** attention for inference to efficiently fetch KV cache pages and handle non-contiguous layouts. - -## **🛠️ Pallas Kernels in MaxText** +## 🛠️ Pallas kernels in MaxText To maximize performance, MaxText uses custom Pallas kernels for memory-bandwidth-bound or structurally irregular operations that a general-purpose compiler cannot optimize as effectively. Below are the key kernels we use. **Note**: Examples evolve; treat this list as guidance. * **Training Attention (Flash/Splash-style):** This kernel is the default for training Transformer models in MaxText, such as DeepSeek, Gemma and Llama. It avoids creating the large \[L,L\] attention matrix to save memory, processing data in smaller, tiled chunks with online softmax accumulation. - * [`src/MaxText/kernels/splash_attention_kernel.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/splash_attention_kernel.py) + * [`src/MaxText/kernels/splash_attention_kernel.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/splash_attention_kernel.py) * **Serving Attention (Paged & Ragged):** For high-throughput inference, this kernel efficiently fetches non-contiguous "pages" of the KV cache from memory. It is a key optimization for our serving stack and is used for models running on MaxText's inference engine. - * [`src/MaxText/inference/paged_attention.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/inference/paged_attention.py) - * [`src/MaxText/inference/paged_attention_kernel_v2.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/inference/paged_attention_kernel_v2.py) -* **MoE Grouped Matmul (Megablox GMM):** Sparse/irregular grouped GEMMs driven by host-built metadata. + * [`src/MaxText/inference/paged_attention.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/inference/paged_attention.py) + * [`src/MaxText/inference/paged_attention_kernel_v2.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/inference/paged_attention_kernel_v2.py) +* **MoE Grouped Matmul (Megablox GMM):** Sparse/irregular grouped GEMMs driven by host-built metadata. > This is an efficient computation method for Mixture-of-Experts (MoE) models like DeepSeek, Llama 4, Mixtral and Qwen-MoE. In MoE, each token is processed by only a few "experts," which is inefficient for standard matrix multiplication. Megablox solves this by having the CPU (**host**) first create a routing plan (**metadata**) that assigns tokens to experts. The accelerator (**device**) then uses this plan to perform many small, dense matrix multiplications in parallel (**Grouped Matrix Multiplication**), avoiding wasted work on unused experts. * [`src/MaxText/kernels/megablox/gmm.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/kernels/megablox/gmm.py) - - **Note:** Megablox accelerates the grouped **matmul**; **routing/gating** is separate code ([`src/MaxText/layers/moe.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/moe.py)). - + **Note:** Megablox accelerates the grouped **matmul**; **routing/gating** is separate code ([`src/MaxText/layers/moe.py`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/moe.py)). -## **🔧 The Pallas Optimization Workflow: Code → Profile → Tune → Repeat** +## 🔧 The Pallas optimization workflow: code → profile → tune → repeat -### 1. High-Level Profiling +### 1. High-level profiling Give the kernel a clear name in traces and capture a profile. Always use [`jax.block_until_ready()`](https://docs.jax.dev/en/latest/_autosummary/jax.block_until_ready.html) when timing your operations. - ``` python -import jax +import jax from jax import profiler -def my_op(...): - # This name shows up in Perfetto/TensorBoard traces - with jax.named_scope("my_custom_kernel"): - out = my_kernel_wrapper(...) - return out - -# Capture a Perfetto/TensorBoard trace -with profiler.trace("/tmp/tb_profile"): - y = my_op(x) - # Stabilize timing for accurate measurement - jax.block_until_ready(y) +def my_op(...): + # This name shows up in Perfetto/TensorBoard traces + with jax.named_scope("my_custom_kernel"): + out = my_kernel_wrapper(...) + return out + +# Capture a Perfetto/TensorBoard trace +with profiler.trace("/tmp/tb_profile"): + y = my_op(x) + # Stabilize timing for accurate measurement + jax.block_until_ready(y) ``` -### **2\. Deeper Compiler View (Optional)** +### 2. Deeper compiler view (optional) For hard cases, inspect compiler dumps (e.g., LLO) to understand schedules, memory moves, and resource usage. Keep this as an advanced tool—most tuning decisions come from traces and microbenchmarks. -### **3\. Systematic Tuning** +### 3. Systematic tuning Performance is a function of interacting hyperparameters, chiefly block shapes (via [`BlockSpec`](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.BlockSpec.html)). Build a small test script (a "harness") to systematically run the kernel with different block sizes. **Record the throughput and latency** for each run, and let data, not rules of thumb, pick the winners. For a more automated approach, consider using libraries like [tune-jax](https://github.com/rdyro/tune-jax). - -## **⚙️ Understanding TPU Memory & Compute** +## ⚙️ Understanding TPU memory & compute Pallas exposes the underlying hardware primitives for you to control. -* **HBM:** High-Bandwidth Memory (standard device memory). -* **VMEM:** On-chip vector SRAM for array tiles; your kernel primarily reads/writes VMEM refs. -* **SMEM:** On-chip scalar SRAM for control/metadata (e.g., counters, small tables). -* **Semaphores:** Available for advanced async/barrier patterns in manual pipelines. -* **MXU:** The Matrix Unit, optimized for large block GEMMs/convolutions. +* **HBM:** High-Bandwidth Memory (standard device memory). +* **VMEM:** On-chip vector SRAM for array tiles; your kernel primarily reads/writes VMEM refs. +* **SMEM:** On-chip scalar SRAM for control/metadata (e.g., counters, small tables). +* **Semaphores:** Available for advanced async/barrier patterns in manual pipelines. +* **MXU:** The Matrix Unit, optimized for large block GEMMs/convolutions. * **VPU:** The Vector Processing Unit, used for elementwise/vector work. **Alignment & Constraints:** Respect TPU BlockSpec constraints (divisibility/shape rules for trailing dimensions and supported block shapes). Start with tile shapes that fit in VMEM and meet these requirements, then sweep different sizes to find the optimum. Let profiling guide you; don't assume powers of two are always best. - -## **🧱 Core Pallas Design Patterns** +## 🧱 Core Pallas design patterns These are the common techniques used in MaxText's Pallas kernels. -* **Tiling & Blocking:** Move just a tile that fits on-chip, compute on it, and write it back. -* **Explicit Pipelining:** Overlap HBM↔VMEM loads with compute to hide latency (e.g., double-buffering). -* **Online Accumulation:** Combine partial results as you go; don’t materialize huge intermediate arrays. -* **Auxiliary Metadata:** Precompute control tables (e.g., token-to-expert ranges) and keep them in fast scalar memory. +* **Tiling & Blocking:** Move just a tile that fits on-chip, compute on it, and write it back. +* **Explicit Pipelining:** Overlap HBM↔VMEM loads with compute to hide latency (e.g., double-buffering). +* **Online Accumulation:** Combine partial results as you go; don’t materialize huge intermediate arrays. +* **Auxiliary Metadata:** Precompute control tables (e.g., token-to-expert ranges) and keep them in fast scalar memory. * **Compute↔Communication Overlap:** In distributed runs, overlap local work with cross-device traffic when possible. - -## **✍️ Writing & Integrating a Kernel** +## ✍️ Writing & integrating a kernel A Pallas kernel is a Python function that operates on Refs (references to array tiles). When this function is passed to [`pl.pallas_call`](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.pallas_call.html), it will be compiled and scheduled by Pallas. -### **Example 1: Minimal Elementwise Add** +### Example 1: Minimal elementwise add Shows the kernel/ref pattern used by `pallas_call`. - ```python -import jax -import jax.numpy as jnp +import jax +import jax.numpy as jnp from jax.experimental import pallas as pl -def add_vectors_kernel(x_ref, y_ref, o_ref): - o_ref[:] = x_ref[:] + y_ref[:] +def add_vectors_kernel(x_ref, y_ref, o_ref): + o_ref[:] = x_ref[:] + y_ref[:] -def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array: - assert x.shape == y.shape - return pl.pallas_call( - add_vectors_kernel, - out_shape=jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype), - )(x, y) +def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array: + assert x.shape == y.shape + return pl.pallas_call( + add_vectors_kernel, + out_shape=jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype), + )(x, y) ``` -### **Example 2: Blocked 2D Add with BlockSpec** +### Example 2: Blocked 2D add with BlockSpec This example shows how to map a grid of blocks over larger arrays. @@ -166,68 +157,64 @@ import jax.numpy as jnp from jax.experimental import pallas as pl def tile_add_kernel(x_ref, y_ref, o_ref): - # Operate on the tile slices handed in by BlockSpecs (already in VMEM on TPU). - o_ref[:, :] = x_ref[:, :] + y_ref[:, :] + # Operate on the tile slices handed in by BlockSpecs (already in VMEM on TPU). + o_ref[:, :] = x_ref[:, :] + y_ref[:, :] def tile_add(x: jax.Array, y: jax.Array) -> jax.Array: - assert x.shape == y.shape and x.ndim == 2 - B0 = min(128, x.shape[0]) # Example choice; tune this with a sweep - B1 = x.shape[1] # Full width tile (for illustration) - - # Map program id (tile index) -> tile origin in the full (HBM) array. - # NOTE: The runtime advances origins by `block_shape`, so `i` is already a tile - # index. Do NOT multiply by B0 here. - in_out_spec = pl.BlockSpec( - block_shape=(B0, B1), - index_map=lambda i: (i, 0), - ) - - # Grid is implied by data & block shape (use ceiling-division helper). - grid = (pl.cdiv(x.shape[0], B0),) - # Note: grid size can also be computed dynamically at runtime. - - return pl.pallas_call( - tile_add_kernel, - out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype), - in_specs=[in_out_spec, in_out_spec], - out_specs=in_out_spec, - grid=grid, - )(x, y) + assert x.shape == y.shape and x.ndim == 2 + B0 = min(128, x.shape[0]) # Example choice; tune this with a sweep + B1 = x.shape[1] # Full width tile (for illustration) + + # Map program id (tile index) -> tile origin in the full (HBM) array. + # NOTE: The runtime advances origins by `block_shape`, so `i` is already a tile + # index. Do NOT multiply by B0 here. + in_out_spec = pl.BlockSpec( + block_shape=(B0, B1), + index_map=lambda i: (i, 0), + ) + + # Grid is implied by data & block shape (use ceiling-division helper). + grid = (pl.cdiv(x.shape[0], B0),) + # Note: grid size can also be computed dynamically at runtime. + + return pl.pallas_call( + tile_add_kernel, + out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype), + in_specs=[in_out_spec, in_out_spec], + out_specs=in_out_spec, + grid=grid, + )(x, y) ``` **Tip:** In practice, you’ll **sweep** `(B0, B1)`-i.e., try a small grid of tile sizes and pick the fastest. Focus tuning on block shapes; treat grid as derived. - -## **⏩ Pipelining Best Practices (TPU)** +## ⏩ Pipelining best practices (TPU) Prefer `pl.pallas_call` with scratch buffers allocated in the appropriate memory space (VMEM/SMEM) and use multi-buffering to overlap HBM loads with compute. Advanced pipelining to consider: custom prefetch block order via a scalar prefetch grid (for details see [here](https://docs.jax.dev/en/latest/pallas/tpu/sparse.html)), which lets you control block execution order based on runtime values. -## **🌐 Distributed Execution** +## 🌐 Distributed execution Dispatch a kernel on multiple devices with `jax.shard_map`. It’s usually simpler and more maintainable than in-kernel cross-device communication. While Pallas supports low-level comms, `shard_map` is the right first choice for multi-device parallelism, and you can **communicate with `shard_map` collectives** when needed. - -## **🐞 Debugging Tips** +## 🐞 Debugging tips * Use `interpret=True` in `pallas_call` to run the kernel body in a Python interpreter backend, simulating device execution on CPU without lowering through XLA. -* Start with a tiny problem size and assert on invariants inside the kernel. +* Start with a tiny problem size and assert on invariants inside the kernel. * Add `jax.named_scope` liberally so kernels are easy to spot in performance traces. +## ✅ Putting it all together (checklist) -## **✅ Putting It All Together (Checklist)** - -1. **Profile** the baseline using `named_scope` and `block_until_ready`. -2. **Tile arrays into smaller chunks using BlockSpecs** (virtually always necessary, even for simple kernels). -3. Build a **sweep harness** for block shapes (and optionally scalar prefetch grid choices). -4. **Validate** end-to-end performance in the model, not just microbenchmarks. +1. **Profile** the baseline using `named_scope` and `block_until_ready`. +2. **Tile arrays into smaller chunks using BlockSpecs** (virtually always necessary, even for simple kernels). +3. Build a **sweep harness** for block shapes (and optionally scalar prefetch grid choices). +4. **Validate** end-to-end performance in the model, not just microbenchmarks. 5. Consider **maintainability** and guard the new kernel with tests. 6. Consider applying **`jax.vmap`** to a Pallas kernel to simplify implementation; think of it as prepending grid dimensions automatically. - ## 📚 References * **Pallas Docs & Quickstart:** [docs.jax.dev/en/latest/pallas/index.html](https://docs.jax.dev/en/latest/pallas/index.html) * **JAX Profiling Guides:** [jax.readthedocs.io/en/latest/profiling.html](https://jax.readthedocs.io/en/latest/profiling.html) * **Manual Parallelism (shard_map):** [docs.jax.dev/en/latest/notebooks/shard_map.html](https://docs.jax.dev/en/latest/notebooks/shard_map.html) -* **Distributed Pallas on TPU:** [docs.jax.dev/en/latest/pallas/tpu/distributed.html](https://docs.jax.dev/en/latest/pallas/tpu/distributed.html) \ No newline at end of file +* **Distributed Pallas on TPU:** [docs.jax.dev/en/latest/pallas/tpu/distributed.html](https://docs.jax.dev/en/latest/pallas/tpu/distributed.html) diff --git a/docs/guides/run_maxtext.md b/docs/guides/run_maxtext.md new file mode 100644 index 0000000000..7e11fc09d5 --- /dev/null +++ b/docs/guides/run_maxtext.md @@ -0,0 +1,10 @@ +# Running MaxText on different environments + +```{toctree} +:maxdepth: 1 + +run_maxtext/run_maxtext_localhost.md +run_maxtext/run_maxtext_single_host_gpu.md +run_maxtext/run_maxtext_via_xpk.md +run_maxtext/run_maxtext_via_pathways.md +``` diff --git a/docs/guides/run_maxtext_localhost.md b/docs/guides/run_maxtext/run_maxtext_localhost.md similarity index 79% rename from docs/guides/run_maxtext_localhost.md rename to docs/guides/run_maxtext/run_maxtext_localhost.md index 725d09eebb..8459df4668 100644 --- a/docs/guides/run_maxtext_localhost.md +++ b/docs/guides/run_maxtext/run_maxtext_localhost.md @@ -1,4 +1,4 @@ -# Run MaxText on Localhost or Single Host VM +# Running MaxText on localhost or single host VM ## Objective This guide provides comprehensive instructions for setting up MaxText on a local machine or single-host environment, covering everything from cloning the repo and dependency installation to building with Docker. By walking through the process of pre-training a small model, you will gain the foundational knowledge to run jobs on TPUs/GPUs. @@ -6,13 +6,13 @@ This guide provides comprehensive instructions for setting up MaxText on a local ## Prerequisites Before you can begin a training run, you need to configure your storage environment and set up the basic MaxText configuration. -### Setup Google Cloud Storage Bucket +### Setup Google Cloud storage bucket You'll need a GCS bucket to store all your training artifacts, such as logs, metrics, and model checkpoints. 1. In your Google Cloud project, create a new storage bucket. 2. Your TPU or GPU VMs require read/write access to this bucket. The simplest way to grant this is by assigning the `Storage Admin` (`roles/storage.admin`) role to the service account associated with your VMs. -### Setup Maxtext +### Setup MaxText MaxText uses a primary YAML file, `configs/base.yml`, to manage its settings. This default configuration sets up a llama2 style decoder-only model with approximately 1 billion parameters. * Before running your first model, take a moment to review this file. Pay special attention to these core settings: @@ -21,28 +21,26 @@ MaxText uses a primary YAML file, `configs/base.yml`, to manage its settings. Th - `max_target_length`: The maximum sequence length for the model. - `learning_rate`: The core hyperparameter for the optimizer. - Mode shape parameters: `base_num_decoder_layers`, `base_emb_dim`, `base_num_query_heads`, `base_num_kv_heads`, and `head_dim`. -* **Override Settings (Optional):** You can modify training parameters in two ways: by editing `configs/base.yml` directly or by passing them as command-line arguments to the training script which is the recommended method. For example, to change the number of training steps, you can pass `--steps=500` when running `train.py`. +* **Override settings (optional):** You can modify training parameters in two ways: by editing `configs/base.yml` directly or by passing them as command-line arguments to the training script which is the recommended method. For example, to change the number of training steps, you can pass `--steps=500` when running `train.py`. * **Note**: You **must** update the variable `base_output_directory` which is initialized in `configs/base.yml` to point to a folder within the GCS bucket you just created (e.g., `gs://your-bucket-name/maxtext-output`). ## Development Local development on a single host TPU/GPU VM is a convenient way to run MaxText on a single host. It doesn't scale to multiple hosts but is a good way to learn about MaxText. The following describes how to run Maxtext on TPU/GPU VMs. -### Run Maxtext on Single Host VM +### Run MaxText on single host VM 1. Create and SSH to the single host VM of your choice. You can use any available single host TPU, such as `v5litepod-8`, `v5p-8`, or `v4-8`. For GPUs, you can use `nvidia-h100-mega-80gb`, `nvidia-h200-141gb`, or `nvidia-b200`. For setting up a TPU VM, use the Cloud TPU documentation available at https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm. For a GPU setup, refer to the guide at https://cloud.google.com/compute/docs/gpus/create-vm-with-gpus. 2. Clone MaxText onto that VM. ```bash - # Clone the repository git clone https://github.com/google/maxtext.git cd maxtext ``` 3. Once you have cloned the repository, you have two primary options for setting up the necessary dependencies on your VM: Installing in a Python Environment, or building a Docker container. For single host workloads, we recommend to install dependencies in a python environment, and for multihost workloads we recommend the containerized approach. -Within the root directory of the cloned repo, install dependencies and the pre-commit hook by running: +Within the root directory of the cloned repo, create a virtual environment and install dependencies and the pre-commit hook by running: ```bash -# Create a virtual environment and install dependencies python3.12 -m venv ~/venv-maxtext source ~/venv-maxtext/bin/activate bash setup.sh DEVICE={tpu|gpu} @@ -52,21 +50,21 @@ bash setup.sh DEVICE={tpu|gpu} After the installation is complete, run a short training job using synthetic data to confirm everything is working correctly. This command trains a model for just 10 steps. Remember to replace `$YOUR_JOB_NAME` with a unique name for your run and `gs://` with the path to the GCS bucket you configured in the prerequisites. ```bash -python3 -m MaxText.train MaxText/configs/base.yml \ +python3 -m MaxText.train src/MaxText/configs/base.yml \ run_name=$YOUR_JOB_NAME \ base_output_directory=gs:// \ dataset_type=synthetic \ steps=10 ``` -**Optional**: If you want to try training on a real dataset, see [Data Input Pipeline](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline.md) for data input options from sources like HuggingFace, Grain, and TFDS. +**Optional**: If you want to try training on a real dataset, see [Data Input Pipeline](data-input-pipeline) for data input options from sources like HuggingFace, Grain, and TFDS. -#### Generate Sample Output (Decoding) +#### Generate sample output (decoding) To demonstrate model output, run the following command: ```bash -python3 -m MaxText.decode MaxText/configs/base.yml \ +python3 -m MaxText.decode src/MaxText/configs/base.yml \ run_name=$YOUR_JOB_NAME \ base_output_directory=gs:// \ per_device_batch_size=1 @@ -74,11 +72,11 @@ python3 -m MaxText.decode MaxText/configs/base.yml \ **Note:** Because the model hasn't been properly trained, the output text will be random. To generate meaningful output, you need to load a trained checkpoint using the `load_parameters_path` argument. -### Running Models Using Provided Configs -Maxtext provides many OSS model configs that you can use directly to run training jobs on those model-specific architectures. These model-specific YAML files are located in `MaxText/configs/models` for TPU-oriented defaults, and `MaxText/configs/models/gpu` for GPU-oriented defaults. +### Running models using provided configs +MaxText provides many OSS model configs that you can use directly to run training jobs on those model-specific architectures. These model-specific YAML files are located in `src/MaxText/configs/models` for TPU-oriented defaults, and `src/MaxText/configs/models/gpu` for GPU-oriented defaults. #### Training on TPUs -To use a pre-configured model for TPUs, you override the `model_name` parameter, and MaxText will automatically load the corresponding configuration from the `MaxText/configs/models` directory and merge it with the settings from `MaxText/configs/base.yml`. +To use a pre-configured model for TPUs, you override the `model_name` parameter, and MaxText will automatically load the corresponding configuration from the `src/MaxText/configs/models` directory and merge it with the settings from `src/MaxText/configs/base.yml`.
llama3-8b (TPU) @@ -107,13 +105,13 @@ python3 -m MaxText.train MaxText/configs/base.yml \
#### Training on GPUs -To use a GPU-optimized configuration, you should specify the path to the model's YAML file within the `MaxText/configs/models/gpu` directory as the main config file in the command. These files typically inherit from `base.yml` and set the appropriate `model_name` internally, as well as GPU-specific settings. +To use a GPU-optimized configuration, you should specify the path to the model's YAML file within the `src/MaxText/configs/models/gpu` directory as the main config file in the command. These files typically inherit from `base.yml` and set the appropriate `model_name` internally, as well as GPU-specific settings.
mixtral-8x7b (GPU) ```bash -python3 -m MaxText.train MaxText/configs/models/gpu/mixtral_8x7b.yml \ +python3 -m MaxText.train src/MaxText/configs/models/gpu/mixtral_8x7b.yml \ run_name=$YOUR_JOB_NAME \ base_output_directory=gs:// \ dataset_type=synthetic \ @@ -126,7 +124,7 @@ This will load `gpu/mixtral_8x7b.yml`, which inherits from `base.yml`. llama3-8b (GPU) ```bash -python3 -m MaxText.train MaxText/configs/models/gpu/llama3-8b.yml \ +python3 -m MaxText.train src/MaxText/configs/models/gpu/llama3-8b.yml \ run_name=$YOUR_JOB_NAME \ base_output_directory=gs:// \ dataset_type=synthetic \ diff --git a/docs/guides/single_host_gpu.md b/docs/guides/run_maxtext/run_maxtext_single_host_gpu.md similarity index 91% rename from docs/guides/single_host_gpu.md rename to docs/guides/run_maxtext/run_maxtext_single_host_gpu.md index f70e5bdf63..229d6d079b 100644 --- a/docs/guides/single_host_gpu.md +++ b/docs/guides/run_maxtext/run_maxtext_single_host_gpu.md @@ -14,7 +14,7 @@ limitations under the License. --> -# Maxtext on Single host GPU +# Running MaxText on single host GPU This is a short guide to run Maxtext on GPU. For this current set of instructions the GPUs used are A3-high. This is a single node 8 H100 instruction. @@ -61,13 +61,13 @@ https://stackoverflow.com/questions/72932940/failed-to-initialize-nvml-unknown-e ## Install MaxText -Clone Maxtext: +Clone MaxText: ```bash git clone https://github.com/AI-Hypercomputer/maxtext.git ``` -## Build Maxtext Docker Image +## Build MaxText Docker image This builds a docker image called `maxtext_base_image`. You can retag to a different name. @@ -81,7 +81,7 @@ cd maxtext ```bash export LOCAL_IMAGE_NAME= -sudo bash docker_build_dependency_image.sh DEVICE=gpu +sudo bash docker_build_dependency_image.sh DEVICE=gpu docker tag maxtext_base_image $LOCAL_IMAGE_NAME docker push $LOCAL_IMAGE_NAME ``` @@ -98,7 +98,7 @@ specify it explicitly: Test the docker, to see if jax can see all the 8 GPUs ```bash -sudo docker run maxtext_base_image:latest python3 -c "import jax; print(jax.devices())" +sudo docker run maxtext_base_image:latest python3 -c "import jax; print(jax.devices())" ``` You should see the following: @@ -114,7 +114,7 @@ We will Run the next commands from inside the docker for convenience. ## SSH into the docker ```bash -sudo docker run --runtime=nvidia --gpus all -it maxtext_base_image:latest bash +sudo docker run --runtime=nvidia --gpus all -it maxtext_base_image:latest bash ``` If you do not wish to ssh execute the next set of commands by prepending the following: @@ -147,8 +147,8 @@ Hardware: GPU ``` ```bash -python3 -m MaxText.train src/MaxText/configs/base.yml run_name=gpu01 base_output_directory=/deps/output \ - dataset_type=synthetic enable_checkpointing=True steps=10 attention=cudnn_flash_te scan_layers=False \ +python3 -m MaxText.train src/MaxText/configs/base.yml run_name=gpu01 base_output_directory=/deps/output \ + dataset_type=synthetic enable_checkpointing=True steps=10 attention=cudnn_flash_te scan_layers=False \ use_iota_embed=True hardware=gpu per_device_batch_size=12 ``` @@ -156,7 +156,7 @@ python3 -m MaxText.train src/MaxText/configs/base.yml run_name=gpu01 base_ou You can find the optimized running of LLama Models for various host configurations here: -https://github.com/AI-Hypercomputer/maxtext/tree/main/MaxText/configs/a3/llama_2_7b +https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/configs/a3/llama_2_7b `1vm.sh` modified script below: diff --git a/docs/guides/run_maxtext_via_multihost_job.md b/docs/guides/run_maxtext/run_maxtext_via_multihost_job.md similarity index 100% rename from docs/guides/run_maxtext_via_multihost_job.md rename to docs/guides/run_maxtext/run_maxtext_via_multihost_job.md diff --git a/docs/guides/run_maxtext_via_multihost_runner.md b/docs/guides/run_maxtext/run_maxtext_via_multihost_runner.md similarity index 100% rename from docs/guides/run_maxtext_via_multihost_runner.md rename to docs/guides/run_maxtext/run_maxtext_via_multihost_runner.md diff --git a/docs/guides/run_maxtext_via_pathways.md b/docs/guides/run_maxtext/run_maxtext_via_pathways.md similarity index 97% rename from docs/guides/run_maxtext_via_pathways.md rename to docs/guides/run_maxtext/run_maxtext_via_pathways.md index bd392a8e68..0e5ff16550 100644 --- a/docs/guides/run_maxtext_via_pathways.md +++ b/docs/guides/run_maxtext/run_maxtext_via_pathways.md @@ -15,7 +15,7 @@ --> (run-pathways)= -# Guide: Running MaxText via Pathways +# Running MaxText via Pathways This guide provides a comprehensive walkthrough for running MaxText workloads on a Google Kubernetes Engine (GKE) cluster using Pathways. Pathways acts as a powerful orchestrator for large-scale JAX jobs on AI Hypercomputer infrastructure. @@ -68,11 +68,11 @@ export RUN_NAME="maxtext-run-1" export DOCKER_IMAGE="gcr.io/${PROJECT}/${USER}_runner" ``` -## 3. Running a Batch workload +## 3. Running a batch workload A batch workload runs entirely within the GKE cluster. You submit the job definition, and Pathways manages its execution. -### Submit the Batch Workload +### Submit the batch workload Use the `xpk workload create-pathways` command to start the job. @@ -133,7 +133,7 @@ kubectl port-forward \ 29000:29000 &> /dev/null & ``` -### Step 3: Run Your MaxText script locally +### Step 3: Run your MaxText script locally With the port forward active, you can now run your MaxText script. The JAX environment variables direct it to connect to the TPUs through the tunnel. @@ -155,7 +155,7 @@ The output streams directly to your terminal, just as if you were running on a l ## Troubleshooting -* **Permission Denied errors for Cloud Storage Bucket**: Check that the service account used by your GKE nodes has "Storage Object Admin" permissions on your GCS bucket. +* **Permission denied errors for Cloud Storage bucket**: Check that the service account used by your GKE nodes has "Storage Object Admin" permissions on your GCS bucket. * **`Image not found` or `ImagePullBackOff`**: * Verify your `DOCKER_IMAGE` variable is correct. * Ensure you have successfully pushed the image to your project's Artifact Registry. diff --git a/docs/guides/run_maxtext_via_xpk.md b/docs/guides/run_maxtext/run_maxtext_via_xpk.md similarity index 86% rename from docs/guides/run_maxtext_via_xpk.md rename to docs/guides/run_maxtext/run_maxtext_via_xpk.md index a0e9a07d5a..e469d9caa2 100644 --- a/docs/guides/run_maxtext_via_xpk.md +++ b/docs/guides/run_maxtext/run_maxtext_via_xpk.md @@ -15,17 +15,17 @@ --> (run-xpk)= -# Running MaxText at Scale with XPK +# Running MaxText at scale with XPK This guide provides the recommended workflow for running MaxText on Google Kubernetes Engine (GKE) using the **Accelerated Processing Kit (XPK)**. For a complete reference on XPK, please see the [official XPK repository](https://github.com/AI-Hypercomputer/xpk). -## Overview of the Workflow +## Overview of the workflow The process involves two main stages. First, you will package the MaxText application and its dependencies into a self-contained Docker image. This is done on your local machine or any environment where Docker is installed. Second, you will use the XPK command-line tool to orchestrate the deployment of this image as a training job on a GKE cluster equipped with accelerators (TPUs or GPUs). XPK abstracts away the complexity of cluster management and job submission, handling tasks like uploading your Docker image to Artifact Registry and scheduling the workload on the cluster. -``` +```none +--------------------------+ +--------------------+ +-------------------+ | | | | | | | Your Development Machine +------> Artifact Registry +------> GKE Cluster | @@ -34,7 +34,6 @@ XPK abstracts away the complexity of cluster management and job submission, hand | 1. Build MaxText Docker | | 2. XPK uploads | | 3. XPK runs job | | Image | | image for you | | using the image| +--------------------------+ +--------------------+ +-------------------+ - ``` * * * * * @@ -43,7 +42,7 @@ XPK abstracts away the complexity of cluster management and job submission, hand Before you begin, you must have the necessary tools installed and permissions configured. -### Required Tools +### Required tools - **Python >= 3.12** with `pip` and `venv`. @@ -53,7 +52,7 @@ Before you begin, you must have the necessary tools installed and permissions co - **Docker:** Follow the [installation instructions](https://docs.docker.com/engine/install/) and complete the [post-install steps](https://docs.docker.com/engine/install/linux-postinstall/) to run Docker without `sudo`. -### GCP Permissions +### GCP permissions Your Google Cloud user account needs the following IAM roles for the project you're using: @@ -75,32 +74,26 @@ Your Google Cloud user account needs the following IAM roles for the project you * * * * * -## 2. One-Time Environment Setup +## 2. One-time environment setup These commands configure your local environment to connect to Google Cloud services. 1. **Authenticate gcloud** - ``` gcloud auth login - ``` -2. **Install GKE Auth Plugin** - +2. **Install GKE auth plugin** ``` sudo apt-get update && sudo apt-get install google-cloud-sdk-gke-gcloud-auth-plugin - ``` -3. **Configure Docker Credentials** - +3. **Configure Docker credentials** ``` gcloud auth configure-docker - ``` * * * * * @@ -109,7 +102,6 @@ These commands configure your local environment to connect to Google Cloud servi It is best practice to install XPK in a dedicated Python virtual environment. - ``` # Create a virtual environment (only needs to be done once) python3 -m venv ~/xpk_venv @@ -119,54 +111,46 @@ source ~/xpk_venv/bin/activate # Install XPK pip install xpk - ``` * * * * * -## 4. Build the MaxText Docker Image +## 4. Build the MaxText Docker image A recommended approach for running MaxText is to build your image from a **JAX AI Image**, which ensures all core libraries are version-matched and stable. -1. **Clone the MaxText Repository** - +1. **Clone the MaxText repository** ``` git clone https://github.com/google/maxtext.git cd maxtext - ``` -2. **Build the Image for your target hardware (TPU or GPU)** This script creates a local Docker image named `maxtext_base_image`. You can find a full list of available base images in the [JAX AI Images documentation](https://cloud.google.com/ai-hypercomputer/docs/images). +2. **Build the image for your target hardware (TPU or GPU)** This script creates a local Docker image named `maxtext_base_image`. You can find a full list of available base images in the [JAX AI Images documentation](https://cloud.google.com/ai-hypercomputer/docs/images). - **For TPUs:** - ``` bash docker_build_dependency_image.sh DEVICE=tpu MODE=jax_ai_image BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.5.2-rev2 - ``` - **For GPUs:** - ``` bash docker_build_dependency_image.sh DEVICE=gpu MODE=jax_ai_image BASEIMAGE=us-central1-docker.pkg.dev/deeplearning-images/jax-ai-image/gpu:jax0.5.1-cuda_dl25.02-rev1 - ``` * * * * * -## 5. Run Your First MaxText Job +## 5. Run your first MaxText job This section assumes you have an existing GKE cluster with either TPU or GPU nodes. -> **IMPORTANT NOTE** -> -> This guide focuses on submitting workloads to an existing cluster. Cluster creation and management is a separate topic. For a comprehensive guide on all `xpk` commands, including `xpk cluster create`, please refer to the **[official XPK documentation](https://github.com/AI-Hypercomputer/xpk)**. - -1. **Set Your Configuration** +```{important} +This guide focuses on submitting workloads to an existing cluster. Cluster creation and management is a separate topic. For a comprehensive guide on all `xpk` commands, including `xpk cluster create`, please refer to the **[official XPK documentation](https://github.com/AI-Hypercomputer/xpk)**. +``` +1. **Set your configuration** ``` export PROJECT_ID="your-gcp-project-id" @@ -174,28 +158,24 @@ This section assumes you have an existing GKE cluster with either TPU or GPU nod export CLUSTER_NAME="your-existing-cluster-name" export BASE_OUTPUT_DIR="gs://your-output-bucket/" export DATASET_PATH="gs://your-dataset-bucket/" - ``` 2. **Configure gcloud CLI** - ``` gcloud config set project $PROJECT_ID gcloud config set compute/zone $ZONE - ``` -### A Note on Multi-Slice and Multi-Node Runs +### A Note on multi-slice and multi-node runs The examples below run on a single TPU slice (`--num-slices=1`) or a small number of GPU nodes (`--num-nodes=2`). To scale your job to a larger, multi-host configuration, you simply increase these values. For instance, to run a job across **four TPU slices**, you would change `--num-slices=1` to `--num-slices=4`. This tells XPK to allocate four `v5litepod-256` slices and orchestrate the training job across all of them as a single workload. Similarly, for GPUs, you would increase the `--num-nodes` value. -3. **Create the Workload (Run the Job)** - - - **On your TPU Cluster:** +3. **Create the workload (run the job)** + - **On your TPU cluster:** ``` xpk workload create\ @@ -205,11 +185,9 @@ For instance, to run a job across **four TPU slices**, you would change `--num-s --tpu-type v5litepod-256\ --num-slices 1\ --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-tpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100" - ``` - - **On your GPU Cluster:** - + - **On your GPU cluster:** ``` xpk workload create\ @@ -219,14 +197,13 @@ For instance, to run a job across **four TPU slices**, you would change `--num-s --device-type h100-80gb-8\ --num-nodes 2\ --command "python3 -m MaxText.train src/MaxText/configs/base.yml run_name=${USER}-gpu-job base_output_directory=${BASE_OUTPUT_DIR} dataset_path=${DATASET_PATH} steps=100" - ``` * * * * * -## 6. Managing and Monitoring Your Job +## 6. Managing and monitoring your job -- **View Logs in Real-Time:** The easiest way to see the output of your training job is through the Google Cloud Console. +- **View logs in real-time:** The easiest way to see the output of your training job is through the Google Cloud Console. 1. Navigate to the **Kubernetes Engine** section. @@ -236,18 +213,15 @@ For instance, to run a job across **four TPU slices**, you would change `--num-s 4. Select the **Logs** tab to view the container logs. -- **List Your Jobs:** - +- **List your jobs:** ``` xpk workload list --cluster ${CLUSTER_NAME} - ``` -- **Analyze Output:** Checkpoints and other artifacts will be saved to the Google Cloud Storage bucket you specified in `BASE_OUTPUT_DIR`. - -- **Delete a Job:** +- **Analyze output:** Checkpoints and other artifacts will be saved to the Google Cloud Storage bucket you specified in `BASE_OUTPUT_DIR`. +- **Delete a job:** ``` xpk workload delete --cluster ${CLUSTER_NAME} --workload diff --git a/docs/guides/understand_logs_and_metrics.md b/docs/guides/understand_logs_and_metrics.md index 5896c2fe65..7ffc71a7fa 100644 --- a/docs/guides/understand_logs_and_metrics.md +++ b/docs/guides/understand_logs_and_metrics.md @@ -14,11 +14,11 @@ limitations under the License. --> -# Understand Logs and Metrics +# Understand logs and metrics When you run a training job, MaxText produces detailed output logs. This guide shows you how to interpret these logs to understand your configuration and monitor performance. -To start, run a simple pretraining job on a single-host TPU. For instance, we can run the following command on TPU v5p-8. The resulting log is used as an example throughout this guide. +To start, run a simple pretraining job on a single-host TPU. For instance, we can run the following command on TPU v5p-8. The resulting log is used as an example throughout this guide. ```bash python3 -m MaxText.train src/MaxText/configs/base.yml \ @@ -27,18 +27,18 @@ model_name=deepseek2-16b \ per_device_batch_size=24 max_target_length=2048 steps=10 dataset_type=synthetic enable_checkpointing=false ``` -## 1 Configuration information +## 1. Configuration information -The first section of the log details the configuration of your run. This is crucial for debugging, as it shows you exactly which parameters were used. +The first section of the log details the configuration of your run. This is crucial for debugging, as it shows you exactly which parameters were used. -MaxText builds its configuration in layers. -- It starts with the **default configuration** from a YAML file. In our example, the file is [MaxText/configs/base.yml](https://github.com/AI-Hypercomputer/maxtext/blob/ffb0d49adcc457e8cbe2872864b4034b21d43326/MaxText/configs/base.yml). +MaxText builds its configuration in layers. +- It starts with the **default configuration** from a YAML file. In our example, the file is [src/MaxText/configs/base.yml](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/configs/base.yml). - Then, it overwrites any of these values with the arguments you provide in the **command line**. ``` Updating keys from env and command line: ['run_name', 'model_name', 'enable_checkpointing', 'base_output_directory', 'per_device_batch_size', 'dataset_type', 'steps', 'max_target_length'] ``` -- It updates keys based on the **model-specific configuration** file. When you specify a model, like `deepseek2-16b`, MaxText reads the corresponding parameters from the [deepseek2-16b.yml](https://github.com/AI-Hypercomputer/maxtext/blob/ffb0d49adcc457e8cbe2872864b4034b21d43326/MaxText/configs/models/deepseek2-16b.yml) file. +- It updates keys based on the **model-specific configuration** file. When you specify a model, like `deepseek2-16b`, MaxText reads the corresponding parameters from the [deepseek2-16b.yml](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/configs/models/deepseek2-16b.yml) file. ``` Running Model: deepseek2-16b @@ -49,15 +49,15 @@ MaxText builds its configuration in layers. ... Updating keys from model: ['base_emb_dim', 'base_num_query_heads', 'base_num_kv_heads', 'base_mlp_dim', 'base_moe_mlp_dim', 'base_num_decoder_layers', 'first_num_dense_layers', 'mlp_activations', 'vocab_size', 'enable_dropout', 'logits_via_embedding', 'normalization_layer_epsilon', 'num_experts', 'num_experts_per_tok', 'shared_experts', 'routed_scaling_factor', 'routed_score_func', 'routed_bias', 'decoder_block', 'attention_type', 'q_lora_rank', 'kv_lora_rank', 'qk_nope_head_dim', 'qk_rope_head_dim', 'v_head_dim', 'rope_type', 'rope_max_timescale', 'max_position_embeddings', 'original_max_position_embeddings', 'rope_factor', 'beta_fast', 'mscale'] ``` - Note that you cannot modify a key from both model config and command line. + Note that you cannot modify a key from both model config and command line. The final, consolidated configuration is printed last. ``` # From base.yml default -Config param opt_type: adamw +Config param opt_type: adamw ... # From model config -Config param base_emb_dim: 2048 +Config param base_emb_dim: 2048 ... # From command line Config param dataset_type: synthetic @@ -114,12 +114,12 @@ enable_tensorboard=True \ profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3 \ gcs_metrics=True \ save_config_to_gcs=True \ -enable_checkpointing=True +enable_checkpointing=True ``` -## 2 Environment information +## 2. Environment information -Next, the log displays the software and hardware environment for your run. This is useful for verifying your setup and understanding how parallelism is being applied. +Next, the log displays the software and hardware environment for your run. This is useful for verifying your setup and understanding how parallelism is being applied. ``` System Information: Jax Version: 0.7.2.dev20250826 @@ -134,14 +134,14 @@ Num_devices: 4, shape (1, 1, 4, 1, 1, 1, 1, 1, 1, 1, 1, 1) - **Parallelism strategy**: The `shape` tuple `(1, 1, 4, 1, 1, 1, 1, 1, 1, 1, 1, 1)` shows how your devices are arranged for parallelism. Recall from Section 1, `Config param data_sharding: (('data', 'stage', 'fsdp', 'fsdp_transpose', 'sequence', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'),)`. This confirms that all 4 devices are being used for Fully Sharded Data Parallelism (FSDP), which is the default behavior. -## 3 Resource accounting +## 3. Resource accounting Before executing training, the program analyzes the resource requirements for your training job, specifically memory and compute (FLOPs). -### 3.1 Memory analysis +### 3.1. Memory analysis We first perform a "dry run" compilation of a training step to [analyze its memory requirement](https://github.com/AI-Hypercomputer/maxtext/blob/f82ce194c490d668b14574a072a0a630c27bbd6e/MaxText/train.py#L630-L632 -). This static analysis is performed by the XLA compiler. The log outputs [memory sizes](https://github.com/AI-Hypercomputer/maxtext/blob/f82ce194c490d668b14574a072a0a630c27bbd6e/MaxText/max_utils.py#L735-L753): +). This static analysis is performed by the XLA compiler. The log outputs [memory sizes](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/max_utils.py#L700-L718): ``` Total memory size: 100.4 GB, Output size: 44.5 GB, Temp size: 55.9 GB, Argument size: 44.5 GB, Host temp size: 0.0 GB. ``` @@ -153,12 +153,11 @@ The most important number is `Total memory size: 100.4 GB`. This is the total Hi In addition, it shows temporary memory used on the host CPU. In this case, `Host temp size: 0.0 GB`, indicating that all the significant memory allocation happens on the accelerator device. - -### 3.2 Memory snapshot +### 3.2. Memory snapshot The previous section is a forecast of memory usage for entire training step, based on static analysis of the compiled code from the XLA compiler. To see the actual memory usage, we now turn to a real-time snapshot from the JAX runtime, captured right after the training state is initialized. -To set the stage for training, we first initialize the training state, which include parameter and optimizer states. At the [beginning](https://github.com/AI-Hypercomputer/maxtext/blob/f82ce194c490d668b14574a072a0a630c27bbd6e/MaxText/train.py#L695), the log shows a real-time snapshot of the [memory statistics](https://github.com/AI-Hypercomputer/maxtext/blob/f82ce194c490d668b14574a072a0a630c27bbd6e/MaxText/max_utils.py#L708-L717) on your TPU devices. +To set the stage for training, we first initialize the training state, which include parameter and optimizer states. At the [beginning](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/train.py#L445), the log shows a real-time snapshot of the [memory statistics](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/max_utils.py#L673-L682) on your TPU devices. ``` number parameters: 15.933 billion @@ -171,31 +170,29 @@ Memstats: After params initialized: ``` This log shows that each of the four TPUs has `95.74 GB` of available High Bandwidth Memory (HBM). The initial training state is evenly distributed across devices, with each using the same amount of `44.63 GB`. +### 3.3. Model TFLOP per device -### 3.3 Model TFLOP per device - -The **model FLOPs** are the floating point operations to perform model computation. For training, the computation includes a single forward and backward pass. -- In MaxText, we estimate model FLOPs by summing operations in matrix multiplications (matmuls); see [calculate_tflops_training_per_device](https://github.com/AI-Hypercomputer/maxtext/blob/e969faabbb571285a51545530f34d8f0a9f237e9/MaxText/maxtext_utils.py#L297). +The **model FLOPs** are the floating point operations to perform model computation. For training, the computation includes a single forward and backward pass. +- In MaxText, we estimate model FLOPs by summing operations in matrix multiplications (matmuls); see [calculate_tflops_training_per_device](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/maxtext_utils.py#L454). - The number of model FLOPs is dependent on model architecture, input size (batch size, sequence length), and gradient accumulation steps. It does not include optimization operations. - We break down the FLOPs into two parts: - "Learnable weight FLOPs" are matmuls between activations and learnable weights. Specifically, this occurs in embedding, feed forward networks, attention-related projections, and unembedding. - - "Attention FLOPs" are matmuls in attention score computation like $\mathrm{softmax}{\left(\frac{QK^\top}{\sqrt{d}}\right)} V$. + - "Attention FLOPs" are matmuls in attention score computation like $\mathrm{softmax}{\left(\frac{QK^\top}{\sqrt{d}}\right)} V$. One **TFLOP** (TeraFLOP) is equal to $10^{12}$ FLOPs. The log shows the theoretical estimate of **model TFLOP per device**: ``` Per train step: - Total TFLOPs: 764.67 + Total TFLOPs: 764.67 split as 94.54% learnable weight flops and 5.46% attention flops ``` -In this example, given `model=deepseek2-16b`, `per_device_batch_size=24`, `max_target_length=2048`, and no gradient accumulation, we have $\text{model tflop per device} \approx 764.67$. -- 94.54% of the TFLOPs are attributed to learnable weight and 5.46% are attributed to attention. -- As you will see next, this number is important for calculating performance metrics, such as TFLOP/s/device and Model FLOPs Utilization (MFU). +In this example, given `model=deepseek2-16b`, `per_device_batch_size=24`, `max_target_length=2048`, and no gradient accumulation, we have $\text{model tflop per device} \approx 764.67$. +- 94.54% of the TFLOPs are attributed to learnable weight and 5.46% are attributed to attention. +- As you will see next, this number is important for calculating performace metrics, such as TFLOP/s/device and Model FLOPs Utilization (MFU). -You can find more information about model FLOPs and MFU in the [Performance Metrics](../guides/performance_metrics.md) topic. +You can find more information about model FLOPs and MFU in the [Performance Metrics](performance-metrics) topic. - -## 4 Training metrics +## 4. Training metrics Finally, we are getting to the training steps! In this section, we introduce performance metrics including TFLOP/s/device, MFU, and Tokens/s/device (throughput). We briefly cover learning metrics including loss and total weights. ``` @@ -215,8 +212,7 @@ Before we dive deep here, recall a few numbers from previous sections: - $\text{max target length} = 2048$, $\text{per device batch size} = 24$ - $\text{model tflop per device} \approx 764.67$ (rounded), $\text{number of devices} = 4$ - -### 4.1 Performance metrics +### 4.1. Performance metrics The performance metrics fluctuate at the beginning, and become stable towards the end. Therefore, we usually read them from the last step. Let's take a closer look at Step 9. ``` @@ -226,39 +222,38 @@ As shown in `seconds: 5.667`, $\text{measured step time in seconds} \approx 5.66 **TFLOP per second per device** -- It is [computed](https://github.com/AI-Hypercomputer/maxtext/blob/e969faabbb571285a51545530f34d8f0a9f237e9/MaxText/metric_logger.py#L193-L194) as +- It is [computed](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/metric_logger.py#L211-L213) as $$\text{tflop/s/device} = \frac{\text{model tflop per device}}{\text{measured step time in seconds}}$$ - Here we have `TFLOP/s/device: 134.924`. Let's try to verify manually: $764.67 / 5.667 = 134.934$. Not exactly the same but close, since the both tflop and time are rounded in the log. - Further, we can calculate **Model FLOPs Utilization (MFU)** from this: - + $$\text{MFU} = \frac{\text{tflop/s/device}}{\text{peak hardware tflop/s}}$$ - - For TPU v5p, $\text{peak hardware tflop/s}=459$. Thus, $134.924 / 459 = 29.40$%. Note this is an example for explanation with small batch size and sequence length, so the MFU is not optimal. + + For TPU v5p, $\text{peak hardware tflop/s}=459$. Thus, $134.924 / 459 = 29.40$%. Note this is an example for explaination with small batch size and sequence length, so the MFU is not optimal. **Tokens per second per device (throughput)** -- It is [computed](https://github.com/AI-Hypercomputer/maxtext/blob/e969faabbb571285a51545530f34d8f0a9f237e9/MaxText/metric_logger.py#L197-L199) as +- It is [computed](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/metric_logger.py#L215-L217) as $$\text{token/s/device} = \frac{\text{number of tokens per device}}{\text{measured step time in seconds}}$$ -- The numerator is from [calculate_tokens_training_per_device](https://github.com/AI-Hypercomputer/maxtext/blob/e969faabbb571285a51545530f34d8f0a9f237e9/MaxText/maxtext_utils.py#L151) +- The numerator is from [calculate_tokens_training_per_device](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/maxtext_utils.py#L122) $$\text{number of tokens per device} = \text{per device batch size} \times \text{max target length}$$ - Here we have `Tokens/s/device: 8672.758`. Let's try to verify manually: $24 \times 2048 / 5.667 = 8673.372$. Not exactly the same but close, since the time is rounded in the log. +### 4.2. Learning metrics -### 4.2 Learning metrics - -**Loss**. The loss is the key indicator of learning progress, which should decrease over training steps. In this example, the loss is `12.038` at Step 0 and decreases to `10.374` at Step 9. Ideally, we want the loss to converge to a small value with sufficiently large training steps. +**Loss**. The loss is the key indicator of learning progress, which should decrease over training steps. In this example, the loss is `12.038` at Step 0 and decreases to `10.374` at Step 9. Ideally, we want the loss to converge to a small value with sufficiently large training steps. -**Total weights**. When discussing the throughput, we have $\text{number of tokens} = \text{per device batch size} \times \text{max target length} \times \text{number of device}$. In this example, $\text{number of tokens} = 24 \times 2048 \times 4 = 196608$. There are two types of tokens: real tokens and pad tokens. The pad tokens are placeholders introduced by data preprocessing: We truncate or pad each sentence to max target length. Only real tokens contribute to the learning signal (i.e., loss). Therefore, we monitor $\text{number of real tokens}$, which is shown as [total weights](https://github.com/AI-Hypercomputer/maxtext/blob/f82ce194c490d668b14574a072a0a630c27bbd6e/MaxText/train.py#L307). -- Here we see `total_weights: 196608` for all steps. This is because we are using `dataset_type=synthetic`, where all sentences are generated with a length of `max_target_length=2048`. As a result, there are no pad tokens and total weights = number of tokens. +**Total weights**. When discussing the throughput, we have $\text{number of tokens} = \text{per device batch size} \times \text{max target length} \times \text{number of device}$. In this example, $\text{number of tokens} = 24 \times 2048 \times 4 = 196608$. There are two types of tokens: real tokens and pad tokens. The pad tokens are placeholders introduced by data preprocessing: We truncate or pad each sentence to max target length. Only real tokens contribute to the learning signal (i.e., loss). Therefore, we monitor $\text{number of real tokens}$, which is shown as [total weights](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/train.py#L151). +- Here we see `total_weights: 196608` for all steps. This is because we are using `dataset_type=synthetic`, where all sentences are generated with a length of `max_target_length=2048`. As a result, there are no pad tokens and total weights = number of tokens. - However, in real datasets, sentences can have variable lengths and total weights < number of tokens. For example, we can set `dataset_type=tfds dataset_path=gs://maxtext-dataset dataset_name='c4/en:3.0.1'`, and will see total weights smaller than `196608`: ``` completed step: 8, seconds: 5.670, TFLOP/s/device: 134.856, Tokens/s/device: 8668.393, total_weights: 163259, loss: 9.596 completed step: 9, seconds: 5.669, TFLOP/s/device: 134.884, Tokens/s/device: 8670.184, total_weights: 155934, loss: 9.580 ``` -- For better convergence, we want to have large total weights. Towards this end, MaxText supports [packing](https://github.com/AI-Hypercomputer/maxtext/blob/f82ce194c490d668b14574a072a0a630c27bbd6e/MaxText/sequence_packing.py#L39) multiple short sequences into one. This is enabled by default with `packing=True` in [base.yml](https://github.com/AI-Hypercomputer/maxtext/blob/eff346c028092c4f4fd421e5c5343308def5de5a/MaxText/configs/base.yml#L454). +- For better convergence, we want to have large total weights. Towards this end, MaxText supports [packing](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/sequence_packing.py#L37) multiple short sequences into one. This is enabled by default with `packing=True` in [base.yml](https://github.com/AI-Hypercomputer/maxtext/blob/fafdeaa14183a8f5ca7b9f7b7542ce1655237574/src/MaxText/configs/base.yml#L477). diff --git a/docs/guides/use_vertex_ai_tensorboard.md b/docs/guides/use_vertex_ai_tensorboard.md index 866cbfd3bd..00daa984af 100644 --- a/docs/guides/use_vertex_ai_tensorboard.md +++ b/docs/guides/use_vertex_ai_tensorboard.md @@ -28,7 +28,7 @@ You can use a single Vertex AI Tensorboard instance to track and compare metrics * Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console. * Assign [Vertex AI User IAM role](https://cloud.google.com/vertex-ai/docs/general/access-control#aiplatform.user) to the service account used by the TPU VMs. This is required to create and access the Vertex AI Tensorboard in Google Cloud console. If you are using XPK for MaxText, the necessary Vertex AI User IAM role will be automatically assigned to your node pools by XPK – no need to assign it manually. -## Upload Logs to Vertex AI Tensorboard +## Upload logs to Vertex AI Tensorboard **Scenario 1: Using XPK to run MaxText on GKE** XPK simplifies MaxText's Vertex AI Tensorboard integration. A Vertex Tensorboard instance and Experiment are automatically created by XPK during workload scheduling. Also, XPK automatically sets the necessary environment variables, eliminating the need to manually configure this in MaxText. Set `use_vertex_tensorboard=False` to avoid setting up Vertex Tensorboard again in MaxText. This is how the configuration will look like for running MaxText via XPK: diff --git a/docs/reference.md b/docs/reference.md index ada3daf799..55f9207a7c 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -19,9 +19,9 @@ ```{toctree} :maxdepth: 1 +reference/config_base.md reference/terminology.md reference/supported_models_and_architectures.md -reference/alternatives.md reference/benchmark_and_performance.md reference/architecture_overview.md reference/jax_xla_and_pallas.md diff --git a/docs/reference/architecture_overview.md b/docs/reference/architecture_overview.md index 1c9d731f6c..9166083e9d 100644 --- a/docs/reference/architecture_overview.md +++ b/docs/reference/architecture_overview.md @@ -15,7 +15,6 @@ The practical application of this principle is evident throughout the codebase, For example, the functional training step in [`train.py`](https://github.com/AI-Hypercomputer/maxtext/blob/01c7137d4e13878e38baae44dc99e588eaa50a70/src/MaxText/train.py#L193) is compiled using `jax.jit` before being executed in the main training loop. ```py - # A simplified representation of the JIT compilation in MaxText/train.py p_train_step = jax.jit( functional_train, @@ -24,7 +23,6 @@ p_train_step = jax.jit( static_argnums=static_argnums_train, donate_argnums=donate_argnums_train, ) - ``` This code snippet demonstrates how much of the complexity of optimizing a training step, including forward pass, loss calculation, and gradient computation, is handed off to the compiler. The developer works with high-level Python and JAX primitives, while the compiler manages the low-level performance details. This strategic decision to trade the fine-grained control of custom kernels for the automated optimization and hardware portability of a compiler is central to MaxText's identity. It shifts the cognitive load from the end-user (the ML engineer) to the compiler development team, making high-performance computing more accessible to a broader audience proficient in Python. @@ -33,7 +31,7 @@ This code snippet demonstrates how much of the complexity of optimizing a traini The control plane of MaxText provides a structured yet flexible interface for users to define, configure, and launch training and inference jobs. It is designed to scale with the user's needs, offering simple command-line execution for local development and sophisticated orchestration tools for production-level runs on large-scale clusters. This system is centered around a primary YAML configuration file and a tiered set of execution scripts. -### base.yml: the central configuration hub +### `base.yml`: the central configuration hub Every MaxText job is governed by the same base YAML configuration file ([`src/MaxText/configs/base.yml`](https://github.com/AI-Hypercomputer/maxtext/blob/01c7137d4e13878e38baae44dc99e588eaa50a70/src/MaxText/configs/base.yml)) with model-specific details and overrides passed through a second config (e.g. [`src/MaxText/configs/models/deepseek3-671b.yml`](https://github.com/AI-Hypercomputer/maxtext/blob/01c7137d4e13878e38baae44dc99e588eaa50a70/src/MaxText/configs/models/deepseek3-671b.yml)). Finally, experiment-specific settings are passed on the command line. The contents of these together comprise all the hyperparameters and settings that define a run: @@ -91,9 +89,9 @@ While the base model implementations are typically simple, MaxText is equipped t The modularity of this design is clearly demonstrated by third-party extensions. For instance, the NVIDIA maxtext-jaxpp fork was able to add support for pipeline parallelism by inserting jaxpp.pipeline\_enter\_stage hooks directly into the \_\_call\_\_ method of the Decoder class, a testament to the codebase's modularity and extensibility. -### Data ingestion (input\_pipeline.py) +### Data ingestion (`input_pipeline.py`) -[The data ingestion pipeline](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline.md) is a critical component for performance at scale. In MaxText, the main training loop interfaces with the data pipeline through the create\_data\_iterator function, which is called from train.py. This function acts as a facade, abstracting the specific data loading implementation from the rest of the training logic. +[The data ingestion pipeline](data-input-pipeline) is a critical component for performance at scale. In MaxText, the main training loop interfaces with the data pipeline through the create\_data\_iterator function, which is called from train.py. This function acts as a facade, abstracting the specific data loading implementation from the rest of the training logic. MaxText supports three primary data loading backends: @@ -104,7 +102,7 @@ MaxText supports three primary data loading backends: While all three are supported, MaxText recommends the use of Grain, particularly for multi-host training scenarios. The rationale stems from performance and determinism considerations, at which Grain excels. Grain uses a data format called ArrayRecord, which supports efficient random access by index. This allows for true global shuffling of data across all hosts and eliminates the performance bottleneck associated with sequential reading. -### State management and persistence (checkpointing.py) +### State management and persistence (`checkpointing.py`) MaxText's state management and persistence layer is built on [Orbax](https://orbax.readthedocs.io/en/latest/), a flexible and powerful open-source checkpointing library for JAX applications. The core logic is encapsulated within the checkpointing.py module, which provides a comprehensive suite of tools for saving and loading training state with high performance and resilience. @@ -112,7 +110,7 @@ checkpointing.py module, which provides a comprehensive suite of tools for savin The central function is create\_orbax\_checkpoint\_manager, which configures and returns an Orbax CheckpointManager instance. This manager handles the core checkpointing operations and is configured with several key features: * Asynchronous checkpointing: By setting the `async_checkpointing` flag to true, users can enable non-blocking checkpoint saves. This is a critical performance optimization. The training loop can proceed with the next step on the accelerators while the CPU on each host handles the process of serializing the previous step's state and writing it to Google Cloud Storage. This effectively hides the I/O latency of checkpointing and maximizes accelerator utilization. -* Flexible state restoration: The load\_state\_if\_possible function implements a sophisticated, prioritized logic for resuming a run. When a job starts, it first attempts to find and load a full checkpoint from the current run's output directory. If that fails, it checks if a path to a full state checkpoint from a different run has been provided via the `load_full_state_from_path` argument. If that also fails, it looks for a parameter-only checkpoint (without training/optimizer state) specified by `load_parameters_from_path`. +* Flexible state restoration: The `load_state_if_possible` function implements a sophisticated, prioritized logic for resuming a run. When a job starts, it first attempts to find and load a full checkpoint from the current run's output directory. If that fails, it checks if a path to a full state checkpoint from a different run has been provided via the `load_full_state_from_path` argument. If that also fails, it looks for a parameter-only checkpoint (without training/optimizer state) specified by `load_parameters_from_path`. * Emergency and replicated checkpointing: For maximum resilience and rapid job resumption in large-scale, production environments like GKE, the module includes support for advanced Orbax features. @@ -123,7 +121,7 @@ A fundamental aspect of the MaxText workflow is the conversion of checkpoints be * There also exist conversion scripts to convert weights to Hugging Face, e.g. `llama_mistral_mixtral_orbax_to_hf.py` -### Utilities and distributed setup (max\_utils.py) +### Utilities and distributed setup (`max_utils.py`) The `max_utils.py` module serves as a collection of common helper functions used across the MaxText codebase, but its most critical function is to abstract away the initialization of the JAX distributed environment. @@ -155,7 +153,7 @@ This logical mesh abstraction enables the implementation of the standard paralle In MaxText, these strategies are implemented by annotating the model's PyTrees (the nested Python structures of arrays that hold the parameters and state) with sharding specifications. This is done using Flax's partitioning utilities, such as nn\_partitioning. These annotations provide requirements and hints to the compiler, telling it how each tensor should be distributed across the axes of the device mesh. The compiler then generates the appropriate collective communication operations (e.g., all-reduce, all-gather) needed to execute the parallel computation correctly and efficiently. -For more information on sharding see [our sharding documentation](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/explanations/sharding.md). +For more information on sharding see [our sharding documentation](sharding). ### Hardware abstraction and performance via XLA diff --git a/docs/reference/benchmark_and_performance.md b/docs/reference/benchmark_and_performance.md index f2f8a8ab53..dcd44386d1 100644 --- a/docs/reference/benchmark_and_performance.md +++ b/docs/reference/benchmark_and_performance.md @@ -1,49 +1,41 @@ -# Benchmark and Performance Tuning -This tutorial guides you through setting up benchmarks and performing performance tuning in MaxText, **focusing on key aspects** like how to set up benchmarks, choose the right config and tuning the benchmark performance. +# Benchmark and performance tuning +This tutorial guides you through setting up benchmarks and performing performance tuning in MaxText, **focusing on key aspects** like how to set up benchmarks, choose the right config and tuning the benchmark performance. -## How to Setup Benchmark +## How to setup benchmark Setting up effective benchmarks is crucial for accurate performance tuning. Here's how to approach it: - ### Benchmark with synthetic data and repeated batches -To efficiently benchmark MaxText without large, real datasets, use synthetic data to eliminate input and pipeline bottlenecks. +To efficiently benchmark MaxText without large, real datasets, use synthetic data to eliminate input and pipeline bottlenecks. set `dataset_type` to "synthetic" and `reuse_example_batch` to 1. - ### Start with arithmetic intensity analysis Begin your benchmarking efforts by performing an arithmetic intensity analysis. This fundamental step helps you determine the ideal batch size and sharding strategy for your specific hardware and workload. Based on your analysis and hardware configuration, select the sharding approach that yields the best results. Arithmetic intensity is calculated as the ratio of floating-point operations (FLOPs) to memory(bytes) or communication(bytes). - - * **Arithmetic Intensity = FLOPs / Bytes** This metric helps determine whether a computation is MXU-bound (high arithmetic intensity) or memory-bound/communication-bound (low arithmetic intensity). [This sharding document](sharding) illustrates various sharding strategies and their roofline analysis, through AI analysis. - -## Metrics for Benchmark Analysis +## Metrics for benchmark analysis For benchmarking purposes, we collect the step time for training. This step time is then used to calculate MFU and throughputs, which provide insights into the utilization achieved for each benchmark workload. - - * **MFU = flops_train_step / step_time / peak HW FLOPS** * **Throughput = global tokens / step_time / number of devices** -More detailed are explained in [performance_metrics ](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/performance_metrics.md) +More detailed are explained in [](performance-metrics). ## Tuning benchmark performance - -### Tuning Remat Policy +### Tuning Remat policy Rematerialization (remat) is a technique used to reduce memory consumption during model training. It works by recomputing activations during the backward pass instead of storing them in memory during the forward pass. This can be particularly beneficial for large models where memory is a bottleneck. @@ -53,25 +45,23 @@ If there is enough host HBM memory, activations can be offloaded to host memory MaxText offers both preset and granular re-materialization policies, allowing you to tailor them to your specific requirements. -**Provided Remat Policies** +**Provided Remat policies** Remat policies can be chosen from: `minimal_with_context`, `minimal`, `save_dot_with_context_except_mlp`, `save_dot_except_mlpwi`, `save_dot_except_mlp`, `save_qkv_proj`, `qkv_proj_offloaded`, `custom`, `minimal_offloaded`, `save_out_proj` and `full`. These options offer a trade-off between speed (fastest to slowest) and HBM usage (highest to lowest) -`minimal_with_context` consumes the most HBM memory, while `full` signifies minimal checkpointing, with everything being rematerialized. [more explanation and latest support](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/decoders.py#L287) +`minimal_with_context` consumes the most HBM memory, while `full` signifies minimal checkpointing, with everything being rematerialized. [More explanation and latest support](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/decoders.py#L287) -**Custom Policy** +**Custom policy** To use a custom policy, set `remat_policy` to `custom` and specify the layers in the decode module as `offload`, `device`, or `remat`. -`offload`: The tensor is offloaded to the CPU host. - -`device`: The activation remains on the TPU device. +- `offload`: The tensor is offloaded to the CPU host. +- `device`: The activation remains on the TPU device. +- `Remat`: Rematerialization is performed during the backward pass. -`Remat`: Rematerialization is performed during the backward pass. - -### **Low Precision Training** +### Low precision training MaxText supports quantization via QWIX. To enable this, set `use_qwix_quantization=true`. @@ -79,22 +69,19 @@ Different quantization recipes are available, including` "int8", "fp8", "fp8_ful For v6e and earlier generation TPUs, use the "int8" recipe. For v7x and later generation TPUs, use "fp8_full". GPUs should use “fp8_gpu” for NVIDIA and "nanoo_fp8" for AMD. -[quantization doc](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/explanations/quantization.md?plain=1) - +See [](../explanations/quantization.md). -### Choose Sharding Strategy +### Choose sharding strategy Sharding is crucial for optimizing model performance. MaxText offers various sharding strategies and hybrid options, including FSDP, TP, EP, CP, and PP, which can be configured through your MaxText settings. -[Scaling & Sharding Concepts in MaxText](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/explanations/sharding.md) This document illustrates in detail how sharding works in maxtext and chooses the write sharding config for your workload. - +[This document](sharding) illustrates in detail how sharding works in maxtext and chooses the write sharding config for your workload. -### Performance Tuning on Custom Pallas call +### Performance tuning on custom Pallas call [Tune-jax](https://github.com/rdyro/tune-jax) offers a tuning tool for Pallas kernels on both GPU and TPU. Users can tune custom kernels for specific model configurations by providing input shapes and defining a tuning search space. -Usage - +Usage: ``` import functools @@ -104,21 +91,15 @@ from tune_jax import tune def my_pallas_function(...): ``` - This example code will benchmark `my_pallas_function` across all combinations of `block_q` and `block_k`, automatically handling any compilation failures. - -### Communication Overlaps and Tuning +### Communication overlaps and tuning There are two methods for asynchronous collective offloading: - 1. Offload Collectives to Sparse Core: - - This method is recommended for v7x. To enable it, set the following flags from[[link](https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/xla_flags_library.py#L70)]: - - + This method is recommended for v7x. To enable it, set the following flags from [[link](https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/xla_flags_library.py#L70)]: * `ENABLE_SPARSECORE_OFFLOADING_FOR_RS_AG_AR` * `ENABLE_SPARSECORE_OFFLOADING_FOR_REDUCE_SCATTER` @@ -127,23 +108,19 @@ There are two methods for asynchronous collective offloading: 2. Overlap Collective Using Continuation Fusion:** - - This method is recommended for v5p and v6e. To enable it, set the following flags[[link](https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/xla_flags_library.py#L39)]: - - + This method is recommended for v5p and v6e. To enable it, set the following flags [[link](https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/xla_flags_library.py#L39)]: * `CF_FOR_ALL_GATHER` * `CF_FOR_ALL_REDUCE` Those XLA can be set via `LIBTPU_INIT_ARGS` - -### Scoped Vmem Tuning +### Scoped Vmem tuning `scoped_vmem` can be tuned using `xla_tpu_scoped_vmem_limit_kib`. The hardware limitations for vmem are 64M for v5e, 128M for v6e, and 64M for v7x. -This can be set via `LIBTPU_INIT_ARGS`, range can be 0 through harward limit. For example: - +This can be set via `LIBTPU_INIT_ARGS`, range can be 0 through hardware limit. For example: ``` LIBTPU_INIT_ARGS="---xla_tpu_scoped_vmem_limit_kib=65536" +``` diff --git a/docs/reference/config_base.md b/docs/reference/config_base.md new file mode 100644 index 0000000000..6247671e83 --- /dev/null +++ b/docs/reference/config_base.md @@ -0,0 +1,6 @@ +# Base configuration file + +MaxText uses a base configuration file to set default values for various parameters. This file is located at [`src/MaxText/configs/base.yml`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/configs/base.yml) in the MaxText repository, and is shown below. + +```{literalinclude} ../../src/MaxText/configs/base.yml +``` diff --git a/docs/reference/supported_models_and_architectures.md b/docs/reference/supported_models_and_architectures.md index 7b5394da4b..038359d4d9 100644 --- a/docs/reference/supported_models_and_architectures.md +++ b/docs/reference/supported_models_and_architectures.md @@ -1,4 +1,4 @@ -# Supported Models & Architectures +# Supported models & architectures > **Purpose**: This page provides detailed, reference-style information about model families supported in MaxText. This page is a technical dictionary for quick lookup, reproducibility, and customization. @@ -17,7 +17,7 @@ MaxText is an open-source, high-performance LLM framework written in Python/JAX. * **Flexible Remat Policy**: Provides fine-grained control over memory-compute trade-offs. Users can select pre-defined policies (like 'full' or 'minimal') or set the policy to **'custom'**. -## Supported Model Families +## Supported model families > _**Note on GPU Coverage**: Support and tested configurations for NVIDIA GPUs can vary by model family. Please see the specific model guides for details._ @@ -54,7 +54,7 @@ MaxText is an open-source, high-performance LLM framework written in Python/JAX. * **Variants**: 20B, 120B * **Notes**: Local–Global interleaved attention, GQA, attention sink; YaRN-style scaling; MoE. -## Parallelism Building Blocks +## Parallelism building blocks MaxText supports a wide range of parallelism strategies for scaling training and inference across TPUs and GPUs: @@ -66,11 +66,11 @@ MaxText supports a wide range of parallelism strategies for scaling training and * **CP (Context Parallelism)**: Splits sequence tokens across devices, complementing tensor parallelism for long-context workloads. * **Hybrid Parallelism**: Allows for flexible combinations of FSDP, TP, EP, DP, PP, and CP to maximize hardware utilization based on model size and topology. -## Performance Characteristics +## Performance characteristics The following summarizes observed runtime efficiency and scaling behaviors of MaxText across different hardware and model types, based on published benchmarks and large-scale runs. -* **High MFU**: MaxText targets high Model FLOPs Utilization across scales; exact numbers vary by model, hardware and config. See [**Performance Metrics → MFU**](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/performance_metrics.md) for the definition and how we calculate it. +* **High MFU**: MaxText targets high Model FLOPs Utilization across scales; exact numbers vary by model, hardware and config. See [**Performance Metrics → MFU**](performance-metrics) for the definition and how we calculate it. * **Quantization**: MaxText supports quantization via both the AQT and Qwix libraries. Qwix is the recommended approach, providing a non-intrusive way to apply various quantization techniques, including Quantization-Aware Training (QAT) and Post-Training Quantization (PTQ). * **MoE**: The Mixture-of-Experts implementation features dropless routing with Megablox and `jax.lax.ragged_dot` kernels for enhanced performance. * **Multi-Token Prediction (MTP)**: This feature improves training efficiency on DeepSeek-style models by adding an auxiliary loss based on predicting multiple future tokens. @@ -82,7 +82,7 @@ The following summarizes observed runtime efficiency and scaling behaviors of Ma * [MaxText Repo](https://github.com/AI-Hypercomputer/maxtext) * **Model Implementation Guides & Source Code:** - * **Llama**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/tutorials/run_llama2.md) | [Llama2 and Llama3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/llama2.py) | [Llama4 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/llama4.py) + * **Llama**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end/tpu/llama2/run_llama2.md) | [Llama2 and Llama3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/llama2.py) | [Llama4 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/llama4.py) * **Gemma**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end/tpu/gemma/Run_Gemma.md) | [Gemma Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/gemma.py) | [Gemma2 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/gemma2.py) | [Gemma3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/gemma3.py) * **Mixtral**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end/tpu/mixtral/Run_Mixtral.md) | [Mixtral Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/mixtral.py) | [Mistral Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/mistral.py) * **DeepSeek**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end/tpu/deepseek/Run_DeepSeek.md) | [DeepSeek Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/deepseek.py) @@ -91,6 +91,6 @@ The following summarizes observed runtime efficiency and scaling behaviors of Ma * **Technical Explanations:** - * [Parallelism & Sharding](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/explanations/sharding.md) - * [Quantization Documentation](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/explanations/quantization.md) - * [AOT Compilation Instructions](https://github.com/AI-Hypercomputer/maxtext/blob/main/README.md#ahead-of-time-compilation-aot) + * [Parallelism & Sharding](../explanations/sharding.md) + * [Quantization Documentation](../explanations/quantization.md) + * [AOT Compilation Instructions](aot-compilation) diff --git a/docs/tutorials.md b/docs/tutorials.md index 7486699787..8b88913bd4 100644 --- a/docs/tutorials.md +++ b/docs/tutorials.md @@ -22,7 +22,7 @@ MaxText supports training and inference of various open models. Some extra helpful guides: * [Gemma](https://ai.google.dev/gemma): a family of open-weights Large Language Model (LLM) by [Google DeepMind](https://deepmind.google/), based on Gemini research and technology. You can run decode and finetuning using [these instructions](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end/tpu/gemma/Run_Gemma.md). -* [Llama2](https://llama.meta.com/llama2/): a family of open-weights Large Language Model (LLM) by Meta. You can run decode and finetuning using [these instructions](run-llama2). +* [Llama2](https://llama.meta.com/llama2/): a family of open-weights Large Language Model (LLM) by Meta. You can run decode and finetuning using [these instructions](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end/tpu/llama2/run_llama2.md). * [Mixtral](https://mistral.ai/news/mixtral-of-experts/): a family of open-weights sparse mixture-of-experts (MoE) model by Mistral AI. You can run decode and finetuning using [these instructions](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end/tpu/mixtral/Run_Mixtral.md) In addition to the getting started guides, there are always other MaxText capabilities that are being constantly being added! The full suite of end-to-end tests is in [end_to_end](https://github.com/AI-Hypercomputer/maxtext/blob/main/end_to_end). We run them with a nightly cadence. They can be a good source for understanding MaxText Alternatively you can see the continuous [unit tests](https://github.com/AI-Hypercomputer/maxtext/blob/main/.github/workflows/RunTests.yml) which are run almost continuously. @@ -41,7 +41,6 @@ You can also find other examples in the [MaxText repository](https://github.com/ tutorials/first_run.md tutorials/pretraining.md tutorials/full_finetuning.md -tutorials/run_llama2.md tutorials/grpo.md tutorials/sft.md tutorials/sft_on_multi_host.md diff --git a/docs/tutorials/first_run.md b/docs/tutorials/first_run.md index 3caf7ba5cd..85d131202d 100644 --- a/docs/tutorials/first_run.md +++ b/docs/tutorials/first_run.md @@ -15,7 +15,7 @@ --> (first-run)= -# Getting Started: First Run +# Getting started: First run This topic provides a basic introduction to get your MaxText workload up and running on single host and multihost environments using Cloud TPUs or NVIDIA GPUs. To help you get familiar with MaxText, we recommend starting with a single host first and then moving to multihost. @@ -27,7 +27,7 @@ This topic provides a basic introduction to get your MaxText workload up and run ## Local development for single host This procedure describes how to run MaxText on a single GPU or TPU host. -### Run MaxText on Cloud TPUs +### Run MaxText on cloud TPUs Local development is a convenient way to run MaxText on a single host. It doesn't scale to multiple hosts but is a good way to learn about MaxText. @@ -60,14 +60,12 @@ python3 -m MaxText.decode src/MaxText/configs/base.yml \ This command uses a model with randomly initialized weights, so the outputs are also random. To get high quality output you need pass in a checkpoint, typically via the `load_parameters_path` argument. -### Run MaxText via Notebook +### Run MaxText via notebook In the same TPU VM where you just installed all the dependencies of MaxText, You can also run training and decoding in MaxText via Notebook (for e.g., via Jupyter or Colab). -#### Decoding in MaxText via Notebook +#### Decoding in MaxText via notebook You can use [demo_decoding.ipynb](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/examples/demo_decoding.ipynb) to try out decoding on MaxText's `Llama3.1-8b` model implementation. In this notebook, we give `"I love to"` as the prompt, and the greedily sampled first output token is `" cook"`. Please remember to provide the path to your `Llama3.1-8b` checkpoint for the `load_parameters_path` argument in the config inside the notebook. You can use [to_maxtext.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/utils/ckpt_conversion/to_maxtext.py) to create a MaxText/Orbax checkpoint from a Huggingface checkpoint. - - ### Run MaxText on NVIDIA GPUs 1. Use `bash docker_build_dependency_image.sh DEVICE=gpu` to build a container with the required dependencies. 2. After installation is complete, run training with the following command on synthetic data: @@ -96,6 +94,6 @@ Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.a Google Kubernetes Engine (GKE) is the recommended way to run MaxText on multiple hosts. It provides a managed environment for deploying and scaling containerized applications, including those that require TPUs or GPUs. See [Running Maxtext with XPK](run-xpk) for details. -## Next steps: Preflight optimizations +## Next steps: preflight optimizations After you get workloads running, there are optimizations you can apply to improve performance. For more information, see [PREFLIGHT.md](https://github.com/google/maxtext/blob/main/PREFLIGHT.md). diff --git a/docs/tutorials/full_finetuning.md b/docs/tutorials/full_finetuning.md index 031f2f6f24..ff1dd3624a 100644 --- a/docs/tutorials/full_finetuning.md +++ b/docs/tutorials/full_finetuning.md @@ -15,7 +15,7 @@ --> (full-finetuning)= -# Full Finetuning LLama3-8B Model +# Example: Full finetuning LLama3-7B model MaxText can perform pre-training and full finetuning. To perform full fine tuning, you need to pass the checkpoint to the training script. @@ -34,7 +34,7 @@ The high level steps involve: - Note: You may need to change the training parameters to fit the model to the TPU or GPU shape and to obtain an optimized performance. -## MaxText Checkpoints +## MaxText checkpoints MaxText checkpoints are in their own format. You can see the format in the script for llama conversion script. @@ -54,7 +54,7 @@ The conversion scripts do not use accelerators but need large host memory to per - For large size model (e.g. 70B model), this script requires large memory VM. - The script load and save weights in a single pass. -### Sample Full Fine tuning Script +### Sample full fine tuning script Below is a sample training script for LLama2-7b. @@ -74,15 +74,15 @@ python3 -m MaxText.train \ You can find some [end to end scripts here](https://github.com/AI-Hypercomputer/maxtext/tree/main/end_to_end/tpu). These scripts can provide a reference point for various scripts. -### Maxtext Checkpoint to Hugging Face +### MaxText checkpoint to Hugging Face -Post finetuning or pre-training, maxtext also provides scripts to convert maxtext format weights back to [hugging face](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/llama_mistral_mixtral_orbax_to_hf.py). +Post finetuning or pre-training, MaxText also provides scripts to convert MaxText format weights back to [Hugging Face](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/llama_mistral_mixtral_orbax_to_hf.py). #### Dataset -Maxtext provides examples to work with [Common Crawl](https://commoncrawl.org/). The dataset is available in TFRecords format in a cloud bucket. Maxtext provides scripts to copy the dataset to a Google Cloud Storage Bucket. +MaxText provides examples to work with [Common Crawl](https://commoncrawl.org/). The dataset is available in TFRecords format in a cloud bucket. MaxText provides scripts to copy the dataset to a Google Cloud Storage Bucket. -##### Common Crawl (c4) Dataset Setup +##### Common Crawl (c4) dataset setup You need to run these steps once per project prior to any local development or cluster experiments. diff --git a/docs/tutorials/grpo.md b/docs/tutorials/grpo.md index 5a86a95f45..4a11426dda 100644 --- a/docs/tutorials/grpo.md +++ b/docs/tutorials/grpo.md @@ -14,15 +14,15 @@ limitations under the License. --> -# Try GRPO! +# Try GRPO This tutorial demonstrates step-by-step instructions for setting up the environment and then training the Llama3.1 8B-IT model on the GSM8K math reasoning benchmark using Group Relative Policy Optimization (GRPO). GRPO can enhance your model's problem-solving skills on mathematical word problems, coding problems, etc. GRPO is an RL algorithm designed to enhance the reasoning abilities of LLMs. It is a variant of Proximal Policy Optimization (PPO) that reduces memory usage by eliminating the need for a separate value function model. GRPO works by generating multiple responses for a given prompt, evaluating these responses using a reward model, and then calculating a relative advantage based on the group's performance to update the policy. -We use Tunix as the library for GRPO. +We use Tunix as the library for GRPO. And we use vLLM as the library for efficient model inference and generation. - + In this tutorial we use a single host TPUVM such as `v6e-8/v5p-8`. Let's get started! ## Setup your virtual environment @@ -39,7 +39,7 @@ venv_name="maxtext_venv" source ~/$venv_name/bin/activate ``` -## vLLM and tpu-commons Installations +## vLLM and tpu-commons installations Next, run the following bash script to get all the necessary installations inside the virtual environment. This will take few minutes. Follow along the installation logs and look out for any issues! @@ -50,7 +50,7 @@ bash ~/maxtext/src/MaxText/examples/install_tunix_vllm_requirement.sh 1. It installs `pip install keyring keyrings.google-artifactregistry-auth` which enables pip to authenticate with Google Artifact Registry automatically. 2. Next, it installs `vLLM` for Jax and TPUs from the artifact registry `https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/` -3. Then, it installs `tpu-commons` from the same artifact registry. +3. Then, it installs `tpu-commons` from the same artifact registry. `tpu_commons` is the TPU backend for vLLM. You will need both libraries to run vLLM on tpus. We use the scheduler code from vLLM, and the model runner code from `tpu_commons` diff --git a/docs/tutorials/pretraining.md b/docs/tutorials/pretraining.md index c16b6f72d4..ff59d184ac 100644 --- a/docs/tutorials/pretraining.md +++ b/docs/tutorials/pretraining.md @@ -29,7 +29,7 @@ To start with, we focus on HuggingFace datasets for convenience. We use the HuggingFace dataset [allenai/c4](https://huggingface.co/datasets/allenai/c4), a processed version of Google's C4. This dataset is organized into subsets (e.g., `en`, `es`), and each subset contains data splits (e.g., `train`, `validation`). -**Data preparation**: You don't need to download data, as the pipeline can stream data directly from the HuggingFace Hub. Alternatively, it can stream from a Cloud Storage bucket; see the [HuggingFace Pipeline](../guides/data_input_hf.md) page. +**Data preparation**: You don't need to download data, as the pipeline can stream data directly from the HuggingFace Hub. Alternatively, it can stream from a Cloud Storage bucket; see the [HuggingFace Pipeline](../guides/data_input_pipeline/data_input_hf.md) page. We can use this **command** for pretraining: @@ -80,10 +80,10 @@ eval metrics after step: 9, loss=9.420, total_weights=75264.0 ## 2. Grain pipeline -Grain is a library for reading data for training and evaluating JAX models. It is the recommended input pipeline for determinism and resilience! It supports data formats like ArrayRecord and Parquet. You can check [Grain pipeline](../guides/data_input_grain.md) for more details. +Grain is a library for reading data for training and evaluating JAX models. It is the recommended input pipeline for determinism and resilience! It supports data formats like ArrayRecord and Parquet. You can check [Grain pipeline](../guides/data_input_pipeline/data_input_grain.md) for more details. -**Data preparation**: You need to download data to a Cloud Storage bucket, and read data via Cloud Storage Fuse with [setup_gcsfuse.sh](https://github.com/AI-Hypercomputer/maxtext/blob/0baff00ac27bb7996c62057f235cc1d2f43d734e/setup_gcsfuse.sh#L18). +**Data preparation**: You need to download data to a Cloud Storage bucket, and read data via Cloud Storage Fuse with [setup_gcsfuse.sh](https://github.com/AI-Hypercomputer/maxtext/blob/0baff00ac27bb7996c62057f235cc1d2f43d734e/setup_gcsfuse.sh#L18). - For example, we can mount the bucket `gs://maxtext-dataset` on the local path `/tmp/gcsfuse` before training ```bash bash setup_gcsfuse.sh DATASET_GCS_BUCKET=maxtext-dataset MOUNT_PATH=/tmp/gcsfuse @@ -122,7 +122,7 @@ tokenizer_type=huggingface tokenizer_path=deepseek-ai/DeepSeek-V2-Lite ## 3. TFDS pipeline -The TensorFlow Datasets (TFDS) pipeline uses dataset in the TFRecord format. You can check [TFDS Pipeline](../guides/data_input_tfds.md) for more details. +The TensorFlow Datasets (TFDS) pipeline uses dataset in the TFRecord format. You can check [TFDS Pipeline](../guides/data_input_pipeline/data_input_tfds.md) for more details. **Data preparation**: You need to download data to a [Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets), and the pipeline streams data from the bucket. - To download the AllenAI C4 dataset to your bucket, you can use [download_dataset.sh](https://github.com/AI-Hypercomputer/maxtext/blob/08d9f20329ab55b9b928543fedd28ad173e1cd97/download_dataset.sh#L19): `bash download_dataset.sh ` diff --git a/docs/tutorials/sft.md b/docs/tutorials/sft.md index f4d52b0deb..3eed56763d 100644 --- a/docs/tutorials/sft.md +++ b/docs/tutorials/sft.md @@ -14,10 +14,10 @@ limitations under the License. --> -# Try SFT! +# Try SFT Supervised fine-tuning (SFT) is a process where a pre-trained large language model is fine-tuned on a labeled dataset to adapt the model to perform better on specific tasks. -This tutorial demonstrates step-by-step instructions for setting up the environment and then training the Llama3.1 8B model on the [HuggingFaceH4/ultrachat_200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) dataset using SFT. If you wish to use a different dataset, you can [update the dataset configurations](https://github.com/AI-Hypercomputer/maxtext/blob/main/MaxText/configs/sft.yml). +This tutorial demonstrates step-by-step instructions for setting up the environment and then training the Llama3.1 8B model on the [HuggingFaceH4/ultrachat_200k](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) dataset using SFT. If you wish to use a different dataset, you can [update the dataset configurations](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/configs/sft.yml). We use [Tunix](https://github.com/google/tunix), a JAX-based library designed for post-training tasks, to perform SFT. diff --git a/docs/tutorials/sft_on_multi_host.md b/docs/tutorials/sft_on_multi_host.md index f960cd3a9c..c9621d002a 100644 --- a/docs/tutorials/sft_on_multi_host.md +++ b/docs/tutorials/sft_on_multi_host.md @@ -14,25 +14,25 @@ # limitations under the License. --> -# Supervised Fine-Tuning (SFT) with Deepseek-V3 Model +# Supervised Fine-Tuning (SFT) with Deepseek-V3 model This guide provides step by step instructions to run SFT with Deepseek-V3 model on TPU v6e-256. Deepseek-V3 is a Mixture-of-Experts (MoE) language model with 671B parameters. -## 1. Build and Upload MaxText Docker Image +## 1. Build and upload MaxText Docker image This section guides you through cloning the MaxText repository, building MaxText Docker image with dependencies, and uploading the docker image to your project's Artifact Registry. -### 1.1. Clone the MaxText Repository +### 1.1. Clone the MaxText repository ```bash git clone https://github.com/google/maxtext.git cd maxtext ``` -### 1.2. Build MaxText Docker Image +### 1.2. Build MaxText Docker image ```bash bash docker_build_dependency_image.sh MODE=jax_ai_image BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest ``` This creates a local Docker image named `maxtext_base_image`. -### 1.3. Upload the Docker Image to Artifact Registry +### 1.3. Upload the Docker image to Artifact Registry ```bash # Replace `$USER_runner` with your desired image name export DOCKER_IMAGE_NAME=${USER}_runner @@ -43,10 +43,10 @@ The `docker_upload_runner.sh` script uploads your Docker image to Artifact Regis ## 2. Install XPK Install XPK by following the instructions in the [official documentation](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#installation-via-pip). -## 3. Create GKE Cluster +## 3. Create GKE cluster If you don't already have a GKE cluster with a `v6e-256` TPU slice available, create one by following the [XPK cluster creation guide](https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#cluster-create). -## 4. Environment Configuration +## 4. Environment configuration ```bash # -- Google Cloud Configuration -- export PROJECT= @@ -66,7 +66,7 @@ export HF_TOKEN= export MODEL_CHECKPOINT_PATH= ``` -## 5. Submit Workload on GKE Cluster +## 5. Submit workload on GKE cluster This section provides the command to run SFT with Deepseek-v3 model on a v6e-256 GKE cluster. ```bash xpk workload create \ @@ -77,7 +77,7 @@ xpk workload create \ --workload=${WORKLOAD_NAME} \ --tpu-type=${TPU_TYPE} \ --num-slices=${TPU_SLICE} \ ---command "python3 -m MaxText.sft.sft_trainer MaxText/configs/sft.yml run_name=$WORKLOAD_NAME base_output_directory=$OUTPUT_PATH model_name=deepseek3-671b load_parameters_path=$MODEL_CHECKPOINT_PATH hf_access_token=$HF_TOKEN tokenizer_path=deepseek-ai/DeepSeek-V3 per_device_batch_size=1 steps=$STEPS profiler=xplane megablox=False sparse_matmul=False ici_expert_parallelism=16 ici_fsdp_parallelism=16 weight_dtype=bfloat16 dtype=bfloat16 remat_policy=full decoder_layer_input=offload sa_block_q=2048 sa_block_q_dkv=2048 sa_block_q_dq=2048 opt_type=sgd attention=flash capacity_factor=1.0 max_target_length=2048" +--command "python3 -m MaxText.sft.sft_trainer src/MaxText/configs/sft.yml run_name=$WORKLOAD_NAME base_output_directory=$OUTPUT_PATH model_name=deepseek3-671b load_parameters_path=$MODEL_CHECKPOINT_PATH hf_access_token=$HF_TOKEN tokenizer_path=deepseek-ai/DeepSeek-V3 per_device_batch_size=1 steps=$STEPS profiler=xplane megablox=False sparse_matmul=False ici_expert_parallelism=16 ici_fsdp_parallelism=16 weight_dtype=bfloat16 dtype=bfloat16 remat_policy=full decoder_layer_input=offload sa_block_q=2048 sa_block_q_dkv=2048 sa_block_q_dq=2048 opt_type=sgd attention=flash capacity_factor=1.0 max_target_length=2048" ``` Once the fine-tuning is completed, you can access your model checkpoint at `${OUTPUT_PATH}/${WORKLOAD_NAME}/checkpoints/${STEPS}/model_params`. diff --git a/docs/tutorials/run_llama2.md b/end_to_end/tpu/llama2/run_llama2.md similarity index 99% rename from docs/tutorials/run_llama2.md rename to end_to_end/tpu/llama2/run_llama2.md index 9d195adab3..5161d232c5 100644 --- a/docs/tutorials/run_llama2.md +++ b/end_to_end/tpu/llama2/run_llama2.md @@ -14,7 +14,6 @@ limitations under the License. --> -(run-llama2)= # About Llama2 MaxText supports [Llama2](https://llama.meta.com/llama2) pretraining, finetuning and decoding for its 7B and 70B flavors. To get started on decoding and finetuning of Llama2, you will first need to download weights along with its tokenizer from [Meta](https://llama.meta.com/llama-downloads). diff --git a/requirements_docs.txt b/requirements_docs.txt index 92c3ee0194..786a73e681 100644 --- a/requirements_docs.txt +++ b/requirements_docs.txt @@ -4,3 +4,4 @@ myst-nb myst-parser[linkify] sphinx-book-theme sphinx-design +sphinx-copybutton