Cloud TPU Autocheckpoint [Public Preview]

Overview

Historically, when a TPU VM requires maintenance, the procedure is initiated immediately, without leaving time for users to perform progress-preserving actions such as saving a checkpoint. This is shown in Figure 1(a).

Fig. 1. Illustration of the Autocheckpoint feature: (a) Without Autocheckpoint, the training progress from the last checkpoint is lost when there is an upcoming maintenance event. (b) With Autocheckpoint, the training progress since the last checkpoint can be preserved when there is an upcoming maintenance event.

You can use Autocheckpoint (Figure 1(b)) to preserve training progress by configuring your code to save a non-scheduled checkpoint when a maintenance event occurs. When a maintenance event occurs, progress since the last checkpoint is automatically saved. The feature works on both single slices and Multislice.

The Autocheckpoint feature works with frameworks that can capture SIGTERM and subsequently save a checkpoint. The supported frameworks include MaxText, Pax, and JAX with Orbax. Support for additional frameworks will be announced as they become available.

Only TPUs (v2-v4, and v5e) created through the Cloud TPU API can use this feature for now. Support for TPUs in GKE will be announced when it becomes available.

Using Autocheckpoint

Autocheckpoint functionality is disabled by default. When you create a TPU or a queued resource, you can enable it by adding the --autocheckpoint-enabled flag when provisioning the TPU. With the feature enabled, Cloud TPU performs the following steps once it receives notification of a maintenance event:

1. Capture SIGTERM sent to the process using the TPU device,
2. Waits until the process exits, or 5 minutes have elapsed, whichever comes first, and performs maintenance on the impacted slices.

Note that the infrastructure used by Autocheckpoint is ML framework-independent. Any ML framework can support Autocheckpoint provided it can capture the SIGTERM signal and initiate a checkpointing process.

In the application code, you need to enable the Autocheckpoint capabilities provided by the ML framework. In Pax, for example, this means enabling command-line flags when launching the training (see the autocheckpoint Quickstart with Pax). Behind the scenes, the frameworks save a non-scheduled checkpoint when a SIGTERM is received and the impacted TPU VM goes through maintenance when the TPU is no longer in use.

Quickstart: Autocheckpoint with MaxText

MaxText is a "high performance, arbitrarily scalable, open source, well-tested LLM written in pure Python/JAX targeting Cloud TPUs". MaxText contains all the necessary setup to use the Autocheckpoint feature.

The MaxText README describes two ways to run MaxText at scale:

• Using multihost_runner.py, recommended for experimentation
• Using multihost_job.job, recommended for production

When using multihost_runner.py, the only change required is to set the autocheckpoint-enabled flag when provisioning the queued resource. When using multihost_job.py, the only change required is to specify the ENABLE_AUTOCHECKPOINT=true command line flag when launching the job.

Quickstart: Autocheckpoint with Pax on single slices

In this section, we provide an example of how to set up and use Autocheckpoint with Pax on a single slice. With the appropriate setup:

• A checkpoint will be saved when a maintenance event occurs.
• Cloud TPU will perform maintenance on the affected TPU VM(s) after the checkpoint is saved.
• When Cloud TPU completes maintenance, you can use the TPU VM as usual.

1. Use the autocheckpoint-enabled flag when creating the TPU VM or queued resource.

   For example:

   PROJECT=your-gcp-project-name
   ZONE=zone-you-want-to-use
   NODE_ID=your-node-id
   ACCELERATOR_TYPE=your-accelerator-type
   gcloud config set project $PROJECT
   gcloud config set compute/zone $ZONE

   gcloud alpha compute tpus tpu-vm create $NODE_ID \
   --accelerator-type $ACCELERATOR_TYPE \
   --version tpu-ubuntu2204-base \
   --autocheckpoint-enabled

2. Install Pax on a single slice

   The Autocheckpoint feature works on Pax versions >= 1.1.0. On the TPU VMs, install jax[tpu] and the latest paxml:

   pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

