Ring All-Reduce

Implement ring_all_reduce: simulate workers passing tensor chunks around a ring until every worker has the elementwise sum. Return one worker's final result. No processes, network connections or multiple devices are needed.

Inputs and output

worker_tensors has shape (N, M, ...). Axis zero selects one of N workers; each worker holds a tensor of shape (M, ...). N, M and any trailing dimensions are positive, and M is divisible by N. Inputs are finite floating-point values.

Return shape (M, ...), preserving dtype and device and leaving the input unchanged. The result is the sum across workers, not the average and not a stack of N copies. With N=1, return that worker's values without running rounds.

Chunk layout and rounds

Split each worker's tensor along its leading axis into N chunks. A chunk has shape (M/N, ...), so keep trailing dimensions intact. Number workers and chunks from 0 to N-1. Worker i sends to its right neighbor (i+1) modulo N.

Each phase has N-1 rounds, numbered r=0 through N-2. Treat a round as simultaneous: read every outgoing chunk from the buffers at the start of that round and write received data into separate next-round buffers. Do not let a worker forward a value it has just received in the same round.

1. Reduce-scatter: add incoming chunks

In round r, worker i sends chunk (i-r) modulo N. Its right neighbor adds the received values to the same chunk slot in its own buffer. All other slots keep their current values.

After this phase, worker i owns the fully summed chunk (i+1) modulo N. The complete chunks are spread across the workers; other slots may still contain partial sums. The completed chunk is not necessarily chunk i.

2. All-gather: copy completed chunks

In round r, worker i sends chunk (i+1-r) modulo N. Its right neighbor replaces that chunk slot with the received values. Copy in this phase; adding would count contributions again.

After N-1 rounds, every worker has all completed chunks in their original positions. Return worker zero's final buffer. Simulate the two phases using chunk reads, additions and copies rather than replacing the exercise with a direct sum across the worker axis or a collective-library call.

Worked example

For worker 0=[1,2] and worker 1=[3,4], each chunk has one element. Reduce-scatter has one round: worker 0 sends chunk 0 to worker 1, while worker 1 sends chunk 1 to worker 0. Their buffers become [1,6] and [4,4]. Worker 0 now owns completed chunk 1=[6]; worker 1 owns completed chunk 0=[4].

During all-gather, they exchange those completed chunks. Both buffers become [4,6], so return [4,6]. Do not divide by two.

With N workers and S elements per worker, this equal-chunk simulation sends S/N elements per worker in each of 2*(N-1) rounds: 2*(N-1)*S/N in total per worker. This counts transmitted elements, not latency or measured speed.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Combine gradients before updating shared parameters

Data-parallel workers start from the same model parameters and process different examples. For a synchronized training step, combine their gradients before the update. An all-reduce gives the combined result to every worker. Check whether the operation returns a sum or a mean.

For equal-sized batches with mean gradients 2 and 6, the combined mean is (2 + 6) / 2 = 4. If the first worker has one example and the second has three, the example-weighted mean is (1*2 + 3*6) / 4 = 5. Averaging worker means gives the correct global mean only when their counts match, or by coincidence.

Pass gradient chunks around a ring

Ring all-reduce splits each tensor into chunks. First, reduce-scatter adds contributions until each worker holds a completed sum for its assigned chunk. Then all-gather distributes those completed chunks so every worker has the whole result. With N workers and the usual equal-chunk scheme, each phase takes N - 1 neighbor exchanges.

In a simulation, all sends for a round must read the previous round's buffers. Updating a buffer and then using its new contents for another send in the same round lets information travel farther than intended. A correct local simulation checks the arithmetic; it does not measure real network performance.

Take one consistent optimizer step

Microbatch accumulation uses the same weighting principle. Combine the gradients, update optimizer state and parameters once, and update any exponential moving average of parameters at the specified point.

Workers also need consistent optimizer state. Running independent Adam updates and averaging the resulting weights is generally different from applying Adam once to the combined gradient, because Adam's update is nonlinear in the gradient.

Detecting runtime…