Distributed training

At some point the model, its gradients and its optimiser state stop fitting on one device, and training becomes a distributed systems problem. The mistake that costs the most is reaching for the most sophisticated answer first. There is an order to try things in, each step is more complicated and slower per device than the last, and most people never need to go past the second one. This section gives that order, the three ways of splitting work and what each costs, and the failure modes that only appear once a run lasts days.

Why one device runs out#

Training needs far more memory than inference, and the weights are the smallest part of it. Every parameter carries a gradient, and an optimiser like Adam carries two more numbers per parameter plus a full-precision master copy. Mixed-precision Adam works out at sixteen bytes per parameter, so a seven-billion-parameter model needs about 104 GiB before storing a single activation — more than an 80 GiB card. The arithmetic is worked in full in the numbers section of this topic; the consequence is that the parameter count alone never tells you whether something fits.

what you have                              what to do
fits on one GPU                            one GPU. stop here.
fits, but the batch is too small           gradient accumulation
model fits, optimiser state does not       ZeRO stage 1 or 2 / FSDP
model does not fit                         ZeRO stage 3 / FSDP full shard
one layer does not fit                     tensor parallelism
too many layers for the interconnect       pipeline parallelism

Three ways to split#

The three parallelisms split different things, and the distinction that matters is what each one has to send across the wire and how often.

                  what is split        what crosses the wire     when
data parallel     the batch            gradients, once            per step
tensor parallel   each layer's maths   activations, twice         per layer
pipeline parallel the layers           activations at boundaries  per stage

Data parallelism gives every device a full copy of the model and a different slice of the batch, then averages the gradients. It is the simplest and the only one that does not change the model code. Its limit is memory: every device holds everything, so it cannot help with a model that does not fit. Sharded data parallelism — FSDP, or ZeRO — keeps the same programming model and splits the optimiser state, gradients and optionally the weights across devices, gathering each piece back only while it is needed. That is usually the whole answer, and it is covered in the next section.

Tensor parallelism splits the arithmetic inside a layer: each device holds part of a weight matrix and computes part of the result, and the parts are combined before the layer finishes. It is what you use when a single layer is too large for one device. Because it communicates inside every layer rather than once per step, it is almost always confined to the devices inside one machine, on NVLink. Pipeline parallelism gives each device a contiguous block of layers and passes activations along the chain, which communicates little but introduces the bubble: while device one works on the first micro-batch, the others have nothing to do. Splitting the batch into many micro-batches shrinks the bubble without removing it.

Large training runs combine all three, which is where the phrase 3D parallelism comes from: tensor parallel within a machine, pipeline parallel across a few machines, data parallel across the rest. You do not need to configure that to be hired, but you should be able to say why each sits at the level it does — fastest links for the chattiest parallelism.

Communication is the limit#

The reason the ordering above is not arbitrary is that every step past the first buys memory with bandwidth. Whether that trade is affordable is arithmetic, not judgement: a seven-billion-parameter gradient sync is about 27 ms over NVLink and about 19.6 seconds over 10 Gb Ethernet, which is 6 per cent of a training step against 98 per cent of it. Same model, same eight GPUs, and the interconnect alone decides whether the cluster is useful. Doing that calculation before provisioning is the single highest-value habit in this topic, and the next section but one works it in full.

Long runs fail#

A run measured in days meets failures a run measured in minutes never does, and none of them are exotic: a node dies, a network partition hangs a collective, a disk fills with checkpoints, a GPU develops memory errors. The distributed part makes each one worse, because collectives are synchronous — one slow or dead device stalls all of them, and the symptom is a job that is not crashed and not progressing.

• Checkpoint on time, not on epochs, and keep the optimiser state: a weights-only checkpoint cannot resume Adam correctly.
• Make resuming the normal path and test it on day one, including the data loader position, or a resumed run silently retrains on the same examples.
• Log throughput per step, not just loss. A straggler shows up as throughput decay long before anything fails.
• Set a timeout on collectives. The default in some stacks is half an hour of silence before anything is reported.
• Keep the last few checkpoints rather than only the newest, because a loss spike is usually discovered after the checkpoint that contains it.
• Record the seed, the data order and the exact configuration with the checkpoint. A run you cannot reproduce is a run you cannot debug.

