How to Parallelize a Transformer for Training

What Do We Mean By Scaling?

The goal of “model scaling” is to be able to increase the number of chips used for training or inference while achieving a proportional, linear increase in throughput (we call this strong scaling). While performance on a single chip depends on the trade-off between memory bandwidth and FLOPs, performance at the cluster level depends on hiding inter-chip communication by overlapping it with useful FLOPs. This is non-trivial, because increasing the number of chips increases the communication load while reducing the amount of per-device computation we can use to hide it. As we saw in Section 3, sharded matrix multiplications often require expensive AllGathers or ReduceScatters that can block the TPUs from doing useful work. The goal of this section is to find out when these become too expensive.

In this section, we'll discuss five common parallelism schemes: (pure) data parallelism, fully-sharded data parallelism (FSDP / ZeRO sharding), tensor parallelism (also known as model parallelism), expert parallelism (for Mixture-of-Experts models), and (briefly) pipeline parallelism. For each, we'll show what communication cost we incur and at what point that cost starts to bottleneck our compute cost.◦We'll focus on communication bounds — since while memory capacity constraints are important, they typically do not bound us when using rematerialization (activation checkpointing) and a very large number of chips during pre-training. (Ed: This edition is expanded to discuss expert parallelism, unlike the original.) For this section, you can focus solely on inter-chip communication costs, since as long as we have a large enough single-chip batch size, the transfer of data from HBM to MXU is already overlapped with computation.

We'll use the following notation to simplify calculations throughout this section.

For simplicity's sake, we'll approximate a Transformer as a stack of MLP blocks — attention is a comparatively small fraction of the FLOPs for larger models as we saw in Section 4. We will also ignore the gating matmul, leaving us with the following simple structure for each layer:✦adaptation With this simplification each layer holds 2·D·E·F weights (E = 1 for a dense model, so simply 2·D·F), and the whole stack has 2·D·E·F·L = parameters at the current settings — the “P” in this page's communication arithmetic. Memory questions are different: a real checkpoint holds the gated MLP's third matrix and the attention stack too, so the memory meters price weights at P_w ≈ 3·D·E·F·L + 2.5·D²·L = , which tracks the model table's published totals to within a few percent (vocab embeddings and MHA-era attention excepted).

Here are the 4 parallelism schemes we will discuss. Each scheme can be thought of as uniquely defined by a sharding for In, W_in, W_out, and Out in the above diagram.✦adaptation A quick reminder of the book's notation: a subscript on an array dimension names the mesh axis it is split over — In[B_DP, D] means the batch dimension is carved into DP pieces, one per chip along axis DP — while a subscript on a · names the dimension being contracted. The explorer below the four descriptions lets you click through the schemes: for each one it draws all four arrays with their shards colored by chip, the per-chip local shapes at the live DP and TP, and the collectives the scheme pays for in the forward and backward pass.

1. Data parallelism: activations sharded along batch, parameters and optimizer state are replicated on each device. Communication only occurs during the backwards pass.

2. Fully-sharded data parallelism (FSDP or ZeRO-3): activations sharded along batch (like pure data parallelism), parameters sharded along same mesh axis and AllGathered just-in-time before use in forward pass. Optimizer state also sharded along batch. Reduces duplicated memory.

3. Tensor parallelism (also called Megatron sharding or model parallelism): activations sharded along D (d_model), parameters sharded along F (d_ff). AllGather and ReduceScatter activations before and after each block. Compatible with FSDP.

4. Pipeline parallelism: weights sharded along the layer dimension, activations microbatched and rolled along the layer dimension. Communication between pipeline stages is minimal (just moving activations over a single hop). To abuse notation:✦adaptation Notice what all four schemes have in common: every one runs the same matmuls — the FLOPs never change, only where the arrays live and which collectives must run between the multiplies. So for each scheme the question is always whether those collectives can hide behind the matmuls. Before the chapter dissects the schemes one by one, this adaptation inserts a short primer — First, Feel the Roofline — building the one picture that answers that question for all four.

First, Feel the Roofline

✦ This entire section is an addition of this adaptation — the chapter's ideas, our framing. The original text resumes at Data Parallelism.

The core DP/FSDP/TP rooflines in this chapter are one picture wearing a few costumes. Mixed sharding combines those clocks; expert and pipeline parallelism add topology- and scheduling-specific activation traffic. Before we meet them, let's get the core picture into your fingers.

When a chip works on one layer of our Transformer, two clocks run at the same time:

The compute clock: the MXU has to chew through this layer's share of FLOPs. With B = tokens split over DP = chips, that's 4 · B · D · k·FDP · C = per layer (k·F because a token only multiplies through its k activated experts).

The network clock: whatever bytes this scheme moves have to squeeze through the interconnect at W_ici = . Crucially, these two clocks can overlap when the implementation schedules the collective successfully: the network carries bytes while the MXU multiplies. Under that explicit assumption, a layer costs the max, not the sum. Communication that fits under the compute clock is hidden; communication that pokes out past it leaves silicon idle.

