Building Blocks for Foundation Model Training and Inference on AWS

Building Blocks for Foundation Model Training and Inference on AWS

Foundation model scaling has evolved from a single focus on pre-training compute to a three-pronged approach encompassing pre-training, post-training (SFT and RL), and test-time compute (search and verification). This shift necessitates a convergent infrastructure strategy based on tightly coupled accelerator compute, high-bandwidth low-latency networking, and distributed storage backends.

Infrastructure: Compute, Network, and Storage

AWS provides a tiered infrastructure designed to minimize bottlenecks in collective communication and memory movement, which often dominate step time more than raw compute throughput.

Accelerated Compute

AWS offers several generations of NVIDIA GPUs via the Amazon EC2 P instance family:

  • P5 Instances: Feature NVIDIA H100 GPUs (p5.48xlarge) and H200 GPUs (p5e.48xlarge/p5en.48xlarge).
  • P6 Instances: Introduce NVIDIA Blackwell B200 (p6-b200.48xlarge) and B300 (p6-b300.48xlarge) architectures.

Networking and Interconnects

GPU communication is split into two regimes to optimize data movement:

  • Internal Scale-up (NVLink/NVSwitch): High-bandwidth, low-latency connectivity within a single node.
  • External Scale-out (EFA): The Elastic Fabric Adapter (EFA) provides OS-bypass RDMA using the Scalable Reliable Datagram (SRD) protocol for cross-node communication. EFAv3 reduces packet latency by ~35% over EFAv2, and EFAv4 provides an additional 18% improvement in collective communication performance over EFAv3.
  • UltraServers: Amazon EC2 UltraServers (e.g., P6e-GB200) extend the NVLink domain beyond a single instance, allowing up to 72 Blackwell GPUs to exist within one NVLink domain, reducing the frequency with which communication must leave the NVLink fabric.

Tiered Storage

To handle multi-terabyte checkpoints and massive corpora, AWS employs a three-tier hierarchy:

  1. Local NVMe SSD: Ephemeral instance store for hot data (up to 30.72 TB raw capacity).
  2. Amazon FSx for Lustre: A managed parallel file system providing high aggregate throughput and sub-millisecond latencies.
  3. Amazon S3: Used for durable persistence and lazy loading of datasets via Data Repository Associations.

Resource Orchestration: Slurm and Kubernetes

Managing thousands of accelerators requires centralized control planes to ensure atomic job scheduling and resource efficiency.

Slurm (HPC-centric)

Slurm is the dominant manager for HPC workloads, scheduling at the job level to ensure all required nodes are allocated before tasks launch. AWS supports Slurm through:

  • AWS ParallelCluster: An open-source tool for automating Slurm cluster deployment.
  • AWS Parallel Computing Service (PCS): A managed control plane for Slurm.
  • Amazon SageMaker HyperPod: Adds continuous node health monitoring and job auto-resume functionality to Slurm mode.

Kubernetes (Cloud-native)

While Kubernetes excels at deployment, it lacks native job-level atomicity and topology awareness. These gaps are filled by:

  • Kueue: Manages job-level gang admission and multi-tenant quotas.
  • Volcano and NVIDIA KAI Scheduler: Provide topology-aware pod placement to optimize NVLink and EFA usage.
  • SageMaker HyperPod (EKS mode): Integrates managed Kueue and Karpenter for just-in-time provisioning. It also introduces checkpointless training, which uses peer-to-peer state replication via EFA to recover from failures without reading multi-terabyte checkpoints from storage.

The ML Software Stack

Performance is determined by a five-layer stack ranging from hardware drivers to high-level frameworks.

Low-Level Enablement and Runtimes

  • Kernel Drivers: NVIDIA GPU drivers support GPUDirect RDMA, while the EFA driver provides OS-bypass networking via libfabric.
  • CUDA and Kernels: CUDA Toolkit 13.x supports the Blackwell architecture. Performance is further driven by fused kernels like FlashAttention and programmable toolchains such as Triton and NVIDIA CuTe.

Communication Substrate

  • NCCL: The NVIDIA Collective Communications Library implements topology-aware algorithms for all-reduce and all-gather. For Mixture-of-Experts (MoE) models, all-to-all collectives are critical for routing tokens between experts.
  • aws-ofi-nccl: A plugin that maps NCCL transport APIs to libfabric, enabling NCCL to use EFA's SRD protocol.
  • NIXL: The NVIDIA Inference Xfer Library provides point-to-point transfers for disaggregated inference architectures (separating prefill and decode phases).

Frameworks

  • PyTorch: The primary framework for distributed workloads, utilizing torch.distributed and FSDP2 to shard parameters and optimizer states.
  • Distributed Frameworks:
    • Hugging Face Transformers/Accelerate: Prioritizes ease of use and compatibility.
    • NVIDIA Megatron Core/NeMo: Targets maximum efficiency via 3D parallelism (tensor, pipeline, and expert).
    • veRL: A reinforcement learning framework that uses HybridFlow to mix training backends (FSDP2, Megatron) with inference engines (vLLM, SGLang) in a single job.
    • vLLM and SGLang: Inference engines utilizing PagedAttention and RadixAttention to optimize KV cache management and request scheduling.

Observability and Fault Detection

Systematic telemetry is required to diagnose bottlenecks across thousands of GPUs.

Telemetry Stack

  • Prometheus and Grafana: The standard for metrics collection and visualization. AWS provides Amazon Managed Service for Prometheus (AMP) and Amazon Managed Grafana (AMG) to remove operational overhead.
  • DCGM-Exporter: Exposes NVIDIA GPU metrics. SM activity (DCGM_FI_PROF_SM_ACTIVE) is cited as a more accurate measure of compute efficiency than basic utilization.

Health Monitoring

Proactive detection of hardware failures prevents training interruptions. Critical indicators include:

  • ECC Errors: Accelerating single-bit error (SBE) rates often precede double-bit errors (DBE).
  • XID Events: Specific errors such as XID 63 (row remap failure), XID 64 (GPU fallen off bus), and XID 94/95 (contained/uncontained errors) trigger immediate node replacement.

Sources