The Transformer series ended with a block I could write from memory. The GPU series put a clock on it. Then an interviewer asked me something that sits across both: “DeepSeek-V3 is a 671 billion parameter model that trains for about the price of a 37 billion parameter one. How?”

The answer is one structural change, applied to one sublayer, and every consequence of it fills five parts. In a dense transformer, parameters and arithmetic per token are the same quantity. You cannot buy one without the other. A Mixture of Experts cuts that link.

Compare two versions of the same feed-forward sublayer. On the left, every token runs one wide FFN. On the right, a router selects two of four smaller FFNs and combines their outputs. Follow the hidden vector for “sat” in “the cat sat down” to see which weights are evaluated and how the mixing equation is built. This is a small teaching example with top-2 softmax routing and ReLU experts, not a simulation of DeepSeek-V3.

What you get is capacity that costs storage but not FLOPs. What you pay is a routing decision on every token in every layer, a load-balancing problem that has no clean solution, and a network bill that turns a layer of matrix multiplies into a distributed systems problem.

This part is the trade in its simplest form: what the layer becomes, where the parameters go, and what the two columns of the ledger actually say.

Why are parameters and FLOPs the same knob in a dense model?

Take the feed-forward sublayer of a dense block. With SwiGLU it is three matrices of shape $d \times d_{ff}$, and a token passes through all three. The parameter count is $3 d\, d_{ff}$ and the arithmetic is one multiply and one add per parameter, so the FLOPs are $6 d\, d_{ff}$.

Twice the parameters, exactly. That relationship holds for every dense matrix in the model, which is where the familiar $6ND$ estimate for training compute comes from: six FLOPs per parameter per token, forward and backward.

So in a dense model there is one dial. Want more capacity? Pay more arithmetic on every token of every batch of the entire training run, and again on every token you ever serve. The scaling laws say capacity is what you want. The compute budget says you cannot have it.

Everything about Mixture of Experts follows from noticing that the token does not actually need all of that width.

What does an MoE layer replace the feed-forward network with?

A bank of $N$ smaller feed-forward networks, called experts, and a router. The router is a single matrix of shape $d \times N$. It turns the token into $N$ scores, the top $k$ are selected, those $k$ experts run on the token, and their outputs are summed with the router’s scores as weights.

That is the entire idea. Total parameters scale with $N$. Arithmetic per token scales with $k$. Two dials where a dense model has one.

The experts are usually much narrower than the layer they replace. DeepSeek-V3’s dense layers use $d_{ff} = 18432$; its experts use $d_{ff} = 2048$, a ninth as wide, and it keeps 256 of them plus one that always runs. That choice has a name, fine-grained expert segmentation, and Part 2 is about why it is the right one.

Which layers get a bank, then? Not all of them. The report is blunt about it: “We substitute all FFNs except for the first three layers with MoE layers.” So the first three blocks are ordinary transformer blocks, the other fifty-eight carry the experts, and attention is left alone in every one of the sixty-one.

That is the box the expandable DeepSeek-V3 walkthrough above stops on, and the two steps after it open the box both ways: as one wide network in a dense layer, and as a bank of narrow ones everywhere else.

Where do DeepSeek-V3’s 671 billion parameters actually go?

Two counts, one model. Total parameters count the weights stored; active parameters count the weights used for one token. Start with one MoE layer below: choose a word and watch the selection change. The unused experts stay in the bank.

The two sliders separate the trade-off: a bigger bank stores more weights; a bigger selection evaluates more weights. At a fixed selection size, adding experts only increases the small router’s contribution to the active count. Across all 58 MoE layers and the rest of the model, the published totals are about 671B stored and 37B active per token. Open the breakdown above to see where they come from.

What does the sparsity actually buy?

Compare at fixed compute, which is the comparison that matters, because compute is the budget. A dense model with DeepSeek-V3’s arithmetic per token has 37 billion parameters, by the identity from two sections ago. V3 has eighteen times that.

The claim that this is worth something is empirical, and the papers are specific. Google’s Switch Transformer reported up to a sevenfold speedup to a fixed quality against dense T5 baselines at equal compute, and about fourfold against T5-XXL. DeepSeek’s own MoE paper reported a 145B sparse model matching their dense 67B while doing a fraction of its arithmetic.

The intuition is the one from Part 5 of the Transformer series: the feed-forward layer behaves like a key-value memory. More slots is more storage. A given token only needs to look up a few of them, and looking up the ones it does not need was never doing any work.

What sparsity does not buy is anything on the attention side. All 11.4 billion parameters of DeepSeek-V3’s attention run on every token. Sparsity lives entirely in the feed-forward half of the block.

What does it cost?

Three bills, and the interview is usually about the second and third.

Memory first, because it is the one people get wrong. An MoE saves FLOPs, not bytes. Every expert has to be resident in HBM somewhere, because routing is decided per token in the middle of a forward pass and there is no time to fetch a weight from anywhere slower.

How many bytes that is depends on the format the weights are stored in. DeepSeek publish V3 only in FP8, with a conversion script for anyone who wants BF16, so a byte a parameter is the right count for serving it: about 671 GB, against the 640 GB that eight H100s hold.

Say that carefully in an interview, because the training claim is not the same claim. V3 was trained in FP8 mixed precision, not in FP8 throughout. The expensive GEMMs run in FP8. The embedding, the output head, the MoE gating, the normalisations and the attention operators are held at BF16 or FP32 on purpose, and the master weights, the gradients and the optimizer states are all kept higher.

Most checkpoints still ship in BF16, two bytes a parameter, which would put the same model at 1.34 TB. That single choice of format is the largest lever on this line, which is why the last step of the expandable DeepSeek-V3 walkthrough puts a toggle on it.

Training is the worse number either way. At the sixteen bytes per parameter the GPU series counted for Adam in mixed precision, 671 billion of them come to 10.7 TB before a single activation. V3’s own choices shave that, keeping optimizer states in BF16 and caching activations in FP8, and it is still a number no single node can hold.

Communication second. The experts a token wants live on other devices, so an MoE layer sends every token across the network and brings the results back, twice per layer, in a collective whose message sizes nobody knows until the router has run. Part 4 does that arithmetic.

Balance third. Nothing in the training objective cares whether the experts get used evenly, and a router left alone will not use them evenly. That is Part 3, and it is the failure mode that quietly deletes the capacity you paid for.

Why the feed-forward layer and not attention?

Two reasons, and having both is the difference between a memorised answer and an understood one.

The feed-forward layer is where the parameters are. In a dense block it is about two thirds of the weights, so it is the only place where sparsifying is worth the machinery.

And it is position-wise. Each token passes through the feed-forward layer independently of every other token, so sending different tokens to different experts changes nothing about what the layer means. Attention is the opposite: it is the one place where tokens interact, and an expert that only saw some of the tokens would be computing a different function, not a cheaper one.

There is research on sparsifying attention heads, and there is sparse attention, which is a different idea about which score entries to compute rather than which weights to use. Neither is what “MoE” means when an interviewer says it.

Is a sparse model just a small model wearing a big number?

This is the sharp version of the question and it deserves a real answer rather than a defensive one.

At fixed pretraining loss, the sparse model is often the worse one. There is measured evidence that at equal perplexity a sparse model underperforms a dense one on reasoning-heavy downstream tasks, which is what you would expect if reasoning consumes inference compute that a sparse model does not spend.

But nobody chooses at fixed loss. They choose at fixed compute, or fixed cost per served token, and on both of those axes the sparse model wins, which is why every frontier open-weight release since 2024 has been one.

So the honest framing is that the sparse model is not a 671B model with a 37B price tag. It is a 37B-compute model with 671B worth of storage behind it, and the reason that is a good trade is that storage got cheaper faster than arithmetic did.

Rapid fire: can you do these from memory?

Part 2 opens the router: what it computes, what the gate value is, and why every large model now has hundreds of small experts instead of eight big ones.

Part 1 established the trade. This part is the mechanism, and it is the part an interviewer will make you write on a whiteboard, because there is nowhere to hide: either you know what shape the router is and what its output multiplies, or you do not.

The whole routing apparatus is one matrix. There is no separate network, no attention over experts, no learned assignment table, and no communication between tokens. A vector goes in, $N$ scores come out, the top $k$ win.

What makes the mechanism interesting is not its complexity. It is that a component holding 0.016% of a layer’s parameters decides where 100% of that layer’s capacity goes, and it makes that decision for one token at a time with no knowledge of what the rest of the batch is doing.

Everything difficult in Parts 3 and 4 is a consequence of that last sentence.

Step through the figure with the sliders. It runs on DeepSeek-V3’s actual shapes, and the two controls in the second step let you see the difference between a softmax gate and a sigmoid gate on identical logits.

What does the router actually compute?

One matrix-vector product. For DeepSeek-V3 the router is $7168 \times 256$, which is 1.84 million parameters and 3.67 megaflops per token.

Set that against what it commits the layer to. The eight chosen experts are 352 million parameters and 705 megaflops. The decision costs half a percent of the arithmetic it authorises.

The router sees one token. It does not see the sequence, it does not see the batch, and it certainly does not see what the other layers decided. Whatever coordination exists between routing decisions is emergent, not designed, and Part 3 is largely about the consequences.

Why did the large models switch from softmax to sigmoid?

The 2017 formulation and everything through Mixtral used a softmax over the expert scores. DeepSeek-V3’s config says scoring_func: sigmoid, and Kimi K2 and several others followed.

The reason is the size of the bank. A softmax spreads one unit of probability across all $N$ experts, so as $N$ grows every individual score shrinks. On the token in the figure, a softmax over 256 experts gives the winner about 5%; a sigmoid gives it about 0.75. The sigmoid scores each expert on its own merits and does not care how many others exist.

