5.3 Distributed Training Interconnects & NCCL Collectives

5.3 Distributed Training Interconnects & NCCL Collectives

Core Concept: Large Language Models (LLMs) and foundation AI models (such as GPT-4, LLaMA-3, and Nemotron) contain tens of billions to trillions of parameters, far exceeding the memory capacity (e.g., 80 GB on H100 or 192 GB on B200) and computational limits of an individual GPU. Training these architectures requires distributed parallelism strategies coordinated by the NVIDIA Collective Communications Library (NCCL), which orchestrates optimized communication collectives across intra-node NVLink and inter-node InfiniBand fabrics.

---

1. Taxonomy of Distributed Parallelism Strategies

To train massive models efficiently, practitioners decompose compute, memory, and data across multiple orthogonal dimensions:

┌────────────────────────────────────────────────────────────────────────┐
│                     DISTRIBUTED PARALLELISM TAXONOMY                   │
│                                                                        │
│  ┌──────────────────────────────────────────────────────────────────┐  │
│  │ 1. DATA PARALLELISM (DDP, FSDP, ZeRO-1/2/3)                      │  │
│  │    Replicates/shards model parameters; partitions input dataset. │  │
│  └──────────────────────────────────────────────────────────────────┘  │
│  ┌──────────────────────────────────────────────────────────────────┐  │
│  │ 2. TENSOR PARALLELISM (TP - Megatron-LM)                         │  │
│  │    Splits individual linear layers / GEMMs across GPUs in a node.│  │
│  └──────────────────────────────────────────────────────────────────┘  │
│  ┌──────────────────────────────────────────────────────────────────┐  │
│  │ 3. PIPELINE PARALLELISM (PP - Inter-Node Pipelining)             │  │
│  │    Partitions consecutive neural layers across separate nodes.   │  │
│  └──────────────────────────────────────────────────────────────────┘  │
│  ┌──────────────────────────────────────────────────────────────────┐  │
│  │ 4. SEQUENCE / CONTEXT PARALLELISM (SP / CP)                      │  │
│  │    Shards sequence dimension (tokens) across GPUs for long contexts│
│  └──────────────────────────────────────────────────────────────────┘  │
│  ┌──────────────────────────────────────────────────────────────────┐  │
│  │ 5. EXPERT PARALLELISM (EP - Mixture of Experts)                  │  │
│  │    Distributes distinct expert networks across GPUs (All-to-All).│  │
│  └──────────────────────────────────────────────────────────────────┘  │
└────────────────────────────────────────────────────────────────────────┘

1. Data Parallelism (DP) & Distributed Data Parallel (DDP)

• Mechanics: Every GPU maintains a complete copy of the model parameters, gradients, and optimizer states. The training batch is partitioned across $N$ GPUs (each GPU receives a micro-batch). During the backward pass, each GPU computes local gradients, which are averaged across all GPUs using an AllReduce collective before updating weights.
• Bottleneck: Parameter redundancy limits maximum model size to what can fit inside a single GPU's HBM.

2. Fully Sharded Data Parallel (FSDP) & DeepSpeed ZeRO

To eliminate redundant memory allocation, Zero Redundancy Optimizer (ZeRO) and PyTorch FSDP shard model states across all data-parallel ranks:

• ZeRO Stage 1: Shards Optimizer States (e.g., FP32 AdamW momentum and variance vectors) across GPUs ($4\times$ memory reduction). Gradients are synchronized via AllReduce.
• ZeRO Stage 2: Shards both Optimizer States and Gradients ($8\times$ memory reduction). Uses ReduceScatter to aggregate and distribute gradients to owner ranks.
• ZeRO Stage 3 / FSDP: Shards Optimizer States, Gradients, and Model Parameters. During the forward pass, parameters are dynamically gathered on demand via AllGather and discarded immediately after layer execution. During the backward pass, parameters are gathered again via AllGather, and computed gradients are reduced and sharded via ReduceScatter.

3. Tensor Parallelism (TP - Megatron-LM)

• Mechanics: Decomposes individual weight matrices within multi-head attention and multi-layer perceptron (MLP) blocks across GPUs. In a Transformer MLP block ($Y = \text{GELU}(X W_1) W_2$):
  • First GEMM ($W_1$) is partitioned column-wise: each GPU computes $X W_{1, i}$ independently.
  • Second GEMM ($W_2$) is partitioned row-wise: each GPU multiplies its intermediate activation by $W_{2, i}$.
  • An AllReduce sum collective is executed across the TP group to produce the final output tensor $Y$.
