NCCL Collective Communication Primitives

Distributed training strategies fundamentally differ in their communication patterns. While Distributed Data Parallel (DDP) relies heavily on AllReduce to average gradients across replicas, Fully Sharded Data Parallel (FSDP) decomposes this operation to optimize memory usage. FSDP operation relies almost exclusively on two NCCL (NVIDIA Collective Communications Library) primitives: AllGather and ReduceScatter. Understanding the mechanics and cost models of these primitives is necessary for diagnosing bottlenecks in multi-node clusters where network bandwidth, not GPU compute, often dictates training throughput.

The FSDP Communication Cycle

In a sharded environment, no single GPU holds the complete model weights. Before a forward or backward pass can occur for a specific layer (or group of layers), the distributed shards must be assembled into full parameters. Once computation is complete, the results must be synchronized and the temporary full parameters discarded to free memory.

This creates a distinctive "accordion" memory pattern, driven by the following cycle during a training step:

1. Forward Pass (AllGather): Collect parameter shards from all ranks to materialize full weights.
2. Backward Pass Input (AllGather): Re-materialize full weights (if discarded after forward) to compute gradients with respect to inputs.
3. Backward Pass Output (ReduceScatter): Synchronize gradients across ranks; simultaneously sum them and scatter the results so each rank ends up with only the gradients corresponding to its specific parameter shard.

The following diagram illustrates the state transitions of a model block during these operations.

Transition of data states on a single GPU during FSDP execution. The system oscillates between low-memory sharded states and high-memory materialized states via collective communications.

AllGather: Parameter Materialization

AllGather is the bandwidth-heavy operation responsible for reconstructing the full model weights. Given $N$ GPUs, where each GPU $i$ holds a parameter shard $P_i$, the goal is for every GPU to possess the set $\{P_0, P_1, \dots, P_{N-1}\}$.

In a standard ring-based implementation used by NCCL, the communication cost is significant. For a model size $M$ (in bytes), each rank must receive the data it does not possess. The volume of data received by each rank is:

$V_{rx} = M \times \frac{N-1}{N}$

As $N$ increases, the transfer volume approaches the total model size $M$. This scaling characteristic is important; adding more nodes does not reduce the AllGather communication volume per GPU for the forward pass. It merely fragments the source data more finely.

The theoretical time to complete an AllGather operation on a ring topology with bus bandwidth $B$ is:

$T_{AllGather} = \frac{M(N-1)}{B \cdot N}$

In multi-node clusters, $B$ effectively becomes the bottleneck bandwidth of the inter-node interconnect (e.g., InfiniBand HDR/NDR or RoCEv2), which is typically an order of magnitude slower than intra-node NVLink.

Hierarchical AllGather

When running Hybrid Sharded Data Parallel (HSDP), the communication pattern changes. PyTorch can utilize hierarchical collectives if supported by the backend, or explicit two-step operations:

1. Intra-node: Ranks within a node aggregate shards.
2. Inter-node: Representatives from each node exchange aggregated data.

This uses the high bandwidth of NVLink for the bulk of the aggregation, reducing the stress on the slower TCP/IP or InfiniBand network.

ReduceScatter: Gradient Synchronization

ReduceScatter acts as the functional inverse of AllGather combined with a reduction operation. During the backward pass, each GPU computes gradients for the full parameters (since it just materialized them). However, each GPU is only responsible for updating its specific shard of the optimizer state.

Therefore, the system must:

1. Reduce (Sum): Aggregate gradients from all $N$ ranks to calculate the global gradient.
2. Scatter: Distribute these summed gradients such that Rank $i$ receives only the gradients corresponding to parameter shard $P_i$.

NCCL implements this as a fused operation. A standard AllReduce (used in DDP) would result in every rank holding a full copy of the global gradients. ReduceScatter avoids this redundancy, ensuring that the memory footprint for gradients remains proportional to $1/N$.

The communication volume cost for ReduceScatter is identical to AllGather:

$T_{ReduceScatter} = \frac{M(N-1)}{B \cdot N}$

Because the backward pass involves both computation and this communication, FSDP implementations aggressively overlap ReduceScatter with the backward computation of the previous layer.

Ring vs. Tree Algorithms in NCCL

While we often model costs using Ring algorithms, NCCL dynamically selects algorithms based on topology and message size.

1. Ring: Optimal for throughput on large messages. Data flows in a logical circle ($0 \to 1 \to \dots \to N \to 0$). This is the default for FSDP's large parameter tensors. Latency scales linearly with $N$.
2. Tree (Double Binary Tree): Optimal for latency on smaller messages or when scaling to thousands of GPUs. It reduces the latency scaling from $O(N)$ to $O(\log N)$.

In FSDP training of Large Language Models, the tensors are sufficiently large that Ring (or segmented Ring) algorithms dominate. However, on extremely large clusters (e.g., >256 GPUs), the linear latency of the ring can prevent effective communication overlap.

The chart below demonstrates the theoretical efficiency drop-off as node count increases, assuming fixed model size and fixed interconnect bandwidth.

As $N$ increases, the term $\frac{N-1}{N}$ rapidly approaches 1. This indicates that after a small number of GPUs, the bandwidth requirement saturates. Adding more nodes does not decrease the communication time per GB of model data; it only distributes the compute and memory capacity.

Tensor Fusion and Bucketing

PyTorch and NCCL do not transmit tensors one by one. Doing so would incur massive kernel launch overheads and fail to saturate the PCIe or NVLink bus. Instead, FSDP utilizes a bucketing strategy.

Tensors are flattened and concatenated into a single contiguous buffer (a "bucket") up to a configured size limit (e.g., 25MB or 50MB). The collective primitive is executed on this bucket.

• Latency-Bandwidth Trade-off:
  • Small Buckets: Higher handshake overhead, dominated by latency. Bad for bandwidth utilization.
  • Large Buckets: High bandwidth utilization, but introduces "start-up latency." The communication cannot begin until the entire bucket is populated by the compute stream.

Configuring bucket_cap_mb in the ShardingStrategy is an optimization vector. For multi-node training over Ethernet or InfiniBand, larger buckets (e.g., 100MB+) often yield better throughput by amortizing the protocol overhead, whereas smaller buckets might be preferred on low-latency NVLink-only setups to tighten the overlap window.

Network Bottlenecks and P2P Transport

In a multi-node cluster, AllGather and ReduceScatter must traverse the inter-node fabric. If the cluster uses GPU Direct RDMA (GDR), NCCL can transfer data directly from GPU memory to the Network Interface Card (NIC), bypassing the CPU.

Without GDR, data follows the path: GPU -> CPU RAM -> NIC -> Network -> NIC -> CPU RAM -> GPU. This introduces significant latency and consumes CPU memory bandwidth.

When profiling FSDP with the PyTorch Profiler, distinct "wait" states in the NCCL streams usually indicate that the collective operations are blocked waiting for data from the slowest link in the ring. In heterogeneous clusters, a single node with a degraded NIC negotiation (e.g., running at 100Gbps instead of 400Gbps) will throttle the AllGather speed of the entire cluster to that lowest speed, regardless of how fast the other GPUs can compute.