Author(s): Athira PT

Originally published on Towards AI.

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

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.

Introduction

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.

Section 1: Why Do We Need Multi-GPU Training?

Deep learning models and datasets have grown exponentially. Consider a standard computer vision task training ResNet-50 on ImageNet:

• Dataset Size: 1,280,000 images
• Batch Size: 256
• Single Iteration (Forward + Backward + Step): ~300 ms on a modern GPU

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:

1. Model Parallelism (Tensor / Pipeline Parallelism): The model is too massive to fit inside a single GPU’s VRAM (e.g., a 70B parameter LLM). Layers or weight matrices are sliced and spread across multiple GPUs.
2. Data Parallelism: The model easily fits inside one GPU’s memory, but the dataset is vast. Each GPU holds an identical copy of the full model and processes a distinct slice (chunk) of the dataset simultaneously.

PyTorch’s DistributedDataParallel (DDP) relies on Data Parallelism.

Section 2: What Happens Inside One GPU?

Before adding multiple GPUs, let’s revisit the canonical single-GPU PyTorch training loop:

In a standard training loop:

1. A Batch of Data is loaded into VRAM.
2. A Forward Pass evaluates predictions and computes a scalar Loss.
3. The Backward Pass computes gradients ∇W via autograd backpropagation.
4. The Optimizer Step updates weights: Wₜ₊₁ = Wₜ − η · ∇Wₜ

Every single parameter update is strictly driven by the gradients calculated from that local batch.

Section 3: The First Wrong Idea (Why Naive Splitting Fails)

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?

The Flaw

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 remain 100% mathematically identical at every step?

Section 4: DDP Big Picture

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:

The Mathematical Guarantee

If:

1. GPU 0 and GPU 1 start with identical initial weights W₀.
2. Both GPUs compute their local gradients (G₀ and G₁.).
3. PyTorch computes the mean gradient: Gₐᵥ₉ = (G₀ + G₁)/2
4. Both GPUs apply the exact same update rule using W₁ = W₀ — η · Gₐᵥ₉

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.

Section 5: Essential DDP Terminology

Before writing code, let me clarify the core terminology:

• Process: An independent Python OS process. In DDP, 1 Process = 1 GPU.
• Rank: A unique global integer assigned to each process across the entire cluster (e.g., 0, 1, 2, … , N-1).
• Local Rank: The process ID relative to a single physical machine/node. If a node has 4 GPUs, local ranks will be 0, 1, 2, 3.
• World Size: The total number of processes participating in training across all nodes.
• Process Group: The communication channel connecting all ranks together.
• Backend: The underlying library handling communication:

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).

Section 6: Why torchrun Exists

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:

1. Spawning N copies of your Python script.
2. Setting critical environment variables in each spawned process:

Instead of executing:

python train.py

You execute:

torchrun --nproc_per_node=2 train.py

Section 7: The Life of One Epoch

Let me trace a single epoch inside a 2-GPU DDP setup:

Section 8: Code Walkthrough: Understanding Every Line

Now let me walk through a complete, runnable DDP script:

pytorch-primer/scripts/ddp_train.py at main · sourangshupal/pytorch-primer

Contribute to sourangshupal/pytorch-primer development by creating an account on GitHub.

github.com

1. Initializing the Process Group (ddp_setup)

import os
import platform
import torch
from torch.distributed import init_process_group, destroy_process_group

def 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)

Key Insights:

• init_process_group(): Establishes the communication socket between all processes. Rank 0 listens on MASTER_ADDR:MASTER_PORT while other ranks connect to it.
• torch.cuda.set_device(rank): Binds this specific Python process strictly to GPU rank. When you execute .to("cuda") or to(rank), it targets this GPU exclusively.

2. Splitting Data (DistributedSampler)

from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler

def 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] Why shuffle=False is mandatory: When using DistributedSampler, the sampler itself handles data shuffling across ranks. If you pass shuffle=True into DataLoader, it conflicts with DistributedSampler and raises a runtime error!

How DistributedSampler Works:

If your dataset has 10 samples ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) and world_size = 2:

• Rank 0 receives: [0, 2, 4, 6, 8]
• Rank 1 receives: [1, 3, 5, 7, 9]

No data point is duplicated across GPUs within the same batch.

3. Model Wrapping (DDP(model))

from torch.nn.parallel import DistributedDataParallel as DDP

# 1. Instantiate local model
model = NeuralNetwork(num_inputs=2, num_outputs=2)

# 2. Move model to process GPU FIRST
model.to(rank)

# 3. Wrap model with DDP
model = DDP(model, device_ids=[rank] if torch.cuda.is_available() else None)

What happens inside DDP(model)?

1. Initial Broadcast: Rank 0 broadcasts its exact initial weights to all other ranks so everyone starts at step 0 with identical parameters.
2. Hook Registration: DDP attaches backward hooks to every parameter tensor in the model.

4. The Backward Pass Hook Magic

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()

What happens during loss.backward()?

1. Autograd computes gradients layer-by-layer backwards.
2. As soon as a gradient for a tensor is ready, DDP’s autograd hook fires.
3. DDP bucketizes gradients and begins Ring-AllReduce in the background while earlier layers are still computing backprop!
4. By the time loss.backward() completes, all gradients across all GPUs have been reduced to their mean.
5. optimizer.step() applies the identical averaged gradients.

5. Why set_epoch(epoch) is Required

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:

• Every epoch uses a new random data ordering.
• All GPUs remain in sync regarding which partition of the shuffled dataset they receive.

6. Cleanup (destroy_process_group)

 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.

7. Main Entry Point

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)

Section 9: What DDP Does Automatically

Here is a summary of what you write vs. what DDP handles under the hood:

Section 10: Common Beginner Pitfalls

Here are 5 common mistakes when implementing DDP:

1. Launching with python script.py

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

2. Setting shuffle=True in DataLoader

Setting shuffle=True inside DataLoader conflicts with DistributedSampler. Always set shuffle=False in DataLoader and let DistributedSampler manage shuffling.

3. Forgetting train_loader.sampler.set_epoch(epoch)

If omitted, your model trains on identical batch ordering every epoch, harming generalization.

4. Printing or Saving Checkpoints on Every Rank

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 CHECKPOINT
if rank == 0:
 torch.save(model.module.state_dict(), "model.pt")

5. Forgetting model.module when Saving/Accessing Custom Methods

# Access original model methods or parameters:
raw_model = model.module

Section 11: How DDP Differs from DataParallel (DP)

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.

Section 12: Final Mental Model

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. During loss.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!

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI