PyTorch Distributed Training

Real-World Scenario

A team trains a large language model on a machine with 4 GPUs. On one GPU an epoch takes days; with DistributedDataParallel each GPU processes a different quarter of every epoch while gradients are averaged across them, so the model sees the same data in far less wall-clock time. Gradient checkpointing then frees enough activation memory to raise the batch size. Every piece of that setup can be written and tested on a CPU first.

The examples below run on a CPU with PyTorch’s gloo backend. On NVIDIA GPUs, the only changes are backend="nccl", torch.cuda.set_device(local_rank), and moving the model to that device before wrapping it in DDP.

DataParallel Is Legacy

nn.DataParallel(model) splits each batch across GPUs inside a single Python process. It needs one line of code, but that single process is bottlenecked by the GIL and copies the model to every GPU on every forward pass. PyTorch’s documentation recommends DistributedDataParallel even on a single machine, and the rest of this lesson uses it.

Processes, Ranks, and Collectives

torchrun starts one process per worker and tells each one who it is through environment variables. init_process_group reads them.

import os
import torch
import torch.distributed as dist

dist.init_process_group(backend="gloo")          # "nccl" on NVIDIA GPUs
rank, world = dist.get_rank(), dist.get_world_size()

t = torch.tensor([float(rank + 1)])
dist.all_reduce(t, op=dist.ReduceOp.SUM)         # every rank ends with the sum

info = f"rank {rank}/{world} local_rank={os.environ['LOCAL_RANK']} sum={t.item()}"
gathered = [None] * world
dist.all_gather_object(gathered, info)           # collect strings on every rank
if rank == 0:                                    # print once, in rank order
    print("\n".join(gathered))
dist.destroy_process_group()

$ torchrun --standalone --nproc_per_node=2 hello_ddp.py 2>/dev/null
rank 0/2 local_rank=0 sum=3.0
rank 1/2 local_rank=1 sum=3.0

Each process has a rank — 0 and 1 here — and all of them know the world_size. all_reduce is a collective: every rank contributes a value and every rank receives the result, so both see 1 + 2 = 3. Gathering the strings and printing only on rank 0 keeps the output in order; if every rank printed, the lines would interleave unpredictably. 2>/dev/null hides the launcher’s log lines, including its notice that it set OMP_NUM_THREADS=1 for each process to avoid oversubscribing the CPU. It also hides tracebacks, so drop it whenever a run fails or prints nothing.

DistributedSampler: One Shard per Rank

import torch
from torch.utils.data import DistributedSampler

data = list(range(10))

for rank in range(3):                             # simulate 3 processes, no launcher needed
    s = DistributedSampler(data, num_replicas=3, rank=rank, shuffle=False)
    print(f"rank {rank}: {list(s)}")

s = DistributedSampler(data, num_replicas=2, rank=0, shuffle=True, seed=0)
print("epoch 0:", list(s))
print("epoch 1:", list(s), "  <- same order: set_epoch not called")
s.set_epoch(1)
print("epoch 1:", list(s), "  <- after set_epoch(1)")

$ python sampler.py
rank 0: [0, 3, 6, 9]
rank 1: [1, 4, 7, 0]
rank 2: [2, 5, 8, 1]
epoch 0: [4, 7, 3, 0, 6]
epoch 1: [4, 7, 3, 0, 6]   <- same order: set_epoch not called
epoch 1: [5, 1, 0, 9, 7]   <- after set_epoch(1)

DistributedSampler gives each rank every world_size-th index, so ranks see disjoint data. Ten items don’t split evenly over three ranks, so the sampler pads by repeating the first indices: 0 and 1 each appear twice. With shuffling, the order is derived from seed + epoch — and the epoch is only what you tell it. Without sampler.set_epoch(epoch) every epoch replays the same order; calling it at the start of each epoch is the fix.

Training with DistributedDataParallel

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset

dist.init_process_group(backend="gloo")           # "nccl" + device_ids=[local_rank] on GPUs
rank = dist.get_rank()

torch.manual_seed(0)                              # same data on every rank
X = torch.randn(4096, 32)
y = (X[:, :4].sum(dim=1) > 0).long()
dataset = TensorDataset(X, y)
sampler = DistributedSampler(dataset, shuffle=True, seed=0)
loader = DataLoader(dataset, batch_size=64, sampler=sampler)

torch.manual_seed(rank)                           # deliberately different init per rank
model = DDP(nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 2)))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

for epoch in range(3):
    sampler.set_epoch(epoch)
    for xb, yb in loader:
        opt.zero_grad()
        loss = nn.functional.cross_entropy(model(xb), yb)
        loss.backward()                           # DDP all-reduces gradients here
        opt.step()
    if rank == 0:
        print(f"epoch {epoch}: {len(loader)} batches per rank, loss {loss.item():.4f}")

checksum = torch.tensor([sum(p.sum().item() for p in model.parameters())])
sums = [torch.zeros(1) for _ in range(dist.get_world_size())]
dist.all_gather(sums, checksum)
if rank == 0:
    print("parameter checksums per rank:", [round(s.item(), 6) for s in sums])
    torch.save(model.module.state_dict(), "ddp_model.pt")   # .module = unwrapped model
dist.destroy_process_group()

$ torchrun --standalone --nproc_per_node=2 train_ddp.py 2>/dev/null
epoch 0: 32 batches per rank, loss 0.5992
epoch 1: 32 batches per rank, loss 0.4699
epoch 2: 32 batches per rank, loss 0.3437
parameter checksums per rank: [-0.553344, -0.553344]

Each rank loads the full dataset but trains only on its sampler’s half, so an epoch is 32 batches of 64 per rank instead of 64 batches on one process. The two ranks deliberately initialized their models with different seeds, yet their parameter checksums are identical at the end. Wrapping in DDP broadcasts rank 0’s weights to every rank, and during every backward() DDP averages the gradients across ranks, so each optimizer takes the same step on the same weights. Save from rank 0 only, and save model.module — the model inside the wrapper — so the checkpoint’s keys don’t carry a module. prefix.

Computing Metrics Across Ranks

Each rank only sees its own shard, so evaluation metrics have to be combined — and averaging them is not always right.

import torch
import torch.distributed as dist

dist.init_process_group(backend="gloo")
rank, world = dist.get_rank(), dist.get_world_size()

# uneven shards: rank 0 scored 90 of 100, rank 1 scored 10 of 20
correct, total = [(90, 100), (10, 20)][rank]

acc = torch.tensor([correct / total])
dist.all_reduce(acc)                              # sum of the per-rank accuracies
counts = torch.tensor([correct, total], dtype=torch.float64)
dist.all_reduce(counts)                           # sum of the raw counts

if rank == 0:
    print(f"mean of per-rank accuracies: {acc.item() / world:.4f}")
    print(f"global accuracy from counts: {counts[0] / counts[1]:.4f}")
dist.destroy_process_group()

$ torchrun --standalone --nproc_per_node=2 metrics.py 2>/dev/null
mean of per-rank accuracies: 0.7000
global accuracy from counts: 0.8333

Rank 0 was right on 90 of 100 examples and rank 1 on 10 of 20. Averaging the two accuracies weights them equally and reports 70%, but the true accuracy is 100 / 120 = 83.3%. All-reducing the raw counts — correct predictions and totals — and dividing at the end is correct whatever the shard sizes, and per-rank counts differ more often than you’d expect: filtering, a final partial batch, or a sampler run without padding all cause it.

Gradient Checkpointing

Activations stored for the backward pass often use more memory than the weights. Checkpointing keeps only each block’s input and recomputes the inside during backward.

import time
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

def block():
    return nn.Sequential(nn.Linear(256, 1024), nn.GELU(), nn.Linear(1024, 256), nn.LayerNorm(256))

torch.manual_seed(0)
blocks = nn.ModuleList(block() for _ in range(12))
x = torch.randn(8, 128, 256)

def run(use_checkpoint):
    saved = 0
    def pack(t):                                   # called for every tensor kept for backward
        nonlocal saved
        saved += t.numel() * t.element_size()
        return t
    t0 = time.perf_counter()
    with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
        h = x
        for b in blocks:
            h = checkpoint(b, h, use_reentrant=False) if use_checkpoint else b(h)
    h.sum().backward()
    return saved / 2**20, time.perf_counter() - t0

