1.9 Parallel Computing & Collective Communication¶

Modern large language models require compute that no single GPU can provide. Training a 70-billion-parameter model in a reasonable amount of time requires hundreds — sometimes thousands — of GPUs working in tight coordination. This chapter explains the substrate that makes that coordination possible: the programming model for parallel computation and the collective communication operations that keep thousands of accelerators synchronized. Everything in Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP and Distributed Training II: Tensor, Pipeline, Sequence & Expert Parallelism builds on the foundations developed here.

We start from first principles — processes, threads, and how GPUs expose parallelism — and build up to the exact cost models engineers use to reason about whether a training job will be network-bound or compute-bound. Even at the small end this matters: the ~100M-parameter model we build in The Pretraining Run has a bf16 gradient buffer of only ~0.24 GB, so a single-node 4-GPU DDP all-reduce over NVLink costs a couple of milliseconds per step — which is exactly why data parallelism alone is the right (and only) parallelism the capstone needs.

Processes, Threads, and the SPMD Model¶

Processes vs. Threads¶

A process is an OS-managed execution context with its own virtual address space, file descriptors, and signal state. Two processes cannot share memory by default; they must use explicit inter-process communication (IPC). A thread lives inside a process and shares its address space; threads communicate through shared memory and synchronize via locks, semaphores, or atomic operations.

GPU distributed training almost universally uses the one-process-per-GPU model rather than threading across GPUs. The reasons are practical:

1. Python’s Global Interpreter Lock (GIL) prevents true CPU-thread-level parallelism, so CPU-side preprocessing and coordination would be serialized.
2. GPU driver libraries (CUDA) are not always thread-safe when the same process owns multiple GPU contexts.
3. Process isolation provides fault containment: a crash on GPU 3 does not corrupt the state of GPU 0.

Each process gets a unique integer rank (from 0 to \(N-1\), where \(N\) is the world size) and belongs to a process group. A process group is just a set of ranks with a shared communicator — think of it as a namespace for collective operations.

Single-Program Multiple-Data (SPMD)¶

The SPMD programming model is the backbone of all distributed deep learning. Every process runs the same program, but operates on different data and uses its rank to branch when needed.

# spmd_demo.py — run with:
#   torchrun --nproc_per_node=4 spmd_demo.py
import os
import torch
import torch.distributed as dist

def main():
    # Initialize the default process group. "nccl" for GPU; "gloo" for CPU.
    dist.init_process_group(backend="nccl")

    rank = dist.get_rank()        # global rank: unique across the whole job
    world_size = dist.get_world_size()  # total GPU count
    # LOCAL_RANK is the index *within this node* — always use it (not RANK) to
    # pick a device, or rank 8 of a 2-node job will ask for a nonexistent cuda:8.
    local_rank = int(os.environ["LOCAL_RANK"])

    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    # Each process creates a tensor whose value equals its rank.
    # This simulates different data shards arriving at different workers.
    x = torch.tensor([float(rank)], device=device)
    print(f"[rank {rank}/{world_size}] before all-reduce: x = {x.item()}")

    # Sum x across all ranks; every rank receives the global sum.
    dist.all_reduce(x, op=dist.ReduceOp.SUM)
    print(f"[rank {rank}/{world_size}] after all-reduce:  x = {x.item()}")
    # Expected: sum(0, 1, 2, 3) = 6.0 on every rank.

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

When you launch with torchrun --nproc_per_node=4, the launcher spawns four processes. Each process discovers its identity through environment variables (RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, MASTER_PORT) that torchrun sets before calling your script. This is the SPMD contract: every process runs main(), but the outcome differs by rank.

torchrun (the torch.distributed.run elastic launcher, successor to the deprecated torch.distributed.launch) is also how you go multi-node. Run the same command on every node and let the c10d rendezvous server on node 0 assign global ranks:

# On EVERY node (identical command; the rendezvous assigns RANK automatically):
torchrun \
  --nnodes=2 --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=node0.cluster.local:29500 \
  --rdzv_id=my_training_job \
  train.py --config configs/base.yaml

# Under Slurm, `srun` supplies one task per node and torchrun fans out to GPUs:
#   srun --nodes=2 --ntasks-per-node=1 torchrun --nnodes=2 --nproc_per_node=8 ...

With --rdzv_backend=c10d you never set --node_rank by hand; with the older static rendezvous you must pass --node_rank, --master_addr, and --master_port explicitly. Elastic rendezvous also lets you specify --nnodes=2:4 (min:max) so the job restarts from the last checkpoint with a smaller world when a node dies — see Checkpointing, Fault Tolerance & Long-Running Jobs.

NCCL: The Collective Communication Library¶

NVIDIA Collective Communications Library (NCCL) is the low-level library that torch.distributed calls under the hood when you use a GPU backend. NCCL implements a set of collective operations: routines where all \(N\) processes participate and the result depends on contributions from all of them. NCCL is highly optimized to exploit NVLink between GPUs on the same node and InfiniBand between nodes.

The Six Core Collectives¶

