Collective Communication

Start with ownership

A GPU computes values locally. The next training operation may need values from other GPUs. Choose a communication pattern from the required output: which values must each GPU own?

A rank identifies one participant in a process group. Here, each rank represents one GPU. The group has four ranks, numbered 0 through 3. A root is the selected source or destination for a rooted collective.

Three questions distinguish the patterns: Do values combine arithmetically? Does every rank receive the same result? Does each destination need a different piece?

Each scene separates input reference from result being built. A send step names both endpoints. The following receive step changes only the receiver. These are original conceptual schedules, not timings or descriptions of a particular library. Real implementations can forward data, use trees or rings, and overlap many transfers.

Broadcast: one input, every rank

Rank 0 owns a tensor. Broadcast gives every rank that same tensor. No values are added. The root also keeps the tensor.

Training example. Initialize replicated parameters from one agreed state. Rank 0 starts with [1, 2, 3, 4]. After Broadcast, all four replicas start with those values. This synchronizes initialization; it does not synchronize future gradient updates. PyTorch DDP design ↗ describes initialization from rank 0.

Reduce: combine at one rank

Each rank owns a contribution to the same logical tensor. Reduce combines matching positions and places the result at one root. The examples use SUM. Other reductions include MIN and MAX.

Here, rank 0 receives [28, 32, 36, 40]. At position 0, the contributions are 1 + 5 + 9 + 13 = 28. Ranks 1–3 still have their input references, but no defined reduction output.

Training example. Aggregate evaluation statistics at a logging coordinator. Reduce the loss sum and valid-token count to rank 0. Divide the global sum by the global count there. Averaging per-rank means gives the wrong global mean when valid-token counts differ.

AllReduce: combine for every rank

AllReduce computes a reduction and gives the full result to every rank. Each output position includes the corresponding contribution from all participants. The inputs need matching logical positions.

All four ranks finish with [28, 32, 36, 40]. The scene first reduces to rank 0, then broadcasts the complete sum. This is one way to demonstrate the contract. AllReduce does not require a central rank in its implementation.

Training example. In data parallelism, replicas compute gradients from different local batches. Gradient synchronization lets replicas apply the same update. Standard DDP averages gradients across ranks; the animation shows a SUM without division. For equal local batch sizes and local mean losses, divide the displayed sum by four to obtain the mean gradient. Unequal batches require an explicit weighting rule. DDP gradient synchronization ↗.

AllGather: collect all shards everywhere

AllGather starts with one shard per rank. Every rank receives the shards in rank order. Values remain separate; no reduction occurs.

The four input shards are A0=1, B1=6, C2=11 and D3=16. Every result is [1, 6, 11, 16]. A shard already local to its owner needs no network transfer.

Training example. With fully sharded data parallelism (FSDP), a layer’s parameter shards can be gathered before that layer computes. This temporarily increases local parameter storage. The full parameter values are reconstructed, rather than added together. FSDP sharding strategies ↗.

ReduceScatter: combine, then keep a shard

ReduceScatter combines matching tensor positions across ranks. It gives each rank only its assigned partition of the reduced result. The examples use four equal partitions, with one element per partition.

Rank 0 gets position 0: 28. Rank 1 gets position 1: 32. Rank 2 gets position 2: 36. Rank 3 gets position 3: 40. Every result includes contributions from all four inputs, but no rank holds the full reduced vector.

FSDP: send each synchronized gradient to its parameter owner

Fully sharded data parallelism (FSDP) assigns each rank a parameter shard and the corresponding optimizer state. During a layer’s computation, ranks temporarily gather its parameters. Each rank processes a different local batch and produces gradient contributions for that gathered parameter group. FSDP parameter and gradient sharding ↗.

Follow a two-rank FSDP step

In the four-rank animation above, use the four positions as four scalar parameters: w0, w1, w2 and w3. Rank 0 owns w0, rank 1 owns w1, and so on. Every input row is a gradient contribution from one local batch. A0=1 and B0=5 therefore contribute to the same parameter w0. They are gradients, not parameter values.

1. Gather parameters for computation. Each rank obtains the parameters needed by the layer. Its local batch can depend on every parameter in that group.
2. Compute local gradients. The example produces [1,2,3,4] on rank 0, [5,6,7,8] on rank 1, [9,10,11,12] on rank 2 and [13,14,15,16] on rank 3.
3. Reduce matching gradient positions. For w0, add 1 + 5 + 9 + 13 = 28. Deliver that sum to rank 0. Repeat for the other parameter owners.
4. Update each local shard. Each owner uses its synchronized gradient and local optimizer state to update its parameter shard.
5. Gather updated parameters when needed. The next computation reconstructs parameters from the owners’ updated shards.

