DistributedDataParallel

DDP is the standard way to scale training beyond one GPU with near-linear speedup.

DistributedDataParallel (DDP) runs one process per GPU. Each process holds a full model replica, trains on its own slice of every batch, and during backward() DDP all-reduces (averages) gradients across processes in buckets, overlapping communication with computation. After step() all replicas hold identical weights.

The moving parts:

Launch with torchrun --nproc_per_node=4 train.py. It starts the processes and sets RANK, LOCAL_RANK and WORLD_SIZE.
Each process calls dist.init_process_group("nccl") (use "gloo" on CPU), picks cuda:LOCAL_RANK and wraps the model: DDP(model, device_ids=[local_rank]).
A DistributedSampler gives each rank a disjoint shard of the data. Call sampler.set_epoch(epoch) each epoch, or every epoch uses the same shuffle.
Save checkpoints, log and evaluate only on rank 0, and save model.module.state_dict().

The global batch is per-GPU batch times world size; scale the learning rate accordingly or retune it. Use nn.SyncBatchNorm.convert_sync_batchnorm(model) if small per-GPU batches hurt BatchNorm.

Prefer DDP over the older nn.DataParallel, which is single-process and slower. When the model or optimizer state no longer fits on one GPU, move on to FSDP.

Example

# Launch: torchrun --nproc_per_node=4 train.py
import os, torch, torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
def main():
    dist.init_process_group("nccl")
    rank, local_rank = dist.get_rank(), int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    model = DDP(nn.Linear(10, 1).cuda(), device_ids=[local_rank])
    ds = TensorDataset(torch.randn(1024, 10), torch.randn(1024, 1))
    sampler = DistributedSampler(ds)                  # disjoint shard per rank
    loader = DataLoader(ds, batch_size=32, sampler=sampler)
    opt = torch.optim.SGD(model.parameters(), lr=0.01)
    for epoch in range(3):
        sampler.set_epoch(epoch)                      # reshuffle each epoch
        for xb, yb in loader:
            loss = nn.functional.mse_loss(model(xb.cuda()), yb.cuda())
            opt.zero_grad(); loss.backward(); opt.step()   # backward all-reduces
    if rank == 0:
        torch.save(model.module.state_dict(), "model.pt")
    dist.destroy_process_group()
if __name__ == "__main__":
    main()

Remember

One process per GPU, launched by torchrun
DistributedSampler + set_epoch; save and log on rank 0 only
Gradients are averaged during backward, so replicas stay in sync

Watch out for

Forgetting sampler.set_epoch(epoch), so every epoch sees the same order.

Check yourself

With DDP on 8 GPUs and batch_size=32 in each DataLoader, what is the effective global batch size?

Official docs: pytorch.org/tutorials/intermediate/ddp_tutorial.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks