Multi-GPU: DDP & FSDP

At some point one GPU isn’t enough — either training is too slow, or the model simply doesn’t fit in memory. Multi-GPU training is now a baseline skill, and the whole subject comes down to one question you can answer with arithmetic: does the model’s training state fit on a single GPU? The answer picks your strategy.

First, the memory math

People underestimate training memory because they only count the weights. Training a model with the Adam optimizer in mixed precision costs roughly 16 bytes per parameter:

bf16 weights        2 bytes
bf16 gradients      2 bytes
fp32 master weights 4 bytes   ← the "master copy" for stable updates
Adam moment  m      4 bytes
Adam moment  v      4 bytes
                   --------
                   16 bytes / parameter

So a 7B-parameter model needs about 7e9 × 16 = 112 GB just for model state — already over an 80GB GPU, before activations. That single number drives the whole decision. Try it:

DDP — replicate the model, split the batch

Distributed Data Parallel is the workhorse when the model does fit. Every GPU holds a full copy of the model; you split each batch across the GPUs, each computes gradients on its shard, and then an all-reduce averages the gradients so every replica stays in sync before the optimizer step.

DDP gives near-linear speedup but no memory savings — each GPU still holds the entire model state. It’s the right tool when the model fits and you just want to train faster on a bigger effective batch.

# torchrun --nproc_per_node=4 train.py
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model.to(local_rank), device_ids=[local_rank])
# ... the normal training loop; DDP all-reduces gradients during backward()

FSDP — shard the model when it won’t fit

When the full state exceeds one GPU, Fully Sharded Data Parallel splits the parameters, gradients, and optimizer states across the GPUs. Each GPU stores only its 1/N shard and gathers the full weights for a layer only momentarily, during that layer’s forward/backward, then frees them. Per-GPU memory drops roughly linearly with the number of GPUs — which is exactly what lets you train a 70B model that could never fit on one card.

Quick check

Next

That completes the training half of the track. Next we move to the architecture that everything modern is built on — starting with how text becomes tensors in tokenization, then self-attention and the transformer block.