Ring All-Reduce

The problem

In data-parallel training, every GPU (graphics processing unit) computes gradients on its own part of the batch, and then all of them need the sum (or the mean) of everyone's gradients. That operation is an all-reduce. The ring algorithm does it with little network traffic: the workers sit in a circle, and each one only ever sends to its right neighbour. You will simulate it.

There are $N$ workers. Worker $i$ sends only to worker $(i + 1) \bmod N$ , where $\bmod N$ wraps around: the last worker sends to worker 0. Each worker's array is cut into $N$ chunks, and chunks[i][c] is chunk $c$ as held by worker $i$ .

Phase 1: reduce-scatter

$N - 1$ steps. In step $s = 0, \dots, N - 2$ , every worker $i$ sends its chunk $c = (i - s) \bmod N$ to worker $i + 1$ , which adds it to its own chunk $c$ . Afterwards, worker $i$ holds the complete sum of chunk $(i + 1) \bmod N$ .

Phase 2: all-gather

$N - 1$ more steps. In step $s$ , every worker $i$ sends its chunk $c = (i + 1 - s) \bmod N$ to worker $i + 1$ , which replaces its chunk $c$ with it. Afterwards every worker holds every complete chunk.

Example. Three workers with [1, 2, 3], [10, 20, 30] and [100, 200, 300]: three chunks of one number each.

after worker 0 worker 1 worker 2

start [1, 2, 3] [10, 20, 30] [100, 200, 300]

reduce-scatter $s = 0$ [1, 2, 303] [11, 20, 30] [100, 220, 300]

reduce-scatter $s = 1$ [1, 222, 303] [11, 20, 333] [111, 220, 300]

all-gather $s = 0$ [111, 222, 303] [11, 222, 333] [111, 220, 333]

all-gather $s = 1$ [111, 222, 333] [111, 222, 333] [111, 222, 333]

Reduce-scatter, $s = 0$ : worker 0 sends chunk $0$ (the 1) to worker 1, which adds it ( $10 + 1 = 11$ ); worker 1 sends chunk 1 to worker 2 ( $200 + 20 = 220$ ); worker 2 sends chunk 2 to worker 0 ( $3 + 300 = 303$ ).

Reduce-scatter, $s = 1$ : worker 0 sends chunk $(0 - 1) \bmod 3 = 2$ , the $303$ , to worker 1, which now has the complete $333$ . Likewise worker 0 completes chunk 1 ( $222$ ) and worker 2 chunk 0 ( $111$ ): each worker $i$ holds the complete chunk $i + 1$ .

All-gather: two steps copy the complete chunks around the ring, replacing the partial ones.

after	worker 0	worker 1	worker 2
start	`[1, 2, 3]`	`[10, 20, 30]`	`[100, 200, 300]`
reduce-scatter $s = 0$	`[1, 2, 303]`	`[11, 20, 30]`	`[100, 220, 300]`
reduce-scatter $s = 1$	`[1, 222, 303]`	`[11, 20, 333]`	`[111, 220, 300]`
all-gather $s = 0$	`[111, 222, 303]`	`[11, 222, 333]`	`[111, 220, 333]`
all-gather $s = 1$	`[111, 222, 333]`	`[111, 222, 333]`	`[111, 222, 333]`

What to implement

reduce_scatter(chunks) and all_gather(chunks): the two phases. Each returns a new nested list and does not modify its input.
ring_allreduce(arrays): split each worker's 1-D array with np.array_split(a, N), run both phases, concatenate each worker's chunks. Return the list of $N$ results, each equal to the sum of all inputs.

Requirements

$N = 1$ works (no steps at all).
Copy the arrays you receive as input before changing them.

Hint

On real hardware all the sends of a step happen at the same time. In the simulation, the safe way is to first collect every message of the step (destination, chunk index, a copy of the data), and only then apply them all. Check your indices against the table above.