EKS上NVRx实现秒级故障恢复的分布式训练

大规模分布式训练任务在数十个节点上运行数小时甚至数天。在此规模和时长下，中断在统计学上不可避免：网络分区、内存错误、软件异常或基础设施事件最终都会导致至少一个 worker 出现故障。单个 GPU 故障会引发级联反应：NVIDIA Collective Communication Library（NCCL）超时传播到健康 worker，Pod 崩溃并以不同步的方式重启，而集群在零训练进度的情况下消耗着昂贵的 GPU 小时数。同步 checkpointing 增加了第二个空闲时间来源：每次保存都会阻塞所有 rank 的 I/O，在本文所述的集群规模下，这部分时间占到了总 wall time 的 40%。

在本文中，我们展示如何将 NVIDIA Resiliency Extension（NVRx）集成到 Amazon Elastic Kubernetes Service（Amazon EKS）上的 PyTorch Fully Sharded Data Parallel（FSDP）训练中，以解决这两个问题。我们将介绍与训练重叠 I/O 的异步 checkpointing、在几秒内从故障中恢复而不触及容器生命周期的进程内重启，以及使用 ft_launcher 进行自动 worker 重spawn 的进程内重启（应对硬崩溃）。我们提供了 H100 GPU 在 2 节点到 8 节点规模下的基准测试结果，所有代码均可复现。

该解决方案将 NVRx 容错原语与专为高性能多节点 GPU 工作负载设计的基于 EKS 的训练环境相结合。NVRx 处理应用级弹性（异步 checkpointing、进程内重启和进程内重启），而 EKS 集群提供基础设施基础：GPU 调度、高带宽网络和用于 checkpoint 持久化的共享存储。

NVIDIA Resiliency Extension（NVRx）是一个可通过 pip 安装的 Python 层（pip install nvidia-resiliency-ext），为 PyTorch 添加了容错原语：无需自定义内核、无需 PyTorch 分支、无需重新编译。这些原语可以作为普通导入插入现有的 FSDP 脚本中。模型和训练代码保持不变。每项功能均可独立采用。我们使用了三个特性：异步 checkpointing、进程内重启和 ft_launcher（进程内重启）。

异步 checkpointing，通过 TorchAsyncCheckpoint 暴露，用 async_save() 调用替换 torch.save，将 state dict 交给后台进程并立即返回。在下一次保存之前配合 finalize_async_save() 来提交之前的写入。结合 FSDP LOCAL_STATE_DICT，每个 rank 直接写入自己的分片，无需 all-gather 且没有 rank-0 瓶颈。

进程内重启，通过 inprocess.Wrapper 暴露，将训练函数包装起来，使得瞬态故障（未处理的异常或 NCCL 挂起）不会杀死 Python 进程。NVRx 中止活动进程组，对每个 rank 运行健康检查（GPU、NVLink、NIC），让幸存者重新 rendezvous，并从最新的 checkpoint 重新进入包装函数。解释器、CUDA 分配器和外部作用域对象得以保留。这针对的是软故障类。

ft_launcher 二进制文件是 NVRx 的进程内重启启动器，处理进程内无法捕获的情况：SIGKILL、内存不足（OOM）杀死和操作系统级挂起。每个 rank 运行一个 RankMonitorClient。启动器根据显式 CLI 设置的超时时间检查心跳，在停滞或死亡时杀死幸存者、回收 GPU 内存，并在同一作业中重新 spawn 新的 worker。恢复的 worker 从最新的 checkpoint 重新加载。每个恢复层覆盖一个独立的故障类：进程内用于软故障、ft_launcher 用于硬故障、集群编排器用于节点丢失。这些层是独立的。选择与你的故障模式相匹配的范围。

