Training at scale: 5D parallelism

What exactly is the problem with one GPU? Let's try to understand. In a simple way. Think of a single GPU as one chip.

It has its own pool of fast memory, and that pool has a fixed size. Say eighty gigabytes on a top card. To train a model, the whole thing has to live inside that memory. The model, plus everything training needs, all at once.

But a modern model is a giant pile of numbers, called parameters. Not millions. Billions of them. And each number takes up space, a couple of bytes.

Do the math. Seventy billion numbers, at two bytes each, is a hundred and forty gigabytes. Your eighty gigabyte card cannot even hold it. It will not load.

And it gets worse. Training needs far more room than the model alone. Room for the work in progress too. So the real number is several times bigger.

Even if it fit, there is time. One chip doing all the math would grind away for hundreds of years. Nobody waits that long. So the answer is simple.

Do not use one GPU. Use many, wired together, sharing the load. The model and the work spread across them. But you can split the work in more than one way.

The data, the layers, even a single layer. Each way has its own tradeoff. That is what this course covers. So, one GPU has a fixed memory and finite speed.

Big models overflow both. The fix is to spread them over many GPUs at once. Quick check now. One question is coming up.

Let's see if it clicked.

What exactly are the two walls of memory? Let's try to understand. In a simple way. To train a model, everything must fit inside one graphics card's memory.

That memory is fixed. Say, eighty gigabytes. Run out, and training stops. So what eats all that space?

Two very different things. People call them the two walls. One is the weights. The other is the activations.

The first wall is the model itself. Its weights, the numbers it has learned. But training stores more than weights. For every weight, it also keeps a gradient, and the optimizer's own memory.

That bookkeeping adds up. A common optimizer keeps two extra numbers per weight. So the weight wall can be four times the raw model size. The second wall is the activations.

During the forward pass, every layer makes intermediate results. These must be held in memory, not thrown away. Why keep them? Because the backward pass needs them.

To nudge each weight, it walks backward through the network, reusing the outputs each layer produced. Here is the key difference. The weight wall is fixed, set by the model's size. The activation wall grows.

Feed a longer prompt, or a bigger batch, and it swells. So a model whose weights fit can still crash. The weights leave room, but the activations pile up and spill over the edge. So, two walls.

The weights, fixed and heavy. The activations, light but growing. Every trick for training at scale targets one wall or the other. Quick check now.

One question is coming up. Let's see if it clicked.

What exactly is data parallelism? Let's try to understand. In a simple way. Say your whole model fits fine on one GPU.

Training still crawls, because there is a mountain of data to feed it. And other GPUs sit idle. So here is the trick. Put a full copy of the model on every GPU.

Same weights, same layers, cloned onto each one. Four GPUs means four identical models. Now take your batch of training examples and cut it into equal slices. One slice goes to each copy.

Every clone trains on different data at the same moment. Each GPU runs its own slice forward, then backward, and works out its own gradients. Because the data differs, these gradients differ too. But here is the catch.

If each copy updated on just its own gradient, the copies would slowly drift apart. They would stop being the same model. So before updating, the GPUs share their gradients and average them together. This averaging step is called an all reduce.

Every copy ends up with one shared gradient. They all apply that same averaged update. So after the step, every copy is identical again. One logical model, many copies, marching in lockstep.

The payoff. With four copies you chew through four times the data every step. The model fits on one GPU, but the work is spread across all of them. So, data parallelism clones the whole model, splits the data across the copies, then averages their gradients to keep them in sync.

Same model everywhere, more data per step. Quick check now. One question is coming up. Let's see if it clicked.

What exactly is the cloning bottleneck? Let's try to understand. In a simple way. Remember data parallelism.

You put a full copy of the model on every GPU, and feed each a different slice of data. Cloning hides two costs. The first cost is memory. Every copy is the whole model.

Its weights, its gradients, the optimizer's numbers. So each GPU holds the entire model, however many you add. So cloning never trains a model bigger than one card can hold. Four GPUs give four times the data, not one extra inch of room.

The second cost is talking. The copies must stay identical, so every step they average their gradients. That all reduce happens after every single step. And the gradient is as big as the model itself.

Billions of numbers. Every step, all of it travels across the links between the GPUs. Here is the squeeze. The GPUs compute fast, but the cable between them is far slower.

So while gradients are shipped and summed, the GPUs sit idle. And it gets worse as you add GPUs. More copies mean more chatter to stay in sync. Past a point, extra GPUs barely help.

