Distributed Training

Distributed training runs one training job on several devices, usually GPUs, so that a network trains faster or a network too large for one device can be trained at all. The devices cooperate on the same model and the same optimization: in the common forms, each step produces the same parameter update that a single device would compute on the whole mini-batch, and the differences lie in what each device stores and computes and in what the devices send each other. The work can be split along three axes: the batch (data parallelism), the training state of parameters, gradients and optimizer moments (sharded data parallelism, as in ZeRO and PyTorch’s FSDP), and the model’s computation (tensor parallelism, which splits individual layers, and pipeline parallelism, which splits the sequence of layers). Large runs combine them. Two single-device techniques belong to the same toolkit: gradient accumulation, which builds a large batch from several small ones, and activation checkpointing, which trades computation for memory.

The PyTorch article describes the library’s wrappers; this article covers the methods, why they give the same gradient, where they do not, and what they cost.

Purpose

Two limits drive it. Time: an image model trained on millions of images for many epochs takes weeks on one device; with $W$ devices each processing a share of every batch, the step time falls if the communication the devices need can be kept small or hidden. Goyal et al. (2017), as the stochastic gradient descent article reports, trained ResNet-50 on ImageNet in one hour with batches of 8,192 images on 256 GPUs. Memory: training stores the parameters, their gradients, the optimizer’s state and the activations of the forward pass. With Adam and mixed precision, the model states alone take 16 bytes per parameter (below), so a 22-billion-parameter vision transformer needs about 355 GB before any activation, far beyond one accelerator. Zhao et al. (2023) give the same limit for plain data parallelism: it requires all parameters, gradients and optimizer states to fit on one GPU, and models above a billion parameters will likely run out of memory on a 40 GB device.

When to Use It

• Data parallelism is the default when the model fits on one device and training is too slow. Li et al. (2020) report that more than 60% of Facebook’s production GPU hours between 11 May and 5 June 2020 were spent in PyTorch’s distributed data-parallel package.
• Sharded data parallelism is the next step when the model states no longer fit, or when sharding frees memory for larger batches. It keeps the data-parallel programming model: each device still runs the whole forward and backward pass.
• Tensor and pipeline parallelism are for models whose layers or whose activations do not fit even when the training state is sharded, and for very large clusters, where they are combined with data parallelism.
• Gradient accumulation gives the gradient of a large batch on fewer devices than the batch needs, at the cost of time.
• Distributed training does not help a model that is fast enough and fits: communication then only adds cost. For a ViT-B-sized network the model states take 1.4 GB, and the reasons to distribute are speed and batch size, not parameter memory.

How It Works

Data parallelism and all-reduce

Each of $W$ processes (one per GPU, called a rank) holds a full copy of the model and reads a different part of every mini-batch. For a loss averaged over a batch $B$ split into equal shards $B_1, \dots, B_W$,

$$\nabla_\theta \frac{1}{|B|}\sum_{i \in B} \ell_i(\theta) = \frac{1}{W}\sum_{r=1}^{W} \nabla_\theta \frac{1}{|B_r|}\sum_{i \in B_r} \ell_i(\theta),$$

so if every rank computes the gradient of its own shard and the ranks then average their gradients, each holds the full-batch gradient and takes the same optimizer step. Li et al. (2020) state the condition this needs as mathematical equivalence: all replicas start from the same state and see the same gradients after every backward pass. PyTorch’s DistributedDataParallel (DDP) therefore broadcasts the model state from rank 0 at construction, broadcasts buffers such as batch-normalization running statistics from rank 0 again before each forward pass (§4.1), and averages gradients during the backward pass.

The averaging is an all-reduce: every rank contributes a tensor of the same size, and every rank receives the elementwise sum (or another reduction). Communication libraries implement it with ring- or tree-based algorithms rather than having every rank broadcast to every other (Li et al., 2020, §2.3). Because an all-reduce cannot complete until every rank joins, it is a synchronizing operation, and the slowest rank sets the pace.

DDP’s engineering is in hiding this communication. Collective operations are much more efficient on large tensors, so DDP packs gradients into buckets (25 MB by default in the version Li et al. describe, PyTorch 1.5) and launches an asynchronous all-reduce on each bucket as soon as all its gradients are ready, while the backward pass continues on earlier layers. Buckets are filled in the reverse order of the model’s parameters, an approximation of the order in which the backward pass produces gradients, and every rank must reduce the buckets in the same order or the results mismatch (Li et al., 2020, §3.2). With 32 GPUs, overlapping communication with the backward pass sped up iterations by 38.0% for ResNet-50 and 35.2% for BERT on the NCCL back end; the authors found the backward pass, which contains the communication, to be the most expensive step, and its communication share to grow with model size (§5.1).

The alternative of averaging parameters after independent local optimizer steps is not equivalent: replicas see different gradients, so optimizer state such as momentum diverges between them, and computation and communication cannot overlap (Li et al., 2020, §2.2).

Synchronous and asynchronous updates

