Elastic Training: TPU Recovery From Failures in Seconds

Ever found yourself deep into a multi-node training run, only for a single machine to unexpectedly flatline? If you’ve trained large models across several machines, you know the drill: communication times out, every worker exits, and you’re left to relaunch the entire job from the last checkpoint. It’s frustrating, costly, and seems like an unavoidable part of distributed training. But what if it wasn’t?

This article dives into a game-changing concept called elastic training. We’ll explore how it works using the JAX AI stack (specifically MaxText and Pathway) on Cloud TPUs. We’ll deliberately cause a worker failure during an LLM training session across multiple TPU chips on Google Kubernetes Engine (GKE) and observe the training process recover in place, without a full restart. Imagine getting back on track in under two minutes, often just waiting for Kubernetes to schedule a replacement pod, all within the same process and PID.

By the end, you’ll grasp the three key components that make this possible, understand its current limitations, and even learn how to replicate the entire process yourself. Get ready to rethink how you handle distributed training failures.

The Problem: A Single Point of Failure

Picture this: you’re training a massive model distributed across several machines, each holding a sharded piece of your model weights. Every training step involves these machines computing gradients on their respective pieces, followed by an “all-reduce” operation. This operation ensures all machines exchange gradients, keeping the model synchronized.

Here’s the critical flaw: an all-reduce operation absolutely requires every participant. If just one machine vanishes, the others are left indefinitely waiting for data that will never arrive. Inevitably, a timeout occurs, the collective operation fails, and every single process exits. One machine’s demise brings down the entire job.

The conventional solution usually lies outside your training code. A scheduler, be it Slurm, Kubernetes, or Ray, detects the job’s failure, reallocates resources, and relaunches everything from scratch. This incurs a significant restart cost: fresh pods need scheduling, containers and Python processes must start anew, connections to accelerators re-established, and dataloaders warmed up. Crucially, you lose all progress made since your last checkpoint.

But what if your training process could gracefully catch the failure and simply continue? That’s precisely the “elastic” approach we’re about to uncover.

Introducing the TPU Ecosystem and Pathways

For those new to the TPU world, it has its own unique terminology. Let’s quickly introduce the hardware and software components powering our elastic training adventure.

Our foundational hardware unit is the TPU chip, an accelerator designed for rapid matrix multiplication. Multiple chips are then clustered into a TPU slice, a group of chips interconnected by a super-fast, dedicated network called ICI (Inter-Chip Interconnect). Within a slice, these chips are attached to standard CPU host VMs (our setup uses four hosts per slice), which are responsible for running the actual worker processes.

To program these powerful chips, we use JAX, a NumPy-like array and automatic differentiation framework that leverages XLA as its compiler. If you’re familiar with PyTorch, JAX serves a similar role. However, we won’t be building a training loop from scratch; instead, we’ll utilize MaxText, an open-source LLM training library built entirely in Python and JAX. Simply provide a model name and configuration, and MaxText delivers a fully sharded training loop.

Two more crucial pieces complete the puzzle. Pathways acts as the orchestration layer, linking our Python script to all the TPU chips in the cluster – and it’s absolutely central to enabling elastic training. Finally, Orbax handles checkpointing, coordinating saves from a central controller while each TPU host writes its model state shard directly to Cloud Storage in parallel. This robust checkpointing mechanism provides a reliable fallback when things inevitably go wrong.

The Power of a Single Controller

The key to understanding elastic training’s recovery mechanism lies in the concept of a single controller. Unlike most distributed training launchers where you start one Python process per node (each running an identical script and coordinating as equals, known as SPMD or Single Program Multiple Data), Pathways operates differently.

With Pathways, there’s exactly one Python process, running on a standard CPU machine, which perceives every TPU chip in the cluster as if it were locally attached. Calling jax.devices() will return all available chips. The TPU machines themselves merely run a lightweight worker binary, receiving compiled programs and executing them. This architecture is vital for failure handling: if a TPU machine fails, a healthy Python process on the CPU node is still alive and capable of taking action.

Elastic Training in Action: Recovery Strategies

Let’s clarify what “elastic” truly means in this context, as the term can be applied to various scenarios. Here, elastic training signifies that when hardware fails, your training loop receives a Python exception instead of a process termination. Because your process is still live, with your configuration and imports already loaded, and surviving TPU slices ready for instructions, you gain powerful recovery options.

Two straightforward yet effective examples are pause and resume and replica resize.

In the pause and resume strategy, the exception is caught, you wait for the failed slice to be replaced, then reload the last viable checkpoint and resume training on the newly complete mesh. With replica resize, you immediately reload the last viable checkpoint onto the surviving slices, allowing training to continue at reduced throughput while the replacement comes online. Once the new slice is ready, training scales back up to full size. Both methods are currently available in MaxText.

This article focuses on the simpler pause and resume approach. While you could write your own exception handler, the pathwaysutils library offers an @elastic_retry decorator that wraps an entire training function, and MaxText conveniently pre-wires it for you. When a failure exception fires, this decorator catches it, cleans up any partial state, restores the last viable checkpoint, and re-invokes the training function—all within the same enduring process.

Why Elastic Recovery is Faster