They just wait. So cloning hits a ceiling. It cannot shrink the model, and its sync tax only grows. To go truly huge, you cannot just copy the model.

You must split it. So, cloning speeds up data, but every copy holds the whole model, and every step pays to sync gradients. Great for more data, useless for a bigger model. Quick check now.

One question is coming up. Let's see if it clicked.

How does slicing a single layer work? Let's try to understand. In a simple way. First, what does a layer do?

It multiplies a list of numbers by a big grid of weights. Out comes a new list of numbers. Sometimes that grid is so huge it cannot fit on one card. Cloning will not help, every copy is still the whole thing.

We must cut the layer itself apart. So read the grid as columns. Each column holds the weights for one output number. Slice the columns into groups, one per GPU.

This is tensor parallelism. Every GPU gets the very same input, the full list of numbers. Nobody splits it. What differs is which slice of weights each card holds.

Now each GPU multiplies that full input by only its own columns. So each card produces just a piece of the output, a different handful of numbers. On their own, these pieces are incomplete. So the cards stitch their slices side by side into one full output, just what a single GPU would make.

Here is the win. Each card stores only its share of the weights, and runs only its share of the multiply. Across four GPUs, that is a quarter each. But splitting has a price.

To gather the slices back, the cards must talk to each other at the seam. That communication is the real cost of slicing. So, slicing a layer, tensor parallelism, cuts one weight grid into column groups. Same input to all, each computes a slice, then stitch back.

Quick check now. One question is coming up. Let's see if it clicked.

How does slicing self attention work? Let's try to understand. In a simple way. Last time, we sliced a plain layer by its columns.

It worked because every output stood alone. Attention is not so kind. What does self attention do? Every token makes a query and compares it to every other key.

Those scores become weights that blend the values. Take one comparison. A query and a key are long lists of numbers. Their score is a dot product: multiply position by position, then sum.

The tempting move: split those lists down the middle. Half of every query and key on one card, half on another. Slice the features. But watch the score.

Each card adds only its own half. Each holds a partial total; neither has the real dot product. Wrong. It gets worse.

Softmax needs the finished score to make its weights. So the cards must rejoin their halves before every softmax. The slice fights the math. Here is the fix.

Attention already comes in many heads, and heads never mix. So do not slice inside a head. Give each card whole heads. Now each card runs its heads start to finish.

Full queries, full keys, complete dot products, its own softmax. Nothing borrowed midway. When the heads finish, the cards line their outputs side by side into one result. One gather at the end, not middle chatter.

So, the dot product and softmax bind a head; you cannot cut it in half. Split by whole heads, then stitch. That is tensor parallelism. Quick check now.

One question is coming up. Let's see if it clicked.

How does splitting a model by its layers work? Let's try to understand. In a simple way. First, picture a deep model as a tall stack of layers.

Your numbers enter at the bottom and climb through every layer to the top. Last time we sliced inside one layer. Now we cut somewhere else, between the layers. A different direction entirely.

So give each GPU its own group of layers. The first card owns the bottom layers, the next card the middle, the last card the top. Now no card holds the whole model. Each stores only its own band of layers, so the giant stack finally fits, spread across the cards.

But the layers still run in order. So the cards form an assembly line. Card one works, then hands its result to card two. What crosses the gap is small.

Only one layer's output, the activations at the boundary, travels to the next card. The weights never move. So a batch flows forward like a relay. Card one, then two, then three, each doing its stretch, until the answer pops out the top.

Splitting a model by depth like this, one band of layers per GPU, is called pipeline parallelism. But there is a catch. While card one runs, the others sit idle, waiting their turn. That idle waiting is the catch the next lesson tackles.

So, pipeline parallelism cuts the model between its layers. Each GPU owns a band and passes its activations along, and the stack that never fit now runs. Quick check now. One question is coming up.

Let's see if it clicked.

What exactly is a pipeline bubble? Let's try to understand. In a simple way. Recall pipeline parallelism.

A deep model is cut into bands of layers, one band per card, wired into an assembly line. But the layers must run in order. Card two cannot start until card one hands over its result. So the cards work one after another.

Follow a single batch. While card one computes, cards two, three, and four have nothing to do. They just wait. Card one finishes and passes its output along.

Now card two works while the others wait. The busy spot slides down the line, one card at a time. Draw this on a timeline. At the start, only the first card is busy.

The rest sit idle, warming up, until the work finally reaches them. That empty stretch, where cards have no work yet, is called a pipeline bubble. It is idle time baked right into the schedule. The same thing happens in reverse.