This synthetic update assumes equal local batch sizes, local mean losses and stochastic gradient descent (SGD) with learning rate 0.1. Divide each sum by four, then subtract 0.1 times that mean gradient from its parameter. The animation stops at the unscaled SUM. Frameworks can apply normalization around the collective.

After the update, a parameter AllGather reconstructs [9.3,19.2,29.1,39.0]. It gathers updated weights. The synchronized gradients remain sharded. With Adam, each owner also updates the first and second moment estimates for its parameter shard.

Each optimizer needs only the gradient positions it owns. ReduceScatter therefore stops with one synchronized gradient shard per rank. An AllReduce would additionally replicate the other owners’ gradient positions, which this local update does not need.

In full sharding, parameter gathers occur before forward and backward computation, with resharding between these phases. Temporary full parameter groups, gradient buffers and activations still contribute to peak memory. Actual shards contain many parameter elements; the scalar table isolates ownership. FSDP communication lifecycle ↗.

AllToAll: a different piece for each destination

Each rank splits its input by destination. Piece j goes to rank j. Each receiver assembles incoming pieces in source-rank order. AllToAll redistributes values; it does not add them.

Read A2 as “from source A (rank 0), addressed to rank 2.” Rank 2 receives A2, B2, C2 and D2: [3, 7, 11, 15]. Rank 0 receives A0, B0, C0 and D0. The piece for the local rank stays local.

Training example. Expert parallelism assigns experts to different ranks. A mixture-of-experts (MoE) router groups token vectors by destination expert. Dispatch sends those groups to their owners. A return exchange sends expert outputs back to the original token owners for combination. Expert-parallel dispatch and return.

The equal-size scene makes ownership visible. Real routing can produce unequal counts or empty destinations. Variable-size exchanges need matching send/receive counts for each peer. Top-k routing can create several token-expert assignments per token. The combine operation must restore token order and apply the routing weights. AllToAll split semantics ↗.

Send / Receive: one peer exchange

Send / Receive transfers data between selected peers. It is point-to-point communication, rather than a collective over the whole group. The sender and receiver must match the payload and peer identities.

Only rank 1 receives rank 0’s [1, 2, 3, 4]. Ranks 2 and 3 do not participate. The source retains its data; a send does not imply deleting the source buffer.

Training example. In pipeline parallelism, a stage sends forward activations to the next stage. During backward, activation gradients travel in the opposite direction. Several microbatches can be in flight, so matching the stage and microbatch matters. PyTorch pipeline tutorial ↗.

Connect the patterns

AllReduce can be understood in two ways: Reduce followed by Broadcast, or ReduceScatter followed by AllGather. The first scene uses a root for clarity. The second factorization explains why gradient reduction can stop at sharded ownership.

The standalone AllGather scene starts from unreduced values [1], [6], [11], [16]. To compose it after ReduceScatter, replace those inputs with [28], [32], [36], [40]. AllGather copies whichever shards it receives; it never supplies the missing arithmetic.

These are equivalences of the required mathematical result. Different floating-point reduction orders can produce slightly different rounded values. Equal results do not imply equal traffic, latency or memory use. Collective contracts and compositions ↗.

Recognize the pattern in a training step

Start from the tensor layout required by the next computation. The same training method can use several patterns in one step.

Tensor parallelism can also use AllReduce for partial output sums, or ReduceScatter when the output stays partitioned. The layer layout determines which result is needed. Infer the collective from matrix ownership.

What must agree across ranks?

All members of a collective’s process group must participate in a compatible order. Agree on the group, root where applicable, reduction, data type and tensor layout. Counts must satisfy the chosen operation; variable-size exchanges must agree per peer.

“Same shape” is insufficient. Two tensors can have identical local dimensions while referring to different global positions. Summing those tensors would combine unrelated values. The rank-to-shard map is part of the contract.

Communication completion must also precede the consuming computation. An asynchronous API can return before the result is ready. Buffer reuse needs its own completion rule. The scenes show logical receipt, rather than GPU memory-order or stream semantics. PyTorch completion semantics ↗.

Once the result contract is clear, investigate scheduling and performance in GPU Communication. That note owns rings, trees, transports and readiness protocols.

Sources and connected notes

The diagrams and integer traces are original teaching models. They show one transfer at a time to expose sender, receiver and ownership. Playback speed and step counts are not performance measurements.

• PyTorch distributed communication reference ↗ — operations, peer exchanges, variable split sizes and asynchronous completion.
• Collective operation definitions ↗ — result contracts and equivalent compositions; no library-specific schedule is assumed here.
• Distributed Data Parallel design ↗ — initialization and gradient synchronization.
• FSDP sharding strategies ↗ — parameter gathering and gradient reduce-scatter.
• Distributed pipeline parallelism ↗ — activation and gradient communication between stages.

Continue with Distributed Training and Parallelism, MoE execution and GPU Communication.