How to Parallelize a Transformer for Training

A new interactive adaptation of the JAX scaling book chapter on transformer training parallelism lets readers drag sliders to see how data, tensor, pipeline, and expert parallelism affect communication bottlenecks, with live computations replacing fixed numbers. The edition, built by Fable (Anthropic) and reviewed by OpenAI Codex, preserves original passages from Maruan Al-Shedivat and Google LLC under the MIT license while adding AI-written editorial notes and interactive figures.

✦ 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.