AI 服務暫時不可用，以下為來源正文，待恢復後補全翻譯。
Large-scale distributed training jobs run for hours or days across dozens of nodes. At that scale and duration, interruptions are statistically inevitable: network partitions, memory errors, software exceptions, or infrastructure events will eventually disrupt at least one worker. A single GPU fault triggers a cascade: NVIDIA Collective Communication Library (NCCL) timeouts propagate to healthy workers, pods crash and restart out of sync, and your cluster burns expensive GPU hours while making zero training progress. Synchronous checkpointing adds a second source of idle time: every save blocks all ranks on I/O, which on the cluster sizes in this post consumed up to 40% of total wall time. In this post, we show how to integrate NVIDIA Resiliency Extension (NVRx) into PyTorch Fully Sharded Data Parallel (FSDP) training on Amazon Elastic Kubernetes Service (Amazon EKS) to solve both problems. You walk through async checkpointing that overlaps I/O with training, in-process restart that recovers from faults in seconds without touching the container lifecycle, and in-job restart using ft_launcher for automatic worker respawn on hard crashes. We include benchmark results on H100 GPUs at 2-node to 8-node scale, with all code available to reproduce. Solution overview The solution combines NVRx fault tolerance primitives with an EKS-based training environment designed for high-performance multi-node GPU workloads. NVRx handles the application-level resilience (async checkpointing, in-process restart, and in-job restart), while the EKS cluster provides the infrastructure foundation: GPU scheduling, high-bandwidth networking, and shared storage for checkpoint persistence. NVRx The NVIDIA Resiliency Extension (NVRx) is a pip-installable Python layer (pip install nvidia-resiliency-ext) that adds fault-tolerance primitives to PyTorch: no custom kernels, no PyTorch fork, no recompile. The primitives drop into an existing FSDP script as ordinary imports. The model and training code stay untouched. Each is independently adoptable. We exercised three features: async checkpointing, in-process restart, and ft_launcher (in-job restart). Async checkpointing, exposed through TorchAsyncCheckpoint, replaces torch.save with an async_save() call that hands the state dict to a background process and returns immediately. A matching finalize_async_save() before the next save commits the prior write. Paired with FSDP LOCAL_STATE_DICT, each rank writes its own shard directly, with no all-gather and no rank-0 bottleneck. In-process restart, exposed through inprocess.Wrapper, wraps the train function so a transient fault (an unhandled exception or an NCCL hang) does not kill the Python process. NVRx aborts the active process group, runs health checks per rank (GPU, NVLink, NIC), re-rendezvouses survivors, and re-enters the wrapped function from the latest checkpoint. The interpreter, CUDA allocator, and outer-scope objects survive. This catches the soft-fault class. The ft_launcher binary, the NVRx in-job restart launcher, handles cases in-process cannot catch: SIGKILL, out-of-memory (OOM) kill, and OS-level hangs. Each rank runs a RankMonitorClient. The launcher checks heartbeats against explicit CLI-set timeouts, and on stall or death it kills survivors, reclaims GPU memory, and respawns fresh workers in the same job. Recovered workers reload from the latest checkpoint. Each recovery layer covers a distinct fault class: in-process for soft faults, ft_launcher for hard faults, and the cluster orchestrator for node loss. The layers are independent. Pick the one whose scope matches your failure modes. Amazon EKS cluster Amazon EKS is a managed Kubernetes service that handles the control plane, upgrades, and API server availability. We run self-managed node groups of p5.48xlarge instances, each with 8 NVIDIA H100 80 GB GPUs and 32 Elastic Fabric Adapter (EFA) network interfaces. Training pods run as Kubernetes Jobs with headless Services for peer discovery, so workers find each other through DNS rather than hardcoded IPs, and pod replacements can rejoin without reconfiguring the job. Each node exposes its GPUs and EFA adapters as extended resources through the NVIDIA device plugin and EFA device plugin. The Kubernetes scheduler places training pods on GPU nodes using node affinity and tolerations, facilitating the full 8-GPU allocation per node. For checkpoint storage, we use Amazon FSx for Lustre (SCRATCH_2, 1.2 TB) mounted into every training pod through the FSx CSI driver. FSx provides the shared filesystem that both async and synchronous checkpointing write to, and critically, it is where recovering workers read their checkpoint state after a fault. Placing FSx in the same Availability Zone as the GPU nodes minimizes read latency during recovery, which matters because checkpoint loading (not the restart mechanism) dominates recovery time at scale. Figure 1: Architecture diagram showing the EKS cluster with 2-8 p5 nodes, EFA interconnect, FSx for Lustre, and NVRx components within training pods The key AWS services involved: Amazon EKS — Kubernetes control plane, pod scheduling, Job lifecycle management. Amazon Elastic Compute Cloud (Amazon EC2) p5.48xlarge — 8x H100 80 GB GPUs, 32x EFA adapters per node. Elastic Fabric Adapter (EFA) — 3,200 Gbps network bandwidth for NCCL all-reduce operations. Amazon FSx for Lustre — Shared POSIX filesystem for distributed checkpoint I/O. Amazon Elastic Container Registry (Amazon ECR) — Container registry for the training image (PyTorch + NVRx + model code) Prerequisites Before deploying this solution, make sure you have the following infrastructure and tooling in place: AWS account with service quota for p5.48xlarge (or p4de.24xlarge) instances. Amazon EKS cluster (v1.28+) with EFA-enabled self-managed GPU node groups and the NVIDIA device plugin installed. Amazon FSx for Lustre filesystem (SCRATCH_2) in the same Availability Zone as the GPU nodes. Container image with PyTorch 2.9+, NVRx 0.4.1 to reproduce the benchmark results shown in this post. Use v0.6.0 with an updated launcher configuration for a current deployment and your training code, pushed to Amazon ECR. kubectl configured for your cluster. HuggingFace account with access to meta-llama/Llama-3.1-8B (or your model of choice) Training dataset pre-downloaded to shared storage (we use 100K samples from the C4 dataset. You can use a dataset of your choice) For EKS cluster creation with GPU nodes and EFA networking, see the infrastructure guides in awsome-distributed-ai/1.architectures. For the complete NVRx-specific setup including Terraform modules, container build, and dataset preparation, see the NVRx test case README. Solution walkthrough NVRx exposes two orthogonal changes to a standard PyTorch FSDP script: async checkpointing (a write-path optimization) and recovery, which comes in two independent layers: in-process restart for soft faults and ft_launcher in-job restart for hard faults. The test case ships a separate script per capability, so you adopt only what you need. We start from a baseline FSDP script with synchronous checkpointing, then introduce each NVRx capability independently (async checkpointing, in-process restart, and the ft_launcher in-job restart) and close with a short note on how the recovery layers compose. Baseline FSDP training script Start with a minimal FSDP loop launched by torchrun using a synchronous distributed checkpoint save (PyTorch’s torch.distributed.checkpoint.save under the hood). The save blocks every rank until per-rank shards land on shared storage, and a worker crash kills the entire training job, requiring a full restart from the last checkpoint. # Baseline FSDP training --- torchrun launches; sync checkpoints block the loop. 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() uses dcp.save() under the hood for FSDP --- collective. # All ranks block here until per-rank shards land on shared storage. save_checkpoint(model, optimizer, step, args.checkpoint_path, rank, "fsdp") Async checkpointing: decouple checkpoint I/O from training Figure 2: NVRx async checkpoint pipeline Replace the synchronous save with the NVRx TorchAsyncCheckpoint: instantiate it once with persistent_queue=True, call async_save(state_dict, path) in place of torch.save, and call finalize_async_save(blocking=True) once as a blocking finalization at job exit, as shown in figure 2. A background process owns the actual write. The main thread continues to the next forward/backward step. Each rank writes its own shard through FSDP LOCAL_STATE_DICT, with no all-gather and no rank-0 bottleneck. import torch import torch.distributed as dist from nvidia_resiliency_ext.checkpointing.async_ckpt.torch_ckpt import TorchAsyncCheckpoint # Initialize the async checkpoint manager once after model/optimizer setup. async_ckpt = TorchAsyncCheckpoint(persistent_queue=True) # Training loop for step, batch in enumerate(dataloader): loss = model(batch) loss.backward() optimizer.step() # Checkpoint every N steps if step % checkpoint_interval == 0: state_dict = build_state_dict(model, optimizer, step) # FSDP LOCAL_STATE_DICT, CPU-staged # torch.save(state_dict, path) # was blocking async_ckpt.async_save(state_dict, path) # returns immediately # Drain any in-flight save at job exit. async_ckpt.finalize_async_save(blocking=True) In-process restart Wrap the train function with inprocess.Wrapper. The wrapper owns the restart loop: an exception in the wrapped function (or a hang the watchdog detects) triggers a re-entry instead of crashing the process. The constructor wires four concerns: timeouts (soft_timeout, hard_timeout, barrier_timeout, completion_timeout), health checks (CudaHealthCheck + FaultCounter), a RetryController that caps the total number of restarts and sets a minimum surviving world size, and rank assignment (ActivateAllRanks + ShiftRanks) that pulls survivors left to keep the world contiguous. Figure 3: NVRx in-process restart architecture 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): # Re-init dist, rebuild FSDP, load latest checkpoint, train. # call_wrapper.ping() each step; # call_wrapper.iteration tells you which restart you're on. ... 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) Behind the scenes (as shown in figure 3): a ProgressWatchdog (progress_watchdog.py:49) uses Py_AddPendingCall to wri [truncated for AI cost control]