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#

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#

Multi-tier checkpointing stores checkpoints across multiple tiers of storage:

RAM (in-memory): Checkpoints are stored in each node’s RAM for the fastest access and lowest latency. This is used for frequent, local saves.
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#

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 the latest checkpoint files to nodes without local checkpoints during restoration.

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.
Multi-tier Checkpointing (Ramdisk + GCS): This approach combines the speed of ramdisk, the resilience of in-cluster replication, and the durability of GCS to offer a more robust and efficient solution. With Multi-Tier Checkpointing, the above blocking issue is resolved because the GCS save is handled at the service level, operating on the checkpoint already saved locally.

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.
Multi-Slice Workload: The training job must be a multi-slice environment, meaning it utilizes more than one node pool.
Orbax Checkpointer: The Orbax library must be used for checkpointing in your training script.
Ramdisk Mounted via Jobset: Each workload pod must have a ramdisk directory mounted by Jobset using the Multi-Tier Checkpointing CSI driver. This provides a high-speed, in-memory storage location for checkpoints.
Supported TPU types: v4, v5e, v5p, and v6e
Cluster version: GKE cluster version needs to be later than 1.32.3-gke.1170000.

Configure the GKE cluster#

To run workloads with Multi-Tier Checkpointing (MTC), use a Google Kubernetes Engine (GKE) cluster with the necessary drivers and features enabled. Follow the Google Cloud Checkpointing Documentation to configure MTC at the cluster level.

The cluster-level settings that must be enabled are:

Setting	Description
Multi-Tier Checkpointing	Enables the required cluster-level checkpointing feature (`HighScaleCheckpointing` addon).
GCS FUSE CSI driver	Installs the required GCS FUSE CSI driver (`GcsFuseCsiDriver` addon).
Ramdisk allocation	Dynamically provisioned by the MTC CSI driver on each workload pod when `--gke-mtc-enabled` and `--gke-mtc-ramdisk-dir` are passed to `gcluster job submit`.
GCS bucket access	Grants the GKE checkpointing service account IAM permissions on your Cloud Storage bucket.

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.

Estimate Total Checkpoint Size: 70,000,000,000 parameters × 12 bytes/parameter = 840,000,000,000 bytes 840,000,000,000 bytes ≈ 840 GB
Calculate Per-Host Checkpoint shard: (Total Checkpoint Size / 32 hosts) = 26.25 GB per host
Calculate Per-Host Ramdisk Size: (Per-Host Checkpoint shard) × 2 = 52.50 GB per host
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.

Cluster configuration values#

Set up environment variables:

export PROJECT_ID=<PROJECT_ID>
export CLUSTER_LOCATION=<ZONE> # example: europe-west4 (region) or us-central1-a (zone)
export BUCKET_LOCATION=<BUCKET_LOCATION>   # example: europe-west4 or us-central1 (must be a region or multi-region, not a zone)
export TPU_ZONE=<ZONE>                     # example: europe-west4-a
export CLUSTER_NAME=<CLUSTER_NAME>
export NODE_POOL_NAME=<TPU_NODE_POOL_NAME> # example: v6e-pool
export COMPUTE_TYPE=<TPU_MACHINE_TYPE>     # example: ct6e-standard-4t
export TOPOLOGY=<TOPOLOGY>             # example: 8x16 (for 32 hosts with ct6e-standard-4t) or 4x8 (8 hosts)
export GKE_VERSION=<GKE_VERSION>           # example: 1.32.4-gke.1415000 (minimum for new clusters)
export GCS_BUCKET=<GCS_BUCKET>             # example: my-checkpoint-bucket

Configure gcloud and Cloud Storage: Configure gcloud defaults and create a Cloud Storage bucket with Hierarchical Namespace (HNS) enabled. HNS provides fast atomic folder renames required for checkpoint finalization and GCS backups:
```
gcloud config set project ${PROJECT_ID?}
gcloud config set compute/zone ${TPU_ZONE?}

gcloud storage buckets create gs://${GCS_BUCKET?} \
  --location=${BUCKET_LOCATION?} \
  --hierarchical-namespace
```

