From 896fc1a70bb88dda7fe9b8b61e099501affb1a1c Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Wed, 3 Sep 2025 01:15:34 -0700 Subject: [PATCH] Add a guide for MaxText checkpointing solutions. PiperOrigin-RevId: 802448405 --- docs/guides.md | 3 + .../emergency_checkpointing.md | 161 +++++++++++++++ .../gcs_checkpointing.md | 49 +++++ .../multi_tier_checkpointing.md | 187 ++++++++++++++++++ 4 files changed, 400 insertions(+) create mode 100644 docs/guides/checkpointing_solutions/emergency_checkpointing.md create mode 100644 docs/guides/checkpointing_solutions/gcs_checkpointing.md create mode 100644 docs/guides/checkpointing_solutions/multi_tier_checkpointing.md diff --git a/docs/guides.md b/docs/guides.md index 8df9678e14..8d61cbfbf0 100644 --- a/docs/guides.md +++ b/docs/guides.md @@ -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 ``` diff --git a/docs/guides/checkpointing_solutions/emergency_checkpointing.md b/docs/guides/checkpointing_solutions/emergency_checkpointing.md new file mode 100644 index 0000000000..0064a150e5 --- /dev/null +++ b/docs/guides/checkpointing_solutions/emergency_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= + PROJECT_ID= + ZONE= + CLUSTER_NAME= + TPU_TYPE= #example: v6e-256 + MACHINE_TYPE= + NUM_SLICES= + RAMDISK_SIZE= #example: 60000Mi + 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= + WORKLOAD_NAME= + TPU_TYPE= + NUM_SLICES= + PROJECT_ID= + LOCAL_CHECKPOINT_PERIOD=<> + CHECKPOINT_PEROID= + STEPS= + DATA_PATH= + 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}" + ``` diff --git a/docs/guides/checkpointing_solutions/gcs_checkpointing.md b/docs/guides/checkpointing_solutions/gcs_checkpointing.md new file mode 100644 index 0000000000..14f92397e6 --- /dev/null +++ b/docs/guides/checkpointing_solutions/gcs_checkpointing.md @@ -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.
**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.
**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.
**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.
**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 diff --git a/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md b/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md new file mode 100644 index 0000000000..b72242130b --- /dev/null +++ b/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md @@ -0,0 +1,187 @@ +# 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 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, 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](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 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 + +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). + +The [xpk script](https://github.com/AI-Hypercomputer/xpk/blob/develop/xpk.py) provides a streamlined way to create a GKE cluster with all the required MTC settings. The key flags used are: + +| Flag | Description | +| :--- | :--- | +| `--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. | + + +### 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= + PROJECT_ID= + ZONE= + CLUSTER_NAME= + TPU_TYPE= #example: v6e-256 + MACHINE_TYPE= + NUM_SLICES= + RAMDISK_SIZE= #example: 60000Mi + 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 + +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 User's perspective all restoration is from local ramdisk, its replicator service responsibility to make the checkpointing available to 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` flags. + +| Flag | Description | Type | Default | +| :--- | :--- | :--- | :--- | +| `enable_checkpointing` | A master switch to enable (`True`) or disable (`False`) saving checkpoints during the training run. | `boolean` | `False` | +| `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` | `0` | + +### 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= + WORKLOAD_NAME= + TPU_TYPE= + NUM_SLICES= + PROJECT_ID= + LOCAL_CHECKPOINT_PERIOD=<> + CHECKPOINT_PEROID= + STEPS= + DATA_PATH= + OUTPUT_PATH= + MULTI_TIER_CHECKPOINTING_BACKUP_INT_MIN= + ``` + +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_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