One GPU
Where we're going
Training a modern AI model takes many GPUs. Here are 512 of them. By the end you'll know what each one is doing. We'll follow one, highlighted here: GPU 1.
In the drawing: A dark field of 512 small chips. One, GPU 1, lights up inside a ring; then the other 511 switch on in a ripple out from it.
Compute cores and memory
Here's GPU 1 up close. The block in the middle holds its compute cores: thousands of them, multiplying and adding in parallel. The towers beside it are memory. Whatever the cores are working on has to fit there.
In the drawing: The view dives into GPU 1 until it fills the picture. Its tiles light from the centre out, three memory blocks trace up on each side, and labels name compute and memory.
Open it up
Let's lift the top off and simplify what's inside so we can watch it work. The two walls become memory; they fill as the GPU stores more. The strip of light becomes the compute cores; wherever it lights up, work is happening.
In the drawing: The lid lifts off and fades. Inside, the memory towers stand as two walls marked with capacity ticks, and a dim strip of light runs between them, swelling once as the compute cores are named.
A model is numbers
An AI model is a huge set of numbers, called weights, arranged in layers. This one has four. Real models can stack dozens. Memory holds every layer, plus the extra room training needs. Watch the walls fill as each layer goes in.
In the drawing: Four slabs, the layers, drop into the box one at a time, and with each the walls fill a notch higher, to near the top tick. Above Layer 4, faint hairlines stand for the rest, under a caption: '4 drawn · real models stack dozens'.
Learning
The next token
The model reads text in small pieces called tokens, which can be whole words or parts of words. Its job is to predict the next token using only the tokens that came before it.
In the drawing: Under the GPU wait four trays of text, the batch. The top one reads 'Once upon a', then a blank, and its tokens begin to enter the GPU's door. A readout appears at the foot: forward · error · backward · update. The chart's title, 'Once upon a ____', appears beside the GPU.
The model's guess
The tokens become lists of numbers, called embeddings, and pass through every layer. That's the forward pass. The correct answer is the next token that actually appears in the text. The error measures how far the model's guess was from that answer: the less probability it assigned to 'time,' the larger the error.
In the drawing: An amber band climbs through the four layers. At the top, a bar for each possible next token, all about the same height. A pink bracket marks the error: how far 'time' falls short of certain.
Work out the gradients
The text said 'time.' Now the model works backward from the error, layer by layer. That's the backward pass. For every weight, it calculates a gradient: which way to move that weight so 'time' becomes a little more likely. For now, it only records the gradients. The weights haven't changed yet.
In the drawing: A pink ring travels back down through the layers, leaving a small tick on each: the gradients, written down but not applied.
One step
It does this for every token in the batch, treating each token as the answer to the tokens before it. Across all the texts in the batch, the gradients pile up and are averaged. Then every model weight is updated. That's one training step. The model is a little less wrong, not suddenly right.
In the drawing: The blank in the top tray fills, 'Once upon a time there was…', and a tick appears under every text in the batch. The trays then fade away, a step counter appears at the foot reading 0, and a loupe opens on Layer 2: 16 of its numbers as dots, shifting as the update lands. The counter reads 1.
Microbatches
A real batch is usually too big to fit in a GPU's memory at once. So it's split into microbatches. The GPU runs the forward pass, the error, and the backward pass on each one, adding up the gradients. Then it updates the weights, once.
In the drawing: Two smaller trays, the microbatches, land under the entry door. Each in turn rises in, climbs the layers as an amber band, and its correction comes back down as a pink ring, its gradient ticks piling on the last. Then one update: the ticks drain and every slab re-tones. The readout gains a loop: forward, error, backward, round again.
Train it
Step after step, its guesses get better on average. The animation below illustrates how a real model's predictions change. It doesn't store the training text like a file; what it learns lives in the values of its weights.
Train it to step 1,000.
or keep scrolling
In the drawing: Five full loops play back to back, an amber band up through the layers, a pink ring back down, every slab re-toned, while the step count runs 1 to 1,000. In the list, 'time' overtakes one word per tenfold until it leads, its bar nearly at the 'certain' line and its error a sliver, and the model fills the blank: 'Once upon a time'.
It's slow
The model can learn this way, but slowly. A single GPU is doing all the work. Twice as much text takes about twice as long.
In the drawing: A time bar fills under the step readout at the foot of the picture: one step's time, four segments, one per text, all of it on GPU 1.
Data parallelism
A new run
The texts in a single training step can be split up. So start a fresh run at step 0 and add a second GPU to take half of every batch.
In the drawing: The step count reads 0 again, and the four trays dim to outlines marked '+ more'. The time bar stays.
Add a second GPU
A second GPU starts empty, so the full model is copied onto it once before training begins. Now both GPUs hold identical copies of the model.
Copy the model onto a second GPU.
or keep scrolling
In the drawing: The view pulls back to make room beside GPU 1, an empty socket is traced there, and GPU 2 arrives in it, empty; the counter reads 2 GPUs. A scan reads GPU 1 one layer at a time; each layer's copy prints, lifts and slides across into the same slot in GPU 2, tied back by a thread. An equals sign settles between the two.
Different text
Each copy reads different microbatches, so each calculates different gradients. If they updated using only their own gradients, the two copies would drift apart into different models.
In the drawing: The trays split, two under each GPU, and each copy's gradient ticks come out a different size. A dashed 'what if' frame surrounds the pair: the equals sign turns to 'not equal', GPU 2's layers shift to another shade, marked 'two different models'. Then the frame lifts and the sign is equals again.
Update together
Instead, the update begins by comparing notes. For every weight, the gradients from both GPUs are averaged, and both copies apply the same update. They stay identical. That's data parallelism: roughly the same batch in half the time, plus a little overhead for comparing.
In the drawing: Beads of gradients leave both GPUs' layers, meet in the gap between them, merge into one and return to both. The step count reads 1. A bracket over the pair reads 'Data parallelism · copies share the work', and the time bar at the foot is about half its old length, two segments, plus a thin pink sliver for the comparing.
Make it bigger
Training is faster, but the model isn't any bigger. Each GPU still holds the entire model. To train a larger model, we need more room.
See if four more layers fit.
or keep scrolling
In the drawing: The view pulls back, and four more layers, dashed, try to stack onto each GPU from above and stop, blocked at the top of the box, under the label 'the model we want · 8 layers'.
Out of memory
It doesn't fit. During training, each weight takes about 18 bytes: 6 for the weight and its gradient, plus 12 for the optimizer's bookkeeping, including a precise copy of the weight and a running record of past gradients. Each GPU would need enough memory for all eight layers.
In the drawing: The walls overflow: a hatched extension climbs far past the top of each box, up to the blocked layers, marked 'doesn't fit'. Each wall's fill splits into two bands as the parts are named: lavender below for the weight and its gradient, sand above for the optimizer's bookkeeping.
Expert parallelism
Inside a layer
Look inside one layer. At the bottom is attention, which every token uses to look back at earlier tokens. Most of the weights sit above it, in 16 experts. For each token, a router chooses two. For this token, it picks experts 3 and 12. Only those two do work. This is a mixture of experts.
In the drawing: The view dives into GPU 1's bottom layer as it opens: attention along its base, the router's notch, and 16 experts above. The token 'cat' goes in, and the router lights experts 3 and 12.
Split the experts
Both GPUs store all 16 experts. So divide them: GPU 1 keeps experts 1–8, GPU 2 keeps 9–16, and attention and the router stay on both. Each GPU now stores half the experts, so memory drops. This is expert parallelism.
Give each GPU half the experts.
or keep scrolling
In the drawing: Back out at both GPUs, every layer now shows its row of 16 experts. A cut runs down every row and each GPU's given-up half turns dashed; the labels read 'GPU 1 keeps experts 1–8' and 'GPU 2 keeps experts 9–16'. The walls' fill drops by about a third, and a bracket under the pair reads 'Expert parallelism · the experts are split'.
Tokens travel
The tradeoff is communication. When a token needs an expert on the other GPU, its numbers travel over and back in every layer, once on the way up and again on the way down. But eight layers still don't fit.
In the drawing: 'Cat's numbers cross the gap to expert 12 on the other GPU and come back, 'sat's do the same from GPU 2, and other tokens' in every layer. Then the pink corrections come back down, the gradients meet in the gap as before, and the step count reads 2. The four extra layers still wait above, too deep to fit.
Pipeline parallelism
Stages
Now split the model by height. GPUs 1 and 2 keep layers 1–4, while two new GPUs above them take layers 5–8. Each level is a stage. This is pipeline parallelism. Now the model fits. All eight layers still start untrained and train together from step 0.
In the drawing: The view pans up. The waiting layers lift into two dashed sockets, GPUs 3 and 4 arrive around them, and the layers turn solid; the counter reads 4 GPUs. A check mark, 'fits', lands on all four, and the rows are labelled 'Stage 1 · layers 1–4' and 'Stage 2 · layers 5–8'.
Hand it up, and wait
Each bottom GPU runs its texts through stage 1, then hands the results up to stage 2. Stage 2 finishes the forward pass and measures the error. During the backward pass, a correction comes back down so stage 1 can calculate its gradients. But one stage is always waiting, shown hatched.
In the drawing: One microbatch climbs stage 1 on both its GPUs, is handed up and climbs stage 2, where the error is measured; a pink correction comes back down. Whichever stage is idle is hatched and labelled 'waiting'. The gradients meet in the gap, and the step count reads 1.
Back to back
So feed the microbatches in one after another. While stage 2 works on the first, stage 1 can start the second. That shrinks the idle time, and more microbatches shrink it further. It also means some GPUs can be running a forward pass while others run a backward pass.
In the drawing: Two microbatches now feed back to back: stage 2 works on the first while stage 1 starts the second. A time gauge beside each stage fills a cell per tick, less of it hatched.
Scale
Deeper
Step back. Those four GPUs were always four cells in the field you saw at the start. Real models go deeper. Imagine 32 layers split across eight stages: each stage holds four layers on GPUs like the one you trained.
In the drawing: The view pulls back, a faint field of empty sockets behind, and the stack grows upward to eight stages of two GPUs, 16 in all, bracketed 'Pipeline parallelism · 8 stages', with a check mark: '32 layers fit'.
Wider
Now make it wider: 64 experts per layer, split eight ways. Each GPU still keeps eight experts, using the room the split created.
In the drawing: Each stage widens to eight GPUs, the columns headed by their experts, 1–8 through 57–64, one machine per stage, 64 in all; then all 64 lids close.
One step, 64 GPUs
Now run one training step on these 64 GPUs. Microbatches set off from every first-stage GPU, and soon every stage is busy: microbatches moving up, corrections coming back down. Then every GPU updates the weights it holds.
In the drawing: Close on the 64-GPU block, the schedule runs: microbatches climb stage by stage, corrections come back down, and each stage waits hatched until every stage is busy, some forward, some backward. The last corrections drain out, then the update re-tones every GPU. A legend names the colours: the model, forward pass, waiting, backward pass.
Copies
Now use the first move again: copy the 64-GPU team eight times. Each stage now has 64 GPUs, each reading its own microbatch, so 64 microbatches move at once. That's 64-way data parallelism, with the experts split eight ways inside it, across eight stages: 64 × 8 = 512.
In the drawing: The view pulls back, and the updated 64-GPU block repeats eight times in a ripple, filling out the 512-GPU field from the opening, cell for cell, as the counter rolls to 512. Brackets read 'Data · 64-way' down the side and 'Expert parallelism · 8 across' along the foot, with '64 × 8 = 512'.
Again, and again
Every 10× more steps, the weights across all 512 GPUs shift a shade, and the writing below improves. That's why it takes 512: 1.2 trillion weights need more memory and computation than one GPU has.
Run it for many more steps.
or keep scrolling
In the drawing: As it runs, rows of amber and pink sweep the field, forward and backward passes, and the whole field re-tones at every tenfold of steps. Below it, sample writing marked as illustrative grows a line per tenfold, from 'the the of a and' to a sentence, as a counter beside it climbs to step 1,000,000.
The GPU we followed, one of 512
Every one of the 512 GPUs is doing a version of what GPU 1 did. GPU 1 holds four of the 32 layers and 8 of the 64 experts in each.
In the drawing: The view dives back into GPU 1, one lit cell among the 512, and its lid lifts once more: four layers, eight experts in each, labelled '4 of 32 layers · experts 1–8 of 64' and '56 more on the 7 GPUs beside it'.