37.3a Case Study: Designing a 128-GPU InfiniBand cluster for LLM pre-training¶

📖 Context Introduction¶

Large Language Model (LLM) pre-training is one of the most demanding workloads in AI infrastructure. It requires massive parallelism, extremely fast interconnects, and careful network topology design. This case study walks through the design decisions for a 128-GPU cluster using InfiniBand, focusing on the key factors that engineers must consider when building such a system.

⚙️ The Scenario: What Are We Building?¶

Your organization needs a cluster to pre-train a 70-billion parameter LLM. The key requirements are:

• 128 NVIDIA GPUs (e.g., H100 or A100)
• InfiniBand interconnect for GPU-to-GPU communication
• Support for distributed training using frameworks like NVIDIA NeMo or Megatron-LM
• High bandwidth and low latency between all GPUs

🧠 Core Design Considerations¶

🔹 GPU Selection¶

• Each GPU must support NVLink for intra-node communication and InfiniBand for inter-node communication.
• Example: NVIDIA H100 GPUs with 900 GB/s NVLink bandwidth and 400 Gb/s InfiniBand per GPU.

🔹 Node Configuration¶

• Each node contains 8 GPUs connected via NVLink.
• Total nodes required: 128 GPUs ÷ 8 GPUs per node = 16 nodes.

🔹 InfiniBand Topology¶

• Fat-tree topology is the standard choice for LLM training.
• Each GPU gets a dedicated InfiniBand adapter (HCA) for maximum throughput.
• Switch hierarchy: Leaf switches connect to nodes, spine switches connect leaf switches.

📊 Network Architecture Breakdown¶

📊 Visual Representation: 128-GPU AI Cluster physical layout¶

This diagram displays a 128-GPU cluster: linking 16 nodes through high-speed InfiniBand leaf/spine fabrics.

🛠️ Key Design Decisions¶

📌 1. Full Bisection Bandwidth¶

• Every GPU should be able to communicate with any other GPU at full speed.
• Achieved by ensuring the spine-to-leaf uplinks match the total leaf-to-node bandwidth.
• Formula: Total leaf downlinks = Total spine uplinks for ideal performance.

📌 2. Network Partitioning (Subnet Manager)¶

• Use a single InfiniBand Subnet Manager (e.g., OpenSM or UFM) to manage routing.
• Partition the network into one large fabric for all 128 GPUs — no segmentation needed for LLM training.

📌 3. Storage Connectivity¶

• Add a separate storage network (e.g., Ethernet or additional InfiniBand) for dataset access.
• Storage traffic must not interfere with GPU communication.
• Use NVIDIA GPUDirect Storage to allow GPUs to read data directly from storage without CPU involvement.

🕵️ Performance Expectations¶

• All-reduce bandwidth: With 128 GPUs, expect ~400 Gb/s per GPU for gradient synchronization.
• Training throughput: A 70B parameter model can achieve ~50-100 teraFLOPs per GPU depending on model parallelism strategy.
• Latency: InfiniBand provides sub-microsecond latency between any two GPUs in the cluster.

📈 Common Pitfalls to Avoid¶

• ❌ Oversubscribing uplinks: If spine-to-leaf bandwidth is less than leaf-to-node bandwidth, performance drops significantly.
• ❌ Mixing GPU generations: H100 and A100 have different NVLink and InfiniBand speeds — always use identical hardware.
• ❌ Ignoring cooling and power: 128 H100 GPUs can consume over 100 kW — plan for liquid cooling or high-density air cooling.
• ❌ Using shared storage for checkpointing: Write checkpoints to local NVMe or a dedicated parallel file system (e.g., Lustre or GPUDirect Storage).

✅ Summary Checklist for Engineers¶

• [ ] 16 nodes, each with 8 GPUs connected via NVLink
• [ ] 128 InfiniBand NDR400 adapters (one per GPU)
• [ ] Fat-tree topology with full bisection bandwidth
• [ ] Single InfiniBand subnet managed by a Subnet Manager
• [ ] Separate storage network with GPUDirect Storage support
• [ ] Power and cooling capacity for 100+ kW
• [ ] Identical GPU and InfiniBand hardware across all nodes

🔍 Final Takeaway¶

Designing a 128-GPU InfiniBand cluster for LLM pre-training is about balancing compute, memory, and network bandwidth. The network is often the bottleneck — so prioritize full bisection bandwidth and low-latency interconnects. With the right topology and hardware choices, this cluster can efficiently train state-of-the-art language models at scale.