Fault tolerant distributed training on Amazon EKS using NVRx
aws ·
This article demonstrates how to integrate NVIDIA Resiliency Extension (NVRx) with Amazon EKS to create fault-tolerant distributed training for large-scale PyTorch models. It addresses two major bottlenecks: synchronous checkpointing, which can block up to 40% of training time, and slow recovery from hardware or software faults. The solution utilizes async checkpointing via TorchAsyncCheckpoint to overlap I/O with computation, achieving over 99% training efficiency compared to the baseline. For fault recovery, NVRx offers in-process restart for soft faults and ft_launcher for hard crashes, reducing recovery times from minutes to seconds. Benchmarks on LLaMA-3.1-8B across 2 to 8 nodes show that these methods significantly outperform standard Kubernetes restarts, which suffer from NCCL timeout cascades. The architecture leverages Amazon FSx for Lustre for shared storage and Elastic Fabric Adapters for high-bandwidth networking. This approach allows teams to checkpoint more frequently without sacrificing GPU utilization, maximizing overall cluster efficiency.
aws