Distributed Training Fundamentals

"My weights weigh a terabyte, my optimizer state weighs four more, and the single H100 you bought me holds eighty gigabytes. Please pass the NVLink."

Tensor, Memory-Constrained Sharding AI Agent

59.1.1 Why One GPU Is Not Enough

The memory accounting for training a large transformer is brutal. For each parameter $\theta$ stored in bfloat16 (2 bytes), the AdamW optimizer keeps an fp32 master copy (4 bytes), an fp32 first moment $m$ (4 bytes), an fp32 second moment $v$ (4 bytes), and an fp32 gradient buffer (4 bytes). That is $18$ bytes per parameter just for state. For a 70-billion-parameter model:

Add activation memory ($\mathcal{O}(BL \cdot d_{\text{model}} \cdot n_{\text{layers}})$ for batch $B$, sequence length $L$) and the working set easily exceeds 2 TB. An H100 holds 80 GB. Even Blackwell B200 with 192 GB is not close. Frontier training therefore must shard state across many devices; the only open question is which axis to shard along.

59.1.2 Three Axes of Parallelism

Every parallelism scheme partitions one of three dimensions of the training computation: the data batch, the layer stack, or the tensor itself. These three axes are orthogonal: production training runs compose all three (often called 3D parallelism, covered in Section 59.4).

59.1.2.1 Data Parallelism

Each of $N$ devices holds a complete model replica and processes a different slice of the global mini-batch. After backward, gradients are all-reduced across devices, then each replica performs an identical optimizer step. This is the workhorse: when it fits in memory, nothing else matches its simplicity. Gradient $g$ on device $i$ updates as:

Because the optimizer state is replicated, data parallelism scales throughput but not model size. A 70B model that fits on $k$ GPUs via tensor-parallel still fits on $k$ GPUs after going data-parallel; data parallelism gives you more samples per second, not more parameters.

59.1.2.2 Model Parallelism (Pipeline)

Partition the layer stack: GPU 0 holds layers $1 \ldots k$, GPU 1 holds layers $k+1 \ldots 2k$, and so on. Activations flow point-to-point between consecutive stages. The challenge is keeping all stages busy simultaneously, which requires micro-batching the work so that stage $i$ processes micro-batch $t$ while stage $i+1$ processes micro-batch $t-1$. Section 59.4 develops this in detail.

59.1.2.3 Tensor Parallelism

Inside a single layer, the weight matrix itself is sharded across devices. For an FFN up-projection $Y = X W$ with $W \in \mathbb{R}^{d \times 4d}$, split $W$ along its columns into $W = [W_1 \mid W_2 \mid \ldots \mid W_T]$ where $T$ is the tensor-parallel degree. Each device computes $Y_t = X W_t$ locally and a collective op stitches the partial outputs into the full result. Communication happens per layer, so tensor parallelism requires extremely fast interconnects (typically NVLink within a single node, never across InfiniBand). Section 59.3 covers Megatron-LM's careful design of which ops are column- versus row-parallel.

59.1.3 Collective Communication Primitives

The currency of distributed training is the collective: an operation that runs in lockstep across a group of devices. NCCL (NVIDIA Collective Communications Library) implements these on NVIDIA hardware; for AMD it is RCCL, and oneCCL on Intel. The four primitives that matter for training:

A useful identity: all-reduce = reduce-scatter + all-gather. This decomposition is what FSDP exploits: instead of one all-reduce that materializes the full gradient on every rank, FSDP does a reduce-scatter so each rank ends up with only its shard of the gradient. The same total bytes move on the wire, but the per-rank peak memory is $N\times$ smaller. We will see this trick again and again.

59.1.3.1 The Ring All-Reduce

The classic NCCL implementation arranges the $N$ devices in a logical ring and performs $2(N-1)$ steps. In the first $N-1$ steps (the reduce-scatter phase), each device sends a $S/N$-byte chunk to the next neighbor and receives a chunk from the previous one, accumulating partial sums. In the next $N-1$ steps (the all-gather phase), the reduced chunks circulate around the ring so every device ends up with the full sum.

The total bandwidth per device is:

