使用NVRx在Amazon EKS上进行容忍分布式训练

Large-scale distributed training jobs run for hours or days across dozens of nodes. At that scale and duration, interruptions are statistically inevitable: network partitions, memory errors, software exceptions, or infrastructure events will eventually disrupt at least one worker. A single GPU fault triggers a cascade: NVIDIA Collective Communication Library (NCCL) timeouts propagate to healthy workers, pods crash and restart out of sync, and your cluster burns expensive GPU hours while making zero training progress. Synchronous checkpointing adds a second source of idle time: every save blocks all ranks on I/O, which on the cluster sizes in this post consumed up to 40% of total wall time.

In this post, we show how to integrate NVIDIA Resiliency Extension (NVRx) into PyTorch Fully Sharded Data Parallel (FSDP) training on Amazon Elastic Kubernetes Service (Amazon EKS) to solve both problems. You walk through async checkpointing that overlaps I/O with training, in-process restart that recovers from faults in seconds without touching the container lifecycle, and in-job restart using ft_launcher for automatic worker respawn on hard crashes. We include benchmark results on H100 GPUs at 2-node to 8-node scale, with all code available to reproduce.

Solution overview

The solution combines NVRx fault tolerance primitives with an EKS-based training environment designed for high-performance multi-node GPU workloads. NVRx handles the application-level resilience (async checkpointing, in-process restart, and in-job restart), while the EKS cluster provides the infrastructure foundation: GPU scheduling, high-bandwidth networking, and shared storage for checkpoint persistence.

NVRx

The NVIDIA Resiliency Extension (NVRx) is a pip-installable Python layer (pip install nvidia-resiliency-ext) that adds fault-tolerance primitives to PyTorch: no custom kernels, no PyTorch fork, no recompile. The primitives drop into an existing FSDP script as ordinary imports. The model and training code stay untouched. Each is independently adoptable. We exercised three features: async checkpointing, in-process restart, and ft_launcher (in-job restart).

Async checkpointing, exposed through TorchAsyncCheckpoint, replaces torch.save with an async_save() call that hands the state dict to a background process and returns immediately. A matching finalize_async_save() before the next save commits the prior write. Paired with FSDP LOCAL_STATE_DICT, each rank writes its own shard directly, with no all-gather and no rank-0 bottleneck.

In-process restart, exposed through inprocess.Wrapper, wraps the train function so a transient fault (an unhandled exception or an NCCL hang) does not kill the Python process. NVRx aborts the active process group, runs health checks per rank (GPU, NVLink, NIC), re-rendezvouses survivors, and re-enters the wrapped function from the latest checkpoint. The interpreter, CUDA allocator, and outer-scope objects survive. This catches the soft-fault class.

The ft_launcher binary, the NVRx in-job restart launcher, handles cases in-process cannot catch: SIGKILL, out-of-memory (OOM) kill, and OS-level hangs. Each rank runs a RankMonitorClient. The launcher checks heartbeats against explicit CLI-set timeouts, and on stall or death it kills survivors, reclaims GPU memory, and respawns fresh workers in the same job. Recovered workers reload from the latest checkpoint. Each recovery layer covers a distinct fault class: in-process for soft faults, ft_launcher for hard faults, and the cluster orchestrator for node loss. The layers are independent. Pick the one whose scope matches your failure modes.

Amazon EKS cluster

Amazon EKS is a managed Kubernetes service that handles the control plane, upgrades, and API server availability. We run self-managed node groups of p5.48xlarge instances, each with 8 NVIDIA H100 80 GB GPUs and 32 Elastic Fabric Adapter (EFA) network interfaces. Training pods run as Kubernetes Jobs with headless Services for peer discovery, so workers find each other through DNS rather than hardcoded IPs, and pod replacements can rejoin without reconfiguring the job.

Each node exposes its GPUs and EFA adapters as extended resources through the NVIDIA device plugin and EFA device plugin. The Kubernetes scheduler places training pods on GPU nodes using node affinity and tolerations, facilitating the full 8-GPU allocation per node.

For checkpoint storage, we use Amazon FSx for Lustre (SCRATCH_2, 1.2 TB) mounted into every training pod through the FSx CSI driver. FSx provides the shared filesystem that both async and synchronous checkpointing write to, and critically, it is where recovering workers read their checkpoint state after a fault. Placing FSx in the same Availability Zone as the GPU nodes minimizes read latency during recovery, which matters because checkpoint loading (not the restart mechanism) dominates recovery time at scale.

Figure 1: Architecture diagram showing the EKS cluster with 2-8 p5 nodes, EFA interconnect, FSx for Lustre, and NVRx components within training pods

The key AWS services involved:

• Amazon EKS — Kubernetes control plane, pod scheduling, Job lifecycle management.
• Amazon Elastic Compute Cloud (Amazon EC2) p5.48xlarge — 8x H100 80 GB GPUs, 32x EFA adapters per node.
• Elastic Fabric Adapter (EFA) — 3,200 Gbps network bandwidth for NCCL all-reduce operations.
• Amazon FSx for Lustre — Shared POSIX filesystem for distributed checkpoint I/O.
• Amazon Elastic Container Registry (Amazon ECR) — Container registry for the training image (PyTorch + NVRx + model code)

