Resilient training with Pathways

This page explains how to configure resiliency features for Pathways workloads on Google Kubernetes Engine (GKE) so that large-scale training jobs can recover from planned interruptions and unplanned hardware failures.

Pathways provides resiliency benefits in the following ways:

  • Suspend-resume: Lets your workload tolerate planned interruptions, such as preemption notices, without requiring custom preemption-handling code.
  • Elastic training: Lets your workload tolerate unplanned hardware failures without crashing the client by using recovery logic in your training code.
  • Multi-Tier Checkpointing: Lets your workload save high-frequency local and peer-replicated checkpoints across TPU workers using colocated Python, with asynchronous Cloud Storage backups to reduce recovery time.
  • Hot-swap: Lets a higher-priority GKE JobSet workload quickly take over spare or lower-priority TPU slices when a hardware failure occurs.

Before you begin

Before you begin, complete the following tasks:

If you use elastic training, install the following packages in your workload container image:

  • pathwaysutils
  • Orbax (orbax-checkpoint): >= 0.12.6 (required for in-memory snapshots with ocp.v1.training.pathways.Snapshotter)
  • Grain (grain): required if you use grain.experimental.ElasticIterator

If you use Multi-Tier Checkpointing, also complete the following prerequisites:

  • Configure your GKE cluster for Multi-Tier Checkpointing by following Train large-scale machine learning models on GKE with Multi-Tier Checkpointing:
    • Enable the Multi-Tier Checkpointing and Cloud Storage FUSE (GcsFuseCsiDriver) add-ons, along with Workload Identity Federation for GKE, on your cluster.
    • Create a Cloud Storage bucket with hierarchical namespace enabled.
    • Grant the roles/storage.objectUser and roles/storage.insightsCollectorService Identity and Access Management (IAM) roles to the Multi-Tier Checkpointing Kubernetes service account.
    • Deploy a CheckpointConfiguration custom resource specifying your Cloud Storage bucket, TPU node selector, and in-memory volume size (inMemoryVolumeSize).
  • Build and push a colocated Python sidecar container image—for example, by using the sample sidecar Dockerfile. In your workload and sidecar images, use matching jax and jaxlib versions that align with your Pathways release and meet the following minimum version requirements:
    • Orbax version: >= 0.12.4
    • JAX version (including colocated Python): >= 0.10.1

Suspend-resume

GKE typically sends a preemption notice to an accelerator Pod before preempting the Pod. Pathways enables preemption tolerance by default on all cloud deployments, and Pathways accelerator jobs listen for these notices.

When a preemption notice arrives, Pathways first determines whether it can transparently save and restore the workload. If it can, Pathways attempts to transparently suspend your ML workload by writing out its current state to persistent storage, such as Cloud Storage, before GKE evicts your accelerator jobs. When GKE reschedules your jobs later, Pathways resumes your ML workload by reading back its persisted state.

If the workload isn't restorable, Pathways shuts down the accelerator job and forwards the failure to your workload if elastic training is configured. If elastic training isn't configured, GKE restarts the entire workload based on the JobSet restart policy.

Typical ML workloads defined using JAX rely on stateless Pathways XLA components, which are restorable using a high-bandwidth memory (HBM) snapshot. Certain ML workloads, such as those defined using the JAX colocated Python API, rely on stateful Pathways components and aren't restorable.

Elastic training

Elastic training lets your training job continue when hardware failures or slice preemptions occur, without restarting the entire workload. In a multi-slice training job, Pathways keeps the single-controller client process and the healthy TPU slices alive when an individual slice becomes unavailable.

Elastic training combines Pathways system capabilities with recovery logic in your training code:

  • Failure detection: When a hardware failure or preemption occurs on a TPU slice, Pathways unassigns the affected slice and raises a jax.errors.JaxRuntimeError the next time your workload accesses data on that slice. Because the client process remains active, your code can catch this error, reconfigure your resources, and continue training.
  • Recovery modes: The pathwaysutils.elastic.manager.Manager class and its elastic_retry decorator support two recovery strategies:
    • Pause and resume: If you set minimum_slice_count to None or to the total slice count, training pauses when a slice fails and resumes after all slices are healthy again. Because the slice count doesn't change across attempts, your workload can reuse its existing device mesh and compiled executables.
    • Replica resize: If you set minimum_slice_count to less than the total slice count, training continues on the remaining active slices when a slice fails by rebuilding the device mesh and recompiling for the smaller topology. A background monitor thread periodically checks inactive slices and records newly healthy slices in elastic_manager.available_inactive_slices. Your training loop can then rebuild the device mesh with the replacement slices, reinitialize for the expanded topology, and place the current state onto the new mesh.
  • State recovery: To resume training after an elastic event, your workload restores model parameters, optimizer state, and data iterator state using in-memory snapshots. The Snapshotter class in orbax.checkpoint asynchronously copies replicated state from TPU HBM to pinned_host CPU memory on each slice. During recovery, Snapshotter reconstructs the state from surviving host replicas and reshards it directly onto the active TPU mesh without reading from remote storage.

