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

TLDR

What occurs when one machine dies in the course of a multi-node coaching run?

If you happen to’ve skilled massive fashions throughout many machines, you already know the reply: the communication occasions out, each employee exits, and also you re-launch the entire job from the final checkpoint. It is painful, and it is simply how distributed coaching works. Or is it?

On this article we’ll discover a attainable reply, referred to as elastic coaching, utilizing the JAX AI stack (MaxText and Pathway) and on Cloud TPUs. We’ll practice a LLM throughout a number of TPU chips on Google Kubernetes Engine (GKE), trigger a employee to fail on objective, and watch the coaching course of get well in place with out restarting. All this utilizing the identical course of, similar PID, no re-launch. In our run, whole downtime from kill to the following coaching step was about lower than two minutes , and most of that was ready for Kubernetes to schedule a alternative pod.

By the top, you will perceive precisely which three elements make that attainable, the place the method nonetheless has tough edges, and find out how to reproduce all the things your self.

Let’s begin with the issue.

Why distributed coaching is so fragile

Think about you are coaching a mannequin throughout a number of machines (or nodes). Your mannequin weights are sharded, so every machine holds a chunk. Each coaching step, the machines compute gradients on their piece after which an all-reduce operation runs the place everybody exchanges gradients so the mannequin stays in sync.

This is the catch: an all-reduce wants each participant. If one machine disappears, the opposite ones sit there ready for knowledge that may by no means arrive. Ultimately a timeout fires, the collective fails, and each course of exits. Because of this, one machine takes down the entire job.

The usual repair is normally outdoors your coaching code. A scheduler (Slurm, Kubernetes, Ray, take your choose) notices the job died, reallocates it and re-launches all the things from scratch. You pay the complete restart value: scheduling recent pods, beginning recent containers and Python processes, reconnecting to the accelerators, and warming up the dataloader. And also you lose each step for the reason that final checkpoint.

However what if the coaching course of may simply catch the failure and preserve going? That is what we’ll have a look at now.

MaxText, Pathways, and Orbax on TPU

If you happen to’re new to the TPU ecosystem, it has its personal vocabulary. Let’s introduce the items we’ll use and the way they match collectively.

Our {hardware} unit is the TPU chip, an accelerator constructed round matrix-multiply models. Chips are grouped right into a TPU slice, a set of chips wired along with a quick devoted interconnect referred to as ICI (Inter-Chip Interconnect). Inside a slice, chips are connected to extraordinary CPU host VMs (4 hosts per slice in our setup), and people hosts are what truly run the employee processes.

To program these chips we use JAX, a NumPy-flavored array and autograd framework that makes use of XLA as its compiler. If you happen to’re coming from PyTorch, it fills the identical position. We can’t write a coaching loop from scratch, although. We’ll use MaxText, an open-source LLM coaching library written in pure Python and JAX. You hand it a mannequin identify and a config, and it offers you a totally sharded coaching loop.

Two extra items full the image. Pathways is the orchestration layer that connects our Python script to all chips, and it is the important thing to this entire story for a cause we’ll get to in a second. And Orbax handles checkpointing: it coordinates the save from the controller, whereas every TPU host writes its personal shard of the mannequin state on to Cloud Storage in parallel. That provides us one thing to fall again to when issues go flawed.

This is the way it appears to be like all collectively:

The element to recollect right here is the single controller.

In most distributed-training launchers, you begin one Python course of per node. Each runs an an identical copy of your script, they usually coordinate as equals (that is referred to as SPMD, or Single Program A number of Knowledge). With Pathways, there’s precisely one Python course of, working on a plain CPU machine, and it sees each TPU chip within the cluster as in the event that they have been native. Name jax.gadgets() and also you get all chips again. The TPU machines themselves simply run a skinny employee binary that receives compiled packages and executes them.

Why does this matter for failure dealing with? As a result of when a TPU machine dies, there’s nonetheless a wholesome Python course of alive on the CPU node that may do one thing about it.

Let’s examine what “doing one thing about it” appears to be like like.

What “elastic coaching” means right here