Everything above is synchronous: all ranks wait for each other at every step. Early large-scale systems were also asynchronous. Google’s DistBelief, the predecessor of TensorFlow, kept parameters on separate parameter servers, and GoogLeNet was trained with it by asynchronous SGD (GoogLeNet paper). In asynchronous training workers send updates through point-to-point communication without waiting for each other (Li et al., 2020, §7), so a worker’s gradient may have been computed on parameters that others have since changed. PipeDream, in Li et al.’s description, “slightly sacrifices accuracy” by applying gradients from several concurrent passes that may not match the current parameters, which its authors found tolerable. Synchronous all-reduce training, which keeps the single-device mathematics, is what PyTorch’s DDP and FSDP implement.

Large batches

Data parallelism multiplies the batch: $W$ ranks with $b$ images each make a global batch of $Wb$. Keeping $b$ fixed as $W$ grows changes the optimization unless the learning rate changes with it. Goyal et al. (2017), as the stochastic gradient descent article reports, proposed the linear scaling rule, which multiplies the learning rate by $k$ when the batch grows by $k$, together with a gradual warmup over the first epochs; that article gives their results and the evidence on large-batch generalization, and the learning-rate schedule article covers warmup. Hyperparameters tuned on one device are therefore tied to the batch size, not to the number of devices.

Gradient accumulation

When the global batch does not fit, a device can process it as $k$ micro-batches in sequence and add the gradients before one optimizer step. Because PyTorch’s backward() adds into each parameter’s .grad, this is a loop with the loss of each micro-batch divided by $k$; the result is the gradient of the mean loss over the whole batch, by the same identity as above. In data-parallel training the micro-batch gradients need not be synchronized every time: DDP’s no_sync() context skips the all-reduce until the last micro-batch (Li et al., 2020, §3.2.4). Skipping synchronization is also used deliberately to reduce communication. With 256 GPUs, synchronizing ResNet-50 every 8 iterations sped up the average iteration by 38% on NCCL and 57% on Gloo; for a ResNet trained on MNIST at batch size 8 it slowed convergence negligibly, but at batch size 256 it raised the final training loss, which the authors explain by the larger accumulated gradient implicitly requiring a smaller learning rate (§5.3). FSDP offers accumulation with communication, which keeps gradients sharded, and without, which keeps unsharded gradients on each rank and trades memory for less communication (Zhao et al., 2023, §3.3.4).

Sharded data parallelism: ZeRO and FSDP

In data parallelism every rank stores the same parameters, gradients and optimizer states: the memory is redundant. ZeRO (Rajbhandari et al., 2020, not read here; described from DeepSpeed’s documentation and the PyTorch papers) partitions this state across the data-parallel ranks in three stages, which DeepSpeed’s tutorial defines for mixed-precision Adam:

• Stage 1 partitions the optimizer states (32-bit weights and Adam’s first and second moments), so each rank updates only its partition.
• Stage 2 also partitions the reduced 16-bit gradients, each rank keeping only those for its part of the optimizer states.
• Stage 3 also partitions the 16-bit parameters, which are gathered when the forward and backward passes need them.

With 2 bytes for the 16-bit parameters, 2 for the 16-bit gradients and 12 for the 32-bit weights and two moments, the model states take $16\Psi$ bytes for $\Psi$ parameters on every rank under data parallelism, and $16\Psi/N$ under stage 3 on $N$ ranks (our arithmetic from the tutorial’s description; it reproduces the tutorial’s example of a 1.5-billion-parameter GPT-2 whose Adam states take 18 GB, cut to 2.25 GB per device on eight ranks by stage 1). PyTorch’s ZeroRedundancyOptimizer shards only the optimizer states, which is stage 1 in DeepSpeed’s terms: each parameter belongs to one rank, whose local optimizer updates it, and the updated parameters are then broadcast to all ranks.

PyTorch’s Fully Sharded Data Parallel (FSDP) was motivated by ZeRO but redesigned around PyTorch’s internals (Zhao et al., 2023). It divides the model into units, for example one per transformer block, flattens each unit’s parameters into one tensor and shards it evenly over the ranks. Before a unit’s forward computation, the ranks all-gather its full parameters, compute, and free the gathered shards; the backward pass gathers them again and then reduce-scatters the gradients, so that each rank ends the step holding only its shard of parameters and gradients, and the optimizer states stay sharded throughout. Peak parameter memory is the sharded model plus the largest unsharded unit, $O(\sum_i \psi_i/F + \max_i \psi_i)$ for units of sizes $\psi_i$ sharded over $F$ ranks; finer units lower the peak but need more collective calls (§3.2.1). The sharding factor $F$ spans a range: $F = 1$ is plain data parallelism with all-reduce, $F = W$ is full sharding, and values in between give hybrid sharding, which shards within a group of ranks (for example the GPUs of one machine) and replicates across groups, cutting slow cross-machine traffic (§3.2.2). The gathers can be prefetched to overlap computation; for a 175-billion-parameter GPT model, backward prefetching gave a speedup of about 18% (§5.2). The paper describes PyTorch’s original FSDP, which shards flattened units; in PyTorch 2.14.1 the documentation recommends FSDP2 (fully_shard), which uses per-parameter sharding and still all-gathers the parameters of each group before its forward and backward passes and reduce-scatters their gradients.

