Distributed Training

Data parallelism is the most common way to scale. You put a full copy of the model on every GPU (each copy is a worker), hand each worker a different slice of the batch, and let them all compute gradients at the same time.

The catch: if every worker just applied its own gradient, the copies would drift apart. So after the backward pass, the workers run an all-reduce — they sum every gradient across all workers and divide by the number of workers. Each GPU then applies that same averaged gradient, so all copies stay identical. The example below does the all-reduce by hand so you can see the math.

# === Data Parallelism: averaging gradients across N workers ===
# In Data Parallel training, every worker (GPU) holds a FULL copy of the
# model. Each worker sees a DIFFERENT slice of the batch, computes its own
# gradients, and then all workers AVERAGE their gradients so every copy
# applies the SAME update and stays identical. That averaging step is the
# "all-reduce".

# Pretend each worker computed a gradient for the same 3 parameters.
# (Real gradients come from backprop; here we hand-pick
...

There are two ways workers can coordinate. In synchronous training, every worker finishes its step and waits at a barrier for the all-reduce before anyone moves on — all copies stay identical, but the slowest worker (a "straggler") holds everyone up.

In asynchronous training, workers push gradients to a central parameter server and grab the latest weights without waiting. It avoids stragglers but workers can train on slightly stale weights, which can hurt accuracy. Modern deep learning almost always uses synchronous all-reduce — it's simpler and more accurate, and fast GPU interconnects (like NVLink) make the waiting cheap.

Big batches train more smoothly, but they don't always fit in GPU memory. Gradient accumulation is the workaround: run several small micro-batches, add up their gradients, and only call the optimizer once at the end. Four micro-batches of 8 give the same effective batch of 32 as one big batch — you trade extra time for less memory.

# === Gradient accumulation: simulate a big batch on a small GPU ===
# Goal: train as if the batch were 32, but your GPU can only hold 8 at a time.
# Trick: run 4 "micro-batches" of 8, ADD UP their gradients, and only update
# the weights once. 4 micro-batches x 8 = an EFFECTIVE batch of 32.

micro_batch_size = 8
accum_steps = 4                       # how many micro-batches to add before updating
effective_batch = micro_batch_size * accum_steps
print(f"Micro-batch: {micro_batch_size}  Accum ste
...

Notice you divide the accumulated gradient by the number of micro-batches so the update has the same scale as a real batch-of-32 step. The same idea stacks with data parallelism: effective batch = micro_batch × accum_steps × num_gpus.

Data parallelism assumes the whole model fits on one GPU. When it doesn't — say a 70B-parameter model that needs hundreds of gigabytes — you have to split the model itself.

A rough memory rule for training in FP16: you need about 4× the model size (weights + gradients + optimizer states + activations). A 7B model is ~14 GB in FP16, so ~56 GB to train — close to filling one 80 GB A100 before you've added the batch.

Mixed precision stores most numbers in 16-bit instead of 32-bit. That halves memory and roughly doubles throughput on modern GPUs. Two 16-bit formats matter:

• FP16 — half precision, fast, but a narrow range, so it needs loss scaling to stop tiny gradients underflowing to zero.
• BF16 (bfloat16) — same wide range as FP32 with less mantissa, so it trains stably without loss scaling. It's the default on newer hardware (A100/H100/TPU).
• FP32 — full precision, kept for the master copy of weights and the optimizer step for stability.

Scaling out (more GPUs, more accumulation) grows the effective batch. Larger batches produce smoother, lower-variance gradients, so you can — and must — take bigger steps. The linear scaling rule: when you multiply the batch by k, multiply the learning rate by k too, and add a few hundred steps of LR warmup so the bigger steps don't blow up early in training.

In practice you rarely write the all-reduce yourself — the framework does it. The most common entry point is PyTorch DistributedDataParallel (DDP): wrap your model in DDP(...) and it averages gradients automatically during loss.backward(). You launch one process per GPU with torchrun. Read this through — it ties together every idea from the lesson.

# === The real thing: PyTorch DistributedDataParallel (DDP) ===
# This is what production code looks like. You launch one process per GPU
# (e.g. "torchrun --nproc_per_node=4 train.py"). DDP does the all-reduce
# of gradients for you, automatically, during backward().

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

def main():
    dist.init_process_group(backend="nccl")          # set up GPU comms
    rank = dist.get_rank()                            # which process am I? (0..N-1)
    world_size = dist.get_world_size()                # how many processes total?
    torch.cuda.set_device(rank)

    model = nn.Linear(1024, 10).cuda(rank)
    model = DDP(model, device_ids=[rank])             # wrap -> auto gradient all-reduce

    # DistributedSampler gives each process a DIFFERENT, non-overlapping shard.
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    # Linear scaling rule: effective batch = 32 * world_size, so scale the LR.
    base_lr = 0.1
    opt = torch.optim.SGD(model.parameters(), lr=base_lr * world_size)

    for epoch in range(epochs):
        sampler.set_epoch(epoch)                      # reshuffle shards each epoch
        for x, y in loader:
            x, y = x.cuda(rank), y.cuda(rank)
            opt.zero_grad()
            loss = loss_fn(model(x), y)
            loss.backward()                           # <-- DDP all-reduces grads HERE
            opt.step()                                # every GPU applies the SAME update

    dist.destroy_process_group()

# Expected output (4 GPUs, conceptually):
#   rank 0..3 each train on 1/4 of the data
#   gradients are averaged across all 4 every step (all-reduce)
#   all 4 model copies stay byte-for-byte identical

Other frameworks fill different gaps: FSDP (built into PyTorch) and DeepSpeed ZeRO shard the model when it no longer fits; Horovod is an older Uber library that adds ring all-reduce to TensorFlow/PyTorch and is still seen in multi-node clusters.

Distributed bugs are sneaky — the code runs but the model trains wrong. Here are the four that bite everyone.

opt = SGD(model.parameters(), lr=0.1)   # ❌ same LR on 8 GPUs -> trains far too slowly

opt = SGD(model.parameters(), lr=0.1 * world_size)   # ✅ scale LR by #GPUs

model = nn.Linear(1024, 10).cuda()
ddp_model = DDP(model, device_ids=[rank])
loss = loss_fn(model(x), y)   # ❌ used the UN-wrapped model -> no all-reduce

loss = loss_fn(ddp_model(x), y)   # ✅ backward() now all-reduces grads

loader = DataLoader(dataset, batch_size=32)   # ❌ every GPU sees ALL the data

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
# remember sampler.set_epoch(epoch) so shuffling differs each epoch

# Symptom: 8 GPUs give only ~3x speedup, GPUs sit at 40% utilisation

# - use a larger per-GPU batch (more compute per all-reduce)
# - use gradient accumulation to all-reduce less often
# - use NCCL + fast interconnect (NVLink/InfiniBand), not Ethernet
# - keep grads in FP16/BF16 so there's half as much to send