Let’s be clear about what we’re constructing, as a result of “elastic” will get used for lots of issues.

At its core, elastic coaching right here implies that when {hardware} fails, your coaching loop receives a Python exception as a substitute of a course of termination. And since you’re nonetheless inside a stay course of, along with your config and imports already loaded and the surviving TPU slices nonetheless up and awaiting directions, you’ve gotten choices.

Two easy, but highly effective, examples are pause and resume and duplicate resize.

In pause and resume, the exception will get caught, you look forward to the failed slice to get replaced, then reload the final viable checkpoint and proceed on the complete mesh. In duplicate resize, you reload the final viable checkpoint instantly onto the surviving slices and coaching retains working at diminished throughput whereas the alternative comes up, then scales again to full measurement as soon as it is prepared.

Each can be found in MaxText as we speak. This submit walks via pause and resume, the less complicated of the 2. You would write that exception handler your self, however you do not have to. The pathwaysutils library offers a decorator referred to as elastic_retry that wraps a complete coaching operate, and MaxText already wires it in for you. When the failure exception fires, the decorator catches it, cleans up any partial state, restores the final viable checkpoint, and calls the coaching operate once more. All inside the identical course of.

It is price being exact about why that is quicker than a restart, as a result of elastic restoration does not skip as a lot as you would possibly suppose. Calling the coaching operate once more from the start means mannequin setup, the dataloader, and the checkpoint restore all run a second time however you’d pay each a type of on a full job restart too. Pod scheduling is not free right here both: the failed employee nonetheless wants a alternative pod scheduled onto the impacted slice, and that wait dominates the wall clock. What elastic restoration truly saves you is all the things round that one slice. A full restart tears down and reschedules the entire workload, from the controller (head) pod, each wholesome employee pod, to the recent controller Python course of that comes with them, whereas elastic restoration leaves all of it working and solely swaps the slice that died.

Compilation, notably, isn’t a part of that saving. Pathways retains a persistent compilation cache in Cloud Storage (on by default), so a full restart reloads the compiled XLA executables from the cache as a substitute of recompiling chilly, and elastic restoration pays a comparable value when it re-enters the coaching operate on the rebuilt mesh. The distinction between the 2 paths is the teardown, not the compile — and in our run, skipping that full-workload teardown was the distinction between about a whole lot of seconds and a number of other minutes. Duplicate resize then provides one thing a restart cannot provide in any respect: coaching retains going even when a number of the TPUs by no means come again.

Earlier than we go on, a fast phrase on a reputation that is simple to confuse with what we have simply described. Droop-resume and the elastic pause and resume sound alike however clear up totally different issues. If you happen to’re working on Spot TPUs, deliberate preemptions take their very own path: Pathways’ suspend-resume characteristic listens for the preemption discover, saves accelerator state to Cloud Storage mechanically, and resumes when the pod is rescheduled — no person code wanted. That is for interruptions that arrive with a warning. Elastic coaching, the mechanism we’re strolling via right here, is the trail for the unplanned failures the place there is no discover in any respect.

Now let us take a look at the three items that truly need to cooperate to make this work.

How elastic restoration works

For elastic restoration, three unbiased elements need to cooperate. This is how they’re laid out on the cluster.

First, Pathways detects the failure. This could floor in two methods. Mostly, an in-flight operation to the useless employee fails and Pathways returns a DATA_LOSS error. If nothing occurs to be in flight, the useful resource supervisor (a container working alongside our script on the CPU node) notices the employee has stopped heartbeating and returns DEADLINE_EXCEEDED after about 10 seconds. Both manner, the error arrives in our coaching step as a jax.errors.JaxRuntimeError. The {hardware} failure has develop into a catchable Python exception.

Subsequent, the elastic_retry decorator catches the exception. This decorator comes from pathwaysutils; MaxText merely applies it round its coaching operate. It catches that particular exception, logs Slice down occasion detected. Retrying. message, and runs the restoration path as a substitute of letting the error crash the method.