Sharding is paid for in communication. The FSDP paper puts full sharding at 1.5 times the communication volume of DDP with a bandwidth-optimal ring algorithm (§3.2.1), and gives the cross-host traffic per GPU for a model of size $M$ as $2M(W-1)/W$ for full replication and $3M(W-1)/W$ for full sharding (§3.2.2). In our reading, these are a ring all-reduce (a reduce-scatter and an all-gather) against two all-gathers and a reduce-scatter, each moving $(W-1)/W$ of the model. On T5 models of 611 million and 2.28 billion parameters FSDP and DDP performed similarly; DDP ran out of memory above 2.28 billion, while FSDP trained the 11-billion-parameter model (§5.2). The DALL·E article describes an earlier vision-scale case: the 12-billion-parameter model’s parameters were sharded across the eight GPUs of each machine, with the next block’s shards prefetched during computation.

Tensor parallelism

Tensor parallelism splits individual layers. For a linear layer $\mathbf{y} = A\mathbf{x}$, each of $k$ devices can hold a block of rows of $A$ and compute its part of $\mathbf{y}$ after gathering all of $\mathbf{x}$, or a block of columns acting on its part of $\mathbf{x}$, producing partial outputs that must be summed across devices. Unlike FSDP, the parameters stay sharded during the computation, and the communicated tensors are activations. ViT-22B (Dehghani et al., 2023, §3) organized its TPU chips into a two-dimensional mesh of $t \times k$ devices, with data parallelism along one axis and this model parallelism along the other: each device of a model group received the same images, kept $1/k$ of the activations and computed $1/k$ of every linear layer’s output. The two sharding modes communicate $(k-1)n/k$ or $(k-1)m/k$ floats for an $m \times n$ matrix, so the authors used column sharding where the output is much smaller than the input (the MLP’s second layer, $n = 4m$) and row sharding elsewhere, and overlapped the matrix products with the communication. They also sharded large parameters along the data axis, gathering them asynchronously one layer ahead. ViT-22B reached a model FLOPs utilization of 54.9% on TPUv4, against 44.0% that the authors measured for the data-parallel-only ViT-e on the same hardware. PyTorch exposes the two linear-layer layouts as the experimental ColwiseParallel and RowwiseParallel styles, which can be composed to shard an MLP or an attention block.

Pipeline parallelism

Pipeline parallelism places consecutive groups of layers (stages) on different devices and passes activations forward and gradients backward between them. Run on one batch, only one stage would work at a time. The batch is therefore split into $m$ micro-batches that flow through the stages in turn. PyTorch’s ScheduleGPipe, named after GPipe (Huang et al., 2019, not read here), goes through the micro-batches in what its docstring calls a fill-drain manner: every stage runs all forward passes and then all backward passes, and the gradients of the micro-batches are accumulated before one synchronous update. Each stage idles while the pipeline fills and drains; with $p$ stages and equal stage times the idle share of the step, the bubble, is $(p-1)/(m+p-1)$ (our derivation, checked by the simulation below), so many micro-batches per stage are needed for high utilization. The 1F1B schedule reorders the same work: after a warm-up of forward passes, each stage alternates one forward and one backward pass, so a stage holds the activations of at most about $p$ micro-batches instead of all $m$; PyTorch’s Schedule1F1B starts the stage at index $s$ with $p - s$ forward passes and requires $m \ge p$. Pipeline stages communicate only activations and their gradients at stage boundaries (Zhao et al., 2023, §1), which suits slower links between machines, but they require the model to be split into stages with static shapes, and the schedule interacts with other parallelism: with FSDP’s default full sharding each stage would gather its parameters again for every micro-batch (Zhao et al., 2023, §7.1.1).

Activation checkpointing

For vision models at high resolution, activations rather than parameters can dominate memory. Activation checkpointing (gradient checkpointing) keeps only some activations in the forward pass and recomputes the rest in the backward pass. The backpropagation article gives its analysis: Chen et al. (2016) reduced feature-map memory to $O(\sqrt{n})$ for an $n$-layer chain at the price of one extra forward pass, about 30% more time in their experiments; the DenseNet article reports a checkpointed implementation. In PyTorch, torch.utils.checkpoint.checkpoint recomputes a wrapped function in the backward pass; its documentation recommends the non-reentrant variant (use_reentrant=False) and warns that if the function behaves differently when recomputed, for example through global state, the gradients may be silently wrong. By default it stashes and restores the random-number state, so that, in our reading, random operations such as dropout repeat identically in the recomputation. FSDP’s largest runs combined checkpointing with sharding and BF16 mixed precision (Zhao et al., 2023, §5.4).

Synchronized batch normalization