Amazon EKS 是一个托管 Kubernetes 服务，负责控制平面、升级和 API 服务器可用性。我们运行自管理的 p5.48xlarge 实例节点组，每个实例配备 8 张 NVIDIA H100 80 GB GPU 和 32 个 Elastic Fabric Adapter（EFA）网络接口。训练 Pod 作为 Kubernetes Job 运行，使用无头 Service 进行对等发现，因此 worker 通过 DNS 而非硬编码 IP 相互查找，Pod 替换可以在不重新配置作业的情况下重新加入。

每个节点通过 NVIDIA 设备插件和 EFA 设备插件将其 GPU 和 EFA 适配器公开为扩展资源。Kubernetes 调度器使用节点亲和性和容忍将训练 Pod 放置在 GPU 节点上，促进每个节点完整的 8-GPU 分配。

对于 checkpoint 存储，我们使用 Amazon FSx for Lustre（SCRATCH_2，1.2 TB），通过 FSx CSI 驱动挂载到每个训练 Pod 中。FSx 提供了异步和同步 checkpointing 都写入的共享 POSIX 文件系统，关键是在故障后恢复的 worker 从中读取 checkpoint 状态。将 FSx 放置在与 GPU 节点相同的可用区可最大程度减少恢复期间的读取延迟，这很重要，因为 checkpoint 加载（而非重启机制）在规模上是恢复时间的主导因素。

图 1：架构图，显示带有 2-8 个 p5 节点的 EKS 集群、EFA 互联、FSx for Lustre 以及训练 Pod 内的 NVRx 组件

涉及的关键 AWS 服务：

Amazon EKS — Kubernetes 控制平面、Pod 调度、Job 生命周期管理。

Amazon Elastic Compute Cloud（Amazon EC2）p5.48xlarge — 每节点 8× H100 80 GB GPU、32× EFA 适配器。

Elastic Fabric Adapter（EFA）— 3200 Gbps 网络带宽，用于 NCCL all-reduce 操作。

Amazon FSx for Lustre — 用于分布式 checkpoint I/O 的共享 POSIX 文件系统。

Amazon Elastic Container Registry（Amazon ECR）— 训练镜像的容器仓库（PyTorch + NVRx + 模型代码）

在部署此解决方案之前，请确保已准备好以下基础设施和工具：

具有 p5.48xlarge（或 p4de.24xlarge）实例服务配额的 AWS 账户。

Amazon EKS 集群（v1.28+），配备启用 EFA 的自管理 GPU 节点组并已安装 NVIDIA 设备插件。

Amazon FSx for Lustre 文件系统（SCRATCH_2），位于与 GPU 节点相同的可用区。

容器镜像，包含 PyTorch 2.9+、NVRx 0.4.1，以复现本文所示的基准测试结果。使用 v0.6.0 和更新的启动器配置进行当前部署，并将你的训练代码推送到 Amazon ECR。

为你的集群配置 kubectl。

HuggingFace 账户，有权访问 meta-llama/Llama-3.1-8B（或你选择的模型）

训练数据集预下载到共享存储（我们使用 C4 数据集的 100K 样本。你可以使用你选择的数据集）

有关使用 GPU 节点和 EFA 网络创建 EKS 集群的信息，请参阅 awsome-distributed-ai/1.architectures 中的基础设施指南。有关完整的 NVRx 特定设置（包括 Terraform 模块、容器构建和数据集准备），请参阅 NVRx 测试用例 README。

NVRx 对标准 PyTorch FSDP 脚本暴露了两个正交变更：异步 checkpointing（写入路径优化）和恢复，恢复有两个独立的层：用于软故障的进程内重启和用于硬故障的 ft_launcher 进程内重启。测试用例为每个功能提供单独的脚本，因此你可以只采用你需要的内容。

我们从带有同步 checkpointing 的基线 FSDP 脚本开始，然后分别引入每个 NVRx 功能（异步 checkpointing、进程内重启和 ft_launcher 进程内重启），最后简要说明恢复层如何组合。

基线 FSDP 训练脚本