Lastly, Orbax decides what’s secure to revive. This half is not distinctive to elastic coaching; it is how Orbax checkpointing works usually, and a full job restart would lean on it in precisely the identical manner. Checkpoints are written to Cloud Storage within the background whereas coaching runs, and a checkpoint is barely thought-about viable when each shard has been flushed and a tiny commit_success marker file is written subsequent to it. Throughout restoration, the cleanup code checks the newest checkpoint listing: if there is no marker (it was mid-write when issues broke), the listing is deleted, and we fall again to the latest checkpoint that does have one. That is what ensures we by no means load half a checkpoint, no matter how we restart.

With the mechanics clear, let’s truly run it and break one thing.

Establishing and submitting the elastic coaching job

We’ll preserve our experiment deliberately small so the failure-and-recovery loop is quick to observe. Right here is our setup:

• {Hardware}: 3 x TPU v5e-16 slices = 48 chips, plus one n2-standard-64 CPU node for the controller.
• Platform: Google Kubernetes Engine. All the pieces runs as pods. The Kubernetes useful resource that ties it collectively is a JobSet, which lays out 1 head pod + 12 employee pods and retains them as one unit.
• Mannequin: qwen3-0.6b. Small on objective so the failure-and-recovery loop is fast to observe and low-cost to run. We’ll cowl what adjustments while you scale to actual mannequin sizes in a second.
• Knowledge: the Glaive function-calling dataset, pre-converted to ArrayRecord shards on Cloud Storage.
• Variations: MaxText at commit 992b4e1, GKE 1.35.3-gke.1993000.

Strolling via the entire demo takes about half-hour finish to finish. At on-demand v5e checklist costs that is roughly $30 with 48 chips at ~$1.20/chip-hour for half an hour, plus the CPU controller node. The coaching job itself will run for so long as you configure it to; we simply want it energetic lengthy sufficient to interrupt it.

As soon as the cluster is up, we’d like two issues: a MaxText command that runs on the top pod, and a JobSet manifest that wires that command to the TPU slices. Let us take a look at every.

The MaxText coaching command

MaxText may be configured totally via command-line flags layered on prime of a base YAML. This is the command our head pod runs, trimmed to the components that matter for this demo:

Most of that is normal MaxText. It picks a mannequin, factors at knowledge, and units a batch measurement. 4 flags activate the elastic conduct:

• enable_single_controller=True routes JAX via the Pathways proxy as a substitute of speaking to native gadgets. That is what makes one Python course of see all 48 chips, and it is a laborious requirement for all the things beneath.
• elastic_enabled=true wraps the coaching operate within the elastic_retry decorator we described earlier and waits for the minimal variety of slices earlier than beginning.
• elastic_timeout_seconds=300 is how lengthy the retry loop will look forward to a failed slice to get replaced earlier than giving up on this try.
• elastic_max_retries=10 is what number of failures we’ll tolerate over the entire run earlier than exiting for actual.

There is a fifth flag we’re counting on with out passing it: elastic_min_slice_count. It controls what number of slices have to be out there earlier than the retry resumes coaching. The default is -1, which suggests all of them, and that is pause and resume, the mode we’re working right here. Setting it to a price between 1 and numSlices – 1 would allow duplicate resize as a substitute, the place coaching retains occurring the surviving slices reasonably than ready.

Yet another flag price calling out: checkpoint_period=100. MaxText’s default is 10,000 steps. At our ~0.16s per step, 100 steps means a brand new checkpoint begins roughly each 16 seconds, so there’s all the time one thing current to fall again to. You possibly can go decrease nonetheless; the proper worth is a trade-off between step time, checkpoint write time, and the way usually you count on failures. One caveat: if a slice fails throughout an energetic checkpoint write, the present model of MaxText exits reasonably than retrying. A frequent-enough checkpoint interval creates secure home windows between writes; alternatively, set enable_continuous_checkpointing=True and let Orbax begin the following save as quickly because the earlier one finishes, so that you’re all the time checkpointing as quick as your storage permits and the mounted interval stops mattering.

Submitting the coaching

