Skip to content

Distributed AI Infrastructure & Parallelism: A Master-Class Deep Dive

As large language models (LLMs) push the boundaries of parameter counts into the trillions, the physical constraints of individual accelerators (GPUs/TPUs) dictate the necessity for distributed, parallelized training and inference architectures. This document dissects the deep engineering realities of distributed AI infrastructure, focusing on hardware interconnect limits, matrix sharding, and communication primitives.

1. Hardware Interconnects and Topology Limits

The fundamental bottleneck in distributed AI is not always compute (FLOPs/s), but data movement (Bytes/s). Understanding interconnect hierarchies is critical.

  • PCIe Gen5 (Host-to-Device): Provides a theoretical maximum of 64 GB/s (128 GB/s bidirectional) on a x16 slot. This is the bottleneck for offloading weights to CPU RAM or NVMe (e.g., in ZeRO-Offload).
  • NVLink (Intra-Node Device-to-Device): Nvidia's proprietary interconnect. NVLink 4.0 (Hopper architecture) provides 900 GB/s bidirectional bandwidth per GPU. A standard 8-GPU node uses NVSwitch to provide non-blocking all-to-all communication. This immense bandwidth is why Tensor Parallelism is typically restricted to a single node.
  • InfiniBand (Inter-Node): Network interconnects like NDR InfiniBand provide 400 Gbps (50 GB/s) or 800 Gbps (100 GB/s) per port. While incredibly fast for networking, it is nearly an order of magnitude slower than intra-node NVLink. Network topologies (Fat-Tree, 3D Torus) dictate the bisection bandwidth available for scaling out across thousands of nodes.

2. Tensor Parallelism (TP)

Tensor Parallelism (TP), popularized by Megatron-LM, shards individual weight matrices across multiple GPUs. Instead of replicating the model, each GPU holds a slice of a tensor and computes a partial result.

Mechanism and Matrix Sharding

Consider a linear layer \(Y = XA\). In 1D Megatron-LM TP, we can slice \(A\) column-wise (\(A = [A_1, A_2]\)). 1. Forward Pass: Each GPU \(i\) computes a partial matrix multiplication \(Y_i = X A_i\). To get the final \(Y\), we don't necessarily need an immediate communication if the next layer is a row-wise parallel layer. 2. If the next layer \(Z = YB\) has \(B\) sharded row-wise (\(B = \begin{bmatrix} B_1 \\ B_2 \end{bmatrix}\)), then each GPU computes \(Z_i = Y_i B_i\). The final result requires an All-Reduce operation: \(Z = Z_1 + Z_2\).

Communication Primitives: All-Reduce and NCCL

An All-Reduce operation sums the partial tensors across all GPUs and broadcasts the result back. Nvidia's Collective Communications Library (NCCL) implements this using algorithms like Ring All-Reduce or Tree All-Reduce. * Ring All-Reduce: Involves \(2(N-1)\) steps of passing data chunks around a logical ring of \(N\) GPUs. The bandwidth utilization is optimal, but latency scales with \(N\). * Why Single Node? Every transformer block in TP requires at least two All-Reduce operations (one in the MLP, one in Attention). Because All-Reduce is heavily latency and bandwidth-dependent, executing it over InfiniBand (inter-node) would devastate utilization (MFU). Thus, TP degree is practically bounded by the number of GPUs sharing an NVLink domain (usually 8).

3. Pipeline Parallelism (PP)

Pipeline Parallelism partitions the model layer-by-layer across multiple GPUs (or nodes). GPU 0 holds layers 1-4, GPU 1 holds layers 5-8, etc.

Micro-batches and Bubble Time

Naive PP is wildly inefficient: GPU 1 sits idle while GPU 0 processes the forward pass, leading to a massive pipeline bubble. To mitigate this, a mini-batch is split into smaller micro-batches.

Scheduling: GPipe vs 1F1B

  • GPipe: Injects all micro-batches of a forward pass sequentially, then executes all backward passes. This requires storing intermediate activations for all micro-batches, leading to peak memory proportional to the number of micro-batches (\(M\)).
  • 1F1B (One Forward, One Backward): Alternates forward and backward passes. Once the first micro-batch completes its forward pass across the pipeline, its backward pass begins immediately. This strictly limits the number of activations in flight, reducing peak memory overhead significantly compared to GPipe, while maintaining the same pipeline bubble size: \(Bubble\_Fraction = \frac{p - 1}{M + p - 1}\) (where \(p\) is pipeline stages).

4. ZeRO and FSDP (Fully Sharded Data Parallel)

Standard Data Parallelism (DP) replicates the entire model state (weights, gradients, optimizer states) on every GPU. For models >10B parameters, this exceeds the VRAM of a single GPU (e.g., 80GB H100).

ZeRO (Zero Redundancy Optimizer) eliminates this redundancy by sharding the model state across the DP group. * ZeRO-1: Shards Optimizer States (momentum, variance). * ZeRO-2: Shards Optimizer States + Gradients. * ZeRO-3 (FSDP): Shards Optimizer States + Gradients + Parameters.

FSDP Execution Flow

In FSDP (ZeRO-3), a GPU only holds a shard of a parameter. 1. Forward Pass: Before computing layer \(L\), the GPU issues an All-Gather to collect the full weights for layer \(L\) from all other GPUs. It computes the forward pass, then immediately discards the gathered weights to save VRAM. 2. Backward Pass: Issues another All-Gather for the weights to compute gradients. Then, a Reduce-Scatter is performed on the gradients so each GPU only holds the accumulated gradients for its specific parameter shard. FSDP overlaps communication (All-Gather) with computation, requiring high inter-node network bandwidth (InfiniBand) but completely solving the VRAM capacity wall.

5. Sequence Parallelism (SP)

While TP shards the hidden dimension, and PP shards the layers, Sequence Parallelism (SP) shards along the sequence dimension (\(s\)). This is critical for scaling to infinite context lengths (e.g., 1M+ tokens), where activation memory (\(s \times b \times h\)) dominates.

Ring Attention and SP Mechanics

Standard Self-Attention requires computing interactions between all tokens, preventing naive sequence sharding. * Megatron-SP: Co-designs with TP. It observes that operations like LayerNorm and Dropout don't need interactions across the hidden dimension. It keeps these operations sharded along the sequence dimension, using an All-Gather before the TP-sharded QKV projection, and a Reduce-Scatter after the Attention output. * Ring Attention / FlashAttention Distributed: To truly shard the attention computation itself, the query, key, and value blocks are distributed across GPUs. GPUs compute local attention blocks and then pass the KV blocks in a ring topology to neighboring GPUs. This enables exact attention computation over context lengths that exceed the aggregate memory of a single node, scaling linearly with the number of devices.

Architectural rule of thumb: Intra-node (NVLink) is for TP. Inter-node (InfiniBand) is for PP, FSDP, and SP.