All-Reduce, All-Gather and Reduce-Scatter: The Collectives Distributed Training Is Built From
The problem, from first principles
01-data-parallelism-gradient-averaging's average_gradients and 03-dataparallel-vs-distributeddataparallel's "ring all-reduce" were both mentioned by name without ever pinning down what an all-reduce actually does, or how it relates to the other collective operations (all-gather, reduce-scatter) that show up constantly in distributed-training documentation. This question implements all three directly, as the exact, general-purpose building blocks the rest of this track's questions are quietly built from.
From theory to code
Implement all_reduce_sum, all_gather, and reduce_scatter_sum — three of the fundamental collective communication operations every distributed deep learning framework (NCCL, MPI, torch.distributed) provides.
Constraints
all_reduce_sum(worker_arrays): every worker holds a same-shape array; returns a list where every worker gets the elementwise sum of all of them.all_gather(worker_shards): every worker holds a (possibly different-sized) shard; returns a list where every worker gets the full concatenation of all shards, in the original order.reduce_scatter_sum(worker_arrays): every worker holds a same-shape array; returns a list where each worker gets only its own chunk of the elementwise sum (the memory-efficient half of all-reduce).- Concatenating
reduce_scatter_sum's output must reconstructall_reduce_sum's (single, shared) result exactly.
Hints
Open one at a time. Each gives away a little more than the last.
Hint 1
All three operations start from the exact same first step: total = sum(worker_arrays). All-reduce hands that whole total back to everyone; reduce-scatter splits it up first and hands each worker only its own piece.
Hint 2
all_gather is conceptually unrelated to summing at all — it's for when every worker holds a genuinely different piece of data (like an embedding-table shard) and needs to see everyone else's piece too. np.concatenate plus copying that result into every worker's slot is the whole implementation.
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.