Integrating NVRx Fault‑Tolerance into PyTorch FSDP on Amazon EKS Reduces Training Downtime

Integrating the NVIDIA Resiliency Extension (NVRx) into a PyTorch Fully Sharded Data Parallel (FSDP) workload running on Amazon EKS introduces three new fault‑tolerance mechanisms—async checkpointing, in‑process restart, and the ft_launcher binary—while keeping the original training code unchanged. Practitioners care because the added mechanisms eliminate the long idle periods caused by synchronous checkpoint I/O and prevent costly GPU hour loss when a single worker fails.

Solution Overview

The architecture layers NVRx’s application‑level resilience on top of an EKS cluster that already provides GPU scheduling, high‑bandwidth EFA networking, and shared storage via FSx for Lustre. Each training pod runs as a Kubernetes Job and discovers peers through a headless Service, allowing pods to be replaced without reconfiguring the job topology.

Key NVRx Primitives

NVRx is distributed as a pip‑installable Python package (pip install nvidia-resiliency-ext) and exposes three independent primitives:

• Async checkpointing – Replaces torch.save with async_save() (via TorchAsyncCheckpoint) and requires a call to finalize_async_save() before the next checkpoint. Each rank writes its own shard directly to FSx, avoiding an all‑gather step and eliminating the rank‑0 bottleneck that previously consumed up to 40% of wall‑time.
• In‑process restart – Wrapped around the training loop with inprocess.Wrapper. When a soft fault such as an unhandled exception or NCCL timeout occurs, NVRx aborts the current process group, runs health checks on GPU, NVLink, and NIC, re‑rendezvous surviving ranks, and resumes from the latest async checkpoint without restarting the container.
• ft_launcher binary – Monitors each rank via RankMonitorClient. It detects hard failures (SIGKILL, OOM kill, OS hangs) by missing heartbeats, then kills any remaining processes, reclaims GPU memory, and respawns fresh workers that reload the most recent checkpoint. This layer handles failures that the in‑process wrapper cannot catch.

EKS Infrastructure Adjustments

The cluster uses self‑managed node groups of p5.48xlarge instances, each providing eight NVIDIA H100 80 GB GPUs and thirty‑two EFA network interfaces. GPU and EFA resources are advertised through the NVIDIA and EFA device plugins, enabling the Kubernetes scheduler to place training pods with node affinity and tolerations that guarantee full‑node GPU allocation.

Checkpoint data resides on an Amazon FSx for Lustre file system (type SCRATCH_2, 1.2 TB) mounted via the FSx CSI driver. Locating the file system in the same Availability Zone as the GPU nodes reduces read latency, which is critical because checkpoint loading dominates recovery time at scale.

Related CloudNinjas coverage: AWS.

What This Means For Practitioners

Adopting NVRx on EKS removes the need to redesign existing FSDP scripts; the primitives drop in as imports and configuration calls. Teams should validate that async checkpoint directories are correctly mounted on FSx and that health‑check scripts cover GPU, NVLink, and NIC status. Monitoring the heartbeat channel used by ft_launcher provides early visibility into hard‑fault conditions. Finally, evaluate the trade‑off between async checkpoint frequency and storage I/O load, and test both soft‑fault and hard‑fault recovery paths before scaling to multi‑node runs.