Use torch.distributed to initialize a process group per GPU, wrap the model in DistributedDataParallel (DDP), and use DistributedSampler so each process gets a shard of the dataset. Synchronization happens inside DDP during backward (all-reduce of gradients) and when using collective operations (barrier). Here's a minimal single-node, multi-process, multi-GPU skeleton:python
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = '127.0.0.1'
os.environ['MASTER_PORT'] = '29500'
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def cleanup():
dist.destroy_process_group()
def train(rank, world_size, dataset, model_fn, epochs=2, batch_size=32, lr=1e-3):
setup(rank, world_size)
device = torch.device(f'cuda:{rank}')
model = model_fn().to(device)
ddp_model = DDP(model, device_ids=[rank]) # sync happens on backward all-reduce
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=2, pin_memory=True)
optimizer = torch.optim.Adam(ddp_model.parameters(), lr=lr)
loss_fn = torch.nn.CrossEntropyLoss()
for epoch in range(epochs):
sampler.set_epoch(epoch) # shuffles differently each epoch
for batch in loader:
inputs, targets = batch
inputs = inputs.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)
optimizer.zero_grad()
outputs = ddp_model(inputs) # forward (no sync)
loss = loss_fn(outputs, targets)
loss.backward() # gradients synchronized across ranks here
optimizer.step() # updated locally with synchronized grads
# optional sync point
dist.barrier()
cleanup()
if __name__ == '__main__':
world_size = torch.cuda.device_count()
# provide dataset and model_fn according to your codebase
# mp.spawn(train, args=(world_size, dataset, model_fn), nprocs=world_size, join=True)
