PyTorch Distributed Training

Real-World Scenario

A team trains a 400M parameter language model on 4 A100 GPUs. Single-GPU training would take 3 weeks. With DDP across 4 GPUs, it trains in 5 days. Gradient checkpointing allows batch_size=16 instead of batch_size=4, further improving GPU utilization and convergence speed.

Single-Process Multi-GPU with DataParallel (Legacy)

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# Simple model for demonstration
class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(512, 2048), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(2048, 2048), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(2048, 512), nn.ReLU(),
            nn.Linear(512, 10),
        )
    def forward(self, x):
        return self.net(x)

# DataParallel — do NOT use for new code, shown for context only
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(LargeModel())    # wraps the model
    # model now splits each batch across all GPUs automatically
    print(f"Using {torch.cuda.device_count()} GPUs (DataParallel — legacy)")
else:
    model = LargeModel()
    print("Single GPU or CPU")

DistributedDataParallel (Recommended)

# train_ddp.py — save this as a standalone script and run with torchrun
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, TensorDataset, DistributedSampler
from torch.amp import autocast, GradScaler


def setup_ddp():
    """Initialize the process group. Called automatically by torchrun."""
    dist.init_process_group(backend="nccl")   # NCCL for GPU, gloo for CPU
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))


def cleanup_ddp():
    dist.destroy_process_group()


class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(512, 2048), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(2048, 2048), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(2048, 512), nn.ReLU(),
            nn.Linear(512, 10),
        )
    def forward(self, x):
        return self.net(x)


def train(rank: int, world_size: int, epochs: int = 5):
    """Train function — each process calls this with its rank."""
    # rank = local GPU index (set by torchrun via LOCAL_RANK env var)
    device = torch.device(f"cuda:{rank}")

    # Wrap model in DDP
    model = LargeModel().to(device)
    model = DDP(model, device_ids=[rank])  # all-reduce gradients across processes

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()
    scaler    = GradScaler(device="cuda")

    # Dataset — DistributedSampler ensures each GPU sees a unique subset
    X = torch.randn(10_000, 512)
    y = torch.randint(0, 10, (10_000,))
    dataset = TensorDataset(X, y)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    loader  = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=4, pin_memory=True)

    for epoch in range(epochs):
        sampler.set_epoch(epoch)   # ensures different shuffling each epoch
        model.train()
        total_loss = 0

        for X_batch, y_batch in loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            optimizer.zero_grad()

            with autocast(device_type="cuda"):
                loss = criterion(model(X_batch), y_batch)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            total_loss += loss.item()

        # Only print from rank 0 to avoid duplicate logs
        if rank == 0:
            avg_loss = total_loss / len(loader)
            print(f"Epoch {epoch+1}: loss={avg_loss:.4f}")

    # Save checkpoint only from rank 0
    if rank == 0:
        torch.save(model.module.state_dict(), "checkpoint_ddp.pt")
        print("Checkpoint saved.")


# Entry point for torchrun
if __name__ == "__main__":
    setup_ddp()
    rank       = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    train(rank, world_size)
    cleanup_ddp()

# Launch command:
# torchrun --nproc_per_node=4 train_ddp.py

Gradient Checkpointing

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint, checkpoint_sequential

class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_model * 4), nn.GELU(),
            nn.Linear(d_model * 4, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        x = self.norm2(x + self.ff(x))
        return x


class LargeTransformer(nn.Module):
    def __init__(self, n_layers: int = 24, d_model: int = 768, n_heads: int = 12):
        super().__init__()
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads) for _ in range(n_layers)
        ])
        self.use_checkpointing = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            if self.use_checkpointing and self.training:
                # checkpoint re-computes activations during backward instead of storing them
                x = checkpoint(block, x, use_reentrant=False)
            else:
                x = block(x)
        return x

    def enable_gradient_checkpointing(self):
        self.use_checkpointing = True
        print("Gradient checkpointing enabled")


# Compare memory usage
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = LargeTransformer(n_layers=12, d_model=256, n_heads=8).to(device)

x = torch.randn(4, 32, 256, device=device, requires_grad=True)

# Without checkpointing
if torch.cuda.is_available():
    torch.cuda.reset_peak_memory_stats()
out = model(x)
out.sum().backward()
if torch.cuda.is_available():
    no_ckpt_mem = torch.cuda.max_memory_allocated() / 1024**2
    print(f"Without checkpointing: {no_ckpt_mem:.1f} MB")

# With checkpointing
model.enable_gradient_checkpointing()
if torch.cuda.is_available():
    torch.cuda.reset_peak_memory_stats()
x = torch.randn(4, 32, 256, device=device, requires_grad=True)
out = model(x)
out.sum().backward()
if torch.cuda.is_available():
    ckpt_mem = torch.cuda.max_memory_allocated() / 1024**2
    print(f"With checkpointing:    {ckpt_mem:.1f} MB")
    print(f"Memory reduction:      {(1 - ckpt_mem/no_ckpt_mem):.0%}")

Synchronizing Metrics Across GPUs

import torch
import torch.distributed as dist
import numpy as np

def reduce_metric(value: float, device: torch.device, op: str = "mean") -> float:
    """Average a metric across all DDP processes."""
    if not dist.is_initialized():
        return value

    tensor = torch.tensor(value, device=device)
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    if op == "mean":
        tensor /= dist.get_world_size()
    return tensor.item()


def gather_predictions(
    preds: torch.Tensor,
    targets: torch.Tensor,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Gather predictions and targets from all processes for evaluation."""
    if not dist.is_initialized():
        return preds, targets

    world_size = dist.get_world_size()

    # Gather all predictions
    all_preds   = [torch.zeros_like(preds)   for _ in range(world_size)]
    all_targets = [torch.zeros_like(targets) for _ in range(world_size)]
    dist.all_gather(all_preds,   preds)
    dist.all_gather(all_targets, targets)

    return torch.cat(all_preds), torch.cat(all_targets)


def evaluate_ddp(model: nn.Module, loader, device: torch.device, rank: int) -> dict:
    """Evaluation that works correctly under DDP."""
    model.eval()
    all_preds, all_targets = [], []

    with torch.no_grad():
        for X, y in loader:
            X, y = X.to(device), y.to(device)
            preds = model(X).argmax(dim=1)
            all_preds.append(preds)
            all_targets.append(y)

    preds   = torch.cat(all_preds)
    targets = torch.cat(all_targets)

    # Gather across all GPUs
    preds, targets = gather_predictions(preds, targets, device)

    if rank == 0:   # compute metrics only on rank 0
        accuracy = (preds == targets).float().mean().item()
        return {"accuracy": accuracy}
    return {}