Prerequisites

Before deploying this solution, make sure you have the following infrastructure and tooling in place:

• AWS account with service quota for p5.48xlarge (or p4de.24xlarge) instances.
• Amazon EKS cluster (v1.28+) with EFA-enabled self-managed GPU node groups and the NVIDIA device plugin installed.
• Amazon FSx for Lustre filesystem (SCRATCH_2) in the same Availability Zone as the GPU nodes.
• Container image with PyTorch 2.9+, NVRx 0.4.1 to reproduce the benchmark results shown in this post. Use v0.6.0 with an updated launcher configuration for a current deployment and your training code, pushed to Amazon ECR.
• kubectl configured for your cluster.
• HuggingFace account with access to meta-llama/Llama-3.1-8B (or your model of choice)
• Training dataset pre-downloaded to shared storage (we use 100K samples from the C4 dataset. You can use a dataset of your choice)

For EKS cluster creation with GPU nodes and EFA networking, see the infrastructure guides in awsome-distributed-ai/1.architectures. For the complete NVRx-specific setup including Terraform modules, container build, and dataset preparation, see the NVRx test case README.

Solution walkthrough

NVRx exposes two orthogonal changes to a standard PyTorch FSDP script: async checkpointing (a write-path optimization) and recovery, which comes in two independent layers: in-process restart for soft faults and ft_launcher in-job restart for hard faults. The test case ships a separate script per capability, so you adopt only what you need.

We start from a baseline FSDP script with synchronous checkpointing, then introduce each NVRx capability independently (async checkpointing, in-process restart, and the ft_launcher in-job restart) and close with a short note on how the recovery layers compose.

Baseline FSDP training script

Start with a minimal FSDP loop launched by torchrun using a synchronous distributed checkpoint save (PyTorch’s torch.distributed.checkpoint.save under the hood). The save blocks every rank until per-rank shards land on shared storage, and a worker crash kills the entire training job, requiring a full restart from the last checkpoint.

# Baseline FSDP training --- torchrun launches; sync checkpoints block the loop.
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")

model, _ = create_model(args.model_name, args.torch_dtype)
model = wrap_model(model, "fsdp", local_rank, args.model_name)
optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
data_iter = iter(create_dataloader(args, tokenizer, rank, world_size))

for step in range(1, args.max_steps + 1):
    loss = train_step(model, next(data_iter), optimizer)  # fwd / bwd / step

    if step % args.checkpoint_interval == 0:
        # save_checkpoint() uses dcp.save() under the hood for FSDP --- collective.
        # All ranks block here until per-rank shards land on shared storage.
        save_checkpoint(model, optimizer, step,
                        args.checkpoint_path, rank, "fsdp")

Async checkpointing: decouple checkpoint I/O from training

Figure 2: NVRx async checkpoint pipeline

Replace the synchronous save with the NVRx TorchAsyncCheckpoint: instantiate it once with persistent_queue=True, call async_save(state_dict, path) in place of torch.save, and call finalize_async_save(blocking=True) once as a blocking finalization at job exit, as shown in figure 2. A background process owns the actual write. The main thread continues to the next forward/backward step. Each rank writes its own shard through FSDP LOCAL_STATE_DICT, with no all-gather and no rank-0 bottleneck.

import torch
import torch.distributed as dist
from nvidia_resiliency_ext.checkpointing.async_ckpt.torch_ckpt import TorchAsyncCheckpoint

# Initialize the async checkpoint manager once after model/optimizer setup.
async_ckpt = TorchAsyncCheckpoint(persistent_queue=True)

# Training loop
for step, batch in enumerate(dataloader):
    loss = model(batch)
    loss.backward()
    optimizer.step()

    # Checkpoint every N steps
    if step % checkpoint_interval == 0:
        state_dict = build_state_dict(model, optimizer, step)  # FSDP LOCAL_STATE_DICT, CPU-staged
        # torch.save(state_dict, path)  # was blocking
        async_ckpt.async_save(state_dict, path)  # returns immediately

# Drain any in-flight save at job exit.
async_ckpt.finalize_async_save(blocking=True)

In-process restart

Wrap the train function with inprocess.Wrapper. The wrapper owns the restart loop: an exception in the wrapped function (or a hang the watchdog detects) triggers a re-entry instead of crashing the process. The constructor wires four concerns: timeouts (soft_timeout, hard_timeout, barrier_timeout, completion_timeout), health checks (CudaHealthCheck + FaultCounter), a RetryController that caps the total number of restarts and sets a minimum surviving world size, and rank assignment (ActivateAllRanks + ShiftRanks) that pulls survivors left to keep the world contiguous.

Figure 3: NVRx in-process restart architecture

import nvidia_resiliency_ext.inprocess as inprocess
from nvidia_resiliency_ext.inprocess import CallWrapper

def train_with_inprocess_restart(args, restart_metrics, base_store=None,
                                 call_wrapper: CallWrapper = None):
    # Re-init dist, rebuild FSDP, load latest checkpoint, train.
    # call_wrapper.ping() each step;
    # call_wrapper.iteration tells you which restart you're on.
    ...