There are six operations you need to know cold. In all of them, let \(n\) be the number of ranks and \(M\) be the data size (bytes or elements) per rank before the operation.

Visually, for \(n=4\) ranks each holding a buffer \([a_0, a_1, a_2, a_3]\):

The key insight used in ZeRO-style data parallelism (see Distributed Training I) is that All-Reduce = Reduce-Scatter followed by All-Gather. This decomposition lets us interleave communication with computation.

torch.distributed Collective API¶

# collectives_demo.py — minimal, heavily-commented examples
import torch
import torch.distributed as dist

def demo_collectives(rank, world_size, device):
    """Demonstrate all six collectives on a simple tensor."""

    # ── 1. Broadcast ──────────────────────────────────────────────────────────
    # Only rank 0 sets meaningful data; others receive it.
    data = torch.zeros(4, device=device)
    if rank == 0:
        data = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
    dist.broadcast(data, src=0)
    # After: every rank holds [1, 2, 3, 4].

    # ── 2. Reduce ─────────────────────────────────────────────────────────────
    x = torch.tensor([float(rank + 1)], device=device)
    dist.reduce(x, dst=0, op=dist.ReduceOp.SUM)
    # After: rank 0 holds sum(1,2,...,world_size); others unchanged.

    # ── 3. All-Reduce ─────────────────────────────────────────────────────────
    x = torch.tensor([float(rank)], device=device)
    dist.all_reduce(x, op=dist.ReduceOp.SUM)
    # After: every rank holds sum(0, 1, ..., world_size-1).

    # ── 4. Reduce-Scatter ─────────────────────────────────────────────────────
    # Each rank contributes a buffer of size world_size;
    # rank i receives the sum of element i from all ranks.
    input_tensor = torch.arange(world_size, dtype=torch.float32, device=device) + rank
    output_tensor = torch.zeros(1, device=device)
    dist.reduce_scatter(output_tensor, list(input_tensor.split(1)), op=dist.ReduceOp.SUM)
    # Element i on rank r equals (i + r), so element i summed across all ranks is
    #   sum over r of (i + r) = world_size*i + world_size*(world_size-1)/2.
    # Rank i receives element i, so the outputs differ per rank (illustrating the scatter).
    # For world_size=4: rank 0 -> 6.0, rank 1 -> 10.0, rank 2 -> 14.0, rank 3 -> 18.0.

    # ── 5. All-Gather ─────────────────────────────────────────────────────────
    my_chunk = torch.tensor([float(rank)], device=device)
    gathered = [torch.zeros(1, device=device) for _ in range(world_size)]
    dist.all_gather(gathered, my_chunk)
    # After: gathered = [tensor(0), tensor(1), ..., tensor(world_size-1)] on every rank.

    # ── 6. All-to-All ─────────────────────────────────────────────────────────
    # Each rank sends a different value to each other rank.
    # input_list[j] goes TO rank j; output_list[j] comes FROM rank j.
    input_list  = [torch.tensor([rank * 10.0 + j], device=device) for j in range(world_size)]
    output_list = [torch.zeros(1, device=device) for _ in range(world_size)]
    dist.all_to_all(output_list, input_list)
    # Rank 0 receives [0*10+0, 1*10+0, 2*10+0, 3*10+0] = [0, 10, 20, 30]
    # (what each other rank sent to slot 0)

Asynchronous Operations¶

Every collective has an async variant that returns a Work handle, allowing you to overlap computation with communication:

# Overlap gradient all-reduce with backward pass (simplified DDP idea)
import torch
import torch.distributed as dist

def overlapped_allreduce_example(grad_tensor, device):
    """
    Fire off the all-reduce without blocking; do other work; then wait.
    This is the heart of DDP's gradient communication overlap.
    """
    # Non-blocking: returns immediately; communication starts in background.
    handle = dist.all_reduce(grad_tensor, op=dist.ReduceOp.SUM, async_op=True)

    # ... CPU work or other GPU kernels can run here ...
    # For example, you could compute loss.backward() for the next micro-batch.

    # Block until the all-reduce is complete before using grad_tensor.
    handle.wait()

    # Scale by 1/world_size to get the average gradient.
    grad_tensor /= dist.get_world_size()
    return grad_tensor

Point-to-Point: send, recv and batch_isend_irecv¶

Collectives are not the whole story. Pipeline parallelism moves activations from stage \(i\) to stage \(i+1\) and gradients back, which is a pure point-to-point (P2P) exchange between two ranks — no collective involved. NCCL exposes this as dist.send/dist.recv (blocking) and dist.isend/dist.irecv (async).

The footgun is ordering: if every rank calls blocking send first and recv second, and every send is large enough to exceed the transport’s internal buffering, the whole ring deadlocks because no one is receiving. The two safe patterns are (a) alternate the order by parity, or (b) use dist.batch_isend_irecv, which hands NCCL the whole set of P2P ops at once so it can group them into a single communication round:

# p2p_ring_exchange.py — each rank sends to rank+1 and receives from rank-1.
import torch
import torch.distributed as dist

