DistributedDataParallel — Multi-GPU Training

The standard for multi-GPU

DistributedDataParallel (DDP) is the way to do multi-GPU training in PyTorch. The mental model: one process per GPU, each process has its own copy of the model and runs its own forward/backward on its slice of the data. After backward, the gradients are synchronized across processes (all-reduce); each optimizer step then applies the same averaged gradient to its local copy of the weights, keeping them in lockstep.

What you have to add to a single-GPU script

Initialize a process group (dist.init_process_group("nccl")).
Set the per-process device (torch.cuda.set_device(rank)).
Wrap the model with DDP (model = DDP(model, device_ids=[rank])).
Use a DistributedSampler so each process sees a different slice of the data.
Launch with torchrun (the modern launcher, not the old torch.distributed.launch).

The DistributedSampler gotcha

Without DistributedSampler, every process iterates the full dataset — you waste compute on redundant work. set_epoch(epoch) on the sampler at the top of each epoch is mandatory; without it, the shuffle is the same every epoch and your training plateaus mysteriously.

torchrun — the modern launcher

torchrun --nproc_per_node=4 train.py spawns four processes, each with environment variables (LOCAL_RANK, WORLD_SIZE, RANK) set so your script can self-identify. The old torch.distributed.launch still works but torchrun is preferred.

Code

DDP training script — the canonical shape·python

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

def setup():
    dist.init_process_group("nccl")
    rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(rank)
    return rank

def train(model, dataset, num_epochs):
    rank = setup()
    model = model.to(rank)
    model = DDP(model, device_ids=[rank])

    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler,
                        num_workers=4, pin_memory=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)               # MANDATORY for shuffle
        for x, y in loader:
            x, y = x.to(rank), y.to(rank)
            optimizer.zero_grad()
            loss = criterion(model(x), y)
            loss.backward()                    # gradients sync here (all-reduce)
            optimizer.step()

    dist.destroy_process_group()

if __name__ == "__main__":
    model = MyModel()
    dataset = MyDataset()
    train(model, dataset, num_epochs=10)

Launch with torchrun·python

# Single-node, 4 GPUs
# torchrun --nproc_per_node=4 train.py

# Single-node, 8 GPUs, with extra args to your script
# torchrun --nproc_per_node=8 train.py --lr 1e-4 --batch_size 32

# Multi-node — each node runs its own torchrun
# Node 0:
#   torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
#            --master_addr=NODE0_IP --master_port=29500 train.py
# Node 1:
#   torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 \
#            --master_addr=NODE0_IP --master_port=29500 train.py

import os
print(os.environ['LOCAL_RANK'], os.environ['RANK'], os.environ['WORLD_SIZE'])

rank-0-only logging — the universal pattern·python

import os
import torch.distributed as dist

def is_main():
    return not dist.is_initialized() or dist.get_rank() == 0

def log(*args, **kwargs):
    if is_main():
        print(*args, **kwargs)

# Use throughout your training
log(f"epoch {epoch}: train_loss={train_loss:.4f}")

# Same trick for checkpoint saving — only rank 0 writes
if is_main():
    torch.save(model.module.state_dict(), 'best.pt')   # .module to unwrap DDP

# And for any setup that should happen ONCE, not N-times
if is_main():
    os.makedirs('outputs', exist_ok=True)

Exercise

Convert your single-GPU training script to DDP. Run it via torchrun --nproc_per_node=N (N=GPU count). Verify: (a) loss prints once per epoch, not N times — that's the is_main() rank-0 trick. (b) The training is genuinely faster with N>1 GPUs (not just N times the cost). (c) The shuffle differs across epochs — sampler.set_epoch is correctly placed.