Gradient synchronization coordinates updates across training processes so replicated models remain consistent. In PyTorch DistributedDataParallel, each process trains on its data subset while sharing gradients across GPUs. section 7.8.2
A typical distributed-training pipeline is:
- Initialize a distributed process group and configure communication, such as NCCL with a host address and port. section 7.8.2
- Assign each process a device and data subset, wrap the model with DistributedDataParallel, then run forward and backward passes. section 7.8.2
- Synchronize gradients across processes before optimizer.step updates the replicated parameters. section 7.8.2
Two-GPU MNIST training
mp.spawn starts one training process per GPU. Each process trains a neural network on its MNIST subset; synchronized gradients keep the model replicas aligned. 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
Inspect the DDP listing to follow process setup, model wrapping, backward propagation, and optimizer updates. section 7.8.2
