AWS recently launched the Amazon EC2 P6e and P6 instance families, powered by NVIDIA’s Blackwell architecture (watch the announcement). The flagship P6e UltraServer packs 72 NVIDIA Blackwell GPUs into a single NVLink domain with 130 TB/s of bisection bandwidth, 13.4 TB of HBM3e, and 360 petaflops of FP8 compute (720 at FP4). These UltraServers are typically used for large-scale distributed training of frontier models at the multi-trillion-parameter scale.

In this post, we focus on improving cold-start TTFT for a single P6 or P5en instance. Specifically, we cover how to get model weights in the correct format to HBM memory as quickly as possible. For UltraClusters with multiple nodes, this same process would be performed in parallel across all nodes in the cluster. Each node in the UltraServer loads the model independently from the shared FSx for Lustre filesystem, taking advantage of the massive scalable GDS-enabled throughput FSx for Lustre can provide.

First, consider the basic difference between GPUDirect Storage and CPU-based model loading. Traditional CPU-based model loading (left) streams the checkpoint through CPU memory and copies weights to each GPU sequentially over PCIe. Sharded GPUDirect Storage loading (right) pre-splits the checkpoint across tensor-parallel ranks on Amazon FSx for Lustre, and all eight GPUs read their shards in parallel directly into HBM through EFA, bypassing the CPU entirely as shown in Figure 1.

Figure 1: CPU-based model loading vs. sharded GPUDirect Storage loading on an 8-GPU instance.

Consider a 405B parameter model like Llama 3.1 405B, that’s roughly 800 GB of checkpoint data in BF16.The traditional model loading path looks like the following: