Amazon EKS NVRx 训练将 GPU 故障恢复时间从数分钟缩短至数秒
Amazon EKS 已将 NVIDIA NVRx 集成到一套可复现的训练栈中,可在约 10 至 17 秒内恢复人为注入的 GPU 故障。在覆盖 16 至 64 块 H100 GPU 的部分测试中,Amazon EKS NVRx 训练方案还将检查点效率维持在 99% 以上。这些结果挑战了一项代价高昂的假设:可靠恢复必须从重启容器或重建整个 Kubernetes 任务开始。
该系统将 PyTorch Fully Sharded Data Parallel(FSDP)与三项独立的 NVRx 能力结合。异步检查点将存储写入移出训练循环。进程内重启无需替换 Python 进程即可重建分布式状态。ft_launcher 组件则会在发生更严重故障后,在现有任务内启动新的工作进程。
因此,关键竞争并非 AWS 与其他云服务商之间的较量,而是应用感知恢复与仅依赖基础设施恢复之间的对比。Kubernetes 仍负责调度和节点级故障,但 NVRx 会处理更靠近训练进程的故障。AWS 表示,这种分工能显著减少健康而昂贵的 GPU 因等待故障节点而闲置的时间。
Amazon EKS NVRx 训练将恢复机制移入任务内部
核心变化在于,工作进程故障不再必须演变为完整的容器生命周期事件。
AWS 与 NVIDIA 围绕 Amazon EKS、自管理 GPU 节点组和 PyTorch FSDP 构建了参考环境。其公开基准测试使用 p5.48xlarge 实例,每个实例配备八块拥有 80 GB 显存的 NVIDIA H100 GPU。
测试集群从两个节点扩展至八个节点,即从 16 块扩展至 64 块 GPU。每个实例还提供 32 个 Elastic Fabric Adapter 接口,用于高带宽通信。Amazon FSx for Lustre 为训练 Pod 提供共享检查点存储。
NVRx 是 NVIDIA Resiliency Extension 的简称,是一个为 PyTorch 工作负载添加恢复和检查点组件的 Python 软件包。它不需要 PyTorch 分支版本、自定义内核或重新编译。团队可以独立添加其功能,而无需替换完整的训练框架。
这种模块化设计很重要,因为检查点性能和故障恢复是两类不同的问题。一个工作负载可能需要更快的保存速度,却不需要进程恢复;另一个工作负载则可能需要防护进程崩溃,同时保留现有的检查点实现。
该参考架构根据故障范围处理各个层级。进程内重启处理 Python 解释器仍存活时发生的异常和通信卡顿。ft_launcher 则处理 SIGKILL、内存不足终止以及部分操作系统级卡顿等事件。
Kubernetes 仍作为最外层,处理导致整个节点失效的故障。这种方法类似一组嵌套的恢复区域:每种机制仅在故障越过其下一层边界时才介入。
训练 Pod 使用无头 Kubernetes Service 进行对等节点发现。工作进程通过 DNS 而非固定 IP 地址相互查找。这种安排使替代工作进程能够重新加入,而无需运维人员改写任务配置。
AWS 此前曾介绍过在 EKS 上使用 PyTorch 工具实现的弹性分布式训练。NVRx 的工作进一步缩小了恢复闭环,重点是在单个 rank 发生故障、停滞或消失时,让活跃任务继续有效运行。
这种差异构成了本文的核心张力。Kubernetes 可以恢复基础设施,但基础设施恢复并不了解模型状态、进程组和检查点时机等细节。NVRx 则将这些决策带入训练应用之中。
阻塞式检查点曾消耗约 40% 的实际运行时间
首要性能问题并非 GPU 计算,而是每个 rank 等待检查点数据写入存储所花费的时间。
同步检查点会暂停训练,直到所需的模型和优化器状态完成写入。在分布式 FSDP 任务中,这一暂停会影响每个参与的 rank。健康的 GPU 仍被占用,但在写入期间不会执行前向或反向计算。
AWS 报告称,在其扩展测试中,同步检查点仅实现了 57% 至 61% 的训练效率。存储写入耗时约 275 秒,在 16 至 64 块 GPU 之间大致保持不变。因此,增加计算资源并不能消除受存储限制的暂停。
NVRx 异步检查点改变了写入路径。训练进程会先在 CPU 上暂存状态,再将工作交给持久化后台进程。主进程则在存储 I/O 持续进行时返回下一步训练。
该实现使用 TorchAsyncCheckpoint 及其 async_save() 方法。在发起下一次保存或退出前,应用会完成尚未完成的操作。这种协调可避免未完成的检查点与下一次检查点悄然发生冲突。
FSDP 本地状态字典进一步强化了这一设计。每个 rank 写入自身的分片,避免 all-gather 操作以及单一 rank-zero 写入瓶颈。PyTorch 的 FSDP 文档介绍了将参数分布到参与工作进程中的更广泛分片模型。
AWS 表示,在每 1,000 步进行一次检查点的条件下,NVRx 异步检查点在两个节点上实现了 99.2% 的训练效率,在八个节点上达到 99.8%。八节点下的同步对比结果为 60.3%。
这些数字是厂商报告的基准测试结果,并不保证适用于每种模型或存储配置。不过,其背后的机制很直接:如果计算时间长于存储写入时间,后台操作几乎可以完全隐藏在有效训练工作之后。
当 AWS 提高检查点频率时,这一边界变得清晰。在八个节点、每 100 步进行一次检查点的情况下,同步效率降至 14.7%。异步效率同样下降,但仍达到较高的 29.6%。
原因在于时序。100 个训练步骤大约耗时 280 秒,而检查点写入约需 275 秒。几乎没有额外的计算窗口可用于隐藏下一次存储操作。
这正是 NVRx 异步检查点的真正限制。异步 I/O 可以将写入隐藏在计算之后,但无法让存储无限加速。当保存请求到达的速度与文件系统完成写入的速度相当时,队列最终会产生压力。
即便存在这一限制,该能力仍改变了团队选择检查点间隔的方式。同步系统会鼓励减少检查点数量,因为每次保存都会带来明显的闲置时间。较低频率的保存则会增加故障后丢失的训练量。
异步检查点削弱了这种权衡。当两次写入之间存在足够计算量时,团队可以更频繁地保存。更短的间隔可减少回滚距离,而重叠 I/O 则能保留更多 GPU 投入价值。
其结果不只是一个更快的检查点 API,而是稳态效率与可恢复进度之间的一种新平衡。随着任务持续时间增长、故障倾向组件增多,这种平衡会愈发重要。
进程内重启挑战仅依赖 Kubernetes 的恢复模型
最快的恢复路径会保留 Python 进程,仅重建因故障受损的分布式资源。
在参考实现中,NVRx 使用进程内重启控制器包装主训练函数。若发生受支持的异常,包装器会中断当前尝试,并准备再次调用。外层 Python 进程在整个过程中保持存活。
NVRx 首先会中止受损的 PyTorch 分布式进程组。它可以收集飞行记录器追踪、停止 NCCL 后端,并销毁无效进程组。NCCL 是 NVIDIA 用于跨 GPU 集体操作通信的库。
随后,健康检查会检查与每个 rank 关联的资源。这些检查可涵盖 GPU、NVLink 连接、网络接口以及 rank 的重复故障。重试控制器会限制重启尝试次数,并确定必须存活的活跃 rank 数量。
系统会将存活的 rank 重新分配为连续组,并启动新的 rendezvous。被包装的训练函数会重新创建 FSDP 模型、加载最新检查点并恢复运行。Python 解释器以及包装函数之外的对象仍然可用。
该技术针对软故障,例如未处理的应用异常,以及看门狗能够检测到的 NCCL 卡顿。它并不假定每一次被阻塞的原生调用都能可靠地产生 Python 异常。
相反,进度看门狗会记录 Python 字节码操作之间的活动。单独的监控线程检查共享状态,并可在某个 rank 停止推进时请求重启。系统随后会在参与的工作进程之间协调中断。
AWS 将这一方法与 ft_launcher 及基础 Kubernetes 恢复进行了比较。实验使用两个 p5.48xlarge 节点、16 块 H100 GPU,以及采用 FSDP 的 Llama 3.1 8B。任务运行 2,000 步,每 500 步保存一次。
研究人员按照相同计划,在每次运行中注入五个确定性故障。NVRx 进程内重启可在不重启容器的情况下,以每次故障约 10 秒的速度完成恢复。AWS 测得训练有效产出率为 31%,基础设施有效产出率为 87%。
训练有效产出率衡量能够产生有效训练进展的时间。基础设施有效产出率衡量已分配基础设施保持运行并可用的时间。两者之间的差距包括无法推动模型进展的工作,如回滚和检查点加载。
在报告的实验中,基础 Kubernetes 恢复每次注入故障约需 270 秒,训练有效产出率为 11.5%,基础设施有效产出率为 35.8%。这意味着该对比远不只是容器启动优化。
AWS 表示,一个失败的 rank 会在存活 rank 上触发通信超时。随后,Pod 会以不同步的方式重启,从而造成反复超时循环和 CrashLoopBackOff 行为。编排器恢复了容器,却不了解分布式训练组需要如何共同恢复。
应用感知恢复可以获得这种缺失的上下文。它知道进度何时停止、哪个进程组已失效,以及哪个检查点能够重新启动训练。Kubernetes 能看到 Pod 状态和节点健康状况,但看不到 FSDP 训练步骤的完整语义。
这并不意味着 Kubernetes 恢复不再必要。失效节点无法保留其解释器、CUDA 状态或本地进程。节点替换仍属于集群层,恢复后的工作进程也仍需要持久化检查点数据。
压力落在那些将 Pod 重启作为唯一容错策略的团队身上。这种方法依然简单,但其恢复窗口可能浪费大量加速器时间。集群越大,成本越会被放大,因为一次故障可能让许多原本健康的工作进程闲置。
NVRx 容错将软故障与硬故障分开处理
没有一种重启机制能够覆盖所有故障,因此 NVRx 容错将保留进程的恢复与工作进程替换分离开来。
ft_launcher 组件负责处理进程内重启无法恢复的故障,包括 SIGKILL、内存不足导致的终止,以及无法保留可用 Python 解释器的失败情形。它取代了 torchrun,同时保留了熟悉的 rendezvous 概念。
每个训练 rank 会在分布式初始化后创建一个 RankMonitorClient。客户端会在训练过程中发送心跳。每个 rank 对应的监控服务器会根据为正常运行和初始启动配置的超时时间,判断这些信号是否正常。
AWS 的配置将 rank 心跳超时设为 900 秒。初始心跳超时为 1,200 秒,为首次加载模型预留更多时间。五秒的监控间隔决定了 launcher 检查 worker 状态的频率。
这些数值只是配置示例,并非通用建议。心跳超时必须长于两次信号之间可能出现的最长正常延迟。若设置过短,缓慢的 checkpoint 或模型初始化可能会被误判为 worker 失效。
当某个 worker 终止或停止响应时,ft_launcher 会终止其余 worker。随后,它会回收 GPU 内存,重新执行 rendezvous,并在同一作业内启动新的进程。新的 worker 将从最近一次 checkpoint 恢复状态。
launcher guide 展示了 NVIDIA NeMo RL 技术栈中的类似通用模式。这种更广泛的集成表明,NVRx 的定位是可复用的韧性层,而非仅面向 EKS 的工具。
AWS 测得,使用 ft_launcher 后,每次注入故障的恢复时间约为 17 秒。报告中的运行实现了 25.5% 的训练 goodput 和 85.9% 的基础设施 goodput。虽然它慢于进程内恢复,但远快于 Kubernetes 基线的 270 秒。
这一差异反映了各种方法能够保留的状态量不同。进程内重启会保留解释器和外层进程;ft_launcher 则必须创建新 worker、初始化分布式状态、重建模型并重新加载 checkpoint。
在更大规模下,checkpoint 加载可能主导恢复时长。更快地创建进程并不能免除读取模型和优化器分片的需求。因此,共享文件系统吞吐量仍是容错设计的一部分。
Amazon EKS NVRx 训练架构使用了与 GPU 节点位于同一可用区的 FSx for Lustre SCRATCH_2 文件系统。这种部署旨在降低 checkpoint 读取延迟。恢复后,所有 worker 都可以访问相同的持久化状态。
NVRx 异步 checkpoint 与两种重启路径相互独立。它减少与写入相关的空闲时间,并控制可能面临风险的训练进度量。重启层决定的是发生故障后 worker 能多快恢复运行。
这种分离为运维人员提供了更多选择,但也增加了策略制定工作。他们必须决定哪些异常可以触发进程内恢复、多少次重试是安全的,以及哪些健康检查应移除一个 rank。同时还必须设置心跳和 rendezvous 超时。
NVIDIA 将 NVRx project 标注为实验性项目,且仍在积极开发中。其文档警告称,功能和接口可能发生变化。生产团队应将版本选择和升级测试纳入韧性计划。
AWS 使用 NVRx 0.4.1 复现其基准测试。该文章建议,在当前部署中使用更新了 launcher 配置的 0.6.0 版本。这一版本差异很重要,因为容错系统直接处于关键的启动和恢复路径上。
如果不可恢复的作业被反复重启,失败的恢复机制可能比没有自动化更糟。重试上限、最小 world size 和故障计数器能够防止无限循环。运维人员仍需要能够区分恢复成功和持续性故障的告警。
99% 的结果存在重要边界
该基准测试证明了一项强有力的机制,但并未说明每种分布式训练工作负载都能达到 99% 效率。
AWS 在基于 H100 的 p5 实例上,测试了一个主要模型配置:采用 PyTorch FSDP 的 Llama 3.1 8B。测试使用 EFA 网络和 FSx for Lustre 存储。不同的模型规模、存储路径、checkpoint 格式和 step 时长都会改变可重叠的时间窗口。
99% 这一数字适用于选定间隔下的异步 checkpoint 效率,并不代表反复发生故障时的端到端 goodput。在注入故障的测试中,进程内恢复的训练 goodput 为 31%,ft_launcher 则为 25.5%。
这些较低的结果并不与 checkpoint 测量相矛盾,因为它们回答的是不同问题。异步效率衡量正常训练期间 checkpoint 的开销,而 goodput 则涵盖故障注入、回滚、加载和其他恢复工作。
Checkpoint 频率也存在不可避免的限制。在每 100 个 step 保存一次 checkpoint 时,异步训练达到的是 29.6% 效率,而非 99%。由于计算和写入耗时相近,存储系统几乎处于持续占用状态。
内存压力同样值得关注。异步 checkpoint 会将数据暂存到即时 GPU 操作之外,并由后台进程承担写入责任。团队应基于实际的状态字典,测量 CPU 内存、队列深度和存储积压情况。
恢复覆盖范围是另一项边界。操作系统终止 worker 或节点消失时,进程内重启无能为力。ft_launcher 可以替换已死亡的进程,但仍依赖于作业、集群网络、rendezvous 服务和 checkpoint 存储。
节点丢失仍是 Kubernetes 需要处理的问题。区域级存储中断或损坏的 checkpoint,可能同时击穿所有恢复层。该架构降低了多类常见故障的成本,但并未消除共享依赖。
故障检测也可能产生误报。较长的编译阶段、数据加载暂停或文件系统停滞,可能超出激进的心跳超时设置。此时 launcher 会重启健康的 worker,并丢弃有效进度。
团队在启用自动恢复前,需要获得工作负载特定的超时数据。他们应记录最长的模型初始化、checkpoint、验证和数据输入间隔。测试还应覆盖这些阶段发生的故障,而不仅是常规训练 step 内部的故障。
AWS 的基准测试使用了确定性的注入故障。这种方法有助于开展可重复的比较,但生产环境中的故障往往并不规律。真实集群可能同时遭遇网络退化、存储变慢、散热问题和进程崩溃。
reference implementation 为团队复现该配置提供了有用起点。在不同实例系列和模型规模上进行复现,将决定报告中这些优势的可移植性。
运维复杂性是最后一项权衡。该技术栈包括 Kubernetes Jobs、基于 DNS 的对等发现、EFA 资源、共享存储、NVRx wrapper、监控客户端和多层超时设置。每个组件都会增加一处配置界面。
当大型 GPU 集群因一个 rank 失败而闲置数分钟时,这种复杂性仍可能合理。然而,较小规模的作业或许可以接受更简单的 pod 重启策略。关键计算应是恢复成本乘以故障频率,而非基准测试的声望。
评估该设计的团队应同时追踪训练和基础设施 goodput。即使模型反复重新加载旧 checkpoint,GPU 利用率也可能看上去正常。一个有用的仪表板必须展示已完成 step、回滚距离、checkpoint 新鲜度和重启原因。
工程组织还需要为这些实验保留可持续的记录。可搜索的工程知识库可以关联超时变更、故障追踪和基准结果。这些上下文有助于团队避免重复采用失败的恢复配置。
Amazon EKS NVRx 基准测试后值得关注的事项
接下来的证据应证明,该设计能否在更大模型、真实故障和持续变化的 NVRx 版本中保持其优势。
第一个信号是在超过 64 块 H100 GPU 的环境中得到独立复现。AWS 测试了从两个节点到八个节点的扩展,但故障频率和协调成本会随集群规模上升。跨越数百个加速器的结果,将更好地揭示 rendezvous 和 checkpoint 加载的限制。
更大规模的测试不应只报告平均恢复时间。恢复时间分布同样重要,因为少数持续五分钟的恢复可能主导长时间运行的经济性。报告应包含尾部延迟、失败的重启尝试,以及每次事件丢失的进度。
第二个信号是主要训练框架对它的采用。NVRx 已经可连接基于 PyTorch 的工作负载,并出现在 NVIDIA 更广泛的软件技术栈中。更多原生集成将减少团队需要维护的自定义 wrapper 和 launcher 代码。
框架采用将强化这样一种观点:NVRx 容错能力可以成为标准应用层。碎片化集成则会削弱这一观点,尤其是在每个框架都需要不同的超时、checkpoint 和 rendezvous 逻辑时。
第三个信号是来自非受控生产故障的证据。确定性的故障注入对于比较必不可少,但现场事故会检验实验室中很少复现的组合。运维人员应分别发布 GPU 错误、NCCL 停滞、OOM 终止和节点丢失的恢复率。
成功恢复的含义不应仅是重启一个进程。作业必须恢复有效状态,继续产生正确的更新,并避免静默的 checkpoint 损坏。经历反复恢复后的模型收敛情况,应与恢复速度一样受到严格审视。
对于正在考虑 Amazon EKS NVRx 训练的团队,务实的第一步是进行受控的影子基准测试。使用真实的模型、checkpoint 大小、文件系统和 step 时长。在同一故障计划下比较同步保存、异步保存、launcher 恢复和仅使用 Kubernetes 的恢复。
接着,根据观测到的计算和存储耗时调整 checkpoint 间隔。如果 checkpoint 所需时间几乎与两次保存之间的间隔一样长,异步重叠仍将不完整。如果计算能够提供更宽的窗口,报告中的 99% 效率便更具可信度。
恢复策略应从保守的重试上限开始。记录每一次重启原因,并保留诊断追踪。若作业在同一个 rank 或 checkpoint 上反复失败,就需要升级处理,而不是陷入无休止的恢复循环。
AWS 和 NVIDIA 的工作使一个结论难以忽视:分布式训练的可靠性不能再只是基础设施问题。应用程序对进度、checkpoint 有效性和进程组状态的理解,是编排器所不具备的。
Amazon EKS 仍然提供了必要的调度和节点替换基础。NVRx 则在这一边界内提供更快的响应。两者结合,为若干重要故障类别提供了一条可信路径:将四分钟的恢复周期缩短至秒级重启。
现在悬而未决的问题是运维层面的,而非概念层面的。团队能否在自身的模型、存储系统和真实故障模式下复现这些收益?这才是下一次 Amazon EKS NVRx 训练部署应遵循的检验标准。