The loss spike deserves its own note, because it is the characteristic failure of a long run and the response is standard: when the loss jumps and does not recover, roll back to the last good checkpoint, skip or reshuffle the batches around the spike, and lower the learning rate for a while. That sequence is what large training reports describe, and being able to say it is what the question is testing.

A model can fit on a GPU and still be untrainable on it, because training also needs gradients and optimiser state, which are three to six times larger than the weights. Sharded data parallelism fixes that without changing how you write the model: each device owns a slice of the state rather than a copy, and the pieces are gathered back only for the moment they are needed. This section covers what each stage shards, what it costs in communication, and activation checkpointing, which is the other half of making something fit and is almost always worth turning on.

Replication is the waste#

Plain data parallelism puts a full copy of the weights, gradients and optimiser state on every device. Eight devices therefore hold eight identical copies of the optimiser state, which for mixed-precision Adam is twelve of the sixteen bytes per parameter. Nothing about that is necessary: each device only ever updates the whole model because it happens to have the whole model. Give each one a slice to own and the memory divides by the device count.

The three stages#

stage   shards                          per GPU, 7B on 8 devices   extra traffic
none    nothing (plain DDP)                      104.3 GiB          all-reduce
1       optimiser state                           35.9 GiB          same
2       + gradients                               24.4 GiB          same (reduce-scatter)
3       + parameters                              13.0 GiB          + all-gather, twice

The stages are ordered by how cheap the sharding is, not by how much it saves. Optimiser state goes first because it is the largest component and is touched exactly once per step, at the update — so each device can hold its slice, update it, and nobody notices. Gradients go second, and this one is genuinely free: the all-reduce that data parallelism already performs is a reduce-scatter followed by an all-gather, and once gradients are sharded the all-gather half is simply not needed. Parameters go last because they are needed by every device on every forward and backward pass, so sharding them means gathering them back twice per step.

Activation checkpointing#

Sharding handles everything that scales with parameter count. Activations do not: they scale with batch size, sequence length and depth, and on a long-context model they can exceed everything else combined. The backward pass needs the tensors the forward pass produced, and a transformer block saves roughly sixteen hidden-sized tensors per token to make its backward pass cheap.

GiB = 1024 ** 3
layers, hidden, seq, batch, dtype = 32, 4096, 4096, 1, 2    # 7B-shaped, bf16

per_layer = batch * seq * hidden * dtype
full = layers * per_layer * 16 / GiB          # ~16 saved tensors per block
checkpointed = layers * per_layer / GiB       # only each block's input

print(f"nothing checkpointed : {full:5.1f} GiB")
print(f"blocks recomputed    : {checkpointed:5.1f} GiB   ({full/checkpointed:.0f}x less)")
print(f"compute cost         : {4/3 - 1:.0%} more FLOPs per step")

# nothing checkpointed :  16.0 GiB
# blocks recomputed    :   1.0 GiB   (16x less)
# compute cost         :  33% more FLOPs per step

Put that beside the sharding table and the decision becomes obvious. At stage 2 the per-device total goes from 40.4 GiB to 25.4 GiB, and at stage 3 from 29.0 to 14.0 — and the batch size is one, so on a realistic batch the activation term is the one that was going to break you. Thirty-three per cent more compute to make the job possible at all is a trade almost everyone takes.

The refinement worth knowing is that it does not have to be all or nothing. Checkpointing every block is the default and the simplest; checkpointing every second block recovers about half the saving for half the recompute. If you are close to fitting, that is the dial, and it is usually a single argument.

Choosing#

In practice the decision is short, and the mistake is starting further down it than necessary.

• If it fits on one device, use one device. Distributed training is slower per device, always.
• If it fits with a full replica per device, use plain data parallelism and nothing else.
• If the optimiser state is what does not fit, stage 2 with activation checkpointing. This is where most models land.
• If the parameters themselves do not fit, stage 3, and measure the throughput cost rather than assuming it.
• If a single layer does not fit on a device, sharding will not save you and you need tensor parallelism.