从一个最小的 FSDP 循环开始，使用 torchrun 启动，使用同步分布式 checkpoint 保存（底层使用 torch.distributed.checkpoint.save）。保存会阻塞每个 rank，直到 per-rank 分片落到共享存储上，而 worker 崩溃会导致整个训练作业终止，需要从上一个 checkpoint 完全重启。

# Baseline FSDP training --- torchrun 启动；同步 checkpoint 阻塞循环。
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")

model, _ = create_model(args.model_name, args.torch_dtype)
model = wrap_model(model, "fsdp", local_rank, args.model_name)
optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
data_iter = iter(create_dataloader(args, tokenizer, rank, world_size))

for step in range(1, args.max_steps + 1):
    loss = train_step(model, next(data_iter), optimizer)  # fwd / bwd / step

if step % args.checkpoint_interval == 0:
        # save_checkpoint() 在底层使用 dcp.save() 用于 FSDP --- 集合操作。
        # 所有 rank 在此阻塞，直到 per-rank 分片落到共享存储上。
        save_checkpoint(model, optimizer, step,
                        args.checkpoint_path, rank, "fsdp")

异步 checkpointing：将 checkpoint I/O 与训练解耦

图 2：NVRx 异步 checkpoint 流水线

用 NVRx TorchAsyncCheckpoint 替换同步保存：在 persistent_queue=True 时实例化一次，在 torch.save 处调用 async_save(state_dict, path)，在作业退出时调用一次 finalize_async_save(blocking=True) 作为阻塞式最终化（如图 2 所示）。后台进程拥有实际写入权。主线程继续到下一个前向/反向步骤。每个 rank 通过 FSDP LOCAL_STATE_DICT 写入自己的分片，无需 all-gather，也无 rank-0 瓶颈。

import torch
import torch.distributed as dist
from nvidia_resiliency_ext.checkpointing.async_ckpt.torch_ckpt import TorchAsyncCheckpoint

# 在模型/优化器设置完成后初始化一次异步检查点管理器
async_ckpt = TorchAsyncCheckpoint(persistent_queue=True)

# 训练循环
for step, batch in enumerate(dataloader):
    loss = model(batch)
    loss.backward()
    optimizer.step()

    # 每 N 步保存一次检查点
    if step % checkpoint_interval == 0:
        state_dict = build_state_dict(model, optimizer, step)  # FSDP LOCAL_STATE_DICT，CPU 暂存
        # torch.save(state_dict, path)  # 原来是阻塞的
        async_ckpt.async_save(state_dict, path)  # 立即返回

# 作业退出时清空所有进行中的保存
async_ckpt.finalize_async_save(blocking=True)

用 inprocess.Wrapper 包装训练函数。该包装器拥有重启循环：被包装函数中的异常（或看门狗检测到的挂起）触发重新进入，而不是让进程崩溃。构造函数接入了四个关注点：超时（soft_timeout、hard_timeout、barrier_timeout、completion_timeout）、健康检查（CudaHealthCheck + FaultCounter）、限制总重启次数并设置最小存活 world size 的 RetryController，以及保持 world 连续的 rank 分配（ActivateAllRanks + ShiftRanks）。

图 3：NVRx 进程内重启架构

import nvidia_resiliency_ext.inprocess as inprocess
from nvidia_resiliency_ext.inprocess import CallWrapper

def train_with_inprocess_restart(args, restart_metrics, base_store=None,
                                 call_wrapper: CallWrapper = None):
    # 重新初始化 dist、重建 FSDP、加载最新检查点、训练
    # 每步调用 call_wrapper.ping()
    # call_wrapper.iteration 告诉你当前是第几次重启
    ...

