How it works
In data parallelism, each GPU receives an identical copy of the model and a different subset of the training batch. Each GPU computes gradients independently. At the end, the gradients are aggregated with all-reduce and the weights are updated in sync.
This scales training throughput linearly with the number of GPUs, up to the limit of communication between them. Synchronization can be synchronous, where all GPUs wait, or asynchronous, with independent updates, which is less common today.
It is also possible to use gradient accumulation to simulate larger batches when memory is limited, processing several mini-batches before updating the weights.
Practical example
Training a model with 8 GPUs: each GPU processes 1/8 of the batch. If the total batch is 256 examples, each GPU sees 32. After forward and backward, the gradients are summed and the model is updated.
import torch
from torch.nn.parallel import DistributedDataParallel as DDP
model = MyModel().to(device)
model = DDP(model, device_ids=[gpu_id])
for i, batch in enumerate(dataloader):
loss = model(batch)
loss.backward()
if (i + 1) % 4 == 0:
optimizer.step()
optimizer.zero_grad()What it is for
Training large models in a viable time. Without data parallelism, training a 70 billion parameter LLM on a single GPU would take months. With dozens of GPUs, it drops to weeks.
It is the simplest parallelism technique and the first to be applied when the model fits in one GPU's memory but training is slow.