*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 ?
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
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, Datafrom torch.utils.data.distributed import DistributedSamplerdef prepare_dataset(): train_ds = ToyDataset(X_train, y_train) test_ds = ToyDataset(X_test, y_test) train_ = Data( 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_ = Data( dataset=test_ds, batch_size=2, shuffle=False, ) return train_, test_
[IMPORTANT]Whyshuffle=False is mandatory:When usingDistributedSampler, the sampler itself handles data shuffling across ranks. If you passshuffle=True intoData, 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.
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_.sampler.set_epoch(epoch) model.train() for features, labels in train_: 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_.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 Data conflicts with DistributedSampler. Always set shuffle=False in Data 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.
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… was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.