Distributed (DDP)

# pip install torch torchvision
# Run with torchrun: torchrun --nproc_per_node=4 train_ddp.py
import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
from torchvision.models import resnet18

# 1) Initialise the process group (one process per GPU)
def setup():
    dist.init_process_group(backend='nccl')             # 'gloo' for CPU; 'nccl' for GPU
    rank = int(os.environ['RANK'])
    local_rank = int(os.environ['LOCAL_RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    torch.cuda.set_device(local_rank)
    return rank, local_rank, world_size

def cleanup():
    dist.destroy_process_group()

# 2) Training loop
def main():
    rank, local_rank, world_size = setup()
    device = torch.device(f'cuda:{local_rank}')

    # Model
    model = resnet18(num_classes=10).to(device)
    model = DDP(model, device_ids=[local_rank], output_device=local_rank)

    # Optimizer + loss
    opt   = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()

    # Data with DistributedSampler
    transform = transforms.Compose([transforms.ToTensor()])
    ds = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(ds, num_replicas=world_size, rank=rank, shuffle=True)
    loader = DataLoader(ds, batch_size=128, sampler=sampler, num_workers=4, pin_memory=True)

    for epoch in range(10):
        sampler.set_epoch(epoch)                         # different shuffle each epoch
        model.train()
        for x, y in loader:
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            loss = loss_fn(model(x), y)
            loss.backward()                              # gradients all-reduced automatically
            opt.step()

        if rank == 0:                                     # save only on rank 0
            torch.save(model.module.state_dict(), f'ckpt-epoch{epoch}.pt')

    cleanup()

if __name__ == '__main__':
    main()

# 3) Launching
# Single-node, multi-GPU:
# torchrun --standalone --nnodes=1 --nproc_per_node=4 train_ddp.py

# Multi-node:
# torchrun --nnodes=2 --node_rank=0 --master_addr=10.0.0.1 --master_port=29500 \\
#          --nproc_per_node=4 train_ddp.py
# (On rank-1 node, use --node_rank=1)

# 4) DistributedSampler
# Each process sees a DIFFERENT subset of the data; gradients are then averaged.
# Without it, each GPU would train on the WHOLE dataset (effectively duplicating work).
# Call sampler.set_epoch(e) so shuffling differs each epoch.

# 5) Saving + loading
# • Save model.module.state_dict() — strips the 'module.' prefix DDP adds
# • Or save with 'module.' prefix and account for it on load:

state = torch.load('ckpt.pt', map_location='cpu')
state = { k.replace('module.', ''): v for k, v in state.items() }
model.load_state_dict(state)

# 6) Validation — only do on rank 0 to avoid duplicate work
if rank == 0:
    model.eval()
    with torch.no_grad():
        for x, y in val_loader:
            preds = model(x.to(device))
            # …

# 7) Mixed precision + DDP — works seamlessly
scaler = torch.cuda.amp.GradScaler()
for x, y in loader:
    x, y = x.to(device), y.to(device)
    opt.zero_grad(set_to_none=True)
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        loss = loss_fn(model(x), y)
    scaler.scale(loss).backward()
    scaler.step(opt)
    scaler.update()

# 8) Synchronising metrics across ranks
def reduce_tensor(tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    return rt / dist.get_world_size()

# Logging on rank 0 only
if rank == 0:
    avg_loss = reduce_tensor(loss).item()
    print(f'epoch {epoch} loss {avg_loss:.4f}')

# 9) DDP-friendly modules
# • Most PyTorch modules work out of the box
# • SyncBatchNorm — replace BatchNorm with sync version (averages stats across ranks)
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = DDP(model, device_ids=[local_rank])

# 10) FSDP — Fully Sharded Data Parallel (for huge models)
# When models don't fit on one GPU, FSDP shards parameters + optimiser state across ranks.
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model)

# Required for very large transformer training.

# 11) Lightning / Accelerate — wrappers
# pytorch-lightning + HuggingFace accelerate hide most DDP plumbing:
#   accelerate launch train.py
#   pl.Trainer(strategy='ddp', devices=4)
# Recommended for new projects.

# 12) Debugging
# • NCCL hangs — check matching world_size + master_addr + reachable port
# • CUDA out of memory — reduce batch_size, gradient_accumulation_steps, FSDP
# • Slow because data loading bottleneck — increase num_workers, pin_memory, prefetch
# • Loss diverges — check learning rate scaling: effective batch = batch * world_size; scale LR proportionally
# • Random initial weights differ across ranks — use seed broadcast or initialise on rank 0 + broadcast

# 13) Gradient accumulation across DDP
for i, (x, y) in enumerate(loader):
    with model.no_sync() if (i + 1) % acc_steps != 0 else nullcontext():
        loss = loss_fn(model(x.to(device)), y.to(device)) / acc_steps
        loss.backward()
    if (i + 1) % acc_steps == 0:
        opt.step()
        opt.zero_grad(set_to_none=True)

# no_sync() defers the all-reduce until the actual optimizer step → less network traffic.

# 14) Common bugs
# • Forgot torchrun → ranks not initialised; manual MASTER_ADDR setup
// • Saving from every rank → race condition or duplicate writes; save only on rank 0
# • Not calling sampler.set_epoch → same shuffle every epoch on each rank
# • BatchNorm in DDP without SyncBN → per-rank statistics; mismatched eval
# • Mismatched random seeds — different initial weights per rank → divergent training; broadcast initial weights
# • Loading checkpoint with 'module.' prefix into non-DDP model → KeyError; strip prefix
# • Validation duplicated across ranks → wasted compute; gate on rank 0
# • LR not scaled with batch — increase LR proportional to world_size (linear scaling rule)
# • Calling barrier() unnecessarily — slows training; remove unless required
# • Mixing CUDA and CPU tensors in cross-rank ops → silent device mismatch

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group('nccl')
model = DDP(model.to(rank), device_ids=[rank])
# Launch
# torchrun --nproc_per_node=4 train.py