wrapped_train = inprocess.Wrapper(
    store_kwargs = {"host_name": master_addr, "port": master_port + 1},
    soft_timeout = datetime.timedelta(seconds=args.soft_timeout_seconds),
    hard_timeout = datetime.timedelta(seconds=args.hard_timeout_seconds),
    barrier_timeout = datetime.timedelta(seconds=args.barrier_timeout_seconds),
    completion_timeout = datetime.timedelta(seconds=args.barrier_timeout_seconds),
    health_check = inprocess.Compose(
        inprocess.health_check.CudaHealthCheck(),
        inprocess.health_check.FaultCounter(max_rank_faults=20)),
    initialize = inprocess.initialize.RetryController(
        max_iterations=args.max_restarts, min_active_world_size=1),
    rank_assignment = inprocess.Compose(
        inprocess.rank_assignment.ActivateAllRanks(),
        inprocess.rank_assignment.ShiftRanks()),
)(train_with_inprocess_restart)

wrapped_train(args, restart_metrics, base_store)

其内部原理（如图 3 所示）：ProgressWatchdog（progress_watchdog.py:49）使用 Py_AddPendingCall 在字节码指令之间写入时间戳，因此即使 NCCL（一个 C 扩展）内部发生挂起也能被检测到。MonitorThread（monitor_thread.py:124）轮询 inprocess TCPStore 以获取中断标志，检测到时向主线程抛出 RankShouldRestart。随后 AbortTorchDistributed（abort.py:62）收集 Flight Recorder 跟踪信息，终止 NCCL 后端，并销毁进程组。健康检查选出存活者。Python 进程保持存活，只有分布式进程组被重建。

进程内重启（ft_launcher）

将启动器从 torchrun 切换到 ft_launcher。它理解相同的 rendezvous 参数，并添加了一些自己的参数来控制重启行为：

# --- 之前 ---
torchrun --nnodes=$NNODES --nproc_per_node=$GPU_PER_NODE \
    --rdzv-backend=c10d --rdzv-endpoint=$MASTER_ADDR:29500 \
    train.py [args]

# --- 之后 ---
ft_launcher --nnodes=$NNODES --nproc_per_node=$GPU_PER_NODE \
    --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29500 \
    --max-restarts=20 --ft-restart-policy=any-failed \
    --ft-rank-heartbeat-timeout=900 \
    --ft-initial-rank-heartbeat-timeout=1200 \
    --monitor-interval=5 \
    train_ft_launcher.py [args]

在训练脚本内部，在分布式初始化后实例化一次 RankMonitorClient，并在每步发送一次心跳：

import nvidia_resiliency_ext.fault_tolerance as fault_tolerance

ft_client = fault_tolerance.RankMonitorClient()
ft_client.init_workload_monitoring()  # 在 dist 初始化后调用一次

for step in range(1, args.max_steps + 1):
    loss = train_step(model, next(data_iter), optimizer)
    ft_client.send_heartbeat()  # 向 RankMonitorServer 发送存活信号

将 --ft-rank-heartbeat-timeout 设置为应用程序心跳之间最长合法间隔以上。示例使用 900 秒。将 --ft-initial-rank-heartbeat-timeout=1200（20 分钟）以容纳首次模型加载。见 kubernetes/training-job-ft-launcher.yaml。

图 4：NVRx 进程内重启（ft_launcher）架构

其内部原理（如图 4 所示）：每个 rank 的 RankMonitorServer 根据上述超时时间跟踪心跳间隔。超时时，启动器向存活者发送 SIGTERM（对落后者发送 SIGKILL），回收 GPU 内存，重新进行 rendezvous，并生成新的 worker。Worker 启动时从最新检查点重新加载。检查点频率限制了丢失的工作量。

各层如何覆盖不同的故障类别

两层恢复机制针对不同的故障模式进行了划分：进程内机制捕获一个 Python 进程内可以处理的问题（瞬态异常、看门狗可见的 NCCL 挂起），ft_launcher 捕获在操作系统层面杀死进程或导致挂起的问题（SIGKILL、OOM、子 Python 死锁），集群调度器则捕获节点丢失。选择爆炸半径与故障模式相匹配的层。异步检查点是正交的：它可以与一个恢复层（或无恢复层）配合使用，并限制丢失工作的爆炸半径。