def ring_exchange(send_buf: torch.Tensor, recv_buf: torch.Tensor):
    """Deadlock-free neighbour exchange — the primitive under 1F1B pipelining."""
    rank, world_size = dist.get_rank(), dist.get_world_size()
    next_rank = (rank + 1) % world_size
    prev_rank = (rank - 1) % world_size

    # Post BOTH ops together. NCCL fuses them into one grouped call, so there is
    # no send-before-recv ordering hazard even if every rank is symmetric.
    ops = [
        dist.P2POp(dist.isend, send_buf, next_rank),
        dist.P2POp(dist.irecv, recv_buf, prev_rank),
    ]
    for req in dist.batch_isend_irecv(ops):
        req.wait()
    return recv_buf

# Blocking alternative — correct ONLY because the order alternates by parity:
#   if rank % 2 == 0:
#       dist.send(send_buf, dst=next_rank); dist.recv(recv_buf, src=prev_rank)
#   else:
#       dist.recv(recv_buf, src=prev_rank); dist.send(send_buf, dst=next_rank)

Two rules that will save you hours: the receiver must allocate a buffer of exactly the right shape and dtype (NCCL P2P carries no metadata, so a shape mismatch silently corrupts memory or hangs), and both ranks must agree on the number and order of P2P ops. Megatron-LM’s p2p_communication.py is essentially this function plus shape negotiation and CUDA-stream bookkeeping; see Distributed Training II.

Ring All-Reduce: The Algorithm That Scaled Deep Learning¶

Naive all-reduce has rank 0 collect everything then broadcast — a star topology that makes rank 0 a bottleneck. Ring all-reduce eliminates this bottleneck and achieves near-optimal bandwidth utilization.

Algorithm¶

Arrange \(n\) ranks in a logical ring. The algorithm runs in two phases, each consisting of \(n-1\) steps:

Phase 1 — Reduce-Scatter: Each rank sends one chunk and receives one chunk per step, accumulating partial sums.

Phase 2 — All-Gather: Each rank sends fully-reduced chunks around the ring until every rank has every chunk.

The total data sent per rank per phase is \(\frac{n-1}{n} \cdot M\) (each step sends \(M/n\) bytes). Both phases together send \(2 \cdot \frac{n-1}{n} \cdot M\) bytes per rank, approaching \(2M\) for large \(n\). This is also the lower bound — ring all-reduce is bandwidth-optimal.

Tree All-Reduce¶

For small messages where the latency of \(2(n-1)\) steps dominates, a binary tree or recursive halving/doubling algorithm is better. In \(\log_2 n\) steps, a tree reduces all values. This is the regime of optimizer-state broadcast or short control messages.

The tradeoff:

The \(O(M \log n)\) row describes a naive binomial-tree reduce followed by a broadcast, where the same bytes traverse every level. NCCL’s actual Tree algorithm is smarter: it is a double binary tree (Sanders et al., 2009), which overlays two complementary trees so that every node is an interior node in one and a leaf in the other. This keeps the \(O(\log n)\) latency and recovers near-optimal bandwidth, which is why NCCL often prefers Tree over Ring for large-message all-reduce at high node counts — the ring’s \(2(n-1)\alpha\) term becomes intolerable long before its bandwidth term does.

NCCL selects among Ring, Tree, NVLS, CollNet and friends automatically, using an internal cost model over message size, topology, and rank count. You can override it for experiments with NCCL_ALGO=Tree / NCCL_ALGO=Ring and force a wire protocol with NCCL_PROTO=Simple|LL|LL128 (LL = low-latency, used for tiny messages; LL128 exploits 128-byte NVLink flits). Overriding is a diagnostic tool, not a tuning strategy — NCCL’s defaults beat hand-picked settings on almost every cluster.

Bandwidth and Latency Cost Models¶

To reason about whether your training run is compute-bound or communication-bound, you need a simple cost model. The alpha-beta model (also called the Hockney model; the LogP family of models is a richer refinement that additionally separates per-message CPU overhead from network occupancy) approximates the time to send a message of \(B\) bytes as:

where \(\alpha\) is the latency (startup cost in seconds, independent of message size) and \(\beta\) is the bandwidth (bytes per second, the asymptotic rate for large messages).

For a ring all-reduce of \(M\) total bytes across \(n\) ranks:

As \(n\) grows, the latency term \(2(n-1)\alpha\) becomes painful for many small messages — this is why gradient bucketing (combining small gradients into large tensors before communicating) is critical in practice.

For very large \(n\), the bandwidth term simplifies:

The bandwidth term is constant in \(n\) — adding more GPUs doesn’t change the bandwidth cost. The latency term grows linearly, which is why ring all-reduce across thousands of GPUs requires hierarchical approaches.

Network Topology: NVLink, NVSwitch, and InfiniBand¶

Communication cost is not uniform across a cluster. You need to understand the physical topology to reason about where collectives should be placed.

Intra-Node: NVLink and NVSwitch¶