Enable elasticity on the Pathways proxy server

To enable elastic device sets on Pathways in GKE, add the --num_elastic_instances flag to the args field of the pathways-proxy container in your JobSet manifest and specify the number of elastic TPU slices in your workload:

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

Replace the following:

  • BUCKET_NAME: the Cloud Storage bucket used for Pathways temporary files.
  • NUM_SLICES: the number of TPU slices in your workload.

Filter active devices and build the device mesh

In a Pathways workload, jax.devices() returns all devices in the initial allocation for the lifetime of the client process. When you use elastic training, use elastic_manager.active_devices so that your device mesh and batch size only use healthy, active slices:

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

When the number of active slices changes in replica resize mode, use an elastic data loader such as grain.experimental.ElasticIterator. Standard per-host data iterators save state per host and can't restore onto a different number of hosts. In contrast, ElasticIterator tracks a single global dataset position that you can restore across any number of active slices.

Elastic training loop

An elastic training loop coordinates slice monitoring, failure recovery, and in-memory snapshotting, as shown in the following code example:

  • Your train function creates a single manager.Manager() instance and an Orbax Snapshotter outside the retry wrapper so slice tracking and host-memory snapshots persist across attempts.
  • Applying @elastic_manager.elastic_retry to run_training_loop configures the function to wait for healthy slices, set jax.default_device to an active slice, keep the background monitor thread active, and retry on a jax.errors.JaxRuntimeError.
  • At the start of a recovery attempt, your code must rebuild the device mesh from elastic_manager.active_devices and call snapshotter.load to restore the latest snapshot onto the updated mesh shardings. Let device buffers from the failed attempt go out of scope so that JAX can free their memory.
  • During each step of the loop, your code periodically calls snapshotter.save to back up the training state, and checks elastic_manager.available_inactive_slices to detect healthy replacement slices. When replacement slices are available, your code adds them to elastic_manager.active_slice_indices, rebuilds the mesh, reinitializes the training loop, and places the state onto the expanded mesh.

The following example shows how to implement an elastic training loop by combining manager.Manager, elastic_retry, and Orbax's in-memory Snapshotter:

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

Configure the elastic manager and retry behavior

You instantiate manager.Manager(devices=None) once per process. If you omit devices, the manager uses jax.devices() and groups devices by their slice_index attribute.

The Manager instance exposes the following properties to inspect cluster topology during training:

  • slice_to_devices: A mapping from each slice_index to the sequence of jax.Device objects belonging to that slice.
  • all_slice_indices: A frozenset of all slice indices in the allocation.
  • active_slice_indices: A frozenset of slice indices participating in the active computation mesh.
  • active_devices: A list of jax.Device objects belonging to the active slices, ordered by slice_index.
  • inactive_slice_indices: A frozenset of slice indices in the allocation that aren't participating in the computation mesh.
  • available_inactive_slices: A frozenset of inactive slice indices that the background monitor thread has detected as healthy and ready to rejoin.
  • total_slice_count and active_slice_count: The total number of allocated slices and the number of active slices.
  • default_device: Device 0 from an active slice, which elastic_retry automatically sets as the default JAX device for each attempt.

You can customize retry and scaling behavior by passing the following arguments to @elastic_manager.elastic_retry:

  • minimum_slice_count: The minimum number of healthy slices required before starting the initial execution or a retry attempt. If set to None (the default), elastic_retry waits for all slices in the allocation (total_slice_count), operating in pause and resume mode. If set to a smaller integer (such as 1), elastic_retry operates in replica resize mode.
  • poll_interval: The interval in seconds between slice health checks performed by the elastic training manager (Manager) during training. Defaults to 10.
  • timeout: The maximum number of seconds to wait for minimum_slice_count slices to become active before raising a TimeoutError. Defaults to None (no timeout).
  • pre_callback: An optional callback that runs at the beginning of each attempt after the manager updates active_slice_indices and before the wrapped training loop function runs.
  • on_elastic_event_callback: An optional callback that runs after elastic_retry catches a slice-down jax.errors.JaxRuntimeError, before the next retry attempt begins. You can use this callback to clean up incomplete checkpoints or reset data iterators.
  • retry_policy: A callable (attempt: int, error: Exception) -> bool that determines whether to continue retrying after an elastic event. Use manager.ElasticRetryLimit(max_attempts) to cap the total number of attempts. If omitted, elastic_retry retries indefinitely.

