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

英文原标题：Fault tolerant distributed training on Amazon EKS using NVRx

本文为英文原文的机器翻译摘要，原文见：AWS Machine Learning Blog ↗

• 文章介绍如何将 NVIDIA Resiliency Extension (NVRx) 集成到 Amazon EKS 上的 PyTorch FSDP 训练中，以解决大规模分布式训练中的故障中断和同步检查点导致的 GPU 空闲问题。
• NVRx 提供三种容错原语：异步检查点（TorchAsyncCheckpoint）将 I/O 与训练重叠；进程内重启（inprocess.Wrapper）处理软故障（如 NCCL 挂起）；ft_launcher 处理硬故障（如 SIGKILL、OOM），实现作业内重启。
• EKS 集群使用 p5.48xlarge 实例（每节点 8 块 H100 GPU 和 32 个 EFA 网卡），通过 headless Service 进行对等发现，并使用 Amazon FSx for Lustre 作为共享检查点存储。
• 同步检查点会阻塞所有 rank，在本文集群规模下消耗高达 40% 的总墙钟时间；异步检查点配合 FSDP LOCAL_STATE_DICT 让每个 rank 直接写入自己的分片，避免 all-gather 和 rank-0 瓶颈。
• 恢复时间主要由检查点加载决定，而非重启机制本身；将 FSx 与 GPU 节点置于同一可用区可最小化恢复时的读取延迟。

关键数据

• 同步检查点消耗高达 40% 的总墙钟时间
• p5.48xlarge 实例每节点配备 8 块 NVIDIA H100 80 GB GPU 和 32 个 EFA 网卡
• EFA 提供 3,200 Gbps 网络带宽
• 使用 PyTorch 2.9+ 和 NVRx 0.4.1 复现基准测试
• FSx for Lustre 使用 SCRATCH_2 类型，容量 1.2 TB
• 训练数据集为 C4 数据集的 100K 样本

引语

• pip install nvidia-resiliency-ext

为什么值得看：本文提供了在 Amazon EKS 上使用 NVRx 实现容错分布式训练的完整方案，包括异步检查点和多层重启机制，可显著减少 GPU 空闲时间并提升大规模训练效率。