Data parallelism is the simplest way to use more than one GPU: every device gets a complete copy of the model and a different slice of each batch, computes gradients on its slice, and the gradients are averaged so all the copies stay identical. PyTorch's implementation is DistributedDataParallel, DDP for short. Turning a single-GPU script into a multi-GPU one with it is about twenty lines, and four of them are the ones people get wrong. This section is the whole conversion: how the processes are started, what wraps the model, how the data is split so two devices do not train on the same examples, how to save without writing the same file twice, how to accumulate gradients without paying for a sync you do not need, and how to tell whether any of it actually made training faster. Every code block here was run on this machine with two processes on the gloo backend, which is the same API you use on GPUs with nccl.

Processes, ranks and torchrun#

DDP is not threads. Each GPU gets its own operating-system process running the same script, and they find each other through a rendezvous that torchrun sets up. Rank is the global index of a process, local rank is its index within one machine and is what you use to pick a device, and world size is the total. You do not set these yourself; torchrun puts them in the environment and the process group reads them.

torchrun --nproc_per_node=4 train.py          # 4 processes on this machine
torchrun --nnodes=2 --nproc_per_node=8 \
         --rdzv_backend=c10d --rdzv_endpoint=host:29500 train.py

rank        0..world_size-1     global index of this process
local_rank  0..gpus_per_node-1  which GPU on this machine
world_size  total processes     usually GPUs in total

The conversion, in full#

import os, torch, torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset, DistributedSampler

def main():
    dist.init_process_group("gloo")          # "nccl" when there are GPUs
    rank, world = dist.get_rank(), dist.get_world_size()
    # On GPUs: torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) before anything else.

    data = TensorDataset(torch.randn(64, 10), torch.randn(64, 1))
    sampler = DistributedSampler(data, num_replicas=world, rank=rank, shuffle=True)
    loader = DataLoader(data, batch_size=8, sampler=sampler)   # batch_size is PER PROCESS

    model = DDP(nn.Linear(10, 1))            # wrap AFTER moving to the device
    opt = torch.optim.SGD(model.parameters(), lr=0.01)

    for epoch in range(2):
        sampler.set_epoch(epoch)             # without this, every epoch has the same order
        seen = 0
        for x, y in loader:
            opt.zero_grad()
            nn.functional.mse_loss(model(x), y).backward()   # gradients sync here
            opt.step()
            seen += len(x)
        if rank == 0:
            print(f"epoch {epoch}: rank 0 saw {seen} of {len(data)} samples")

    if rank == 0:                            # save from one process only
        torch.save(model.module.state_dict(), "model.pt")     # .module, not the wrapper
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

# torchrun --nproc_per_node=2 train.py
# epoch 0: rank 0 saw 32 of 64 samples
# epoch 1: rank 0 saw 32 of 64 samples

Four of those lines are the ones that go wrong, and none of them fail loudly. Forgetting the DistributedSampler means every process iterates the whole dataset, so an epoch trains on everything twice and the results look mysteriously worse than the single-GPU run. Forgetting set_epoch means the shuffle is identical every epoch, which quietly removes shuffling. Saving state_dict from the wrapper rather than from .module writes keys prefixed with module., which then fail to load into the plain model at inference. And saving from every rank has four processes writing one path at once.

Gradient accumulation without extra syncs#

DDP synchronises gradients during every backward call. When you accumulate over several micro-batches, all but the last of those syncs are wasted — the gradients are not being used yet. no_sync suppresses them, so a four-step accumulation costs one all-reduce instead of four.

from contextlib import nullcontext

accum = 4
for i, (x, y) in enumerate(batches):
    last = (i == accum - 1)
    ctx = nullcontext() if last else model.no_sync()   # sync only on the last one
    with ctx:
        (nn.functional.mse_loss(model(x), y) / accum).backward()
opt.step()
opt.zero_grad()