部署和运行实验

我们使用围绕 kubectl 的精简包装脚本来部署训练任务，该脚本处理清单模板化、任务清理和环境变量替换。环境变量定义了实例类型、GPU 数量、EFA 设备数量和其他硬件特定参数，因此通过更改单个配置文件即可在不同 GPU 类型（p5、p4de）上运行相同的训练代码和清单。为了在相同条件下比较恢复机制，我们使用确定性故障注入，使用带种子的 RNG 预先生成恰好 N 个固定训练步和 rank 的故障：

--fault_count=5 --fault_seed=42 --fault_types=exception,hang --fault_type_weights=0.6,0.4

相同的种子在实验之间产生相同的故障计划，使得可以直接比较基线 K8s 重启、ft_launcher 和 NVRx 进程内重启。我们对每种机制运行相同的 5 故障模式，然后单独运行异步与同步检查点（无故障注入），以隔离检查点开销。

本节介绍两个关键能力的基准测试结果：异步检查点和故障恢复。

我们在 LLaMA-3.1-8B FSDP 训练上对异步（NVRx）与同步（torch.save）检查点进行了基准测试，涵盖两个维度：从 2 节点扩展到 8 节点（16-64 个 H100 GPU），检查点间隔固定为每 1000 步一次；以及在 8 节点规模下扫描检查点频率（每 100 到 1000 步）。

图 5：同步与异步检查点对比

异步检查点在每个规模下都保持了 99%+ 的训练效率（2 节点 99.2%，8 节点 99.8%），而同步检查点保持在 57-61% 的水平，无论节点数量如何都会因阻塞 I/O 损失约 40% 的墙钟时间。差距持续存在是因为检查点开销是存储绑定的：FSx Lustre 写入时间（约 275 秒）在 16 个或 64 个 GPU 时是恒定的。异步通过将 I/O 与下一个训练段重叠完全隐藏了此延迟。

当我们在 8 节点规模下扫描检查点频率时，异步优势变得更加明显：

图 6：检查点频率的影响

在每 1000 步时，异步达到 99.8% 的效率，而同步仅为 60.3%。在每 100 步时，同步骤降至 14.7%，而异步则优雅地退化到 29.6%，仍然比同步高效 2 倍。交叉点是检查点之间训练差距（约 100 步时约 280 秒）接近 FSx I/O 时间（约 275 秒）的地方。高于该阈值时，异步完全重叠。低于该阈值时，部分重叠仍然优于阻塞写入。

这意味着什么：大规模训练的团队通常会减少检查点频率以避免阻塞惩罚，但在发生故障时会承受数小时的进度损失。异步检查点消除了这种权衡：你可以高频率（每隔几分钟）做检查点，几乎零开销，既减少了回滚距离又不牺牲 GPU 利用率。

我们在 2 个 p5.48xlarge 节点（共计 16 块 H100 GPU）上运行了故障恢复实验，使用 FSDP 训练 LLaMA-3.1-8B，每 500 步检查一次，共运行 2000 步。我们在每次训练中注入 5 个确定性故障（相同种子、相同故障调度），并比较了三种恢复机制：

图 7：故障恢复方法对比：进程内重启 vs 作业内重启 vs Kubernetes 重启

NVRx 进程内重启实现了 31% 的训练吞吐量（in-process restart achieves 31% training goodput），基础设施吞吐量达 87%，恢复时间约 10 秒，零容器重启。ft_launcher 实现了 25.5% 的训练吞吐量（85.9% 基础设施吞吐量），每次故障恢复时间 17 秒。基线 Kubernetes 重启则是灾难性的：11.5% 的训练吞吐量、35.8% 的基础设施吞吐量，每次故障恢复时间 270 秒。基线的失败模式是级联故障：一个 rank 崩溃，存活 rank 触发 NCCL 超时（60 秒），然后两个 pod 异步重启，触发 CrashLoopBackOff 和额外的超时循环，之后才能恢复训练。

