The usual advice for scaling deep learning is to throw more GPUs at the problem. That skips the part where moving from one machine to a cluster changes how the training loop has to be written. A PyTorch DDP training pipeline only pays off if the process groups, the rank-aware logging and the sampler seeding are all handled properly. Get those wrong and the extra GPUs burn electricity in parallel while the loss curve looks the same as it did on one box.
What NCCL and all-reduce actually do
DDP is a communication pattern built on collective operations, specifically NCCL, the NVIDIA Collective Communications Library. Every process holds an identical copy of the model and sees a different slice of the data. The interesting part happens inside backward(). DDP registers hooks on every parameter, and as soon as a gradient is computed it fires an all-reduce across the process group. That averages the gradient over all ranks, which is what keeps the replicas in sync.
Wrapping the model in DDP and hoping for the best usually runs into the host memory bottleneck instead. Idle GPUs sitting next to a CPU pegged at 100% is a data loading problem, not a model problem. I went through that one in more detail in solving the host memory bottleneck.
Setting up the PyTorch DDP training pipeline
The distributed lifecycle runs in three phases: initialize, run, tear down. Miss one and the job either hangs with no output or leaves zombie processes behind. The initialize step reads the environment variables torchrun sets for you (RANK, LOCAL_RANK, WORLD_SIZE) and brings up the process group.
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def bbioon_setup_distributed():
# torchrun sets these variables automatically
rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
# Pin the correct GPU
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")
return rank, local_rank, world_sizeThe distributed sampler gotcha
This is where most custom pipelines break convergence, and they break it quietly. A DistributedSampler is what gives each GPU its own unique slice of the data, and it needs sampler.set_epoch(epoch) called at the start of every training epoch. Skip that call and the shuffling logic stays identical from one epoch to the next, so the model keeps seeing the same batches in the same order, generalizes worse, and costs you the same compute either way.
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
def bbioon_get_dataloader(dataset, config, rank, world_size):
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True
)
loader = DataLoader(
dataset,
batch_size=config.batch_size,
sampler=sampler,
num_workers=4,
pin_memory=True # Critical for async CPU-to-GPU transfers
)
return loader, samplerRank-aware checkpointing
Saving a checkpoint from every rank is a reliable way to corrupt one. Four processes write to checkpoint.pt at the same moment and the file is garbage. Guard the I/O so only rank 0 touches the disk, then park every other rank at a dist.barrier() before they try to load that state back in.
Serving the model once it is trained is a separate problem with its own latency budget. There is a companion post here on building FastAPI APIs for low-latency inference.
If this pipeline work is eating your dev hours, I can take it off your plate. I have been wrestling with WordPress and high-performance backend infrastructure since the 4.x days.
The short version
A production PyTorch DDP training pipeline is mostly plumbing, not modelling. Get the distributed lifecycle right, keep pin_memory=True for throughput, and call set_epoch() every epoch. Once you scale past a single node, hand process management to torchrun so it deals with the messy environment variable injection instead of you.