# accumulated 4 micro-batches with 1 sync instead of 4

Measuring whether it scaled#

Adding GPUs is an assumption until it is measured, and the number to report is scaling efficiency: throughput on N devices divided by N times the throughput on one. Measure samples per second over a steady-state window, skipping the first few steps so compilation and warmup do not count.

devices   samples/s   ideal   efficiency
  1          412        412      100%
  2          790        824       96%
  4         1510       1648       92%
  8         2560       3296       78%      <- look here

78% on eight means roughly two of the eight are doing nothing useful.
Usual causes, in the order worth checking:
  the data loader cannot feed them      (workers, prefetch, decode cost)
  the batch per device is too small     (GPU idle between kernels)
  the interconnect                      (see the numbers section)
  one straggler                         (a slower card, or a busy host)

The habit worth forming is running this table once when you set the job up and once more after any change to the data pipeline. Efficiency falling from 92 to 60 per cent after someone adds a heavier augmentation is a common and entirely invisible regression — the loss curve looks the same, and the bill doubles.

FSDP is the API for the sharding described in the previous section: each device owns a slice of the parameters, gradients and optimiser state, and the pieces are gathered back for the moment they are used. The mental model that makes it behave predictably is that sharding happens per wrapped unit — you choose the units, and that choice sets both the memory saving and the communication cost. This section covers choosing them, the mixed-precision settings, combining it with activation checkpointing, saving a model whose weights are spread across eight machines, and what to look at when it is slower than you expected.

Sharding per block#

Wrap the repeating unit of the model — a transformer block — rather than the whole thing. Wrapping only the root gathers every parameter at once, which defeats the purpose, since peak memory then includes the entire model. Wrapping each block means only one block's parameters are gathered at a time, and the next block's gather can overlap with the current block's compute.

import torch, torch.nn as nn
from torch.distributed.fsdp import fully_shard

for block in model.transformer_blocks:   # the repeating unit, innermost first
    fully_shard(block)
fully_shard(model)                       # then the root, for anything left over

# Peak memory now holds one block's worth of gathered parameters, not the model.

The one tuning knob worth knowing is reshard_after_forward. Left on, parameters are freed after the forward pass and gathered again for the backward, which is the lowest-memory setting and costs one extra all-gather. Turned off, they stay resident between the two, which is faster and uses more memory. The usual configuration is to reshard everywhere except the last block, which is gathered again almost immediately.

Mixed precision settings#

from torch.distributed.fsdp import MixedPrecisionPolicy

policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,     # what the forward and backward compute in
    reduce_dtype=torch.float32,     # what gradients are reduced in
)
for block in model.transformer_blocks:
    fully_shard(block, mp_policy=policy)

That split is the answer to a common interview question about why bf16 is safe for training when fp16 needs loss scaling. bf16 has the same exponent range as fp32, so it does not underflow the way fp16 does and needs no scaler; what it has instead is fewer mantissa bits, so the place to be careful is accumulation — which is exactly what reduce_dtype controls.

Activation checkpointing alongside#

Sharding and checkpointing solve different problems and are almost always used together: sharding handles everything that scales with parameter count, checkpointing handles activations, which scale with batch and sequence length. For the 7B-shaped model in the previous section that was 16 GiB of activations reduced to 1 GiB, for 33 per cent more compute.

from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper, CheckpointImpl,
)

for i, block in enumerate(model.transformer_blocks):
    if i % 2 == 0:                       # every block, or every second one
        model.transformer_blocks[i] = checkpoint_wrapper(
            block, checkpoint_impl=CheckpointImpl.NO_REENTRANT,
        )
# Apply checkpointing BEFORE fully_shard, so FSDP shards the wrapped module.

Checkpoints for a sharded model#

This is where people lose a training run. Calling state_dict on a sharded model gives you that rank's shard, so saving it from rank zero saves an eighth of the model, and the failure is only discovered at load time. There are two correct approaches and the right one depends on what the checkpoint is for.