Critically, this is independent of $N$ in the asymptotic limit: doubling the cluster size does not double the wire time for an all-reduce. That is why data-parallel training scales gracefully to thousands of GPUs; the per-step communication is bounded.

59.1.3.2 Tree-Based All-Reduce

For very small messages (typical at startup or in inference), the ring's $\mathcal{O}(N)$ latency dominates. NCCL switches to a tree topology at small message sizes: the data is reduced up a binary tree in $\mathcal{O}(\log N)$ steps and broadcast back down. The crossover point depends on the latency-to-bandwidth ratio of the link, which NCCL auto-tunes per fabric.

59.1.4 The Interconnect Hierarchy

Not all communication paths are equal. A modern training cluster has at least three tiers, and the parallelism plan must respect this hierarchy.

59.1.4.1 NVLink and NVSwitch

Within a single node (DGX, HGX, or similar), NVIDIA's NVLink interconnect provides ~900 GB/s per direction between GPUs on H100 (NVLink 4) and 1.8 TB/s on B200 (NVLink 5). An NVSwitch crossbar gives all-to-all connectivity at full bandwidth within an 8-GPU server. This is fast enough that the per-layer all-reduces of tensor parallelism are tolerable.

59.1.4.2 InfiniBand and RoCE

Across nodes, you use a separate fabric: either Mellanox / NVIDIA InfiniBand (NDR 400 Gb/s per port in 2026, XDR 800 Gb/s on the horizon) or RDMA-over-Converged-Ethernet (RoCE v2). Both support remote direct memory access (RDMA), which lets one GPU's network card read or write another GPU's HBM directly without CPU involvement. NCCL transparently uses RDMA when available.

Bandwidth between any two GPUs in different nodes is roughly $400 \text{ Gb/s} = 50 \text{ GB/s}$ per direction, or about $20\times$ slower than NVLink. This factor is the single largest determinant of how to lay out your parallelism plan: tensor-parallel groups must fit inside a node; pipeline and data-parallel groups can cross nodes.

59.1.4.3 Spine and Rail Topologies

Above 256 GPUs, you cross a second tier of switches called the spine. Many production fabrics oversubscribe the spine (2:1 or 4:1) because intra-pod traffic is far more common than inter-pod, so cross-cluster bandwidth is meaningfully lower per pair. Designs like the "rail-optimized" topology (favored by Meta, Microsoft) dedicate one rail per GPU index so that all_reduce on GPU-3-of-each-node never contends with GPU-4's traffic.

59.1.5 The BSP Model and Sync vs Async

The mathematical contract of standard distributed training is the bulk synchronous parallel (BSP) model, due to Leslie Valiant in 1990. A BSP step consists of:

1. Local computation: each device runs forward + backward, producing local gradients.
2. Communication: all devices participate in a collective (all-reduce, or reduce-scatter + all-gather).
3. Synchronization barrier: no device proceeds until the collective is complete.
4. Optimizer step: each device applies the now-identical gradient.

BSP guarantees that all replicas remain bitwise identical after each step, which makes distributed training mathematically equivalent to a single-device run with a larger batch size. The price is the barrier: if device $i$ is slow this step, every other device idles.

59.1.5.1 The Asynchronous Alternative

For decades, distributed ML used the asynchronous parameter server model: workers push gradients to a central server whenever they finish, and pull the latest parameters whenever they start a new step. Workers see slightly stale parameters; the optimizer must tolerate the staleness. Hogwild! (Niu et al., 2011) showed this can work for sparse SGD, and DistBelief / TensorFlow's PS architecture (Dean et al., 2012) scaled it to thousands of CPUs at Google.

Async lost to BSP for dense neural network training, for three reasons that still hold in 2026:

1. Staleness hurts large-batch optimizers. AdamW and LION are extraordinarily sensitive to gradient direction; even a few-step delay shifts the moment estimates enough to slow convergence by 10-30%.
2. Reproducibility. Async training is non-deterministic at the bit level. Debugging a divergence at step 25,000 is hard enough; debugging a non-reproducible one is hopeless.
3. Hardware uniformity. Modern clusters are homogeneous: stragglers are rare, and when they happen you want to find and fix them, not paper over them with async.

