Understanding PyTorch DistributedDataParallel (DDP) From Scratch: Why Every GPU Trains the Entire…

PyTorch's DistributedDataParallel (DDP) enables multi-GPU training by synchronizing gradients across all GPUs via AllReduce, ensuring each GPU maintains identical model weights. The article explains the mechanics of DDP, including the necessity of one Python process per GPU, the use of torchrun for launching, and the requirement of shuffle=False with DistributedSampler. It provides a first-principles guide to help developers troubleshoot common distributed training issues.

Why does every GPU need its own copy of the model? How do gradients get averaged without a central master worker? Why is shuffle=False mandatory when using a DistributedSampler? Here is the first-principles guide to PyTorch DDP that explains what actually happens under the hood before walking line-by-line through a complete training script. If you have ever attempted to scale a PyTorch training pipeline from a single GPU to a cluster or a multi-GPU workstation, you have likely encountered DistributedDataParallel DDP . Most tutorials start by telling you: “Just wrap your model in DDP model and launch with torchrun " While that sounds delightfully simple, it often leaves developers stranded when subtle bugs appear. Why did your training freeze? Why are all GPUs receiving identical data batches? Why is loss diverging? Why does PyTorch insist on setting shuffle=False in the data loader? To troubleshoot or build robust distributed deep learning systems, you need a solid mental model of what every GPU is doing at every microsecond of the training loop . In this article, we will unpack DDP from first principles: starting from single-GPU mechanics, exposing the intuition behind gradient synchronization, clarifying distributed terminology, and finally walking line-by-line through a clean, working PyTorch DDP script. Deep learning models and datasets have grown exponentially. Consider a standard computer vision task training ResNet-50 on ImageNet: At 300 ms per batch, 5,000 batches per epoch over 90 epochs takes days on a single card. To accelerate training, we need multiple GPUs working concurrently. But how should GPUs divide the work? Broadly, distributed training falls into two paradigms: PyTorch’s DistributedDataParallel DDP relies on Data Parallelism . Before adding multiple GPUs, let’s revisit the canonical single-GPU PyTorch training loop: In a standard training loop: Every single parameter update is strictly driven by the gradients calculated from that local batch. Suppose we have 2 GPUs GPU 0 and GPU 1 . Naive Idea: What if we simply split the batch of 256 images into two chunks of 128 images, load them on GPU 0 and GPU 1, and run independent training loops on both? Because GPU 0 and GPU 1 processed completely different images, their loss outputs differed, producing different gradients . Applying those gradients independently causes Model Copy A and Model Copy B to diverge immediately. Within a few steps, you no longer have one trained model—you have two competing models with completely different weights The Core Challenge:How do we let GPUs train on different data batches while ensuring all GPU copies remain100% mathematically identicalat every step? DDP resolves this challenge through Gradient Synchronization via AllReduce . Instead of letting each GPU update its weights using its own local gradients, DDP averages the gradients across all GPUs before the optimizer step occurs. Here is the master blueprint of DDP: If: Then GPU 0 and GPU 1 will arrive at identical weights W₁. This mathematical invariant holds true for step 1, step 2, and step 100,000. Before writing code, let me clarify the core terminology: NCCL NVIDIA Collective Communication Library :Optimized for GPU-to-GPU transfers via NVLink or PCIe.Always use NCCL for CUDA GPUs. GLOO:A cross-platform fallback library used for CPU distributed training or Windows environments . When you type python train.py, your operating system launches a single Python process . However, DDP requires N independent Python processes running simultaneously — one for each GPU. torchrun is PyTorch’s official process launcher. It automates: Instead of executing: python train.py You execute: torchrun --nproc per node=2 train.py Let me trace a single epoch inside a 2-GPU DDP setup: Now let me walk through a complete, runnable DDP script: pytorch-primer/scripts/ddp train.py at main · sourangshupal/pytorch-primer https://github.com/sourangshupal/pytorch-primer/blob/main/scripts/ddp train.py python import osimport platformimport torchfrom torch.distributed import init process group, destroy process groupdef ddp setup rank: int, world size: int - None: Only set MASTER ADDR / MASTER PORT if torchrun hasn't already set them. if "MASTER ADDR" not in os.environ: os.environ "MASTER ADDR" = "localhost" if "MASTER PORT" not in os.environ: os.environ "MASTER PORT" = "12345" if platform.system == "Windows": Windows fallback gloo os.environ "USE LIBUV" = "0" init process group backend="gloo", rank=rank, world size=world size else: Standard Linux GPU setup nccl init process group backend="nccl", rank=rank, world size=world size if torch.cuda.is available : torch.cuda.set device rank python from torch.utils.data import Dataset, DataLoaderfrom torch.utils.data.distributed import DistributedSamplerdef prepare dataset : train ds = ToyDataset X train, y train test ds = ToyDataset X test, y test train loader = DataLoader dataset=train ds, batch size=2, shuffle=False, CRITICAL: shuffle MUST be False pin memory=True, drop last=True, sampler=DistributedSampler train ds , Handles chunking across GPUs test loader = DataLoader dataset=test ds, batch size=2, shuffle=False, return train loader, test loader IMPORTANT Whyshuffle=False is mandatory:When usingDistributedSampler, the sampler itself handles data shuffling across ranks. If you passshuffle=True intoDataLoader, it conflicts withDistributedSampler and raises a runtime error If your dataset has 10 samples 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 and world size = 2: No data point is duplicated across GPUs within the same batch. python from torch.nn.parallel import DistributedDataParallel as DDP 1. Instantiate local modelmodel = NeuralNetwork num inputs=2, num outputs=2 2. Move model to process GPU FIRSTmodel.to rank 3. Wrap model with DDPmodel = DDP model, device ids= rank if torch.cuda.is available else None When you invoke loss.backward in your training loop: for epoch in range num epochs : CRITICAL: Reshuffle sampler seed for each epoch train loader.sampler.set epoch epoch model.train for features, labels in train loader: features, labels = features.to rank , labels.to rank logits = model features loss = F.cross entropy logits, labels optimizer.zero grad loss.backward <--- HERE IS WHERE DDP ALLREDUCE HAPPENS optimizer.step Notice line 3 in the loop above: train loader.sampler.set epoch epoch Without this line, DistributedSampler will generate the exact same dataset ordering for every single epoch . set epoch epoch seeds the sampler deterministically per epoch so that: destroy process group At the end of training, destroy process group cleanly closes the socket connection between ranks. Forgetting this can lead to orphaned Python processes holding onto VRAM. if name == " main ": Environment variables populated by torchrun world size = int os.environ.get "WORLD SIZE", 1 rank = int os.environ.get "LOCAL RANK", os.environ.get "RANK", 0 if rank == 0: print "PyTorch version:", torch. version print "CUDA available:", torch.cuda.is available print "Number of GPUs available:", torch.cuda.device count torch.manual seed 123 main rank, world size, num epochs=3 Here is a summary of what you write vs. what DDP handles under the hood: Here are 5 common mistakes when implementing DDP: Running python script.py only starts 1 process on 1 GPU. You must launch distributed jobs with torchrun: torchrun --nproc per node=2 scripts/ddp train.py Setting shuffle=True inside DataLoader conflicts with DistributedSampler. Always set shuffle=False in DataLoader and let DistributedSampler manage shuffling. If omitted, your model trains on identical batch ordering every epoch, harming generalization. If you log output or save model checkpoints inside the training loop without checking if rank == 0:, all N processes will attempt to write to disk simultaneously, causing corrupt files or cluttered terminal output. CORRECT WAY TO SAVE CHECKPOINTif rank == 0: torch.save model.module.state dict , "model.pt" Access original model methods or parameters:raw model = model.module PyTorch legacy codebase tutorials sometimes mention torch.nn.DataParallel DP . Do not use DP for training. Here is how they compare: DataParallel uses a single master thread to scatter input data and gather outputs across GPUs. The single Python Global Interpreter Lock GIL creates a massive bottleneck. DDP avoids this entirely by giving every GPU its own independent Python interpreter process. If you retain only one paragraph from this article, let it be this: The Golden Rule of DDP:Every GPU holds a complete, independent copy of the model. Every GPU processes a unique batch of data. Duringloss.backward , DDP uses Ring-AllReduce to average gradients across all GPUs. Because every GPU starts from identical weights and applies identical averaged gradients, every model copy remains in perfect mathematical synchronization after every optimization step. Happy distributed training Understanding PyTorch DistributedDataParallel DDP From Scratch: Why Every GPU Trains the Entire… https://pub.towardsai.net/understanding-pytorch-distributeddataparallel-ddp-from-scratch-why-every-gpu-trains-the-entire-2b80f89bed3d was originally published in Towards AI https://pub.towardsai.net on Medium, where people are continuing the conversation by highlighting and responding to this story.