PyTorch Distributed Training
Distributed training splits a model's training workload across multiple GPUs or multiple machines. It reduces training time from days to hours for large models and datasets. PyTorch provides built-in tools — DataParallel and DistributedDataParallel — that handle the complexity of splitting data, syncing gradients, and aggregating results.
Why Distributed Training Is Needed
Single GPU: 16GB VRAM, batch_size=32 Training ResNet50 on ImageNet: ~5 days 4 × GPUs (DataParallel): Effective batch_size=128 Training time: ~30 hours 8 × GPUs across 2 machines (DDP): Effective batch_size=256 Training time: ~10 hours Scaling compute reduces training time proportionally.
Two Approaches in PyTorch
DataParallel (DP): Single machine, multiple GPUs Easiest to set up — one line of code Less efficient: bottleneck on the master GPU Uses one Python process DistributedDataParallel (DDP): Single or multiple machines, multiple GPUs More complex setup More efficient: each GPU has its own process Industry standard for serious training
DataParallel — Quick Multi-GPU
Wrap your model in nn.DataParallel and PyTorch replicates the model on all available GPUs. Each GPU processes a chunk of the batch in parallel, and results are gathered back on GPU 0.
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(1024, 512),
nn.ReLU(),
nn.Linear(512, 10)
)
# Wrap for multi-GPU (uses all available GPUs by default)
if torch.cuda.device_count() > 1:
print(f"Using {torch.cuda.device_count()} GPUs")
model = nn.DataParallel(model)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
x = torch.rand(128, 1024).to(device)
output = model(x)
print(output.shape) # torch.Size([128, 10])DataParallel — How the Batch Is Split
batch_size = 128, using 4 GPUs:
GPU 0: processes samples 0– 31 (32 samples)
GPU 1: processes samples 32– 63 (32 samples)
GPU 2: processes samples 64– 95 (32 samples)
GPU 3: processes samples 96–127 (32 samples)
↓ all results gathered to GPU 0
GPU 0: computes loss, backward, optimizer step
Bottleneck: GPU 0 does extra work (gathering + optimizer).
DistributedDataParallel — Production Standard
DDP launches one process per GPU. Each process trains on its own data shard and its own model copy. After every backward pass, DDP uses an all-reduce operation to average gradients across all processes — so all models stay identical. This approach eliminates the GPU 0 bottleneck and scales linearly.
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import os
def train(rank, world_size):
# Initialize process group
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
dist.init_process_group('nccl', rank=rank, world_size=world_size)
# Create model on this GPU
model = nn.Linear(10, 2).to(rank)
model = DDP(model, device_ids=[rank]) # wrap with DDP
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss()
for epoch in range(10):
x = torch.rand(32, 10).to(rank)
y = torch.rand(32, 2).to(rank)
optimizer.zero_grad()
loss = loss_fn(model(x), y)
loss.backward() # DDP auto-syncs gradients here
optimizer.step()
if rank == 0: # only print from main process
print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
dist.destroy_process_group()
# Launch with torch.multiprocessing
if __name__ == '__main__':
world_size = torch.cuda.device_count()
torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size)DDP Data Splitting with DistributedSampler
Each process must receive a different, non-overlapping subset of the training data. The DistributedSampler handles this automatically.
from torch.utils.data import DataLoader, DistributedSampler
sampler = DistributedSampler(
dataset,
num_replicas=world_size, # total number of processes
rank=rank # this process's rank
)
loader = DataLoader(
dataset,
batch_size=32,
sampler=sampler, # replaces shuffle=True
num_workers=4,
pin_memory=True
)
# At each epoch, reshuffle:
for epoch in range(num_epochs):
sampler.set_epoch(epoch) # ensures different shuffle per epoch
for batch in loader:
... # training stepDDP Gradient Synchronization Diagram
Process 0 (GPU 0): Process 1 (GPU 1):
Forward pass Forward pass
Backward → local grads Backward → local grads
│ │
└─────── All-Reduce ────────────┘
(average grads across GPUs)
│ │
Optimizer.step() Optimizer.step()
(identical update) (identical update)
Saving and Loading with DDP
# Save only from rank 0 to avoid file conflicts
if rank == 0:
torch.save(model.module.state_dict(), 'model.pth')
# model.module gives the underlying model (not the DDP wrapper)
# Load on any rank
model = MyModel()
model.load_state_dict(torch.load('model.pth', map_location=f'cuda:{rank}'))
model = DDP(model, device_ids=[rank])Launching DDP with torchrun
The simplest way to launch a DDP training script across multiple GPUs is with the torchrun command.
# 4 GPUs on one machine:
torchrun --nproc_per_node=4 train.py
# 2 machines, 4 GPUs each (8 total):
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \
--master_addr=192.168.1.1 --master_port=12355 train.py
Summary
Distributed training reduces training time by running computation across multiple GPUs or machines in parallel. nn.DataParallel is the quickest multi-GPU setup but bottlenecks on the master GPU. DistributedDataParallel is the production standard — one process per GPU, gradients averaged via all-reduce, scales efficiently. Use DistributedSampler to ensure each process sees a different data shard. Save checkpoints only from rank 0. Launch with torchrun for the simplest multi-GPU and multi-node setup. DDP is the approach used in training every large-scale model in industry today.