{"slug": "how-to-parallelize-a-transformer-for-training", "title": "How to Parallelize a Transformer for Training", "summary": "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.", "body_md": "✦ 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.)\n\n**This page is a working model, not a description of one.**\nEvery 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\nand watch the per-chip batch\nfollow (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 —\nrestores 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.\n\n**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:\nwherever the chapter printed a fixed number, this page computes it live (these in-place swaps aren't individually marked);\nthe interactive figures and their captions replace the original static figures;\n✦ 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\n[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;\nthe 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);\nunder 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);\ncontent 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);\nwhere a sentence had to be altered to host a live element, a Δ margin note quotes the original and states the change;\nand where this edition's additions make a chapter statement inaccurate as written, an italic (Ed: …) interjection corrects it in place.\n\nThe goal of “model scaling” is to be able to increase the number of\nchips used for training or inference while achieving a proportional, linear\nincrease in throughput (we call this *strong scaling*). While performance\non a single chip depends on the trade-off between memory bandwidth and FLOPs,\nperformance at the cluster level depends on hiding inter-chip communication by\noverlapping it with useful FLOPs. This is non-trivial, because increasing the\nnumber of chips increases the communication load while reducing the amount of\nper-device computation we can use to hide it. As we saw in\n[Section 3](https://jax-ml.github.io/scaling-book/sharding/), sharded\nmatrix multiplications often require expensive\nAllGathers or\nReduceScatters that can block the TPUs from doing\nuseful work. The goal of this section is to find out when these become\n*too expensive.*\n\nIn this section, we'll discuss five common parallelism schemes: (pure)\n**data parallelism, fully-sharded data parallelism** (FSDP / ZeRO\nsharding), **tensor parallelism** (also known as model parallelism),\n**expert parallelism** (for Mixture-of-Experts models),\nand (briefly) **pipeline parallelism**. For each, we'll show what\ncommunication cost we incur and at what point that cost starts to bottleneck our\ncompute cost.We'll\nfocus on communication bounds — since while memory capacity constraints are\nimportant, they typically do not bound us when using rematerialization\n(activation checkpointing) and a very large number of chips during pre-training.\n(Ed: This edition is expanded to discuss\n[expert parallelism](#expert-parallelism), unlike the\noriginal.) For this section, you can focus solely on inter-chip\ncommunication costs, since as long as we have a large enough single-chip batch\nsize, the transfer of data from HBM to MXU is already overlapped with\ncomputation.\n\nWe'll use the following notation to simplify calculations throughout this section.\n\n| Notation | Meaning (model parameters) | Live value |\n|---|---|---|\nD | dmodel (the hidden dimension/residual stream dim) |\n|\nF | dff (the feed-forward dimension)adaptation\nF convention (everywhere): the width of one\nexpert (= dff when dense); math runs through\nk·F, weights hold E·F, and the\nchapter's equations are the E = k = 1 case\n(Chapter 12's resolution). One honest limitation: models that mix\ndense and MoE blocks have two genuinely different F's —\nDeepSeek-V3 runs its first three layers dense at a much wider width —\nand this page approximates such models as uniformly MoE. Hover any\nF for the live widths. |\n|\nB | Batch dimension (number of tokens in the batch; total, not per-device) | |\n| T | Sequence length | — |\nL | Number of layers in the model |\n\n| Notation | Meaning (hardware characteristic) | Live value |\n|---|---|---|\nC | FLOPS/s per chip | |\nW | Network bandwidth (bidirectional per TPU mesh axisone-way GPU or node egress, often subscripted as e.g. W or iciW)dcn |\nici · dcn |\nDP | Number of chips along the data-parallel mesh axis (the chapter's X) | |\nTP | Number of chips along an alternate, tensor-parallel mesh axis (the chapter's Y) | |\nZ | Number of chips along a third mesh axis, labeled Z | — |\nPP | Pipeline stages (the pipelining section's Z) | |\nEP | Expert-parallel degree (chapter 12's Z; see the expert-parallelism section) |\n\nThe chapter's examples are dense LLaMA-era models; the frontier has since gone\nMixture-of-Experts.Shapes\nfrom each model's published `config.json`\n\non Hugging Face; parameter\ntotals from its safetensors metadata. Retrieved August 2026.\nE and k count shared experts, so k·*F* is the activated\nwidth for the architectures represented by the live presets; column headers\nexplain each field. The dense models from the\ntop-bar dropdown lead the table for contrast, and whichever model is loaded\nshows its row in live green — scrub it right here.\n\n| Model | params | D | F | act. k·F | L | E | k |\n|---|---|---|---|---|---|---|---|\n| (chapter default) | 70.6B | 8,192 | 28,672 | 28,672 | 80 | 1 | 1 |\n| 13.0B | 5,120 | 13,824 | 13,824 | 40 | 1 | 1 | |\n| 8.54B | 3,072 | 24,576 | 24,576 | 28 | 1 | 1 | |\n| 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 |\n| 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 |\n| 753B | 6,144 | 2,048 | 18,432 | 78 | 257 | 8+1 | |\n| 1.60T | 7,168 | 3,072 | 21,504 | 61 | 385 | 6+1 | |\n| 2.45T | 8,192 | 2,048 | 22,528 | 92 | 513 | 10+1 | |\n| 952B | 6,144 | 3,072 | 24,576 | 66 | 258 | 6+2 | |\n| 427B | 6,144 | 3,072 | 15,360 | 60 | 129 | 4+1 |\n\nClick 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.\n\nEvery hardware number this page computes with, spec and sustained, with its\nsource.Full\ncitations live in [ SOURCES.md](SOURCES.md) alongside this\npage — every value traces to a vendor spec sheet, a published measurement, or the\nbook's own benchmarks; retrieved 2026-08-17; click any cell to pin its citation\nand follow the source link. Methodology for the synthesized numbers: NVIDIA\ndatasheets headline\n\n| Hardware | C (dense bf16) | × sust. | W link | × achv. | W scale-out | HBM |\n|---|---|---|---|---|---|---|\n| 459 TF | ≈0.72 | 180 GB/s | ≈0.95 | 6.25 GB/s | 96 GB | |\n| 197 TF | ≈0.67 | 90 GB/s | ≈0.95 | 3.13 GB/s | 16 GB | |\n| 989 TF | 0.73 | 450 GB/s | 0.82 | 50 GB/s | 80 GB | |\n| 2.25 PF | 0.69 | 900 GB/s | ≈0.82 | 50 GB/s | 180 GB | |\n| 2.5 PF | ≈0.70 | 900 GB/s | ≈0.82 | ≈50 GB/s | 186 GB | |\n| 2.5 PF | ≈0.70 | 900 GB/s | ≈0.82 | 100 GB/s | 288 GB | |\n| 989 TF | ≈0.73 | 200 GB/s | 0.80 | 50 GB/s | 80 GB |\n\nFor simplicity's sake, **we'll approximate a Transformer as a stack of\nMLP blocks** — attention is a comparatively small fraction of the FLOPs\nfor larger models as we saw in\n[Section 4](https://jax-ml.github.io/scaling-book/transformers/).\nWe will also ignore the gating matmul, leaving us with the following simple\nstructure for each layer:adaptation\nWith this simplification each layer holds\n2·*D*·E·*F* weights (E = 1 for a\ndense model, so simply 2·D·F), and the whole stack has\n2·*D*·E·*F*·*L* =\nparameters at the current\nsettings — the “P” in this page's communication arithmetic. Memory\nquestions are different: a real checkpoint holds the gated MLP's third matrix\nand the attention stack too, so the memory meters price weights at\nPw ≈ 3·D·E·F·L + 2.5·D²·L =\n, which tracks the\nmodel table's published totals to within a few percent (vocab embeddings and\nMHA-era attention excepted).\n\n**Forward pass:** need to compute Loss[B]\n\n**Backward pass:** need to compute dWout[F, D], dWin[D, F]\n\nWe provide this for comparison to the algorithms with communication added.\n\nHere are the 4 parallelism schemes we will discuss. Each scheme can be thought\nof as uniquely defined by a sharding for **In**,\n**W in, Wout, and Out** in the above\ndiagram.adaptation\nA quick reminder of the book's notation: a subscript on an array dimension names\nthe mesh axis it is split over — In[B\n\n**1. Data parallelism:** *activations sharded along batch,\nparameters and optimizer state are replicated on each device. Communication only\noccurs during the backwards pass.*\n\n**2. Fully-sharded data parallelism (FSDP or ZeRO-3):**\n*activations sharded along batch (like pure data parallelism), parameters\nsharded along same mesh axis and AllGathered\njust-in-time before use in forward pass. Optimizer state also sharded along\nbatch. Reduces duplicated memory.*\n\n**3. Tensor parallelism (also called Megatron sharding or model\nparallelism):** *activations sharded along D (d model),\nparameters sharded along F (dff).\nAllGather and\nReduceScatter activations before and after each\nblock. Compatible with FSDP.*\n\n**4. Pipeline parallelism:** *weights sharded along the layer\ndimension, activations microbatched and rolled along the layer dimension.\nCommunication between pipeline stages is minimal (just moving activations over a\nsingle hop). To abuse notation:*adaptation\nNotice what all four schemes have in common: every one runs the *same*\nmatmuls — the FLOPs never change, only where the arrays live and which\ncollectives must run between the multiplies. So for each scheme the question is\nalways whether those collectives can hide behind the matmuls. Before the chapter\ndissects the schemes one by one, this adaptation inserts a short primer —\n[First, Feel the Roofline](#roofline) — building the one picture that\nanswers that question for all four.\n\n✦ This entire section is an addition of this adaptation — the chapter's ideas, our framing. The original text resumes at [Data Parallelism](#data-parallelism).\n\nThe 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.\n\nWhen a chip works on one layer of our Transformer, two clocks run *at the same time*:\n\nThe **compute clock**: the MXU has to chew through this layer's share of FLOPs. With\nB = tokens\nsplit over DP = chips,\nthat's 4 · *B* · *D* · k·*F**DP* · *C* =\nper layer (k·F because a token only multiplies through its k activated experts).\n\nThe **network clock**: whatever bytes this scheme moves have to squeeze through the interconnect at\nWici = .\nCrucially, these two clocks *can* overlap when the implementation schedules\nthe collective successfully: the network carries bytes while the MXU\nmultiplies. Under that explicit assumption, a layer costs the **max**,\nnot the sum. Communication that fits under the compute clock is hidden;\ncommunication that pokes out past it leaves silicon idle.\n\nTry: 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.\n\nThat toggle captures the core DP/FSDP/TP distinction: *what travels?*\n\nNow 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:\n\nTry: 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.\n\n**Syntax:**\n\nWhen your model fits on a single chip with even a tiny batch size (>240\ntokens, so as to be compute-bound), **you should always use simple data\nparallelism.** Pure data parallelism splits our activations across any\nnumber of TPUs so long as the number of TPUs is smaller than our batch size.\nThe forward pass involves no communication, but at the end of every step,\n**each TPU performs an AllReduce on its\nlocal gradients to synchronize them before updating the\nparameters.**\n\n**Pure Data Parallelism Algorithm:**\n\n**Forward pass:** need to compute Loss[BDP]\n\n**Backward pass:** need to compute\ndWout[F, D],\ndWin[D, F]adaptation\nThe {UDP} annotation below marks a\nresult that is *unreduced* over the *DP* axis: each chip holds a\npartial sum from its own slice of the batch.\n\nWe ignore the details of the loss function and abbreviate\nTmp = Win · In. Note that, although our\nfinal loss is the average\nAllReduce(Loss[BDP]),\nwe only need to compute the AllReduce on the backward pass when averaging\nweight gradients.\n\nNote that the forward pass has no communication — **it's all in the\nbackward pass**! The backward pass also has the great property that the\nAllReduces aren't in the “critical path”, meaning that each\nAllReduce can be performed whenever it's convenient and doesn't block you from\nperforming subsequent operations. The overall communication cost *can still\nbottleneck us* if it exceeds our total compute cost, but it is much more\nforgiving from an implementation standpoint. We'll see that model/tensor\nparallelism doesn't have this\nproperty.adaptation\nIn [the primer](#roofline)'s terms: because the AllReduce can be\nlaunched whenever convenient, the only question left is whether the total\ncomms fits under the total compute — the roofline question, answered just\nbelow. Tensor parallelism's collectives, by contrast, block the very next\nmatmul.\n\n**Why do this?** Pure data parallelism reduces activation\nmemory pressure by splitting our activations over the batch dimension, allowing\nus to almost arbitrarily increase batch size as long as we have more chips to\nsplit the batch dimension over. Especially during training when our activations\noften dominate our memory usage, this is very helpful.\n\n**Why not do this?** Pure data parallelism does nothing to\nreduce memory pressure from model parameters or optimizer states, which means\npure data parallelism is rarely useful for interesting models at scale where\nour parameters + optimizer state don't fit in a single TPU. To give a sense of\nscale, if we train with parameters in bf16 and optimizer state in fp32 with\nAdamAdam\nstores parameters, first order and second order accumulators. Since the params\nare in bfloat16 and optimizer state is in float32, this gives us\n`2 + 8 = 10`\n\nbytes per parameters., the largest model\nwe can fit has TPU memory / 10 parameters, so e.g. on\na TPUv5p chip with\nof HBM and pure data parallelism this is about\nparameters.\n\n*To make this useful for real models during training, we'll need to at\nleast partly shard the model parameters or optimizer.*\n\n**When do we become bottlenecked by communication?** As we can\nsee above, we have two AllReduces per layer, each of size\n2*D**F* (for\nbf16 weights). When does data parallelism make us communication\nbound?adaptation\nThe network here carries weight *gradients* —\n2 · *D* · E · *F* =\nper matrix (all E experts' gradients, not just the k a token used) —\nwhose size doesn't care about the batch. This is the weight-moving roofline\nfrom [the primer](#roofline): a batch-blind comms cost that a big\nenough per-chip batch can always hide.\n\nAs in the table above, let *C* = per-chip FLOPs,\n*W ici* =\n\n*Communication time:* From a previous section we know that the time\nrequired to perform an AllReduce in a 1D mesh depends only on the total bytes\nof the array being AllReduced and the ICI bandwidth\n*W ici*; specifically the AllReduce time is\n2 · total bytes / W\n\n**✦ Adaptation:** This DP collective spans\nmore than one NVLink domain, so the live clock uses scale-out-limited bandwidth\n, not the faster\nlocal NVLink egress.\n\n*Matmul time:* Each layer comprises two matmuls in the forward pass,\nor four matmuls in the backwards pass, each of which requires\n2(*B*/*DP*)*D**F*\nFLOPs. Thus, for a single layer in the backward pass, we have\n\nSince we overlap, the total time per layer is the max of these two quantities:\n\nWe become compute-bound when\nTmath/Tcomms > 1, or\nwhen\n\nThe upshot is that, to remain compute-bound with data parallelism, we need\nthe per-device batch size *B*/*DP* to\nexceed the ICI operational intensity,\n*C*/*W ici*. This is ultimately\na consequence of the fact that the computation time scales with the per-device\nbatch size, while the communication time is independent of this quantity (since\nwe are transferring model weights). Note the resemblance of the\n\nFor a Mixture of Experts (MoE) model, where we have *E* experts\nand *k* experts per token, this increases to\n\nwhich inflates the per-GPU token batch size by a factor of\n*E*/*k*, i.e.\n\nFor example, the new OpenAI OSS model with k=4 and\nE=128, this increases to\n32 · 2475 = 79,200 across nodes, a kind of ridiculously\nhigh number.adaptation\nKept as the chapter's static example (its 2475 is the H100 cross-node ridge from\n[the GPU section](#gpus)). At your current hardware and model, the\nsame computation reads (E/k) · C/Wcollective =\ntokens per\nchip. Expert parallelism — sharding the experts themselves, so gradients stop\ncrossing the whole DP axis — is the standard escape; it gets\n[its own section](#expert-parallelism) below.\n\nLet's put in some real numbers to get a sense of scale. For TPUv5p,\n`C`\n\n=\nand `W`\n\n=\nfor 1D data parallelism over ICI, so **our batch size per chip must be at\nleast to avoid\nbeing\ncommunication-bound**.adaptation\nThe famous 2,550 — the same constant [the primer](#roofline)\npromised would keep reappearing. It's live here: change the hardware in the\nmachine bar and this floor moves with it. Since we can do data\nparallelism over multiple axes, if we dedicate all three axes of a TPUv5p pod\nto pure data parallelism, we 3x our bandwidth\n*W ici* and can scale down to only\nBS= per TPU or\ntokens per batch per pod (of\nchips)!\n\n**Syntax:**\n\nFully-sharded data parallelism (often called FSDP or\n[ZeRO-sharding](https://arxiv.org/abs/1910.02054)) splits the\nmodel optimizer states and weights across the data parallel shards and\nefficiently gathers and scatters them as needed. **Compared to pure\ndata parallelism, FSDP drastically reduces per-device memory usage and saves\non backward pass FLOPs, with very minimal overhead.**\n\nYou'll remember (from\n[Section 3](https://jax-ml.github.io/scaling-book/sharding/))\nthat an AllReduce can be decomposed into an\nAllGather and a\nReduceScatter. This means that, instead of\ndoing the full gradient AllReduce for standard data parallelism, we can\nshard the weights and optimizer states across chips,\nAllGather them at each layer during the\nforward pass and ReduceScatter across the\nweights during the backward pass at no extra cost.\n\n**Fully-Sharded Data Parallelism (FSDP):**\n\n**Forward pass:** need to compute Loss[BDP]\n\n**Backward pass:** need to compute dWout[F, DDP], dWin[DDP, F]\n\nThis 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.\n\n**Why would we do this?** Standard data parallelism involves\na lot of duplicated work. Each TPU AllReduces\nthe full gradient, then updates the full optimizer state (identical work on\nall TPUs), then updates the parameters (again, fully duplicated). For ZeRO\nsharding (sharding the gradients/optimizer state), instead of an AllReduce,\nyou can ReduceScatter the gradients, update\nonly your shard of the optimizer state, update a shard of the parameters,\nthen AllGather the parameters as needed for\nyour forward pass.adaptation Try:\nin the memory meter below, drag *DP* =\ndown toward 1 and watch the meter fill up and overflow — at\n*DP* = 1 you're just pure DP on one chip's memory budget.\nEvery doubling of *DP* halves the parameter and optimizer\nstripes.\n\n**When do we become bottlenecked by communication?** Our\nrelative FLOPs and comms costs are exactly the same as pure data\nparallelism, since each AllReduce in the\nbackward pass has become an AllGather +\nReduceScatter. Recall that an AllReduce is\nimplemented as an AllGather and a ReduceScatter, each with half the cost.\nHere we model the forward pass since it has the same FLOPs-to-comms ratio as\nthe backward pass:adaptation The\nchapter writes these equations for one mesh axis. The live line beneath them\n(and every meter on this page) spreads the collective over\n*M DP* =\nmesh axes — the M\n\n**✦ Adaptation:** This FSDP collective\ncrosses NVLink domains, so Wcollective is the scale-out-limited\n.\n\nTherefore, as with pure data-parallelism, we are compute bound when\n*B* / *DP* >\n*C* / *W collective*, i.e.\nwhen the per-device batch size\n\nFor example, borrowing only DeepSeek-V2's reported batch size as a\n*dense-model thought experiment* (this calculation does not model its\nexpert 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.\n**This would allow us to scale to roughly\nchips,\nor around\nTPUv5 pods, before we hit a bandwidth\nlimit.**adaptation Load\nthe DeepSeek scenario with the button below, then make the batch your own:\nthe mini-calculator that follows is an addition of this edition. With\n*B* =\ntokens, FSDP scales to *DP max* chips before hitting the bandwidth\nlimit.\n\nFor LLaMA-3 70B, which was trained for approximately\n(15e12 · 70e9 · 6) FLOPs, we could split a batch of\ntokens over roughly *B* / (α / 3) =\nchips\n(roughly\npods of chips), each with\nFLOPs running at\npeak FLOPs utilization (often called MFU), and **train it in\napproximately\n**.\nNot bad! But let's explore how we can do\nbetter.adaptation The\nchapter's numbers (16M tokens, 18,823 chips, 17 days) are one point of this\nlive sentence — the recipe button below restores them. Then drag\n*B* and watch chips and wall-clock trade off: a bigger\nbatch rides the same ridge on more chips and finishes sooner, which is\nexactly why the labs fight for every doubling of critical batch size. The\nequation below, also an addition, shows the wall-clock\narithmetic.\n\n**Syntax:**\n\n(we use *TP* to eventually combine with FSDP)\n\nIn a fully-sharded data-parallel AllReduce we\nmove the weights across chips. We can also shard the feedforward dimension of\nthe model and move the activations during the layer — this is called\n“1D model parallelism” or Megatron sharding\n([Shoeybi et al. 2019](https://arxiv.org/abs/1909.08053)). This can\nunlock a smaller efficient batch size per pod. The figure below shows an example of a single matrix\nsharded in this way:\n\nAs noted, **In[B, D TP] ·D\nWin[D, FTP] ·F\nWout[FTP, D] →\nOut[B, DTP] means we have to gather our\nactivations before the first matmul. This is cheaper than ZeRO sharding when\nthe activations are smaller than the weights.**adaptation\nCompare the two freights live, per layer in bf16: gathering activations moves\n2 ·\n\n**Tensor Parallelism:** adaptation\nWatch the phrase *on critical path*. With pure data parallelism the\nAllReduce happened after the loss was already\ncomputed, so the network could grind away while the chips moved on. Here\nthe matmuls cannot start until the gathers finish — these collectives sit\nsquarely in the layer’s serial path.\n\n**Forward pass:** need to compute Loss[B]\n\n**Backward pass:** need to compute\ndWout[FTP, D],\ndWin[D, FTP]\n\nOne nice thing about tensor parallelism is that it interacts nicely with\nthe two matrices in our Transformer forward pass. Naively, we would do an\nAllReduce after each of the two matrices. But\nhere we first do **In[B, D TP] ·\nWin[D, FTP] →\nTmp[B, FTP]** and then\n\n**How costly is this?** Let's only model the forward pass -\nthe backwards pass is just the transpose of each operation here. In 1D tensor\nparallelism we AllGather the activations before the first matmul, and\nReduceScatter them after the second, sending two bytes at a time (bf16). Let's\nfigure out when we're bottlenecked by communication.\n\nNoting that we want compute cost to be greater than comms cost, we\nget:adaptation\nNotice that *B* · *D* appears in\n*both* clocks, so the batch cancels out of the ratio. This is the flat\norange line from [the primer](#roofline): tensor parallelism’s\ncompute-to-comms ratio is pinned at\n*k* · *F* / (*TP* · αTP) =\nno matter\nthe batch — a weight-moving scheme can hide its comms behind more tokens per\nchip, but no batch size can raise this bar.\n\n**✦ Adaptation:** This TP collective\nspans NVLink domains. The live clock therefore uses\nscale-out-limited\nbandwidth rather than the local .\n\nThus for instance, for TPUv5p,\n*C*/*W ici* =\nin bf16, so we can\nonly do tensor parallelism up to\n\n**Note that this doesn't depend on the precision of the\ncomputation**, since e.g. for int8, on TPUv5p,\nCint8/*W ici* is\ninstead of\nbut the comms\nvolume is also halved, so the two factors of two cancel.\n\n**Let's think about some examples:** adaptation\nThe buttons below load each model's real shape into the page's state — every\nnumber, meter, and verdict recomputes when you click one. Try: with a model\nloaded, set *TP* =\nto 8, then 16, then 32, and watch the verdicts. Or hold the model fixed and\nscrub *C* =\n:\nfaster chips shrink *TP max* =\non the\nfabric carrying the current TP collective, which is\nwhy each hardware generation makes tensor parallelism a little harder to\nhide.\n\n**Syntax:**\n\nThe nice thing about FSDP and tensor parallelism is that they can be\ncombined. By sharding **W in** and\n\n**Forward pass:** need to compute\nLoss[B]adaptation\nCount what actually sits on the critical path: on the\n*TP* axis, one AllGather\nin (step 1) and one ReduceScatter out\n(step 6) of activation bytes; the weight gathers on the\n*DP* axis (steps 2 and 4) can be prefetched. Those two\n*TP*-axis activation hops and two *DP*-axis weight hops are exactly the\n2 · 2 factors in TTP comms and TFSDP comms\nbelow.\n\n**Backward pass:** need to compute\ndWout[FTP, DDP],\ndWin[DDP, FTP]\n\n**What's the right combination of FSDP and TP?** A simple but\nkey maxim is that FSDP moves weights and tensor parallelism moves\nactivations. That means as our batch size shrinks (especially as we do more\ndata parallelism), tensor parallelism becomes cheaper because our activations\nper-shard are smaller.adaptation\nThis maxim is [the primer](#roofline)'s weights-move vs\nactivations-move toggle made load-bearing: per layer in bf16, FSDP's freight\nis 2 · *D* · *E* · *F*/*TP* =\nof weights\nwhile TP's is 2 · *B* · *D*/*DP* =\nof\nactivations — each scheme shrinks the *other's* bill.\n\nThus 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:\n\n**TPU closed form.** Let *DP* be the number of chips dedicated to FSDP and\n*TP* be the number of chips dedicated to tensor\nparallelism. Let *N* be the total number of chips in our\nslice with *N* = *DP**TP*.\nLet *M DP* and\n\nAnd likewise our total FLOPs time is\n\n**✦ Adaptation — GPU topology correction:**\nthe TPU equations immediately above are not valid GPU substitutions. An outer\nFSDP reduction does not become TP times faster while TP remains inside one\nNVLink domain; the scale-out link still carries the reduction. The live clocks,\nmeter, and explorers use Chapter 12's\nmax(Tdomain, Tscale-out) model.\nIn the equations below, bytes = 4·D·E·F and G is the selected NVLink-domain size.\nThe closed-form optimum below is therefore shown only on TPU; on GPU the\nexplorer finds the topology-aware minimum directly.\n\nTo simplify the analysis, we make two assumptions: first, we allow\n*DP* and *TP* to take on non-integer\nvalues (as long as they are positive and satisfy\n*DP**TP* = *N*);\nsecond, we assume that we can fully overlap comms on the\n*DP* and *TP* axis with each other. Under\nthe second assumption, the total comms time is\n\nBefore we ask under what conditions we'll be compute-bound, let's find the\noptimal values for *DP* and *TP* to\nminimize our total communication. Since our FLOPs is independent of\n*DP* and *TP*, the optimal settings are\nthose that simply minimize comms. To do this, let's write\nTcomms above in terms of *DP* and\n*N* (which is held fixed, as it's the number of chips in\nour system) rather than *DP* and *TP*:\n\nBecause TFSDP comms is monotonically increasing in\n*DP*, and TTP comms is monotonically decreasing\nin *DP*, the maximum must be minimized when\nTFSDP comms = TTP comms,\nwhich occurs when\n\nThis is super useful! This tells us, for a given *B*,\n*F*, and *N*, what amount of FSDP is\noptimal. Let's get a sense of scale. Plugging in realistic values, namely\n*N* = 64 (corresponding to a 4x4x4 array of chips),\n*B* = 48,000, *F* = 32768, gives\nroughly *DP* ≈ .\nSo we would choose *DP* to be 16 and\n*TP* to be 4, close to our calculated\noptimum.adaptation\nThe chapter rounds this to ≈13.9; the pinned live value here is\n√(48,000 · 2 · 64 / 32,768) exactly. And at whatever is loaded right now,\n*DP opt* =\n.\nPress the first button below to load the chapter's exact scenario into the\nwhole page.\n\nNow let's return to the question we've been asking of all our parallelism\nstrategies: **under what conditions will we be\ncompute-bound?** Since we can overlap FLOPs and comms, we are\ncompute-bound whenadaptation\nSame question as [the primer](#roofline)'s: does the slower of the\ntwo comms clocks fit under the compute clock?\n\nBy letting\nα ≡ *C* / *W ici*,\nthe ICI arithmetic intensity, we can simplify:\n\nSince we calculated *DP opt* to make the LHS maximum equal, we can\njust plug it into either side (noting that\n\nFurther simplifying, we find that\n\nwhere 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:\n\nTo ensure that this ratio is greater than one so we are compute bound, we require\n\nTo get approximate numbers, again plug in *F* = 32,768,\nα = 2550, and\n*M DP*\n\nBelow we plot the ratio of FLOPs to comms time for mixed FSDP + TP,\ncomparing it both to only tensor parallelism (TP) and only data parallelism\n(FSDP), on a representative 4x4x4 chip array. While pure FSDP parallelism\ndominates for very large batch sizes, in the regime where batch size over\nnumber of chips is between roughly 100 and 850, a mixed FSDP + TP strategy is\nrequired in order to be\ncompute-bound.adaptation\nThe live chart below plays this figure's role: flip its view toggle to\n*ratio* to see Tmath/Tcomms for all three\nschemes, where any curve above 1 is compute-bound. It is drawn at the page's\ncurrent *N* =\nchips — press the\n4×4×4 chapter-example preset above to reproduce the chapter's exact\nframe.\n\nHere's another example of TPU v5p 16x16x16 showing the FLOPs and comms\ntime as a function of batch size for different sharding\nschemes.adaptation\nThat second figure is the same chart in *absolute-times* view. A\n16x16x16 slice is 4096 chips — exactly the page's default\n*DP*·*TP* = 512 · 8, so the\n“back to page defaults” preset above reproduces it.\n\nThe black curve is the amount of time spent on model FLOPs, meaning any\nbatch size where this is lower than all comms costs is strictly comms bound.\nYou'll notice the black curve intersects the aqua curve at about\n4e5, as\npredicted.adaptation\nOn the live chart that crossing sits at\n*B* = *N*·α²·*E*/(*M DP*·\n\nHere's an interactive animation to play with this, showing the total compute time and communication time for different batch sizes:\n\nYou'll notice this generally agrees with the above (minimum around\nFSDP=256, TP=16), plus or minus some wiggle factor for some slight\ndifferences in the number of axes for\neach.adaptation\nThe chapter's animation swept the FSDP/TP split itself; on this page that\nsweep is the earlier *DP*-axis explorer, whose optimum at the current state is\n*DP opt* =\n→ nearest power of two\n-way\nFSDP. The ±wiggle from mesh-axis bookkeeping is exactly the\nM\n\n✦ This section is\ndrawn from [Chapter 12\n(GPUs)](https://jax-ml.github.io/scaling-book/gpus/) of the same book and merged into this chapter's flow by this\nadaptation; condensed passages are marked. Its cost model is Chapter 12's\nswitched GPU fabric (NVLink node + InfiniBand scale-out — see\n[the GPU network model](#gpus) below). One naming change\nthroughout: Chapter 12 calls the expert-parallel axis *Z*; this\nedition names every parallelism degree after its scheme, so that axis is\nrendered *EP* here. The routed/shared-expert split and the\nhardware-domain generalization of the H100-specific formula are AI-written\nadaptation material, labeled again at the live estimate.\n\nAs we've already noted above, Mixture of Expert (MoE) models come with\n*E* times more model weights with only *k* times\nmore FLOPs, making data parallelism significantly\nharder.adaptation In\nChapter 12 \"noted above\" pointed at its Data Parallelism section; on this page that\npassage lives at the end of [Data Parallelism](#data-parallelism).\nThis page's *E* and *k* count all experts, including\nalways-on shared experts. Expert routing instead uses\nEr = E − s routed experts and\nkr = k − s routed selections, where\n*s* =\nshared experts. Right now, *E*r =\nand\n*k*r = .\nWe can mitigate the routed weight cost by sharding along the expert dimension, i.e.\nWin[EEP, D, F]. To do\nthe MLP block, we need to introduce 2x\nAllToAll to send our activations to the\ncorresponding experts.\n\n**What does an AllToAll cost here?** GPUs within a node have\nall-to-all connectivity, which makes AllToAlls, well, quite easy: each GPU just\nsends directly to the destination. For Mixture of Expert (MoE) models, we\nfrequently want to do a *sparse or ragged AllToAll*, where we guarantee at\nmost *k*r of *N* shards on the output dimension\nare non-zero; the cost is reduced by\nkr/N.adaptation Condensed\nfrom Chapter 12's intra-node collectives discussion (two paragraphs on dense and\nragged AllToAlls, with the exact expected-occupancy footnote) — see Chapter 12 for\nthe full derivation. The takeaway below is carried verbatim.\n\nFor the eight-GPU H100 node used in Chapter 12, the cost of this\nAllToAllEP→k r([B, D, k]) if it\nspans multiple nodes is roughly\nT\n\n**✦ Adaptation:** The live estimate is the substantive mash-up here. Within one\nNVLink domain it uses the finite ragged AllToAll cost from Chapter 12 rather\nthan calling that transfer free. Beyond the domain it takes the slower of the\nlocal switched-fabric component and the chapter's scale-out component, replacing\nthe H100-specific 8 with the selected hardware's domain size. Shared experts\nremain in the total *k*-wide compute, but never become routed AllToAll\ndestinations.\n\nFor that H100 case, Chapter 12 concludes that we either need\nkr > EP/8 with\n*F* > α · (EP − 8)/kr\nor EP ≫ kr and\n*F* > 8 · α, where\nα = *C*/*W*. This\ngives you two domains in which expert parallelism is possible, one with a small\namount of expert parallelism (roughly 2-node) and small *F*,\nor one with large *F* and EP arbitrarily large (up\nto *E*r-way expert parallelism).\n\nYou'll see both cases in practice, either a small amount of expert-parallelism\n(like DeepSeek v3 which has very small *F* and relatively\nsmall, restricted cross-node expert parallelism), or models with large\n*F*, in which case we can do significant cross-node EP\nalongside TP.\n\nYou'll probably notice we've avoided talking about pipelining at all in the\nprevious sections. Pipelining is a dominant strategy for GPU parallelism that is\nsomewhat less essential on TPUs. Briefly, pipelined training involves splitting the\nlayers of a model across multiple devices and passing the activations between\npipeline stages during the forward and backward pass.adaptation On\nthis page the split is live: with *L* =\nlayers over\n*PP* =\npipeline stages (scrubbable below), each device owns about\nconsecutive\nlayers. The algorithm is something like:\n\nThis 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.\n\n```\nbatch_size = 32\nd_model = 128\nd_ff = 4 * d_model\n\nnum_layers = len(jax.devices())\n\nkey = jax.random.PRNGKey(0)\n\n# Pretend each layer is just a single matmul.\nx = jax.random.normal(key, (batch_size, d_model))\nweights = jax.random.normal(key, (num_layers, d_model, d_model))\n\ndef layer_fn(x, weight):\n  return x @ weight\n\n# Assume we have num_layers == num_pipeline_stages\nintermediates = [x]\nfor i in range(num_layers):\n  x = layer_fn(x, weights[i])\n  intermediates.append(x)\n\n  if i != num_layers - 1:\n    x = jax.device_put(x, jax.devices()[i+1])\n\ndef loss_fn(batch):\n  return jnp.mean(batch ** 2)  # make up some fake loss function\n\nloss, dx = jax.value_and_grad(loss_fn)(x)\n\nfor i in range(num_layers - 1, -1, -1):\n  _, f_vjp = jax.vjp(layer_fn, intermediates[i], weights[i])\n  dx, dw = f_vjp(dx)  # compute the jvp dx @ J(L)(x[i], W[i])\n  weights[i] = weights[i] - 0.01 * dw  # update our weights\n\n  if i != 0:\n    dx = jax.device_put(dx, jax.devices()[i-1])\n```\n\n**Why is this a good idea?** Pipelining is great for many reasons:\nit has a low communication cost between pipeline stages, meaning you can train very\nlarge models even with low bandwidth interconnects. This is often very useful on\nGPUs since they are not densely connected by ICI in the way TPUs\nare.adaptation The\nchapter doesn't quantify \"low communication cost,\" so the check below is ours. A\nstage-boundary hop is a single point-to-point copy of one activation block —\n2*D* =\nper token in bf16 — and it's the [same roofline question](#roofline) as\never: does the hop fit under one stage's compute clock? The line below runs the\nnumbers for one microbatch.\n\n**Why is this difficult/annoying?** You might have noticed in the\npseudocode above that TPU 0 is almost always idle! It's only doing work on the very\nfirst and last step of the pipeline. The period of idleness is called a pipeline\nbubble and is very annoying to deal with. Typically we try to mitigate this first\nwith microbatching, which sends\n*M micro* =\nsmall batches through the\n\nThe overall communication cost of pipelining is tiny: with\n*N MB* microbatches and\n\nSince we are dividing by *N layers*, this is vastly\nsmaller than any of the other costs. In other words, from a communication\nstandpoint, pipelining is basically free. So why don't we just do pipelining?\nThere are a few reasons:\n\n(1) **Code complexity:** pipelining fits poorly into automatic\nparallelism frameworks (like XLA's GSPMD), because microbatching and custom\nzero-bubble schedules change the structure of the\nprogram.adaptation Condensed\nto one sentence — see Chapter 12 for the full paragraph.\n\n(2) **Pipelining makes data parallelism and FSDP hard:** probably\nthe biggest reason not to do pipelining is that it plays badly with FSDP and data\nparallelism. ZeRO-3 sharding in particular works badly, since it requires us to\nAllGather the weights on every microbatch which\ndoesn't work when we have only\n*B* / Nmicrobatches tokens\nto amortize the AllGather cost. Furthermore, during the backward pass, *we\ncan't AllReduce or ReduceScatter the gradients until the last microbatch has\npassed a given stage, which means we have significant non-overlapped\ncommunication time.*\n\n(3) **Pipeline bubbles and step imbalance:** naive pipeline\nschedules leave stages idle in bubbles, and passing activations from stage to\nstage on the critical path shifts stages relative to each other and adds\noverhead.adaptation Condensed\nto one sentence — see Chapter 12, and the live bubble math just above.\n\nThere 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.\n\nA second approach is to carefully overlap the forward matmul\nWi @ xi, the backward\ndx matmul\nWi @ ∂L/∂xi+1, and the\ndW matmul\n∂L/∂xi+1 @ xi. Since each of these\nrequires some FLOPs, we can overlap them to fully hide the bubble. Here's our live\nstand-in for the plot from the recent\n[DeepSeek v3 paper](https://arxiv.org/abs/2412.19437) showing their\n\"bubble-free\" pipeline\nschedule:adaptation Toggle\nthe widget's mode: *naive* is GPipe, *1F1B* interleaves one forward\nwith one backward (same bubble, far less activation memory held live), and\n*overlap-dW* is the DeepSeek-v3-style schedule — rush every ∂x result down\nthe pipeline to unblock neighbors, and drop the deferred ∂W matmuls into slots that\nwould otherwise sit idle.\n\nBecause 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.\n\nThe largest possible TPU slice is a TPU v5p SuperPod with 8960 chips (and 2240\nhosts). When we want to scale beyond this size, we need to cross the Data-Center\nNetworking (DCN) boundary. Each TPU host comes equipped with one or several NICs\n(Network Interface Cards) that connect the host to other TPU v5p pods over Ethernet.\nAs noted in the [TPU\nSection](https://jax-ml.github.io/scaling-book/tpus/), each host has about 200Gbps (25GB/s) of full-duplex DCN bandwidth,\nwhich is about\nfull-duplex (egress) bandwidth per\nTPU.adaptation Per\nthe TPU chapter of the original book: each v5p host serves 4 chips, so 25 GB/s per\nhost ÷ 4 ≈ 6.25 GB/s of egress per\nchip.adaptation The\nchapter printed 6.25GB/s; here *W dcn* is scrubbable —\ndrag it and this whole section (ridge included) recomputes. The hardware presets in\nthe top bar set it per machine.\n\nTypically, when scaling beyond a single pod, we do some form of model parallelism\nor FSDP within the ICI domain, and then pure data parallelism across multiple pods.\nLet *N* =\nbe the number of TPUs we want to scale to and\n*M* =\nbe the number of TPUs per ICI-connected slice. To do an\nAllReduce over DCN, we\ncan do a ring-reduction over the set of pods, giving us (in the backward pass):\n\n**✦ Adaptation:** The printed derivation assumes full, equal-size slices\n(*N* is a multiple of *M*). The live model\nbalances the chips across\nslices, so a partial final slice cannot silently receive a full slice's aggregate\nNIC bandwidth.\n\nThe comms bandwidth scales with *M*, since unlike ICI the total\nbandwidth grows as we grow our ICI domain and acquire more NICs. Simplifying, we\nfind that Tmath > Tcomms when\n\nFor TPU v5p, the\n*C*/*W dcn*\nis about\n/\n=\n. This tells us\nthat to efficiently scale over DCN, there is a minimum batch size per ICI domain\nneeded to egress each\nnode.adaptation This\nis the\n\n**How much of a problem is this?** To take a specific example, say we\nwant to train LLaMA-3 70B on TPU v5p with a BS of\ntokens. LLaMA-3 70B has *F* ≈\n.\nFrom the above sections, we know the following:\n\nThe TLDR is that we have a nice recipe for training with BS=1M, using roughly\n*DP* (FSDP) = 1024 and *TP* (TP) = 8, but with\nBS=2M we need to use DCN. As noted above, we have a DCN arithmetic intensity of\n, so we just need to\nmake sure our batch size per ICI domain is greater than this. This is trivial for\nus, since with 2 pods we'd have a per-pod BS of\n, and a per TPU\nbatch size of , which is\ngreat (maybe cutting it a bit close, but theoretically\nsound).adaptation The\nchapter's printed values (per-pod BS of 1M, per-TPU batch of 111) appear when you\nload the two-pod preset below; everything is recomputed from the live state, so try\nthe one-pod recipe first and watch both numbers move.\n\n✦ This section is an\naddition of this adaptation, drawing its text from\n[Chapter 12 (GPUs)](https://jax-ml.github.io/scaling-book/gpus/) of the\nsame book; condensed passages are marked. Chapter 12's per-scheme roofline\nderivations are not repeated here — they re-derive what this chapter already\nderived, so they are merged into the scheme sections above (the MoE penalty into\n[Data Parallelism](#data-parallelism), the TP bound into\n[Tensor Parallelism](#tensor-parallelism), expert parallelism into\n[its own section](#expert-parallelism), and the pipelining reasons into\n[Pipelining](#pipelining)). What remains here is the network model\nitself: the fabric, its bandwidths, what collectives cost on it, and the worked\nexamples.\n\nNow let's look at what this has all been building towards: understanding\nrooflines for LLM scaling on GPU. This is to complement the TPU training chapter\n[here](#scaling). As we did there, the goal here is to look at the total\nTmath and Tcomms for different parallelism strategies and\nunderstand at what point Tcomms > Tmath. As before, we\nconsider only the MLP block with operations\n\nwhere *B* is the global batch size **in tokens**\n(i.e. *B* = batch size · sequence length).\n\nHere we'll reproduce the table from Chapter 12 showing effective bandwidths at both the GPU and node level:\n\n| Node Type | GPUs per node | GPU egress bandwidth | Node egress bandwidth |\n|---|---|---|---|\n| H100 | 8 | 450e9 | 400e9 |\n| B200 | 8 | 900e9 | 400e9 |\n| GB200 NVL72 | 72 | 900e9 | 3600e9 |\n| 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 |\n\nLet's look at the compute communication rooflines as we did for TPUs for\n**data parallelism, tensor parallelism, pipeline parallelism, expert\nparallelism,** and combinations thereof. For the rest of this section we'll\nfocus on H100 rooflines for specific calculations. GB200-NVL72 has the same general\nrooflines but because we have a larger node egress bandwidth, we can sometimes be\nbottlenecked at the node level instead. The scheme derivations are merged into\ntheir corresponding sections above; below are the bounds they land on here.\n\nHere is the mapping used by the live GPU rooflines: read\n*W ici* as the per-GPU\n\nFor data parallelism and ZeRO sharding, the compute-bound rule derived in\n[Data Parallelism](#data-parallelism) —\n*B*/*DP* >\n*C*/*W collective* — is\nreused unchanged, where\n\nThis is quite a bit higher than on a TPU, where the number is 850 with all\nthree axes. On the H100 scale-out fabric the dense asymptotic floor is\n990e12/400e9 = 2,475 tokens per GPU, so 16,384 GPUs would require about\n40.6M tokens before the small-ring and model-parallel refinements; Llama 3.1\n405B used 16M. Chapter 12 quoted a 3,300-token H800 baseline from an unsupported\n300 GB/s figure. The reconciled H800 *local-link dense baseline* is 4,950\nin spec mode (990e12/200e9) and about 4,517 in this page's measured mode. Those\nare not a model of DeepSeek's full sparse run: its EP, PP, and 2-way DP alter the\nouter reduction. DeepSeek reports a pretraining batch schedule from 3,072 to\n15,360 sequences at a 4K maximum sequence length — about 12.6M to 62.9M tokens,\nwith 62.9M at steady state.edited\nThe source says H800 has 300 GB/s and “in practice, they used\n4M”. H800 is 200 GB/s per direction by the reconciled spec, DeepSeek\nreports 160 GB/s measured, and its report gives the sequence-batch schedule\nabove. See the [hardware table](#hardware-table).\n\n**Small-DP correction.** The asymptotic ridge above omits the\nring factor. With *X* scale-out domains, the exact dense condition is\n*B*/*N* >\n(*C*/*W collective*) · (X−1)/X\n(and ×E/k for the equal-width MoE model). For exactly two\ndomains the floor is halved, which is why 2-way data parallelism appears so\noften.\n\nFor tensor parallelism, the bound from\n[Tensor Parallelism](#tensor-parallelism) —\n*TP* < *F* · *W collective* /\n\nBeyond the node level: to a first approximation, because we have full bisection\nbandwidth, the cost of an AllGather or\nReduceScatter is roughly the buffer size in bytes\ndivided by the node egress bandwidth (400GB/s on H100) *regardless of any of\nthe details of the tree reduction.*\n\nWith 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.\n\n**What does DeepSeek do?** For reference,\n[DeepSeek V3](https://arxiv.org/abs/2412.19437) is trained with 2048\nH800 GPUs with:adaptation These\ntwo worked examples are Chapter 12's. The buttons load their cited hardware,\nmodel shape, batch, and listed parallelism degrees. Because the page's generic\nchip count is DP·TP, widgets that do not explicitly model EP or PP should be read\nas component-level views, not as a reconstruction of the full training run.\n\nThey 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.\n\n**What does Llama 3.1 405B do?** Llama 3.1 405B trains with a BS of\n16M tokens on 16,384 H100 GPUs, or about 977 tokens per GPU. They do:\n\nThe 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.\n\n**Practical recipe:** smaller dense models can use aggressive\nFSDP when batch permits; larger dense models commonly combine one- or two-domain\nTP with many-stage PP and DP; MoEs add EP, generally preferring it to TP while\nkeeping latency-sensitive collectives within as few domains as practical.\n\n| Strategy | Description |\n|---|---|\nData Parallelism |\nActivations are batch sharded, everything else is fully-replicated, we all-reduce gradients during the backward pass. |\nFSDP |\nActivations, weights, and optimizer are batch sharded, weights are gathered just before use, gradients are reduce-scattered. |\nTensor Parallelism (aka Megatron, Model) |\nActivations are sharded along dmodel,\nweights are sharded along dff,\nactivations are gathered before\nWin, the result reduce-scattered\nafter Wout. |\nMixed FSDP + Tensor Parallelism |\nBoth of the above, where FSDP gathers the model sharded weights. |\n\nAnd here are the \"formulas\" for each method:\n\n| Strategy | Formula |\n|---|---|\n| DP | In[BDP, D] ·D\nWin[D, F] ·F\nWout[F, D] →\nOut[BDP, D] |\n| FSDP | In[BDP, D] ·D\nWin[DDP, F] ·F\nWout[F, DDP] →\nOut[BDP, D] |\n| TP | In[B, DTP] ·D\nWin[D, FTP] ·F\nWout[FTP, D] →\nOut[B, DTP] |\n| TP + FSDP | In[BDP, DTP] ·D\nWin[DDP, FTP] ·F\nWout[FTP, DDP] →\nOut[BDP, DTP] |\n\n| Strategy | Compute per layer (ignoring gating einsum) |\nComms per layer (bytes, forward + backward pass) |\n|---|---|---|\n| DP | 4BDF/DP + 8BDF/DP |\n0 + 8DF |\n| FSDP | 4BDF/DP + 8BDF/DP |\n4DF + 8DF |\n| TP | 4BDF/TP + 8BDF/TP |\n4BD + 4BD |\n| FSDP + TP | 4BDF/(DPTP) + 8BDF/(DPTP) |\n(4BD/DP + 4DF/TP) + (8BD/DP + 8DF/TP) |\n\n**✦ Adaptation:** The source table below this\nanchor is the dense TPU summary, so it is hidden for the current state rather\nthan allowed to display false MoE or GPU formulas. For GPU guidance, use the\n[GPU TLDR and practical recipe](#gpus); for MoE routing, use\n[Expert Parallelism](#expert-parallelism). The scheme-level meters\nremain live for the selected model and hardware.\n\nLet'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.\n\n| hyperparam | value |\n|---|---|\nL |\n|\nD |\n|\nF |\n|\n| N | |\n| K | |\n| H | |\n| V |\n\nLLaMA-2 has separate embedding and output matrices and a gated MLP block.\n\n**Question 1:** How many parameters does LLaMA-2 13B have (I know\nthat's silly but do the math)? *Note that, as in\nTransformer Math,\nLLaMA-3 has 3 big FFW matrices, two up-projection and one down-projection. We\nignored the two \"gating\" einsum matrices in this section, but they behave the same\nas W in in this section.*\n\n**Question 2:** Let's assume we're training with BS=\ntokens and using Adam. Ignoring parallelism for a moment, how much total memory is\nused by the model's parameters, optimizer state, and activations? *Assume we\nstore the parameters in bf16 and the optimizer state in fp32 and checkpoint\nactivations three times per layer (after the three big matmuls).*\n\nThe total memory used for the parameters (bf16) and the two optimizer states\n(fp32, the first and second moment accumulators) is (2 + 4 + 4) ·\n≈\n.\nThe activations after the first two matmuls are shaped\nBF and after the last one BD\n(per the Transformer diagram above), so the total memory for bf16 is\n2 · *L* · (*B**D* + 2 · *B**F*) =\n2*L**B* · (*D* + 2*F*) or\n2 · ·\n·\n·\n(1 + 2 · ) ≈\n=\n,\nsince B=. All other\nactivations are more or less negligible.adaptation Try:\ndrag the batch\nand watch: the parameter + optimizer term\n()\nnever moves, while the activation term scales linearly with it. That memory\nmonster is what FSDP-style activation sharding exists to slay.\n\n**Question 3:** Assume we want to train with 32k sequence length\nand a total batch size of 3M tokens on a TPUv5p 16x16x16 slice. Assume we want to\nuse bfloat16 weights and a float32 optimizer, as above.\n\nFirst, let's write down some numbers. With 32k sequence length and a 3M batch\nsize, we have a sequence batch size of\n.adaptation The\nchapter says 96, which is 3·220/32,768; the live math here uses a\nliteral 3e6, which gives\n. Either way:\nsmall! Long contexts eat a token budget fast. On a TPU v5p\n16x16x16 slice, we have\nof HBM.\n\nWe 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 ().\n\nLet'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.\n\nNext let's look at whether we'll be comms or compute-bound. With 4096\nchips and 3 axes of parallelism, we can do a minimum batch size of\n· 4096 =\ntokens. That's slightly above our 3M batch size. So we're actually\ncomms-bound, which is sad. So the general answer is **no, we cannot do\nFSDP alone**.\n\nNow we know our primary concern is being comms-bound, so let's plug in\nsome numbers. First of all, we know from above that our per-chip batch size\nwith mixed FSDP + tensor parallelism needs to be above\n² / 2*F* =\nhere. That means we can in theory do this! Let's figure out how much of\neach.\n\nWe have the rule\n\nso 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\n\nAbove, we simplified the Transformer layer forward pass as\nOut[B, D] = In[B, D] ·D Win[D, F] ·F Wout[F, D].\nHow do we derive the comms necessary for the backwards pass?\n\nThis follows fairly naturally from the rule in the previous section for a\nsingle matmul Y = X · A:adaptation\nIn this appendix X and Y are the input and output *matrices* of a generic\nmatmul — the chapter's letters, kept as-is since this edition's mesh axes go by\n*DP* and *TP*, so nothing collides.\n\nUsing this, we get the following formulas (letting\nTmp[B, F] stand for\nIn[B, D] · Win[D, F]):\n\nNote that these formulas are mathematical statements, with no mention of\nsharding. The job of the backwards pass is to compute these four quantities. So\nto figure out the comms necessary, we just take the shardings of all the\nquantities which are to be matmulled in the four equations above (Tmp, dOut,\nWout, Win), which are specified by our parallelization\nscheme, and use the rules of sharded matmuls to figure out what comms we have to\ndo. Note that dOut is sharded in the same way as Out.\n\nLook back at\n[Part 4: Transformer Math](https://jax-ml.github.io/scaling-book/transformers/),\ncontinue to\n[Part 6: Applied Training](https://jax-ml.github.io/scaling-book/applied-training/),\nwhich works this content through real LLaMA models, or revisit the\n[original chapter](https://jax-ml.github.io/scaling-book/training/) this\npage adapts.", "url": "https://wpnews.pro/news/how-to-parallelize-a-transformer-for-training", "canonical_source": "https://ezyang.github.io/interactive-parallelize-transformer/", "published_at": "2026-08-18 19:55:18+00:00", "updated_at": "2026-08-18 20:11:05.183051+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-research", "ai-tools"], "entities": ["Google LLC", "Maruan Al-Shedivat", "Fable", "Anthropic", "OpenAI Codex", "JAX"], "alternates": {"html": "https://wpnews.pro/news/how-to-parallelize-a-transformer-for-training", "markdown": "https://wpnews.pro/news/how-to-parallelize-a-transformer-for-training.md", "text": "https://wpnews.pro/news/how-to-parallelize-a-transformer-for-training.txt", "jsonld": "https://wpnews.pro/news/how-to-parallelize-a-transformer-for-training.jsonld"}}