Distributed Training with Multiple GPUs

When people first hear that a model is being trained on several GPUs at once, a natural picture forms: the network is cut into pieces, each GPU holds one piece, and the pieces pass data back and forth like an assembly line. That picture is a fair description of some other parallel techniques. It is not how PyTorch's DistributedDataParallel strategy works. This article starts from that picture on purpose, so you can test it against the mechanism.

DDP uses replication. The model starts on the CPU, and a copy of it is placed on each GPU. Every copy is the full network, with all of its layers and all of its weights. Each GPU runs its own process, and that process works with its own copy. Nothing about the model is cut apart to fit across devices.

The two copies start out identical. Keeping them identical during training takes extra work, which the later sections explain. First, the data side of the strategy needs to be clear.

Each model copy is fed a different minibatch from the data loader in every training iteration. The copies therefore never process the same examples at the same time. A DistributedSampler is the tool that keeps those minibatches from overlapping, so two GPUs are not both working on the same samples. Replication gives each device a full model, and the sampler gives each device a distinct slice of the data.

Because each copy sees a different slice of the data, its outputs differ from the others. Each copy produces its own logits and computes its own gradients during the backward pass. If nothing were shared, the copies would gradually drift apart and stop being the same model. DDP prevents this by combining the gradients from all GPUs and using the combined result to update every copy in the same way.

Trace one iteration in order. Each GPU runs its forward pass on its own minibatch, producing its own logits. Each GPU then runs its backward pass, producing its own local gradients. The forward and backward passes run independently, and only after both finish do the GPUs exchange gradient information. Once the gradients are synchronized, every replica applies the same update, so every copy ends the iteration with the same weights.

The ideal speedup from DDP is near-linear in the number of GPUs. If a single GPU needs a given amount of time for a full training epoch, the idealized estimate divides that time by the number of GPUs.

DDP code is written as a script that is launched outside the notebook. The process model explains why. DDP runs a separate process on each GPU, and those processes run independently while coordinating at the gradient step. A notebook, by contrast, gives you a single interactive session. The DDP structure of one process per GPU does not fit inside that one session.

Moving from a single-GPU script to a DDP script is mostly a matter of setup. You import the additional distributed submodules, classes, and functions from PyTorch, and you arrange for each process to operate on its own GPU and its own data subset. The core training logic stays familiar.