# How to Parallelize a Transformer for Training > Source: > Published: 2026-08-18 19:55:18+00:00 ✦ We begin with the original dense TPU schemes — data parallelism, FSDP, tensor parallelism, their mixed form, and pipelining — then splice in the GPU fabric model and expert parallelism for MoEs. For each, we ask when communication becomes the bottleneck. (This summary is the adaptation’s; the chapter’s own dek described its four dense schemes.) **This page is a working model, not a description of one.** Every green number can be dragged left or right, or double-clicked to type an exact value. Every blue number is computed live from the green ones — try it here: drag the batch and watch the per-chip batch follow (hover any blue number for its formula). They share one model-and-hardware state, so a change made anywhere propagates everywhere. Parallelism degrees remain scheme-local: the dense mixed group uses N = DP·TP, while EP and PP are modeled in their own sections; composite worked examples state their full product explicitly. And scrub without fear — restores every scrubbed number to its default while keeping your model, hardware, and spec/measured picks (it's the same button as in the top bar, which lights up orange whenever a scrub has strayed), any single number reverts on its own when you double-click it and commit it blank, and the browser's back button walks through your earlier configurations. **Whose words are you reading?** Source passages come from the original TPU and GPU chapters (© 2022 Maruan Al-Shedivat, © 2025 Google LLC, [MIT license](LICENSE-scaling-book.txt)); AI-authored departures are explicitly labeled, with these conventions: wherever the chapter printed a fixed number, this page computes it live (these in-place swaps aren't individually marked); the interactive figures and their captions replace the original static figures; ✦ margin notes and passages explicitly labeled as adaptation are AI-written editorial voice — the initial edition was built by Fable (Anthropic) and this adversarial review and its corrections were performed by OpenAI Codex — including instructions, asides, and the new [roofline primer](#roofline); the [expert parallelism](#expert-parallelism) and [GPU network](#gpus) sections instead mash up Chapter 12 source passages, with their AI-written connective and adaptation prose labeled by the same convention; the chapter's single-letter mesh-axis names are rendered as named parallelism degrees throughout — its X is DP, its Y is TP, the pipelining section's Z is PP, and chapter 12's expert axis Z is EP (a global substitution; each is its own scrubbable variable, adjusted in the text where its section uses it); under a GPU preset the hardware vocabulary follows suit — TPU→GPU, ICI→NVLink, DCN→InfiniBand, pod→node, MXU→tensor core — so the article reads as one consistent machine, and any TPU preset restores the chapter's exact words (sentences that deliberately *compare* the two never swap); content woven into the chapter's text by this edition carries a dotted underline like that (and splice edits beside it that standard quotation practice would allow — a bracket or an ellipsis — go unmarked); where a sentence had to be altered to host a live element, a Δ margin note quotes the original and states the change; and where this edition's additions make a chapter statement inaccurate as written, an italic (Ed: …) interjection corrects it in place. 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](https://jax-ml.github.io/scaling-book/sharding/), 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](#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. | Notation | Meaning (model parameters) | Live value | |---|---|---| D | dmodel (the hidden dimension/residual stream dim) | | F | dff (the feed-forward dimension)adaptation F convention (everywhere): the width of one expert (= dff when dense); math runs through k·F, weights hold E·F, and the chapter's equations are the E = k = 1 case (Chapter 12's resolution). One honest limitation: models that mix dense and MoE blocks have two genuinely different F's — DeepSeek-V3 runs its first three layers dense at a much wider width — and this page approximates such models as uniformly MoE. Hover any F for the live widths. | | B | Batch dimension (number of tokens in the batch; total, not per-device) | | | T | Sequence length | — | L | Number of layers in the model | | Notation | Meaning (hardware characteristic) | Live value | |---|---|---| C | FLOPS/s per chip | | W | Network bandwidth (bidirectional per TPU mesh axisone-way GPU or node egress, often subscripted as e.g. W or iciW)dcn | ici · dcn | DP | Number of chips along the data-parallel mesh axis (the chapter's X) | | TP | Number of chips along an alternate, tensor-parallel mesh axis (the chapter's Y) | | Z | Number of chips along a third mesh axis, labeled Z | — | PP | Pipeline stages (the pipelining section's Z) | | EP | Expert-parallel degree (chapter 12's Z; see the expert-parallelism section) | The chapter's examples are dense LLaMA-era models; the frontier has since gone Mixture-of-Experts.Shapes from each model's published `config.json` on Hugging Face; parameter totals from its safetensors metadata. Retrieved August 2026. E and k count shared experts, so k·*F* is the activated width for the architectures represented by the live presets; column headers explain each field. The dense models from the top-bar dropdown lead the table for contrast, and whichever model is loaded shows its row in live green — scrub it right here. | Model | params | D | F | act. k·F | L | E | k | |---|---|---|---|---|---|---|---| | (chapter default) | 70.6B | 8,192 | 28,672 | 28,672 | 80 | 1 | 1 | | 13.0B | 5,120 | 13,824 | 13,824 | 40 | 1 | 1 | | | 8.54B | 3,072 | 24,576 | 24,576 | 28 | 1 | 1 | | | Counting example: 256 routed + 1 shared expert → E 257; top-8 + shared → k 9. Its first three layers are actually dense (see the F-convention note above). | 685B | 7,168 | 2,048 | 18,432 | 61 | 257 | 8+1 | | Kimi K3 (reference only)K3 is not a live preset because its routed experts operate after a projection from residual D = 7,168 into a 3,584-wide latent space. Its routed-expert intermediate width is F = 3,072. The page's single D×F expert model cannot represent both dimensions faithfully. | 2.78T | 7,168 | 3,072 | 55,296 | 93 | 896+2 | 16+2 | | 753B | 6,144 | 2,048 | 18,432 | 78 | 257 | 8+1 | | | 1.60T | 7,168 | 3,072 | 21,504 | 61 | 385 | 6+1 | | | 2.45T | 8,192 | 2,048 | 22,528 | 92 | 513 | 10+1 | | | 952B | 6,144 | 3,072 | 24,576 | 66 | 258 | 6+2 | | | 427B | 6,144 | 3,072 | 15,360 | 60 | 129 | 4+1 | Click a supported model to load its shape (D, F, L, E, k) into the page's shared state (the top bar follows); click a column header to sort. F = per-expert width; act. k·F = activated width per token; E / k = total / activated experts, counting shared. Across the supported live MoE presets, per-expert F is just 2,048 or 3,072, and activated width k·F clusters between 15k and 25k even as total parameter counts span hundreds of billions to trillions. Since the tensor-parallelism bound later in this chapter scales with the activated width k·F, that clustering is why the TP limits look so similar across the supported frontier presets. K3 is retained as a reference row, but its latent-MoE shape is deliberately not loaded into these formulas. Every hardware number this page computes with, spec and sustained, with its source.Full citations live in [ SOURCES.md](SOURCES.md) alongside this page — every value traces to a vendor spec sheet, a published measurement, or the book's own benchmarks; retrieved 2026-08-17; click any cell to pin its citation and follow the source link. Methodology for the synthesized numbers: NVIDIA datasheets headline | Hardware | C (dense bf16) | × sust. | W link | × achv. | W scale-out | HBM | |---|---|---|---|---|---|---| | 459 TF | ≈0.72 | 180 GB/s | ≈0.95 | 6.25 GB/s | 96 GB | | | 197 TF | ≈0.67 | 90 GB/s | ≈0.95 | 3.13 GB/s | 16 GB | | | 989 TF | 0.73 | 450 GB/s | 0.82 | 50 GB/s | 80 GB | | | 2.25 PF | 0.69 | 900 GB/s | ≈0.82 | 50 GB/s | 180 GB | | | 2.5 PF | ≈0.70 | 900 GB/s | ≈0.82 | ≈50 GB/s | 186 GB | | | 2.5 PF | ≈0.70 | 900 GB/s | ≈0.82 | 100 GB/s | 288 GB | | | 989 TF | ≈0.73 | 200 GB/s | 0.80 | 50 GB/s | 80 GB | 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](https://jax-ml.github.io/scaling-book/transformers/). 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 Pw ≈ 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). **Forward pass:** need to compute Loss[B] **Backward pass:** need to compute dWout[F, D], dWin[D, F] We provide this for comparison to the algorithms with communication added. 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, Wout, 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 **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 (dff). 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](#roofline) — building the one picture that answers that question for all four. ✦ This entire section is an addition of this adaptation — the chapter's ideas, our framing. The original text resumes at [Data Parallelism](#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·*F**DP* · *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 Wici = . 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?* Now the roofline itself. In [Part 1 of the original book](https://jax-ml.github.io/scaling-book/roofline/), 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 Wici = 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. **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.** **Pure Data Parallelism Algorithm:** **Forward pass:** need to compute Loss[BDP] **Backward pass:** need to compute dWout[F, D], dWin[D, F]adaptation The {UDP} annotation below marks a result that is *unreduced* over the *DP* axis: each chip holds a partial sum from its own slice of the batch. We ignore the details of the loss function and abbreviate Tmp = Win · In. Note that, although our final loss is the average AllReduce(Loss[BDP]), we only need to compute the AllReduce on the backward pass when averaging weight gradients. 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](#roofline)'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 AdamAdam 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 2*D**F* (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](#roofline): 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* = *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 **✦ 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*)*D**F* 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 Tmath/Tcomms > 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 For a Mixture of Experts (MoE) model, where we have *E* experts and *k* experts per token, this increases to which inflates the per-GPU token batch size by a factor of *E*/*k*, i.e. For example, the new OpenAI OSS model with k=4 and E=128, this increases to 32 · 2475 = 79,200 across nodes, a kind of ridiculously high number.adaptation Kept as the chapter's static example (its 2475 is the H100 cross-node ridge from [the GPU section](#gpus)). At your current hardware and model, the same computation reads (E/k) · C/Wcollective = tokens per chip. Expert parallelism — sharding the experts themselves, so gradients stop crossing the whole DP axis — is the standard escape; it gets [its own section](#expert-parallelism) below. 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](#roofline) 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)! **Syntax:** Fully-sharded data parallelism (often called FSDP or [ZeRO-sharding](https://arxiv.org/abs/1910.02054)) 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](https://jax-ml.github.io/scaling-book/sharding/)) 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. **Fully-Sharded Data Parallelism (FSDP):** **Forward pass:** need to compute Loss[BDP] **Backward pass:** need to compute dWout[F, DDP], dWin[DDP, F] 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 costTechnically, 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 **✦ Adaptation:** This FSDP collective crosses NVLink domains, so Wcollective 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 For example, borrowing only DeepSeek-V2's reported batch size as a *dense-model thought experiment* (this calculation does not model its expert parallelism), take a batch size of ~40M tokens.adaptation This qualifier is added because the source imports DeepSeek-V2's batch into a dense FSDP calculation; it does not model that MoE's expert parallelism. **This would allow us to scale to roughly chips, or around TPUv5 pods, before we hit a bandwidth limit.**adaptation Load the DeepSeek scenario with the button below, then make the batch your own: the mini-calculator that follows is an addition of this edition. With *B* = tokens, FSDP scales to *DP max* chips before hitting the bandwidth limit. For LLaMA-3 70B, which was trained for approximately (15e12 · 70e9 · 6) FLOPs, we could split a batch of tokens over roughly *B* / (α / 3) = chips (roughly pods of chips), each with FLOPs running at peak FLOPs utilization (often called MFU), and **train it in approximately **. Not bad! But let's explore how we can do better.adaptation The chapter's numbers (16M tokens, 18,823 chips, 17 days) are one point of this live sentence — the recipe button below restores them. Then drag *B* and watch chips and wall-clock trade off: a bigger batch rides the same ridge on more chips and finishes sooner, which is exactly why the labs fight for every doubling of critical batch size. The equation below, also an addition, shows the wall-clock arithmetic. **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](https://arxiv.org/abs/1909.08053)). 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 Win[D, FTP] ·F Wout[FTP, D] → Out[B, DTP] 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 · **Tensor Parallelism:** adaptation Watch the phrase *on critical path*. With pure data parallelism the AllReduce happened after the loss was already computed, so the network could grind away while the chips moved on. Here the matmuls cannot start until the gathers finish — these collectives sit squarely in the layer’s serial path. **Forward pass:** need to compute Loss[B] **Backward pass:** need to compute dWout[FTP, D], dWin[D, FTP] 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] · Win[D, FTP] → Tmp[B, FTP]** and then **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](#roofline): 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 **Note that this doesn't depend on the precision of the computation**, since e.g. for int8, on TPUv5p, Cint8/*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. **Syntax:** The nice thing about FSDP and tensor parallelism is that they can be combined. By sharding **W in** and **Forward pass:** need to compute Loss[B]adaptation Count what actually sits on the critical path: on the *TP* axis, one AllGather in (step 1) and one ReduceScatter out (step 6) of activation bytes; the weight gathers on the *DP* axis (steps 2 and 4) can be prefetched. Those two *TP*-axis activation hops and two *DP*-axis weight hops are exactly the 2 · 2 factors in TTP comms and TFSDP comms below. **Backward pass:** need to compute dWout[FTP, DDP], dWin[DDP, FTP] **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](#roofline)'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. 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* = *DP**TP*. Let *M DP* and 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(Tdomain, Tscale-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. To simplify the analysis, we make two assumptions: first, we allow *DP* and *TP* to take on non-integer values (as long as they are positive and satisfy *DP**TP* = *N*); second, we assume that we can fully overlap comms on the *DP* and *TP* axis with each other. Under the second assumption, the total comms time is Before we ask under what conditions we'll be compute-bound, let's find the optimal values for *DP* and *TP* to minimize our total communication. Since our FLOPs is independent of *DP* and *TP*, the optimal settings are those that simply minimize comms. To do this, let's write Tcomms above in terms of *DP* and *N* (which is held fixed, as it's the number of chips in our system) rather than *DP* and *TP*: Because TFSDP comms is monotonically increasing in *DP*, and TTP comms is monotonically decreasing in *DP*, the maximum must be minimized when TFSDP comms = TTP comms, which occurs when This is super useful! This tells us, for a given *B*, *F*, and *N*, what amount of FSDP is optimal. Let's get a sense of scale. Plugging in realistic values, namely *N* = 64 (corresponding to a 4x4x4 array of chips), *B* = 48,000, *F* = 32768, gives roughly *DP* ≈ . So we would choose *DP* to be 16 and *TP* to be 4, close to our calculated optimum.adaptation The chapter rounds this to ≈13.9; the pinned live value here is √(48,000 · 2 · 64 / 32,768) exactly. And at whatever is loaded right now, *DP opt* = . Press the first button below to load the chapter's exact scenario into the whole page. Now let's return to the question we've been asking of all our parallelism strategies: **under what conditions will we be compute-bound?** Since we can overlap FLOPs and comms, we are compute-bound whenadaptation Same question as [the primer](#roofline)'s: does the slower of the two comms clocks fit under the compute clock? By letting α ≡ *C* / *W ici*, the ICI arithmetic intensity, we can simplify: Since we calculated *DP opt* to make the LHS maximum equal, we can just plug it into either side (noting that Further simplifying, we find that where the left-hand-side is proportional to the communication time and the right-hand-side is proportional to the computation time. Note that while the computation time scales linearly with the batch size (as it does regardless of parallelism), the communication time scales as the square root of the batch size. The ratio of the computation to communication time thus also scales as the square root of the batch size: To ensure that this ratio is greater than one so we are compute bound, we require To get approximate numbers, again plug in *F* = 32,768, α = 2550, and *M DP* Below we plot the ratio of FLOPs to comms time for mixed FSDP + TP, comparing it both to only tensor parallelism (TP) and only data parallelism (FSDP), on a representative 4x4x4 chip array. While pure FSDP parallelism dominates for very large batch sizes, in the regime where batch size over number of chips is between roughly 100 and 850, a mixed FSDP + TP strategy is required in order to be compute-bound.adaptation The live chart below plays this figure's role: flip its view toggle to *ratio* to see Tmath/Tcomms for all three schemes, where any curve above 1 is compute-bound. It is drawn at the page's current *N* = chips — press the 4×4×4 chapter-example preset above to reproduce the chapter's exact frame. Here's another example of TPU v5p 16x16x16 showing the FLOPs and comms time as a function of batch size for different sharding schemes.adaptation That second figure is the same chart in *absolute-times* view. A 16x16x16 slice is 4096 chips — exactly the page's default *DP*·*TP* = 512 · 8, so the “back to page defaults” preset above reproduces it. The black curve is the amount of time spent on model FLOPs, meaning any batch size where this is lower than all comms costs is strictly comms bound. You'll notice the black curve intersects the aqua curve at about 4e5, as predicted.adaptation On the live chart that crossing sits at *B* = *N*·α²·*E*/(*M DP*· 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 ✦ This section is drawn from [Chapter 12 (GPUs)](https://jax-ml.github.io/scaling-book/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](#gpus) 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](#data-parallelism). This page's *E* and *k* count all experts, including always-on shared experts. Expert routing instead uses Er = E − s routed experts and kr = 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. Win[EEP, 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 kr/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 AllToAllEP→k r([B, D, k]) if it spans multiple nodes is roughly T **✦ 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 kr > EP/8 with *F* > α · (EP − 8)/kr or EP ≫ kr 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. 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: This pseudocode should run on a Cloud TPU VM. While it's not very efficient or realistic, it gives you a sense how data is being propagated across devices. ``` 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 — 2*D* = per token in bf16 — and it's the [same roofline question](#roofline) 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 The overall communication cost of pipelining is tiny: with *N MB* microbatches and Since we are dividing by *N layers*, this is vastly smaller than any of the other costs. In other words, from a communication standpoint, pipelining is basically free. So why don't we just do pipelining? There are a few reasons: (1) **Code complexity:** pipelining fits poorly into automatic parallelism frameworks (like XLA's GSPMD), because microbatching and custom zero-bubble schedules change the structure of the program.adaptation Condensed to one sentence — see Chapter 12 for the full paragraph. (2) **Pipelining makes data parallelism and FSDP hard:** probably the biggest reason not to do pipelining is that it plays badly with FSDP and data parallelism. ZeRO-3 sharding in particular works badly, since it requires us to AllGather the weights on every microbatch which doesn't work when we have only *B* / Nmicrobatches tokens to amortize the AllGather cost. Furthermore, during the backward pass, *we can't AllReduce or ReduceScatter the gradients until the last microbatch has passed a given stage, which means we have significant non-overlapped communication time.* (3) **Pipeline bubbles and step imbalance:** naive pipeline schedules leave stages idle in bubbles, and passing activations from stage to stage on the critical path shifts stages relative to each other and adds overhead.adaptation Condensed to one sentence — see Chapter 12, and the live bubble math just above. There are workarounds for each of these issues, but they tend to be complicated to implement and difficult to maintain; pipelining remains a technique with low communication cost relative to other methods. A second approach is to carefully overlap the forward matmul Wi @ xi, the backward dx matmul Wi @ ∂L/∂xi+1, and the dW matmul ∂L/∂xi+1 @ xi. 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](https://arxiv.org/abs/2412.19437) 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 largest possible TPU slice is a TPU v5p SuperPod with 8960 chips (and 2240 hosts). When we want to scale beyond this size, we need to cross the Data-Center Networking (DCN) boundary. Each TPU host comes equipped with one or several NICs (Network Interface Cards) that connect the host to other TPU v5p pods over Ethernet. As noted in the [TPU Section](https://jax-ml.github.io/scaling-book/tpus/), each host has about 200Gbps (25GB/s) of full-duplex DCN bandwidth, which is about full-duplex (egress) bandwidth per TPU.adaptation Per the TPU chapter of the original book: each v5p host serves 4 chips, so 25 GB/s per host ÷ 4 ≈ 6.25 GB/s of egress per chip.adaptation The chapter printed 6.25GB/s; here *W dcn* is scrubbable — drag it and this whole section (ridge included) recomputes. The hardware presets in the top bar set it per machine. Typically, when scaling beyond a single pod, we do some form of model parallelism or FSDP within the ICI domain, and then pure data parallelism across multiple pods. Let *N* = be the number of TPUs we want to scale to and *M* = be the number of TPUs per ICI-connected slice. To do an AllReduce over DCN, we can do a ring-reduction over the set of pods, giving us (in the backward pass): **✦ Adaptation:** The printed derivation assumes full, equal-size slices (*N* is a multiple of *M*). The live model balances the chips across slices, so a partial final slice cannot silently receive a full slice's aggregate NIC bandwidth. The comms bandwidth scales with *M*, since unlike ICI the total bandwidth grows as we grow our ICI domain and acquire more NICs. Simplifying, we find that Tmath > Tcomms when For TPU v5p, the *C*/*W dcn* is about / = . This tells us that to efficiently scale over DCN, there is a minimum batch size per ICI domain needed to egress each node.adaptation This is the **How much of a problem is this?** To take a specific example, say we want to train LLaMA-3 70B on TPU v5p with a BS of tokens. LLaMA-3 70B has *F* ≈ . From the above sections, we know the following: The TLDR is that we have a nice recipe for training with BS=1M, using roughly *DP* (FSDP) = 1024 and *TP* (TP) = 8, but with BS=2M we need to use DCN. As noted above, we have a DCN arithmetic intensity of , so we just need to make sure our batch size per ICI domain is greater than this. This is trivial for us, since with 2 pods we'd have a per-pod BS of , and a per TPU batch size of , which is great (maybe cutting it a bit close, but theoretically sound).adaptation The chapter's printed values (per-pod BS of 1M, per-TPU batch of 111) appear when you load the two-pod preset below; everything is recomputed from the live state, so try the one-pod recipe first and watch both numbers move. ✦ This section is an addition of this adaptation, drawing its text from [Chapter 12 (GPUs)](https://jax-ml.github.io/scaling-book/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](#data-parallelism), the TP bound into [Tensor Parallelism](#tensor-parallelism), expert parallelism into [its own section](#expert-parallelism), and the pipelining reasons into [Pipelining](#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](#scaling). As we did there, the goal here is to look at the total Tmath and Tcomms for different parallelism strategies and understand at what point Tcomms > Tmath. 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: | Node Type | GPUs per node | GPU egress bandwidth | Node egress bandwidth | |---|---|---|---| | H100 | 8 | 450e9 | 400e9 | | B200 | 8 | 900e9 | 400e9 | | GB200 NVL72 | 72 | 900e9 | 3600e9 | | GB300 NVL72adaptation This row is the adaptation's, not Chapter 12's — from NVIDIA's published GB300 NVL72 specs (dense BF16 = 180 PFLOPS/rack ÷ 72 = 2.5 PFLOP/s per GPU; ConnectX-8 at 800 Gb/s per GPU doubles the scale-out egress to 7200e9 per domain). | 72 | 900e9 | 7200e9 | 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 For data parallelism and ZeRO sharding, the compute-bound rule derived in [Data Parallelism](#data-parallelism) — *B*/*DP* > *C*/*W collective* — is reused unchanged, where 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](#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](#tensor-parallelism) — *TP* < *F* · *W collective* / 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. **What does DeepSeek do?** For reference, [DeepSeek V3](https://arxiv.org/abs/2412.19437) 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. 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: 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. **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. | Strategy | Description | |---|---| Data Parallelism | Activations are batch sharded, everything else is fully-replicated, we all-reduce gradients during the backward pass. | FSDP | Activations, weights, and optimizer are batch sharded, weights are gathered just before use, gradients are reduce-scattered. | Tensor Parallelism (aka Megatron, Model) | Activations are sharded along dmodel, weights are sharded along dff, activations are gathered before Win, the result reduce-scattered after Wout. | Mixed FSDP + Tensor Parallelism | Both of the above, where FSDP gathers the model sharded weights. | And here are the "formulas" for each method: | Strategy | Formula | |---|---| | DP | In[BDP, D] ·D Win[D, F] ·F Wout[F, D] → Out[BDP, D] | | FSDP | In[BDP, D] ·D Win[DDP, F] ·F Wout[F, DDP] → Out[BDP, D] | | TP | In[B, DTP] ·D Win[D, FTP] ·F Wout[FTP, D] → Out[B, DTP] | | TP + FSDP | In[BDP, DTP] ·D Win[DDP, FTP] ·F Wout[FTP, DDP] → Out[BDP, DTP] | | Strategy | Compute per layer (ignoring gating einsum) | Comms per layer (bytes, forward + backward pass) | |---|---|---| | DP | 4BDF/DP + 8BDF/DP | 0 + 8DF | | FSDP | 4BDF/DP + 8BDF/DP | 4DF + 8DF | | TP | 4BDF/TP + 8BDF/TP | 4BD + 4BD | | FSDP + TP | 4BDF/(DPTP) + 8BDF/(DPTP) | (4BD/DP + 4DF/TP) + (8BD/DP + 8DF/TP) | **✦ Adaptation:** The source table below this anchor is the dense TPU summary, so it is hidden for the current state rather than allowed to display false MoE or GPU formulas. For GPU guidance, use the [GPU TLDR and practical recipe](#gpus); for MoE routing, use [Expert Parallelism](#expert-parallelism). The scheme-level meters remain live for the selected model and hardware. 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. | hyperparam | value | |---|---| L | | D | | F | | | N | | | K | | | H | | | V | 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).* The total memory used for the parameters (bf16) and the two optimizer states (fp32, the first and second moment accumulators) is (2 + 4 + 4) · ≈ . The activations after the first two matmuls are shaped BF and after the last one BD (per the Transformer diagram above), so the total memory for bf16 is 2 · *L* · (*B**D* + 2 · *B**F*) = 2*L**B* · (*D* + 2*F*) or 2 · · · · (1 + 2 · ) ≈ = , since B=. All other activations are more or less negligible.adaptation Try: drag the batch and watch: the parameter + optimizer term () never moves, while the activation term scales linearly with it. That memory monster is what FSDP-style activation sharding exists to slay. **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. First, let's write down some numbers. With 32k sequence length and a 3M batch size, we have a sequence batch size of .adaptation The chapter says 96, which is 3·220/32,768; the live math here uses a literal 3e6, which gives . Either way: small! Long contexts eat a token budget fast. On a TPU v5p 16x16x16 slice, we have of HBM. We can't use pure data parallelism, because it replicates the parameters and optimizer states on each chip, which are already around (from Q2) which is more HBM than we have per-chip (). Let's start by looking purely at memory. Replacing BS= with 3M in Q2, we get ~ total checkpoint activations, and with the optimizer state this brings us to almost exactly = . The TPUv5p slice has of HBM in total, so we are safely under the HBM limit. Next let's look at whether we'll be comms or compute-bound. With 4096 chips and 3 axes of parallelism, we can do a minimum batch size of · 4096 = tokens. That's slightly above our 3M batch size. So we're actually comms-bound, which is sad. So the general answer is **no, we cannot do FSDP alone**. Now we know our primary concern is being comms-bound, so let's plug in some numbers. First of all, we know from above that our per-chip batch size with mixed FSDP + tensor parallelism needs to be above ² / 2*F* = here. That means we can in theory do this! Let's figure out how much of each. We have the rule so here we have sqrt(3e6 · 2 · 4096 / ) = , meaning we'll do roughly way DP and way TP. Per TPU memory will be as in (2), and step time will just be Above, we simplified the Transformer layer forward pass as Out[B, D] = In[B, D] ·D Win[D, F] ·F Wout[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] · Win[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, Wout, Win), 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](https://jax-ml.github.io/scaling-book/transformers/), continue to [Part 6: Applied Training](https://jax-ml.github.io/scaling-book/applied-training/), which works this content through real LLaMA models, or revisit the [original chapter](https://jax-ml.github.io/scaling-book/training/) this page adapts.