Batch normalization computes its statistics over the batch it sees. Under data parallelism each rank sees only its shard, so with a few images per GPU, as in detection and segmentation, the statistics are noisy, and the step is no longer the single-device step on the global batch. Synchronized batch normalization pools the statistics over the ranks. In PyTorch 2.14.1, SyncBatchNorm’s forward pass all-gathers each rank’s per-channel mean, inverse standard deviation and count, combines them into global statistics (so ranks with different numbers of images are weighted correctly), and its backward pass all-reduces the two per-channel sums that the input gradient needs; synchronization happens only in training mode, and a process group can restrict it to a subset of ranks. It supports DDP with one GPU per process, and when it synchronizes it accepts only accelerator inputs (CUDA, HPU, XPU or a custom back end), not CPU tensors, which is why the example below simulates it. The batch normalization article describes its use in the MegDet detector and its alternatives that do not depend on the batch, such as group normalization and layer normalization.

Communication costs

A step’s communication depends on the method, not on the batch: data parallelism moves about two model sizes of gradients per device per step, full sharding three, tensor parallelism moves activations inside every layer, and pipeline parallelism moves activations between stages. Li et al. (2020) measured communication as the main contributor to iteration time, and per-iteration latency growing as the number of GPUs grew: on 256 GPUs with NCCL, a ResNet-50 iteration took twice as long as local training, an effective speedup of 128 rather than 256, and on Gloo the slowdown was about three times for ResNet-50 and six times for BERT (§5.3). Network topology matters as much as volume: links within a machine are much faster than links between machines, which is why hybrid sharding keeps gathers inside a machine, why the FSDP authors keep tensor-parallel communication within a node (§7.1.2), and why the DALL·E authors compressed gradients for the cross-machine all-reduce. Collectives that move larger tensors per call use the bandwidth better, which explains DDP’s buckets and FSDP’s flattened units.

How to Apply It

The example checks the equivalences above without a cluster. One process plays four ranks: a small convolutional network with fixed random weights computes gradients on synthetic images, and the all-reduce is replaced by its result, an average. Part 1 compares the full-batch gradient with data-parallel and accumulated gradients, using PyTorch’s own DistributedSampler (with an explicit number of replicas and rank, so no process group is needed) to decide which images each rank reads, including a dataset that does not divide evenly. Part 2 repeats the comparison with batch normalization, once with per-shard statistics and once with statistics pooled over the shards as SyncBatchNorm pools them. Part 3 computes model-state memory per device for ViT-B/16 (counted from a model built on the meta device, which allocates no memory), DeepSpeed’s 1.5-billion-parameter example and ViT-22B at its reported 22,165 million parameters, the communication of replication against full sharding, and the idle fraction of GPipe and 1F1B pipeline schedules simulated with a forward time of 1 and a backward time of 2. The measure in parts 1 and 2 is the relative difference $\lVert \mathbf{g} - \mathbf{g}_{\text{ref}} \rVert / \lVert \mathbf{g}_{\text{ref}} \rVert$ over all parameter gradients; the gradients are in float64 except one float32 line. Memory counts the 2 + 2 + 12 bytes per parameter described above, excludes activations, and uses 1 GB = $10^9$ bytes.

The script ran with PyTorch 2.14.1 and NumPy 2.5.4 on a CPU with one thread in 1 to 2 seconds (approximate, measured while other agents shared the machine), and gave identical output on two runs.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DistributedSampler

torch.set_num_threads(1)
W = 4                                     # simulated number of processes (ranks)

def render(n, size=32, seed=0):
    """n grayscale images of a disk (0), square (1) or triangle (2) with noise."""
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[0:size, 0:size] + 0.5
    imgs, masks = np.zeros((n, 1, size, size)), np.zeros((n, size, size), bool)
    labels = np.arange(n) % 3
    for i in range(n):
        r = rng.uniform(5, 9)
        cx, cy = rng.uniform(r + 1, size - r - 1, 2)
        if labels[i] == 0:
            m = (xx - cx) ** 2 + (yy - cy) ** 2 <= r ** 2
        elif labels[i] == 1:
            m = (np.abs(xx - cx) <= r) & (np.abs(yy - cy) <= r)
        else:
            m = (yy <= cy + r) & (yy - (cy - r) >= 2 * np.abs(xx - cx))
        masks[i] = m
        imgs[i, 0] = rng.uniform(0.5, 1.0) * m + rng.normal(0, 0.1, (size, size))
    return torch.tensor(imgs), torch.tensor(labels), masks

X, Y, masks = render(32)
border = masks[:, 0].any() | masks[:, -1].any() | masks[:, :, 0].any() | masks[:, :, -1].any()
print("data: class counts", np.bincount(Y.numpy()).tolist(),
      f"| values {X.min().item():.2f}..{X.max().item():.2f} | shape touches border: {border}")