That has two consequences. Adding experts to the bank does not rescale every existing gate, which matters if you ever want to grow a model. And the scores stay in a range where the numerics are comfortable, rather than drifting toward a regime where an exponential of a drifting logit becomes a precision problem.

After the top-$k$ renormalisation the two look similar, because both get divided by their own sum. The differences live in the gradients, in what happens when $N$ changes, and in the fact that the softmax’s exponential exaggerates whatever gap the logits had.

What is the gate value and what does it multiply?

The gate is one scalar per chosen expert. It multiplies that expert’s entire $d$-dimensional output before the sum. DeepSeek-V3 renormalises over the chosen $k$ and then multiplies by a fixed routed scaling factor of 2.5, both of which are in the published config.

Being precise about the order matters, because it is a common interview trip. Some models take the softmax over all $N$ and use those probabilities directly as gates; others take the top $k$ first and normalise within them. The second gives gates that sum to one regardless of how confident the router was, which throws away a signal and gains a scale that does not drift.

Notice what the gate is not. It is not a probability that gets sampled. It is not a mask. It is a deterministic weight in a weighted sum, and its entire job during the backward pass is to be the road the router’s gradient travels down.

Why does the router get so little gradient?

Because the only differentiable thing in the whole selection is those $k$ scalars.

The top-$k$ operation is a hard argmax. It has no useful derivative, so no gradient flows through the choice itself. What does flow is $\partial L / \partial g_i$ for the $k$ experts that were chosen, and from there back into the router’s weights. The 248 experts that were not chosen contribute nothing to the output and receive nothing from the backward pass, and the router learns nothing about whether one of them would have been better.

This is the exact reason routing is hard, and it has a sharp corollary worth having ready. If you use top-1 and renormalise the gate over the chosen set, the gate is identically 1, a constant, and the router receives no gradient at all from the output. It would learn only from the balancing loss.

Switch Transformer, which is the model that made top-1 respectable, avoids this by keeping the raw softmax probability as the multiplier rather than renormalising. Shazeer’s original 2017 formulation had argued you needed $k \ge 2$ for the router to get a meaningful signal; Switch’s answer was that you need $k \ge 2$ only if you throw the magnitude away.

Why 256 tiny experts instead of eight large ones?

This is the design move that separates the 2024 generation from Mixtral, and it comes from the DeepSeekMoE paper, where it is called fine-grained expert segmentation.

Hold the arithmetic fixed and cut each expert into $m$ narrower pieces, then select $m$ times as many. Active parameters are unchanged to the digit. FLOPs are unchanged. What changes is the number of distinct expert teams a token can assemble, and it changes by a lot: $\binom{8}{2}$ is 28, while $\binom{256}{8}$ is about $4 \times 10^{14}$.

The bet is that expressiveness in the combination is worth something, and the published models all took it. Mixtral had 8 experts and used 2. DeepSeek-V3 has 256 and uses 8. Kimi K2’s config lists 384 and uses 8. Qwen3-235B lists 128 and uses 8.

What it costs is not FLOPs. It is a router matrix that grows with $N$, a longer sort in the kernel, and $mk$ network destinations per token instead of $k$, which is the whole subject of Part 4. Fine-grained routing is cheap in arithmetic and expensive in systems work.

What is a shared expert for, and does it help?

DeepSeekMoE’s other idea: hold out one or two experts that run on every token, outside the top-$k$. The argument is that whatever every token needs, common syntax, the shape of the residual stream, does not need to be learned redundantly in all 256 experts, so isolating it frees the routed experts to be different from each other.

DeepSeek, Moonshot and Z.ai all ship one. Their published configs list n_shared_experts: 1 alongside 256, 384 and 160 routed experts respectively.

And it is genuinely contested. AI2 ran the matched-compute ablation for OLMoE and reported that sharing an expert performed slightly worse than using two routed ones, with a neat combinatorial argument: with 32 experts and 4 active, making one of them shared cuts the choices from $\binom{32}{4} = 35{,}960$ to $\binom{31}{3} = 4{,}495$, removing almost 90% of the combinations. Alibaba shipped shared experts in Qwen2.5-MoE and removed them in Qwen3.

If asked, the honest answer is that it is a small effect either way, that the labs disagree, and that it probably interacts with granularity, because the more experts you have the less a single always-on one costs you in combinations.

What happens when an expert runs out of room?

To run the experts as batched matrix multiplies you need fixed-size buffers, and you have to size them before you know the routing. That is what the capacity factor is for. Switch Transformer defines expert capacity as tokens per batch divided by the number of experts, times a capacity factor, and a token arriving at a full expert is simply dropped: its computation is skipped and it passes to the next layer through the residual connection.

Note what that means. A dropped token is not an error and does not raise an exception. It is a token that silently did not get a feed-forward layer, and it happens more the more skewed the routing is.

The capacity factor is the knob and it is a bad one. Too low and you drop tokens; too high and your buffers sit half empty and the matrix multiply is mostly padding.

The modern answer is to stop doing it. Block-sparse kernels, of which MegaBlocks is the reference, run variable-size expert groups directly and remove the capacity factor as a hyperparameter entirely. DeepSeek state plainly that V3 does not drop any tokens during training, and describe a deployment strategy that avoids dropping at inference too. If you are asked about capacity factors in 2026, the useful answer explains what they were for and why the frontier stopped needing them.

What does the kernel actually run?

Not $k$ forward passes. This is the misconception that a whiteboard question is designed to catch.

The sequence is: compute the router for the whole batch, take the top $k$, sort the token rows by expert index into one contiguous buffer, run a single grouped matrix multiply whose group boundaries are the expert offsets, then scatter the results back to token order and weight them by the gates.

The two permutations are pure data movement. For a batch of 8192 tokens with $k = 8$ and $d = 7168$ they move a few gigabytes per layer in BF16 and do zero arithmetic, which puts them at the far left of the roofline from Part 2 of the GPU series: bandwidth-bound, nothing else.

That is where a naive implementation loses most of its throughput, and it is why MegaBlocks, the Triton grouped-GEMM kernels and DeepSeek’s own DeepEP exist. It is also why the group sizes being unknown until the router has run is such a persistent problem: a kernel that wants static shapes has to wait for a number that arrives late, every layer, every step.

Rapid fire: can you do these from memory?

Part 3 is the failure mode: what happens to a router that nobody is supervising, and the three different mechanisms the field uses to stop it.

Part 2 ended on the sentence that makes this part necessary: the router’s only gradient is the gate values of the experts it already chose. An expert it stops choosing sends back nothing, and the router learns nothing about it.

That is a positive feedback loop with no damping term. An expert that wins a few extra tokens early gets a few extra gradient updates, becomes slightly better at the tokens it sees, and wins more of them. Left alone, a Mixture of Experts is a winner-take-all system.

The failure is quiet. The loss curve looks fine for a long time, because a bank in which four experts do all the work is still a perfectly good four-expert model. What is being destroyed is the capacity you paid the memory bill for, and nothing in the objective mentions it.

Three mechanisms exist to stop it, and a good answer distinguishes them by what they change: the loss, the selection, or the numerics.

The figure runs the same sixteen-expert router three times on the same token stream: once with nothing, once with the auxiliary loss, once with the bias. Press play on the first step and watch it fall over.

Why does a router collapse?

Two rules are enough, and they are both true of any Mixture of Experts trained by gradient descent.

Tokens improve an expert. An expert that receives more tokens takes more gradient steps on the task, so it gets better at it.

A better expert scores higher. The router is trained to send tokens where the loss goes down, so an expert that is better at a token attracts more of them.

Compose those and you have compounding. The figure’s first step models exactly this and nothing else, with the two rules written at the top of its source, and it collapses to a single expert in about a hundred steps. Nothing exotic is required: no bad data, no bad initialisation, no bug.

What makes it dangerous is the shape of the curve. The load imbalance sits flat for the first forty steps, drifts for thirty more, then falls off a cliff. If your only instrument is the training loss, the first sign of trouble arrives long after the mechanism has started.

What does collapse actually cost you?

Nothing you can see on a memory profile, which is the point.

The parameters are all still there. The checkpoint is the same size, the HBM bill from Part 1 has not moved by a byte, and the FLOPs per token are unchanged because $k$ experts still run. What has moved is how much of the bank is learning anything. An expert with no tokens gets no gradient, so it stops improving, and a sixteen-expert bank in which one expert does everything is a one-expert model with sixteen experts’ worth of storage.

Scale that to a frontier model. If a quarter of DeepSeek-V3’s routed experts went quiet, about 163 billion parameters would stop learning while continuing to occupy HBM and continuing to be sharded, checkpointed and shipped.

That is why the load histogram, logged per layer per step, is the cheapest instrument in an MoE run and the first thing to add. It costs an $N$-element reduction and it is the only place collapse is visible early.

What does the auxiliary loss actually penalise?

The formulation from GShard, made standard by Switch Transformer:

where $f_i$ is the fraction of tokens in the batch dispatched to expert $i$, and $P_i$ is the mean router probability assigned to expert $i$ over those tokens. Both vectors sum to one over the $N$ experts.

Their dot product is smallest when both are flat, and the $N$ in front normalises the minimum to exactly 1. Its maximum is $N$, when one expert takes everything. So the number is readable: 1 means perfectly even, 16 on a sixteen-expert bank means total collapse, and in practice it hovers a little above 1 because routing noise never lets it settle.

Switch used $\alpha = 10^{-2}$, having swept it from $10^{-1}$ down to $10^{-5}$ and found that value balanced the load quickly without interfering with the training loss. That sweep is worth remembering: too small and the mechanism does nothing, too large and you are optimising for uniformity rather than for the task.

Why the product of two vectors and not the variance of one?

Because you need something differentiable, and $f$ is not.

$f_i$ is a histogram. It counts how many tokens the hard top-$k$ actually sent to expert $i$, so it is piecewise constant in the router’s weights and its gradient is zero almost everywhere. Penalising the variance of $f$ alone would give you a number that describes the problem and a gradient that does nothing about it.

