copybara-service · Sep 3, 2025
@@ -32,4 +32,7 @@ guides/use_vertex_ai_tensorboard.md
 guides/features_and_diagnostics.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
 ```
@@ -0,0 +1,161 @@
+# 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.
+
+## Assumptions
+
+* **GKE Environment**: A **Google Kubernetes Engine (GKE)** cluster must be used. GCE infrastructure solutions like QueuedResources are not supported.
+* **Multi-Tier Checkpointing Enabled on GKE cluster level**: The Multi-Tier Checkpointing feature must be enabled and configured on your GKE cluster. This involves setting up the necessary CSI drivers and configurations as outlined in the [Google Cloud Checkpointing Documentation](https://cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing).
+* **Multi-Slice Workload**: The training job must be a [multi-slice environment](https://cloud.google.com/kubernetes-engine/docs/how-to/tpu-multislice), meaning it utilizes more than one node pool.
+* **Orbax Checkpointer**: The [Orbax library](https://orbax.readthedocs.io) must be used for checkpointing in your training script.
+* **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
+
+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.
+* `--enable-gcsfuse-csi-driver`: Installs the required GCS FUSE CSI driver.
+* `--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 
+
+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:
+**Total Checkpoint Size ≈ Number of Parameters × 12 bytes**
+
+For example, a 1 billion parameter model would require approximately **1B × 12 bytes = 12 GB** for a full checkpoint.
+
+In a distributed training environment, the checkpoint is **sharded**, or split, across all the hosts in a slice. Each host is only responsible for saving its portion of the total checkpoint. Therefore, the ramdisk on a single pod only needs to be large enough for its local shard.
+
+The formula is:
+**Required Ramdisk Size per Pod ≈ 2 * ( Total Checkpoint Size / Number of Hosts in the Slice**)
+
+It's a good practice to add a **10-15% buffer** .
+
+### Example Calculation
+
+Let's walk through an example for a large model.
+
+* **Model**: A 70 billion parameter language model.
+* **Training Slice**: A nodepool with **32 hosts**.
+
+1.  **Estimate Total Checkpoint Size**:
+    `70,000,000,000 parameters × 12 bytes/parameter = 840,000,000,000 bytes`
+    `840,000,000,000 bytes ≈ 840 GB`
+2.  **Calculate Per-Host Checkpoint shard**:
+    `(Total Checkpoint Size / 32 hosts) = 26.25 GB per host`
+
+3.  **Calculate Per-Host Ramdisk Size**:
+    `(Per-Host Checkpoint shard) * 2 = 52.50 GB per host`
+
+4.  **Add a Safety Buffer (e.g., 15%)**:
+    `(Per-Host Ramdisk Size) × 1.15 ≈ 60.3 GB`
+
+In this scenario, you should configure each pod in that slice with a ramdisk of at least **60 GB**.
+
+### Example XPK cluster creation command
+
+1.  **Set up environment variables:**
+    ```bash
+    OUTPUT_PATH=<gcs bucket output path>
+    PROJECT_ID=<project id>
+    ZONE=<your zone>
+    CLUSTER_NAME=<cluster name>
+    TPU_TYPE=<tpu-type> #example: v6e-256
+    MACHINE_TYPE=<cpu machine-type>
+    NUM_SLICES=<number of slices>
+    RAMDISK_SIZE=<ramdisk size> #example: 60000Mi
+    GKE_VERSION=<gke version> #example: 1.32.3-gke.1785000
+    ```
+2.  **Configure gcloud:**
+    ```bash
+    gcloud config set project ${PROJECT_ID}
+    gcloud config set compute/zone ${ZONE}
+    ```
+3.  **Clone the XPK repository:**
+    ```bash
+    git clone [https://github.com/AI-Hypercomputer/xpk.git](https://github.com/AI-Hypercomputer/xpk.git)
+    ```
+4.  **Run the cluster creation command:**
+    ```bash
+    python3 xpk/xpk.py cluster create \
+    --cluster ${CLUSTER_NAME} \
+    --cluster-cpu-machine-type=${MACHINE_TYPE} \
+    --num-slices=${NUM_SLICES} \
+    --tpu-type=${TPU_TYPE} \
+    --enable-mtc \
+    --enable-gcsfuse-csi-driver \
+    --mtc-ramdisk-size=${RAMDISK_SIZE} \
+    --mtc-gcs-bucket=${OUTPUT_PATH} \
+    --gke-version=${GKE_VERSION}
+    ```
+
+## 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.
+
+* **Local Emergency Checkpoints**: It saves checkpoints much more frequently to a fast, local directory on each host (i.e. a ramdisk). If a preemption or failure occurs, the job can restore from this recent local copy, minimizing lost work without needing to download from slower persistent storage. This feature is enabled by setting `enable_checkpointing`, `enable_emergency_checkpoint`, `local_checkpoint_directory` and a non-zero `local_checkpoint_period`.
+
+* **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` |
+| `enable_emergency_checkpoint` | When set to (`True`), this flag enables the two-tiered emergency checkpointing feature. | `boolean` | `False` |
+| `async_checkpointing` | When set to (`True`), this flag makes checkpoint saving asynchronous. The training step is only blocked for the minimal time needed to capture the model's state, and the actual writing to storage happens in a background thread. This is highly recommended for performance. It's enabled by default. | `boolean` | `True` |
+| `local_checkpoint_directory` | The high-speed local filesystem path(i.e. ramdisk) where **emergency checkpoints** are saved. Setting this path, along with a non-zero `local_checkpoint_period`, enables the emergency checkpointing feature. | `string` | `""` |
+| `local_checkpoint_period` | The interval, in training steps, for how often a **local checkpoint** is saved. This should be set to a much smaller value than `checkpoint_period` for frequent, low-overhead saves. | `integer` | `0` |
+| `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
+
+The flags below would give the user access to the ramdisk in their workload:
+
+| Flag | Description |
+| :--- | :--- |
+| `--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:**
+    ```bash
+    RAMDISK_DIRECTORY=<your ramdisk directory>
+    WORKLOAD_NAME=<YOUR WORKLOAD NAME>
+    TPU_TYPE=<tpu-type>
+    NUM_SLICES=<number of slices>
+    PROJECT_ID=<project-id>
+    LOCAL_CHECKPOINT_PERIOD=<>
+    CHECKPOINT_PEROID=<checkpoint_period>
+    STEPS=<steps>
+    DATA_PATH=<dataset path>
+    OUTPUT_PATH=<gcs bucket output path>
+    ```
+
+2.  **Define the Docker image:**
+    ```bash
+    DOCKER_IMAGE=gcr.io/${PROJECT_ID}/${USER}_mtc_runner:latest
+    ```
+
+3.  **Run the workload creation command:**
+    ```bash
+    python3 xpk/xpk.py workload create \
+    --cluster ${CLUSTER_NAME} \
+    --docker-image ${DOCKER_IMAGE} \
+    --workload ${WORKLOAD_NAME} \
+    --tpu-type=${TPU_TYPE} \
+    --num-slices=${NUM_SLICES} \
+    --ramdisk-directory=${RAMDISK_DIRECTORY} \
+    --mtc-enabled \
+    --command "python3 MaxText/train.py 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_emergency_checkpoint=True local_checkpoint_period=${LOCAL_CHECKPOINT_PERIOD} local_checkpoint_directory=/${RAMDISK_DIRECTORY}"
+    ```
@@ -0,0 +1,49 @@
+# GCS bucket-based checkpointing
+
+In this type of checkpointing, the checkpoints are saved directly to a Google
+Cloud Storage (GCS) bucket. This is the most straightforward and common method.
+The training process briefly pauses, for capturing the model's state and keeping
+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
+
+The system follows a specific order when deciding which checkpoint to load at startup. The first valid condition met is the one executed:
+
+1.  **Resume Current Run**: If a checkpoint already exists for the current
+    `run_name`, the system loads the latest fully-saved checkpoint. This is the
+    default behavior to ensure minimal state loss when resuming after an
+    interruption.
+2.  **Load from Specific Path**: The system checks for a user-specified path.
+    * If `load_parameters_path` is set, we load a parameter only checkpoint from that path..
+    * If `load_full_state_path` is set, we load a full state checkpoint from that path.
+    * **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
+
+| Flag | Description | Type | Default |
+| :--- | :--- | :--- | :--- |
+| `enable_checkpointing` | A master switch to enable (`True`) or disable (`False`) saving checkpoints during the training run. | `boolean` | `False` |
+| `async_checkpointing` | When set to (`True`), this flag makes checkpoint saving asynchronous. The training step is only blocked for the minimal time needed to capture the model's state, and the actual writing to storage happens in a background thread. This is highly recommended for performance. It's enabled by default. | `boolean` | `True` |
+| `checkpoint_period` | The interval, in training steps, for how often a checkpoint is saved. | `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.<br>**Note**: This feature is only compatible with training jobs that utilize a Distributed Data Parallel (DDP) strategy. | `boolean` | `False` |
+| `load_parameters_path` | Specifies a path to a checkpoint directory to load a parameter only checkpoint.<br>**Example**: `"gs://my-bucket/my-previous-run/checkpoints/items/1000"` | `string` | `""` (disabled) |
+| `load_full_state_path` | Specifies a path to a checkpoint directory to load a full checkpoint including optimizer state and step count from a specific directory.<br>**Example**: `"gs://my-bucket/my-interrupted-run/checkpoints/items/500"` | `string` | `""` (disabled) |
+| `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
+
+These settings control the underlying storage mechanism ([Orbax](https://orbax.readthedocs.io)) for performance and compatibility.
+
+| Flag | Description | Type | Default |
+| :--- | :--- | :--- | :--- |
+| `checkpoint_storage_target_data_file_size_bytes` | Sets a target file size for Orbax to chunk large arrays into smaller physical files. This can dramatically speed up loading over a network and in distributed environments. | `integer` | `2147483648` (2 GB) |
+| `checkpoint_storage_use_ocdbt` | If `True`, uses the TensorStore **OCDBT** (Optionally-Cooperative Distributed B+ Tree)) key-value store as the underlying storage format for checkpointing. | `boolean` | `True` |
+| `checkpoint_storage_use_zarr3` | If `True`, uses the Zarr v3 storage format within Orbax, which is optimized for chunked, compressed, N-dimensional arrays. | `boolean` | `True` |
+| `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.<br>**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` |