AWS发布NVRx在EKS上做容错分布式训练方案：异步检查点保持 99%+ 效率
AWS Machine Learning博客介绍如何将NVIDIA Resiliency Extension（NVRx）集成到Amazon EKS上的PyTorch FSDP训练：异步检查点把I/O与训练重叠，进程内重启在数秒内恢复软故障，ft_launcher处理硬故障，H100上 2 至 8 节点测试显示 99%+ 训练效率。
AI解读：AWS这篇博客解决的是大规模分布式训练里两个很实际的浪费：GPU一坏，NCCL超时会把健康worker一起拖下水；同步检查点又会卡住所有rank等I/O写完。按文中数据，后者在测试集群里最多吃掉 40% 的总墙钟时间。
NVRx是一个pip安装的Python层（pip install nvidia-resiliency-ext），不改PyTorch源码、不重编译。异步检查点用TorchAsyncCheckpoint替代torch.save，配合FSDP的LOCAL_STATE_DICT让每个rank直接写自己的分片，避免all-gather和rank-0瓶颈。
恢复分两层：进程内重启（inprocess.Wrapper）处理瞬时异常和NCCL挂起，进程和CUDA分配器都活着，只重建通信组；ft_launcher处理SIGKILL、OOM、OS级挂起这类进程级故障，通过心跳超时检测后杀掉幸存者、回收显存、重新拉起worker。两层可按故障类型独立选用。
实测H100上 2 到 8 节点（16 至 64 卡）训练LLaMA-3.1-8B：异步检查点各规模训练效率 99.2% 至 99.8%，同步始终停在 57% 至 61%。差距来自FSx Lustre写入时间约 275 秒基本恒定，异步把这段完全藏到训练后面。
检查点越频繁优势越大：每 1000 步时异步 99.8% 对同步 60.3%；每 100 步时同步掉到 14.7%，异步仍有 29.6%。对需要频繁存点、缩短回滚距离的团队，这改变了同步检查点的成本取舍。
AWS Machine Learning博客发布了一篇技术文章，介绍如何把NVIDIA Resiliency Extension（NVRx）集成到Amazon EKS上运行的PyTorch Fully Sharded Data Parallel（FSDP）训练中，解决大规模分布式训练的两类中断问题：GPU故障导致NCCL超时扩散，以及同步检查点阻塞I/O。
文章给出的基准测试针对LLaMA-3.1-8B，使用H100 GPU，规模从 2 节点到 8 节点（16 至 64 卡）。异步检查点在所有规模下保持 99%+ 训练效率（2 节点 99.2%，8 节点 99.8%），同步检查点则稳定在 57% 至 61%，约 40% 墙钟时间花在阻塞I/O上。
NVRx是可通过pip install nvidia-resiliency-ext安装的Python层，不需要自定义kernel、不需要fork PyTorch、不需要重新编译，以普通import方式加入现有FSDP脚本，模型和训练代码保持不变。文章覆盖三个特性：异步检查点、进程内重启和ft_launcher任务内重启。
异步检查点：把写入从训练主线程移走
NVRx的TorchAsyncCheckpoint用async_save(state_dict, path) 替代torch.save，把状态字典交给后台进程写入后立即返回；下一次保存前调用finalize_async_save() 提交上一次写入，任务退出时再调用一次finalize_async_save(blocking=True) 做阻塞收尾。
配合FSDP的LOCAL_STATE_DICT，每个rank直接写自己的分片，没有all-gather，也没有rank-0瓶颈。文章给出的示例在模型和优化器初始化后实例化一次TorchAsyncCheckpoint(persistent_queue=True)，然后在训练循环中每N步调用async_save。
两层恢复：进程内处理软故障，ft_launcher处理硬故障
进程内重启通过inprocess.Wrapper包裹train函数。包装器拥有重启循环：被包裹函数抛异常或看门狗检测到挂起时，触发重新进入而不是让进程崩溃。构造参数覆盖四类关注点：超时（soft_timeout、hard_timeout、barrier_timeout、completion_timeout）、健康检查（CudaHealthCheck加FaultCounter）、限制重启总次数和最小存活world size的RetryController，以及把幸存rank左移保持连续的rank分配（ActivateAllRanks加ShiftRanks）。
ft_launcher是NVRx的任务内重启启动器，处理进程内无法捕获的情况：SIGKILL、OOM kill、OS级挂起。每个rank运行一个RankMonitorClient，启动器对照命令行设定的超时检查心跳，发现停滞或死亡后杀掉幸存者、回收GPU内存，并在同一任务里重新拉起全新worker，恢复的worker从最新检查点重新加载。
文章明确两类恢复层覆盖不同的故障类别：进程内针对瞬时异常和看门狗可见的NCCL挂起，ft_launcher针对杀死进程或在OS级挂起的故障，集群调度器则针对节点丢失。异步检查点与恢复层正交，可以和任一恢复层搭配，也可以单独使用，作用是限制丢失工作的范围。
实测数据：异步检查点效率与检查点频率的关系
在固定每 1000 步检查点的场景下，异步检查点 2 节点效率 99.2%，8 节点 99.8%；同步检查点各规模稳定在 57% 至 61%。文章解释差距在于检查点开销受存储限制：FSx Lustre写入时间约 275 秒，无论 16 卡还是 64 卡都是这个量级，异步通过把I/O与下一段训练重叠完全隐藏了这段延迟。
在 8 节点上扫描检查点频率：每 1000 步时异步 99.8% 对同步 60.3%；每 100 步时同步跌到 14.7%，异步降到 29.6%，仍是同步的两倍。文章指出交叉点是检查点之间的训练间隔（100 步约 280 秒）接近FSx I/O时间（约 275 秒）的位置：高于该阈值时异步完全重叠，低于时部分重叠仍优于阻塞写入。
部署前提与实验方法
方案使用自管理节点组的p5.48xlarge实例，每节点 8 块NVIDIA H100 80 GB GPU和 32 个Elastic Fabric Adapter（EFA）网卡，EFA提供 3,200 Gbps网络带宽用于NCCL all-reduce。训练pod以Kubernetes Job运行，用headless Service做对等发现，通过DNS互相找到，替换的pod无需重新配置即可重新加入。
检查点存储使用Amazon FSx for Lustre（SCRATCH_2，1.2 TB），通过FSx CSI驱动挂载到每个训练pod，放在与GPU节点相同的可用区以降低故障恢复时的读取延迟。文章强调规模化恢复时，决定恢复时间的是检查点加载而不是重启机制本身。
前提条件包括：Amazon EKS集群v1.28+且节点组支持EFA、安装NVIDIA device plugin；PyTorch 2.9+与NVRx 0.4.1（复现博客基准）或 0.6.0（配合更新后的启动器配置用于当前部署）；HuggingFace账号可访问meta-llama/Llama-3.1-8B；训练数据预先下载到共享存储，文章使用C4数据集的 100K样本。
为对比恢复机制，文章使用确定性故障注入：在固定训练步和rank上按种子随机数预生成正好N个故障，示例命令为 --fault_count=5 --fault_seed=42 --fault_types=exception,hang --fault_type_weights=0.6,0.4。同一个种子产生相同的故障时间表，便于对基线K8s重启、ft_launcher和NVRx进程内重启做直接比较。
恢复对比实验设置
故障恢复实验在 2 台p5.48xlarge（共 16 块H100）上进行，训练LLaMA-3.1-8B加FSDP，每 500 步检查点一次，总运行 2000 步。每次训练运行注入 5 个确定性故障（相同种子、相同故障时间表），比较三种恢复机制：NVRx进程内重启、ft_launcher任务内重启和Kubernetes重启。文章摘要称进程内重启可在数秒内从故障恢复（来源为博客摘要与正文，正文中该对比的具体数值在提供的材料中未完整给出）。