$P_i$, the mean router probability, is smooth. Multiplying the two gives a surrogate that reads as “wherever a lot of tokens actually went, push the probability down”, and its gradient is

which is positive for over-used experts and negative for under-used ones. Being able to say which factor carries the gradient is the standard follow-up question, and it separates people who have read the formula from people who have implemented it.

There is a second-order version of this worth knowing, because it is an implementation detail that changes the model rather than the speed. The batch you compute $f_i$ over is a choice.

Most frameworks compute it per micro-batch, which at frontier scale is a handful of sequences, so the term effectively demands that every individual sequence spread itself evenly over all the experts. A micro-batch of pure code is then forced to use the whole bank.

Qiu and colleagues showed in January 2025 that computing $f_i$ over the global batch instead, which costs one extra all-reduce of an $N$-element vector, improves both perplexity and the domain specialisation of the experts, at scales up to 42.8 billion parameters and 400 billion tokens.

What is DeepSeek’s auxiliary-loss-free balancing?

The observation behind it is that the auxiliary loss is a foreign term. Its gradient is not trying to make the model better at the task; it is trying to make the histogram flat, and those two gradients are added together and pull in different directions.

So remove it. Keep a per-expert bias $b_i$, add it to the score used for the top-$k$ comparison, and leave the gate value alone:

After each step, decrease $b_i$ by a small $\gamma$ for experts that were over-subscribed and increase it for those that were under-subscribed. No term is added to the loss. No gradient is distorted. It is a controller bolted to the side of the router rather than an objective, and it works because it changes which experts are chosen without changing how much they count.

DeepSeek’s numbers for it, on models up to 3 billion parameters, are a perplexity of 9.50 against the auxiliary loss’s 9.56 at 1B parameters and 100B tokens, with a maximum load violation of 0.04 against 0.72. Better balance and slightly better perplexity, which is the claim.

Two footnotes, both worth having. In DeepSeek-V3 the bias update rate was $\gamma = 0.001$ for the first 14.3 trillion tokens and zero for the rest, and they kept a complementary sequence-wise auxiliary loss at $\alpha = 0.0001$ to stop any single sequence from piling onto one expert. So the “loss-free” method ships alongside a small loss, with its own controller mostly switched off.

And the paper was rejected at ICLR 2025 and shipped in V3 anyway. In the public reviews the authors conceded that the interference-gradient motivation was based on intuition and needed more rigorous validation. Knowing that is a better answer than either endorsing or dismissing the method.

What is the z-loss for, and why is it a different problem?

The router z-loss, from the ST-MoE paper, is not about balance at all.

It penalises the squared log-sum-exp of the router logits. Squaring means the penalty grows with the size of the logits and says nothing about which expert is winning, so it constrains magnitude and leaves the ranking alone.

The reason it exists is numerical. Router logits drift upward over a long run, the way attention logits do, and large numbers have large rounding errors in low precision, which an exponential then amplifies.

ST-MoE chose a coefficient of 0.001 by sweeping for the best model quality after pretraining, and they pair it with a second recommendation: cast the router’s input to float32 before the softmax and cast the dispatch tensors back to bfloat16 afterwards. The router tensor is tiny, so the fp32 cast costs nothing and removes the problem at the source.

This is the router’s version of the QK-norm story from the Transformer series. An exponential of a drifting logit is a stability hazard wherever it appears, and the fix is always to bound the logit rather than to fix the exponential.

Does a balanced router mean specialised experts?

No, and this is the question that separates a memorised answer from an understood one.

Every metric in this part counts how many tokens each expert received. None of them says anything about whether the assignment means anything. Two routings can have identical perfect balance, zero maximum violation and an auxiliary loss of exactly 1, with one assigning tokens by content and the other assigning them at random. The figure’s last step draws both.

What the evidence says is genuinely mixed, and the disagreement is informative. Mistral’s own routing analysis of Mixtral found no obvious pattern by topic: the expert assignment distribution for arXiv papers, PubMed abstracts and philosophy papers is very similar at every layer. What they did find is syntactic and positional structure. Python’s self and indentation tokens route consistently, and consecutive tokens pick the same first-choice expert about 28% of the time at layer 15 against a 12.5% random baseline.

ST-MoE found clear specialisation in their encoder experts, punctuation, verbs, proper nouns, and explicitly did not find it in the decoder. They also looked for language specialisation in a multilingual model and found none: experts handled English, Japanese, French and Chinese indiscriminately.

AI2’s OLMoE, trained from scratch rather than upcycled, does report domain and vocabulary specialisation, with one layer-0 expert nearly 100% specialised to arXiv, and they hypothesise that Mixtral shows less of it because it was upcycled from a dense model, so all its experts started from the same optimum.

So the honest position: specialisation at the level of token class and vocabulary is real and measurable; specialisation by subject matter is contested and may depend on how the model was initialised; specialisation by language does not appear. “Which expert handles medicine” is the wrong question. “Is any expert idle” is the right one.

What do you actually watch during a run?

Four things, and none of them is the training loss.

The per-layer load histogram, every step. Collapse starts in one layer and spreads, so a model-wide average hides it. The single number to derive from it is the maximum violation: the busiest expert’s load over the even load, minus one.

The number of experts receiving zero tokens. It is the metric that turns into lost capacity, and it is a step function: fine, fine, fine, then several.

The router logit magnitude, which is what the z-loss is defending, and which tells you whether the numerics are drifting before a spike appears.

And the auxiliary loss value itself, if you are using one, remembering that it never reaches its minimum and is not supposed to. What matters is the trend, not the level.

Rapid fire: can you do these from memory?

Part 4 leaves the loss function and asks what it costs to actually run this: two all-to-alls a layer, and the reason a sparse model is harder to batch during decoding than a dense one.

Parts 2 and 3 were about the router as a piece of mathematics. This part is about what happens when you actually have to run it, and the answer is that a Mixture of Experts layer stops being a matrix multiply and becomes a distributed systems problem.

The reason is simple to state. A dense layer’s weights are the same for every token, so you can shard them however you like and the tokens never move. An MoE layer’s weights are chosen per token, the chosen ones are on other devices, and so the tokens move instead.

That single change swaps an all-reduce for two all-to-alls, replaces static message sizes with sizes the router decides at runtime, and, during generation, quietly undoes most of the arithmetic-intensity argument that made sparsity attractive in the first place.

The GPU series drew expert parallelism as one axis on the parallelism map. This part is what happens along that axis.

The figure computes all of that from DeepSeek-V3’s shapes and H100-class hardware, with every assumption on the screen so you can disagree with it.

Why is expert parallelism the natural split?

Because a bank of experts has a seam that a dense matrix does not. Expert parallelism sits alongside the other four axes in the Parallelism series; here it is the sparse model’s own bill.

Tensor parallelism cuts every matrix by rows or columns, so each device computes a partial result and the partials have to be summed: an all-reduce, every layer, with a message size known at compile time. That works, and for the attention half of an MoE model it is still what you do.

Expert parallelism gives whole experts to whole devices. Nothing needs summing across devices, because each expert’s output is complete where it was computed. What has to move instead is the tokens, out to their experts and back with the results.

For DeepSeek-V3 the arithmetic is stark. One MoE layer’s experts are 11.3 billion parameters. Spread 257 experts over 64 GPUs and each holds four of them, 177 million parameters per layer instead of 11.3 billion. That is not an optimisation, it is the only reason 671 billion parameters fit anywhere.

What travels, and how much of it?

Per token, per layer, in each direction: the hidden state out to each destination, and one result vector back per expert.

Dispatch sends $d$ values per destination. Combine brings back $d$ values per expert. DeepSeek’s DeepEP dispatches in FP8 and combines in BF16, so the return journey moves twice the bytes of the outward one. With $d = 7168$ and $k = 8$ that is about 57 kB out and 115 kB back per token per layer, before any optimisation.

Now multiply. Fifty-eight MoE layers, two collectives each in the forward pass and two more in the backward, is 232 all-to-alls per training step, each of which is a full every-GPU-to-every-GPU exchange with $P^2$ flows.

The lever that matters most is node-limited routing. Cap the number of distinct nodes a token is allowed to reach, send it across the slow network once per node rather than once per expert, and fan it out over NVLink inside. DeepSeek-V3 capped it at four nodes, which turns eight network destinations into at most four.

Is an MoE layer a FLOP problem or a network problem?

Do the arithmetic and it is not close.

Take the figure’s default: expert parallelism 64 across eight nodes, 8192 tokens per GPU, node-limited routing on, an H100 at 989 TFLOP/s with 450 GB/s of NVLink and a 400-gigabit port that moves 50 GB/s. Each GPU computes about 65,000 token-expert pairs, which is 6.5 TFLOP and 6.6 milliseconds at peak. The traffic it has to move is 705 MB across the network, which is 14.1 milliseconds, plus 1.4 GB over NVLink, which is 3.1.

Communication is 2.6 times the arithmetic. That is the whole story of why MoE systems work exists, and it is why an interviewer asking “what is the bottleneck in an MoE layer” is not asking about matrix multiplies.

Change any input and the ratio moves, which is why the sliders are there. What does not move is the shape of the answer. The arithmetic and the traffic both grow linearly with tokens, so batching does not help the ratio. What helps is the interconnect, the precision of the payload, and how many nodes a token is allowed to touch.

How do you hide the all-to-all?

A ratio above one only hurts if the two things happen in sequence, and essentially all the systems work since GShard is about making them happen at the same time.

Pipeline schedules put a different micro-batch’s arithmetic underneath this one’s collective. DeepSeek’s DualPipe runs the pipeline in both directions at once so there is always compute travelling the other way, which the GPU series covered as the reason their 64-way expert parallelism was affordable at all.

