Skip to content

All-Reduce 3: Ring — Constant Per-Step Size

The two-phase shape from step 09, but each rank only ever moves data with its left neighbour. Instead of reading P-1 peers per stage, the chunks rotate around the ring: 2 * (P-1) steps, each moving one N/P-sized chunk — so the per-step transfer stays N/P as P grows under weak scaling (for a fixed N the chunk shrinks instead). Synchronization is neighbour-local too: each rank notifies its right neighbour after a store and waits on its left neighbour before a read — no full barrier per round.

Prerequisites: 14-allreduce_two_phase. Four sim devices recommended (the constant per-step size is only observable at P≥4); SIZE must divide evenly by the rank count.

Suggested reading order: 01 → … → 09 → 10 — this page is step 10.

The idea

Two-phase halves mesh's traffic but still reads every peer per stage — O(P) peers, each a N/P chunk. The ring removes the "every peer" part: arrange the ranks in a circle and pass chunks to the left neighbour. The same reduce-scatter + all-gather split now takes P-1 steps each:

  • Reduce-scatter (P-1 steps): in each step a chunk travels one hop and is added at its destination. After P-1 steps every rank holds the reduced chunk it owns.
  • All-gather (P-1 steps): the reduced chunks keep circulating, and each rank copies every chunk as it passes. After P-1 more steps every rank has the full result.

Total bytes equal two-phase (2 * (P-1) / P * N), but each step moves only N/P — under weak scaling (a workload that grows with P) the per-step size stays constant as P grows, which is what keeps the ring efficient at large world sizes; for a fixed N the chunk would shrink instead.

Run it

# P=4 (the constant per-step size shows here) and P=2:
python examples/distributed/10_allreduce_ring.py -p a2a3sim -d 0,1,2,3
python examples/distributed/10_allreduce_ring.py -p a2a3sim -d 0,1

Expected output:

OK

Walkthrough

The kernel is monolithic (one InCore function) and the signal generalises the two-phase idea: [2 * (nr-1), nr] — one row per round. The chunk index arithmetic is the heart of the schedule:

left = (my_rank - 1 + nranks) % nranks      # never negative

# Reduce-scatter: (nr-1) steps.
for s in pl.range(nranks - 1):
    step = s + 1
    recv_add_idx = (my_rank - step - 1 + nranks) % nranks
    left_send_idx = (left - step + nranks) % nranks
    # Wait for the left neighbour's round-s chunk (signal row s), then:
    pld.system.wait(signal, offsets=[s, left], expected=1, cmp=pld.WaitCmp.Ge)
    recv = pld.tile.remote_load(scratch, peer=left,
                                offsets=[0, left_send_idx * chunk],
                                shape=[1, chunk])
    acc = pl.load(scratch, [0, recv_add_idx * chunk], [1, chunk])
    acc = pl.add(acc, recv)
    scratch = pl.store(acc, [0, recv_add_idx * chunk], scratch)
    # The store stages next round's send: signal the right neighbour (row s+1).
    pld.system.notify(signal, peer=right, offsets=[s + 1, my_rank],
                      value=1, op=pld.NotifyOp.AtomicAdd)

# All-gather: (nr-1) steps (rows nranks-1 .. 2*(nranks-1)-1), copying the
# left neighbour's send chunk into the local chunk.

Reduce-scatter, P=4: each round, rank r adds its left neighbour's copy of one chunk into its own copy of that same chunk; which chunk index shifts by one every round:

round s=0:  R0 += C2(R3)   R1 += C3(R0)   R2 += C0(R1)   R3 += C1(R2)
round s=1:  R0 += C1(R3)   R1 += C2(R0)   R2 += C3(R1)   R3 += C0(R2)
round s=2:  R0 += C0(R3)   R1 += C1(R0)   R2 += C2(R1)   R3 += C3(R2)
            (last round: each rank's chunk index matches its own rank number)

After P-1 = 3 rounds, rank r owns the fully-reduced chunk r
(R0 -> C0, R1 -> C1, R2 -> C2, R3 -> C3) — all-gather then spreads every
chunk to every rank.
  • left = (my_rank - 1 + nranks) % nranks. The + nranks keeps the dividend non-negative — a bare (my_rank - 1) % nranks yields -1 at rank 0 under truncating modulo (the step-06 lesson, now on the index side).
  • Chunks rotate, not ranks. In round s, the chunk you add from the left and the chunk you forward both shift by one (- step), so every chunk visits every rank exactly once per phase.
  • Neighbour-ready handshakes, not barriers. Each round gets its own signal row — notify the right neighbour after a store, wait on the left neighbour before a remote_load. The monotonic Ge(1) counters never leak between rounds (the step-09 discipline, now at 2*(P-1) rows), and only the two adjacent ranks ever synchronize: O(P) signals per rank (O(P²) system-wide) — versus the O(P²) per rank a full-mesh barrier per round would cost.
  • Row 0 isn't written by the RS loop itself. The excerpt above starts at left = ..., but round s=0's wait (signal, offsets=[0, left]) depends on a notify the RS loop never sends: a stage-in phase, right before it, copies each local chunk into scratch and then notifies the right neighbour at row 0 — the same stage-before-notify ordering every hand-rolled collective in this ladder uses (only the ring notifies a single neighbour instead of every peer), just easy to miss when only the loop body is quoted.
  • The last all-gather round sends no notify. Its own loop guards the final round with if s < nranks - 2: — the row that round's notify would use, 2*(nranks-1), is past the signal's last valid row (2*(nranks-1) - 1), and nothing downstream needs to wait on it anyway. Look for this exact guard if you're diffing the doc against the source and wondering why the last round's notify is missing.

Cost card (per rank): 2 * (P-1) / P * N total — the same as two-phase — but in 2*(P-1) steps of N/P bytes each. Under weak scaling the per-step size stays constant as P grows — unlike mesh's (P-1) * N per step — which is the ring's reason to exist.

Edge cases

Fatal pitfall — a negative left-neighbour index. (my_rank - 1) % nranks is -1 at rank 0 under truncating modulo — an invalid peer the remote_load then targets. Fix: always write (my_rank - 1 + nranks) % nranks (at P=2 rank 0's left neighbour is rank 1; the + nranks keeps the dividend non-negative).

Symptom Likely cause Fix
Hang at P=2 Negative dividend in the left-neighbour index (my_rank - 1 + nranks) % nranks
Wrong chunks in the result Chunk index arithmetic off by one in a round Trace recv_add_idx / left_send_idx for s=0 on paper
Two handshakes share a signal row Row index reused across RS and AG RS rows 0..P-2, AG rows P-1..2(P-1)-1
Result correct at P=2 only P=2 has a single round, hiding rotation bugs Run P=4 and check every chunk position
Same result on every rank but ≠ torch sum Reduction order differs (not a bug) Compare with a tolerance

See also