On an elastic event, JAX frees device buffers from the failed attempt once the stack frame unwinds and no variables reference those arrays. Avoid keeping references to device arrays from a failed attempt inside global variables or long-lived callbacks so that memory is released before the next attempt starts.

In-memory snapshots with Orbax

The ocp.v1.training.pathways.Snapshotter class (available in orbax-checkpoint version 0.12.6 and later) maintains an in-memory backup of your training state in pinned_host memory across slices:

  • Snapshotter(replica_axis_index=0): Initializes the snapshotter and specifies which mesh axis represents data-parallel slice replicas. All JAX arrays in the saved state must be replicated along this mesh axis so that surviving slices retain a complete copy of the model and optimizer state when a slice fails. If your state contains single-device arrays, replicate them across the mesh before calling save.
  • save(step, state): Asynchronously transfers state to pinned_host memory and splits it into per-replica trees along replica_axis_index. If a previous snapshot transfer is still in progress, save skips the new request to avoid stalling the training loop.
  • load(abstract_state, reset_snapshot_state=True): Probes the saved host replicas to identify surviving slices, concatenates the healthy replicas, and reshards the state onto the target shardings specified by abstract_state.
  • join(): Blocks until any queued snapshot operation completes.
  • latest: Returns metadata (including step) for the most recently completed snapshot, or None if no snapshot has completed.

Reduce sharding on slice failure

When a TPU slice fails or is preempted during execution, Pathways raises a jax.errors.JaxRuntimeError (such as DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND, or INTERNAL). In a function decorated with @elastic_manager.elastic_retry, elastic_retry catches this exception and checks elastic.is_error_due_to_slice_down(error):

  • If the error isn't caused by a slice-down event, elastic_retry re-raises the exception immediately.
  • If the error is caused by a slice-down event, elastic_retry stops any active JAX profiler trace, runs on_elastic_event_callback, and evaluates retry_policy(attempt, error).
  • Before starting the next attempt, elastic_retry calls elastic.wait_for_slices once to poll slice health until at least minimum_slice_count slices are active. It then updates elastic_manager.active_slice_indices and runs your training loop function on the surviving slices.

Increase sharding when slices rejoin

When your workload runs on a reduced number of slices in replica resize mode, the background monitor thread periodically runs health checks on elastic_manager.inactive_slice_indices. When repaired or rescheduled slices pass the health check, the monitor thread adds their slice indices to elastic_manager.available_inactive_slices.

When elastic_manager.available_inactive_slices is non-empty at the end of a training step, the workload scales back up in place without losing training progress:

  • Adding the slice indices in elastic_manager.available_inactive_slices to elastic_manager.active_slice_indices incorporates the rejoined slices into elastic_manager.active_devices and clears them from available_inactive_slices.
  • Rebuilding the device mesh from elastic_manager.active_devices and reinitializing the training step and data iterator adapts the computation to the expanded topology.
  • Calling jax.device_put places the current training state onto the new mesh shardings without restoring from a snapshot.

Multi-Tier Checkpointing

Large-scale training jobs experience frequent interruptions, and writing or restoring checkpoints directly to and from Cloud Storage can bottleneck training progress. Multi-Tier Checkpointing reduces checkpointing overhead and recovery time by coordinating three storage tiers on GKE:

  • Local RAM disk: TPU workers write high-frequency checkpoints to an in-memory CSI volume (multitier-checkpoint.csi.storage.gke.io) mounted on each worker node. This volume persists across Pod restarts on the same node.
  • In-cluster peer replication: The GKE replicator service replicates local checkpoint shards across peer TPU workers in the cluster.
  • Cloud Storage backup: The replicator service asynchronously backs up local checkpoints to a Cloud Storage bucket with hierarchical namespace enabled, without blocking TPU training steps.