NVLink is NVIDIA’s proprietary high-bandwidth GPU-to-GPU interconnect. On an H100 SXM server, each GPU has 18 NVLink 4.0 lanes, providing roughly 900 GB/s bidirectional bandwidth per GPU — vastly higher than PCIe Gen 5 (on the order of 128 GB/s bidirectional). H100/H200 (NVLink 4) remain common training and inference hardware as of 2026, but the current frontier generation is Blackwell, whose fifth-generation NVLink roughly doubles per-GPU bandwidth to about 1.8 TB/s.

NVSwitch is a crossbar switch that connects all GPUs on a node with full NVLink bandwidth — every GPU can communicate with every other GPU simultaneously at full speed, rather than routing through a chain. An 8-GPU DGX H100 uses four NVSwitch 3.0 chips, providing an effective 3.6 TB/s of all-to-all bandwidth within the node. NVIDIA’s Blackwell-generation GB200 NVL72 rack extends this idea to rack scale: NVSwitch fabric ties 72 GPUs into a single NVLink domain with roughly 130 TB/s of aggregate GPU-to-GPU bandwidth, letting an all-to-all or all-reduce span far more GPUs before ever touching the slower inter-node network.

This topology means:

NCCL exploits NVSwitch by using its own all-reduce algorithm that leverages the all-to-all connectivity, avoiding a sequential ring in favor of a one-shot reduce pattern across the switch fabric. On Hopper and Blackwell systems this is the NVLS algorithm (NVLink SHARP): the NVSwitch chips contain arithmetic units that perform the reduction inside the switch, so each GPU pushes its buffer once and pulls the reduced result once instead of relaying \(2\frac{n-1}{n}M\) bytes around a ring. InfiniBand has the analogous SHARP in-network reduction in Quantum switches for the inter-node hop. You will see NVLS and NVLS Tree appear in NCCL_DEBUG=INFO logs when NCCL picks these paths.

Inter-Node: InfiniBand¶

Between nodes, current clusters use InfiniBand (IB). Common configurations:

• HDR (200 Gb/s): ~25 GB/s effective unidirectional per port
• NDR (400 Gb/s): ~50 GB/s effective unidirectional per port, the mainstream choice in most H100/H200-era clusters
• XDR (800 Gb/s): shipping since 2024 (e.g., NVIDIA Quantum-X800 switches) and increasingly common in new Blackwell-generation deployments as of 2026

A cluster of nodes is connected through an IB fabric, often organized as a fat-tree or dragonfly topology, providing full bisection bandwidth in principle (but subject to hotspots in practice). The IB Host Channel Adapter (HCA) on each node connects the CPUs and GPUs to the fabric; GPU Direct RDMA (Remote Direct Memory Access) allows the NIC to read/write GPU HBM directly, bypassing the CPU.

Hierarchical Communication¶

Because intra-node bandwidth is ~10–50× higher than inter-node bandwidth, production NCCL jobs use hierarchical collectives: first reduce within a node (fast, NVLink), then reduce across nodes (slower, IB), then broadcast back. NCCL implements this automatically via its topology detection.

A Complete torch.distributed Training Step¶

The following implements a minimal DDP-style training loop from scratch, showing exactly where each collective fires and why.

# minimal_ddp.py
# Run: torchrun --nproc_per_node=4 minimal_ddp.py
import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.utils.data import DataLoader, TensorDataset, DistributedSampler

# ──────────────────────────────────────────────────────────────────────────────
# 1. Initialize distributed environment
# ──────────────────────────────────────────────────────────────────────────────
dist.init_process_group(backend="nccl")
rank       = dist.get_rank()
world_size = dist.get_world_size()
local_rank = int(os.environ["LOCAL_RANK"])   # node-local device index
device     = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)

# ──────────────────────────────────────────────────────────────────────────────
# 2. Build a toy model. Every rank starts with a random init;
#    we must broadcast weights from rank 0 so all ranks start identically.
# ──────────────────────────────────────────────────────────────────────────────
class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(128, 10)

    def forward(self, x):
        return self.linear(x)

model = TinyModel().to(device)

# Broadcast initial parameters so all ranks are identical.
for param in model.parameters():
    dist.broadcast(param.data, src=0)

# ──────────────────────────────────────────────────────────────────────────────
# 3. Build a toy dataset. DistributedSampler ensures each rank gets
#    a non-overlapping shard of the data, which is the "data parallel" part.
# ──────────────────────────────────────────────────────────────────────────────
N = 1024
dataset = TensorDataset(
    torch.randn(N, 128),
    torch.randint(0, 10, (N,))
)
sampler    = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# ──────────────────────────────────────────────────────────────────────────────
# 4. Training loop: forward → backward → all-reduce gradients → optimizer step
# ──────────────────────────────────────────────────────────────────────────────
for epoch in range(2):
    sampler.set_epoch(epoch)  # ensures different shuffling per epoch

    for x_batch, y_batch in dataloader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        optimizer.zero_grad()
        logits = model(x_batch)
        loss   = criterion(logits, y_batch)
        loss.backward()

        # ── All-Reduce: average gradients across all ranks ──────────────────
        # Without this, each rank would update with its own local gradient,
        # causing model divergence after the first step.
        for param in model.parameters():
            if param.grad is not None:
                # SUM then divide → average. The NCCL backend also supports
                # op=dist.ReduceOp.AVG directly, which is both fewer kernels and
                # numerically safer in bf16 (it divides inside the reduction
                # tree instead of letting the sum grow by a factor of world_size
                # before scaling). Gloo/CPU does not support AVG.
                dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
                param.grad.data /= world_size
        # ────────────────────────────────────────────────────────────────────

        optimizer.step()

    if rank == 0:
        print(f"Epoch {epoch}: loss = {loss.item():.4f}")

