Communication Primitives

Before understanding distributed training or inference, you need the vocabulary of multi-GPU communication. There are only 7 operations — learn these and you can analyze any parallelism strategy.

1. Point-to-Point: The Foundation

The simplest operation: one GPU sends data to another.

# GPU 0 sends tensor to GPU 1
if rank == 0:
    dist.send(tensor, dst=1)
elif rank == 1:
    dist.recv(tensor, src=0)

All collective operations are built from point-to-point sends under the hood. But you rarely use point-to-point directly — collectives are more efficient.

2. The Five Key Collectives

Collective Operations

Broadcast: One → All

One GPU’s data is copied to all GPUs:

Before:  GPU0=[A]  GPU1=[_]  GPU2=[_]  GPU3=[_]
After:   GPU0=[A]  GPU1=[A]  GPU2=[A]  GPU3=[A]

Use case: Distributing model weights from GPU 0 to all GPUs at initialization.

Data moved: $N$ bytes × $(P-1)$ copies, where $N$ = tensor size, $P$ = GPU count.

Reduce: All → One

All GPUs contribute data, reduced (summed/averaged) to one GPU:

Before:  GPU0=[1]  GPU1=[2]  GPU2=[3]  GPU3=[4]
After:   GPU0=[10] GPU1=[_]  GPU2=[_]  GPU3=[_]   (sum = 10)

Use case: Rare in practice — AllReduce is almost always preferred.

AllReduce: All → Reduce → All ⭐

The most important collective in AI. All GPUs contribute, result available on all GPUs:

Before:  GPU0=[1]  GPU1=[2]  GPU2=[3]  GPU3=[4]
After:   GPU0=[10] GPU1=[10] GPU2=[10] GPU3=[10]  (sum on every GPU)

Use case: Gradient synchronization in data parallelism. Each GPU computes gradients on its data shard, AllReduce averages gradients across all GPUs, then every GPU has the same averaged gradient for the optimizer step.

The fundamental identity:

\text{AllReduce} = \text{ReduceScatter} + \text{AllGather}

This decomposition is how NCCL actually implements AllReduce, and it’s the basis for ZeRO.

ReduceScatter: All → Reduce → Scatter

Reduce all data, then each GPU keeps only its shard of the result:

Before:  GPU0=[a0,a1,a2,a3]  GPU1=[b0,b1,b2,b3]  GPU2=[c0,c1,c2,c3]  GPU3=[d0,d1,d2,d3]
After:   GPU0=[a0+b0+c0+d0]  GPU1=[a1+b1+c1+d1]  GPU2=[a2+b2+c2+d2]  GPU3=[a3+b3+c3+d3]

Each GPU gets 1/P of the total reduced result.

Use case: ZeRO Stage 1-3. After backward pass, ReduceScatter gives each GPU the averaged gradient shard it needs for its optimizer partition.

AllGather: Scattered → All

The inverse of ReduceScatter — each GPU has a shard, AllGather assembles the full tensor on all GPUs:

Before:  GPU0=[s0]  GPU1=[s1]  GPU2=[s2]  GPU3=[s3]
After:   GPU0=[s0,s1,s2,s3]  GPU1=[s0,s1,s2,s3]  GPU2=[s0,s1,s2,s3]  GPU3=[s0,s1,s2,s3]

Use case: FSDP/ZeRO-3 forward pass. Parameters are sharded — before each layer, AllGather reconstructs the full parameters from all GPUs’ shards.

All-to-All: Everyone → Everyone (Different)

Each GPU sends a different piece to each other GPU:

Before:  GPU0=[a0,a1,a2,a3]  GPU1=[b0,b1,b2,b3]
After:   GPU0=[a0,b0]        GPU1=[a1,b1]
         GPU2=[a2,b2]        GPU3=[a3,b3]

GPU $i$ receives the $i$ -th piece from every GPU.

Use case: Expert parallelism (MoE). Tokens must be dispatched to the GPU holding their selected expert. Each GPU sends different tokens to different expert-GPUs.

3. Communication Volume Summary

For $P$ GPUs, each with $N$ bytes of data:

Collective	Per-GPU send	Per-GPU receive	Total data moved
Broadcast	$N$ (from root)	$N$ (to each)	$N \times (P-1)$
Reduce	$N$ (from each)	$N$ (to root)	$N \times (P-1)$
AllReduce	$2N \times \frac{P-1}{P}$	$2N \times \frac{P-1}{P}$	$2N \times (P-1)$
ReduceScatter	$N \times \frac{P-1}{P}$	$N/P$	$N \times (P-1)$
AllGather	$N/P$	$N \times \frac{P-1}{P}$	$N \times (P-1)$
All-to-All	$N \times \frac{P-1}{P}$	$N \times \frac{P-1}{P}$	$N \times (P-1)$

Key insight: AllReduce moves twice the data of ReduceScatter or AllGather individually — because it’s ReduceScatter + AllGather combined. But for large messages, the per-GPU volume scales as $\frac{2N(P-1)}{P} \approx 2N$ — nearly independent of GPU count. This is why distributed training scales.

4. Where Each Collective Appears in AI

Parallelism strategy	Primary collective	When
Data Parallel (DDP)	AllReduce	Gradient sync after backward
ZeRO Stage 1	ReduceScatter	Gradient sharding
ZeRO Stage 3 / FSDP	AllGather + ReduceScatter	Parameter gather (forward) + gradient scatter (backward)
Tensor Parallel	AllReduce	After each TP layer (column/row parallel)
Pipeline Parallel	Point-to-point	Activation transfer between stages
Expert Parallel (MoE)	All-to-All	Token dispatch to/from experts
Context Parallel	AllGather or Ring	KV sharing across sequence chunks

Understanding which collective each strategy uses tells you the communication cost — and where the bottleneck will be.

Self-Check

Name the 5 key collectives and draw the before/after data distribution for each
Explain why AllReduce = ReduceScatter + AllGather
Calculate per-GPU communication volume for AllReduce with P GPUs and N bytes
Explain why AllReduce volume is nearly independent of GPU count for large messages
Map each parallelism strategy to its primary collective operation
Explain why All-to-All is needed for MoE but not for dense models