After the last card, the backward pass empties the pipe from the far end, and cards fall idle again. With four stages and one batch, about three cards sit idle at every moment. Expensive chips, bought and powered, doing nothing. So splitting by depth saves memory, but it wastes time.

And the longer the pipeline, the bigger the bubble grows. So, a pipeline bubble is the idle time when cards wait their turn. Layers run in order, so most of the line sits still while one card works. Quick check now.

One question is coming up. Let's see if it clicked.

How does shrinking the bubble with microbatches work? Let's try to understand. In a simple way. Recall the pipeline.

A deep model is split across cards, and one big batch flows through in order. Only one card works, the rest sit idle. Here is the fix. Take that one big batch and slice it into many small pieces.

Each piece is called a microbatch. Now feed the pieces in, one after another. The moment card one finishes the first, it hands it on and starts the second. So the cards overlap.

Card one works on piece two while card two works on piece one. Every card is busy on a different piece. Draw it on the timeline. The work packs into a tight diagonal, and in the middle all four cards compute together.

The bubble does not vanish. The pipe still fills at the start and drains at the end. But that idle sliver is tiny beside the long busy middle. Here is the lever.

The bubble is stages minus one, over microbatches plus stages minus one. More microbatches, and the wasted fraction shrinks toward zero. But do not slice forever. Pieces too small give each card too little to do, and each microbatch in flight still holds activations in memory.

There is a sweet spot. This staggered schedule, slicing one batch to keep the pipe full, is called micro-batch pipelining. Same chips, far less idle. So, microbatches shrink the pipeline bubble.

Slice the batch, stream the pieces, and the cards overlap instead of waiting. Idle time nearly vanishes. Quick check now. One question is coming up.

Let's see if it clicked.

What exactly is the long context problem? Let's try to understand. In a simple way. By now you have split the model many ways.

Batch, layers, even single tensors. Then someone pastes in a whole book, and something new breaks. Weights are not the only thing in memory. Every token you feed in leaves a trail of activations.

Intermediate numbers at each layer, saved for the backward pass. This trail grows with the prompt. Twice the tokens, twice the activations to store. A short question is cheap.

A book length prompt piles up an enormous amount. Attention makes it sharper. Every token looks at every other token. That forms a grid, length by length.

Double the tokens, four times the grid. Can the other splits rescue you? Data parallelism clones the model, one example per clone. But each clone still holds a whole sequence, the full prompt on it.

Tensor parallelism cuts across features inside a layer. Pipeline parallelism cuts between layers. Neither one reduces how many tokens land on a single card. So one very long sequence, with all its activations, still stacks up on one GPU.

That card fills, and training stops. A fresh wall, made of sequence length alone. The way out is obvious. If the sequence is too big, cut the sequence.

Give each GPU a slice, not the whole prompt. That is the next lesson. So, the long context problem. A long prompt explodes activation and attention memory on one device, and the model splits cannot shrink it.

Split the tokens instead. Quick check now. One question is coming up. Let's see if it clicked.

How does splitting attention across devices work? Let's try to understand. In a simple way. Last lesson we cut the sequence.

Each card now holds just a slice of the tokens, with its own queries, keys, and values. But attention is greedy. Every token must see every other one. A query on card one needs keys sitting on card four.

The lazy fix is to copy every key onto every card. But then the whole sequence lands back on one card. The memory wall returns. So we get clever.

Arrange the cards in a ring. Each keeps its queries in place, and passes its keys and values onward. At each step a card compares its queries to whichever block of keys is visiting. It scores that piece, then the blocks rotate one hop.

You cannot just add the scores. Softmax needs the whole row at once. So each card keeps a running max and sum, folding partials in. Once the blocks travel the full circle, every query has seen every key.

Full attention, yet no card ever held the whole sequence. One neat trick. While a card scores the current block, it pulls in the next. The passing hides behind the math, costing little extra time.

This traveling pattern has a name. Ring attention, a form of context parallelism. More cards means a longer sequence you can handle. So, spin the keys and values around a ring, keep running totals, and every query meets every key.

No card holds the whole sequence. Quick check now. One question is coming up. Let's see if it clicked.

What exactly is sparse activation? Let's try to understand. In a simple way. So far, every trick split one model across many cards.

But a dense model has a catch. Every token you push through wakes up all of its weights. Every time. The mixture of experts idea breaks that big block into many smaller ones.

Each one is called an expert. Picture thirty two of them, side by side. Here is the twist. For each token, only a couple of experts switch on.