dist.destroy_process_group()

In production, torch.nn.parallel.DistributedDataParallel (DDP) wraps this loop and uses a more sophisticated gradient bucketing scheme to maximize overlap between backward computation and all-reduce communication. The principle is identical to the manual loop above.

Process Groups and Custom Communicators¶

So far we’ve used the default process group containing all ranks. For pipeline parallelism and tensor parallelism (see Distributed Training II), you need sub-groups where only a subset of ranks communicate.

# process_groups.py — illustrating tensor-parallel sub-groups
import torch.distributed as dist

def build_tp_dp_groups(tp_size: int):
    """
    Construct tensor-parallel (TP) and data-parallel (DP) process groups.

    For world_size=8 and tp_size=4:
      TP groups (communicate tensor shards):  [0,1,2,3] and [4,5,6,7]
      DP groups (communicate gradients):      [0,4], [1,5], [2,6], [3,7]

    Within a TP group, all-gather/reduce-scatter move activation shards.
    Across DP groups, all-reduce moves gradient updates.
    """
    world_size = dist.get_world_size()
    rank       = dist.get_rank()
    assert world_size % tp_size == 0
    dp_size = world_size // tp_size

    # ── Tensor-parallel groups ────────────────────────────────────────────────
    tp_group = None
    for i in range(dp_size):
        ranks_in_group = list(range(i * tp_size, (i + 1) * tp_size))
        group = dist.new_group(ranks=ranks_in_group)
        if rank in ranks_in_group:
            tp_group = group

    # ── Data-parallel groups ──────────────────────────────────────────────────
    dp_group = None
    for j in range(tp_size):
        ranks_in_group = list(range(j, world_size, tp_size))
        group = dist.new_group(ranks=ranks_in_group)
        if rank in ranks_in_group:
            dp_group = group

    return tp_group, dp_group

# Usage:
# tp_group, dp_group = build_tp_dp_groups(tp_size=4)
# dist.all_reduce(tensor_shard, group=tp_group)   # fast, NVLink
# dist.all_reduce(grad,         group=dp_group)   # slower, IB

NCCL treats each new_group call as a new communicator; it internally runs its topology detection and algorithm selection within that group.

DeviceMesh: The Library API for the Same Idea¶

Hand-rolling nested new_group loops is exactly the code that gets subtly wrong when you go from 2D (DP × TP) to 4D (DP × TP × PP × CP). PyTorch’s DeviceMesh (torch.distributed.device_mesh) is the supported abstraction for it: you declare an \(n\)-dimensional logical grid of devices with named dimensions, and it builds and caches every sub-group for you. It is the substrate that FSDP2, TensorParallel, and DTensor are all written against.

# device_mesh_demo.py — the same TP/DP decomposition as above, in 3 lines.
from torch.distributed.device_mesh import init_device_mesh

# world_size = 8, laid out as dp=2 x tp=4.
# Mesh dimensions vary fastest on the RIGHT, so this yields exactly:
#   tp groups: [0,1,2,3], [4,5,6,7]      (contiguous → lands on NVLink)
#   dp groups: [0,4], [1,5], [2,6], [3,7]
mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))

tp_group = mesh["tp"].get_group()   # a real ProcessGroup, usable with dist.*
dp_group = mesh["dp"].get_group()

# dist.all_reduce(activation_shard, group=tp_group)  # fast, intra-node
# dist.all_reduce(grad,             group=dp_group)  # slower, inter-node

# Higher-level APIs consume the mesh directly:
#   from torch.distributed.fsdp import fully_shard
#   fully_shard(model, mesh=mesh["dp"])
#   from torch.distributed.tensor.parallel import parallelize_module
#   parallelize_module(block, mesh["tp"], {...})

Note the ordering rule, which is the whole reason meshes are worth using: the last mesh dimension is the fastest-varying, so putting tp last guarantees tensor-parallel ranks are contiguous and therefore co-located on one node behind NVLink, while the data-parallel dimension is the one that straddles the slow InfiniBand hop. Getting that backwards puts your highest-volume collective on your slowest link. Megatron-LM encodes the same convention in its parallel_state.py rank-ordering logic (see Megatron-LM, DeepSpeed & Parallelism in Practice).

Collective Performance: Benchmarking and Profiling¶

Understanding actual achieved bandwidth versus theoretical peak is essential for diagnosing training slowdowns.

