Collective communication

Deep dive

How it works

In: n ranks, each with S bytes (NCCL convention: S is the full array) · Out: per operation, below · State: none between calls (the communicator keeps the topology) · Cost: steps·α + bytes/B

1. The family. Minimum bytes per rank; these are the bus-bandwidth factors of nccl-tests.

2. Identities. All-reduce = reduce-scatter + all-gather (ring all-reduce); broadcast = scatter + all-gather; reduce = reduce-scatter + gather. All-to-all is not built from the others.
3. Cost per algorithm (counting arguments; Thakur et al. choose among these families by message size):

4. Pick by size. A ring's two terms are equal at S ≈ n·α·B (the 16 MB above). Bruck beats pairwise all-to-all below about 7 MB (arithmetic).
5. NCCL. NCCL_ALGO accepts Ring, Tree, CollnetChain/Direct, NVLS, NVLSTree (NVLink SHARP offload) and PAT, a log-step all-gather and reduce-scatter for one GPU per node (2.23+). Double binary trees (2.4) give all-reduce full bandwidth with log latency. All-to-all is grouped send/recv or ncclAlltoAll.
6. Topology. In a node NVLink/NVSwitch links the GPUs; between nodes each GPU has a NIC. Hierarchical all-gather over m nodes of g GPUs: each NIC carries (m−1)/(m·g) · S instead of (m−1)/m · S (arithmetic). NCCL_CROSS_NIC controls whether rings and trees may switch NICs between nodes.

Code

Threads on a channel mesh run three collectives; asserts check results and traffic per rank against the formulas.

use std::ops::Range;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread;

// One rank on a full mesh: a channel to and from every peer, plus traffic counters.
struct Comm { rank: usize, n: usize, tx: Vec<Sender<Vec<f32>>>, rx: Vec<Receiver<Vec<f32>>>, bytes: usize, msgs: usize }

impl Comm {
    fn send(&mut self, to: usize, d: &[f32]) { self.bytes += d.len() * 4; self.msgs += 1; self.tx[to].send(d.to_vec()).unwrap(); }
    fn recv(&self, from: usize) -> Vec<f32> { self.rx[from].recv().unwrap() }
}
fn chunk(c: usize, n: usize, len: usize) -> Range<usize> { c * len / n..(c + 1) * len / n }

// Ring reduce-scatter: afterwards rank r holds the full sum of chunk r.
fn reduce_scatter(c: &mut Comm, d: &mut [f32]) {
    let (r, n, len) = (c.rank, c.n, d.len());
    for s in 0..n - 1 {
        c.send((r + 1) % n, &d[chunk((r + 2 * n - s - 1) % n, n, len)]);
        let inc = c.recv((r + n - 1) % n);
        for (x, y) in d[chunk((r + 3 * n - s - 2) % n, n, len)].iter_mut().zip(inc) { *x += y; }
    }
}
// Ring all-gather: rank r starts with chunk r; afterwards every rank has all chunks.
fn all_gather(c: &mut Comm, d: &mut [f32]) {
    let (r, n, len) = (c.rank, c.n, d.len());
    for s in 0..n - 1 {
        c.send((r + 1) % n, &d[chunk((r + n - s) % n, n, len)]);
        let inc = c.recv((r + n - 1) % n);
        d[chunk((r + 2 * n - s - 1) % n, n, len)].copy_from_slice(&inc);
    }
}
// Pairwise all-to-all: chunk j goes to rank j; chunk i of the result came from rank i.
fn all_to_all(c: &mut Comm, d: &mut [f32]) {
    let (r, n, len) = (c.rank, c.n, d.len());
    let mut out = d.to_vec();                                  // chunk r stays in place
    for s in 1..n {                                            // step s: send to r+s, receive from r-s
        c.send((r + s) % n, &d[chunk((r + s) % n, n, len)]);
        let from = (r + n - s) % n;
        out[chunk(from, n, len)].copy_from_slice(&c.recv(from));
    }
    d.copy_from_slice(&out);
}
type Out = (Vec<Vec<f32>>, Vec<usize>, Vec<usize>); // buffers, bytes and messages sent per rank
fn run(n: usize, bufs: Vec<Vec<f32>>, f: fn(&mut Comm, &mut [f32])) -> Out {
    let (mut txs, mut rxs): (Vec<Vec<_>>, Vec<Vec<_>>) = ((0..n).map(|_| vec![]).collect(), (0..n).map(|_| vec![]).collect());
    for a in 0..n { for b in 0..n { let (t, r) = channel(); txs[a].push(t); rxs[b].push(r); } } // a→b: txs[a][b], rxs[b][a]
    let hs: Vec<_> = bufs.into_iter().zip(txs).zip(rxs).enumerate().map(|(rank, ((mut d, tx), rx))| thread::spawn(move || {
        let mut c = Comm { rank, n, tx, rx, bytes: 0, msgs: 0 };
        f(&mut c, &mut d);
        (d, c.bytes, c.msgs)
    })).collect();
    let res: Vec<_> = hs.into_iter().map(|h| h.join().unwrap()).collect();
    (res.iter().map(|x| x.0.clone()).collect(), res.iter().map(|x| x.1).collect(), res.iter().map(|x| x.2).collect())
}

