Build an LLMAdvancedLesson 456 min read

Step 9 · Training across many GPUs

One model no longer fits on one card. Four ways to split it, and the rule for choosing between them.

🧪 Lab · 20 min📝 Lesson test
Lesson in motion

In 60 seconds

Step 9 · Training across many GPUs

One model no longer fits on one card. Four ways to split it, and the rule for choosing between them.

1/5
In simple words
One person cannot carry a piano. So you either give everyone their own small piano, or you cut one piano into pieces and each person carries a piece. Both work; they cost different amounts of shouting.
StrategyWhat gets splitCommunication costUse when
Data parallel (DDP)The batchGradient all-reduce each stepThe model fits on one GPU
FSDP / ZeRO-3Weights, gradients, optimiser stateGather and scatter each layerThe model does not fit — the default choice today
Tensor parallelIndividual matrices, within a layerEvery layer, twice — needs fast interconnectWithin one node with NVLink
Pipeline parallelLayers across devicesOnly at stage boundariesAcross nodes on slower links

Start here: FSDP

Fully Sharded Data Parallel is the practical default. Each GPU holds a slice of the weights, gradients and optimiser state. Before computing a layer it gathers the full weights for that layer, uses them, then throws the copy away. Memory per GPU drops close to linearly with device count.
FSDP in practicepython
import torch, functools
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

torch.distributed.init_process_group("nccl")
torch.cuda.set_device(local_rank)

policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={Block},          # shard at block granularity
)

model = FSDP(
    model,
    auto_wrap_policy=policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,     # ZeRO-3
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,          # reduce gradients in fp32
    ),
    device_id=torch.cuda.current_device(),
    limit_all_gathers=True,
)
# the training loop is otherwise unchanged

The rule for combining them

  1. 1

    Tensor parallel inside a node

    Splitting a matrix means talking twice per layer. Only viable over NVLink-class bandwidth. Typically 2–8 way.
  2. 2

    Pipeline parallel across nodes

    Only communicates at stage boundaries, so it tolerates slower links. Costs you a pipeline "bubble" — idle time while stages fill and drain.
  3. 3

    Data parallel on top of everything

    Replicate the whole tensor-and-pipeline arrangement, and all-reduce gradients between replicas.
  4. 4

    3D parallelism

    All three combined. This is how frontier runs are organised, and the configuration search is a real engineering job in itself.

What actually goes wrong at scale

  • Stragglers. Every GPU waits for the slowest. One throttling card slows the entire cluster.
  • Hardware failure. At thousands of GPUs running for weeks, failures are routine. Frequent checkpointing and automatic restart are mandatory, not nice-to-have.
  • Loss spikes. Large runs hit sudden instabilities. Standard response: roll back to the last good checkpoint, skip the offending data shard, resume.
  • Silent numerical drift. Different reduction orders give slightly different results across runs. Do not chase exact reproducibility across cluster sizes; chase stable statistics.
  • Data loader starvation. Frequently the bottleneck is not the GPUs at all. Pre-tokenise, shard, and use memory-mapped reads.
Do this
Measure Model FLOPs Utilisation — the fraction of theoretical peak FLOPs you actually achieve. Under 30% means something is badly wrong. 40–55% is a well-tuned run. That single number tells you more about your setup than anything else.
Watch out
Do not build this yourself for a first project. Use a mature framework and spend your effort on data and evaluation instead — that is where the quality actually comes from.

Watch and read more

Distributed training explainedYouTube searchPyTorch FSDP documentationpytorch.org · reading

Lab

A multi-GPU run with measured MFU.

~20 min

The problem

Run the same model on 1 GPU, then with FSDP on 2+. Measure tokens/second and MFU for each. Find the bottleneck — if MFU is under 30%, prove whether it is the data loader, communication, or the model.
Starter codepython
mfu = (6 * n_params * tokens_per_sec) / (peak_flops * n_gpus)
print(f"{tokens_per_sec:,.0f} tok/s  MFU {mfu:.1%}")
# Under 30%: time the data loader in isolation first. It is usually the loader.

You are done when

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Two GPUs give 1.4x throughput, not 2x. List the causes in the order you would check them.Reveal
One: data loader saturation — one loader now feeding two devices; time it in isolation first because it is the most common and the cheapest to fix. Two: gradient all-reduce over a slow interconnect — check whether communication overlaps with backward compute. Three: batch size too small per device, so kernels are launch-bound rather than compute-bound. Four: stragglers or thermal throttling on one card. Check in that order; the first two explain most cases.

Please sign in to continue.

Questions people ask

FSDP or DeepSpeed?

Both implement the same sharding ideas. FSDP is native to PyTorch and increasingly the default; DeepSpeed has a longer feature list and strong CPU-offload support. Either is fine — pick the one your team already knows.

How many GPUs do I need for a 7B model?

Full training needs roughly 100 GB of state, so at least two 80 GB cards with FSDP, and realistically 8–64 for a run that finishes in reasonable time. A LoRA fine-tune of the same model fits on one 24 GB card.

What is the pipeline bubble?

Idle time while the first micro-batch works its way through the stages and the last one drains out. More micro-batches shrink it. Interleaved schedules shrink it further at the cost of complexity.

Does gradient checkpointing help?

Yes — recompute activations during the backward pass instead of storing them. It typically trades about 30% extra compute for a large memory saving, which often lets you raise batch size enough to come out ahead.

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.

Sign in with your phone number to take the test and save your progress