Kernel libraries attack it from below. DeepEP’s dispatch and combine kernels are written for zero or minimal streaming-multiprocessor occupation, which matters more than it sounds: an all-to-all implemented as an ordinary kernel steals the multiprocessors the expert matrix multiplies want, so a communication kernel that uses no SMs is not just faster, it stops competing.

DeepSeek’s own April 2026 report puts a number on how far this goes. They split the experts into waves, so that the computation of the current wave, the token transfer for the next, and the result-sending of the completed ones all proceed at once.

And they state the balance condition directly: because a token-expert pair costs $6hd$ FLOPs but only $3h$ bytes, each gigabyte per second of interconnect bandwidth suffices to hide the communication for 6.1 teraflop per second of compute. Above that threshold, more bandwidth buys nothing. They report 1.50 to 1.73 times speedups on general inference and up to 1.96 on latency-sensitive work, and they shipped the fused kernel as MegaMoE.

That result is also why V4 removed the node-limited routing constraint that V3 needed. Once the communication is genuinely hidden, the cap that bounded it stops earning its cost in routing quality.

Why is decoding a sparse model harder than decoding a dense one?

This is the part that surprises people, and it is the single best thing to have on your fingertips in this whole series.

Training moves thousands of tokens together. Generation moves one per sequence per step. Each decoding token picks its $k$ experts independently, so with $B$ tokens decoding together the expected number of distinct experts touched in a layer is

For DeepSeek-V3, $N = 256$ and $k = 8$. At $B = 1$ that is 8 experts, 3% of the bank. At $B = 32$ it is 163. At $B = 128$ it is 252, which is 98% of the bank.

So the bytes you have to read stop scaling with $k$ and start scaling with $N$, while the FLOPs keep scaling with $k$. Arithmetic intensity is the ratio of exactly those two things, and it collapses.

At $B = 128$, one MoE layer reads about 11 GB of expert weights to do about 100 gigaflops: an intensity of 9 FLOP per byte, against the H100 ridge of 295 from Part 2 of the GPU series. A dense model holding the same weights would be at 256 at that batch, because it reads its weights once and every token uses all of them.

As $B$ grows the ratio settles at exactly $N/k$, which for V3 is 32. A sparse layer needs thirty-two times the concurrent batch of a dense one to sit at the same point on the roofline.

What does that do to how you serve it?

It splits the machine in two, which is roughly what every 2026 inference stack converged on.

Prefill is compute-bound. The matrix multiplies are already large, so the all-to-all is pure overhead and narrow expert parallelism wins.

Decode is bandwidth-bound. Spreading the experts over more devices shrinks the weights each one has to read and frees the HBM that a larger batch needs, so wide expert parallelism wins.

DeepSeek published their own serving configuration and it is exactly this split: expert parallelism 32 across four nodes for prefilling, and expert parallelism 144 across eighteen nodes for decoding, with two routed experts and one shared expert per GPU. Their stated reason is the arithmetic above, in their words: only 8 of 256 experts are activated per layer, so the system needs an extremely large overall batch to give each expert a batch size worth having.

The two phases want opposite settings of the same knob. That is the systems argument for running them on separate machines, and it is the answer to “how would you deploy this” that shows you have thought past the model card.

The other levers are all about the skew from Part 3, which does not go away at inference. The same DeepSeek writeup runs 32 redundant routed experts in both phases and an expert-parallel load balancer whose objective is to minimise the maximum dispatch load across GPUs, because the all-to-all is synchronous and the slowest rank sets the step time for everyone.

Why is temperature zero not determinism?

Because on a capacity-limited serving stack your token’s routing depends on the other tokens in the batch.

If expert capacity is enforced, an expert that fills up drops whatever arrives next, and what arrives next depends on which other requests happened to be batched with yours. Two identical prompts at temperature zero can take different paths through the model because someone else’s traffic changed at the same moment.

The dropless kernels remove that particular version of the problem, but a subtler one remains: routing is a top-$k$ over floating-point scores, so any numerical difference at all, a different kernel, a different batch shape, a different accumulation order, can flip which expert wins. Discrete selection turns a rounding difference into a completely different computation.

That is a good thing to raise unprompted in a systems interview, because it explains a class of production bug that looks impossible: the model gives different answers, the seed is fixed, and nothing in the request changed.

What is the memory floor, really?

Total parameters, always. There is no way around it, and this is the most common misconception about MoE in the whole subject.

The reason is timing. The router decides mid-forward-pass, per token, which experts are needed. There is no window in which to fetch a weight from host memory or from disk, so every expert has to be resident in HBM somewhere before the layer starts.

Offloading the cold experts to the host is a real technique and it works exactly where you would expect: at batch one, where a token touches eight of 256 experts, most of the bank genuinely is idle. As soon as you batch, the coverage formula above says you touch nearly all of them, and offloading turns into streaming the whole model over PCIe every step.

So the honest sentence is: an MoE saves FLOPs, not bytes. It is a way to buy capacity with memory instead of with arithmetic, and memory is the resource you have to be able to afford.

Rapid fire: can you do these from memory?

Part 5 leaves the mechanism behind and asks what is actually unsettled: which models are sparse in September 2026, what changed in the last year, and the list of questions where a confident answer is a bad sign.

The first four parts described a mechanism the field has largely converged on. A bank of narrow experts, sigmoid affinities, top-8, one shared expert or none, balanced by a bias rather than a loss, sharded by expert parallelism, dropless.

The interesting thing about September 2026 is that the convergence has started to come apart, and the lab that broke it first is the one that set most of it. DeepSeek-V4, published in April, changed the gate, removed the routing constraint that V3 needed, replaced the dense layers at the bottom of the stack with hash-routed experts, and moved the expert weights to four bits.

This part is the state of play, dated, with the questions where a confident answer is a bad sign. I have tried to keep the distinction visible throughout between what a lab measured, what a lab asserted, and what nobody has settled.

The figure is the landscape: eleven open-weight models with published numbers, the sparsity trend, what routing alternatives exist, and the upcycling question.

Which models are sparse, and how sparse?

Every open-weight frontier release since 2024 is a Mixture of Experts, and their configs are public, so this is one of the few places in the subject where you can be exact.

Mixtral 8x7B, in December 2023, stored 47 billion parameters and used 13. DeepSeek-V3, a year later, stored 671 and used 37. Meta’s Llama 4 Maverick stores 400 and uses 17, with 128 experts and an unusual routing scheme: each token goes to one routed expert plus a shared one, in layers that alternate with dense ones.

Qwen3-235B-A22B has 128 experts, top-8, no shared expert, across 94 layers. Kimi K2’s config lists 384 routed experts plus one shared. GLM-4.5 lists 160 plus one. OpenAI’s gpt-oss-120b has 128 experts with top-4 routing and the expert weights quantised to MXFP4.

DeepSeek-V4, in April 2026, comes in two sizes: V4-Pro at 1.6 trillion total and 49 billion active, and V4-Flash at 284 billion and 13.

Two trends fall out of that list. The sparsity ratio went from 3.6 to more than 30 in under three years. And the active-parameter axis barely moved: the frontier settled somewhere between 13 and 50 billion active parameters and spent all its growth on the other axis.

That second observation is the one worth being able to state, because it says what the binding constraint actually is. Compute per token is set by the inference budget and by what a serving fleet can afford. Total parameters are set by how much HBM you can put behind it. Sparsity is the only knob that moves capacity without moving the first number.

What changed in DeepSeek-V4?

More than the parameter count, and the specific changes are the best available snapshot of where the recipe is going.

The affinity function moved. V3 computed expert scores with a sigmoid; V4 changed it to the square root of a softplus. That is a small edit to one line and it is the first time a frontier lab has publicly moved off the gate everyone had copied.

Node-limited routing is gone. V3 capped each token at four destination nodes to bound the network fan-out, which was the lever Part 4 spent a section on.

V4 removes the constraint, and the reason is in their own systems section: with the wave-based overlap they describe, the communication within a single MoE layer takes less time than the computation, so once it is fused into one pipeline the compute is the bottleneck and the system tolerates lower interconnect bandwidth. Their threshold is that each gigabyte per second of interconnect suffices to hide 6.1 teraflop per second of compute, and past it more bandwidth buys nothing.

The dense layers at the bottom are gone too. V3 kept plain feed-forward networks in its first three blocks. V4 replaces them with MoE layers using hash routing, where a predefined hash of the token id decides the experts. No router, no gradient, no balancing problem, in exactly the layers where the earlier literature found routing looks most like a partition of the vocabulary anyway.

And the routed expert parameters are stored in FP4. Balancing is still auxiliary-loss-free with a small sequence-wise loss alongside, which is the one part of the V3 recipe that survived untouched.

Has routing moved past top-k?

Not in what ships, and that is itself the interesting answer.

The alternatives are all well known. Expert choice inverts the argmax: instead of every token picking $k$ experts, every expert picks its best $c$ tokens. Load is then balanced by construction, with no auxiliary loss and no bias controller at all. The catch is that an expert’s choice depends on the other tokens in the batch, including later ones, which in a causal decoder leaks information backwards in time, and that a token can end up chosen by nobody.

AI2 ran the comparison directly for OLMoE and reported that dropless token choice beat expert choice on every task at the same token budget, while expert choice ran about 20% faster per device. That is the trade in one line: expert choice buys throughput and balance and pays in quality and in causality.

Hash routing goes the other way and does not learn at all. The Hash Layers work in 2021 reported that routing tokens by a fixed hash was competitive with learned routing, which is an uncomfortable result if you believe the router is doing something clever. It sat as a curiosity for five years and then reappeared in DeepSeek-V4’s first blocks.

So the honest summary is that top-$k$ token choice is still the default, that nothing has beaten it convincingly at scale, and that the two most credible alternatives are one that cannot be used in a decoder and one that does not learn.

Should you upcycle a dense checkpoint?