# benchmark_allreduce.py — measure achieved bandwidth for a ring all-reduce
import os
import time
import torch
import torch.distributed as dist

def benchmark_all_reduce(message_bytes: int, n_iters: int = 50):
    """
    Measure the achieved bandwidth of all-reduce for a given message size.
    Returns (latency_ms, bandwidth_GBps).
    """
    rank       = dist.get_rank()
    world_size = dist.get_world_size()
    device     = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")

    # Allocate a float32 buffer of the requested size.
    n_elements = message_bytes // 4
    buf = torch.randn(n_elements, device=device)

    # Warm up: let NCCL initialize its internal state.
    for _ in range(5):
        dist.all_reduce(buf, op=dist.ReduceOp.SUM)
    torch.cuda.synchronize()
    dist.barrier()

    # Timed runs.
    start = time.perf_counter()
    for _ in range(n_iters):
        dist.all_reduce(buf, op=dist.ReduceOp.SUM)
    torch.cuda.synchronize()
    dist.barrier()
    elapsed = time.perf_counter() - start

    latency_ms = elapsed / n_iters * 1000
    # Busbw formula: 2*(n-1)/n * message_bytes / time  (ring all-reduce)
    busbw_GBps = (2 * (world_size - 1) / world_size * message_bytes) / (elapsed / n_iters) / 1e9

    if rank == 0:
        print(f"Message: {message_bytes/1e6:.1f} MB | "
              f"Latency: {latency_ms:.2f} ms | "
              f"Bus BW: {busbw_GBps:.1f} GB/s")
    return latency_ms, busbw_GBps


# Example results on an 8-GPU DGX H100 (NVLink):
# Message:   1.0 MB | Latency:  0.15 ms | Bus BW:  52.3 GB/s
# Message:  64.0 MB | Latency:  1.23 ms | Bus BW: 413.8 GB/s
# Message: 512.0 MB | Latency:  7.84 ms | Bus BW: 521.0 GB/s
# (Illustrative figures; actual results depend on driver version and cluster state)

The busbw (bus bandwidth) formula \(\frac{2(n-1)}{n} \cdot M / T\) is the standard metric because it accounts for the ring algorithm’s traffic pattern and is comparable across cluster sizes. Compare busbw to the theoretical NVLink bandwidth to understand efficiency. Note the correction factor is collective-specific: all-reduce uses \(\frac{2(n-1)}{n}\), while reduce-scatter and all-gather each use \(\frac{n-1}{n}\), and broadcast/reduce use \(1\). Using the wrong factor is the most common way people convince themselves a cluster is twice as fast as it is.

nccl-tests: The Standard Cluster Benchmark¶

Before blaming your training code, benchmark the fabric with the vendor tool. NVIDIA/nccl-tests is what every cluster acceptance test uses, and it reports both algbw (raw bytes/time) and busbw (the corrected metric above) at every message size:

# Build (MPI=1 if you want to sweep multiple nodes via mpirun).
git clone https://github.com/NVIDIA/nccl-tests && cd nccl-tests
make MPI=0 CUDA_HOME=/usr/local/cuda -j

# Sweep 8 B → 8 GB, doubling each step, across 8 GPUs on this node.
#   -b begin size   -e end size   -f size multiplier   -g GPUs per process
./build/all_reduce_perf -b 8 -e 8G -f 2 -g 8
# Columns:  size  count  type  redop  time(us)  algbw(GB/s)  busbw(GB/s)  #wrong
# Read the plateau of the busbw column: that is your real all-reduce bandwidth.

# Same sweep for the MoE-critical collective:
./build/alltoall_perf -b 1M -e 1G -f 2 -g 8

The environment variables you will actually reach for:

export NCCL_DEBUG=INFO                     # print topology + chosen algorithm
export NCCL_DEBUG_SUBSYS=INIT,GRAPH,ENV    # narrow the (very verbose) output
export NCCL_SOCKET_IFNAME=eth0             # which NIC carries bootstrap traffic
export NCCL_IB_HCA=mlx5_0,mlx5_1           # which InfiniBand HCAs to use
export NCCL_IB_DISABLE=0                   # set 1 to force TCP (diagnostic only)
export NCCL_P2P_DISABLE=0                  # set 1 to disable NVLink/P2P (diagnostic)

# PyTorch-side hang debugging (names are TORCH_*-prefixed since PyTorch 2.2):
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1   # crash the rank instead of hanging
export TORCH_NCCL_TRACE_BUFFER_SIZE=2000   # NCCL "flight recorder": on a watchdog
                                           # timeout, dump the last N collectives
                                           # per rank so you can see WHICH rank
                                           # failed to enter WHICH collective.

The flight recorder is the single highest-value debugging tool for a hung multi-node job: a collective mismatch (one rank taking a different code path, as in the deadlock pitfall below) shows up immediately as one rank stuck on ALLREDUCE sequence number \(k\) while everyone else is at \(k+1\).

For production profiling, use PyTorch Profiler with NCCL tracing enabled:

from torch.profiler import profile, ProfilerActivity, schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=5),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./logs/profiler"),
    record_shapes=True,
    with_stack=True,
) as prof:
    for step, (x, y) in enumerate(dataloader):
        train_step(x, y, model, optimizer)
        prof.step()

# In TensorBoard → Distributed → you will see the NCCL ops timeline,
# showing exactly how much time is spent in AllReduce vs compute.

All-to-All and Expert Parallelism¶

The all-to-all collective is less common in vanilla data parallelism but is critical for Mixture-of-Experts (MoE) models (see Mixture-of-Experts (MoE) Architectures). In MoE, tokens are routed to expert sub-networks that may live on different GPUs. Sending each token to its assigned expert is exactly an all-to-all operation.

# all_to_all_moe_sketch.py — illustrating token dispatch in MoE
import torch
import torch.distributed as dist

def moe_dispatch(tokens: torch.Tensor, expert_ids: torch.Tensor, n_experts: int):
    """
    tokens: [seq_len, hidden_dim]  — tokens on this GPU
    expert_ids: [seq_len]          — which expert each token should go to
    n_experts: total number of experts, one per GPU

    Returns: tokens sorted and dispatched to the correct expert GPU.
    """
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    assert n_experts == world_size  # one expert per GPU for simplicity

    seq_len, hidden = tokens.shape

    # Count how many tokens go to each expert.
    counts = torch.bincount(expert_ids, minlength=world_size)  # [world_size]

    # All-to-all the counts so each rank knows what it will receive.
    recv_counts = torch.zeros_like(counts)
    dist.all_to_all_single(recv_counts, counts)

    # Sort tokens by destination expert.
    sort_idx = expert_ids.argsort()
    sorted_tokens = tokens[sort_idx]  # [seq_len, hidden], grouped by expert

    # Build send/recv lists for all-to-all.
    send_splits = counts.tolist()
    recv_splits = recv_counts.tolist()

    # Allocate receive buffer.
    total_recv = recv_counts.sum().item()
    recv_buf = torch.zeros(total_recv, hidden, device=tokens.device, dtype=tokens.dtype)

    # Perform the actual all-to-all data transfer.
    dist.all_to_all_single(
        recv_buf,
        sorted_tokens,
        output_split_sizes=recv_splits,
        input_split_sizes=send_splits,
    )

    # recv_buf now contains tokens that belong to this GPU's expert.
    return recv_buf, recv_counts

The all-to-all cost is \(O(M)\) where \(M\) is the total token volume, and unlike all-reduce it does not benefit from the ring structure — it is inherently limited by the bisection bandwidth of the network.

Key Considerations for Choosing Collectives¶

Different distributed training strategies use different collectives as their communication primitive:

Understanding this table is what separates an engineer who can debug a distributed training job from one who cannot. If your profiler shows an all-gather is slow, you know FSDP is materializing parameters and the bottleneck is memory bandwidth for the parameter gather, not compute. If all-to-all is slow, you have an MoE routing or load-imbalance problem.

# WRONG — only rank 0 calls all-reduce; all others hang forever.
if rank == 0:
    dist.all_reduce(tensor)  # deadlock!

# CORRECT — every rank calls every collective, always.
dist.all_reduce(tensor)
if rank == 0:
    process(tensor)  # branch after the collective, not before.

Further Reading¶

• Thakur, Rabenseifner, Gropp (2005): “Optimization of Collective Communication Operations in MPICH” — the foundational paper on ring and tree all-reduce algorithms and their cost models.
• Li et al. (2020): “PyTorch Distributed: Experiences on Accelerating Data Parallel Training” — covers DDP bucketing, hook-based gradient compression, and the design choices in torch.distributed.
• NVIDIA NCCL Documentation and Source (github.com/NVIDIA/nccl) — the authoritative reference for NCCL algorithm selection, topology detection, and tuning knobs.
• Rajbhandari et al. (2020): “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models” (DeepSpeed) — explains how reduce-scatter + all-gather enables full optimizer/gradient/parameter sharding.
• Jiang et al. (2022): “Megascale: Scaling Large Language Model Training to More Than 10,000 GPUs” — describes hierarchical collectives, network topology design, and reliability engineering at cluster scale.
• Shoeybi et al. (2019): “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism” — the canonical reference for tensor-parallel all-reduce patterns within a transformer layer.

Exercises¶

1. (Conceptual) Why one process per GPU? The chapter states that GPU distributed training “almost universally uses the one-process-per-GPU model rather than threading across GPUs.” Give the three reasons the chapter offers, and explain which of them would not be fixed simply by removing Python’s GIL.

2. (Conceptual) The forbidden conditional collective. Consider this snippet, intended to log the global gradient norm only from rank 0:

```python
if rank == 0:
    dist.all_reduce(grad_norm, op=dist.ReduceOp.SUM)
    log(grad_norm)
```

On a 4-GPU job, describe exactly what happens at runtime and why. Then rewrite the snippet so it is correct.