• Interconnect Requirement: Because an AllReduce is required for every single transformer layer in both the forward and backward passes, TP generates extreme communication frequency. TP must strictly run over high-speed NVLink / NVSwitch fabrics within a single node (or within an NVL72 rack domain).

4. Pipeline Parallelism (PP)

• Mechanics: Partitions the network vertically by assigning consecutive groups of layers to different GPUs/nodes (e.g., Layers 1–8 on Node 0, Layers 9–16 on Node 1). Micro-batches are pipelined through stages using 1F1B (One Forward, One Backward) scheduling to minimize pipeline bubbles.
• Interconnect Requirement: Communication occurs only at stage boundaries via point-to-point (P2P Send/Recv) transfers, making PP well-suited for inter-node communication across InfiniBand networks.

5. Sequence Parallelism (SP) & Context Parallelism (CP)

• Mechanics: Shards long token sequences (e.g., 32k to 1M+ context windows) along the sequence dimension across GPUs, distributing LayerNorm and Dropout activations that are normally replicated in TP.

6. Expert Parallelism (EP - Mixture of Experts)

• Mechanics: In MoE models (e.g., Mixtral 8x7B, Grok-1), distinct "expert" feed-forward networks reside on different GPUs. A routing gating network directs tokens to their top-$k$ assigned experts. Distributing tokens to experts and collecting results requires massive All-to-All crossbar communications.

---

2. Collective Communication Operations in AI

Collective communications represent synchronized data movement patterns involving an entire group of participating GPU processes (called ranks in an MPI/NCCL communicator):

┌────────────────────────────────────────────────────────────────────────┐
│                     CORE COLLECTIVE OPERATIONS PATTERNS                │
│                                                                        │
│  1. ALLREDUCE: (Combine & Distribute)                                  │
│     Rank 0: [A]   Rank 1: [B]   Rank 2: [C]   Rank 3: [D]              │
│     -----------------------------------------------------              │
│     All Ranks receive: [A + B + C + D]                                 │
│                                                                        │
│  2. ALLGATHER: (Concatenate & Distribute)                              │
│     Rank 0: [A]   Rank 1: [B]   Rank 2: [C]   Rank 3: [D]              │
│     -----------------------------------------------------              │
│     All Ranks receive: [A | B | C | D]                                 │
│                                                                        │
│  3. REDUCESCATTER: (Combine & Shard)                                   │
│     Rank 0: [A0,A1]  Rank 1: [B0,B1]                                   │
│     -----------------------------------------------------              │
│     Rank 0 receives: [A0 + B0]    Rank 1 receives: [A1 + B1]           │
│                                                                        │
│  4. ALL-TO-ALL: (Matrix Transpose / Token Routing)                     │
│     Each rank sends distinct custom chunks to every other rank.        │
└────────────────────────────────────────────────────────────────────────┘

Mathematical Formulation & Data Volume

Let $S$ represent the total message size in bytes, and $N$ represent the total number of GPU ranks in the communication communicator:

---

3. NCCL: Architecture, Topology Detection & Algorithms

The NVIDIA Collective Communications Library (NCCL) is a specialized, open-source library that provides multi-GPU and multi-node collective primitives tuned to extract 100% of the physical performance from NVIDIA hardware.

Hardware-Aware Automatic Topology Discovery

Upon initialization (ncclCommInitRank), NCCL automatically scans the physical server hardware graph using NVIDIA Management Library (NVML) and PCIe subsystem queries:

• Detects direct NVLink connections and NVSwitch crossbar configurations.
• Maps PCIe bus hierarchies, host bridges, and NUMA CPU affinities.
• Identifies InfiniBand Host Channel Adapters (HCAs) and RoCE network devices, mapping each GPU to its closest network rail (Rail-Optimized Topology).
• Automatically constructs an internal directed communication graph to route data along the highest-bandwidth paths (NVLink $\to$ PCIe $\to$ InfiniBand).

┌────────────────────────────────────────────────────────────────────────┐
│                     RING VS. TREE COLLECTIVE ALGORITHMS                │
│                                                                        │
│    RING ALLREDUCE (Bandwidth-Optimal)   TREE ALLREDUCE (Latency-Optimal)│
│                                                                        │
│         [Rank 0] ────► [Rank 1]                    [Root]              │
│            ▲              │                        /    \              │
│            │              ▼                  [Node 0]  [Node 1]        │
│         [Rank 3] ◄──── [Rank 2]               /    \    /    \         │
│                                             [R0]  [R1] [R2]  [R3]      │
│    - Transfer Time: Independent of N   - Latency: O(log N)             │
│    - Ideal for: Large message sizes    - Ideal for: Small/Med messages │
└────────────────────────────────────────────────────────────────────────┘

