Data parallelism speeds up training by copying the same model onto multiple devices and giving each device a different portion of the data. Each copy computes its own gradients. Before updating the model, the devices must synchronize those gradients so their copies do not drift apart. section 16.4
A training step follows this sequence:
- Run the forward pass, compute the loss, and call loss.backward() on each device. section 16.4
- Average gradients across devices, then call opt.step(). In a manual implementation, dist.all_reduce performs the reduction; DistributedDataParallel (DDP) handles synchronization when it wraps the model. section 16.3 section 16.4
Listing 16.7 Simple example of gradient synchronization
input_data = torch.randn(3) + rank #1
output = model(input_data) #2
loss = output.mean() #2
opt.zero_grad() #3
loss.backward() #4
for param in model.parameters(): #5
dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM) #5
param.grad.data /= world_size #6
opt.step() #7
The manual training loop shows the critical ordering: backward pass, all-reduce and average each parameter’s gradient, then update the parameters. section 16.4
Let PyTorch handle the replicas
With a distributed process group initialized and the model, dataset, device, rank, world_size, and batch_size defined: from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler model = DDP(model.to(device)) sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler) DDP handles gradient synchronization. DistributedSampler assigns each process its portion of the dataset rather than having every loader repeat the same entries. section 16.4
The trade-off is memory: every device must hold a model copy. If the model cannot fit on one device, data parallelism alone is not enough; model parallelism splits the model instead. For larger systems, data parallelism can be combined with other strategies, although all-reduce can become a bottleneck at scale. section 16.5 section 16.6
