Data Parallelism

The gradient of a sum is the sum of gradients, so the batch can be split across devices and the results added. This is gradient accumulation with the micro-batches running concurrently rather than in sequence.

Every rank starts from identical weights, computes on its own shard of the batch, all-reduces, and steps. Because all ranks apply the same averaged gradient to the same weights, they remain identical without anything being synchronised except that one tensor.

Why it scales

The all-reduce moves $2P$ bytes per rank regardless of how many ranks there are. Compute per rank is constant. So the communication-to-compute ratio does not degrade with scale, and that is the property no other parallelism strategy has.

Larger micro-batches make the left side bigger without changing the right, which is why data parallelism wants as much local batch as memory allows — and why activation memory is the constraint that binds first.

Overlapping the collective

The gradient for the last layer is ready long before the first layer’s. DDP exploits this: it registers autograd hooks, buckets gradients as they become available, and launches an all-reduce per bucket while the backward pass is still running. By the time the backward finishes, most of the communication is done.

Batch size and the learning rate

Doubling the ranks doubles the effective batch, which changes the optimisation problem. Goyal et al.’s two rules are still the practical answer:

Linear scaling. Multiply the learning rate by the number of ranks. The gradient’s variance falls as $1/N$, so a proportionally larger step is justified — up to a point.

Warmup. Ramp to the scaled rate over a few epochs. The linear rule fails at the start of training, when the loss surface changes fastest.

Both hold to a limit. Past a critical batch size, additional examples stop reducing gradient variance meaningfully and the extra compute buys nothing — the quantity Kaplan et al. estimated and the reason data parallelism eventually stops being the answer.

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


def setup(model, dataset, batch: int, accum: int):
    dist.init_process_group("nccl")
    rank, world = dist.get_rank(), dist.get_world_size()
    torch.cuda.set_device(rank)

    model = DDP(model.cuda(rank), device_ids=[rank], gradient_as_bucket_view=True)

    # Without the sampler every rank sees the same data and the whole
    # exercise averages N copies of one gradient.
    sampler = DistributedSampler(dataset, num_replicas=world, rank=rank)
    loader = DataLoader(dataset, batch_size=batch, sampler=sampler)
    return model, loader, sampler


def epoch(model, loader, sampler, opt, epoch_idx: int, accum: int):
    # Reshuffles per epoch, identically on every rank. Omit it and every
    # epoch sees the same order.
    sampler.set_epoch(epoch_idx)

    for i, batch in enumerate(loader):
        sync = (i + 1) % accum == 0
        with model.no_sync() if not sync else contextlib.nullcontext():
            (model(**batch).loss / accum).backward()
        if sync:
            opt.step()
            opt.zero_grad(set_to_none=True)

set_epoch is the line everyone forgets. Without it DistributedSampler uses the same permutation every epoch, so the model sees an identical data order throughout training — which degrades the result and produces no error.

Every rank holding a full copy of the parameters, gradients and optimiser state is $16P$ bytes per device to store one model’s worth of information $N$ times. ZeRO shards all three across the data-parallel group, and FSDP is the same idea in PyTorch. Both keep the structure above — split the batch, average the gradients — and change only what is resident between steps, which is why “data parallelism” now nearly always means one of them.