GPU 训练失败并不一定意味着整个任务要从头开始。将 NVIDIA Resiliency Extension（NVRx）接入运行在 Amazon EKS 上的 PyTorch FSDP 训练后，可以把故障恢复从“重建整个 Job、重新加载全部状态”缩短到秒级，同时通过异步 checkpoint 把存储 I/O 与训练计算重叠起来。
这套方案的核心不只是保存模型，而是把训练任务设计成可以在进程内或作业内恢复：正常运行时持续产生可恢复的 checkpoint；发生 GPU 故障时，优先尝试 in-process restart；如果进程无法继续，再由 ft_launcher 在同一个训练 Job 内重启训练进程。
三个能力分别解决什么问题
异步 checkpoint：不要让保存动作卡住训练
传统 checkpoint 通常由训练主线程执行：收集 FSDP 状态、序列化、写入共享存储，整个过程中 GPU 可能处于等待状态。模型和优化器状态越大，这个停顿越明显。
异步 checkpoint 的思路是把“生成可保存状态”和“写入持久化存储”拆开。训练进程继续推进下一个 step，后台线程或进程负责把 checkpoint 写入 Amazon S3、FSx for Lustre 或其他持久化存储。实际部署时，需要确认底层存储具备足够的吞吐，并避免多个 worker 同时写入同一个目标路径。
一个实用的 checkpoint 目录可以按训练任务、全局 step 和版本组织：
s3://my-training-checkpoints/
└── job-20250301/
├── step-000100/
├── step-000200/
├── step-000300/
└── latest.json
latest.json 只应该在一个完整 checkpoint 写入成功后更新。这样即使节点在写盘期间崩溃，恢复逻辑也不会误选到半成品目录。
In-process restart：优先在当前 Job 内恢复
如果故障可以被训练进程和 NVRx 的故障处理逻辑捕获，in-process restart 可以避免重新申请 Kubernetes Job。恢复时通常需要完成以下动作：
- 停止或隔离失效的通信组；
- 清理旧的 CUDA、NCCL 和分布式状态；
- 从最近的完整 checkpoint 恢复模型、优化器、学习率调度器和随机数状态；
- 重新建立 FSDP 和 distributed process group；
- 从保存的全局 step 继续训练。
恢复粒度取决于 checkpoint 中保存了什么。如果只保存模型权重而没有保存 optimizer、scheduler、随机数和数据迭代器状态，任务可能可以继续，但训练轨迹和严格可复现性会发生变化。
ft_launcher：进程无法自救时，在 Job 内重启
更严重的 GPU 故障可能会让 Python 进程直接退出。这时可以使用 NVRx 的 ft_launcher 管理训练入口，在同一个 Kubernetes Job 中重新启动训练进程。它的价值在于把“训练进程失败”和“训练任务永久失败”分开处理。
容器入口不要直接调用训练脚本，而是让 launcher 成为主进程。下面的命令示例假设镜像中已经安装了 NVRx，并且训练脚本位于 /workspace/train.py：
#!/usr/bin/env bash
set -euo pipefail
export MASTER_ADDR="${MASTER_ADDR:?MASTER_ADDR is required}"
export MASTER_PORT="${MASTER_PORT:-29500}"
export CHECKPOINT_DIR="${CHECKPOINT_DIR:-/shared/checkpoints/job-20250301}"
# 参数名可能随 NVRx 版本变化；请以镜像中 ft_launcher --help 为准。
exec ft_launcher \
--nproc-per-node="${NPROC_PER_NODE:-8}" \
--max-restarts="${MAX_RESTARTS:-10}" \
/workspace/train.py \
--config /workspace/configs/fsdp.yaml \
--checkpoint-dir "$CHECKPOINT_DIR"
运行前建议在目标镜像中确认入口和参数：
python -m pip show nvidia-resiliency-ext
ft_launcher --help
上面的 launcher 参数是部署示例，不应直接假定适用于所有 NVRx 版本。升级 NVRx 或基础镜像时，应把 --help 输出纳入集成测试，避免参数变化导致 Pod 启动即失败。
在 EKS 中安排训练 Pod
Kubernetes 层需要解决两个问题：让多个 Pod 能够互相发现并建立分布式通信，以及让 checkpoint 在 Pod 重建后仍然可见。下面是一个可改造的 StatefulSet 示例。它使用 headless Service 提供稳定 DNS，并把共享 checkpoint 路径挂载到 /shared/checkpoints。
apiVersion: v1
kind: Service
metadata:
name: fsdp-train
spec:
clusterIP: None
selector:
app: fsdp-train
ports:
- name: master
port: 29500
targetPort: 29500
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
name: fsdp-train
spec:
serviceName: fsdp-train
replicas: 2
selector:
matchLabels:
app: fsdp-train
template:
metadata:
labels:
app: fsdp-train
spec:
restartPolicy: Always
containers:
- name: trainer
image: <account>.dkr.ecr.<region>.amazonaws.com/fsdp-nvrx:latest
command: ["/bin/bash", "/workspace/entrypoint.sh"]
env:
- name: MASTER_ADDR
value: "fsdp-train-0.fsdp-train"
- name: MASTER_PORT
value: "29500"
- name: NNODES
value: "2"
- name: NODE_RANK
valueFrom:
fieldRef:
fieldPath: metadata.name
- name: CHECKPOINT_DIR
value: "/shared/checkpoints/job-20250301"
resources:
requests:
nvidia.com/gpu: "8"
limits:
nvidia.com/gpu: "8"
volumeMounts:
- name: checkpoints
mountPath: /shared/checkpoints
volumes:
- name: checkpoints
persistentVolumeClaim:
claimName: fsdp-checkpoints
示例中的 NODE_RANK 仍需要在 entrypoint 中从 StatefulSet Pod 名称解析为整数，或改用 Job/自定义控制器注入 rank。真实环境还需要补充 GPU 节点选择器、Topology Manager、EFA/NCCL 配置、Pod 反亲和性以及适合训练规模的存储类。若 checkpoint 使用 S3，则可以把写入逻辑放在训练代码中，并通过 IAM Roles for Service Accounts（IRSA）授予最小权限，而不是把 AWS 密钥放进 Secret。
FSDP checkpoint 设计要点
FSDP 会把参数和梯度分片到不同 rank。checkpoint 设计必须明确保存格式和恢复方式：
- 需要快速恢复训练时，保存模型、optimizer、scheduler、全局 step 和必要的随机数状态；
- 使用与 FSDP 配置匹配的 state-dict 类型，避免保存端和恢复端的分片语义不一致；
- 由 rank 0 或专门的 checkpoint 协调逻辑负责提交最终 manifest；
- 写临时目录，完成后再原子地标记为可用；
- 清理旧 checkpoint，但至少保留两个可恢复版本，防止最新版本损坏。
下面是一个与框架无关的恢复流程示例。它不是 NVRx 的固定 API，而是训练代码应具备的最小控制结构：
from pathlib import Path
import torch
def find_latest_checkpoint(root: str) -> Path | None:
candidates = sorted(Path(root).glob("step-*/READY"))
return candidates[-1].parent if candidates else None
def restore_training_state(model, optimizer, scheduler, root: str) -> int:
checkpoint = find_latest_checkpoint(root)
if checkpoint is None:
return 0
state = torch.load(checkpoint / "training_state.pt", map_location="cpu")
model.load_state_dict(state["model"])
optimizer.load_state_dict(state["optimizer"])
scheduler.load_state_dict(state["scheduler"])
return int(state["global_step"])
生产代码需要根据 FSDP 的 sharded state-dict 配置替换 load_state_dict 细节，并处理 rank 间同步。这里刻意使用 READY 标记，避免恢复逻辑把未完成的异步写入当成有效 checkpoint。
故障注入比“成功跑完”更重要
训练效率不能只看无故障时的 samples/sec。应该至少覆盖以下测试：
# 查看训练 Pod 与 GPU 状态
kubectl get pods -l app=fsdp-train -o wide
kubectl logs -f statefulset/fsdp-train -c trainer
# 在测试集群中模拟一个训练 Pod 被删除
kubectl delete pod fsdp-train-1
# 检查 Pod 是否重新加入，以及是否从最近 checkpoint 继续
kubectl get pods -l app=fsdp-train -w
kubectl logs fsdp-train-1 -c trainer | grep -E "restart|checkpoint|global_step"
测试指标至少包括：恢复耗时、丢失的训练 step、checkpoint 写入期间的吞吐、重启次数、最终 loss 曲线是否异常，以及失败后是否留下需要人工清理的 NCCL 或临时资源。对于 H100 集群，摘要中的 2 到 8 节点基准显示，在合理的异步 checkpoint 和恢复配置下，训练效率可以达到 99% 以上，恢复时间达到秒级；但这个结果不应直接外推到所有网络、存储和模型规模。
落地时的检查清单
- 先在单节点上验证 FSDP checkpoint 的保存和恢复，再扩展到多节点；
- 固定并记录 PyTorch、CUDA、NCCL、NVRx、GPU 驱动和容器版本；
- 使用
ft_launcher --help验证当前镜像支持的参数，不复制其他版本的命令行； - 确认 checkpoint 存储在 Pod 重启后仍可访问，并具备足够的 I/O 吞吐；
- 让 checkpoint 具备完整性标记和保留策略；
- 通过故障注入验证 in-process restart 和 ft_launcher restart 两条路径；
- 监控恢复次数、恢复时长、checkpoint 延迟和训练有效吞吐；
- 把“能够恢复”与“严格复现”分开验收，尤其是数据迭代器和随机数状态。
NVRx 的价值不在于让 GPU 故障消失，而在于把故障从灾难性事件变成可观测、可恢复的运行状态。对于长时间运行的 FSDP 任务，异步 checkpoint 提供持续保护，in-process restart 提供低开销恢复，ft_launcher 则为进程级崩溃提供兜底。三者配合，才构成适合 EKS 生产训练的容错闭环。