Maybe two out of thirty two. The rest stay asleep. So a single token only ever touches a thin slice of the whole model. That is the name.

Sparse activation. Most weights sit idle for that token. Now the payoff. Add up every expert and the model is enormous.

But the work for one token stays tiny, only two experts run. Big model, small bill. And this is decided fresh for every token. The word math might wake two experts, the word poem two others.

Across a sentence many experts light up, never all together. It is also a new way to split the model. Park different experts on different cards. A token's two fire, the rest stay quiet.

One question is left open. Who decides which experts a token visits? A tiny part called the router. How it chooses is the next lesson.

So, using only a fraction of the model. Break the block into experts, wake just a couple per token, and a giant model stays cheap to run. Quick check now. One question is coming up.

Let's see if it clicked.

How does expert routing work? Let's try to understand. In a simple way. A mixture of experts model holds many small experts side by side.

For each token, only a couple wake up. But who decides which ones? The router. The router is tiny.

Just a single small layer sitting in front of the experts. It does one job. Look at a token, and score every expert. It takes the token's vector and gives each expert a number.

A score. High means a good fit. Low means probably not. Those raw scores get squeezed through softmax.

Now they are clean weights that add up to one. How confident the router is in each expert. Here is the key move. Keep only the highest few.

Maybe the top two. Every other expert is ignored for this token, no matter how close it came. The two winners run. Then their outputs are blended, weighted by those scores.

The expert the router trusted more gets a bigger say. And the router is not fixed. Its weights are trained along with everything else. Over time it learns to send each token to the experts that handle it best.

One danger. Left alone, the router can play favorites. A few experts get every token while the rest starve. So a small balancing penalty pushes it to spread the load.

So, the router. Score every expert, soften the scores, keep the top few, and blend their work. A tiny chooser steering every token. Quick check now.

One question is coming up. Let's see if it clicked.

How does combining all five splits work? Let's try to understand. In a simple way. By now you have five ways to cut a model.

Data, tensor, pipeline, sequence, and expert. At real scale you do not pick one. You use them all. They stack because each cuts a different thing.

Data splits the batch. Tensor splits inside a layer. Pipeline splits across layers. Sequence splits the tokens.

Expert splits the experts. Picture your GPUs not as a line, but as a grid. One axis per split. That grid is what people mean by five D parallelism.

Every card has an address. One number per axis. Data rank three, pipeline stage two, tensor slice zero. That address decides which piece of work it holds.

And the axes multiply. Eight tensor slices, times four pipeline stages, times sixty four data copies. That is two thousand GPUs, from three small numbers. Now the part that decides your speed.

Each axis talks a different amount. Tensor trades numbers inside every layer. It is the chattiest by far. So tensor stays inside one machine, on the quickest links.

Pipeline hands off only at the cuts, so it can span machines. Data syncs once a step, and goes widest. Each axis also brings a bill. Pipeline brings bubbles.

Tensor brings constant chatter. Data copies the weights onto every clone. Tuning means trading one against another. So, five splits, one grid, every card with an address.

Multiply the axes to fill the cluster. Put the chattiest split on the fastest wires. Quick check now. One question is coming up.

Let's see if it clicked.

What exactly is forcing frontier models into five D parallelism? Let's try to understand. In a simple way. Start with the size.

A frontier model can carry close to a trillion parameters, and trains for months on tens of thousands of GPUs. Now look at one card. Eighty gigabytes. But a trillion parameters, with gradients and optimizer state, need terabytes.

One card holds a rounding error of it. So splitting is not a choice. It is the only way the model exists at all. And here is the catch.

No single cut carries the whole load. Every axis hits a ceiling, long before ten thousand cards. Tensor parallel is the chattiest. It trades numbers inside every layer, so it only works on the fast links inside one machine.

That caps it near eight. Pipeline parallel is capped by the model itself. Too many stages, and each one gets thin while the bubble eats your step. Tens, not thousands.

Data parallel goes widest, but every clone still holds a full copy of the model. And the batch cannot grow forever. None of them reaches the cluster alone. But stacked, they multiply.

Eight tensor, times sixteen pipeline, times sixty four data. Small ceilings, eight thousand GPUs. The last two axes are demanded by the model itself. A very long context forces the sequence split.

A mixture of experts forces the expert split. So, five D is not showing off. Every axis caps out early, so the frontier multiplies all five. It is the only shape that fits.

Quick check now. One question is coming up. Let's see if it clicked.