All-reduce and other collectives

Summary

Choose a collective by each rank’s required output: concatenate pieces, combine aligned values, or distribute a completed result.

Problem

The process group contains ranks 0 and 1, holding [1, 2] and [3, 4]. Both buffers have shape (2,), dtype int64 and the same index meanings. Specify what each rank holds after a sum all-reduce, reduce-scatter and all-gather. Use rank order for concatenation and equal, index-aligned chunks for scattering.

Picture: Add Aligned Positions or Join Pieces

Sum reduction combines corresponding positions: [1+3, 2+4] becomes [4, 6]. All-reduce gives that result to both ranks. Reduce-scatter keeps [4] on rank 0 and [6] on rank 1. All-gather of those reduced shards restores [4, 6] on both.

All-gather of the original buffers instead produces [1, 2, 3, 4]. It concatenates; it does not add. The all-reduce identity therefore needs compatible reduction, chunk boundaries and rank order.

Six Output Contracts

Sum does not divide by rank count. Equal-sized data-parallel batches can average by dividing the summed gradients appropriately. Unequal local batches require weighting by example counts, not blindly averaging local means.

Reference the Values Before the Network

This reference models values and order, not transport, dtype enforcement or synchronization.

def sum_by_index(buffers):
	if not buffers:
		raise ValueError("the process group must be nonempty")
	width = len(buffers[0])
	if any(len(buffer) != width for buffer in buffers):
		raise ValueError("rank inputs must have matching shapes")
	return [sum(contributions) for contributions in zip(*buffers)]


def reduce_scatter_reference(buffers):
	reduced = sum_by_index(buffers)
	ranks = len(buffers)
	if len(reduced) % ranks:
		raise ValueError("equal chunks require a divisible width")
	width = len(reduced) // ranks
	return [reduced[rank * width:(rank + 1) * width] for rank in range(ranks)]


def all_gather_reference(shards):
	if not shards or any(len(shard) != len(shards[0]) for shard in shards):
		raise ValueError("this reference requires equal-sized rank shards")
	return [value for shard in shards for value in shard]


buffers = [[1, 2], [3, 4]]
assert sum_by_index(buffers) == [4, 6]
assert reduce_scatter_reference(buffers) == [[4], [6]]
assert all_gather_reference(reduce_scatter_reference(buffers)) == [4, 6]
assert all_gather_reference(buffers) == [1, 2, 3, 4]

The Ring All-Reduce

Now use four ranks and full input size $B=8$ MB, with decimal MB. Split each input into four equal, aligned chunks. For chunk A, rank-local contributions are $a_0,a_1,a_2,a_3$. Three receive-and-add steps finish $A^{\ast }$; three forwarding steps distribute it without further addition. Other chunks follow the corresponding schedule concurrently.

Each rank sends $B/4=2$ MB per step: six steps send 12 MB. It also receives 12 MB; do not count sent-plus-received bytes as sent traffic. Generally, ring sent traffic is $2B(N-1)/N$ per rank over $2(N-1)$ steps.

Change 1: The Consumer Needs Only Its Shard

New task. The optimizer needs its own portion of the summed gradient, not the full sum everywhere.

Picture and change. Stop after reduce-scatter and preserve the optimizer’s shard mapping. Do not all-gather a result merely to discard the other shards.

Check. The two-rank outputs remain [4] and [6], not two full buffers. In the four-rank ring model, this phase sends 6 MB per rank instead of 12 MB.

Use when. Downstream work can consume sharded results.

Change 2: Thousands of Tiny Buffers

New task. Replace one 8 MB buffer with 1,000 separate 8 KB buffers.

Picture and change. Total ring bytes stay fixed, but sequential communication steps multiply. For effective bandwidth $b$ and per-step latency $\alpha$, use the simplified unoverlapped model

At four ranks, $\alpha =5$ microseconds and $b=1$ GB/s give 12.03 ms for the single buffer. The 1,000 serial collectives take 42 ms in this model, despite sending the same 12 MB per rank.

Check. Measure backend launch/hop costs, bucketing, tree alternatives and actual topology before selecting an algorithm. These calculations are not hardware measurements.

Use when. Fixed overhead dominates small messages.

Failure: Ranks Disagree About the Operation

A skipped call or mismatched shape, dtype, chunk order or collective sequence can hang, fail or invalidate results. All participating ranks must follow the process group’s matching communication contract.

Transfer Check

Does a long total collective duration establish the training bottleneck?

Reasoned Answer

No. Measure its exposed critical-path tail after overlap and contention. Summed durations can include work already hidden behind computation.