The command above does not know something about Kubernetes. The simplest approach to run it throughout three TPU slices is xpk, which takes a cluster, a TPU kind, and your coaching command, and submits the workload for you. The elastic settings are simply two flags:

That is the entire submission. --elastic-slices=3 tells Pathways what number of slices could also be lacking earlier than GKE offers up and restarts the entire JobSet — distinct from the MaxText elastic_min_slice_count above, which is what number of slices have to be current for the retry to try coaching. --max-slice-restarts is the restart funds. The official elastic coaching information makes use of this xpk path finish to finish.

Underneath the hood, xpk wraps your command in a JobSet and fingers it to GKE. A JobSet is a single Kubernetes useful resource that teams a number of Jobs collectively and offers them a shared restart coverage and a shared headless Service so the pods can discover one another by identify. That is precisely what a Pathways cluster wants: one head Job on the CPU node, and one employee Job per TPU slice. With elasticity, restoration is extra surgical than a full JobSet restart: when a slice fails, solely that slice’s employee Job is recreated, on the Job stage — the top Job and the opposite employee Jobs preserve working untouched.

You usually by no means see this, however it’s price as soon as, as a result of one line in it bit me later. The complete manifest is about 230 traces (principally env-var plumbing); here is the form of it with the components that matter for elastic coaching left in:

Let’s break it down. The pathways-head Job is the CPU aspect. Its pod runs three containers: our primary container with the MaxText command from above, plus two Pathways containers. The pathways-rm container is the useful resource supervisor (it assigns slices to shoppers, compiles XLA capabilities, and tracks slice well being, amongst different issues), and pathways-proxy is the IFRT proxy that JAX talks to once we set JAX_PLATFORMS=proxy.

The employee Job is the TPU aspect. replicas: 3 offers us three copies of the Job, one per slice, and completions: 4 / parallelism: 4 places 4 pods in every (one per TPU host). These pods do not run our code in any respect. They run the Pathways employee binary, which connects again to the useful resource supervisor and waits to be handed compiled XLA packages.

The only most essential line for elasticity is backoffLimit: 20 on the employee Job. It lets a failed slice’s pods restart in place — on the Job stage — as much as 20 occasions earlier than the employee Job itself is marked failed, which is what would escalate right into a full JobSet restart. In different phrases, a excessive backoffLimit is what retains a slice failure native: the slice’s pods come again, the top and the opposite slices preserve working, and also you keep away from the costly whole-workload restart. (Sidenote: newer JobSet variations are including a devoted job-level restart coverage that may categorical this conduct extra immediately than backoffLimit does as we speak.)

The one argument that is particular to elastic coaching is --num_elastic_slices=3 on the proxy (the in-manifest type of --elastic-slices). We set it equal to the slice depend, which tells Pathways it could be lacking any variety of slices, even all of them, earlier than GKE would hand over and restart the JobSet. Shedding the entire slices is secure in pause-and-resume mode as a result of restoration state comes from the GCS checkpoint, not from the surviving slices.

Yet another area price a look is limits: {reminiscence: 100G} on the pathways-proxy container. xpk picks a default you’ll be able to override, although for actual mannequin sizes the higher reply is not a much bigger quantity right here — it is turning on checkpoint persistence, which we get to within the scaling part.

As soon as the workload is submitted, you’ll be able to watch it the standard Kubernetes manner:

If you happen to’d reasonably not tail within the terminal, Cloud Logging offers you an identical output with search and historical past that persists via the pod restarts. You possibly can filter on useful resource.labels.container_name="primary". A minute or so later the log begins scrolling: coaching is working on all 48 chips, loss is dropping, ~43 TFLOP/s per system as proven beneath.

Now the enjoyable half.

Flip down a employee and watch what occurs

As soon as coaching has been going for some time and there are a couple of checkpoints on disk, we choose a employee pod on slice 2 and force-kill it:

The --grace-period=0 --force means SIGKILL now. No swish shutdown, no cleanup. That is how we imitate actual {hardware} failures that do not give anybody an opportunity to organize.

This is what occurs behind the scenes:

