Distributed Training: Data Parallel vs Model Parallel
When your model doesn't fit on one GPU, you need model parallelism. When it fits but training is too slow, you need data parallelism. When it's both, you need pipeline parallelism or tensor parallelism — and the choices interact in non-obvious ways.
Data Parallelism:
Each GPU gets a full copy of the model and a different mini-batch. Gradients are synchronised across GPUs after each backward pass. This is the default and works well when the model fits in a single GPU's memory.
AllReduce: the standard gradient synchronisation primitive. Each GPU computes its gradients, then AllReduce averages them across all GPUs. With ring-AllReduce (used by NCCL), communication cost is proportional to gradient size × 2, independent of the number of GPUs.
PyTorch DDP (DistributedDataParallel): hooks into the backward pass to overlap gradient communication with backward computation. As soon as a gradient is ready, it starts being communicated while other gradients are still being computed. Near-linear scaling up to ~64 GPUs for large models.
ZeRO (Zero Redundancy Optimizer): addresses memory inefficiency in data parallelism. With naive DDP, each GPU stores the full model parameters, gradients, AND optimiser state. ZeRO-1: shard optimiser state. ZeRO-2: shard gradients. ZeRO-3: shard parameters. ZeRO-3 reduces per-GPU memory by the number of GPUs but introduces communication overhead.
Model Parallelism:
When the model doesn't fit on one GPU. Two main variants:
Tensor parallelism: split individual layers across GPUs. A matrix multiply A × B is split so each GPU holds a column shard of B and computes a partial result. Requires all-to-all communication for each layer. Used in Megatron-LM for large transformer models.
Pipeline parallelism: split the model into stages, each stage on a different GPU. GPU 1 processes layers 1–8, GPU 2 processes layers 9–16, etc. The key challenge: GPU 2 is idle while GPU 1 processes the first micro-batch. Solved with micro-batching: break each batch into 8 micro-batches, pipeline them through the stages. GPU utilisation approaches ~(n-1)/n where n is the number of stages.
3D Parallelism:
Large language models (GPT-3, PaLM scale) use all three simultaneously: data parallelism across groups of GPUs, pipeline parallelism across stages within a group, tensor parallelism within each stage. Megatron-DeepSpeed uses this configuration for models with hundreds of billions of parameters.
Gradient accumulation:
Simulates a larger batch size by computing gradients over multiple micro-batches before updating. Useful when you can't increase batch size due to memory but want the stability benefits of larger batches. gradient_accumulation_steps=8 with micro_batch=16 is equivalent to batch=128.
Practical guidance:
Start with DDP. Add ZeRO stages if you need memory relief. Only add model parallelism if the model genuinely doesn't fit on a single node. Communication costs scale super-linearly with node count — profile before scaling.
```python import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP
def train_ddp(rank, world_size, model, dataset, accum_steps=4): """Minimal DDP loop with gradient accumulation. accum_steps=4 with micro_batch=8 → effective batch = 32 × world_size.""" dist.init_process_group('nccl', rank=rank, world_size=world_size) model = model.to(rank) model = DDP(model, device_ids=[rank]) optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
loader = torch.utils.data.DataLoader( dataset, sampler=torch.utils.data.distributed.DistributedSampler(dataset), batch_size=8, )
for step, batch in enumerate(loader): loss = model(batch.to(rank)).loss / accum_steps # scale loss loss.backward() if (step + 1) % accum_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() optimizer.zero_grad()
dist.destroy_process_group()
# Production note: unused params in DDP cause a hang — set find_unused_parameters=True # only as a debug fallback; fix the architecture instead. ```