基于 NVRx 在 Amazon EKS 上实现具备容错能力的分散式训练
- 作者
- 姓名
- Nino
- 职业
- Senior Tech Editor
在大规模超大语言模型 (LLM) 训练任务中，集群通常包含数百甚至数千张 GPU。随着节点数量增加，硬件故障的出现几乎无法避免。在传统的分布式训练架构中，一旦出现节点故障，受限于同步检查点 (Checkpoint) 的写入延迟以及繁重的集群级重新初始化流程，整体训练吞吐量会受到严重影响。
通过在托管于 Amazon Elastic Kubernetes Service (Amazon EKS) 的 PyTorch Fully Sharded Data Parallel (FSDP) 训练工程中集成 NVIDIA Resiliency Extension (NVRx)，技术团队能够成功解耦 I/O 阻塞瓶颈，并在数秒内完成 GPU 硬件故障恢复。在 2 至 8个节点的 NVIDIA H100 Tensor Core GPU 集群上，该方案成功实现了 99%+ 的训练效率。
无论是在本地集群中进行大规模模型微调，还是通过 n1n.ai 等聚合 API 平台快速调用基础大模型进行性能对比与评估，底层计算架构的稳定性都是确保 AI 项目按时交付的关键所在。
分布式训练的不稳定性瓶颈分析
在跨 Kubernetes 节点训练 70B+ 参数规模的模型时，主要有两个核心痛点影响 GPU 的实际利用率：
- 同步检查点阻塞 (Synchronous Checkpoint Latency)：传统 PyTorch 检查点机制要求所有 GPU 计算 Worker 暂停当前的 Forward/Backward 迭代，等待模型权重和优化器状态完全写入远程存储（如 Amazon S3 或 AWS FSx for Lustre）。随着模型参数体量飙升至几百 GB，这种同步等待会浪费大量的 GPU 算力。
- 冷启动恢复耗时极长 (Cold-Restart Delay)：当发生无法自动纠正的 GPU 显存错误或 NVLink 通信异常时，常规 Kubernetes Pod 会直接被终止。重建容器、重新拉取镜像、建立 TCP/NCCL 通信环路并将最新的检查点文件重新加载回 GPU 显存，通常需要 10 到 30分钟。
NVRx 专门针对上述瓶颈进行了重构，推出了异步非阻塞检查点 (Async Checkpointing)、进程内 Rank 恢复 (In-Process Restart) 以及 作业级重启协调器 (ft_launcher)。
架构设计：Amazon EKS 上的 NVRx + PyTorch FSDP
为了消除等待时间，NVRx 将持久化过程划分为两个阶段：极速内存暂存（或本地临时存储副本）与后台多线程异步传输。此外，NVRx 深入 Hook 了 PyTorch 的分布式进程组，使得在不销毁父级 Python 容器环境的前提下，动态重建 Rank 链接。
+-----------------------------------------------------------------------------------+
| Amazon EKS 集群 |
| |
| +----------------------------------+ +----------------------------------+ |
| | EKS Worker Node 1 (8x H100) | | EKS Worker Node 2 (8x H100) | |
| | | | | |
| | +----------------------------+ | | +----------------------------+ | |
| | | PyTorch FSDP + NVRx Engine| | NVLink| | PyTorch FSDP + NVRx Engine| | |
| | +----------------------------+ |<----->| +----------------------------+ | |
| | | 内存暂存区 (Host RAM / SHM) | | 互联 | | 内存暂存区 (Host RAM / SHM) | | |
| | +--------------+-------------+ | | +--------------+-------------+ | |
| +-----------------|----------------+ +-----------------|----------------+ |
| | | |
+--------------------|------------------------------------------|-------------------+
v v
+---------------------------------------------------------------------------+
| 后台异步 I/O 线程写入引擎 |
+---------------------------------------------------------------------------+
|
v
+-------------------------------------------------------+
| AWS 共享文件存储服务 (Amazon S3 / FSx for Lustre) |
+-------------------------------------------------------+
在开发复杂的生成式 AI 业务时，企业通常需要在自主训练专属模型与调用 n1n.ai 提供的商业化高可用 LLM API 之间进行权衡与评估，从而选择最具成本效益的方案。
代码实现：在 PyTorch FSDP 中集成 NVRx 异步检查点
在标准 PyTorch Distributed 脚本中引入 NVRx 极其简便。以下代码展示了如何通过异步 Save 逻辑实现零等待的权重落盘：
import os
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
import nvrx.checkpoint.async_engine as nvrx_ckpt
def setup_distributed():
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
return local_rank
def train_with_nvrx():
local_rank = setup_distributed()
# 构建模型并挂载至 GPU
model = BuildLargeTransformerModel().to(local_rank)
model = FSDP(model)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
# 配置 NVRx 异步检查点管理器
async_ckpt_mgr = nvrx_ckpt.AsyncCheckpointManager(
checkpoint_dir="/mnt/fsx/checkpoints