You already have a good dense model. Copy its feed-forward layer $N$ times, perturb the copies, bolt on a router, and keep training. That is sparse upcycling, and it is the cheapest way to get an MoE.

Google’s paper on it put the crossover at roughly 120% of the original dense training budget: below that, upcycle, above it, start over.

AI2 tried it for OLMoE and got a different number. They upcycled their OLMo-1B checkpoint after 2 trillion tokens and trained for 610 billion more. An otherwise equivalent MoE trained from scratch caught up after 500 billion extra tokens and started beating it around 600 billion. That is 25% of the original budget, not 120%, and they rejected upcycling for their final model.

Both are right about their own setting, and the useful thing is the reason they differ. Beyond the compute bracket, AI2’s stated objection was that an upcycled model inherits hyperparameters tuned for the dense model it came from, and the previous section says those do not transfer. There is a second reason, and it connects to the specialisation question below: experts that all start from the same weights start correlated, and it takes training to pull them apart.

The practical reading is that upcycling is a way to spend a small budget well, not a way to reach a frontier model.

Do the hyperparameters transfer?

This is the newest result on the list, dated 8 September 2026, and it is the one I would raise unprompted.

Muon, muP and the whole hyperparameter-transfer literature rest on the idea that you can tune a small proxy and scale the settings up along a known law. The paper “Hyperparameter Scaling Laws Across MoE Sparsity” tests that across sparsity levels with 1,800 pretraining runs, six activated-parameter scales, models up to 6 billion total non-embedding parameters, about 20 trillion tokens, and roughly 200,000 H800-hour equivalents.

The finding is that conventional hyperparameter scaling laws are insufficient for ultra-sparse models. Optimal learning rate and optimal batch size both shift with the activation ratio, and those shifts cannot be explained by either total or activated parameter count alone. Their resolution is that at fixed sparsity the optimal batch size is a power law in the token count and the optimal learning rate a power law in compute, and that across sparsity levels the activation ratio enters both as an additional multiplicative factor.

The interview version of this: if you tune on a dense proxy and transfer to a model activating a sixty-fourth of its parameters, you should expect the transfer to be wrong, and neither of the two parameter counts you might reach for will tell you by how much.

Do experts specialise?

The popular claim is that a Mixture of Experts learns a maths expert, a code expert, a biology expert. The evidence does not support that, and being able to say why is a genuinely differentiating answer.

Mistral looked for it in Mixtral and did not find it. Their own routing analysis says the assignment distribution is very similar across arXiv papers, PubMed abstracts and philosophy papers at every layer. What they did find is structure of a different kind: Python’s self and indentation tokens route consistently, and consecutive tokens pick the same first-choice expert about 28% of the time at layer 15 against a 12.5% random baseline. That is syntax and temporal locality, not subject matter.

ST-MoE found clear specialisation in their encoder experts, on punctuation and verbs and proper nouns, and explicitly did not find it in the decoder. They also looked for language specialisation in a multilingual model and reported that experts handled English, Japanese, French and Chinese indiscriminately.

OLMoE, trained from scratch rather than upcycled, does report domain and vocabulary specialisation, with one layer-0 expert nearly 100% specialised to arXiv, and their hypothesis for the disagreement is that Mixtral was upcycled from Mistral, so its experts all started from the same optimum and had less room to diverge.

There is a mechanism underneath that makes the whole thing less mysterious. The balancing scope from Part 3 acts directly against specialisation: if you compute the dispatch fractions over a micro-batch of a few sequences, you are demanding that a batch of pure code spread itself over the whole bank.

Computing them over the global batch instead improves both perplexity and measured domain specialisation, at scales up to 42.8 billion parameters. Some of what looks like a fact about Mixture of Experts is a fact about a reduction in a framework.

So: token class and vocabulary, yes. Subject matter, contested and initialisation-dependent. Language, no.

What breaks when you reinforcement-learn a sparse model?

The routing, and it is the sharpest modern failure mode.

Top-$k$ is a discrete function of continuous scores, so any numerical difference at all can flip which expert wins. In reinforcement learning that matters more than anywhere else, because the algorithms depend on comparing the probability a token had under the old policy against its probability under the new one, and if the two policies ran different experts you are comparing two different networks.

The Qwen team put a number on it while introducing GSPO. On the 48-layer Qwen3-30B-A3B base model, after each gradient update and for the same rollout sample, roughly 10% of the activated experts differ between the new policy and the old one, and the effect becomes more prominent in deeper models.

Their workaround was routing replay: cache the experts the old policy activated and replay those routing decisions when computing the importance ratios. They call it essential for normal convergence of GRPO on MoE models, and they note it costs memory, communication, and some of the model’s effective capacity.

GSPO’s own contribution is to sidestep it by working at the sequence level rather than the token level, so it is not sensitive to individual token likelihoods and does not need the replay at all.

If someone tells you their reinforcement learning run diverges on a sparse model and is fine on a dense one of the same size, this is the first thing to check, and it is not a bug in their code.

Does sparsity change what else you have to change?

Yes, and the combinations are where the 2026 work is.

The obvious pairing is sparse attention. Parts 1 and 4 established that sparsity lives entirely in the feed-forward half of the block: all of DeepSeek-V3’s attention runs on every token, and the KV cache is untouched by any of it.

So a model that is sparse in the feed-forward layer and dense in attention has moved its bottleneck rather than removed it. That is why DeepSeek pairs MoE with multi-head latent attention to shrink the cache and, from V3.2 onward, with a sparse attention mechanism to cut the score computation.

Their natively trainable sparse attention work in February 2025 laid the groundwork. V4’s hybrid attention combines a compressed sparse mechanism with a heavily compressed one, and they report V4-Pro needing 27% of V3.2’s single-token inference FLOPs and 10% of its KV cache at million-token context.

The other pairing is quantisation, and it interacts with sparsity in a way dense models do not have to think about. Experts see wildly different amounts of data. A hot expert is calibrated on a diverse stream; a cold one is under-represented in whatever calibration set you use, and a shared expert sees everything.

Quantising them all identically is the obvious thing and the wrong one. gpt-oss quantises the MoE weights to MXFP4 and leaves attention, the router and the embeddings alone, which is the shape of the answer: the router is tiny and precision-sensitive, the experts are enormous and are where the bits are.

What would I actually say if asked what comes next?

One prediction and one refusal.

The prediction: the active-parameter axis stays roughly where it is and the total-parameter axis keeps climbing. The binding constraints are inference FLOPs and HBM bandwidth, and sparsity is the only knob that moves capacity without moving either. Nothing in the public sweeps has found a loss-based ceiling on how sparse a model can be; the ceilings that have actually been hit are systems ceilings and numerical ones.

The refusal: I would not predict what the router looks like. In eighteen months the field has kept top-$k$ token choice, moved from softmax to sigmoid, moved from an auxiliary loss to a bias controller, moved from sigmoid to the square root of a softplus, and put hash routing back into production.

Every one of those changes was small and none was forced by a theory. The honest summary is that the router is the least understood component of the most successful architecture of the decade.

That is also the reason the questions in Part 6 exist. The mechanism is learnable in an afternoon. Knowing which of its parts are load-bearing and which are conventions nobody has re-examined is the thing that takes longer.

Rapid fire: can you do these from memory?

Part 6 is the self-test: three levels of questions with answer sketches, built from everything above.

Five parts of prose is a good way to build an understanding and a bad way to test one. This part is the test.

Three levels. The first asks whether you can state the mechanism precisely: shapes, what is computed, what is stored, where the bytes and the FLOPs go. The second asks whether you can reason about the trade-offs, which is where most interviews actually live. The third is the frontier, where the useful answer is usually a question about what the interviewer means.

Every answer sketch is what a strong answer contains, not an essay. Nothing in them is a number this series has not already established, and if one of them surprises you, the part it came from is the one to reread.

Level 1: can you state the mechanism?

Level 2: can you reason about the trade-offs?

Level 3: can you defend a design decision?

What do you do with a list like this?

Not memorise the answers. The sketches are deliberately short because the point of each one is a single load-bearing idea, and if you have that idea you can build the sentences around it in the room.

The pattern across all three levels is the same. Level 1 is arithmetic you should be able to do on a whiteboard from the config alone. Level 2 is a question about which quantity is binding, and almost every one of them is answered by asking “at fixed what?” before answering. Level 3 has no settled answer, and the strong response names the axis it turns on and then commits to a default anyway.

That last habit is worth practising specifically. “It depends” is a bad answer on its own and a good one when it is followed by what it depends on and what you would do today.

Six parts ago the question was how a 671 billion parameter model trains for the price of a 37 billion parameter one. The answer turned out to be one edit to one sublayer, and five parts of consequences: a router with almost no gradient, a bank that collapses without supervision, a layer that spends more time on the network than on arithmetic, and a decode path that gives back most of its sparsity as soon as you batch.

None of that is hard. It is just longer than the one sentence the architecture is usually described in, and the gap between the sentence and the consequences is where the interview happens.

The GPU series ended with an inventory: sixteen bytes per parameter, an activation formula, a KV cache, and four clocks that decide which of them is the bottleneck. That inventory is for one GPU. The next question an interviewer asks is the one that follows from it, and it is almost always phrased as a task rather than a question: “you have 512 H100s and a 70 billion parameter model. Lay it out.”

There are five ways to cut a model across GPUs, and they are not interchangeable. Each one divides a different line of that inventory, each one buys the division with a different collective, and each one runs out of road for a different reason. Getting the list right is easy. Knowing which line each one divides, and what it costs, is the part that separates a candidate who has read about this from one who has run it.

This series derives the five, one at a time, on one model and one cluster. Everything is arithmetic you can do on a whiteboard, and every number in every figure is computed from the same constants.

What does one GPU actually have to hold?

