We terminated a TPU mid-training and it recovered in seconds:Introduction to elastic training with MaxText

We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText

Distributed AI training is notoriously fragile because losing a single machine typically crashes the entire multi-node job, forcing a time-consuming, full-workload infrastructure restart. To address this, Google's JAX ecosystem utilizes elastic training via Pathways, which converts a hardware failure into a catchable Python exception so the running process can survive. When an unplanned failure occurs, the system automatically replaces only the broken worker, restores the last viable checkpoint from Cloud Storage, and resumes training in place-minimizing total downtime to under two minutes without ever restarting the main controller process.

What happens when one machine dies in the middle of a multi-node training run? If you've trained large models across many machines, you already know the answer: the communication times out, every worker exits, and you re-launch the whole job from the last checkpoint. It's painful, and it's just how distributed training works. Or is it? In this article we'll explore a possible answer, called elastic training, using the JAX AI stack (MaxText and Pathways) and on Cloud TPUs. We'll train a LLM across multiple TPU chips on Google Kubernetes Engine (GKE), cause a worker to fail on purpose, and watch the training process recover in place without restarting. All this using the same process, same PID, no re-launch. In our run, total downtime from kill to the next training step was about less than two minutes, and most of that was waiting for Kubernetes to schedule a replacement pod. By the end, you'll understand exactly which three components make that possible, where the approach still has rough edges, and how to reproduce everything yourself.

The problem with distributed training

Let's start with the problem. Imagine you're training a model across multiple machines (or nodes). Your model weights are sharded, so each machine holds a piece. Every training step, the machines compute gradients on their piece and then an all-reduce operation runs where everyone exchanges gradients so the model stays in sync.

Here's the catch: an all-reduce needs every participant. If one machine disappears, the other ones sit there waiting for data that will never arrive. Eventually a timeout fires, the collective fails, and every process exits. As a result, one machine takes down the whole job.

The standard fix is usually outside your training code. A scheduler (Slurm, Kubernetes, Ray, take your pick) notices the job died, reallocates it and re-launches everything from scratch. You pay the full restart cost: scheduling fresh pods, starting fresh containers and Python processes, reconnecting to the accelerators, and warming up the dataloader. And you lose every step since the last checkpoint.

But what if the training process could just catch the failure and keep going? That's what we're going to look at now.

The TPU ecosystem

If you're new to the TPU ecosystem, it has its own vocabulary. Let's introduce the pieces we'll use and how they fit together.

Our hardware unit is the TPU chip, an accelerator built around matrix-multiply units. Chips are grouped into a TPU slice, a set of chips wired together with a fast dedicated interconnect called ICI (Inter-Chip Interconnect). Within a slice, chips are attached to ordinary CPU host VMs (four hosts per slice in our setup), and those hosts are what actually run the worker processes.

To program those chips we use JAX, a NumPy-flavored array and autograd framework that uses XLA as its compiler. If you're coming from PyTorch, it fills the same role. We won't write a training loop from scratch, though. We'll use MaxText, an open-source LLM training library written in pure Python and JAX. You hand it a model name and a config, and it gives you a fully sharded training loop.

Two more pieces complete the picture. Pathways is the orchestration layer that connects our Python script to all chips, and it's the key to this whole story for a reason we'll get to in a moment. And Orbax handles checkpointing: it coordinates the save from the controller, while each TPU host writes its own shard of the model state directly to Cloud Storage in parallel. That gives us something to fall back to when things go wrong.

Here's how it looks all together:

The component to remember here is the single controller. In most distributed-training launchers, you start one Python process per node. Each one runs an identical copy of your script, and they coordinate as equals (this is called SPMD, or Single Program Multiple Data). With Pathways, there is exactly one Python process, running on a plain CPU machine, and it sees every TPU chip in the cluster as if they were local. Call jax.devices() and you get all chips back. The TPU machines themselves just run a thin worker binary that receives compiled programs and executes them.

Why does this matter for failure handling? Because when a TPU machine dies, there is still a healthy Python process alive on the CPU node that can do something about it.

What "elastic training" means

Let's see what "doing something about it" looks like. Let's be clear about what we're building, because "elastic" gets used for a lot of things. At its core, elastic training here means that when hardware fails, your training loop receives a Python exception instead of a process termination. And because you're still inside a live process, with your config and imports already loaded and the surviving TPU slices still up and awaiting instructions, you have options.

Two simple, yet powerful, examples are:

• Pause and resume: The exception gets caught, you wait for the failed slice to be replaced, then reload the last viable checkpoint and continue on the full mesh.
• Replica resize: You reload the last viable checkpoint immediately onto the surviving slices and training keeps running at reduced throughput while the replacement comes up, then scales back to full size once it's ready.

Both are available in MaxText today. This post walks through pause and resume, the simpler of the two.

You could write that exception handler yourself, but you don't have to. The pathwaysutils library provides a decorator called elastic_retry that wraps an entire training function, and MaxText already wires it in for you. When the failure exception fires, the decorator catches it, cleans up any partial state, restores the last viable checkpoint, and calls the training function again. All inside the same process.

Why elastic recovery is faster

It's worth being precise about why this is faster than a restart, because elastic recovery doesn't skip as much as you might think. Calling the training function again from the beginning means model setup, the dataloader, and the checkpoint restore all run a second time-but you'd pay every one of those on a full job restart too. Pod scheduling isn't free here either: the failed worker still needs a replacement pod scheduled onto the impacted slice, and that wait dominates the wall clock.

