Live GPU Migration for Resilient Distributed Training

TorchPass is an infrastructure software solution for live GPU migration. It sits between the training framework (PyTorch) and the cluster scheduler. It automatically detects when a node is experiencing problems or has failed, provisions a healthy replacement, and migrates the training state from the problematic node to the replacement, minimizing restore and restart time and eliminating lost training progress. From a training job's perspective, there is just a short pause after which training resumes with no lost progress.

Common errors in large AI training clusters

  • GPU memory errors
  • GPUs falling off the bus
  • NVLink errors
  • Network errors
  • Driver and firmware bugs

TorchPass is particularly well suited to NVL72 rack-scale systems where every GPU is topologically equivalent to every other GPU in the rack. When TorchPass migrates a training worker from a failing GPU or from a failed tray (node) to a spare in the same rack, there is no topology penalty. Operators are already adopting a practice of setting aside 1 or 2 trays as spares that TorchPass can leverage to enable continuous training.

Without TorchPass, the state of the art is to restart the entire training job from the most recent persistent checkpoint - losing GPU hours and increasing job completion time due to restart, restore and recompute. TorchPass enables a radically different paradigm - live GPU migration to a spare node, often already in place and waiting, delivering a dramatically different SLA!

TorchPass Mechanisms

TorchPass provides 3 migration mechanisms:

  1. Planned Migration (model-aware): Offers the fastest migration path for planned events. Since it uses the model's own checkpoint save and load functions, minor code changes are required to register the functions with TorchPass at training startup.

  2. Planned Migration (model-transparent): Offers a migration path for planned events that does not require any model code changes. Rather than using the model's checkpoint load and save functions, it uses CRIU (Checkpoint Restore in User Space) and NVIDIA's CUDA-checkpoint technology to capture the complete execution state of the training worker process and restore it on a different node.

  3. Unplanned Migration (model-aware): Offers a migration path for sudden failures where a node, GPU or other component has already crashed and the affected workers can no longer be checkpointed. Instead, JIT (just-in-time) checkpoint recovery reconstructs the lost state surgically from surviving data-parallel peers and places it on a spare resource. It requires only minor code changes since it leverages the model's checkpoint save and load functions.

Whether it's planned or unplanned, the training application is completely unaware that a migration has occurred and, after a brief pause during migration, training resumes at the exact same iteration step. Note that no code changes or library imports are required for model-transparent migration, making it well-suited to cloud operators that serve diverse tenant workloads and cannot mandate changes to customer training code.

By contrast, model-aware migration requires only very small, well-bounded changes to training code - importing the TorchPass Python library and registering existing checkpoint save/load functions. The remainder of this document provides a technical overview and practical considerations for implementing TorchPass in a production environment.

TorchPass Control Plane Components

TorchPass is designed as a minimal-footprint addition to existing cluster infrastructure. The same control plane is installed regardless of which migration mode is used. What changes between migration mechanisms is how state is captured and restored, not what gets installed in the cluster. The TorchPass binaries are built for both X86 and ARM, so the same images can be deployed on x86 hosts and Grace ARM CPUs in NVL72 racks.

Active Training Nodes

  1. Orchestrator tracks spares and jobs in progress through a scheduler plugin.
  2. Scheduler Plugin updates the main scheduler's job definition to reflect the swap.

The control plane has the following components:

TorchPass Orchestrator

The orchestrator is the central coordination service for TorchPass. It runs as a single long-lived process and is responsible for tracking all active training jobs, monitoring node health through the scheduler plugin, deciding when migrations should occur and coordinating the migration process across all affected nodes.

Scheduler Plugin

The scheduler plugin is a thin adapter layer between TorchPass and the cluster's native scheduling infrastructure. Its job is to translate between TorchPass concepts (spare nodes, taint events, migration decisions) and the scheduler's native mechanisms (Kubernetes pod lifecycle, Slurm job management). It reports node health status, responds to spare provisioning requests, and notifies the orchestrator when jobs start or exit. It injects the Launchpad, a torchrun replacement, as a per-node agent to manage PyTorch worker processes and communicate with the TorchPass orchestrator.

Launchpad

Launchpad runs as the per-node parent of the PyTorch worker processes. It is transparently swapped in by the scheduler plugin as a replacement for torchrun. Launchpad supervises the local workers, sets up their environment, and stays in contact with the Orchestrator throughout the job.

The Migration Process

Part I: Planned Migration

Planned migration covers two use cases. The first is preemptive where the node is still operational, but telemetry indicates a failure is imminent (e.g. rising correctable-ECC rates, climbing temperatures, fan slowdowns, accumulating NVLink errors, etc.). The second is proactive where the operator wants to vacate a node for maintenance such as firmware updates, security patching, workload rebalancing, or removing a known straggler.

How Planned Migration Works

  1. Detections: Hardware degradation in GPU clusters manifests through a variety of signals: rising ECC error counts in GPU memory, a GPU becoming unresponsive to the PCIe bus, temperatures exceeding safe thresholds, etc.
  2. Scheduler Plugin Reaction: The TorchPass scheduler plugin continuously polls the cluster state. When it detects that a node participating in an active training job has been tainted or drained, it reports this event to the orchestrator.
  3. Orchestrator Decision: The orchestrator receives the taint report and determines the appropriate response. It identifies which workers on the tainted node need to be moved and calculates how many spare nodes are needed.
  4. Coordinated Pause of All Workers: TorchPass ensures all workers in the training job must be paused to achieve a globally consistent point before capturing state.
  5. State Capture: In model-transparent planned migration, the orchestrator captures state from outside the process through mechanisms like CRIU.
  6. State Transfer to the Spare Node: The captured state must now be moved from the leaver node to the healthy spare node using Launchpad's wormhole state-transfer engine.
  7. State Restoration: On the spare node, its Launchpad agent receives the model state dictionary and reconstructs model parameters and other state variables.

Part II: Unplanned Migration

Unplanned migration covers cases where failures have already occurred. The orchestrator usually detects these failures chiefly through the Launchpad's heartbeat and through scheduler-plugin events.

How Model-aware Unplanned Migration Works

  1. Detection: Hard failures present differently from taintable conditions. The orchestrator detects these failures within seconds.
  2. Orchestrator Decision: It identifies the missing rank(s) and determines whether model-aware unplanned migration is feasible.
  3. Coordinated Pause of Surviving Workers: Surviving workers need to be coordinating paused to reconstruct the missing rank's state.
  4. Just-in-time Checkpoint Capture from Surviving Peers: The surviving workers take a JIT checkpoint of their state by invoking the registered save function.
  5. State Reconstruction on the Spare: The joining worker receives the relevant slices from the surviving peers and invokes the load function on the spare.
  6. NCCL Repair and Resume: All workers are instructed to restart collective operations and training resumes at the exact iteration step where the failure occurred, with no recompute of completed work.

The key insight is that individual GPU failures are the most common failure mode, followed by tray-level failures. With 18 compute trays in the rack, designating 2 trays as spares yields 16 active trays and 2 spare trays. There is opportunity to reduce the number of spare trays when using TorchPass.

Summary

TorchPass eliminates one of the largest sources of wasted GPU-hours in distributed training by offering three migration mechanisms - planned migration (model-transparent/zero code changes), planned migration (model-aware), and unplanned migration - covering the full spectrum of failure events. The result is a fundamentally improved SLA: training that doesn't stop when hardware misbehaves, with recovery times measured in seconds rather than the hours typical of full job restart.