Simulate a ring all-reduce and count the elements each rank sends

mediumPremium edge casesimplement ring_all_reduce

Before you code

After one ring step, does each rank hold a complete reduction or only progress on some chunks? Trace the ownership before coding.

Review the prerequisite concept

The concept page shows its own access level. The statement and visible tests below are free.

Simulate a sum all-reduce of p equal-length vectors over a ring, the way a tensor-parallel layer combines its partial sums.

ring_all_reduce(bufs) -> (results, sent)
  bufs     list of p 1-D numpy arrays of the same length n (do not modify them)
  results  list of p arrays, each the elementwise sum of all inputs
  sent     list of p ints, the elements each rank sent in total
  split each buffer into p chunks with sizes differing by at most one;
  reduce-scatter for p - 1 steps (rank r sends chunk (r - step) mod p to rank r + 1,
  which adds it), then all-gather for p - 1 steps, starting from the chunk each rank
  has finished

The visible tests check the result and the byte count on four ranks. The full suite adds one rank, two ranks, a length that does not divide by p, the inputs left untouched, and the step structure: every rank sends exactly 2(p - 1) messages.

YOUR IMPLEMENTATION · PYTHON 3.14 + NUMPY

Visible tests run for everyone; the edge cases and the solution are Premium.

The visible tests

import numpy as np


def test_every_rank_ends_with_the_sum(ns):
    rng = np.random.default_rng(0)
    bufs = [rng.normal(size=64) for _ in range(4)]
    out, _ = ns["ring_all_reduce"](bufs)
    total = np.sum(bufs, axis=0)
    for r, o in enumerate(out):
        assert np.allclose(o, total), f"rank {r} does not hold the full sum"


def test_each_rank_sends_two_p_minus_one_over_p_of_the_buffer(ns):
    bufs = [np.ones(8192) for _ in range(4)]
    _, sent = ns["ring_all_reduce"](bufs)
    assert list(sent) == [12288] * 4, f"2 x 3/4 x 8,192 = 12,288 elements per rank, got {list(sent)}"

The reasoning behind this component: the companion page.