wrapped_train = inprocess.Wrapper(
    store_kwargs = {"host_name": master_addr, "port": master_port + 1},
    soft_timeout = datetime.timedelta(seconds=args.soft_timeout_seconds),
    hard_timeout = datetime.timedelta(seconds=args.hard_timeout_seconds),
    barrier_timeout = datetime.timedelta(seconds=args.barrier_timeout_seconds),
    completion_timeout = datetime.timedelta(seconds=args.barrier_timeout_seconds),
    health_check = inprocess.Compose(
        inprocess.health_check.CudaHealthCheck(),
        inprocess.health_check.FaultCounter(max_rank_faults=20)),
    initialize = inprocess.initialize.RetryController(
        max_iterations=args.max_restarts, min_active_world_size=1),
    rank_assignment = inprocess.Compose(
        inprocess.rank_assignment.ActivateAllRanks(),
        inprocess.rank_assignment.ShiftRanks()),
)(train_with_inprocess_restart)

wrapped_train(args, restart_metrics, base_store)

Behind the scenes (as shown in figure 3): a ProgressWatchdog (progress_watchdog.py:49) uses Py_AddPendingCall to write timestamps between bytecode instructions, so a hang inside NCCL (a C extension) is still detectable. A MonitorThread (monitor_thread.py:124) polls the inprocess TCPStore for an interrupted flag and raises RankShouldRestart into the main thread on detection. AbortTorchDistributed (abort.py:62) then collects Flight Recorder traces, aborts NCCL backends, and destroys the process group. Health checks pick survivors. The Python process stays alive, and only the distributed process group is rebuilt.

In-job restart (ft_launcher)

Switch the launcher from torchrun to ft_launcher. It understands the same rendezvous flags and adds a few of its own to control restart behavior:

# --- before ---
torchrun --nnodes=$NNODES --nproc_per_node=$GPU_PER_NODE \
    --rdzv-backend=c10d --rdzv-endpoint=$MASTER_ADDR:29500 \
    train.py [args]

# --- after ---
ft_launcher --nnodes=$NNODES --nproc_per_node=$GPU_PER_NODE \
    --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29500 \
    --max-restarts=20 --ft-restart-policy=any-failed \
    --ft-rank-heartbeat-timeout=900 \
    --ft-initial-rank-heartbeat-timeout=1200 \
    --monitor-interval=5 \
    train_ft_launcher.py [args]

Inside the train script, instantiate RankMonitorClient once after distributed init and send a heartbeat each step:

import nvidia_resiliency_ext.fault_tolerance as fault_tolerance

ft_client = fault_tolerance.RankMonitorClient()
ft_client.init_workload_monitoring()  # once, after dist init

for step in range(1, args.max_steps + 1):
    loss = train_step(model, next(data_iter), optimizer)
    ft_client.send_heartbeat()  # liveness signal to RankMonitorServer

Set --ft-rank-heartbeat-timeout above the longest legitimate interval between application heartbeats. The example uses 900 seconds. Setting --ft-initial-rank-heartbeat-timeout=1200 (20 minutes) accommodates first-time model loading. See kubernetes/training-job-ft-launcher.yaml.

Figure 4: NVRx in-job restart (ft_launcher) architecture

Behind the scenes (as shown in figure 4): a RankMonitorServer per rank tracks heartbeat intervals against the timeouts above. On timeout, the launcher SIGTERMs survivors (SIGKILL for stragglers), reclaims GPU memory, re-rendezvouses, and spawns fresh workers. Workers reload from the latest checkpoint on startup. Checkpoint frequency caps the amount of lost work.

How the layers cover distinct fault classes

The two recovery layers are scoped to different failure modes: in-process catches what fits inside one Python process (transient exceptions, watchdog-visible NCCL hangs), ft_launcher catches what kills the process or hangs at the OS level (SIGKILL, OOM, sub-Python deadlocks), and the cluster scheduler catches node loss. Pick the layer whose blast radius matches your failure modes. Async checkpointing is orthogonal: it pairs with a recovery layer (or none) and caps the lost-work blast radius.

Deploy and run experiments

We deploy training jobs using a thin wrapper script around kubectl that handles manifest templating, job cleanup, and environment variable substitution. Environment variables define the instance type, GPU count, number of EFA devices, and other hardware-specific parameters, so the same training code and manifests run on different GPU types (p5, p4de) by changing a single configuration file. To compare recovery mechanisms under identical conditions, we use deterministic fault injection, pre-generating exactly N faults at fixed training steps and ranks using a seeded RNG:

--fault_count=5 --fault_seed=42 --fault_types=exception,hang --fault_type_weights=0.6,0.4

The same seed produces the same fault schedule across experiments, enabling direct comparison between baseline K8s restart, ft_launcher, and NVRx in-process restart. We run each mechanism against the same 5-fault pattern, then run async vs sync checkpointing separately without fault injection to isolate checkpoint overhead.

Results

This section presents benchmarking results for two key capabilities: async checkpointing and fault recovery.

Async checkpointin