Data parallel training in PyTorch

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.