Configure the cluster: Multi-Tier Checkpointing requires the HighScaleCheckpointing and GcsFuseCsiDriver addons, as well as Workload Identity Federation, to be enabled on your GKE cluster.

For an existing cluster, update the cluster workload pool and addons in two separate commands (since --workload-pool and --update-addons cannot be specified together in gcloud):

# Step 1: Enable Workload Identity Federation
gcloud container clusters update ${CLUSTER_NAME?} \
  --workload-pool=${PROJECT_ID?}.svc.id.goog \
  --location=${CLUSTER_LOCATION?}

# Step 2: Enable MTC and GCS FUSE addons
gcloud container clusters update ${CLUSTER_NAME?} \
  --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
  --location=${CLUSTER_LOCATION?}

For a new cluster, include the addons and workload pool during cluster creation:

gcloud container clusters create ${CLUSTER_NAME?} \
  --workload-pool=${PROJECT_ID?}.svc.id.goog \
  --addons=HighScaleCheckpointing,GcsFuseCsiDriver \
  --location=${CLUSTER_LOCATION?} \
  --cluster-version=${GKE_VERSION?}

Verify that HighScaleCheckpointing is enabled:
```
gcloud container clusters describe ${CLUSTER_NAME?} \
  --location=${CLUSTER_LOCATION?} \
  --format="yaml(addonsConfig.highScaleCheckpointingConfig)"
```
The output should confirm enabled: true.

Note: If HighScaleCheckpointing is not enabled, Cluster Toolkit (gcluster job submit) will reject the workload submission with: Error: Multi-Tier Checkpointing (MTC) requires the HighScaleCheckpointing addon to be enabled on the target GKE cluster.

Create the TPU node pool: If your cluster does not already have a TPU node pool, create one with the target compute type and topology:

gcloud container node-pools create ${NODE_POOL_NAME?} \
  --cluster=${CLUSTER_NAME?} \
  --location=${CLUSTER_LOCATION?} \
  --node-locations=${TPU_ZONE?} \
  --machine-type=${COMPUTE_TYPE?} \
  --tpu-topology=${TOPOLOGY?}

Authenticate kubectl: Authenticate kubectl with the cluster credentials:

gcloud container clusters get-credentials ${CLUSTER_NAME?} \
  --location=${CLUSTER_LOCATION?} \
  --project=${PROJECT_ID?}

Grant access to Cloud Storage buckets: Multi-Tier Checkpointing requires access to the Cloud Storage bucket to store checkpoints. Grant the Storage Object User (roles/storage.objectUser) IAM role on the bucket to both the GKE checkpointing service account and the workload service account:

GKE Checkpointing Service Account: Grant roles/storage.objectUser on the bucket to the GKE checkpointing daemon:

PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID?} --format="value(projectNumber)")
gcloud storage buckets add-iam-policy-binding gs://${GCS_BUCKET?} \
  --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER?}/locations/global/workloadIdentityPools/${PROJECT_ID?}.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
  --role="roles/storage.objectUser"

Workload Service Account: MaxText pods write TensorBoard metrics, configs, and training logs to OUTPUT_PATH. Grant roles/storage.objectUser to the Kubernetes service account (KSA) used by your training workload pod (e.g. default in namespace default):

WORKLOAD_NAMESPACE="default"
WORKLOAD_KSA="default"
gcloud storage buckets add-iam-policy-binding gs://${GCS_BUCKET?} \
  --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER?}/locations/global/workloadIdentityPools/${PROJECT_ID?}.svc.id.goog/subject/ns/${WORKLOAD_NAMESPACE?}/sa/${WORKLOAD_KSA?}" \
  --role="roles/storage.objectUser"

MaxText configuration#

This configuration manages a multi-tiered checkpointing system designed for both durability and rapid recovery.

