Google DeepMind、Transformer 学習の並列化手法を解説

Model

D=8,192

F=28.7k

L=80

E=1·k=1

Hardware

C=459 TFLOP/s

Wici=180 GB/s

Batch

B=2M tokens

An *explorable* adaptation of Part 5 of “How to Scale Your Model” by Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya, Charlie Chen, Sharad Vikram, Federico Lebron, Peter Choy, Vinay Ramasesh, Albert Webson & Reiner Pope (Google DeepMind).

✦ 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

2M and watch the per-chip batch

3,906 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); 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; the expert parallelism and GPU network 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.

What Do We Mean By Scaling?

The goal of “model scaling” is to be able to increase the number of

chips used for training or inference while achieving a proportional, linear

increase in throughput (we call this *strong scaling*). While performance

on a single chip depends on the trade-off between memory bandwidth and FLOPs,

performance at the cluster level depends on hiding inter-chip communication by

overlapping it with useful FLOPs. This is non-trivial, because increasing the

number of chips increases the communication load while reducing the amount of

per-device computation we can use to hide it. As we saw in

Section 3, sharded

matrix multiplications often require expensive

AllGathers or

ReduceScatters that can block the TPUs from doing

useful work. The goal of this section is to find out when these become

*too expensive.*

In this section, we'll discuss five common parallelism schemes: (pure)

data parallelism, fully-sharded data parallelism (FSDP / ZeRO

sharding), tensor parallelism (also known as model parallelism),

expert parallelism (for Mixture-of-Experts models),

and (briefly) pipeline parallelism. For each, we'll show what

communication cost we incur and at what point that cost starts to bottleneck our

compute cost.We'll

focus on communication bounds — since while memory capacity constraints are

important, they typically do not bound us when using rematerialization

(activation checkpointing) and a very large number of chips during pre-training.

(Ed: This edition is expanded to discuss

expert parallelism, unlike the

original.) For this section, you can focus solely on inter-chip

communication costs, since as long as we have a large enough single-chip batch

size, the transfer of data from HBM to MXU is already overlapped with

computation.

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

section.

live values shown for:

(mirrors the top bar)

✦ adaptation — this notation, worn by the chapter's dense models and today's frontier open-source ones (supported rows are clickable)

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.

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.

✦ adaptation — the hardware, with receipts (click a row to load it)

Every hardware number this page computes with, spec and sustained, with its

source.Full

citations live in 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 *sparsity* FLOP/s, halved here to dense; "bidirectional"

bandwidths are halved to per-direction; per-GPU scale-out is the node's NIC total

divided by its GPUs. ≈ marks factors that are *estimates* (stated basis, no

direct public measurement — e.g. Blackwell collectives inherit H100's measured

NCCL ratio until independent nccl-tests exist) rather than measured.

Hover any cell for its citation. The *sustained* and *achieved*

columns are the measured fractions of spec: flip the top bar's

spec / measured control and every equation on the page derates

by them (compute × sustained, bandwidth × achieved — the loaded hardware's row

shows them as live green scrubs). Wall-clock estimates that already assume an

MFU keep using spec peak, so nothing double-counts. Notice the punchline the

citations force: TPUs sustain far closer to their paper numbers than the

power-throttled NVIDIA parts.

For simplicity's sake, we'll approx