NCCL Communication Algorithms: Ring vs. Tree

1. Ring AllReduce (Bandwidth-Optimal for Large Payloads):
   • Decomposes the AllReduce operation into two sequential phases: a ReduceScatter followed by an AllGather.
   • Data arrays are divided into $N$ equal chunks. Each GPU rank sends chunk $i$ to its downstream neighbor in a logical ring while receiving chunk $i-1$ from its upstream neighbor.
   • Performance: In $2(N-1)$ communication steps, all ranks complete the AllReduce. The total data transferred per rank is $2 \cdot \frac{N-1}{N} \cdot S$. As $N$ grows large, $\frac{N-1}{N} \approx 1$, meaning total transfer volume is $2S$ bytes per rank, independent of the number of GPUs.
   • Trade-off: Ring latency scales linearly ($O(N)$), making it suboptimal for thousands of GPUs when message sizes are small.
2. Double Binary Tree AllReduce (Latency-Optimal for Large Cluster Scale):
   • Introduced in NCCL 2.4, this algorithm organizes GPU ranks into two complementary balanced binary trees spanning the cluster.
   • Communication latency scales logarithmically with rank count: $O(\log N)$ instead of $O(N)$.
   • Enables massive multi-node clusters (thousands of GPUs) to synchronize small-to-medium gradient tensors without suffering from high ring latency.
3. CollNet (Hierarchical Tree & Network Offload):
   • Hierarchical algorithm that executes intra-node reductions over NVLink crossbars simultaneously on every node, and uses dedicated network collectives (such as SHARP) across InfiniBand to reduce cross-node data in a single hop.

---

4. NVIDIA SHARP: In-Network Reduction Offload

In standard distributed training, reduction operations (such as floating-point addition) are computed exclusively on GPU Streaming Multiprocessors (SMs). As clusters scale to thousands of nodes, traversing network switches to reach remote GPU SMs creates severe packet serialization and latency overhead.

┌────────────────────────────────────────────────────────────────────────┐
│                     NVIDIA SHARP IN-NETWORK COMPUTING                  │
│                                                                        │
│       [ Node 0 ]      [ Node 1 ]      [ Node 2 ]      [ Node 3 ]       │
│       (GPU Grads)     (GPU Grads)     (GPU Grads)     (GPU Grads)      │
│            │               │               │               │           │
│            └───────┬───────┴───────┬───────┴───────┬───────┘           │
│                    │ Packets Streamed to Network   │                   │
│                    ▼                               ▼                   │
│     ┌────────────────────────────────────────────────────────────┐     │
│     │            QUANTUM INFINIBAND SWITCH ASIC (SHARP)          │     │
│     │   - Embedded Arithmetic Logic Units (ALUs)                 │     │
│     │   - Executes Floating-Point SUM/MIN/MAX In-Flight          │     │
│     │   - Halves Network Data Traffic & Drops Latency            │     │
│     └──────────────────────────────┬─────────────────────────────┘     │
│                                    │                                   │
│       ┌────────────────────────────┴────────────────────────────┐      │
│       ▼                                                         ▼      │
│  Single Reduced Result Broadcasted Directly Back to All Nodes!         │
└────────────────────────────────────────────────────────────────────────┘

How NVIDIA SHARP Operates

NVIDIA Scalable Hierarchical Aggregation and Reduction Protocol (SHARP) moves collective communication computation directly into the physical network switch ASICs (Quantum-2 InfiniBand and Quantum-X800):

• In-Switch Arithmetic Execution: InfiniBand switch chips contain dedicated embedded arithmetic processing units capable of performing vector additions (FP16, BF16, FP32, FP64, INT32) directly on packet payloads as they transit through the switch crossbar.
• Elimination of Multiple Network Traverses: Rather than streaming data between GPUs multiple times across ring hops, each node sends its local gradient partition once to the Top-of-Rack (ToR) switch. The switch computes the reduction in-flight and broadcasts the final aggregated vector back to all nodes.
• 50% Reduction in Fabric Traffic: Halves the physical data volume traversing the network core, cutting AllReduce communication time by up to 2$\times$ and freeing GPU SMs to focus purely on forward and backward neural network compute.