Distributed Training

Training GPT-3 took 355 GPU-years on a single GPU — but it was done in weeks using thousands of GPUs in parallel. Distributed training is how you scale deep learning beyond what a single GPU can handle.

Why Distribute Training?

Strategy 1: Data Parallelism (DDP)

The simplest and most common approach. Each GPU gets a complete copy of the model and a different mini-batch of data. After the backward pass, gradients are averaged across all GPUs (all-reduce), then each GPU updates its copy.

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

# Initialise process group
dist.init_process_group(backend='nccl')  # nccl = NVIDIA's collective comms library
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)

# Wrap model in DDP
model = MyModel().to(local_rank)
model = DDP(model, device_ids=[local_rank])

# DistributedSampler ensures each GPU gets different data
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler, batch_size=32)

# Training loop — gradients automatically averaged!
for batch in dataloader:
    loss = model(batch).mean()
    loss.backward()  # All-reduce happens here
    optimizer.step()

# Launch: torchrun --nproc_per_node=4 train.py

Strategy 2: Model Parallelism (Tensor Parallelism)

When the model is too large for a single GPU, split the model itself across GPUs. Tensor parallelism shards individual weight matrices — each GPU holds a slice of each layer and they communicate after each matrix multiply.

Strategy 3: Pipeline Parallelism

Split the model layer-wise (Layer 1–4 on GPU0, Layer 5–8 on GPU1, etc.). Use micro-batches so multiple batches are in flight simultaneously — like a CPU pipeline.

ZeRO — DeepSpeed's Memory Miracle

ZeRO (Zero Redundancy Optimizer) eliminates the memory redundancy in data parallel training. Instead of each GPU holding a full copy of weights + gradients + optimizer states, ZeRO shards them across GPUs:

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" },
    "offload_param": { "device": "cpu" }
  },
  "bf16": { "enabled": true },
  "train_batch_size": 256,
  "gradient_accumulation_steps": 4
}

# Launch:
# deepspeed --num_gpus=8 train.py --deepspeed ds_config.json

Communication: NVLink & InfiniBand

Frequently Asked Questions