Unlike multi-controller JAX, where each TPU host runs its own training script, Pathways runs a single controller on the pathways-head Pod and dispatches work to pathways-worker Pods. To read and write checkpoint shards directly on each TPU host's local RAM disk without routing checkpoint data through the single controller, Pathways uses colocated Python. Each worker Pod runs a colocated-python-sidecar container alongside the pathways-worker container, and both containers mount the Multi-Tier Checkpointing CSI volume and a shared memory directory (/tmp/sidecar). During initialization and training, Orbax dispatches checkpoint setup, saves, and restores from the Pathways controller to the worker sidecars on each TPU host. Multi-Tier Checkpointing on Pathways also supports elastic training when scaling down or scaling up the number of active TPU slices.

Pathways supports Multi-Tier Checkpointing with MaxText. For setup and configuration instructions, see the MaxText multi-tier checkpointing guide.

Configure your Pathways workload

You can configure a Pathways workload for Multi-Tier Checkpointing either by using Cluster Toolkit (gcluster job submit) or by configuring a custom Pathways JobSet manifest.

Use Cluster Toolkit

When you submit a job with gcluster job submit --pathways, pass the --gke-mtc-enabled, --gke-mtc-ramdisk-dir, and --pathways-colocated-python-sidecar-image flags:

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

Replace SIDECAR_IMAGE with the URI of your colocated Python sidecar container image.

When you set these flags, gcluster automatically configures the following:

  • Passes --sidecar_name=external to the pathways-proxy container.
  • Runs the colocated-python-sidecar container in each worker Pod and mounts a shared /tmp/sidecar directory (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) on both pathways-worker and colocated-python-sidecar.
  • Mounts the multitier-checkpoint.csi.storage.gke.io CSI volume at the path specified by --gke-mtc-ramdisk-dir in both pathways-worker and colocated-python-sidecar on the worker Pods, but not on the pathways-head Pod.

Use a custom Pathways JobSet manifest

If you manage your own Pathways JobSet manifest instead of using gcluster, pass --sidecar_name=external to pathways-proxy. In your worker Pod specification, mount both the shared sidecar memory volume and the multitier-checkpoint.csi.storage.gke.io CSI volume on the pathways-worker and colocated-python-sidecar containers:

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Use Orbax Multi-Tier Checkpointing in a custom JAX workload

In a custom JAX training script running on Pathways, use the orbax.checkpoint.experimental.emergency.multi_tier_checkpointing package and set use_colocated_python=True in both initialize_multi_tier_checkpointing and ReplicatorCheckpointManagerOptions:

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

Hot-swap

Hot-swap is a feature of the GKE JobSet API that lets a higher-priority job quickly take over resources from a lower-priority job, minimizing downtime and speeding up recovery.

When you create a JobSet, GKE schedules the workload across multiple slices, as specified in the JobSet configuration. If a hardware failure occurs on one or more slices, GKE marks the affected Pods as failed. If your GKE cluster has a spare slice running a lower-priority job, the JobSet controller remaps the failed slice's workload from the higher-priority job onto that spare slice. This remapping typically takes less than a minute.

When a JobSet restarts, hot-swap can occur in the following situations:

  • Default mode: If spare, idle TPU slices are available within the same cluster, the Kubernetes scheduler prioritizes scheduling the restarted jobs onto these slices rather than waiting for the failed slices to be repaired. This scheduling behavior provides faster recovery.
  • Heterogeneous workloads: In clusters running multiple workloads with a configured Kubernetes PriorityClass resource, a restarted JobSet can trigger a hot-swap. If the restarted job's affinity matches a lower-priority job's resources, Kubernetes preempts the lower-priority job so that the higher-priority job can start without waiting for slice repair. For example, you can configure your Pathways worker Pods with different priorities by using a PriorityClass resource.

To use priorities in your cluster, define a PriorityClass manifest and save it as high-prior-job.yaml:

apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

Apply this YAML to your GKE cluster:

kubectl apply -f high-prior-job.yaml

Next, attach the new priority class to your Pathways workload. If you submit your workload by using Cluster Toolkit, pass --priority=high-prior-job to the gcluster job submit command, which sets priorityClassName on both the pathways-head and worker Pods. If you write your own JobSet manifest, set priorityClassName in the Pod template of the worker replicated job, which runs the pathways-worker container:

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

If you created your cluster by using a Cluster Toolkit TPU blueprint with Pathways enabled, the cluster already includes the very-low, low, medium, high, and very-high priority classes, so you can pass one of them to --priority instead of creating your own.

What's next