Take Llama 3 70B, whose shape Meta published: 70.5 billion parameters, 80 layers, hidden size 8192, an FFN width of 28,672, 64 query heads and 8 key-value heads of dimension 128.

Training one parameter with Adam in mixed precision costs sixteen bytes: a BF16 weight, a BF16 gradient, an FP32 master copy, and two FP32 optimizer moments. That is 1.13 TB of model state, and none of it is optional while the optimizer is Adam.

Then the activations. Megatron’s accounting for one transformer layer, with FlashAttention removing the quadratic attention term, is $34\,s\,b\,h$ bytes. At a sequence of 8192 tokens and $h = 8192$ that is 2.28 GB for one layer, and 182.5 GB across all eighty.

Add them and one GPU would need 1.31 TB. An H100 has 80 GB. The first figure draws that as sixteen and a bit GPUs’ worth of boxes, and then empties them one axis at a time.

How long would one GPU take anyway?

The second budget is time, and it does not care about the first. A dense transformer costs about six FLOPs per parameter per token: two in the forward pass, four in the backward.

An H100 does 989 TFLOP/s of dense BF16 at peak, and a well-tuned training run achieves about 40% of that. So one GPU delivers 396 TFLOP/s of useful arithmetic and finishes the run in roughly five hundred years.

Five hundred and twelve of them take 363 days. Sixteen thousand take eleven. That is the shape of the second budget: it is satisfied only by throwing hardware at it, and the hardware is only useful if the model is split in a way that keeps it busy.

Notice that the two budgets give different answers. Memory says you need at least seventeen GPUs. Time says you want thousands. The interesting engineering lives in the gap.

What are the five things you can split?

Five, and it helps to name them by what gets cut rather than by their acronyms.

Split the batch. Every GPU holds the model and processes different tokens. Gradients are averaged at the end of the step. ZeRO and FSDP are refinements that stop replicating the states.

Split the tensors. Each matrix multiply inside a block is cut across GPUs, so every GPU holds a slice of every layer and works on the same tokens.

Split the layers. Each GPU holds a contiguous group of layers and passes activations to the next group, like a factory line.

Split the sequence. Each GPU holds part of the token sequence, which only matters because attention couples the whole sequence together.

Split the experts. In a mixture-of-experts model, the experts live on different GPUs and tokens travel to them.

Data parallelism is the only one of the five where the model is never cut. That is why it is the default, and why it is the first to run out.

Why does each axis divide a different thing?

This is the table I would draw on the whiteboard before answering anything else, and the fourth step of the figure builds it.

Four things occupy memory during training: the weights, the gradients, the optimizer states, and the activations saved for the backward pass. Every axis divides some of them and leaves the rest alone.

Plain data parallelism divides nothing. ZeRO stage 1 divides the optimizer states, stage 2 the gradients too, stage 3 the weights as well. None of the three touches the activations, because each rank still has to keep the activations of its own tokens.

Tensor parallelism divides all four, which is what makes it valuable and expensive at the same time. Pipeline parallelism divides the three model-state lines by the number of stages, but keeps several micro-batches in flight, so the activations barely move. Context parallelism divides the activations and nothing else. Expert parallelism divides the expert weights and nothing else.

What does each axis cost on the wire?

Numbers, for the layout this series arrives at: tensor parallelism of 8 inside each node, fully sharded data parallelism across the 64 nodes, one micro-batch of 8192 tokens.

One GPU owns 8.81 billion parameters and does 0.44 seconds of arithmetic per micro-batch at peak. Against that:

The gradient all-reduce of plain data parallelism moves 34.7 GB per rank per optimizer step. Fully sharded data parallelism moves 52.9 GB, because stage 3 adds an all-gather in the backward pass. Both ride the 400 Gb/s network at 50 GB/s, so they take of order a second, and both can hide behind the backward pass if the framework overlaps them properly.

Tensor parallelism moves 75.2 GB per micro-batch, not per step, because it fires four times per layer. It gets away with that only because NVLink inside a node is 450 GB/s per direction, nine times the network. On the network the same traffic takes 1.5 seconds against 0.44 seconds of arithmetic.

Pipeline parallelism moves 33.6 MB per stage boundary per micro-batch, forward and backward. It is the cheapest axis on the wire by three orders of magnitude, and it is the only one whose cost is idle time rather than bytes.

What does a good layout look like on this cluster?

The last step of the figure sets it out, and every line in it is derived rather than remembered.

Tensor parallelism of 8, filling one node, divides both the weights and the activations. Fully sharded data parallelism across the 64 nodes divides the states again, so no byte of the model is stored twice anywhere in the cluster. Micro-batch of one 8192-token sequence per tensor-parallel group. Gradient accumulation sets the global batch without changing a single collective.

That gives 2.2 GB of model state per GPU, 22.8 GB of activations, about 4 GB of context and buffers and workspace, and 29 GB of an 80 GB card in use. The arithmetic is 438 ms per micro-batch at peak, and the whole run is 363 days at 40% MFU.

There is a lot of headroom in that 29 GB, and the rest of the series is about what to spend it on: a larger micro-batch, a deeper pipeline, a longer sequence, or nothing at all.

What is this series, and what is it not?

It is the derivation, not the walkthrough. Every part asks the question an interviewer opens with, works out the mechanism from the shapes, and puts a number on it.

For the practical side of the same subject, running jobs on a real cluster, I wrote Distributed Training from Scratch earlier this year: the batch-size hierarchy, DDP against FSDP with real config, and what actually breaks at 256 GPUs. That series is the one to read for launch scripts and the shape of a bad night. This one is the one to read before someone asks you to derive a bubble formula on a whiteboard.

For the hardware underneath, the GPU series has the four clocks, the roofline, the memory hierarchy and the byte-by-byte inventory this part starts from. I will point at it rather than repeat it.

Six parts of derivation, and then a graded question set to check whether it stuck.

Rapid fire: can you do these from memory?

Part 2 takes the axis everyone starts with, splitting the batch, and follows it until it stops working.

Data parallelism is the axis nobody has to be talked into. Every GPU runs the ordinary single-device code on its own slice of the batch, the framework averages the gradients behind your back, and the result is mathematically identical to one enormous GPU with one enormous batch. Nothing about the model is cut.

Which is exactly why it runs out. The model is not cut, so every GPU still needs room for all of it, and the activations of its own tokens on top. ZeRO fixes the first half of that sentence and cannot touch the second.

This part follows the axis from a plain all-reduce to fully sharded data parallelism, counts the bytes at every step, and then works out the three walls it hits.

What actually happens in a data-parallel step?

Four phases, and only one of them talks to anyone.

Each rank takes a different slice of the global batch, runs the forward pass, runs the backward pass, and ends up with a gradient for every parameter. Those gradients disagree, because each rank saw different tokens, and the optimizer wants the average over the whole batch.

One collective fixes that. After it, every rank holds the identical averaged gradient, applies the identical update to its identical weights, and the invariant holds for the next step.

For our running layout, where tensor parallelism has already cut the weights eight ways, each data-parallel rank is responsible for $\Psi = 8.81$ billion parameters. In BF16 that is 17.6 GB of gradient, and the all-reduce moves about twice that per rank.

Why does the all-reduce not get more expensive with more ranks?

This is the first thing interviewers check, and the answer is a two-line derivation.

A ring all-reduce is two halves. In the reduce-scatter, the message of $M$ bytes is cut into $d$ chunks, and each chunk walks around the ring being summed as it goes. After $d-1$ hops, rank $i$ owns the completely reduced version of one chunk. In the all-gather, those finished chunks walk around again until everyone has all of them.

Every hop moves one chunk, $M/d$ bytes, and there are $2(d-1)$ hops:

So bandwidth cost is flat in $d$. What is not flat is latency: there are $2(d-1)$ sequential hops, so a ring across 64 ranks pays 126 link latencies before the last byte lands. That is why large clusters use tree or hierarchical algorithms for small messages and rings for large ones, and why NCCL picks between them at runtime.

The second step of the figure animates the chunk grid and counts the bytes as they move. At $d = 8$ with a 17.6 GB gradient, each rank sends 30.8 GB, and it would send less than 35.2 GB no matter how large the ring got.

How does the collective hide behind the backward pass?

The reduction does not have to wait for the backward pass to finish, because gradients become final layer by layer, from the top of the network down.

So the framework buckets them. PyTorch’s DDP fills a bucket, 25 MB by default, and launches an asynchronous all-reduce the moment it is full. The reduction of the last layer overlaps the backward pass of the one before it, and only the tail sticks out past the end.

That works beautifully when there is a lot of compute per byte, and stops working when there is not. The collective is fixed by the parameter count. The backward pass shrinks with the tokens on the GPU.

At 8192 tokens per GPU the backward arithmetic is 292 ms and the all-reduce over 64 ranks is 694 ms, so more than half the step is exposed. The third step of the figure lets you drag the token count and watch the tail appear.

What is redundant about holding the optimizer state 512 times?

Nothing about a data-parallel step needs every rank to hold every optimizer moment. The ZeRO paper’s observation is that after the reduction, rank $i$ only ever updates its own slice of the parameters, so it only ever needs its own slice of the states.

Stage 1 shards the twelve bytes of FP32 master weight and Adam moments across the $d$ ranks. Stage 2 shards the two bytes of gradient too, which turns the all-reduce into a reduce-scatter. Stage 3 shards the two bytes of BF16 weight as well, which means gathering each layer’s parameters just before it runs and freeing them straight after.

The floors matter more than the slopes. Stage 1 never gets below $4\Psi$ and stage 2 never below $2\Psi$, whatever $d$ is, because the BF16 weight and then the gradient stay replicated on every rank. For our $\Psi$, that is 35.3 GB and 17.6 GB of permanent floor. Only stage 3 goes to zero.

What does each stage cost on the wire?

The ZeRO paper counts communication in elements, which is worth copying because it makes the comparison clean.