What elastic recovery actually saves you is everything around that one slice. A full restart tears down and reschedules the whole workload, from the controller (head) pod, every healthy worker pod, to the fresh controller Python process that comes with them, while elastic recovery leaves all of it running and only swaps the slice that died.

Compilation, notably, is not part of that saving. Pathways keeps a persistent compilation cache in Cloud Storage (on by default), so a full restart reloads the compiled XLA executables from the cache instead of recompiling cold, and elastic recovery pays a comparable cost when it re-enters the training function on the rebuilt mesh. The difference between the two paths is the teardown, not the compile-and in our run, skipping that full-workload teardown was the difference between about hundreds of seconds and several minutes.

Replica resize then adds something a restart can't offer at all: training keeps going even if some of the TPUs never come back.

Suspend-resume vs. elastic training

Before we go on, a quick word on a name that's easy to confuse with what we've just described. Suspend-resume and the elastic pause and resume sound alike but solve different problems. If you're running on Spot TPUs, planned preemptions take their own path: Pathways' suspend-resume feature listens for the preemption notice, saves accelerator state to Cloud Storage automatically, and resumes when the pod is rescheduled-no user code needed. That's for interruptions that arrive with a warning. Elastic training, the mechanism we're walking through here, is the path for the unplanned failures where there's no notice at all.

The three cooperating components

Now let's look at the three pieces that actually have to cooperate to make this work. For elastic recovery, three independent components have to cooperate. Here's how they're laid out on the cluster.

1. Pathways detects the failure

This can surface in two ways. Most commonly, an in-flight operation to the dead worker fails and Pathways returns a DATA_LOSS error. If nothing happens to be in flight, the resource manager (a container running alongside our script on the CPU node) notices the worker has stopped heartbeating and returns DEADLINE_EXCEEDED after about 10 seconds. Either way, the error arrives in our training step as a jax.errors.JaxRuntimeError. The hardware failure has become a catchable Python exception.

2. The elastic_retry decorator catches the exception

This decorator comes from pathwaysutils; MaxText simply applies it around its training function. It catches that specific exception, logs a "Slice down event detected. Retrying." message, and runs the recovery path instead of letting the error crash the process.

3. Orbax decides what's safe to restore

This part isn't unique to elastic training; it's how Orbax checkpointing works in general, and a full job restart would lean on it in exactly the same way. Checkpoints are written to Cloud Storage in the background while training runs, and a checkpoint is only considered viable when every shard has been flushed and a tiny commit_success marker file is written next to it. During recovery, the cleanup code checks the latest checkpoint directory: if there's no marker (it was mid-write when things broke), the directory is deleted, and we fall back to the newest checkpoint that does have one. That's what guarantees we never load half a checkpoint, regardless of how we restart.

Running the experiment

With the mechanics clear, let's actually run it and break something. We'll keep our experiment intentionally small so the failure-and-recovery loop is fast to watch.

Here is our setup: 1.35.3-gke.1993000.

Walking through the whole demo takes about 30 minutes end to end. At on-demand v5e list prices that's roughly $30 with 48 chips at ~$1.20/chip-hour for half an hour, plus the CPU controller node. The training job itself will run for as long as you configure it to; we just need it active long enough to break it.

Once the cluster is up, we need two things: a MaxText command that runs on the head pod, and a JobSet manifest that wires that command to the TPU slices. Let's look at each.

The MaxText command

MaxText can be configured entirely through command-line flags layered on top of a base YAML. Here's the command our head pod runs, trimmed to the parts that matter 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 is standard MaxText. It picks a model, points at data, and sets a batch size. Four flags turn on the elastic behavior:

• enable_single_controller=True routes JAX through the Pathways proxy instead of talking to local devices. This is what makes one Python process see all 48 chips, and it's a hard requirement for everything below.
• elastic_enabled=true wraps the training function in the elastic_retry decorator we described earlier and waits for the minimum number of slices before starting.
• elastic_timeout_seconds=300 is how long the retry loop will wait for a failed slice to be replaced before giving up on this attempt.
• elastic_max_retries=10 is how many failures we'll tolerate over the whole run before exiting for real.

There's a fifth flag we're relying on without passing it: elastic_min_slice_count. It controls how many slices must be available before the retry resumes training. The default is -1, which means all of them, and that's pause and resume, the mode we're running here. Setting it to a value between 1 and numSlices - 1 would enable replica resize instead, where training keeps going on the surviving slices rather than waiting.

One more flag worth calling out: checkpoint_period=100. MaxText's default is 10,000 steps. At our ~0.16s per step, 100 steps means a new checkpoint starts roughly every 16 seconds, so there's always something recent to fall back to. You can go lower still; the right value is a trade-off between step time, checkpoint write time, and how often you expect failures.

One caveat: if a slice fails during an active checkpoint write, the current version of MaxText exits rather than retrying. A frequent-enough checkpoint period creates safe windows between writes; alternatively, set enable_continuous_checkpointing=True and let Orbax start the next save as soon as the previous one finishes, so you're always checkpointing as fast as your storage allows and the fixed period stops mattering.

Submitting the workload

The command above doesn't know anything about Kubernetes. The easiest way to run it across three TPU slices is xpk, which takes a cluster, a TPU type, and your training command, and submits the workload for you. The elastic settings are just two flags:

xpk workload create-pathways \
  --workload=${RUN_NAME} --cluster=${GKE_CLUSTER} \
  --tpu-type=v5litepod-16 --num-slices=3