PyTorch’s DistributedDataParallel (DDP) trains a model across multiple GPUs by giving processes different portions of the data and synchronizing their gradients. It is useful when one GPU is too slow for a large training workload. Unlike DataParallel, which replicates the model on each forward pass, DDP generally incurs less overhead and can also span multiple machines. section 7.8 section 7.8.1
A single-machine training workflow has five parts:
- Initialize a process group, assign each process a GPU, and divide the data yourself; DDP does not split it automatically. section 7.8.1 section 7.8.2
- Create the model and wrap it with DistributedDataParallel. section 7.8.1
- In each process, run batches through the usual forward pass, loss calculation, backpropagation, and optimizer step. DDP synchronizes gradients so the processes update the same model. section 7.8.1 section 7.8.2
- Coordinate completion where needed—for example, with dist.barrier() before the next epoch—and destroy the process group when training ends. section 7.8.1
- Guard the script’s entry point when launching training processes. For multi-machine runs, also plan model loading and saving carefully. section 7.8.1
Two-GPU training examples
A PyTorch example trains a Flatten-and-Linear network on MNIST. It initializes communication with NCCL, uses mp.spawn to launch a process for each GPU, and assigns devices by process rank. A GCN variation uses the Amazon Products dataset and NeighborLoader on a dual-GPU system. section 7.8.2
Listing 7.2 Training using DDP
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel #1
from torch.utils.data import DataLoader #2
from torchvision import datasets, transforms
class Flatten(nn.Module):
def forward(self, input):
return input.view(input.size(0), -1)
def train(model, trainloader,
criterion,
optimizer,
device): #3
model.train()
for batch_idx, (data, target) in enumerate(trainloader):
print(f'Process {device}, Batch {batch_idx}')
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
def main(rank, world_size): #4
filepath = '~/.pytorch/MNIST_data/'
dist.init_process_group( #5
backend='nccl',
init_method='tcp://localhost:23456',
rank=rank,
world_size=world_size #6
)
torch.manual_seed(0) #7
device = torch.device(f'cuda:{rank}') #8
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5,),
(0.5,))]
)
trainset = datasets.MNIST(filepath , #9
download=True, #9
train=True, #9
transform=transform) #9
train_loader = DataLoader(trainset, #10
batch_size=64, #10
shuffle=True, #10
num_workers=2) #10
model = nn.Sequential(Flatten(), nn.Linear(784, 10)).to(device)
model = DistributedDataParallel(model, device_ids=[rank]) #11
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
train(model, train_loader, criterion, optimizer, device) #12
The training listing shows where the model is wrapped, how each process runs the training loop, and how distributed execution is launched. section 7.8.2
Choose DDP with the available hardware in mind. Moving from one GPU to multiple devices—or to a small cluster—adds planning, setup, and synchronization work alongside the potential training-speed benefit. section 7.4.1 section 7.8.1