# For resuming: every rank writes its own shard, in parallel. Fast, and the
# only sane option at scale -- no rank ever holds the whole model.
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.save({"model": model_sd, "optim": optim_sd}, checkpoint_id="step_1000")

# Resuming: load into the same shape, then push back into the live objects.
state = {"model": model_sd, "optim": optim_sd}
dcp.load(state, checkpoint_id="step_1000")
set_state_dict(model, optimizer, model_state_dict=state["model"],
               optim_state_dict=state["optim"])

For the other purpose — a single file to hand to inference — you need the full unsharded weights gathered onto one rank, which requires that rank to have enough memory for the whole model. Do that once at the end of training rather than every few hundred steps, and keep it separate from the resume checkpoints, which are the ones that must be fast.

When it is slow#

FSDP can run correctly and slowly, and the symptom is always the same — the memory fits, the loss falls, and the throughput is a fraction of what the hardware should give. The causes are a short list, in the order worth checking.

symptom / cause                          what to do
wrapped only the root                    wrap each repeating block
wrapped too finely, every Linear         wrap blocks, not layers
stage 3 on a slow interconnect           try stage 2, or accumulate gradients
reshard_after_forward on everywhere      leave the last block resident
gathering full state_dict every save     use dcp.save for resume checkpoints
CPU offload left on                      turn it off unless you truly need it

Measure it the way the previous section measures DDP: samples per second on one device against samples per second on all of them. Sharding that has taken throughput below what a smaller batch on plain data parallelism would have given is a configuration to abandon, not to tune further — and knowing that boundary is most of what the FSDP question in an interview is really asking.

Four numbers decide a distributed training design, and all four are arithmetic you can do before requesting a single GPU: how much memory each device needs, how long a gradient sync takes, what fraction of the step that sync is, and how much of the hardware you are actually using. This section works all four for a concrete model — seven billion parameters on eight GPUs — because the numbers, not the vocabulary, are what an interviewer is probing and what stops you renting a cluster that cannot help.

The three collectives#

Almost all distributed training is built from three operations, and knowing which one a framework is doing tells you what it costs. On NVIDIA hardware they are implemented by NCCL, the library PyTorch calls when you select the nccl backend, and it is NCCL that decides the ring or tree the data actually travels. All-reduce combines a value across every device and leaves every device with the combined result — this is the gradient sync in plain data parallelism. Reduce-scatter combines and then gives each device only its slice. All-gather is the reverse: each device contributes its slice and everyone ends up with the whole. The identity worth remembering is that an all-reduce is exactly a reduce-scatter followed by an all-gather, which is why sharded training can replace one with the other two and pay roughly the same total.

operation        each device starts with   each device ends with
all-reduce       its own full gradient    the summed full gradient
reduce-scatter   its own full gradient    the summed slice it owns
all-gather       the slice it owns        every slice, assembled

How long a gradient sync takes#

A ring all-reduce does not send everything to everyone. Each device passes slices around a ring, and the traffic per device works out at twice the payload times N minus one over N — approaching twice the payload however many devices you add, which is the useful property. Seven billion bf16 gradients is 14 GB, so on eight devices each one moves 24.5 GB per step.

params, N = 7e9, 8
grad_bytes = params * 2                     # bf16 gradients
payload = grad_bytes * 2 * (N - 1) / N      # ring all-reduce traffic per device

print(f"{grad_bytes/1e9:.0f} GB of gradients -> {payload/1e9:.1f} GB moved per device")

for link, bw in [("NVLink", 900e9), ("PCIe 4.0 x16", 25e9),
                 ("100 Gb Ethernet", 12.5e9), ("10 Gb Ethernet", 1.25e9)]:
    print(f"  {link:18} {payload/bw*1000:9.1f} ms per step")

# 14 GB of gradients -> 24.5 GB moved per device
#   NVLink                  27.2 ms per step
#   PCIe 4.0 x16           980.0 ms per step
#   100 Gb Ethernet       1960.0 ms per step
#   10 Gb Ethernet       19600.0 ms per step

A time in milliseconds means nothing until you compare it with the step it sits inside. A transformer forward and backward pass costs about six FLOPs per parameter per token, so at 4,096 tokens per device that is 1,376 TFLOPs a step, and eight H100s sustaining a realistic forty per cent of their 989 TFLOP/s bf16 peak take 435 ms of compute.

interconnect        sync      step        sync as a share
NVLink            27 ms     462 ms          5.9%
PCIe 4.0         980 ms    1415 ms         69.3%
10 Gb Ethernet 19600 ms   20035 ms         97.8%

Memory per GPU, worked#

The memory question is the one people get wrong by a factor of five, because they count the parameters and stop. Mixed-precision training with Adam holds five things per parameter, and the parameters are the smallest of them.

params = 7e9

GiB = 1024 ** 3          # memory is sold and reported in binary units
costs = [("bf16 weights", 2), ("bf16 gradients", 2), ("fp32 master copy", 4),
         ("Adam momentum", 4), ("Adam variance", 4)]
for name, b in costs:
    print(f"{name:18} {b:2d} bytes/param  {params*b/GiB:6.1f} GiB")

total = params * sum(b for _, b in costs) / GiB
print(f"{'TOTAL':18} {sum(b for _,b in costs):2d} bytes/param  {total:6.1f} GiB")

# bf16 weights        2 bytes/param    13.0 GiB
# bf16 gradients      2 bytes/param    13.0 GiB
# fp32 master copy    4 bytes/param    26.1 GiB
# Adam momentum       4 bytes/param    26.1 GiB
# Adam variance       4 bytes/param    26.1 GiB
# TOTAL              16 bytes/param   104.3 GiB

That total is also the rule of thumb worth memorising, because it answers the interview question directly: mixed-precision Adam costs roughly sixteen bytes per parameter, so multiply your parameter count by sixteen and compare it with your device. Plain SGD with momentum costs eight and full fp32 Adam costs twenty. Activations sit on top and scale with batch size rather than parameter count, which is why they are the part you can trade.

weights = params * 2 / GiB         # bf16 copy every device keeps resident
grads   = params * 2 / GiB
opt     = params * 12 / GiB        # fp32 master + Adam momentum + Adam variance

for name, per_gpu in [
    ("none (plain DDP)",      weights + grads + opt),
    ("stage 1 (optimiser)",   weights + grads + opt / N),
    ("stage 2 (+ gradients)", weights + grads / N + opt / N),
    ("stage 3 (+ parameters)", (weights + grads + opt) / N),
]:
    print(f"{name:23} {per_gpu:6.1f} GiB  fits an 80 GiB card: {per_gpu < 80}")

# none (plain DDP)        104.3 GiB  fits an 80 GiB card: False
# stage 1 (optimiser)      35.9 GiB  fits an 80 GiB card: True
# stage 2 (+ gradients)    24.4 GiB  fits an 80 GiB card: True
# stage 3 (+ parameters)   13.0 GiB  fits an 80 GiB card: True

Model FLOPs utilisation#

MFU is the single number that says whether a training run is healthy: the FLOPs your model actually needed, divided by the FLOPs the hardware could have delivered in the same wall-clock. It ignores everything you cannot bill for — recomputation from activation checkpointing, padding, waiting on the network — which is exactly what makes it honest.

tokens_per_step = N * 4096
flops = 6 * params * tokens_per_step        # 6 FLOPs per parameter per token

compute_s = flops / (989e12 * 0.4 * N)      # 40% of bf16 peak, sustained
comm_s = payload / 900e9                    # NVLink
step = compute_s + comm_s

achieved = flops / step / N
print(f"step {step*1000:.0f} ms   achieved {achieved/1e12:.0f} TFLOP/s per GPU")
print(f"MFU {achieved/989e12*100:.1f}%")

# step 462 ms   achieved 372 TFLOP/s per GPU
# MFU 37.6%

MFU you measure    what it usually means
40-50%             healthy; tune elsewhere
25-35%             communication or activation checkpointing overhead
10-20%             the data loader, or a batch too small to fill the GPU
under 10%          something is on the CPU, or synchronising every step