Local checkpointing: 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 almost instantly, minimizing lost work without needing to download from slower persistent storage. This feature is enabled by setting enable_checkpointing, enable_multi_tier_checkpointing, local_checkpoint_directory, and a non-zero local_checkpoint_period flags.
Backup checkpointing: These are checkpoints saved periodically to persistent storage (i.e., GCS bucket). They ensure that you can recover your training state even after a complete job failure (repair of all nodepools). From the user’s perspective, all restoration is from the local ramdisk; it is the replicator service’s responsibility to make the checkpoints available in local storage in case of job restart. The interval for backup can be enabled by setting a non-zero multi_tier_checkpointing_backup_interval_minutes or multi_tier_checkpointing_backup_interval_steps flags (but not both).
Restoring from a specific backup: By default, the replicator service restores the most recent complete checkpoint, preferring in-cluster copies over GCS backups. If the latest checkpoints are not trustworthy, you can force the job to restore from a specific older GCS backup by creating a restore.version file in the job’s backup directory before starting the job. No MaxText flags are involved; see Restore from a specific backup in the Google Cloud documentation.

Flag	Description	Type	Default
`enable_checkpointing`	A master switch to enable (`True`) or disable (`False`) saving checkpoints during the training run.	`boolean`	`True`
`enable_multi_tier_checkpointing`	When set to (`True`), this flag enables the multi-tier checkpointing feature on maxtext level.	`boolean`	`False`
`local_checkpoint_directory`	The high-speed local filesystem path (i.e., ramdisk) where Multi-tier checkpoints are saved. Setting this path, along with a non-zero `local_checkpoint_period`, enables the Multi-tier Checkpointing feature.	`string`	`""`
`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`	`None`
`multi_tier_checkpointing_backup_interval_steps`	The interval, in steps, for how often a Multi-tier checkpoint is saved to backup from local ramdisks.	`integer`	`None`

Workload submission using Cluster Toolkit#

The flags below would give the user access to the ramdisk in their workload:

Flag	Description
`--gke-mtc-enabled`	Enables the Multi-Tier Checkpointing feature for the workload pods.
`--gke-mtc-ramdisk-dir`	Specifies the ramdisk mount path inside each workload pod.

Example Cluster Toolkit workload submission#

Set up environment variables:

# If you skipped the cluster setup section above, also export PROJECT_ID, CLUSTER_NAME, CLUSTER_LOCATION, COMPUTE_TYPE, TOPOLOGY, and GCS_BUCKET
export RAMDISK_DIRECTORY=<RAMDISK_DIRECTORY> # example: /tmp/ramdisk
export WORKLOAD_NAME=<RUN_NAME>
export NUM_SLICES=<NUM_SLICES> # example: 2
export LOCAL_CHECKPOINT_PERIOD=10
export STEPS=<STEPS>
export OUTPUT_PATH="gs://${GCS_BUCKET?}/checkpoints"
export MULTI_TIER_CHECKPOINTING_BACKUP_INT_MIN=<BACKUP_INTERVAL_MINUTES>
export DATA_PATH=<DATASET_PATH>    # optional: only required if dataset_type is not synthetic

Define the Docker image:

# Official release pre-training image (recommended)
export DOCKER_IMAGE="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_pre_training:latest"
# Or your custom runner image:
# export DOCKER_IMAGE="${CLUSTER_LOCATION?}-docker.pkg.dev/${PROJECT_ID?}/<REPO>/${USER?}_mtc_runner:latest"

Run the workload creation command:

gcluster job config set project ${PROJECT_ID?}
gcluster job config set cluster ${CLUSTER_NAME?}
gcluster job config set location ${CLUSTER_LOCATION?}

gcluster job submit \
  --image=${DOCKER_IMAGE?} \
  --name=${WORKLOAD_NAME?} \
  --compute-type=${COMPUTE_TYPE?} \
  --topology=${TOPOLOGY?} \
  --num-slices=${NUM_SLICES?} \
  --gke-mtc-enabled \
  --gke-mtc-ramdisk-dir=${RAMDISK_DIRECTORY?} \
  --command="python3 -m maxtext.trainers.pre_train.train run_name=${WORKLOAD_NAME?} base_output_directory=${OUTPUT_PATH?} model_name=default dataset_type=synthetic steps=${STEPS?} per_device_batch_size=6 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?} num_slices=${NUM_SLICES?}"