Plain data parallelism moves $2\Psi$ elements per rank per step: an all-reduce is a reduce-scatter of $\Psi$ followed by an all-gather of $\Psi$. Stage 1 moves the same $2\Psi$. Stage 2 moves the same $2\Psi$, because the reduce-scatter it uses is exactly the first half of the all-reduce and the all-gather of updated parameters is the second half.

Stage 3 moves $3\Psi$. The extra $\Psi$ is a second all-gather: the weights have to be gathered once for the forward pass and once again for the backward, because they were freed in between.

In BF16 those are four and six bytes per parameter. For our rank, 35.3 GB and 52.9 GB per step. A 50% increase in traffic for memory that falls linearly in $d$ is the trade, and at any reasonable scale it is a good one.

There is a middle option that shows up in real runs. Meta’s Llama 3 used FSDP to shard the optimizer states and the gradients, but deliberately did not reshard the parameters after the forward pass, so the backward pass reuses the gathered weights instead of paying for a second all-gather. That is stage 3’s memory for the forward and stage 2’s traffic overall, and it is the right choice when memory is not the binding constraint.

What does stage 3 actually do per layer?

Stage 3 is a promise that a layer’s weights will exist by the time the layer runs.

Before layer $i$ computes, its shard is all-gathered from the $d$ ranks into a full copy. The layer computes. The full copy is freed. In the backward pass the same gather happens again, the gradient is computed, a reduce-scatter sends each rank its own slice, and the copy is freed again.

Done naively, the GPU idles through every gather. Done with prefetch, the gather for layer $i+1$ is issued while layer $i$ is still computing, and the only exposed gather in the whole model is the first one. The fifth step of the figure has a prefetch toggle: turn it off and the compute row fills with holes.

The knob is wrapping granularity. Wrap too coarsely and the gathered chunk is enormous, which defeats the memory saving. Wrap too finely and each gather is too small to reach line rate on the network.

PyTorch’s FSDP2 shards each parameter on its own with DTensor rather than flattening a module’s parameters into one buffer, which is what turns wrapping into a per-module decision and makes it compose with tensor parallelism. In practice you wrap one transformer block per unit, and one block’s weights on our layout are 220 MB gathered.

What does gradient accumulation actually buy?

It buys arithmetic to hide the collective behind, and nothing else.

Run $k$ micro-batches, accumulate the gradients locally, reduce once. The collective is unchanged, because it is sized by the parameter count. The compute per optimizer step is $k$ times larger. So the exposed fraction falls as $1/k$ without changing the global batch’s learning dynamics at all, as long as you were going to run those tokens anyway.

The catch is numerical. Accumulating $k$ BF16 gradients loses the small contributions, so the accumulator has to be FP32. Meta reported doing exactly that for Llama 3: FP32 gradient accumulation across micro-batches, and an FP32 reduce-scatter across data-parallel workers.

The other thing to know is what accumulation does not buy. It does not reduce memory, because the accumulator is the same size as the gradient. And it does not help the pipeline, which wants micro-batches for a completely different reason that Part 4 gets to.

Where does the batch axis stop?

Three walls, and none of them is about bandwidth.

The activations never shard along this axis. After ZeRO-3 has taken the model states on our layout from 141 GB to 2.2 GB, the 182.5 GB of activations for a single 8192-token sequence is exactly where it was. That is the wall that sends you to tensor, pipeline or context parallelism, and it is the one that matters most.

The global batch has a ceiling. Global batch is tokens per rank times $d$, so at fixed tokens per rank, adding ranks grows the batch, and batch size is a hyperparameter that stops helping. Hold the batch fixed instead and the tokens per rank fall. Meta held Llama 3 405B at 16 million tokens per batch and halved the per-rank batch when the data-parallel degree went from 64 to 128; the reported MFU fell from 43% to 41%, and they attribute the drop to exactly that.

And falling tokens per rank means less arithmetic to hide the same fixed collective behind, which is the third wall and a consequence of the second.

All three say the same thing. Past a point you have to cut the model, not the batch.

Rapid fire: can you do these from memory?

Part 3 cuts the model for the first time, and starts with the smallest unit worth cutting: a single matrix multiply.

Part 2 ended at a wall: the activations of a sequence do not shard along the batch axis, no matter which ZeRO stage you run. To divide them you have to cut the model itself, and the smallest thing worth cutting is one matrix multiply.

There are exactly two ways to cut a weight matrix, and which one you choose is forced by the nonlinearity that follows it, not by the bandwidth. That single observation is most of Megatron’s tensor parallelism, and it is the thing to be able to derive rather than recall.

The rest of the part is arithmetic: how many collectives per layer, how many bytes each one moves, what sequence parallelism adds for free, and why the degree stops at the number of GPUs in a node.

Which way do you cut the weight matrix?

Take one GEMM followed by a nonlinearity, $Y = f(XA)$, and two GPUs.

Cut $A$ along its rows and you must cut $X$ along its columns to match. Each GPU computes a partial product, and $f$ cannot be applied to a partial product, because $f(u + v) \neq f(u) + f(v)$ for anything worth calling a nonlinearity. So the partials have to be summed first, which means a synchronisation in the middle of the block.

Cut $A$ along its columns instead and each GPU owns whole output features:

The nonlinearity applies column by column, locally, and nothing is communicated at all. That is the asymmetry the whole scheme is built on.

Megatron pairs them. First GEMM column-parallel, second GEMM row-parallel so it consumes the sharded output directly, and one reduction at the end of the pair. The reduction sits exactly where the residual stream has to be whole anyway.

What does that look like in the MLP block?

Llama’s MLP is three matrices, not two: a gate projection and an up projection, both $8192 \times 28672$, and a down projection of $28672 \times 8192$. SwiGLU multiplies the gated half by the up half elementwise.

Split gate and up along their columns. At $t = 8$ each GPU owns 3584 of the 28,672 intermediate features, computes both halves for those features, and applies SwiGLU to its own slice with no communication, because the product is elementwise in the feature index.

Split down along its rows over the same partition. Each GPU multiplies its 3584 features by its 3584 rows and produces a full-width partial sum of the output. One all-reduce adds the eight partial sums.

Each GPU holds 176 MB of this block instead of 1.4 GB, and every activation inside the parallel region is a factor of eight smaller.

What are f and g, exactly?

Two autograd hooks, and they are conjugates of each other. This is the detail that tells an interviewer whether you have implemented it or only read about it.

$f$ sits at the entry to the parallel region. In the forward pass it does nothing, because the input was already broadcast to every GPU in the group. In the backward pass it all-reduces the gradient, because every GPU produced a gradient with respect to that same shared input and they have to be summed.

$g$ sits at the exit. In the forward pass it all-reduces, because the output was a partial sum. In the backward pass it does nothing, because the incoming gradient is already identical on every GPU.

Identity forward, all-reduce backward. All-reduce forward, identity backward. Two lines of PyTorch each.

How does attention split?

More naturally than the MLP, because the heads are already independent.

Give each GPU $64/t$ query heads and the key-value heads that go with them. The query, key and value projections are column-parallel over the head dimension, so each GPU makes only its own heads. The scores, the softmax and the weighted sum for a head happen entirely on the device that owns it: no head ever needs another head.

The output projection is row-parallel over the same partition, so it consumes the local head outputs and emits a partial sum. Same shape as the MLP: column, row, one all-reduce.

There is a hard limit hiding in there, and it is grouped-query attention. Llama 3 70B has 8 key-value heads. At $t = 8$ each GPU owns exactly one, which is tidy. At $t = 16$ the key-value heads have to be duplicated across pairs of GPUs, and duplicated weights mean duplicated gradients and an extra reduction to keep them in step. The KV head count is a ceiling on the tensor-parallel degree that has nothing to do with bandwidth.

How many collectives is that, and how many bytes?

Two all-reduces per layer in the forward pass, one after attention and one after the MLP. Two more in the backward pass from the conjugate operators. Four per layer, and for 80 layers that is 320 collectives per micro-batch.

Each message is $s\,b\,h$ elements. At 8192 tokens and $h = 8192$ in BF16, 134.2 MB. A ring all-reduce over $t$ ranks moves $2(t-1)/t$ of that per rank, which at $t = 8$ is 234.9 MB.

Against 438 ms of arithmetic on this GPU, that is 167 ms on NVLink at 450 GB/s per direction, so 38% overhead. On a 400 Gb/s network port at 50 GB/s it would be 1.50 seconds, three and a half times the arithmetic.

And note what it is per. Data parallelism’s collective fires once per optimizer step and amortises over gradient accumulation. This one fires per micro-batch, four times per layer, and amortises over nothing.

What is sequence parallelism, and why is it free?

Tensor parallelism leaves two things replicated on every GPU in the group: the layer norms and the dropouts, and the residual stream they act on. They are cheap to compute and expensive to store.

In Megatron’s activation accounting they are the $10\,s\,b\,h$ that stubbornly refuses to divide:

The fix is that those operations are independent along the sequence, so they can be split that way instead. The all-reduce at the end of a parallel region becomes a reduce-scatter, which lands each GPU with its own slice of the sequence, and the entry to the next region becomes an all-gather.

The bytes on the wire do not change, because a ring all-reduce was a reduce-scatter followed by an all-gather all along. Korthikanti and colleagues make that point explicitly: four all-reduces per forward and backward becomes four all-gathers and four reduce-scatters, at identical bandwidth.

What changes is memory. At $t = 8$ and our shape, a layer goes from 872 MB to 285 MB, a factor of three, and the model’s activations from 69.8 GB to 22.8 GB. It is the rare optimisation with no downside, which is why every framework bundles it with tensor parallelism rather than exposing it as a separate switch.

What does tensor parallelism do to each kernel?

Two things, and they both get worse as $t$ grows.

