在 Amazon EKS 上使用 NVRx 进行容错分布式训练
aws ·
本文展示了如何将 NVIDIA Resiliency Extension (NVRx) 与 Amazon EKS 集成,以构建针对大规模 PyTorch 模型的容错分布式训练。该方案解决了两个主要瓶颈:同步检查点可能导致高达 40% 的训练时间被阻塞,以及从硬件或软件故障中恢复缓慢。通过利用 TorchAsyncCheckpoint 实现异步检查点,将 I/O 与计算重叠,训练效率相比基线提升了 99% 以上。在故障恢复方面,NVRx 提供针对软故障的进程内重启和针对硬崩溃的 ft_launcher,将恢复时间从分钟级缩短至秒级。在 LLaMA-3.1-8B 模型上进行的 2 到 8 个节点的基准测试表明,这些方法显著优于标准 Kubernetes 重启方式,后者常受限于 NCCL 超时级联效应。该架构利用 Amazon FSx for Lustre 进行共享存储,并使用弹性 fabric 适配器(EFA)实现高带宽网络。这种方法使团队能够在不牺牲 GPU 利用率的情况下更频繁地进行检查点保存,从而最大化整体集群效率。
aws