How to Parallelize a Transformer for Training — an explorable explanation

Model

LLaMA-3 70B<br>LLaMA-2 13B<br>Gemma 7B<br>DeepSeek-V3<br>GLM-5.2<br>DeepSeek-V4-Pro<br>Qwen3.8-2.4T-A95B<br>Inkling<br>MiniMax-M3<br>custom…

D=<br>F=<br>L=<br>E=·k=1">·shared=

Hardware

TPU v5p<br>TPU v5e<br>H100 (8-GPU node)<br>B200 (8-GPU node)<br>GB200 NVL72<br>GB300 NVL72<br>H800 (DeepSeek)<br>custom…

C=<br>Wici=

specmeasured

Batch<br>B= tokens

reset

What Do We Mean By Scaling?

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

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

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

live values shown for:

(mirrors the top bar)

NotationMeaning (model parameters)Live value

Dd model (the hidden dimension/residual stream dim)

Fd ff (the feed-forward dimension)✦adaptation<br>F convention (everywhere): the width of one<br>expert (= dff when dense); math runs through<br>k·F, weights hold E·F, and the<br>chapter's equations are the E = k = 1 case<br>(Chapter 12's resolution). One honest limitation: models that mix<br>dense and MoE blocks have two genuinely different F's —<br>DeepSeek-V3 runs its first three layers dense at a much wider width —<br>and this page approximates such models as uniformly MoE. Hover any<br>F for the live widths.

BBatch dimension (number of tokens in the batch; total, not per-device)

TSequence length—<br>LNumber of layers in the model

NotationMeaning (hardware characteristic)Live value

CFLOPS/s per chip

WNetwork bandwidth (bidirectional per TPU mesh axisone-way GPU or node egress, often subscripted as e.g. Wici or Wdcn)<br>ici · dcn<br>DPNumber of chips along the data-parallel mesh axis (the chapter's X)

TPNumber of chips along an alternate, tensor-parallel mesh axis (the chapter's Y)

ZNumber of chips along a third mesh axis, labeled Z<br>PPPipeline stages (the pipelining section's Z)

EPExpert-parallel degree (chapter 12's Z; see the expert-parallelism section)

✦ adaptation — this notation, worn by the chapter's dense models and today's frontier open-source ones (supported rows are clickable)<br>The chapter's examples are dense LLaMA-era models; the frontier has since gone<br>Mixture-of-Experts.✦Shapes<br>from each model's published config.json on Hugging Face; parameter<br>totals from its safetensors metadata. Retrieved August 2026.<br>E and k count shared experts, so k·F is the activated<br>width for the architectures represented by the live presets; column headers<br>explain each field. The dense models from the<br>top-bar dropdown lead the table for contrast, and whichever model is loaded<br>shows its row in live green — scrub it right here.

Model<br>params<br>act. k·F

LLaMA-3 70B (chapter default)<br>70.6B8,19228,67228,6728011<br>LLaMA-2 13B<br>13.0B5,12013,82413,8244011<br>Gemma 7B<br>8.54B3,07224,57624,5762811<br>DeepSeek-V3✦Counting<br>example: 256 routed + 1 shared expert → E 257; top-8 + shared →<br>k 9. Its first three layers are actually dense (see the F-convention<br>note above).<br>685B7,1682,04818,432612578+1<br>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...