NVRx

NVRx distributed training on AWS

This test case benchmarks NVIDIA Resiliency Extension (NVRx) features for PyTorch distributed training on AWS. NVRx provides fault tolerance and resiliency primitives for large-scale training, integrating as a simple pip install layer on top of standard PyTorch — no custom CUDA builds or kernel modifications required.

NVRx Features Tested

Note: Local checkpointing with NVRx LocalCheckpointManager has been tested with GPT-2 on g5 instances only. FSDP sharded state dicts (ShardedTensor) are not yet compatible with NVRx BasicTensorAwareStateDict on larger models.

Parallelism Support

All training scripts support both FSDP and DDP via the --parallel_strategy argument:

• --parallel_strategy=fsdp (default) — Full Sharded Data Parallel, recommended for large models (LLaMA-3.1-8B)
• --parallel_strategy=ddp — Distributed Data Parallel, suitable for smaller models (GPT-2)

Model Support

Set the model via MODEL_NAME in env_vars (e.g., MODEL_NAME=gpt2 or MODEL_NAME=meta-llama/Llama-3.1-8B).

Fault Injection

The test case includes a configurable fault injection framework (src/failure_simulator.py) that simulates real-world training failures:

• Exception faults — RuntimeError caught by NVRx in-process Wrapper for sub-second recovery
• Hang faults — Simulates NCCL deadlocks, detected by ft_launcher heartbeat timeout

Two injection modes are supported:

• Deterministic (recommended for benchmarks): --fault_count=5 --fault_seed=42 injects exactly N faults at pre-determined steps for reproducible, apples-to-apples comparisons across recovery mechanisms
• Stochastic (for exploratory testing): --fault_probability=0.005 injects faults randomly per step per rank

Prerequisites

To run these tests, you need a training cluster with GPU nodes and shared storage. Instructions for creating a cluster can be found in architectures, the aws-do-eks project, or EKS Blueprints.

1. Build Container Image

From the nvrx/ directory, build a container image with PyTorch 2.9, NVRx 0.4.1, and the training scripts:

cd examples/training/nvrx

export AWS_REGION=$(aws ec2 describe-availability-zones --output text --query 'AvailabilityZones[0].[RegionName]')
export ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
export REGISTRY=${ACCOUNT}.dkr.ecr.${AWS_REGION}.amazonaws.com

docker build -f Dockerfile -t ${REGISTRY}/nvrx-fsdp-training:latest .

2. Push Container Image to Amazon ECR

# Create registry if needed
REGISTRY_COUNT=$(aws ecr describe-repositories | grep \"nvrx-fsdp-training\" | wc -l)
if [ "$REGISTRY_COUNT" == "0" ]; then
    aws ecr create-repository --repository-name nvrx-fsdp-training
fi

# Login and push
aws ecr get-login-password | docker login --username AWS --password-stdin $REGISTRY
docker image push ${REGISTRY}/nvrx-fsdp-training:latest

Alternatively, use the provided build_and_push.sh script which handles ECR login, repository creation, and pushing in one step:

cp env_vars.template env_vars
# Edit env_vars with your settings
source env_vars
./build_and_push.sh

Set BUILD_HOST in env_vars to build remotely via SSH (e.g., on an EC2 instance with Docker). Leave it empty for local builds.

3. Prepare Dataset (Optional)

Training scripts support two dataset modes:

Option A: Stream from HuggingFace Hub (default) — No setup required. Works well for experiments without frequent restarts.

Option B: Pre-download to local storage (recommended for fault recovery experiments) — Eliminates HuggingFace API rate limiting (429 errors) during rapid restarts:

# Run inside the training container on your cluster
python prepare_dataset.py --output_path /checkpoints/c4_subset --num_samples 100000

This downloads 100K samples (~227 MB) in about 20 seconds. Then add --dataset_path=/checkpoints/c4_subset to the training script arguments. See the platform-specific README for detailed instructions on running this on your cluster.

Platform-Specific Instructions

• Amazon EKS: See kubernetes/README.md
• Slurm: Coming soon

Source Files

References

• NVIDIA Resiliency Extension (NVRx) GitHub
• NVRx Documentation
• NVRx on PyPI
• PyTorch FSDP Tutorial
• Amazon EKS User Guide
• AWS EFA Documentation