Try: with the toggle on weights, drag B = down and watch only the compute bar shrink — the network bar doesn't hear the batch size at all, so at some point the comms must poke out. Then flip to activations and drag again: now both bars move together, and no amount of batch will save you.

That toggle captures the core DP/FSDP/TP distinction: what travels?

• If what travels is weights — 2·D·E·F bytes (every one of the E experts, active or not), as in data parallelism and FSDP — the network clock is fixed, while the compute clock scales with your per-chip batch. Big batch ⇒ compute wins ⇒ comms hidden.
• If what travels is activations — 2·B·D bytes, as in tensor parallelism — batch size appears on both clocks and cancels. Hiding then depends on the model's shape (the width a token multiplies through: k·F) instead.

Now the roofline itself. In Part 1 of the original book, a single chip was compute-bound only when its arithmetic intensity — FLOPs per byte touched — beat the ratio of FLOP speed to memory bandwidth. The identical logic applies here, one level up, with the interconnect playing the role of memory. For weight-moving schemes, your FLOPs scale with B/DP and your bytes don't, so:

Try: drag the dot up the slope and feel where the ridge is. Then make the interconnect worse — drag W_ici = down — and watch the ridge slide right: a slower network demands a bigger per-chip batch before it can hide. Faster chips (drag C = up) do the same thing, which is why each hardware generation makes parallelism harder, not easier.

Data Parallelism

Syntax:

When your model fits on a single chip with even a tiny batch size (>240 tokens, so as to be compute-bound), you should always use simple data parallelism. Pure data parallelism splits our activations across any number of TPUs so long as the number of TPUs is smaller than our batch size. The forward pass involves no communication, but at the end of every step, each TPU performs an AllReduce on its local gradients to synchronize them before updating the parameters.

Note that the forward pass has no communication — it's all in the backward pass! The backward pass also has the great property that the AllReduces aren't in the “critical path”, meaning that each AllReduce can be performed whenever it's convenient and doesn't block you from performing subsequent operations. The overall communication cost can still bottleneck us if it exceeds our total compute cost, but it is much more forgiving from an implementation standpoint. We'll see that model/tensor parallelism doesn't have this property.✦adaptation In the primer's terms: because the AllReduce can be launched whenever convenient, the only question left is whether the total comms fits under the total compute — the roofline question, answered just below. Tensor parallelism's collectives, by contrast, block the very next matmul.

Why do this? Pure data parallelism reduces activation memory pressure by splitting our activations over the batch dimension, allowing us to almost arbitrarily increase batch size as long as we have more chips to split the batch dimension over. Especially during training when our activations often dominate our memory usage, this is very helpful.

Why not do this? Pure data parallelism does nothing to reduce memory pressure from model parameters or optimizer states, which means pure data parallelism is rarely useful for interesting models at scale where our parameters + optimizer state don't fit in a single TPU. To give a sense of scale, if we train with parameters in bf16 and optimizer state in fp32 with Adam◦Adam stores parameters, first order and second order accumulators. Since the params are in bfloat16 and optimizer state is in float32, this gives us 2 + 8 = 10 bytes per parameters., the largest model we can fit has TPU memory / 10 parameters, so e.g. on a TPUv5p chip with of HBM and pure data parallelism this is about parameters.

To make this useful for real models during training, we'll need to at least partly shard the model parameters or optimizer.