It’s important to understand precisely why elastic recovery outpaces a full job restart. While re-entering the training function means model setup, dataloader initialization, and checkpoint restoration all run again (costs also present in a full restart), elastic recovery saves you a crucial step: the full workload teardown and rescheduling. A complete restart tears down and reschedules the entire job, including the controller (head) pod and every healthy worker pod, along with a fresh controller Python process. Elastic recovery, however, leaves all these running, only replacing the specific slice that died.

Compilation isn’t a primary saving point here because Pathways utilizes a persistent compilation cache in Cloud Storage (enabled by default). Both full restarts and elastic recovery reload compiled XLA executables from this cache, avoiding cold recompilations. The real efficiency gain comes from skipping the full-workload teardown. In our tests, this difference translated to recovering in around hundreds of seconds compared to several minutes. Furthermore, replica resize offers a unique advantage: training can continue even if some TPUs are permanently unavailable, a capability a standard restart simply cannot provide.

Before moving on, a quick distinction: suspend-resume and elastic pause and resume sound similar but address different issues. Suspend-resume, a Pathways feature for Spot TPUs, handles planned preemptions by automatically saving accelerator state to Cloud Storage upon notice and resuming when rescheduled. No user code is needed. Elastic training, conversely, is designed for unplanned failures—those interruptions that strike without warning.

The Three Pillars of Elastic Recovery

For truly elastic recovery to work, three distinct components must seamlessly cooperate. Here’s a breakdown of how they operate within the cluster:

• Pathways detects the failure: This detection can manifest in two ways. Most commonly, an in-flight operation to the dead worker will fail, prompting Pathways to return a DATA_LOSS error. If no operations are active, the resource manager (a container co-located with our script on the CPU node) will notice the worker has stopped heartbeating, returning a DEADLINE_EXCEEDED error after approximately 10 seconds. In either scenario, the hardware failure is elevated to a catchable Python exception: jax.errors.JaxRuntimeError.
• The @elastic_retry decorator catches the exception: This decorator, sourced from pathwaysutils, is applied by MaxText around its training function. It specifically intercepts the JaxRuntimeError, logs a “Slice down event detected. Retrying.” message, and initiates the recovery path instead of allowing the error to crash the entire process.
• Orbax ensures safe checkpoint restoration: This functionality is inherent to Orbax’s general checkpointing mechanism and would be leveraged identically in a full job restart. Checkpoints are asynchronously written to Cloud Storage during training. A checkpoint is only deemed viable once all its shards have been flushed and a small commit_success marker file is written alongside it. During recovery, the cleanup code inspects the latest checkpoint directory: if the marker is absent (indicating an in-progress write at the time of failure), that directory is deleted, and the system falls back to the newest fully committed checkpoint. This guarantees that we never load a partial or corrupted checkpoint, regardless of the recovery method.

Setting Up and Triggering a Failure

With the mechanics now clear, let’s put it into practice and intentionally break something. Our experiment will be kept small to make the failure-and-recovery loop quick to observe.

The entire demonstration takes about 30 minutes. At on-demand v5e list prices, this equates to roughly $30 for 48 chips (approximately $1.20 per chip-hour for half an hour), plus the cost of the CPU controller node. The training job itself will run for as long as you configure it; we just need it active long enough to simulate a failure and observe recovery.

Once the cluster is provisioned, we need two main things: a MaxText command to run on the head pod and a JobSet manifest to link this command to the TPU slices. Let’s examine each.

MaxText Configuration for Elasticity

MaxText can be fully configured via command-line flags layered over a base YAML file. Here’s the command our head pod executes, trimmed to highlight the essential components for this demo:

python3 -m maxtext.trainers.pre_train.train \
src/maxtext/configs/base.yml \
base_output_directory=gs://${BUCKET_NAME}/output \
run_name=${RUN_NAME} \
model_name=qwen3-0.6b \
per_device_batch_size=1 \
steps=5000 \
enable_checkpointing=true \
checkpoint_period=100 \
enable_single_controller=true \
elastic_enabled=true \
elastic_timeout_seconds=300 \
elastic_max_retries=10 \
dataset_type=grain \
grain_file_type=arrayrecord \
grain_train_files=gs://${BUCKET_NAME}/data/glaive-fc-v2/train.array_record*

Most of this command is standard MaxText: it selects a model, specifies data paths, and sets a batch size. However, four critical flags activate the elastic behavior:

• enable_single_controller=true: This tells MaxText to use the Pathways single-controller architecture, essential for intercepting failures.
• elastic_enabled=true: This activates the elastic retry mechanism within MaxText.
• elastic_timeout_seconds=300: This sets the maximum time (in seconds) the system will wait for a replacement slice after a failure before giving up.
• elastic_max_retries=10: This specifies how many times the system will attempt to recover from failures.

There’s a fifth flag we’re implicitly relying on: elastic_min_slice_count. Its default value of -1 means that all slices must be available before training resumes, enabling the “pause and resume” mode demonstrated here. Setting it to a value between 1 and numSlices - 1 would activate “replica resize,” allowing training to continue on the surviving slices while waiting for replacements. Another notable flag is checkpoint_period=100. MaxText’s default is 10,000 steps; at our approximate 0.16 seconds per step, 100 steps ensures a new checkpoint begins roughly every…

Source: Google Developers Blog