dist.all_reduce(grad_norm, op=dist.ReduceOp.SUM)  # all ranks participate
if rank == 0:
    log(grad_norm)                                 # branch AFTER the collective

3. (Quantitative) Ring all-reduce with the alpha-beta model. You all-reduce a gradient buffer of \(M = 512\) MB across \(n = 16\) GPUs on a fabric with latency \(\alpha = 5\ \mu\text{s}\) and per-direction bandwidth \(\beta = 50\) GB/s. Using the chapter’s ring all-reduce cost model, compute (a) the latency term, (b) the bandwidth term, and © the total time. (d) Which term dominates, and what practical technique from the chapter attacks the other term? Use \(1\text{ GB} = 10^9\) bytes.

4. (Quantitative) Reduce-scatter output values. Reproduce the chapter’s reduce-scatter example by hand for \(n = 4\). Each rank \(r\) builds input = torch.arange(4) + r, i.e. rank \(r\) holds the vector \([r, r{+}1, r{+}2, r{+}3]\). The collective sums element \(i\) across all ranks and delivers element \(i\) to rank \(i\). Derive the closed form for the value rank \(i\) receives, then list the four outputs.

5. (Implementation) All-reduce built from reduce-scatter + all-gather. The chapter’s key identity is All-Reduce = Reduce-Scatter followed by All-Gather. Implement a function manual_all_reduce_sum(x) that reproduces dist.all_reduce(x, op=SUM) using only dist.reduce_scatter and dist.all_gather (no dist.all_reduce). Assume x is a 1-D tensor whose length is divisible by world_size. Add an assertion that verifies your result against the real all_reduce.

import torch
import torch.distributed as dist

def manual_all_reduce_sum(x: torch.Tensor) -> torch.Tensor:
    """All-Reduce(SUM) via Reduce-Scatter + All-Gather.
    x: 1-D tensor with len(x) % world_size == 0. Returns the summed tensor."""
    world_size = dist.get_world_size()
    n = x.numel()
    assert n % world_size == 0, "length must be divisible by world_size"
    chunk = n // world_size

    # --- Phase 1: Reduce-Scatter ---
    # Split x into world_size chunks; rank i receives the SUM of chunk i.
    input_chunks = list(x.split(chunk))
    reduced_chunk = torch.zeros(chunk, device=x.device, dtype=x.dtype)
    dist.reduce_scatter(reduced_chunk, input_chunks, op=dist.ReduceOp.SUM)

    # --- Phase 2: All-Gather ---
    # Every rank collects all reduced chunks and concatenates them.
    gathered = [torch.zeros(chunk, device=x.device, dtype=x.dtype)
                for _ in range(world_size)]
    dist.all_gather(gathered, reduced_chunk)
    return torch.cat(gathered)

# ---- Verification (run under torchrun --nproc_per_node=N) ----
def check():
    rank = dist.get_rank()
    device = torch.device(f"cuda:{rank}")
    x = torch.arange(8, dtype=torch.float32, device=device) + rank

    mine = manual_all_reduce_sum(x.clone())

    reference = x.clone()
    dist.all_reduce(reference, op=dist.ReduceOp.SUM)

    assert torch.allclose(mine, reference), (rank, mine, reference)
    if rank == 0:
        print("OK: manual all-reduce matches dist.all_reduce")

6. (Implementation / Conceptual) Hierarchical all-reduce over sub-groups. Using dist.new_group, sketch a two-level all-reduce for 8 GPUs on 2 nodes (ranks 0-3 on node A, 4-7 on node B): first reduce within each node over fast NVLink, then reduce across nodes over slow InfiniBand, then broadcast back. Explain why this moves less data over the slow inter-node link than a flat 8-way ring all-reduce would.

import torch
import torch.distributed as dist

def hierarchical_all_reduce_sum(x, gpus_per_node=4):
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    node_id = rank // gpus_per_node
    n_nodes = world_size // gpus_per_node

    # Intra-node groups: [0,1,2,3] and [4,5,6,7]
    intra_group = None
    for nid in range(n_nodes):
        ranks = list(range(nid * gpus_per_node, (nid + 1) * gpus_per_node))
        g = dist.new_group(ranks=ranks)
        if rank in ranks:
            intra_group = g

    # Inter-node "leaders" group: one rank per node, e.g. [0, 4]
    leader_ranks = list(range(0, world_size, gpus_per_node))
    leader_group = dist.new_group(ranks=leader_ranks)  # every rank must call
    is_leader = (rank % gpus_per_node == 0)

    # 1. Reduce within node onto the local leader (dst is a GLOBAL rank).
    local_leader = node_id * gpus_per_node
    dist.reduce(x, dst=local_leader, op=dist.ReduceOp.SUM, group=intra_group)

    # 2. Leaders all-reduce across nodes over InfiniBand.
    if is_leader:
        dist.all_reduce(x, op=dist.ReduceOp.SUM, group=leader_group)

    # 3. Broadcast the global sum back down within each node.
    dist.broadcast(x, src=local_leader, group=intra_group)
    return x