Let’s stroll via what the diagram is displaying, with a stopwatch working from the second of the kill.

The very first thing to note is that failure is not on the spot. For about 13 seconds the coaching loop has no thought something is flawed: the employee pod is already gone, however the useful resource supervisor’s heartbeat window hasn’t closed but, and JAX dispatch is asynchronous, so steps preserve touchdown all the best way as much as step 3388. Solely when the heartbeat occasions out does Pathways increase the JaxRuntimeError into our Python course of, and elastic_retry catches it with a single log line: Slice down occasion detected. Retrying.

The handler’s first transfer is housekeeping. It lists the checkpoint listing on Cloud Storage, sees that step 3300 has its commit_success marker, and confirms there’s nothing half-written to delete. That takes below a second.

Then it waits on infrastructure, not on our code. Kubernetes has to schedule a alternative employee pod onto slice 2, and that pod has to start out its container and rejoin the Pathways mesh. In our run that took about 50 seconds, and it is the place a lot of the wall-clock time goes. At roughly the 64-second mark the log prints Ample slices energetic: 3 >= 3 and the handler re-enters the coaching operate from the highest.

Up so far, what we’re doing appears to be like loads like a JobSet restart — besides for 3 variations that matter. We solely reschedule the one failed slice, not the entire workload; the controller’s Python state can persist throughout the occasion if we would like it to; and we may be selective about what will get reinitialized (as we speak we reinitialize all the things, however that is a alternative, not a requirement).

Now the precise restore, and it is quick. Orbax pulls the ~7 GiB of mannequin and optimizer state again from Cloud Storage and pushes it out to the TPUs in 5.39 seconds — that is the complete wall-clock path, GCS learn plus the push to the gadgets. After a brief warm-up — the coaching operate re-entering and working its first step on the rebuilt mesh — the log prints accomplished step: 3301. That first step took 12.7 seconds, versus ~0.2 as soon as it is again at regular state. Complete time from kill to that line: about 1 minute 50 seconds.

This is the place that point went:

Under you’ve gotten logs you’d see within the Cloud logging view.

That final line is the entire story. The step counter goes 3388 to 3301 within the similar log stream: we misplaced 88 steps of progress, rewound to the final dedicated checkpoint at step 3300, and saved going.

Lastly, to show this was in-process restoration reasonably than Kubernetes quietly restarting all the things for us:

Zero restarts. Identical course of, similar PID, below two minutes of downtime, and nearly all of it was ready for Kubernetes to schedule the alternative employee. For comparability, a full job restart on this similar cluster pays for all the things we simply skipped: tearing down and rescheduling the complete workload — controller and all staff, not simply the failed slice — generally a full termination grace interval, then beginning recent containers and Python processes earlier than it even will get to the checkpoint load. (Compilation is a wash both manner — Pathways’ persistent cache means a restart reloads the compiled executables from Cloud Storage as a substitute of recompiling chilly.) And an actual preemption would add node reprovisioning on prime of that.

There’s an much more direct approach to persuade your self it is the identical course of: preserve a plain Python object on the controller, say an inventory that appends each step’s wall-clock time, and examine it after restoration. It nonetheless has each entry from earlier than the failure. That object lives in CPU reminiscence on the controller, not on any TPU, so an elastic occasion cannot contact it. If Kubernetes had restarted the pod, it might be empty.

Okay, that is the joyful path. Earlier than we wrap up, let’s speak about what to be careful for while you take this to your personal fashions.

Scaling past the demo

Our demo makes use of Qwen3-0.6B to maintain the restoration loop quick to observe and the associated fee low not as a result of bigger fashions do not work. They do, however the first time we constructed this demo on a bigger mannequin (Qwen3-4B) we ran right into a checkpoint-through-proxy bottleneck that is price understanding earlier than you scale up. This is what occurred, and find out how to keep away from it.