3. Launch the training with the appropriate configuration

   The following example shows how to configure the LmCloudSpmd2B model to save checkpoints triggered by Autocheckpoint to a Google Cloud Storage bucket:

   JOB_LOG_DIR=gs://your-storage-bucket
   
   { python3 .local/lib/python3.10/site-packages/paxml/main.py
   --jax_fully_async_checkpoint=1 \
   --exit_after_ondemand_checkpoint=1 \
   --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2B \
   --job_log_dir=$JOB_LOG_DIR; } 2>&1 | tee pax_logs.txt

   Note the two flags that are passed to the command:

   • jax_fully_async_checkpoint: With this flag on, orbax.checkpoint.AsyncCheckpointer will be used. The AsyncCheckpointer class automatically saves a checkpoint when the training script receives a SIGTERM signal.
   • exit_after_ondemand_checkpoint: With this flag on, the TPU processes exits after the Autocheckpoint is successfully saved, which triggers the maintenance to be performed immediately. If you do not use this flag, the training will continue after the checkpoint is saved and Cloud TPU will wait for a timeout to occur (5 minutes) before performing the required maintenance..

Quickstart: Autocheckpoint with Pax on Multislice

Autocheckpoint works not only for single slices, but also for Multislice. This section details the steps needed to use Autocheckpoint with Multislice.

1. Specify Autocheckpoint during queued resource creation.

   A Multislice environment can only be provisioned through a queued resource request. Similar to the single-slice case, use the autocheckpoint-enabled flag in the call to create a queued resource.

   QR_ID=your-qr-id
   NODE_COUNT=your-node-count
   ACCELERATOR_TYPE=your-accelerator-type
   
   gcloud compute tpus queued-resources create $QR_ID \
   --node-count $NODE_COUNT \
   --accelerator-type $ACCELERATOR_TYPE \
   --runtime-version tpu-ubuntu2204-base \
   --autocheckpoint-enabled

   Refer to the Multislice User Guide for details on all available options. Once the queued resource request is created and in the ACTIVE state, follow the next steps to run Pax with Autocheckpoint.

2. Install Pax on all VMs in the Multislice environment.

   On the TPU VMs, install jax[tpu] and the latest paxml on all of the TPU VMs in your Multislice environment:

   pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

3. Launch the training with the appropriate configuration

   This example shows how to configure the model LmCloudSpmd2B for Autocheckpoint when training in a Multislice environment. Before running the training script, set DCN_MESH_SHAPE to [2, 1, 1] as shown in the following code:

   @experiment_registry.register
   class LmCloudSpmd2B(LmCloudSpmd):
   
   """SPMD model with 2B params.
   
   Global batch size = 2 * 2 * 1 * 32 = 128
   """
   PERCORE_BATCH_SIZE = 8
   
   NUM_LAYERS = 18
   MODEL_DIMS = 3072
   HIDDEN_DIMS = MODEL_DIMS * 4
   
   CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_NOTHING
   ICI_MESH_SHAPE = [1, 4, 1]
   DCN_MESH_SHAPE = [2, 1, 1]

   When launching the training, in addition to the command line flags discussed in the single-slice case, three more are required:

   • num_hosts: the total number of hosts. In this case, it is 2.
   • host_index: the index of the host launching the training. It varies from 0 to N-1 where N is the total number of hosts.
   • server_addr: the IP address of worker 0 of node 0, with an unused port (for example, 8476). To find this information, use hostname -i on worker 0 of node 0.

Autocheckpoint with Orbax

The Autocheckpoint feature is not limited to MaxText or Pax. Any framework that can capture the SIGTERM signal and initiate a checkpointing process works with the infrastructure provided by Autocheckpoint. Orbax, a namespace that provides common utility libraries for JAX users, provides these capabilities.

As explained in the Orbax documentation, these capabilities are enabled by default for users of orbax.checkpoint.CheckpointManager. The save method that is called after every step automatically checks whether a maintenance event is impending, and if so, saves a checkpoint even if the step number is not a multiple of save_interval_steps. The GitHub documentation also illustrates how to make the training exit after saving an Autocheckpoint, with a modification in the user code.