run(False), run(True)                            # warm-up: first calls are slower
for flag in (False, True):
    mb, secs = run(flag)
    print(f"checkpointing={flag!s:5}  saved for backward: {mb:6.1f} MB  time: {secs:.2f} s")

$ python grad_ckpt.py
checkpointing=False  saved for backward:  143.1 MB  time: 0.14 s
checkpointing=True   saved for backward:   12.0 MB  time: 0.23 s

saved_tensors_hooks calls pack for every tensor autograd keeps for the backward pass, which gives an exact measure of what the forward pass holds on to. Without checkpointing, the twelve blocks keep 143 MB; with it, they keep only each block’s 1 MB input — 12 MB — and rebuild the rest one block at a time during backward. The price is that recomputation: here the step took about 60% longer. use_reentrant=False selects the implementation PyTorch recommends; the older reentrant variant has restrictions around inputs that don’t require gradients. The warm-up calls keep one-time setup costs out of the timing, since the first pass through each variant is slower than later ones.

Practice

from torch.utils.data import DistributedSampler

data = list(range(10))
for rank in range(3):
    s = DistributedSampler(data, num_replicas=3, rank=rank, shuffle=False, drop_last=True)
    print(f"rank {rank}: {list(s)}")

$ python p1_drop_last.py
rank 0: [0, 3, 6]
rank 1: [1, 4, 7]
rank 2: [2, 5, 8]

import torch
import torch.distributed as dist

dist.init_process_group(backend="gloo")
rank, world = dist.get_rank(), dist.get_world_size()

local = [0.42, 3.9][rank]                         # rank 1 has a diverging batch
loss = torch.tensor([local])
dist.all_reduce(loss, op=dist.ReduceOp.MAX)       # every rank now holds the maximum
skip = loss.item() > 2.0

decisions = [None] * world
dist.all_gather_object(decisions, (rank, local, skip))
if rank == 0:
    for r, l, s in decisions:
        print(f"rank {r}: local loss {l:.2f}  max {loss.item():.2f}  skip step: {s}")
dist.destroy_process_group()

$ torchrun --standalone --nproc_per_node=2 p2_max_loss.py 2>/dev/null
rank 0: local loss 0.42  max 3.90  skip step: True
rank 1: local loss 3.90  max 3.90  skip step: True

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="gloo")
rank, world = dist.get_rank(), dist.get_world_size()
model = DDP(nn.Linear(4, 1))

def grads_agree():
    g = model.module.weight.grad.clone()
    all_g = [torch.zeros_like(g) for _ in range(world)]
    dist.all_gather(all_g, g)
    return all(torch.equal(all_g[0], other) for other in all_g)

torch.manual_seed(rank)                           # each rank gets different data
with model.no_sync():                             # micro-batch 1: no communication
    model(torch.randn(8, 4)).sum().backward()
first = grads_agree()
model(torch.randn(8, 4)).sum().backward()         # micro-batch 2: all-reduce happens
second = grads_agree()                            # collectives must run on every rank
if rank == 0:
    print("after no_sync backward:  grads agree =", first)
    print("after normal backward:   grads agree =", second)
dist.destroy_process_group()

$ torchrun --standalone --nproc_per_node=2 p3_no_sync.py 2>/dev/null
after no_sync backward:  grads agree = False
after normal backward:   grads agree = True

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

torch.manual_seed(0)
blocks = nn.ModuleList(
    nn.Sequential(nn.Linear(256, 1024), nn.GELU(), nn.Linear(1024, 256), nn.LayerNorm(256))
    for _ in range(12))
x = torch.randn(8, 128, 256)

saved = 0
def pack(t):
    global saved
    saved += t.numel() * t.element_size()
    return t

with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t):
    h = x
    for i, b in enumerate(blocks):
        h = checkpoint(b, h, use_reentrant=False) if i % 2 == 0 else b(h)
h.sum().backward()
print(f"saved for backward: {saved / 2**20:.1f} MB")

$ python p4_every_other.py
saved for backward: 78.1 MB