The matmuls get narrower. The MLP GEMM on one GPU is $8192 \times 3584$ at $t = 8$, and $8192 \times 896$ at $t = 32$. A narrow GEMM has fewer tiles to spread over 132 streaming multiprocessors and a worse ratio of epilogue to arithmetic, so it sits further below the roofline. This is the Part 2 of the GPU series problem arriving through a different door.

And every all-reduce is a synchronisation. Whatever the fastest GPU in the group finishes early is wasted, four times per layer, so a straggler tax that would be invisible once per step becomes 320 taxes per micro-batch.

There is a live line of work on the first problem. Asynchronous tensor parallelism decomposes the GEMM and the collective into chunks so that a slice of the matmul overlaps the transfer of the previous slice, using symmetric memory to write directly into a peer’s buffer. It is in torchtitan under compile, and it is the main reason the practical ceiling on $t$ has moved at all since 2021.

Why does the degree stop at eight?

Three limits arrive at almost the same place, which is unusual and convenient.

Bandwidth. Arithmetic per GPU falls as $1/t$, while communication climbs towards a ceiling, because $2(t-1)/t \to 2$. So the ratio of communication to computation grows roughly as $t$: 5% at $t=2$, 16% at $t=4$, 38% at $t=8$, 82% at $t=16$ on NVLink. Across a network it is nine times worse and hopeless everywhere.

Heads. Eight key-value heads means eight is the largest degree that does not duplicate them.

Tiles. At $t = 32$ the FFN matmul is 896 wide, which is a handful of tile columns.

Megatron’s own guidance from the other direction is the same conclusion: use tensor parallelism up to the number of GPUs in a server, then use pipeline parallelism to go further. Every published dense recipe in Part 6 has $t = 8$, and the reason is that servers have eight GPUs.

Rapid fire: can you do these from memory?

Part 4 takes the other way of cutting the model: not across each layer, but between them.

Tensor parallelism cut across every layer and paid in bandwidth. The other way to cut a model is between layers: give the first twenty to one group of GPUs, the next twenty to the next, and pass activations down the line.

It is the cheapest axis on the wire by a very wide margin. One activation tensor per stage boundary per micro-batch, point to point, no collective at all. It is also the only axis whose cost is idle time, and idle time does not overlap with anything.

Everything about pipeline parallelism is a fight with that idle time. There are four schedules worth being able to draw, and the difference between them is entirely in when the backward passes run.

What crosses a stage boundary?

Almost nothing, which is the point.

The forward pass sends one activation tensor from the last layer of stage $i$ to the first layer of stage $i+1$. With tensor and sequence parallelism already applied, that tensor is $s/t \times h$ in BF16, which for our shape is 16.8 MB. The backward pass sends the same shape the other way.

Compare that with the 75.2 GB per micro-batch that tensor parallelism moves in the same pass. Pipeline traffic is three orders of magnitude smaller, and it is point to point rather than a collective, so it does not synchronise anyone who is not directly involved.

The weights, gradients and optimizer states all divide cleanly by $p$. At four stages our 70B model’s states go from 141 GB per tensor-parallel rank to 35.3 GB.

And then you run one batch through it and every stage but one is idle.

Where does the bubble come from?

Cut the batch into $m$ micro-batches and push them in one after another. Stage 0 starts on micro-batch 1 at time zero, stage 1 starts on it one slot later, and stage $p-1$ cannot start until $p-1$ slots have passed.

The same thing happens in reverse at the end, when the last micro-batch drains back down the pipeline. Between those two, every stage is busy.

At $p = 4$ and $m = 16$ that is 18.8%, and the figure measures exactly that off the drawn schedule rather than asserting it. Drag $m$ and the number falls as $1/m$. Drag $p$ and it climbs.

A flush at the end of every batch is what forces this. Asynchronous and bounded-staleness pipelines avoid it by letting different micro-batches see different weight versions, and every production framework rejected that trade, because reproducing a training curve matters more than the last few percent.

Why does 1F1B not make it faster?

Because the bubble does not care what order the passes run in. It cares how deep the pipeline is.

In GPipe, each stage runs all $m$ forward passes, then all $m$ backward passes. In 1F1B, a stage warms up with a few forwards and then alternates one forward and one backward for the rest of the batch. Put the two schedules next to each other on the same axis, as the third step of the figure does, and the wall clocks are identical.

What changes is the number of micro-batches a stage has in flight, meaning micro-batches whose forward pass has run and whose backward pass has not, so whose activations must still be resident. GPipe holds all $m$. 1F1B holds at most $p - i$ on stage $i$.

At $m = 16$ and four stages, that is 91.3 GB of activations on stage 0 against 22.8 GB. The first number does not fit in an H100 and the second does.

So 1F1B does not make the pipeline faster. It makes a large $m$ affordable, and a large $m$ is what makes the pipeline fast. Megatron calls it PipeDream-Flush and it has been the default everywhere for years.

How does interleaving cut the bubble?

The bubble is $p-1$ slots long, and a slot is however long one stage’s forward pass takes. So make the slot shorter.

Give each device $v$ non-contiguous chunks of the model instead of one contiguous block. With $v = 2$ and eight stages, device 0 owns the first sixteenth of the layers and the ninth sixteenth. The pipeline now has $vp$ virtual stages, each a $v$-th as long.

The last device still starts after $p-1$ forwards, because virtual stages 0 through $p-1$ are the first chunk on devices 0 through $p-1$. But each of those forwards is now $1/v$ of a full one. The warm-up costs $v$ times less time:

The fourth step of the figure zooms into the warm-up of both schedules, which is where the entire difference lives, and measures the bubble off each.

The price is $v$ times as many point-to-point messages, over the slowest links in the job. At 16.8 MB a message that is a trade almost anyone would take. Meta used interleaving for Llama 3 and quoted the bubble in the paper as $(PP-1)/(V \cdot M)$.

What does splitting the backward pass buy?

A backward pass is two computations glued together. One computes the gradient with respect to the layer’s input, which the previous stage is blocked on. The other computes the gradient with respect to the layer’s weights, which nobody is blocked on.

The Zero Bubble paper’s whole idea is to stop treating them as one unit. Call them $B$ and $W$. Now $B$ is half as long, so the critical path through the pipeline is shorter, and $W$ can be dropped into any gap after its own $B$.

Their ZB-H1 schedule keeps 1F1B’s peak memory exactly and takes the bubble from $(p-1)(F+B+W)$ to $(p-1)(F+B-W)$, which is a third of the size when the three passes take equal time. The fifth step of the figure shows the ochre $W$ blocks filling what were holes, and measures 18.8% falling to 6.3%.

ZB-H2 goes further, adding forwards to the warm-up and reordering the $W$s at the tail so the shape becomes a parallelogram rather than a trapezoid. That reaches $(p-1)(F+B-2W)$, which is zero when the three take the same time, and it costs about twice the activations in flight. Getting the last of the bubble also required removing the synchronisation in the optimizer step, which the paper handles with a post-hoc validation rather than an up-front all-reduce.

DeepSeek-V3 went further again with DualPipe: two pipelines running in opposite directions at once, with a bubble of $(p/2 - 1)(F\&B + B - 3W)$, paid for with two copies of the parameters. They built it because their expert-parallel all-to-all needed something to hide behind, which is Part 5’s problem.

How much memory is actually in flight?

This is the part people get wrong, and it is a good interview question because the naive answer is so tempting.

Deepening the pipeline divides the weights by $p$. It does not divide the activations, because with a bubble-minimising schedule the first stage has to keep $p$ micro-batches in flight to stay busy, and each of those carries $L/p$ layers of activations.

Korthikanti and colleagues state it exactly that way: the total activation memory on the first stage is $\frac{s\,b\,h\,L}{t}\left(34 + \frac{5as}{h}\right)$, with no $p$ in it at all. On our shape that is 22.8 GB regardless of whether the pipeline is four stages deep or sixteen.

Interleaving makes it slightly worse, by a factor of $1 + \frac{p-1}{pm}$, because each device holds pieces of two places in the network. Nobody notices that term; everyone notices the $L$.

How do you choose the number of micro-batches?

You do not choose it directly. It falls out of a batch identity:

Larger $m$ shrinks the bubble as $1/m$. But at a fixed global batch, every GPU added to the data-parallel dimension takes micro-batches away from the pipeline. That is why the pipeline degree and the data-parallel degree cannot be chosen independently, and why frontier runs hold tokens per batch constant and accept a falling MFU rather than letting the bubble grow.

The last step of the figure plots the bubble against $m$ for one, two and four-way interleaving. At four stages you need about thirty micro-batches for a 10% bubble with plain 1F1B, or eight with four-way interleaving.

And a practical constraint that a lot of implementations impose and Meta had to remove: many schedules require the micro-batch count to be divisible by the number of stages. Their fix for Llama 3 was to make the number of consecutive micro-batches per stage a free parameter, so they could run fewer micro-batches than stages when the batch was tight, or more when they wanted to hide point-to-point latency.

What else goes wrong in a real pipeline?

Two imbalances that the clean picture does not show, and Meta documented both.

The first stage carries the embedding table on top of its layers, and carries the most warm-up micro-batches, so it uses the most memory. The last stage computes the output projection and the loss, so it is the slowest, and in a pipeline the slowest stage sets the pace for everyone. Their fix was to take one transformer layer away from each end: the first chunk holds only the embedding, the last only the projection and the loss.

Beyond that, the list is unglamorous. Asynchronous point-to-point so a stage does not block on a send. Proactive deallocation of stage input and output tensors that nothing will read again. And a document mask that makes different micro-batches cost different amounts, which is a straggler problem with a pipeline shape.

With all of that, Meta reported pre-training Llama 3 on 8192-token sequences with no activation checkpointing at all, which is the outcome the whole exercise is for.

Rapid fire: can you do these from memory?

Part 5 takes the two axes people reach for last, and the two that interviews probe hardest: splitting the sequence, and splitting the experts.