NVIDIA Resiliency Extension (NVRx) integrates with PyTorch FSDP on Amazon EKS for fault-tolerant distributed training

Addressing Distributed Training Challenges

Large-scale distributed training jobs, which can run for days across many nodes, are prone to interruptions from network issues, memory errors, software exceptions, or infrastructure events. These disruptions can lead to significant idle time and wasted GPU resources, as a single GPU fault can cause a cascade of failures and out-of-sync restarts. Synchronous checkpointing further contributes to idle time by blocking all ranks on I/O operations.

NVRx Integration on Amazon EKS

NVIDIA Resiliency Extension (NVRx) has been integrated with PyTorch Fully Sharded Data Parallel (FSDP) training on Amazon Elastic Kubernetes Service (Amazon EKS) to mitigate these issues. This integration provides solutions for both interruption recovery and idle time reduction. The EKS cluster provides the underlying infrastructure for GPU scheduling, high-bandwidth networking, and shared storage for checkpoint persistence, while NVRx handles application-level resilience.

Key Features of NVRx

NVRx is a pip-installable Python layer that adds fault-tolerance primitives to PyTorch without requiring custom kernels, PyTorch forks, or recompilation. It integrates into existing FSDP scripts as ordinary imports, leaving model and training code untouched. The three main features exercised are async checkpointing, in-process restart, and in-job restart using ft_launcher.

Async checkpointing, exposed through TorchAsyncCheckpoint, replaces torch.save with an async_save() call that offloads state dict saving to a background process, returning immediately. A subsequent finalize_async_save() commits the prior write, allowing I/O to overlap with training.

Impact on Training Efficiency

This integration allows for recovery from faults in seconds through in-process restart, avoiding container lifecycle disruptions. In-job restart, facilitated by ft_launcher, provides automatic worker respawn for hard crashes. These capabilities reduce the impact of inevitable interruptions in large-scale distributed training, improving overall efficiency and resource utilization.