From bae5f450c7aeb1e3736b2aa53a9daa1e39ace6aa Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Mon, 10 Aug 2026 19:49:34 -0700 Subject: [PATCH] Add `multi_tier_checkpointing_backup_interval_steps` param It's similar to `multi_tier_checkpointing_backup_interval_minutes`, but uses steps instead of time as a signal to backup. These 2 params are mutually exclusive. PiperOrigin-RevId: 962511853 --- .../multi_tier_checkpointing.md | 493 +++++++++++------- src/maxtext/configs/base.yml | 17 +- src/maxtext/configs/types.py | 79 +-- src/maxtext/utils/max_utils.py | 153 +++--- tests/unit/max_utils_test.py | 64 ++- 5 files changed, 506 insertions(+), 300 deletions(-) diff --git a/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md b/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md index 3af6fb880c..9fa5348b05 100644 --- a/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md +++ b/docs/guides/checkpointing_solutions/multi_tier_checkpointing.md @@ -1,69 +1,131 @@ # 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. +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. +- **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. +- **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. +- **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. +- **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://docs.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://docs.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://docs.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://docs.cloud.google.com/tpu/docs/v4), [v5e](https://docs.cloud.google.com/tpu/docs/v5e), [v5p](https://docs.cloud.google.com/tpu/docs/v5p), and [v6e](https://docs.cloud.google.com/tpu/docs/v6e) -- **Cluster version**: Gke cluster version needs to be later than [1.32.3-gke.1170000](https://docs.cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing#existing-cluster). +- **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://docs.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://docs.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://docs.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://docs.cloud.google.com/tpu/docs/v4), + [v5e](https://docs.cloud.google.com/tpu/docs/v5e), + [v5p](https://docs.cloud.google.com/tpu/docs/v5p), and + [v6e](https://docs.cloud.google.com/tpu/docs/v6e) +- **Cluster version**: Gke cluster version needs to be later than + [1.32.3-gke.1170000](https://docs.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://docs.cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing). - -The [xpk script](https://github.com/AI-Hypercomputer/xpk/blob/main/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. | +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://docs.cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing). + +The [xpk script](https://github.com/AI-Hypercomputer/xpk/blob/main/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** +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. +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. +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)** +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** . @@ -71,190 +133,223 @@ It's a good practice to add a **10-15% buffer** . Let's walk through an example for a large model. -- **Model**: A 70 billion parameter language model. -- **Training Slice**: A nodepool with **32 hosts**. +- **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` +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` +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` +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` +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**. +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?} - ``` +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` | `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` | `0` | +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 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). + +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 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. | +| 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 src/maxtext/trainers/pre_train/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?}" - ``` +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 src/maxtext/trainers/pre_train/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?}" + ``` ## 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: +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=` | Specifies the ramdisk mount path. This must match the MaxText `local_checkpoint_directory` configuration. | -| `--pathways-colocated-python-sidecar-image=` | Specifies the Colocated Python sidecar image used for worker-local checkpoint operations. | +Flag | Description +:-------------------------------------------------- | :---------- +`--gke-mtc-enabled` | Configures the MTC service and mounts the ramdisk on the workload pods. +`--gke-mtc-ramdisk-dir=` | Specifies the ramdisk mount path. This must match the MaxText `local_checkpoint_directory` configuration. +`--pathways-colocated-python-sidecar-image=` | Specifies the Colocated Python sidecar image used for worker-local checkpoint operations. ### Example Cluster Toolkit workload submission -1. **Set up environment variables:** - - ```bash - JOB_NAME="" - COMPUTE_TYPE="" - TOPOLOGY="" - NUM_SLICES="" - OUTPUT_PATH="gs://" - MAXTEXT_IMAGE="" - COLOCATED_PYTHON_IMAGE="" - MAXTEXT_CONFIG="/deps/src/maxtext/configs/base.yml" - RAMDISK_DIRECTORY="/tmp/mtc_checkpoints" - LOCAL_CHECKPOINT_PERIOD=10 - BACKUP_INTERVAL_MINUTES=30 - TRAINING_ARGS="" - OPTIONAL_ELASTICITY_ARGS="" - ``` - - To enable elastic replica resizing, set `OPTIONAL_ELASTICITY_ARGS` to the appropriate MaxText elasticity flags and configure the corresponding Pathways elastic slice settings. See [Elastic training with Pathways](../../run_maxtext/run_maxtext_elastic_training.md) for configuration details. - - ```{warning} - **Use compatible MaxText head and Colocated Python sidecar images.** Pathways MTC uses [Colocated Python](https://docs.jax.dev/en/latest/notebooks/colocated-python.html) 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. - ``` - -2. **Define the MaxText command:** - - ```bash - COMMAND="python3 -m maxtext.trainers.pre_train.train ${MAXTEXT_CONFIG} \ - ${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}" - ``` - -3. **Submit the workload:** - - ```bash - ./gcluster job submit \ - --name="${JOB_NAME}" \ - --pathways \ - --compute-type="${COMPUTE_TYPE}" \ - --topology="${TOPOLOGY}" \ - --num-slices="${NUM_SLICES}" \ - --image="${MAXTEXT_IMAGE}" \ - --pathways-colocated-python-sidecar-image="${COLOCATED_PYTHON_IMAGE}" \ - --pathways-gcs-location="${OUTPUT_PATH}" \ - --gke-mtc-enabled \ - --gke-mtc-ramdisk-dir="${RAMDISK_DIRECTORY}" \ - --command="${COMMAND}" - ``` +1. **Set up environment variables:** + + ```bash + JOB_NAME="" + COMPUTE_TYPE="" + TOPOLOGY="" + NUM_SLICES="" + OUTPUT_PATH="gs://" + MAXTEXT_IMAGE="" + COLOCATED_PYTHON_IMAGE="" + MAXTEXT_CONFIG="/deps/src/maxtext/configs/base.yml" + RAMDISK_DIRECTORY="/tmp/mtc_checkpoints" + LOCAL_CHECKPOINT_PERIOD=10 + BACKUP_INTERVAL_MINUTES=30 + TRAINING_ARGS="" + OPTIONAL_ELASTICITY_ARGS="" + ``` + + To enable elastic replica resizing, set `OPTIONAL_ELASTICITY_ARGS` to the + appropriate MaxText elasticity flags and configure the corresponding + Pathways elastic slice settings. See + [Elastic training with Pathways](../../run_maxtext/run_maxtext_elastic_training.md) + for configuration details. + + ```{warning} + **Use compatible MaxText head and Colocated Python sidecar images.** Pathways MTC uses [Colocated Python](https://docs.jax.dev/en/latest/notebooks/colocated-python.html) 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. + ``` + +2. **Define the MaxText command:** + + ```bash + COMMAND="python3 -m maxtext.trainers.pre_train.train ${MAXTEXT_CONFIG} \ + ${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}" + ``` + +3. **Submit the workload:** + + ```bash + ./gcluster job submit \ + --name="${JOB_NAME}" \ + --pathways \ + --compute-type="${COMPUTE_TYPE}" \ + --topology="${TOPOLOGY}" \ + --num-slices="${NUM_SLICES}" \ + --image="${MAXTEXT_IMAGE}" \ + --pathways-colocated-python-sidecar-image="${COLOCATED_PYTHON_IMAGE}" \ + --pathways-gcs-location="${OUTPUT_PATH}" \ + --gke-mtc-enabled \ + --gke-mtc-ramdisk-dir="${RAMDISK_DIRECTORY}" \ + --command="${COMMAND}" + ``` diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 08d22bae24..aa523a386d 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -479,9 +479,21 @@ base_output_directory: "" # enable_multi_tier_checkpointing=true local_checkpoint_directory="/local" local_checkpoint_period=20 multi_tier_checkpointing_backup_interval_minutes=20 enable_multi_tier_checkpointing: false -# The interval to backup local checkpoints to the persistent storage(GCS bucket) in minutes. +# The interval to backup local checkpoints to the persistent storage +# (GCS bucket) in minutes. # It should be a positive number when enabling multi-tier checkpointing. -multi_tier_checkpointing_backup_interval_minutes: 0 +# Note: This parameter and `multi_tier_checkpointing_backup_interval_steps` +# are mutually exclusive. Exactly one must be specified when enabling +# multi-tier checkpointing. +multi_tier_checkpointing_backup_interval_minutes: null + +# The interval to backup local checkpoints to the persistent storage +# (GCS bucket) in steps. +# It should be a positive number when enabling multi-tier checkpointing. +# Note: This parameter and `multi_tier_checkpointing_backup_interval_minutes` +# are mutually exclusive. Exactly one must be specified when enabling +# multi-tier checkpointing. +multi_tier_checkpointing_backup_interval_steps: null # Number of identical pipelines in job, should be equal to ICI data parallelism * DCN data parallelism. # It should be a positive number when enabling multi-tier checkpointing. If set to 0, it will be set to num of slices. @@ -1355,4 +1367,3 @@ elastic_backup_kind: "snapshot" elastic_timeout_seconds: 300 elastic_max_retries: 10 elastic_min_slice_count: -1 - diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 68684cb38a..80fb31b3cc 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -19,27 +19,27 @@ import datetime import enum from enum import Enum -from jinja2 import Environment, TemplateSyntaxError import logging import math from math import prod import os from tempfile import gettempdir -import yaml from typing import Any, Literal, NewType, Optional import jax -from maxtext.common.common_types import AttentionType, DecoderBlockType, ReorderStrategy, ShardMode, CustomRule, VisionEncoderBlockType +from jinja2 import Environment, TemplateSyntaxError +from maxtext.common.common_types import AttentionType, CustomRule, DecoderBlockType, ReorderStrategy, ShardMode, VisionEncoderBlockType +from maxtext.utils import accelerator_to_spec_map +from maxtext.utils import elastic_utils from maxtext.utils import gcs_utils from maxtext.utils import max_utils -from maxtext.utils import elastic_utils -from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, HF_IDS -from maxtext.utils import accelerator_to_spec_map +from maxtext.utils.globals import HF_IDS, MAXTEXT_ASSETS_ROOT from pydantic.config import ConfigDict from pydantic.fields import Field from pydantic.functional_validators import field_validator, model_validator from pydantic.main import BaseModel from pydantic.types import NonNegativeFloat, NonNegativeInt, PositiveInt +import yaml class XProfTPUPowerTraceMode(enum.IntEnum): # pylint: disable=invalid-name @@ -399,9 +399,13 @@ class EmergencyCheckpointing(BaseModel): ) local_checkpoint_directory: PathStr = Field("", description="Local directory for emergency checkpoints.") local_checkpoint_period: NonNegativeInt = Field(0, description="Frequency (in steps) for local emergency checkpoints.") - multi_tier_checkpointing_backup_interval_minutes: NonNegativeInt = Field( - 0, - description="Interval in minutes to back up local checkpoints to persistent storage.", + multi_tier_checkpointing_backup_interval_minutes: PositiveInt | None = Field( + None, + description=("Interval in minutes to back up local checkpoints to persistent" " storage."), + ) + multi_tier_checkpointing_backup_interval_steps: PositiveInt | None = Field( + None, + description=("Interval in steps to back up local checkpoints to persistent" " storage."), ) mtc_data_parallelism: int = Field( 0, @@ -2632,8 +2636,8 @@ def infer_cp_axes(logical_axis_rules: list) -> tuple[str, ...]: returns the physical axis/axes it is mapped to. Args: - logical_axis_rules: The list of ``[logical_name, physical_axes]`` pairs - (the ``logical_axis_rules`` config field). + logical_axis_rules: The list of ``[logical_name, physical_axes]`` pairs (the + ``logical_axis_rules`` config field). Returns: A tuple of physical axis name strings that act as CP. Empty if the @@ -2652,8 +2656,8 @@ def infer_ep_axes(logical_axis_rules: list) -> tuple[str, ...]: physical axis/axes it is mapped to. Args: - logical_axis_rules: The list of ``[logical_name, physical_axes]`` pairs - (the ``logical_axis_rules`` config field). + logical_axis_rules: The list of ``[logical_name, physical_axes]`` pairs (the + ``logical_axis_rules`` config field). Returns: A tuple of physical axis name strings that act as EP. Empty if the @@ -2786,12 +2790,13 @@ class MaxTextConfig( # Derived DerivedValues, ): - """ - The main configuration object for MaxText. + """The main configuration object for MaxText. - This class aggregates all configuration options from modular `BaseModel` classes + This class aggregates all configuration options from modular `BaseModel` + classes into a single, validated object. It is populated by the `initialize` function. - Every field is explicitly defined to prevent misconfigurations (`extra='forbid'`). + Every field is explicitly defined to prevent misconfigurations + (`extra='forbid'`). """ dpo: DPO = Field( @@ -2862,9 +2867,7 @@ def _validate_use_te_comm_gemm_overlap(self): ) def validate_num_moe_emb_chunks(self): - """ - Validates that num_moe_emb_chunks is used with supported settings. - """ + """Validates that num_moe_emb_chunks is used with supported settings.""" if self.num_moe_emb_chunks > 0: if not self.use_gmm_v2 or not self.use_ring_of_experts: raise ValueError( @@ -2891,9 +2894,10 @@ def _load_mesh_config_from_yaml(rule_value: str) -> dict: @model_validator(mode="after") def set_derived_and_validate_values(self) -> "MaxTextConfig": - """ - Computes all derived values and runs all cross-field validations after initial parsing. - This logic is ported from the legacy pyconfig_deprecated.py system and adapted for Pydantic. + """Computes all derived values and runs all cross-field validations after initial parsing. + + This logic is ported from the legacy pyconfig_deprecated.py system and + adapted for Pydantic. """ # Handle primary custom mesh and rule if self.custom_mesh_and_rule is not CustomRule.DEFAULT: @@ -3470,8 +3474,23 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("`local_checkpoint_directory` must be set for multi-tier checkpointing.") if self.local_checkpoint_period <= 0: raise ValueError("`local_checkpoint_period` must be > 0 for multi-tier checkpointing.") - if self.multi_tier_checkpointing_backup_interval_minutes <= 0: - raise ValueError("`multi_tier_checkpointing_backup_interval_minutes` must be > 0.") + if (self.multi_tier_checkpointing_backup_interval_minutes is None) == ( + self.multi_tier_checkpointing_backup_interval_steps is None + ): + raise ValueError( + "Exactly one of `multi_tier_checkpointing_backup_interval_minutes`" + " or `multi_tier_checkpointing_backup_interval_steps` must be" + " specified." + ) + if ( + self.multi_tier_checkpointing_backup_interval_steps is not None + and self.multi_tier_checkpointing_backup_interval_steps < self.local_checkpoint_period + ): + raise ValueError( + "`multi_tier_checkpointing_backup_interval_steps`" + f" ({self.multi_tier_checkpointing_backup_interval_steps}) must be" + f" >= `local_checkpoint_period` ({self.local_checkpoint_period})." + ) if self.colocated_python_checkpointing and not self.enable_single_controller: raise ValueError("`colocated_python_checkpointing` is only supported with `enable_single_controller` set to True.") if self.enable_emergency_checkpoint: @@ -4144,9 +4163,7 @@ class RLConfig( TrainingLoop, DerivedValues, ): - """ - Configuration for Reinforcement Learning in MaxText. - """ + """Configuration for Reinforcement Learning in MaxText.""" num_epoch: int = Field(1, ge=1, description="Number of epochs to train for.") eval_interval: int = Field( @@ -4266,10 +4283,8 @@ def set_derived_values_and_validate(self) -> "RLConfig": self.tokenizer_path = HF_IDS[model_name] self.tokenizer_type = TokenizerType.HUGGINGFACE else: - raise ValueError( - "model_name not found in HF_IDS in maxtext/src/maxtext/utils/globals.py. \ - Please pass tokenizer_path in your command." - ) + raise ValueError("model_name not found in HF_IDS in maxtext/src/maxtext/utils/globals.py. \ + Please pass tokenizer_path in your command.") if self.optimizer_memory_host_offload: raise ValueError( diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index e6aebce059..735cb9c249 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -13,41 +13,41 @@ # limitations under the License. """Common Max Utils needed by multiple modules. -All the functions include MaxText modules, such as Pyconfig, should be moved to MaxText utils file. + +All the functions include MaxText modules, such as Pyconfig, should be moved to +MaxText utils file. """ import collections from collections.abc import Sequence +from contextlib import contextmanager import functools from functools import partial import os -import socket +from pathlib import Path import re +import socket import subprocess import time from typing import Any -from packaging.version import Version - from etils import epath import flax import jax -from pathlib import Path -from contextlib import contextmanager from jax.experimental import mesh_utils -from jax.sharding import PartitionSpec as P import jax.numpy as jnp +from jax.sharding import PartitionSpec as P +from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN +from maxtext.common.gcloud_stub import is_decoupled +from maxtext.common.gcloud_stub import StubSummaryWriter, _TENSORBOARDX_AVAILABLE, writer +from maxtext.utils import elastic_utils +from maxtext.utils import max_logging import numpy as np import orbax.checkpoint as ocp from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import initialization +from packaging.version import Version import psutil -from maxtext.utils import elastic_utils -from maxtext.common.gcloud_stub import is_decoupled -from maxtext.common.gcloud_stub import writer, _TENSORBOARDX_AVAILABLE, StubSummaryWriter -from maxtext.utils import max_logging -from maxtext.common.common_types import MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN - initialize_multi_tier_checkpointing = initialization.initialize_multi_tier_checkpointing HYBRID_RING_64X4 = "hybrid_ring_64x4" HYBRID_RING_32X8 = "hybrid_ring_32x8" @@ -56,8 +56,8 @@ def parse_libtpu_flags_to_dict(flags_str: str) -> dict: - """ - Parses a string of XLA flags into a dictionary of compilation options. + """Parses a string of XLA flags into a dictionary of compilation options. + This function is only for compilation usage. """ if not flags_str or not flags_str.strip(): @@ -144,6 +144,7 @@ def calculate_leaf_params_per_chip(arr): def _bytes_of(x): """Return the number of bytes used by a single leaf in a pytree. + Handles concrete arrays (NumPy/JAX), abstract shapes, scalars, and None. Unknown types default to 0. """ @@ -232,12 +233,15 @@ def add_text_to_summary_writer(key, value, summary_writer): def maybe_initialize_jax_distributed_system(raw_keys): - """The best recipe to initialize the Jax Distributed System has varied over time. We keep a layer of - indirection in MaxText to avoid breaking the call sites unnecessarily. + """The best recipe to initialize the Jax Distributed System has varied over time. + + We keep a layer of indirection in MaxText to avoid breaking the call sites + unnecessarily. Currently jax.distributed.initialize() fully works as expected! - For CPUs, we call jax.distributed.initialize() explicitly, with the specified arguments. + For CPUs, we call jax.distributed.initialize() explicitly, with the specified + arguments. """ # Early exit for cases where we don't need to initialize the jax distributed system. @@ -252,6 +256,7 @@ def maybe_initialize_jax_distributed_system(raw_keys): initialize_multi_tier_checkpointing( local_checkpoint_directory=raw_keys["local_checkpoint_directory"], backup_interval_minutes=raw_keys["multi_tier_checkpointing_backup_interval_minutes"], + backup_interval_steps=raw_keys["multi_tier_checkpointing_backup_interval_steps"], run_name=raw_keys["run_name"], jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"], use_colocated_python=True, @@ -298,6 +303,7 @@ def maybe_initialize_jax_distributed_system(raw_keys): initialize_multi_tier_checkpointing( local_checkpoint_directory=raw_keys["local_checkpoint_directory"], backup_interval_minutes=raw_keys["multi_tier_checkpointing_backup_interval_minutes"], + backup_interval_steps=raw_keys["multi_tier_checkpointing_backup_interval_steps"], run_name=raw_keys["run_name"], jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"], data_parallelism=raw_keys["mtc_data_parallelism"], @@ -345,7 +351,10 @@ def initialize_jax_for_gpu(raw_keys): def initialize_jax_for_cpu(raw_keys): - """Jax distributed initialize for CPUs. Includes retries until the coordinator is ready.""" + """Jax distributed initialize for CPUs. + + Includes retries until the coordinator is ready. + """ coordinator_ip_address = get_coordinator_ip_address() coordinator_address = coordinator_ip_address + ":1234" # JAX coordinator port used in XPK # Env variables to be set in XPK or otherwise @@ -365,9 +374,10 @@ def initialize_jax_for_cpu(raw_keys): def initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys): """Initialize JAX distributed runtime for TPUs when emergency checkpointing is used. - The information required to initialize JAX distributed runtime will be written by GKE to - the local checkpoint directory. This function retrieves that information and initializes - JAX distributed runtime. + + The information required to initialize JAX distributed runtime will be written + by GKE to the local checkpoint directory. This function retrieves that + information and initializes JAX distributed runtime. """ process_id, coordinator_address = _retrieve_jax_init_info(raw_keys) @@ -459,9 +469,7 @@ def get_coordinator_ip_address(): def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_type): """Evaluates unspecified DCN/ICI parallelism values""" if -1 in parallelism_vals: - assert ( - parallelism_vals.count(-1) == 1 - ), f"Found unspecified values (-1) for more than one {parallelism_type}\ + assert parallelism_vals.count(-1) == 1, f"Found unspecified values (-1) for more than one {parallelism_type}\ parallelism axis. At most one axis can be unspecified." determined_val = target_product / np.prod(parallelism_vals) * -1 @@ -606,8 +614,8 @@ def unbox_logicallypartioned(boxed_pytree): """Unboxes the flax.LogicallyPartitioned pieces Args: - boxed_pytree: a pytree that includes LogicallyPartitioned - leaves. + boxed_pytree: a pytree that includes LogicallyPartitioned leaves. + Returns: a pytree where all all LogicallyPartitioned leaves have been unboxed. """ @@ -627,7 +635,9 @@ def cross_entropy_with_logits( z_loss: float = 0.0, ) -> tuple[jnp.ndarray, jnp.ndarray]: """Computes cross entropy loss with stable custom gradient. + Computes a stabilized-gradient version of: + -jnp.sum(targets * nn.log_softmax(logits), axis=-1) If z_loss > 0, then an auxiliary loss equal to z_loss*log(z)^2 will be added to the cross entropy loss (z = softmax normalization constant). @@ -640,6 +650,7 @@ def cross_entropy_with_logits( targets: categorical one-hot targets [batch, length, num_classes] float array. z_loss: coefficient for auxiliary z-loss loss term. + Returns: tuple with the total loss and the z_loss, both float arrays with shape [batch, length]. @@ -826,7 +837,9 @@ def bytes_to_gb(num_bytes): def print_system_information(): """Print system information of the current environment. - Note that this will initialize the JAX backend.""" + + Note that this will initialize the JAX backend. + """ max_logging.log(f"System Information: Jax Version: {jax.__version__}") max_logging.log(f"System Information: Jaxlib Version: {jax.lib.__version__}") max_logging.log(f"System Information: Jax Backend: {jax.extend.backend.get_backend().platform_version}") @@ -843,9 +856,7 @@ def permute_to_match_maxtext_rope(arr): def unpermute_from_match_maxtext_rope(arr, model_size): - """ - Function to get the RoPE values in correct ordering - """ + """Function to get the RoPE values in correct ordering""" if model_size[:8] != "llama3.1": return arr evens = arr[..., ::2] @@ -855,10 +866,10 @@ def unpermute_from_match_maxtext_rope(arr, model_size): @partial(jax.jit, static_argnames=("cp_size", "seq_dim", "to_contiguous")) def reorder_sequence(tensor, cp_size: int, seq_dim: int = 1, to_contiguous: bool = False): - """Reorders the sequence of the tensor. For example, with cp_size=2, - [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 6, 7, 2, 3, 4, 5] - and backward - [0, 1, 6, 7, 2, 3, 4, 5] -> [0, 1, 2, 3, 4, 5, 6, 7] + """Reorders the sequence of the tensor. + + For example, with cp_size=2, [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 6, 7, 2, 3, 4, + 5] and backward [0, 1, 6, 7, 2, 3, 4, 5] -> [0, 1, 2, 3, 4, 5, 6, 7] """ if tensor is None: @@ -919,33 +930,45 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu """Reorders the example batch sequences using a hardware-appropriate backend. On GPU (hardware="gpu" or "gpu_multiprocess"), uses Transformer Engine's - reorder_causal_load_balancing which supports both DUAL_CHUNK_SWAP and STRIPED strategies. - On TPU/CPU, falls back to the pure-JAX reorder_sequence (DUAL_CHUNK_SWAP only). + reorder_causal_load_balancing which supports both DUAL_CHUNK_SWAP and STRIPED + strategies. + On TPU/CPU, falls back to the pure-JAX reorder_sequence (DUAL_CHUNK_SWAP + only). Args: batch: The batch to reorder. cp_size: The size of the compute parallelism. - reorder_strategy: The ReorderStrategy enum value (DUAL_CHUNK_SWAP or STRIPED). - hardware: The hardware type string ("tpu", "gpu", "gpu_multiprocess", "cpu"). + reorder_strategy: The ReorderStrategy enum value (DUAL_CHUNK_SWAP or + STRIPED). + hardware: The hardware type string ("tpu", "gpu", "gpu_multiprocess", + "cpu"). Returns: The reordered batch. Reorder Strategy: - - DUAL_CHUNK_SWAP: This strategy splits each query into two chunks and do the mirror swap between - GPUs. This is currently used for non-THD load balance. It requires the max_seqlens be the + - DUAL_CHUNK_SWAP: This strategy splits each query into two chunks and do the + mirror swap between + GPUs. This is currently used for non-THD load balance. It requires the + max_seqlens be the multiple of 2 * cp_size. Examples: - - Before reorder: GPU0: [0, 1, 2, 3]; GPU1: [4, 5, 6, 7]; GPU2: [8, 9, 10, 11]; GPU3: [12, 13, 14, 15]; - - After reorder: GPU0: [0, 1, 14, 15]; GPU1: [4, 5, 10, 11]; GPU2: [8, 9, 6, 7]; GPU3: [12, 13, 2, 3] + - Before reorder: GPU0: [0, 1, 2, 3]; GPU1: [4, 5, 6, 7]; GPU2: [8, 9, 10, + 11]; GPU3: [12, 13, 14, 15]; + - After reorder: GPU0: [0, 1, 14, 15]; GPU1: [4, 5, 10, 11]; GPU2: [8, 9, 6, + 7]; GPU3: [12, 13, 2, 3] - - STRIPED: This strategy distributes the tokens in a striped (interleaved) manner across + - STRIPED: This strategy distributes the tokens in a striped (interleaved) + manner across the sequence. This is currently used for THD load balance. Example: Consider 4 GPUs with seqlens=16. - - Before reorder: GPU0: [0, 1, 2, 3]; GPU1: [4, 5, 6, 7]; ...; GPU3: [12, 13, 14, 15] - - After reorder: GPU0: [0, 4, 8, 12]; GPU1: [1, 5, 9, 13]; ...; GPU3: [3, 7, 11, 15] + - Before reorder: GPU0: [0, 1, 2, 3]; GPU1: [4, 5, 6, 7]; ...; GPU3: [12, + 13, 14, 15] + - After reorder: GPU0: [0, 4, 8, 12]; GPU1: [1, 5, 9, 13]; ...; GPU3: [3, 7, + 11, 15] - See: https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/attention.py + See: + https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/attention.py """ # pylint: disable=import-outside-toplevel from maxtext.common.common_types import ReorderStrategy @@ -997,12 +1020,12 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu @staticmethod def reorder_mask_load_balancing(tensor, cp_size: int, seq_dim: int): - """ - Reorders a tensor for load balancing the compute of causal attention. - This function works on numpy arrays instead of jax.numpy arrays. - This is needed because we need the mask to be statically computable. - So, we need to redefine the same logic as reorder_causal_load_balancing. - We are still doing [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 6, 7, 2, 3, 4, 5] + """Reorders a tensor for load balancing the compute of causal attention. + + This function works on numpy arrays instead of jax.numpy arrays. This is + needed because we need the mask to be statically computable. So, we need to + redefine the same logic as reorder_causal_load_balancing. We are still doing + [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 6, 7, 2, 3, 4, 5] Args: tensor: The tensor to reorder. @@ -1062,15 +1085,13 @@ def parse_custom_args(argv): def unscan_train_state_params(params, sharding, mesh, scan_axis, layer_groups): - """ - Unrolls scanned parameter groups into per-layer entries. + """Unrolls scanned parameter groups into per-layer entries. Args: train_state: training state with scanned `params` mesh: the mesh to use for sharding output scan_axis: axis along which scanning was applied (usually 0) - layer_groups: list of tuples like: - [("dense_layers", 4), ("moe_layers", 12)] + layer_groups: list of tuples like: [("dense_layers", 4), ("moe_layers", 12)] """ params_copy = params.unfreeze() if hasattr(params, "unfreeze") else params decoder = params_copy["params"]["decoder"] @@ -1110,8 +1131,7 @@ def strip_scan_axis(pspec: P) -> P: def rescan_train_state_params(params, source_shardings, scan_axis, layer_groups): - """ - Reconstruct scanned layers from per-layer entries using minimal HBM. + """Reconstruct scanned layers from per-layer entries using minimal HBM. Args: train_state: training state with unrolled {layer_name}_{i} entries @@ -1145,13 +1165,12 @@ def stack_layers(*layers): def get_batch_seq_len_for_mode(config, model_mode): - """ - Resolves the batch size and sequence length based on the model's operational mode. + """Resolves the batch size and sequence length based on the model's operational mode. Args: config: A configuration object with model parameters. - model_mode: The current operational mode - (e.g., PREFILL, AUTOREGRESSIVE, TRAIN). + model_mode: The current operational mode (e.g., PREFILL, AUTOREGRESSIVE, + TRAIN). Returns: A tuple of (batch_size, seq_len). @@ -1187,7 +1206,9 @@ def print_non_trivial_mesh_axis(mesh): def bootstrap_transformer_engine_cgemm(config): """Potentially initialize NCCL communicators for Collective GEMM operations if - the environment is distributed and has the appropriate config.""" + + the environment is distributed and has the appropriate config. + """ import transformer_engine.jax.cpp_extensions as tex # pylint: disable=import-outside-toplevel # pytype: disable=import-error tsp_size = config.ici_tensor_sequence_parallelism * config.dcn_tensor_sequence_parallelism @@ -1221,7 +1242,9 @@ def dummy_context_manager(): @contextmanager def transformer_engine_context(): """If TransformerEngine is available, this context manager will provide - the library with MaxText-specific details needed for correcct operation.""" + + the library with MaxText-specific details needed for correcct operation. + """ try: from transformer_engine.jax.sharding import global_shard_guard, MeshResource # pylint: disable=import-outside-toplevel # Inform TransformerEngine of MaxText's physical mesh resources. diff --git a/tests/unit/max_utils_test.py b/tests/unit/max_utils_test.py index 097d85d491..04da4c6428 100644 --- a/tests/unit/max_utils_test.py +++ b/tests/unit/max_utils_test.py @@ -13,8 +13,19 @@ # limitations under the License. """Tests for the common Max Utils""" + import os import sys + +try: + from maxtext.utils import max_utils as max_utils_module + from maxtext.utils import max_logging as max_logging_module + + sys.modules["maxtext.utils"] = max_utils_module + sys.modules["maxtext.utils.max_utils"] = max_utils_module + sys.modules["maxtext.utils.max_logging"] = max_logging_module +except ImportError: + pass import time import unittest from unittest import mock @@ -289,6 +300,7 @@ def test_initialize_jax_for_gpu_invalid_devices(self, _mock_log, _mock_devices, @mock.patch("maxtext.utils.max_logging.log") def test_initialize_jax_for_gpu_no_devices(self, _mock_log, _mock_devices, mock_init, mock_config_update): """When coordinator env is set but neither CUDA_VISIBLE_DEVICES nor SLURM_STEP_GPUS is set, JAX uses all devices + (config) and init gets no local ids. """ raw_keys = {"jax_distributed_initialization_timeout": 300} @@ -399,6 +411,7 @@ def _base_keys(self, **overrides): "enable_multi_tier_checkpointing": False, "local_checkpoint_directory": "/tmp/ckpt", "multi_tier_checkpointing_backup_interval_minutes": 5, + "multi_tier_checkpointing_backup_interval_steps": None, "run_name": "test_run", "mtc_data_parallelism": 1, "num_slices": 2, @@ -470,6 +483,7 @@ def test_tpu_multi_tier_checkpointing(self, mock_mtc): mock_mtc.assert_called_once_with( local_checkpoint_directory=self._base_keys()["local_checkpoint_directory"], backup_interval_minutes=self._base_keys()["multi_tier_checkpointing_backup_interval_minutes"], + backup_interval_steps=self._base_keys()["multi_tier_checkpointing_backup_interval_steps"], run_name=self._base_keys()["run_name"], jax_initialization_timeout_seconds=self._base_keys()["jax_distributed_initialization_timeout"], data_parallelism=self._base_keys()["mtc_data_parallelism"], @@ -485,6 +499,7 @@ def test_single_controller_multi_tier_checkpointing_uses_colocated_python(self, mock_mtc.assert_called_once_with( local_checkpoint_directory=self._base_keys()["local_checkpoint_directory"], backup_interval_minutes=self._base_keys()["multi_tier_checkpointing_backup_interval_minutes"], + backup_interval_steps=self._base_keys()["multi_tier_checkpointing_backup_interval_steps"], run_name=self._base_keys()["run_name"], jax_initialization_timeout_seconds=self._base_keys()["jax_distributed_initialization_timeout"], data_parallelism=self._base_keys()["mtc_data_parallelism"], @@ -522,6 +537,7 @@ def test_single_controller_multi_tier_checkpointing_uses_elastic_utils_kwargs( mock_mtc.assert_called_once_with( local_checkpoint_directory=self._base_keys()["local_checkpoint_directory"], backup_interval_minutes=self._base_keys()["multi_tier_checkpointing_backup_interval_minutes"], + backup_interval_steps=self._base_keys()["multi_tier_checkpointing_backup_interval_steps"], run_name=self._base_keys()["run_name"], jax_initialization_timeout_seconds=self._base_keys()["jax_distributed_initialization_timeout"], data_parallelism=1, @@ -530,6 +546,47 @@ def test_single_controller_multi_tier_checkpointing_uses_elastic_utils_kwargs( devices=active_devices, ) + @mock.patch("maxtext.utils.max_utils.elastic_utils.single_controller_mtc_init_kwargs") + @mock.patch("maxtext.utils.max_utils.initialize_multi_tier_checkpointing") + @mock.patch("jax.distributed.initialize") + def test_single_controller_multi_tier_checkpointing_with_steps_override( + self, mock_init, mock_mtc, mock_mtc_init_kwargs + ): + active_devices = ( + mock.Mock(slice_index=0), + mock.Mock(slice_index=0), + ) + mock_mtc_init_kwargs.return_value = { + "data_parallelism": 1, + "num_slices": 1, + "devices": active_devices, + } + raw_keys = self._base_keys( + enable_single_controller=True, + enable_multi_tier_checkpointing=True, + elastic_enabled=True, + mtc_data_parallelism=0, + num_slices=2, + multi_tier_checkpointing_backup_interval_minutes=None, + multi_tier_checkpointing_backup_interval_steps=100, + ) + + max_utils.maybe_initialize_jax_distributed_system(raw_keys) + + mock_init.assert_not_called() + mock_mtc_init_kwargs.assert_called_once_with(raw_keys) + mock_mtc.assert_called_once_with( + local_checkpoint_directory=raw_keys["local_checkpoint_directory"], + backup_interval_minutes=None, + backup_interval_steps=100, + run_name=raw_keys["run_name"], + jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"], + data_parallelism=1, + num_slices=1, + use_colocated_python=True, + devices=active_devices, + ) + @mock.patch("jax.distributed.initialize") def test_tpu_checkpointing_no_emergency_calls_jax_init(self, mock_init): raw_keys = self._base_keys(enable_checkpointing=True, compile_topology_num_slices=-1) @@ -657,4 +714,9 @@ def test_reorder_roundtrip(self): if __name__ == "__main__": - unittest.main() + try: + from absl.testing import absltest + + absltest.main() + except ImportError: + unittest.main()