When do we become bottlenecked by communication? As we can see above, we have two AllReduces per layer, each of size 2DF (for bf16 weights). When does data parallelism make us communication bound?✦adaptation The network here carries weight gradients — 2 · D · E · F = per matrix (all E experts' gradients, not just the k a token used) — whose size doesn't care about the batch. This is the weight-moving roofline from the primer: a batch-blind comms cost that a big enough per-chip batch can always hide.

As in the table above, let C = per-chip FLOPs, W_ici = bidirectional per-axis ICI bandwidthone-way per-GPU NVLink egress, bandwidth, and DP = number of shards across which the batch is partitioned◦We assume this partitioning is done over an ICI mesh, so the relevant network bandwidth is W_ici.. Let's calculate the time required to perform the relevant matmuls, T_math, and the required communication time T_comms. Since this parallelism scheme requires no communication in the forward pass, we only need to calculate these quantities for the backwards pass.

Communication time: From a previous section we know that the time required to perform an AllReduce in a 1D mesh depends only on the total bytes of the array being AllReduced and the ICI bandwidth W_ici; specifically the AllReduce time is 2 · total bytes / W_ici. Since we need to AllReduce for both W_in and W_out, we have 2 AllReduces per layer. Each AllReduce is for a weight matrix, i.e. an array of DF parameters, or 2DF bytes. Putting this all together, the total time for the AllReduce in a single layer is

✦ Adaptation: This DP collective spans more than one NVLink domain, so the live clock uses scale-out-limited bandwidth , not the faster local NVLink egress.

Matmul time: Each layer comprises two matmuls in the forward pass, or four matmuls in the backwards pass, each of which requires 2(B/DP)DF FLOPs. Thus, for a single layer in the backward pass, we have

Since we overlap, the total time per layer is the max of these two quantities:

We become compute-bound when T_math/T_comms > 1, or when

The upshot is that, to remain compute-bound with data parallelism, we need the per-device batch size B/DP to exceed the ICI operational intensity, C/W_ici. This is ultimately a consequence of the fact that the computation time scales with the per-device batch size, while the communication time is independent of this quantity (since we are transferring model weights). Note the resemblance of the B/DP > C/W_ici condition to the single-device compute-bound rule B > 240; in that case as well, the rule came from the fact that computation time scaled with batch size while data-transfer size was (in the B ≪ F, D regime) independent of batch size.✦adaptation Try: drag DP = up and watch the compute bar in the meter below shrink toward the frozen comms bar — every doubling of chips halves T_math and does nothing to T_comms. The verdict flips the moment the per-chip batch dips below .adaptation With a Mixture-of-Experts model loaded (E > 1 — click e.g. DeepSeek-V3 in the intro's frontier-models table), a penalty paragraph from Chapter 12 appears just below the meter: the gradients move all E experts while the FLOPs touch only the activated width, which inflates this per-chip floor by E/k = .

Let's put in some real numbers to get a sense of scale. For TPUv5p, C = and W = for 1D data parallelism over ICI, so our batch size per chip must be at least to avoid being communication-bound.✦adaptation The famous 2,550 — the same constant the primer promised would keep reappearing. It's live here: change the hardware in the machine bar and this floor moves with it. Since we can do data parallelism over multiple axes, if we dedicate all three axes of a TPUv5p pod to pure data parallelism, we 3x our bandwidth W_ici and can scale down to only BS= per TPU or tokens per batch per pod (of chips)! This tells us that it's fairly hard to become bottlenecked by pure data parallelism!

Fully-Sharded Data Parallelism (FSDP)

Syntax:

Fully-sharded data parallelism (often called FSDP or ZeRO-sharding) splits the model optimizer states and weights across the data parallel shards and efficiently gathers and scatters them as needed. Compared to pure data parallelism, FSDP drastically reduces per-device memory usage and saves on backward pass FLOPs, with very minimal overhead.

You'll remember (from Section 3) that an AllReduce can be decomposed into an AllGather and a ReduceScatter. This means that, instead of doing the full gradient AllReduce for standard data parallelism, we can shard the weights and optimizer states across chips, AllGather them at each layer during the forward pass and ReduceScatter across the weights during the backward pass at no extra cost.

This is also called "ZeRO Sharding", from "Zero Redundancy Optimizer" since we don't perform any unnecessary compute or store any unnecessary state. ZeRO-{1,2,3} are used to refer to sharding the optimizer states, gradients, and weights in this way, respectively. Since all have the same communication cost◦Technically, FSDP adds communication in the forward pass that pure DP doesn't have, but this is in the same proportion as the backward pass so it should have no effect on the comms roofline. The key here is that ZeRO-3 turns a backward-pass AllReduce into an AllGather and a ReduceScatter, which have the same total comms volume., we can basically always do ZeRO-3 sharding, which shards the parameters, gradients, and optimizer states across a set of devices.

Why would we do this? Standard data parallelism involves a lot of duplicated work. Each TPU AllReduces the full gradient, then updates the full optimizer state (identical work on all TPUs), then updates the parameters (again, fully duplicated). For ZeRO sharding (sharding the gradients/optimizer state), instead of an AllReduce, you can ReduceScatter the gradients, update only your shard of the optimizer state, update a shard of the parameters, then AllGather the parameters as needed for your forward pass.✦adaptation Try: in the memory meter below, drag DP = down toward 1 and watch the meter fill up and overflow — at DP = 1 you're just pure DP on one chip's memory budget. Every doubling of DP halves the parameter and optimizer stripes.

When do we become bottlenecked by communication? Our relative FLOPs and comms costs are exactly the same as pure data parallelism, since each AllReduce in the backward pass has become an AllGather + ReduceScatter. Recall that an AllReduce is implemented as an AllGather and a ReduceScatter, each with half the cost. Here we model the forward pass since it has the same FLOPs-to-comms ratio as the backward pass:✦adaptation The chapter writes these equations for one mesh axis. The live line beneath them (and every meter on this page) spreads the collective over M_DP = mesh axes — the M_DP the chapter's takeaway below refers to — so T_comms is divided by M_DP.

✦ Adaptation: This FSDP collective crosses NVLink domains, so W_collective is the scale-out-limited .

Therefore, as with pure data-parallelism, we are compute bound when B / DP > C / W_collective, i.e. when the per-device batch size B/DP exceeds the collective fabric's operational intensity C/W_collective ( / = right now). This is great for us, because it means if our per-device batch size is big enough to be compute-bound for pure data-parallelism, we can — without worrying about leaving the compute-bound regime — simply upgrade to FSDP, saving ourselves a massive amount of parameter and optimizer state memory! Though we did have to add communication to the forward pass, this cost is immaterial since it just overlaps with forward-pass FLOPs.

Tensor Parallelism

Syntax:

(we use TP to eventually combine with FSDP)

In a fully-sharded data-parallel AllReduce we move the weights across chips. We can also shard the feedforward dimension of the model and move the activations during the layer — this is called “1D model parallelism” or Megatron sharding (Shoeybi et al. 2019). This can unlock a smaller efficient batch size per pod. The figure below shows an example of a single matrix sharded in this way:

As noted, In[B, D_TP] ·_D W_in[D, F_TP] ·_F W_out[F_TP, D] → Out[B, D_TP] means we have to gather our activations before the first matmul. This is cheaper than ZeRO sharding when the activations are smaller than the weights.✦adaptation Compare the two freights live, per layer in bf16: gathering activations moves 2 · B · D = , gathering weights moves 2 · D · E · F = — right now the lighter freight is . Try: drag B = down toward the per-group batches a mixed scheme would see (a few thousand tokens) and watch the verdict flip in the activations’ favor. This is typically true only with some amount of ZeRO sharding added (which reduces the size of the gather). This is one of the reasons we tend to mix ZeRO sharding and tensor parallelism.

One nice thing about tensor parallelism is that it interacts nicely with the two matrices in our Transformer forward pass. Naively, we would do an AllReduce after each of the two matrices. But here we first do In[B, D_TP] · W_in[D, F_TP] → Tmp[B, F_TP] and then Tmp[B, F_TP] · W_out[F_TP, D] → Out[B, D_TP]. This means we AllGather In at the beginning, and ReduceScatter Out at the end, rather than doing an AllReduce.✦adaptation And since an AllReduce is itself an AllGather plus a ReduceScatter, one AllGather in and one ReduceScatter out is half the bytes of the naive two-AllReduce plan.

How costly is this? Let's only model the forward pass - the backwards pass is just the transpose of each operation here. In 1D tensor parallelism we AllGather the activations before the first matmul, and ReduceScatter them after the second, sending two bytes at a time (bf16). Let's figure out when we're bottlenecked by communication.

Noting that we want compute cost to be greater than comms cost, we get:✦adaptation Notice that B · D appears in both clocks, so the batch cancels out of the ratio. This is the flat orange line from the primer: tensor parallelism’s compute-to-comms ratio is pinned at k · F / (TP · α_TP) = no matter the batch — a weight-moving scheme can hide its comms behind more tokens per chip, but no batch size can raise this bar.

✦ Adaptation: This TP collective spans NVLink domains. The live clock therefore uses scale-out-limited bandwidth rather than the local .

Thus for instance, for TPUv5p, C/W_ici = in bf16, so we can only do tensor parallelism up to TP < F / . When we have multiple ICI axes, our T_comms is reduced by a factor of M_TP, so we get TP < M_TP · F / .

Note that this doesn't depend on the precision of the computation, since e.g. for int8, on TPUv5p, C_int8/W_ici is instead of but the comms volume is also halved, so the two factors of two cancel.

Let's think about some examples:✦adaptation The buttons below load each model's real shape into the page's state — every number, meter, and verdict recomputes when you click one. Try: with a model loaded, set TP = to 8, then 16, then 32, and watch the verdicts. Or hold the model fixed and scrub C = : faster chips shrink TP_max = on the fabric carrying the current TP collective, which is why each hardware generation makes tensor parallelism a little harder to hide.

• On TPUv5p with LLaMA 3-70B with D = 8192, F ≈ 30,000, we can comfortably do 8-way tensor parallelism, but will be communication bound on 16 way tensor parallelism. The required F for 8-way model sharding is 20k.✦adaptation Check it live at the current hardware: with LLaMA 3-70B's F = 28,672 pinned, TP_max = -way; 8-way is , 16-way is , and the required F for 8-way sharding is 8 · α / M_TP = .
• For Gemma 7B, this page's ignore-gating convention uses the config's F = 24,576. On TPU v5p that gives TP_max ≈ 9.64, so 8-way fits but 16-way is communication-bound.Δedited The chapter reads “For Gemma 7B, F ≈ 50k, so we become communication bound with 19-way tensor parallelism. That means we could likely do 16-way and still see good performance.” Its ≈50k adds Gemma's two gated input projections together, which is inconsistent with this chapter's two-main-matmul model. At the currently selected TPU, the corrected TP_max is -way, and 16-way is .

Combining FSDP and Tensor Parallelism

Syntax:

The nice thing about FSDP and tensor parallelism is that they can be combined. By sharding W_in and W_out along both axes we both save memory and compute. Because we shard B along DP, we reduce the size of the model-parallel AllGathers, and because we shard F along TP, we reduce the communication overhead of FSDP. This means a combination of the two can get us to an even lower effective batch size than we saw above.

What's the right combination of FSDP and TP? A simple but key maxim is that FSDP moves weights and tensor parallelism moves activations. That means as our batch size shrinks (especially as we do more data parallelism), tensor parallelism becomes cheaper because our activations per-shard are smaller.✦adaptation This maxim is the primer's weights-move vs activations-move toggle made load-bearing: per layer in bf16, FSDP's freight is 2 · D · E · F/TP = of weights while TP's is 2 · B · D/DP = of activations — each scheme shrinks the other's bill.

• Tensor parallelism performs AllGather_TP([B_DP, D_TP]) which shrinks as DP grows.
• FSDP performs AllGather_DP([D_DP, F_TP]) which shrinks as TP grows.

Thus by combining both we can push our minimum batch size per replica down even more. We can calculate the optimal amount of FSDP and TP in the same way as above:

TPU closed form. Let DP be the number of chips dedicated to FSDP and TP be the number of chips dedicated to tensor parallelism. Let N be the total number of chips in our slice with N = DPTP. Let M_DP and M_TP be the number of mesh axes over which we do FSDP and TP respectively (these should roughly sum to 3). We'll purely model the forward pass since it has the most communication per FLOP. Then adding up the comms in the algorithm above, we have

And likewise our total FLOPs time is

✦ Adaptation — GPU topology correction: the TPU equations immediately above are not valid GPU substitutions. An outer FSDP reduction does not become TP times faster while TP remains inside one NVLink domain; the scale-out link still carries the reduction. The live clocks, meter, and explorers use Chapter 12's max(T_domain, T_scale-out) model. In the equations below, bytes = 4·D·E·F and G is the selected NVLink-domain size. The closed-form optimum below is therefore shown only on TPU; on GPU the explorer finds the topology-aware minimum directly.

Here's an interactive animation to play with this, showing the total compute time and communication time for different batch sizes:

You'll notice this generally agrees with the above (minimum around FSDP=256, TP=16), plus or minus some wiggle factor for some slight differences in the number of axes for each.✦adaptation The chapter's animation swept the FSDP/TP split itself; on this page that sweep is the earlier DP-axis explorer, whose optimum at the current state is DP_opt = → nearest power of two -way FSDP. The ±wiggle from mesh-axis bookkeeping is exactly the M_DP, M_TP factors you can scrub above.

Expert Parallelism

✦ This section is drawn from Chapter 12 (GPUs) of the same book and merged into this chapter's flow by this adaptation; condensed passages are marked. Its cost model is Chapter 12's switched GPU fabric (NVLink node + InfiniBand scale-out — see the GPU network model below). One naming change throughout: Chapter 12 calls the expert-parallel axis Z; this edition names every parallelism degree after its scheme, so that axis is rendered EP here. The routed/shared-expert split and the hardware-domain generalization of the H100-specific formula are AI-written adaptation material, labeled again at the live estimate.

As we've already noted above, Mixture of Expert (MoE) models come with E times more model weights with only k times more FLOPs, making data parallelism significantly harder.✦adaptation In Chapter 12 "noted above" pointed at its Data Parallelism section; on this page that passage lives at the end of Data Parallelism. This page's E and k count all experts, including always-on shared experts. Expert routing instead uses E_r = E − s routed experts and k_r = k − s routed selections, where s = shared experts. Right now, E_r = and k_r = . We can mitigate the routed weight cost by sharding along the expert dimension, i.e. W_in[E_EP, D, F]. To do the MLP block, we need to introduce 2x AllToAll to send our activations to the corresponding experts.

What does an AllToAll cost here? GPUs within a node have all-to-all connectivity, which makes AllToAlls, well, quite easy: each GPU just sends directly to the destination. For Mixture of Expert (MoE) models, we frequently want to do a sparse or ragged AllToAll, where we guarantee at most k_r of N shards on the output dimension are non-zero; the cost is reduced by k_r/N.✦adaptation Condensed from Chapter 12's intra-node collectives discussion (two paragraphs on dense and ragged AllToAlls, with the exact expected-occupancy footnote) — see Chapter 12 for the full derivation. The takeaway below is carried verbatim.

For the eight-GPU H100 node used in Chapter 12, the cost of this AllToAll_EP→kr([B, D, k]) if it spans multiple nodes is roughly T_AllToAll = 2 · B · D · (EP − 8)/(W · EP) · min(8 · k_r / EP, 1), so for pure expert parallelism we needΔedited The chapter reads: “the cost of this AllToAll_Z->k([B, D, k]) if it spans multiple nodes is roughly $T_\text{AllToAll} = 2 \cdot B \cdot D \cdot (Z-8)/Z \min(8 * k / Z, 1)$” — its inline expression omits the division by W, evidently a typo (the chapter's own displayed T_comms just below includes it); restored here.

✦ Adaptation: The live estimate is the substantive mash-up here. Within one NVLink domain it uses the finite ragged AllToAll cost from Chapter 12 rather than calling that transfer free. Beyond the domain it takes the slower of the local switched-fabric component and the chapter's scale-out component, replacing the H100-specific 8 with the selected hardware's domain size. Shared experts remain in the total k-wide compute, but never become routed AllToAll destinations.

For that H100 case, Chapter 12 concludes that we either need k_r > EP/8 with F > α · (EP − 8)/k_r or EP ≫ k_r and F > 8 · α, where α = C/W. This gives you two domains in which expert parallelism is possible, one with a small amount of expert parallelism (roughly 2-node) and small F, or one with large F and EP arbitrarily large (up to E_r-way expert parallelism).

You'll see both cases in practice, either a small amount of expert-parallelism (like DeepSeek v3 which has very small F and relatively small, restricted cross-node expert parallelism), or models with large F, in which case we can do significant cross-node EP alongside TP.

Pipelining

You'll probably notice we've avoided talking about pipelining at all in the previous sections. Pipelining is a dominant strategy for GPU parallelism that is somewhat less essential on TPUs. Briefly, pipelined training involves splitting the layers of a model across multiple devices and passing the activations between pipeline stages during the forward and backward pass.✦adaptation On this page the split is live: with L = layers over PP = pipeline stages (scrubbable below), each device owns about consecutive layers. The algorithm is something like:

batch_size = 32
d_model = 128
d_ff = 4 * d_model

num_layers = len(jax.devices())

key = jax.random.PRNGKey(0)

# Pretend each layer is just a single matmul.
x = jax.random.normal(key, (batch_size, d_model))
weights = jax.random.normal(key, (num_layers, d_model, d_model))

def layer_fn(x, weight):
  return x @ weight

# Assume we have num_layers == num_pipeline_stages
intermediates = [x]
for i in range(num_layers):
  x = layer_fn(x, weights[i])
  intermediates.append(x)

  if i != num_layers - 1:
    x = jax.device_put(x, jax.devices()[i+1])

def loss_fn(batch):
  return jnp.mean(batch ** 2)  # make up some fake loss function

loss, dx = jax.value_and_grad(loss_fn)(x)

for i in range(num_layers - 1, -1, -1):
  _, f_vjp = jax.vjp(layer_fn, intermediates[i], weights[i])
  dx, dw = f_vjp(dx)  # compute the jvp dx @ J(L)(x[i], W[i])
  weights[i] = weights[i] - 0.01 * dw  # update our weights

  if i != 0:
    dx = jax.device_put(dx, jax.devices()[i-1])

Why is this a good idea? Pipelining is great for many reasons: it has a low communication cost between pipeline stages, meaning you can train very large models even with low bandwidth interconnects. This is often very useful on GPUs since they are not densely connected by ICI in the way TPUs are.✦adaptation The chapter doesn't quantify "low communication cost," so the check below is ours. A stage-boundary hop is a single point-to-point copy of one activation block — 2D = per token in bf16 — and it's the same roofline question as ever: does the hop fit under one stage's compute clock? The line below runs the numbers for one microbatch.

Why is this difficult/annoying? You might have noticed in the pseudocode above that TPU 0 is almost always idle! It's only doing work on the very first and last step of the pipeline. The period of idleness is called a pipeline bubble and is very annoying to deal with. Typically we try to mitigate this first with microbatching, which sends M_micro = small batches through the PP = -stage pipeline, keeping TPU 0 utilized for at least a larger fraction of the total step time.

A second approach is to carefully overlap the forward matmul W_i @ x_i, the backward dx matmul W_i @ ∂L/∂x_i+1, and the dW matmul ∂L/∂x_i+1 @ x_i. Since each of these requires some FLOPs, we can overlap them to fully hide the bubble. Here's our live stand-in for the plot from the recent DeepSeek v3 paper showing their "bubble-free" pipeline schedule:✦adaptation Toggle the widget's mode: naive is GPipe, 1F1B interleaves one forward with one backward (same bubble, far less activation memory held live), and overlap-dW is the DeepSeek-v3-style schedule — rush every ∂x result down the pipeline to unblock neighbors, and drop the deferred ∂W matmuls into slots that would otherwise sit idle.

Because it is less critical for TPUs (which have larger interconnected pods), we won't delve into this as deeply, but it's a good exercise to understand the key pipelining bottlenecks.✦adaptation The condensed picture: pipelining's communication is one activation hop per stage boundary, so it thrives on weak interconnects and dominates GPU training; the price is the bubble — currently of each device's time — which microbatching shrinks and careful ∂x/∂W overlap can erase.

The GPU Network Model

✦ This section is an addition of this adaptation, drawing its text from Chapter 12 (GPUs) of the same book; condensed passages are marked. Chapter 12's per-scheme roofline derivations are not repeated here — they re-derive what this chapter already derived, so they are merged into the scheme sections above (the MoE penalty into Data Parallelism, the TP bound into Tensor Parallelism, expert parallelism into its own section, and the pipelining reasons into Pipelining). What remains here is the network model itself: the fabric, its bandwidths, what collectives cost on it, and the worked examples.

Now let's look at what this has all been building towards: understanding rooflines for LLM scaling on GPU. This is to complement the TPU training chapter here. As we did there, the goal here is to look at the total T_math and T_comms for different parallelism strategies and understand at what point T_comms > T_math. As before, we consider only the MLP block with operations

where B is the global batch size in tokens (i.e. B = batch size · sequence length).

Here we'll reproduce the table from Chapter 12 showing effective bandwidths at both the GPU and node level:

Let's look at the compute communication rooflines as we did for TPUs for data parallelism, tensor parallelism, pipeline parallelism, expert parallelism, and combinations thereof. For the rest of this section we'll focus on H100 rooflines for specific calculations. GB200-NVL72 has the same general rooflines but because we have a larger node egress bandwidth, we can sometimes be bottlenecked at the node level instead. The scheme derivations are merged into their corresponding sections above; below are the bounds they land on here.

Here is the mapping used by the live GPU rooflines: read W_ici as the per-GPU egress bandwidth into the NVLink switch fabric (450 GB/s on H100, 900 GB/s on B200); read a pod as an NVLink domain (8 GPUs per node, or 72 on GB200 NVL72); read W_dcn as each GPU's share of its node's InfiniBand egress (400e9/8 = 50 GB/s on H100 and B200, and likewise 3600e9/72 on GB200); and, because a switched fabric has no mesh axes, set M_DP = M_TP = 1. Then C/W_ici = is the in-node ridge, C/W_dcn = the per-GPU scale-out ratio. DP/FSDP/TP select the local or scale-out collective bandwidth from their degree. Mixed FSDP+TP needs one extra correction: Chapter 12's hierarchical max(T_domain, T_scale-out) rule, because inner TP does not accelerate an outer reduction until it spans more than one NVLink domain. The mixed meters and explorers implement that rule.

Where the bounds land on this fabric

For data parallelism and ZeRO sharding, the compute-bound rule derived in Data Parallelism — B/DP > C/W_collective — is reused unchanged, where W_collective is either the GPU or node level egress bandwidth depending on whether we're sharding within a node or across nodes. Thus:

• Within a node, we just need the per-GPU token batch size > / = .
• Within an SU or at the spine level, BS > / = .

This is quite a bit higher than on a TPU, where the number is 850 with all three axes. On the H100 scale-out fabric the dense asymptotic floor is 990e12/400e9 = 2,475 tokens per GPU, so 16,384 GPUs would require about 40.6M tokens before the small-ring and model-parallel refinements; Llama 3.1 405B used 16M. Chapter 12 quoted a 3,300-token H800 baseline from an unsupported 300 GB/s figure. The reconciled H800 local-link dense baseline is 4,950 in spec mode (990e12/200e9) and about 4,517 in this page's measured mode. Those are not a model of DeepSeek's full sparse run: its EP, PP, and 2-way DP alter the outer reduction. DeepSeek reports a pretraining batch schedule from 3,072 to 15,360 sequences at a 4K maximum sequence length — about 12.6M to 62.9M tokens, with 62.9M at steady state.Δedited The source says H800 has 300 GB/s and “in practice, they used 4M”. H800 is 200 GB/s per direction by the reconciled spec, DeepSeek reports 160 GB/s measured, and its report gives the sequence-batch schedule above. See the hardware table.

Small-DP correction. The asymptotic ridge above omits the ring factor. With X scale-out domains, the exact dense condition is B/N > (C/W_collective) · (X−1)/X (and ×E/k for the equal-width MoE model). For exactly two domains the floor is halved, which is why 2-way data parallelism appears so often.

For tensor parallelism, the bound from Tensor Parallelism — TP < F · W_collective / C — gives about F/-way TP within a node, and pipelining — whose Chapter 12 treatment now lives in Pipelining under GPU presets — is basically free from a communication standpoint. As with above, we get an extra 2X bandwidth when we span exactly 2 fast-fabric domains. For the chapter's eight-GPU H100 node this usually permits 16-way TP; generalized to a domain of G = , the refinement is F > · (TP − G).

What collectives cost on this fabric

Beyond the node level: to a first approximation, because we have full bisection bandwidth, the cost of an AllGather or ReduceScatter is roughly the buffer size in bytes divided by the node egress bandwidth (400GB/s on H100) regardless of any of the details of the tree reduction.

With in-network reductions enabled and using pure data parallelism, theoretically we have 2x the AllReduce bandwidth, which would halve both of these numbers. However, in practice the benefit is closer to 30%, which only really makes up for the fact that we typically struggle to reach the reported numbers. Furthermore, because pure data parallelism is rarely useful, this basically doesn't matter in practice.

Examples

What does DeepSeek do? For reference, DeepSeek V3 is trained with 2048 H800 GPUs with:✦adaptation These two worked examples are Chapter 12's. The buttons load their cited hardware, model shape, batch, and listed parallelism degrees. Because the page's generic chip count is DP·TP, widgets that do not explicitly model EP or PP should be read as component-level views, not as a reconstruction of the full training run.

• 64-way Expert Parallelism (EP) spanning 8 nodes
• 16-way Pipeline Parallelism (PP)
• 2-way ZeRO-1 Data Parallelism (DP)

They had a steady state batch size of 4096 · 15360 = 62,914,560 tokens, or 30k tokens per GPU. You can see that this is already quite large, but their model is also very sparse (k=8, E=256) so you need a fairly large batch size. You can see that with 64-way EP and 16-way PP, we end up with 1024-way model parallelism in total, which means the AllReduce is done at the spine level, and because it's only 2-way, we end up with 2 / (2 − 1) = 2 times more bandwidth in practice. This also helps reduce the cost of the final data-parallel AllReduce overlapping with the final pipeline stages.

What does Llama 3.1 405B do? Llama 3.1 405B trains with a BS of 16M tokens on 16,384 H100 GPUs, or about 977 tokens per GPU. They do:

• 8-way Tensor Parallelism within a node (TP)
• 16-way Pipeline Parallelism (PP)
• 128-way ZeRO-1 Data Parallelism

The decomposition is 8 TP · 16 PP · 128 DP = 16,384 GPUs. This is also a dense model so in general these things are pretty trivial. The 16-way PP reduces the cost of the data parallel AllReduce by 16x, which helps us reduce the critical batch size.

TLDR of LLM scaling on GPUs✦adaptation Restored from Chapter 12's omitted “TLDR of LLM Scaling on GPUs” and practical recipe; condensed to avoid repeating the full derivations already merged into the scheme sections.

• DP or FSDP needs a local batch of roughly 2,500 dense-model tokens per GPU on H100-scale fabrics; MoEs multiply that floor by E/k, while a small DP degree benefits from the ring correction above.
• TP is usually compute-bound only within one NVLink domain, or at most roughly two. NVL72 expands that local domain but does not remove the topology check.
• Model parallelism that spans domains can reduce the outer FSDP cost, but on GPUs the reduction tracks domains spanned, not merely the inner sharding degree.
• PP has low byte cost if its scheduling complexity, bubbles, and delayed gradient reductions are handled successfully.

Practical recipe: smaller dense models can use aggressive FSDP when batch permits; larger dense models commonly combine one- or two-domain TP with many-stage PP and DP; MoEs add EP, generally preferring it to TP while keeping latency-sensitive collectives within as few domains as practical.

Some Problems to Work

Let's use LLaMA-2 13B as a basic model for this section. Here are the model details:✦adaptation Every value in this table is scrubbable, and every answer below is computed from it live: the exercises grade themselves against whatever model you dial in; each question's stated givens (a batch size, a chip count) stay pinned, the way a problem set's givens should. The preset button under the table restores the chapter's LLaMA-2 13B.

LLaMA-2 has separate embedding and output matrices and a gated MLP block.

Question 1: How many parameters does LLaMA-2 13B have (I know that's silly but do the math)? Note that, as in Transformer Math, LLaMA-3 has 3 big FFW matrices, two up-projection and one down-projection. We ignored the two "gating" einsum matrices in this section, but they behave the same as W_in in this section.

Question 2: Let's assume we're training with BS= tokens and using Adam. Ignoring parallelism for a moment, how much total memory is used by the model's parameters, optimizer state, and activations? Assume we store the parameters in bf16 and the optimizer state in fp32 and checkpoint activations three times per layer (after the three big matmuls).

Question 3: Assume we want to train with 32k sequence length and a total batch size of 3M tokens on a TPUv5p 16x16x16 slice. Assume we want to use bfloat16 weights and a float32 optimizer, as above.

1. Can we use pure data parallelism? Why or why not?
2. Can we use pure FSDP? Why or why not? With pure FSDP, how much memory will be used per device (assume we do gradient checkpointing only after the 3 big FFW matrices).
3. Can we use mixed FSDP + tensor parallelism? Why or why not? If so, what should DP and TP be? How much memory will be stored per device? Using only roofline FLOPs estimates and ignoring attention, how long will each training step take at MFU?

That's it for Part 5! For Part 6, which applies this content to real LLaMA models, click here!✦adaptation Try: open answer (3) above and drag F = down. A skinnier MLP raises the mixed-scheme threshold α²/2F (tensor parallelism has less F to hide behind) and shifts DP_opt — watch the verdict pill flip when the model gets too narrow for even the mixed scheme to save this small batch.

Appendix

Appendix A: Deriving the backward pass comms

Above, we simplified the Transformer layer forward pass as Out[B, D] = In[B, D] ·_D W_in[D, F] ·_F W_out[F, D]. How do we derive the comms necessary for the backwards pass?

This follows fairly naturally from the rule in the previous section for a single matmul Y = X · A:✦adaptation In this appendix X and Y are the input and output matrices of a generic matmul — the chapter's letters, kept as-is since this edition's mesh axes go by DP and TP, so nothing collides.

Using this, we get the following formulas (letting Tmp[B, F] stand for In[B, D] · W_in[D, F]):

Note that these formulas are mathematical statements, with no mention of sharding. The job of the backwards pass is to compute these four quantities. So to figure out the comms necessary, we just take the shardings of all the quantities which are to be matmulled in the four equations above (Tmp, dOut, W_out, W_in), which are specified by our parallelization scheme, and use the rules of sharded matmuls to figure out what comms we have to do. Note that dOut is sharded in the same way as Out.

Look back at Part 4: Transformer Math, continue to Part 6: Applied Training, which works this content through real LLaMA models, or revisit the original chapter this page adapts.