By default, when elastic_retry recovers from a failure, Orbax restores the checkpoint into host RAM on the CPU controller after which JAX pushes these arrays via the pathways-proxy container out to the TPUs. This controller-routed path works high quality when the checkpoint is small. However with a bigger mannequin, like Qwen3-4B the place params plus Adam optimizer moments add as much as about 135 GiB, the proxy has to buffer all of that in flight, and the proxy OOM-kills and restoration fails. In our case the proxy was set to 100 GB, the checkpoint was 135 GiB, and restoration would fail each time.

The repair is to not give the proxy extra reminiscence. A proxy OOM is a sign that you simply’re routing the entire checkpoint via the controller — which you do not need to be doing at scale within the first place. The fitting reply is to take the controller out of the info path: set ENABLE_PATHWAYS_PERSISTENCE=1 within the controller’s surroundings. This tells the Pathways Persistence API to have every TPU host learn and write its personal checkpoint shards on to Cloud Storage. The bytes by no means funnel via the controller’s proxy in any respect, so the proxy’s reminiscence restrict stops mattering — and also you get quicker, parallel checkpoint I/O as a bonus. That is the beneficial path and it really works as we speak. (You would as a substitute simply increase limits: {reminiscence: ...} on the proxy to clear the OOM, however that solely masks it: you retain the sluggish controller-routed path and miss the efficiency win, so deal with it as a stopgap, not a repair.)

Trying additional forward, MaxText is including Colocated Python, which runs the Orbax save and restore code on the TPU hosts themselves — combining direct-to-storage I/O with the pliability of working arbitrary Python on the accelerator VMs. That is nonetheless in preview: the bottom photographs want guide allowlisting as we speak, so it isn’t one thing you’ll be able to choose up out of the field but. We point out it right here so you already know the place issues are heading, not as a step to comply with.

Our demo sidesteps all of this as a result of qwen3-0.6b’s checkpoint is barely ~6.7 GiB, nicely below the proxy restrict. With persistence enabled, elastic restoration works the identical manner at a lot bigger checkpoint sizes.

What’s subsequent: snapshot-based elasticity

All the pieces on this submit recovers from a checkpoint on Cloud Storage. That is why the rewind is sure by your checkpoint interval. We misplaced 88 steps as a result of the final dedicated checkpoint was at step 3300 and why a bit of the restoration time goes into studying state again from GCS.

The following iteration of elasticity removes that round-trip. As a substitute of rebuilding from the newest GCS checkpoint, the elastic supervisor periodically saves a snapshot of the coaching state into host reminiscence, saved alongside the stay course of. When a slice drops, restoration reloads from that in-memory snapshot reasonably than from Cloud Storage. Two issues get higher: you rewind solely to the final snapshot (which may be much more frequent than a full checkpoint, because it by no means touches storage), and also you skip the GCS learn on the best way again. The identical equipment additionally handles resharding down onto the surviving slices and again up when replacements arrive with the replica-resize path, pushed from snapshots reasonably than checkpoints.

It is the identical thought we have walked via right here with a slice failure changing into a brief rewind, not a full restart and with the rewind made smaller and the restoration made quicker. That is the topic of the following submit, so preserve an eye fixed out.

Conclusion

On the scale and period of actual coaching runs, {hardware} failure is not an edge case. It is a certainty it’s important to plan for. The same old plan is to checkpoint usually and eat a full restart when a slice dies. Elastic coaching adjustments that: the failure turns into a catchable exception inside a course of that by no means exited, so restoration is a brief rewind on the identical controller as a substitute of a chilly relaunch of the entire workload.

That is what we walked via. A single employee force-killed mid-run, coaching again on the subsequent step in below two minutes, zero JobSet restarts utilizing pause and resume, the less complicated of the 2 restoration paths Pathways presents. Duplicate resize, the opposite path, retains coaching occurring the surviving slices as a substitute of ready, and is on the market as we speak as nicely.

You possibly can take a look at this workflow your self as we speak. All the mandatory code and step-by-step directions can be found in our elastic coaching information within the MaxText docs.

For now, thanks for studying! I hope this made you slightly interested in what’s attainable with JAX and MaxText on Cloud TPUs. If you happen to run the demo and it breaks, let me know on LinkedIn or X.

Comfortable coaching!