torch.distributed is PyTorch’s package for initializing processes and coordinating communication during distributed training. It is useful when multiple GPUs or machines should train in parallel, with synchronized model updates. section 7.8.1 section 7.8.2
A typical distributed-training workflow proceeds as follows:
- Initialize a process group with dist.init_process_group and configure an address, port, and communication backend. section 7.8.2
- Create the model, wrap it with DistributedDataParallel, and assign each process its device and data subset. section 7.8.1 section 7.8.2
- Run forward and backward passes, update parameters, and synchronize processes with dist.barrier() when required. section 7.8.1 section 7.8.2
- Launch one main function per process with mp.spawn, then destroy the process group after training. section 7.8.1 section 7.8.2
MNIST with two GPUs
A PyTorch script initializes NCCL communication, assigns each process a CUDA device, wraps a neural network in DistributedDataParallel, and trains each process on its subset of MNIST while synchronizing gradients. 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 DDP listing provides a concrete implementation skeleton covering imports, training, process setup, device selection, and multiprocessing launch. section 7.8.2
