Distributed TPU training gives you more compute by spreading one training job across multiple TPU hosts. The problem is that the training job also becomes dependent on those hosts working together.
If one GKE worker disappears, the training process usually cannot just continue with the remaining workers. Distributed JAX programs rely on collective operations, communication between processes, and a known device topology. A missing worker can leave other workers waiting for communication that will never complete.
For production training, the goal should therefore not be to assume that a worker will never fail. The better approach is to design the training job so that a failure becomes a recoverable event: detect the failure, stop the broken training attempt, recreate the required TPU resources, restore the latest valid checkpoint, and continue from a known training state.
Google Cloud also distinguishes between individual TPU instance recovery and multi-host TPU slice recovery. For multi-host slices, the instances in the slice must be recovered or rescheduled together because the TPU network topology is created across the slice.
What Does It Mean When a GKE Worker Goes Down?
A worker going down can mean several different things.
The Kubernetes Pod may terminate. The underlying GKE node may become unavailable. A TPU host may experience a hardware failure. A network problem may prevent one process from communicating with the others. The workload may be preempted or restarted. The process itself may crash because of an application error.
These failures look different from the infrastructure perspective, but the training program often sees the same basic problem: one participant in a distributed computation is no longer available.
For example, imagine a training job using four TPU workers:
Worker 0 ──┐
Worker 1 ──┤
Worker 2 ──┼── Distributed training
Worker 3 ──┘
If Worker 2 stops responding, Workers 0, 1, and 3 may still be running, but that does not mean the training step can safely continue.
The next collective operation may require all four participants. Instead of producing a partially updated model, the distributed runtime may fail, hang, or report a communication/runtime error.
This is why worker failure should be treated as a distributed job failure rather than simply a normal Kubernetes Pod restart.
Why Restarting Only One Worker Usually Is Not Enough
A common assumption is:
Worker 2 crashes
↓
Kubernetes restarts Worker 2
↓
Training continues
That model is usually too simplistic for a tightly coupled multi-host TPU training job.
The other workers may already have moved to a different execution point, lost synchronization, or terminated because a collective operation failed. Even if Kubernetes successfully starts a replacement Pod, the new process does not automatically know the exact model state, optimizer state, random state, dataset position, and distributed execution state that the other workers were using.
There is another issue at the TPU infrastructure level. Multi-host TPU slices have a fixed physical topology. Google Cloud documents that a failed instance in a multi-host slice cannot simply be rescheduled independently; the slice must be recovered or rescheduled together so that the required topology can be recreated.
The practical recovery model is therefore closer to this:
Worker failure
↓
Detect distributed job failure
↓
Stop the failed training attempt
↓
Recover/recreate the TPU slice
↓
Start all training workers again
↓
Load the latest valid checkpoint
↓
Restore model + optimizer + training state
↓
Continue training
This is much safer than trying to patch one worker back into a running distributed process.
Checkpointing Is the Main Recovery Mechanism
Checkpointing is what turns a failed training job from a complete restart into a recoverable job.
A checkpoint should contain enough state to continue training correctly. Depending on the training framework, that can include:
Model parameters
Optimizer state
Current training step
Learning-rate scheduler state
Random-number generator state
Dataset or input position where practical
Mixed-precision or gradient-scaler state
Training configuration or metadata
Checkpoint version information
Saving only the model weights is often insufficient.
For example, suppose training reaches step 100,000 and the worker fails at step 100,350. If the checkpoint contains only model weights from step 100,000 but not optimizer state, restarting from those weights can produce different training behavior.
The recovery target should therefore be a complete training state, not simply a copy of the model.
Google recommends frequent checkpointing for TPU workloads because work performed since the most recent checkpoint can be lost when a workload is interrupted.
Save Checkpoints to Durable Storage
Do not make the only copy of a checkpoint live on the failed TPU worker.
A local checkpoint can disappear together with the worker or TPU slice. For production training, checkpoints should be stored somewhere that survives worker replacement.
Cloud Storage is a common choice:
TPU Workers
│
│ checkpoint
▼
Local checkpoint
│
│ backup
▼
Cloud Storage
The training job should be able to start on new workers and retrieve the last known good checkpoint.
This is particularly important when the recovery process recreates the entire multi-host TPU slice.
For large-scale GKE training, Google Cloud provides Multi-Tier Checkpointing for JAX workloads using Orbax and other supported training workloads. It can maintain local and peer copies while backing up checkpoints to Cloud Storage, reducing the amount of work that must be restored from remote storage after a failure.
Do Not Assume the Latest Checkpoint Is Valid
A checkpoint must be treated as a transaction.
Consider this sequence:
Worker 0 ── writes checkpoint
Worker 1 ── writes checkpoint
Worker 2 ── fails
Worker 3 ── writes checkpoint
If the checkpoint operation was not coordinated, the resulting files may represent different training states.
The next job startup could then restore an inconsistent state.
A safer approach is:
Train for a defined number of steps.
Synchronize the workers.
Create the checkpoint from a consistent global state.
Complete the checkpoint write.
Record the checkpoint as valid.
Continue training.
The recovery process should select a checkpoint that is known to be complete.
For example, you might maintain:
checkpoint/
step-100000/
state
metadata
COMPLETE
The COMPLETE marker is created only after the complete checkpoint has been successfully written.
A recovery process should ignore incomplete checkpoint directories.
Use a Single Source of Truth for the Training Step
A distributed training job needs a clear definition of where it is in training.
Suppose one worker believes the job is at step 101,000 while another believes it is at step 100,999.
That discrepancy can create problems during recovery.
The checkpoint should therefore contain the global training step:
state = {
"step": 100000,
"model": model_state,
"optimizer": optimizer_state,
}
When the job restarts, every worker should restore the same logical training state.
Do not calculate the recovery step independently on each worker from local counters.
The worker's process index identifies the worker. It should not determine the training progress.
Let Kubernetes Restart the Job, Not the Training State
Kubernetes is good at restarting containers and scheduling workloads.
It does not automatically understand the semantic state of your machine-learning training loop.
For example:
restartPolicy: Never
can be appropriate for a Job where the failure should cause the overall Job attempt to fail and an external controller should create the next attempt.
The exact restart strategy depends on the controller you use.
The important distinction is:
Kubernetes recovery
=
restart infrastructure/processes
Training recovery
=
restore distributed application state
You need both.
A Pod restart without checkpoint restoration can simply start training from the beginning.
A checkpoint without a controller that recreates the required workers leaves the application with no infrastructure on which to resume.
Detect Worker Failure Quickly
Recovery time depends partly on how quickly the system identifies the failure.
Useful signals include:
Pod termination
Node NotReady state
Container exit codes
TPU runtime errors
JAX runtime errors
Collective communication failures
Training process heartbeat
Job controller failure status
TPU slice health
Missing worker logs
For example, your monitoring system might detect:
Expected workers: 4
Active workers: 3
Training state: FAILED
That should trigger the recovery workflow instead of allowing the remaining workers to sit indefinitely.
A distributed training job should also log the worker identity:
import jax
print("process index:", jax.process_index())
print("process count:", jax.process_count())
print("local devices:", jax.local_device_count())
print("global devices:", jax.device_count())
These values become extremely useful when diagnosing failures.
If the logs only say "training failed," it can be difficult to determine whether the problem affected one worker, the entire TPU slice, or the application itself.
Avoid Infinite Restart Loops
Automatic recovery can introduce another production problem.
Imagine this sequence:
Training fails
↓
Restart
↓
Training fails
↓
Restart
↓
Training fails
↓
Restart forever
This can consume TPU capacity and create a large bill without making progress.
Use a bounded retry policy.
For example:
Attempt 1 → failure → retry
Attempt 2 → failure → retry
Attempt 3 → failure → retry
Attempt 4 → failure → stop and alert
The retry counter should be associated with the training attempt, not with an individual container restart.
You should also record the reason for each failure.
A transient TPU infrastructure failure and a deterministic Python exception should not necessarily receive the same recovery policy.
Separate Infrastructure Failures From Application Failures
Not every training failure should trigger the same recovery process.
Consider these two examples.
Infrastructure failure:
TPU worker becomes unavailable
↓
Distributed communication fails
↓
Recover TPU resources
↓
Restore checkpoint
↓
Resume training
Application failure:
Python exception
↓
Same code fails again
↓
Restart
↓
Same exception
The second case should usually stop after a small number of attempts and alert the team.
Otherwise, the system can repeatedly recreate expensive TPU resources while executing the same broken code.
The recovery controller should therefore capture the failure reason and classify it before deciding whether to retry.
Preserve the Training Configuration
A checkpoint from one training configuration should not automatically be loaded into a different distributed configuration.
For example, changing these values can affect recovery:
A useful checkpoint directory can contain metadata such as:
{
"step": 100000,
"model_version": "v12",
"dataset_version": "2026-11-10",
"training_config": "large-tpu-v3",
"code_version": "git-sha-example"
}
Before restoring the checkpoint, the new job should validate that the environment is compatible.
This prevents a subtle class of errors where recovery succeeds technically but training continues with a different configuration.
Handle Dataset Position Carefully
Restoring the model state is only part of the problem.
The input pipeline also matters.
Suppose the training job consumed records through batch 50,000 before failing. After restarting, the new workers may begin reading the dataset from the beginning.
That does not always make the training mathematically invalid, but it changes the training trajectory.
For some workloads, repeating data is acceptable. For others, maintaining a deterministic data position is important.
A production training design should explicitly decide whether recovery means:
Restore model state and continue with a new data sequence
or:
Restore model state and continue from the same logical dataset position
The decision should be documented rather than left to accidental behavior in the input pipeline.
Keep Worker Startup Deterministic
After a TPU slice is recreated, all workers should start from the same expected configuration.
For example, each worker should know:
process index
process count
TPU topology
checkpoint location
dataset configuration
model configuration
training version
A worker that starts with a different configuration can make the entire distributed job fail.
Startup validation is therefore useful.
You can perform checks before entering the training loop:
import jax
expected_processes = 4
if jax.process_count() != expected_processes:
raise RuntimeError(
f"Expected {expected_processes} processes, "
f"but found {jax.process_count()}"
)
print("Process:", jax.process_index())
print("Processes:", jax.process_count())
print("Devices:", jax.device_count())
This is better than discovering the mismatch several minutes into training.
Use the Job Controller as the Recovery Boundary
For a tightly coupled multi-host training job, the distributed Job should generally be the recovery boundary.
Conceptually:
GKE
│
├── Worker 0
├── Worker 1
├── Worker 2
└── Worker 3
│
▼
Distributed training job
│
▼
Checkpoint in Cloud Storage
If Worker 2 fails:
Worker failure
↓
Job becomes unhealthy
↓
Stop the failed distributed attempt
↓
Recreate required workers/resources
↓
Start the training program
↓
Restore checkpoint
↓
Resume
This approach avoids trying to repair distributed state inside an already-broken process group.
For some modern TPU workloads, however, there are newer resiliency mechanisms that can avoid restarting the entire workload. Google Cloud Pathways supports elastic training for certain workloads, allowing recovery logic to reconfigure resources after a slice failure. This is workload- and runtime-specific, so it should not be assumed to apply to every JAX-on-GKE training job.
Consider Multi-Tier Checkpointing for Long Jobs
For short training jobs, periodically writing a checkpoint to Cloud Storage may be sufficient.
For large, long-running jobs, checkpoint recovery can itself become a bottleneck.
If restoring hundreds of gigabytes takes several minutes, the TPU workers sit idle while the checkpoint is downloaded.
Multi-Tier Checkpointing addresses this by keeping checkpoint copies closer to the training workers and asynchronously backing them up to Cloud Storage. Google documents support for JAX training on TPUs using Orbax.
The basic idea is:
Training workers
│
├── Local checkpoint
│
├── Peer checkpoint
│
└── Cloud Storage backup
During recovery, the system can prefer an available complete checkpoint closer to the workers.
This becomes particularly useful when failures are frequent or checkpoints are large.
Test Worker Failure Before Production
A recovery system that has never been tested is not a recovery system.
Do not wait for a real TPU failure to discover that your checkpoint cannot be restored.
Create failure tests such as:
Test 1: Kill one training Pod
Test 2: Restart one worker process
Test 3: Interrupt the training Job
Test 4: Simulate a network failure
Test 5: Delete the training Job and recreate it
Test 6: Restore the latest checkpoint
Test 7: Restore an older checkpoint
Test 8: Corrupt an incomplete checkpoint
For each test, measure:
Time to detect failure
Time to stop the failed job
Time to provision replacement TPU resources
Time to start all workers
Time to restore checkpoint
Time to recompile the workload
Time to resume training
Number of training steps lost
Whether the recovered model matches expectations
Google specifically recommends checkpointing and resuming from the last checkpoint for TPU workloads because infrastructure interruptions can still occur even when mitigation features are enabled.
Account for JAX Compilation Time
Recovery may take longer than expected because the restarted training process may need to compile again.
The recovery timeline can therefore look like:
Failure detection
↓
TPU recovery
↓
Worker startup
↓
Container initialization
↓
Checkpoint restore
↓
JAX/XLA compilation
↓
Training resumes
If compilation takes several minutes, restarting the training process for every small interruption can become expensive.
Where appropriate, use JAX compilation caching and preloaded artifacts to reduce startup time. Google Cloud's AI/ML performance guidance specifically recommends compilation caching and faster artifact loading for interrupted TPU workloads.
Monitor Recovery, Not Just Training
A dashboard showing training loss and throughput is not enough.
Production monitoring should also track:
Training step
Checkpoint step
Checkpoint age
Checkpoint duration
Active workers
Expected workers
Training job status
Restart count
Last failure reason
Recovery duration
Time since last successful step
One particularly useful metric is checkpoint age.
If the latest successful checkpoint is 30 minutes old and the training job fails, you already know approximately how much progress may be lost.
You can also alert when:
current_step - checkpoint_step
becomes larger than your accepted recovery window.
Common Mistakes
Restarting only the failed Pod
A new Pod does not automatically rejoin an existing distributed process group with the same application state.
Treat the distributed training attempt as a coordinated unit unless your framework explicitly supports elastic recovery.
Saving only model weights
Optimizer state, scheduler state, training step, and other state can be required for correct continuation.
Save the complete training state.
Keeping checkpoints only on local storage
The worker that owns the checkpoint can be the same worker that fails.
Use durable storage and, for large workloads, consider a multi-tier checkpoint strategy.
Assuming every TPU failure is recoverable by the same mechanism
A Kubernetes Pod failure, TPU hardware failure, network partition, and application exception have different recovery paths.
Classify the failure before applying retry logic.
Restarting forever
Repeated failures can waste large amounts of TPU capacity.
Use bounded retries and alert after the retry budget is exhausted.
Restoring an incompatible checkpoint
Changing the model, dataset, sharding, or runtime configuration can make an old checkpoint unsuitable.
Store version metadata with the checkpoint and validate it during startup.
Testing only graceful shutdown
A graceful shutdown may allow the application to save state.
A hardware or network failure may not provide that opportunity.
Test abrupt failures as well.
A Practical Recovery Architecture
A production design can be organized into five layers.
1. Training layer
The JAX application periodically creates complete checkpoints and records the global training step.
2. Storage layer
Checkpoints are stored durably, with incomplete checkpoints excluded from recovery.
3. Orchestration layer
GKE and the training controller detect failed workers and recreate the required distributed training resources.
4. Recovery layer
A new training attempt restores the latest known-good checkpoint and validates its compatibility.
5. Monitoring layer
Metrics and alerts track worker health, checkpoint age, restart count, and recovery duration.
The resulting workflow looks like this:
┌──────────────────────┐
│ GKE Training Job │
└──────────┬───────────┘
│
┌─────────────┴─────────────┐
│ TPU Worker Group │
│ W0 W1 W2 W3 │
└─────────────┬─────────────┘
│
Training State
│
▼
┌──────────────────────┐
│ Checkpoint Manager │
└──────────┬───────────┘
│
▼
┌──────────────────────┐
│ Cloud Storage │
│ Last Known Good State │
└──────────┬───────────┘
│
Worker Failure
│
▼
┌──────────────────────┐
│ Detect + Stop Attempt │
└──────────┬───────────┘
│
▼
┌──────────────────────┐
│ Recover TPU Resources │
└──────────┬───────────┘
│
▼
┌──────────────────────┐
│ Start All Workers │
└──────────┬───────────┘
│
▼
┌──────────────────────┐
│ Restore Checkpoint │
└──────────┬───────────┘
│
▼
┌──────────────────────┐
│ Resume Training │
└──────────────────────┘
Production Checklist
Before running a long TPU training job on GKE, verify the following:
Every worker participates in the same distributed configuration.
Worker failures are detected automatically.
The training process can exit cleanly when the distributed job becomes invalid.
Checkpoints contain complete training state.
Checkpoints are stored outside the individual TPU worker.
Incomplete checkpoints cannot be selected during recovery.
The global training step is stored in the checkpoint.
Model and optimizer state are restored together.
Dataset recovery behavior is documented.
Training configuration versions are recorded.
The recovery process recreates the required worker topology.
Retry attempts are limited.
Application failures are distinguished from infrastructure failures.
JAX compilation and startup time are measured.
Checkpoint age is monitored.
Recovery time is measured.
Worker failure has been tested deliberately.
Recovery from an abrupt failure has been tested.
The team has a runbook for failed training jobs.
Advantages and Disadvantages
The checkpoint-and-restart approach is relatively straightforward and works well for tightly coupled distributed training.
Advantages include:
Simple recovery model
Clear training state
Easy operational boundaries
Works with long-running training jobs
Allows infrastructure to be recreated
Limits the amount of lost training work
There are also costs:
Checkpoint storage consumes space
Frequent checkpoints can affect performance
Recovery requires restarting the distributed training attempt
TPU provisioning can take time
Large checkpoints can take time to restore
Dataset position may be difficult to reproduce exactly
Some failures can still require manual investigation
More advanced systems can reduce some of these limitations. For example, Google Cloud's current TPU resiliency tooling includes elastic training for supported workloads, where recovery can happen without restarting the entire client workload. That capability depends on the runtime and workload architecture, so it should be evaluated separately rather than treated as a universal replacement for checkpoint-based recovery.
Summary
A worker failure in a multi-host TPU training job is not the same as restarting an ordinary Kubernetes Pod. Distributed training depends on coordinated workers, collective communication, and a consistent TPU topology. When one participant disappears, the safest recovery strategy is usually to treat the distributed training attempt as failed and recreate the required workers together.
Checkpointing is the foundation of that recovery process. Save complete training state to durable storage, record the global training step, verify checkpoint consistency, and restore the latest known-good state when the new training attempt starts.
GKE provides the infrastructure needed to schedule and manage TPU workloads, while the training application and orchestration layer must handle application-level state recovery. For multi-host TPU slices, the underlying TPU infrastructure also has specific recovery behavior because the hosts share a fixed topology.
For long-running production training, the practical target is not zero failures. It is predictable recovery: detect the failure quickly, avoid wasting TPU capacity, restore a valid checkpoint, recreate the distributed environment, and continue training with as little lost work as possible.
Join the conversation! Your thoughts help the community grow.