Deploying MTC on Pathways using Cluster Toolkit#

To run a Pathways workload with Multi-Tier Checkpointing, use Cluster Toolkit with an MTC-enabled GKE cluster. The following Cluster Toolkit flags enable MTC for the workload:

Flag	Description
`--gke-mtc-enabled`	Configures the MTC service and mounts the ramdisk on the workload pods.
`--gke-mtc-ramdisk-dir=<path>`	Specifies the ramdisk mount path. This must match the MaxText `local_checkpoint_directory` configuration.
`--pathways-colocated-python-sidecar-image=<image>`	Specifies the container image for the Colocated Python sidecar running on workers. Required when `colocated_python_checkpointing=True`.

Example Pathways MTC workload submission#

Set up environment variables:

# If you skipped the cluster setup section above, also export PROJECT_ID, CLUSTER_NAME, CLUSTER_LOCATION, COMPUTE_TYPE, TOPOLOGY, and GCS_BUCKET
export JOB_NAME=<RUN_NAME>
export NUM_SLICES=<NUM_SLICES> # example: 2
export OUTPUT_PATH="gs://${GCS_BUCKET?}/checkpoints"
export MAXTEXT_IMAGE="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_pre_training:latest"
export RAMDISK_DIRECTORY="/tmp/mtc_checkpoints"
export LOCAL_CHECKPOINT_PERIOD=10
export BACKUP_INTERVAL_MINUTES=30
export TRAINING_ARGS=<TRAINING_ARGS>
export OPTIONAL_ELASTICITY_ARGS=""

To enable elastic replica resizing, set OPTIONAL_ELASTICITY_ARGS to the appropriate MaxText elasticity flags and configure the corresponding workload settings. See Elastic training for the Cluster Toolkit recovery workflow.

Warning

Use compatible MaxText head and Colocated Python sidecar images. Pathways MTC uses Colocated Python for worker-local checkpoint operations. Both images must use compatible MaxText and Orbax revisions and exactly the same jax and jaxlib versions. Version skew can cause initialization or restore failures.

Define the MaxText command:

COMMAND="python3 -m maxtext.trainers.pre_train.train \
${TRAINING_ARGS?} \
run_name=${JOB_NAME?} \
base_output_directory=${OUTPUT_PATH?} \
num_slices=${NUM_SLICES?} \
enable_single_controller=True \
enable_multi_tier_checkpointing=True \
colocated_python_checkpointing=True \
local_checkpoint_directory=${RAMDISK_DIRECTORY?} \
local_checkpoint_period=${LOCAL_CHECKPOINT_PERIOD?} \
multi_tier_checkpointing_backup_interval_minutes=${BACKUP_INTERVAL_MINUTES?} \
${OPTIONAL_ELASTICITY_ARGS?}"

Submit the workload:

gcluster job submit \
   --project="${PROJECT_ID?}" \
   --cluster="${CLUSTER_NAME?}" \
   --location="${CLUSTER_LOCATION?}" \
   --name="${JOB_NAME?}" \
   --pathways \
   --compute-type="${COMPUTE_TYPE?}" \
   --topology="${TOPOLOGY?}" \
   --num-slices="${NUM_SLICES?}" \
   --image="${MAXTEXT_IMAGE?}" \
   --pathways-colocated-python-sidecar-image="${MAXTEXT_IMAGE?}" \
   --pathways-gcs-location="${OUTPUT_PATH?}" \
   --gke-mtc-enabled \
   --gke-mtc-ramdisk-dir="${RAMDISK_DIRECTORY?}" \
   --command="${COMMAND?}"

Multi-tier checkpointing

Contents

Multi-tier checkpointing#

Purpose and benefits#

Architecture and tiers#

Implementation details#

Comparison with other checkpointing methods#

Assumptions#

Configure the GKE cluster#

Calculating ramdisk size per host#

Example calculation#

Cluster configuration values#

MaxText configuration#

Workload submission using Cluster Toolkit#

Example Cluster Toolkit workload submission#

Deploying MTC on Pathways using Cluster Toolkit#

Example Pathways MTC workload submission#