fn main() {
    for n in [4usize, 8] {
        let (len, s) = (n * 1000, n * 4000);                   // S = s bytes, divisible by n: formulas exact
        let input: Vec<Vec<f32>> = (0..n).map(|r| (0..len).map(|i| (r * 7 + i % 13) as f32).collect()).collect();
        let sum: Vec<f32> = (0..len).map(|i| input.iter().map(|b| b[i]).sum()).collect();
        let per_rank = (n - 1) * s / n;                        // (n-1)/n * S

        let (rs, by, ms) = run(n, input.clone(), reduce_scatter);
        for r in 0..n { assert_eq!(rs[r][chunk(r, n, len)], sum[chunk(r, n, len)]); }
        assert!(by.iter().all(|&b| b == per_rank) && ms.iter().all(|&m| m == n - 1));

        let (ag, by, _) = run(n, rs, all_gather);              // reduce-scatter + all-gather = all-reduce
        assert!(ag.iter().all(|b| *b == sum) && by.iter().all(|&b| b == per_rank));

        let (a2a, by, ms) = run(n, input.clone(), all_to_all);
        for r in 0..n { for i in 0..n { assert_eq!(a2a[r][chunk(i, n, len)], input[i][chunk(r, n, len)]); } }
        assert!(by.iter().all(|&b| b == per_rank) && ms.iter().all(|&m| m == n - 1));

        println!("n={n} S={s} B: RS/AG/A2A {per_rank} B in {} msgs per rank, all-reduce {} B", n - 1, 2 * per_rank);
    }
}

Complexity and constants

• Measured on this Mac (Apple M4 Pro), bytes not time: n = 8, S = 32,000 B; each of reduce-scatter, all-gather and all-to-all sends 28,000 B in 7 messages per rank (= (n−1)/n · S). Asserts also passed for n = 2…9.
• Arithmetic (same α, B): 1 GB ring all-gather on 64 ranks is 315 µs latency plus 19.7 ms transfer.
• MoE all-to-all (assumed 8,192 tokens per rank, top-2, hidden 4,096, bf16, 64 ranks): 132 MB per rank, about 2.6 ms, for dispatch and again for combine.

Efficiency

• Cost: about 20 ms (1 GB all-gather) and 2.6 ms (MoE all-to-all) above, arithmetic. Bound: bandwidth for large buffers, latency for small ones, the slowest rank for uneven all-to-all.
• Levers:
  • double binary tree instead of ring — up to 180x lower all-reduce latency at 24,576 GPUs (NCCL 2.4 blog) — small and medium messages at scale
  • hierarchical decomposition — NIC bytes ÷ g (arithmetic) — multi-GPU nodes
  • Bruck all-to-all — n−1−⌈log₂n⌉ fewer latency steps, more bytes (arithmetic) — small messages
  • PAT or NVLS — log-step or in-switch offload (NCCL docs) — one GPU per node, SHARP hardware
  • overlap with compute — prefetch the next layer's all-gather (ZeRO/FSDP) — layered models
• Naive → standard → state of the art:
  • gather to one rank, then broadcast · root link carries (n−1)·S
  • ring or pairwise by size · (n−1)/n · S per rank
  • hierarchy, trees, in-network reduction · optimal bytes, log-depth latency

Implementation notes

• Chunk by c·len/n so uneven lengths work; pipeline large chunks as smaller messages.
• Variable-size all-to-all first exchanges the counts.

Pitfalls

• Every rank must issue the same collectives in the same order, or the job hangs silently.
• NCCL's root is a rank, not a device number.
• Skewed all-to-all (uneven expert routing or shuffle keys): the slowest rank sets the time.
• Forcing NCCL_PROTO: the docs warn LL128 on an unsupported platform can corrupt data.

Variants and alternatives

In production

• ZeRO/FSDP: all-gather weights, reduce-scatter gradients; tensor parallelism: all-reduce per block; pipeline parallelism uses point-to-point sends instead.
• DDP broadcasts model state from rank 0 at construction (Li et al.). GShard reshards expert dispatch with all-to-all (mixture of experts); DLRM moves embedding lookups with one (embedding sharding).
• FFTW's MPI transforms transpose the whole array between passes; Spark's shuffle redistributes data across executors, like the MapReduce shuffle.

Implementations

Sources

• Thakur, Rabenseifner, Gropp, Optimization of Collective Communication Operations in MPICH, IJHPCA 2005. paper
• Sanders, Speck, Träff, Two-tree algorithms for full bandwidth broadcast, reduction and scan, Parallel Computing 2009. paper
• NVIDIA, NCCL user guide. collectives · environment variables · code
• NVIDIA, nccl-tests: performance (algbw vs busbw). doc · 2.23.4 release notes (PAT)
• Jeaugey, Massively Scale Your Deep Learning Training with NCCL 2.4, NVIDIA blog 2019. blog
• Li et al., PyTorch Distributed: Experiences on Accelerating Data Parallel Training, VLDB 2020. paper
• Lepikhin et al., GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2020. paper
• Naumov et al., Deep Learning Recommendation Model for Personalization and Recommendation Systems, 2019. paper
• FFTW, Transposed distributions (MPI). docs
• Apache Spark, RDD programming guide: shuffle operations. docs