这意味着什么：在大规模训练（数百到数千块 GPU）场景下，硬件故障每隔几小时就会发生一次。如果使用基线 Kubernetes 恢复，每次故障消耗 4 分钟以上，那么一个每小时发生 3-4 次故障的集群，花在恢复上的时间会比训练时间还多。NVRx 将恢复窗口缩短到秒级，即使故障频率随规模增加而上升，也能保持高 GPU 利用率。

为避免持续产生费用，请删除你创建的资源。顺序很重要：挂载 FSx 的 Kubernetes 对象必须先于文件系统本身删除，否则删除操作会被卷终结器（finalizers）阻塞。

删除训练任务。

./deploy.sh --delete training-job-inprocess.yaml
kubectl get jobs -n nvrx-training

释放持久卷声明。

kubectl delete -f fsx-storage.yaml

PersistentVolume 使用 persistentVolumeReclaimPolicy: Retain，因此只会分离文件系统而不会删除它，在第 4 步删除 FSx 之前它会继续计费。先把需要保留的结果从 /checkpoints 复制出来。删除 SCRATCH_2 文件系统是永久性的，不会做最终备份。

将 GPU 节点组缩容至零。

aws autoscaling update-auto-scaling-group \
    --auto-scaling-group-name <your-gpu-asg-name> \
    --min-size 0 --max-size 0 --desired-capacity 0

按需实例在终止后停止计费。对于 EC2 Capacity Blocks，由于预订费用是预付的，提前释放不会退还任何费用，但在块结束前至少 30 分钟开始缩容，因为 EC2 会在结束时间之前就开始终止 Capacity Block 实例。

删除文件系统、镜像仓库和集群。

aws fsx delete-file-system --file-system-id <your-fsx-id>
aws ecr delete-repository --repository-name <your-repo-name> --force

然后使用创建时所用的工具删除 EKS 集群和 GPU 节点组（terraform destroy、eksctl delete cluster 或 CloudFormation），以保持基础设施状态一致。

有关上下文中的作业级清理，请参阅测试用例 README 中的停止训练（Stop Training）。

在这篇文章中，我们演示了如何将 NVIDIA Resiliency Extension（NVRx）与 Amazon EKS 相结合，通过缓解两个常见瓶颈来解锁更高效的大规模分布式 PyTorch 训练：检查点阻塞和故障恢复。

在报告的运行中，异步检查点减少了阻塞 I/O，在 16 到 64 块 H100 GPU 规模下测量到 99%+ 的训练效率。对于故障恢复，NVRx 进程内重启可以在约 10 秒内恢复，零容器重启，而 ft_launcher 可以处理导致进程完全崩溃的硬故障（SIGKILL、挂起）。两者都与基线 Kubernetes 恢复相比有质的飞跃——NCCL 超时级联和 CrashLoopBackOff 风暴每次故障消耗 4 分钟以上。

这些功能的组合意味着团队可以高频率做检查点并快速恢复，最大化数据保护和 GPU 利用率。

在 awsome-distributed-ai 仓库中尝试完整的可复现示例。

探索将异步检查点与故障恢复相结合，以同时最大化训练效率，减少回滚距离和恢复时间。

NVIDIA Resiliency Extension

Awsome-distributed-ai 展示 NVRx 与 EKS 的集成

GTC 演讲：构建大规模容错分布式 AI 训练

Amazon Elastic Kubernetes Service（Amazon EKS）

Amazon FSx for Lustre

Amazon Elastic Container Registry（Amazon ECR）

Elastic Fabric Adapter（EFA）

使用 Amazon EKS 和 Torch Distributed Elastic 进行分布式训练