def make_net(bn, dtype=torch.float64):
    torch.manual_seed(0)                  # the same fixed random weights on every "rank"
    layers = [nn.Conv2d(1, 8, 3, padding=1)] + ([nn.BatchNorm2d(8)] if bn else [])
    layers += [nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(8, 16, 3, padding=1)]
    layers += ([nn.BatchNorm2d(16)] if bn else [])
    layers += [nn.ReLU(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(16, 3)]
    return nn.Sequential(*layers).to(dtype).train()

def flat_grad(net):
    return torch.cat([p.grad.flatten() for p in net.parameters()]).clone()

def grad(net, x, y):
    net.zero_grad()
    F.cross_entropy(net(x), y).backward()
    return flat_grad(net)

def rel(a, b):
    return ((a - b).norm() / b.norm()).item()

def shards(n, drop_last=False):           # the indices each rank would load
    ds = list(range(n))
    return [list(DistributedSampler(ds, num_replicas=W, rank=r, shuffle=False,
                                    drop_last=drop_last)) for r in range(W)]

# ---- Part 1: data parallelism and accumulation (no batch normalization) ----
net = make_net(bn=False)
g_full = grad(net, X, Y)
g_dp = torch.stack([grad(net, X[i], Y[i]) for i in shards(32)]).mean(0)   # "all-reduce"
print(f"\n1. data parallel, {W} shards of 8, mean of shard gradients: rel diff {rel(g_dp, g_full):.1e}")

net.zero_grad()
for i in shards(32):                       # micro-batches processed one after another
    (F.cross_entropy(net(X[i]), Y[i]) / W).backward()   # .grad accumulates
g_acc = flat_grad(net)
print(f"   accumulation, 4 micro-batches of 8:          rel diff {rel(g_acc, g_full):.1e}")
net.zero_grad()
for i in shards(32):
    F.cross_entropy(net(X[i]), Y[i]).backward()          # forgot to divide by 4
print(f"   accumulation without /4: gradient norm x {(flat_grad(net).norm() / g_full.norm()).item():.12f}")

g30 = grad(net, X[:30], Y[:30])
for drop_last in (False, True):
    idx = shards(30, drop_last)
    g = torch.stack([grad(net, X[i], Y[i]) for i in idx]).mean(0)
    used = sorted(sum(idx, []))
    line = f"   30 images, drop_last={drop_last!s:5}: shard sizes {[len(i) for i in idx]},"
    line += f" {len(used)} indices, {len(set(used))} distinct; vs 30-image gradient {rel(g, g30):.1e}"
    if drop_last:
        line += f", vs its {len(set(used))} images {rel(g, grad(net, X[used], Y[used])):.1e}"
    print(line)

net32 = make_net(bn=False, dtype=torch.float32)
X32 = X.float()
g32 = grad(net32, X32, Y)
g32_dp = torch.stack([grad(net32, X32[i], Y[i]) for i in shards(32)]).mean(0)
print(f"   float32, data parallel vs full batch:        rel diff {rel(g32_dp, g32):.1e}")

# ---- Part 2: batch normalization with per-shard and synchronized statistics ----
def bn_sync(xs, bn):
    """Normalize every shard with statistics pooled over all shards, as SyncBatchNorm does."""
    n = sum(x.shape[0] * x.shape[2] * x.shape[3] for x in xs)
    s1 = sum(x.sum((0, 2, 3)) for x in xs)          # all-reduce of per-shard sums
    s2 = sum((x * x).sum((0, 2, 3)) for x in xs)    # and of sums of squares
    mean, var = s1 / n, s2 / n - (s1 / n) ** 2      # biased variance, as in training
    return [bn.weight[:, None, None] * (x - mean[:, None, None])
            / torch.sqrt(var[:, None, None] + bn.eps) + bn.bias[:, None, None] for x in xs]

def forward_shards(net, xs, sync):
    for layer in net:
        if isinstance(layer, nn.BatchNorm2d):
            xs = bn_sync(xs, layer) if sync else [
                F.batch_norm(x, None, None, layer.weight, layer.bias, True, 0.0, layer.eps)
                for x in xs]
        else:
            xs = [layer(x) for x in xs]
    return xs

netbn = make_net(bn=True)
g_full = grad(netbn, X, Y)
idx = shards(32)
h = netbn[0](X)
check = (torch.cat(bn_sync([h[i] for i in idx], netbn[1]))
         - F.batch_norm(h[sum(idx, [])], None, None, netbn[1].weight, netbn[1].bias,
                        True, 0.0, netbn[1].eps)).abs().max().item()
print(f"\n2. check: pooled-statistics BN vs BatchNorm on the whole batch, max abs diff {check:.1e}")
for sync in (False, True):
    netbn.zero_grad()
    outs = forward_shards(netbn, [X[i] for i in idx], sync)
    # each rank's mean loss; the all-reduce averages the ranks' gradients
    sum(F.cross_entropy(o, Y[i]) for o, i in zip(outs, idx)).div(W).backward()
    name = "synchronized statistics" if sync else "per-shard statistics   "
    print(f"   BN with {name}: rel diff {rel(flat_grad(netbn), g_full):.1e}")

# ---- Part 3: memory, communication and pipeline arithmetic ----
D, L, MLP, P, C, T = 768, 12, 3072, 16, 1000, 197      # ViT-B/16, 224 x 224 input

class ViT(nn.Module):
    def __init__(self):
        super().__init__()
        self.patch = nn.Conv2d(3, D, P, stride=P)
        self.cls = nn.Parameter(torch.zeros(1, 1, D))
        self.pos = nn.Parameter(torch.zeros(1, T, D))
        self.blocks = nn.ModuleList(nn.TransformerEncoderLayer(D, 12, MLP, norm_first=True,
                                                               batch_first=True) for _ in range(L))
        self.norm, self.head = nn.LayerNorm(D), nn.Linear(D, C)

with torch.device("meta"):                 # shapes only, no memory allocated
    vit = ViT()
n_vit = sum(p.numel() for p in vit.parameters())
block = 2 * 2 * D + (3 * D * D + 3 * D) + (D * D + D) + (D * MLP + MLP) + (MLP * D + D)
closed = (3 * P * P * D + D) + D + T * D + L * block + 2 * D + (D * C + C)
print(f"\n3. ViT-B/16 parameters: {n_vit:,} counted, {closed:,} closed form")

def per_device_gb(psi, n):
    """Model states per device in GB: 16-bit params (2 B), 16-bit grads (2 B),
    fp32 weights + Adam moments (12 B); ZeRO stages shard more of them over n devices."""
    p, g, o = 2 * psi, 2 * psi, 12 * psi
    return [x / 1e9 for x in (p + g + o, p + g + o / n, p + (g + o) / n, (p + g + o) / n)]

print("   model             N   DDP    stage1  stage2  stage3   (GB per device, activations excluded)")
for name, psi in (("ViT-B/16", n_vit), ("1.5B (DeepSpeed)", 1.5e9), ("ViT-22B", 22_165e6)):
    for n in (1, 8, 64):
        print(f"   {name:16} {n:3} " + " ".join(f"{v:7.2f}" for v in per_device_gb(psi, n)))
print(f"   DeepSpeed's example: Adam states {12 * 1.5e9 / 1e9:.2f} GB, over 8 ranks"
      f" {12 * 1.5e9 / 8 / 1e9:.2f} GB")
for w in (8, 64):
    rep, full = 2 * (w - 1) / w, 3 * (w - 1) / w
    print(f"   traffic per device per step, W={w:2}: replicated {rep:.3f}, fully sharded {full:.3f}"
          f" model sizes (ratio {full / rep:.2f})")

def simulate(p, m, order, tf=1.0, tb=2.0):
    """Run per-stage op lists [('F'|'B', microbatch)] respecting data dependencies."""
    done, free, pos = {}, [0.0] * p, [0] * p
    held, peak = 0, 0
    while any(pos[s] < len(order[s]) for s in range(p)):
        for s in range(p):
            if pos[s] == len(order[s]):
                continue
            kind, j = order[s][pos[s]]
            dep = (('F', s - 1, j) if s > 0 else None) if kind == 'F' else \
                  (('B', s + 1, j) if s < p - 1 else ('F', s, j))
            if dep is not None and dep not in done:
                continue
            start = max(free[s], done[dep] if dep else 0.0)
            done[(kind, s, j)] = free[s] = start + (tf if kind == 'F' else tb)
            pos[s] += 1
            if s == 0:
                held += 1 if kind == 'F' else -1
                peak = max(peak, held)
    makespan = max(free)
    return 1 - m * (tf + tb) / makespan, peak

def gpipe(p, m):
    return [[('F', j) for j in range(m)] + [('B', j) for j in range(m)] for _ in range(p)]

def one_f_one_b(p, m):
    order = []
    for s in range(p):
        k = min(m, p - s)                    # warm-up forwards, as in Schedule1F1B
        ops, f, b = [('F', j) for j in range(k)], k, 0
        while f < m:
            ops += [('B', b), ('F', f)]
            b, f = b + 1, f + 1
        ops += [('B', j) for j in range(b, m)]
        order.append(ops)
    return order

p = 4
print(f"\n   pipeline, {p} stages: m, idle fraction GPipe / 1F1B / (p-1)/(m+p-1),"
      " micro-batches held by stage 0 GPipe / 1F1B")
for m in (1, 2, 4, 8, 16, 32):
    (ig, hg), (i1, h1) = simulate(p, m, gpipe(p, m)), simulate(p, m, one_f_one_b(p, m))
    print(f"   m={m:2}  {ig:.4f} {i1:.4f} {(p - 1) / (m + p - 1):.4f}   {hg:2d} {h1:2d}")

data: class counts [11, 11, 10] | values -0.45..1.26 | shape touches border: False

1. data parallel, 4 shards of 8, mean of shard gradients: rel diff 6.7e-16
   accumulation, 4 micro-batches of 8:          rel diff 6.7e-16
   accumulation without /4: gradient norm x 4.000000000000
   30 images, drop_last=False: shard sizes [8, 8, 8, 8], 32 indices, 30 distinct; vs 30-image gradient 2.9e-01
   30 images, drop_last=True : shard sizes [7, 7, 7, 7], 28 indices, 28 distinct; vs 30-image gradient 3.8e-01, vs its 28 images 1.2e-15
   float32, data parallel vs full batch:        rel diff 5.2e-07

2. check: pooled-statistics BN vs BatchNorm on the whole batch, max abs diff 3.3e-14
   BN with per-shard statistics   : rel diff 5.6e-02
   BN with synchronized statistics: rel diff 1.0e-14

3. ViT-B/16 parameters: 86,567,656 counted, 86,567,656 closed form
   model             N   DDP    stage1  stage2  stage3   (GB per device, activations excluded)
   ViT-B/16           1    1.39    1.39    1.39    1.39
   ViT-B/16           8    1.39    0.48    0.32    0.17
   ViT-B/16          64    1.39    0.36    0.19    0.02
   1.5B (DeepSpeed)   1   24.00   24.00   24.00   24.00
   1.5B (DeepSpeed)   8   24.00    8.25    5.62    3.00
   1.5B (DeepSpeed)  64   24.00    6.28    3.33    0.38
   ViT-22B            1  354.64  354.64  354.64  354.64
   ViT-22B            8  354.64  121.91   83.12   44.33
   ViT-22B           64  354.64   92.82   49.18    5.54
   DeepSpeed's example: Adam states 18.00 GB, over 8 ranks 2.25 GB
   traffic per device per step, W= 8: replicated 1.750, fully sharded 2.625 model sizes (ratio 1.50)
   traffic per device per step, W=64: replicated 1.969, fully sharded 2.953 model sizes (ratio 1.50)

   pipeline, 4 stages: m, idle fraction GPipe / 1F1B / (p-1)/(m+p-1), micro-batches held by stage 0 GPipe / 1F1B
   m= 1  0.7500 0.7500 0.7500    1  1
   m= 2  0.6000 0.6000 0.6000    2  2
   m= 4  0.4286 0.4286 0.4286    4  4
   m= 8  0.2727 0.2727 0.2727    8  4
   m=16  0.1579 0.1579 0.1579   16  4
   m=32  0.0857 0.0857 0.0857   32  4

Equivalence. Data parallelism and accumulation reproduce the full-batch gradient to float64 round-off, because a mean over equal shards of means is the mean; forgetting to divide the accumulated loss by the number of micro-batches multiplies the gradient by exactly that number, which with SGD is the same as multiplying the learning rate. In float32 the same computation differs from the full batch in the seventh digit: the sums are taken in a different order, and floating-point addition is not associative, so distributed and single-device runs agree mathematically but not bit for bit, as PyTorch’s numerical-accuracy notes warn.

Uneven data. Thirty images do not split into four equal shards. With the default drop_last=False, DistributedSampler pads the index list by repeating its first two indices, so two images count twice and the averaged gradient is not the gradient of the dataset; with drop_last=True it drops two images, and the average is exactly the gradient of the 28 that remain. Either way the step differs from the single-device step on all 30 images, here by 29% and 38% of the gradient’s norm, a large share because the batch is tiny. For training the difference is a slightly different epoch; for evaluation, padding means that two images are scored twice, so validation metrics computed this way should be corrected or computed on all ranks’ unpadded data.

Batch normalization. With each shard normalized by its own statistics, the averaged gradient differs from the full-batch gradient by about 6%: each shard’s normalization, and therefore its loss, depends on the other images in that shard only. Pooling the per-shard sums and sums of squares, the quantities an all-reduce or all-gather can combine, restores the full-batch gradient to round-off; the separate check confirms that the pooled normalization equals batch normalization on the whole batch. The difference here comes from shards of eight images; with fewer images per device the statistics get noisier still.

Memory. For ViT-B/16, whose 86.6 million parameters the meta-device model and the closed form count identically, the model states take 1.39 GB per device, so sharding saves little that matters; activations, which this table excludes, decide whether a batch fits. The ViT-22B rows show why sharding exists: 355 GB per device under data parallelism, still 122 GB with stage 1 and 83 GB with stage 2 on eight devices, and 44 GB only when the parameters are partitioned too. Stage 1 gives most of the saving at small $N$, since the optimizer states are 12 of the 16 bytes, but the 4 bytes of replicated parameters and gradients then set a floor that only stage 3 removes. The 1.5-billion-parameter rows reproduce DeepSpeed’s 18 GB of Adam states and 2.25 GB per device after partitioning. The traffic lines give the price of full sharding: 1.5 times the volume of an all-reduce, at any number of devices.

Pipelines. The simulated idle fraction equals $(p-1)/(m+p-1)$ for both schedules: 1F1B does not shrink the bubble, it only reorders the work. What it changes is memory: stage 0 holds the activations of all $m$ micro-batches under GPipe but never more than $p = 4$ under 1F1B, so more micro-batches, the remedy for the bubble, no longer cost activation memory. With 4 stages, 32 micro-batches still leave each device idle 8.6% of the step.

Variants

• 3D parallelism combines data, tensor and pipeline parallelism, as Megatron-LM’s training of transformers on large clusters does (cited by Zhao et al., 2023); PyTorch builds such combinations on device meshes, for example FSDP on one mesh dimension and tensor parallelism on another, with tensor-parallel communication kept within a node (§7.1.2).
• Hybrid sharding (above) shards within groups and replicates across them, between DDP and full sharding in both memory and communication.
• Offloading moves optimizer states or parameters to CPU memory or NVMe storage; DeepSpeed’s ZeRO-3 includes such an engine (ZeRO-Infinity), at the cost of host–device transfers.
• Gradient compression sends fewer bits per gradient. Li et al. (2020) list it as a direction for DDP, and the DALL·E authors compressed gradients with PowerSGD, as that article describes.
• Sequence and context parallelism split the token dimension of transformers to reduce activation memory, used together with tensor parallelism (Zhao et al., 2023, §6).

Best Practices

• Scale the learning rate with the global batch, not the device count, with warmup for large batches, and retune when the global batch changes (learning-rate schedule article).
• Seed and synchronize everything that must match. DDP broadcasts parameters and buffers from rank 0; DistributedSampler takes one seed for all ranks and needs set_epoch before each epoch, or every epoch repeats the same order.
• Choose the least parallelism that fits. Plain DDP first; sharding when memory runs out; hybrid sharding before full sharding on slow inter-machine networks; tensor and pipeline parallelism for models that need them. Zhao et al. report that FSDP and DDP performed similarly on models DDP could hold.
• Use the fast collective back end (NCCL on NVIDIA GPUs), and measure bucket sizes and unit sizes rather than guessing: Li et al. found the best bucket size to depend on the model and back end.
• Combine memory techniques: mixed-precision training, activation checkpointing and sharding address different parts of memory and were used together in FSDP’s largest experiments.
• Use synchronized batch normalization, or a normalization that does not depend on the batch, when each device sees only a few images.

Pitfalls

• Unscaled accumulation. Summing micro-batch losses without dividing multiplies the gradient, as the example shows; pipeline schedules in PyTorch scale gradients by the number of micro-batches by default (scale_grads=True), which must match whether the loss function averages or sums.
• Padded or dropped samples in distributed evaluation change the metric (example above).
• Per-device batch normalization with small shards gives different gradients from the global batch and running statistics taken from rank 0’s shards only.
• Divergent replicas. Anything that differs between ranks outside the synchronized gradients, such as an unseeded random choice that changes the computation graph or a parameter used on some ranks only, breaks the equivalence or hangs the all-reduce; DDP has a find_unused_parameters mode for models that use only part of their parameters per step.
• Sharding breaks some optimizers. The FSDP of the paper shards flattened parameters without respecting parameter boundaries, so optimizer computations that need a whole parameter, such as its norm, or global state are invalid without extra work (Zhao et al., 2023, §7.2.1); its mixed precision needs a sharded gradient scaler.
• Reproducibility. PyTorch does not guarantee bitwise-identical results across releases, platforms, or CPU and GPU, and reduction order differs with the number of devices, as the float32 line shows. torch.use_deterministic_algorithms(True) removes some nondeterministic kernels, but a run on 8 GPUs and a run on 64 will still differ in the last digits, and those differences can grow over a long training run (our reading). Report the number of devices, the global batch and the per-device batch together.
• Asynchrony and staleness. Asynchronous and pipelined schemes that apply gradients computed on older parameters change the optimization; Li et al. note that PipeDream trades a little accuracy for speed this way.

Evidence and Impact

Synchronous data parallelism made large-batch training of image models routine: Goyal et al.’s one-hour ImageNet training (the stochastic gradient descent article reports its accuracy against batch size) and MegDet’s 256-image detection batches with synchronized statistics (the batch normalization article) both rest on it. Li et al. (2020) reported near-linear scalability for DDP on 256 GPUs when bucket sizes and synchronization frequency are configured appropriately, while their measurements also show per-iteration latency rising with the number of GPUs, communication taking more than half of the backward pass, and a slower back end (Gloo) losing most of the gain.

Sharding moved the memory limit. In Zhao et al. (2023), FSDP trained a 175-billion-parameter GPT model on 128 to 512 A100 GPUs at more than 173 and 186 TFLOPS per GPU (batch sizes 1 and 2), about 55% and 60% of the GPU’s BF16 peak, with near-linear scaling in TFLOPS; T5-11B lost 7% of per-GPU throughput between 8 and 512 GPUs, which the authors take as a sign that communication begins to outweigh computation at scale. Its study of communication throttling included an 8-billion-parameter vision transformer, DeepViT, for which throttling added about 5% overhead rather than helping (§5.3). In vision, ViT-22B (Dehghani et al., 2023) combined data parallelism, model parallelism of its linear layers with overlapped communication, and parameter sharding to train a 22-billion-parameter vision transformer on about 4 billion JFT images with batches of 65k; such models depend on these methods, since the model states alone exceed any single accelerator’s memory.