Async survives at the edges: federated learning (Chapter 60 covers this), continual / online learning, and recently in decentralized training over the public internet (DiLoCo, OpenDiLoCo). But for frontier pretraining, BSP rules.

59.1.6 A Minimal DDP Training Loop

To make all of this concrete, the canonical PyTorch DistributedDataParallel (DDP) loop fits in fewer than 40 lines. Every more advanced scheme in this chapter starts from this skeleton.

# minimal_ddp.py: torchrun --nproc_per_node=8 minimal_ddp.py
import os
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup():
    # torchrun sets RANK, LOCAL_RANK, WORLD_SIZE automatically
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup():
    dist.destroy_process_group()

def train(local_rank, model, dataset, epochs=1, lr=3e-4):
    # 1) Build a sampler that gives each rank a disjoint slice of the data
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=8, sampler=sampler)

    # 2) Move the model to this rank's GPU and wrap with DDP
    model = model.to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # so shuffling differs across epochs
        for batch in loader:
            x, y = batch["input"].to(local_rank), batch["label"].to(local_rank)
            optimizer.zero_grad(set_to_none=True)
            loss = nn.functional.cross_entropy(model(x), y)
            loss.backward()         # DDP hooks fire here, all-reducing grads
            optimizer.step()        # bitwise-identical update on every rank

if __name__ == "__main__":
    local_rank = setup()
    model = nn.Linear(1024, 10)            # stand-in for the real model
    dataset = MyDataset()                  # any torch.utils.data.Dataset
    train(local_rank, model, dataset)
    cleanup()

Three things to notice:

• DistributedSampler is essential. Without it, every rank trains on the same data and the all-reduced gradient is identical to a single-rank gradient on a smaller batch; you have spent N GPUs to achieve 1 GPU's effective step.
• DDP's gradient bucketing overlaps communication with compute. The first parameters whose backward finishes have their gradients all-reduced while later layers are still computing. This is why DDP can sustain > 90% scaling efficiency even on modest interconnects.
• The optimizer step is local but produces an identical result on every rank, because every rank applies the same averaged gradient to the same starting parameters. This invariant is the BSP contract; the rest of this chapter is about preserving it under sharding.

59.1.7 Gradient Accumulation and the Effective Batch

Frontier training runs care about a quantity called the effective batch size: the number of independent samples that contribute to a single optimizer step. Effective batch determines the noise scale of the gradient estimate and, indirectly, the optimal learning rate (McCandlish et al., 2018, formalized this as the gradient noise scale). For a 100B+ model, the empirical sweet spot is roughly $4$M to $16$M tokens per step.

You cannot necessarily fit that many tokens through the model at once. Two levers expand the effective batch without changing the parallelism plan:

• Data parallelism. Each of $N$ replicas processes its own slice of the global batch; gradients are averaged before the step. Effective batch scales linearly with $N$.
• Gradient accumulation. Within a single replica, $K$ micro-steps run forward + backward without calling optimizer.step(); their gradients accumulate in place. The all-reduce happens once at the end of the $K$ micro-steps. Effective batch scales by $K$ at zero communication cost.

Gradient accumulation is the operational secret to running with massive effective batches on modest cluster sizes: a $256$-GPU cluster with $K=16$ accumulation steps and per-rank micro-batch $4$ achieves effective batch $256 \cdot 16 \cdot 4 = 16{,}384$ sequences without ever materializing all 16k in memory at once. PyTorch's no_sync() context manager (covered in Section 59.2 for FSDP) is the canonical way to implement this; in plain DDP, you simply omit the optimizer.step() until the final micro-step.

59.1.8 The Cost of Correctness: Why Synchronous Wins

To close this section, a worked example of why frontier training stayed synchronous even at thousand-GPU scale. Consider three competing patterns:

The 5-15% throughput penalty of sync DP versus async is small compared to the engineering cost of debugging a non-deterministic loss curve. The cases where async approaches win are precisely the cases where reproducibility is already lost: cross-cluster training (DiLoCo) where you cannot make hundreds of nodes finish a step in lock-step anyway, and federated learning where the participants are heterogeneous by design. For the standard "one cluster, one frontier model" case, sync wins.