the path · 0/15
start the path

the kernel path · ICI ⇄ chips · lesson 01 of 6

The TPU fabric

TPUs scale by wiring each chip to its neighbors and to nothing else. Everything about distributed kernels follows from that one wiring decision.

the goal Given a slice shape and a collective, count the hops and put a lower bound on its time using only link constants.

mastery work · this chapter0/3
  1. go →
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
§ 01

Neighbors, not switches

Ask how 4096 chips talk to each other and the obvious answer is a network: switches in the middle, every chip a cable away from every other. TPUs refuse the switch. Each chip carries direct links, ICIThe inter-chip links (4.5e10 bytes per second each way per link on v5e); every collective resolves to hops over these.taught in /l/ici →, to its physical neighbors: four of them on v5e and v6e, six on v4p, v5p, and 7x. A message to a distant chip hops chip to chip, forwarded by the chips between. What the design buys is brutal simplicity of scale: doubling the pod doubles the chips and the cables in the same proportion, with no switch tier growing in the middle to pay for.

What it costs is distance. Two chips across the pod are many hops apart, and every hop spends link bandwidth on forwarding. The entire discipline of distributed TPU work, deciding which axis of your mesh maps to which axis of the machine, is the management of that distance.

§ 02

The torus and the wraparound

Four neighbors make a grid; six make a 3D lattice. Then the edges fold: the last chip in a row links back to the first, turning each straight axis into a ring. A grid whose every axis is a ring is a torus, and the fold is worth exactly a factor of two, because the farthest chip on a ring of n is n/2 hops away instead of n-1. The fold is physical cabling, so it exists only when a slice spans a full axis: on 3D generations a slice gets wraparound as full 4x4x4 cubes and their multiples, on 2D generations when an axis reaches the full 16.

The maximum shapes from the generation table are torus dimensions: 16x16 on v5e and v6e, 16x16x16 on v4p, 16x20x28 on v5p, and the tall 4x4x576 of a full 7x pod, 9216 chips. Pick two chips in the stepper below and watch what the fold does to the hop count; then take an axis away and watch it come back.

EX·20 the torus, hop by hop
7 hops3 on x · 4 on yworst pair on this grid:8 hops with wrapeach hop re-sends over one9.0e+10 B/s linkclick chips to move the pair
an 8x8 axis pair standing for 16x16 · hop counts exact for the drawn grid · link constants from jax-ml.github.io/scaling-book/tpus/

The links themselves are constants you can hold: 4.5e10 bytes per second one-way per link on v4p and v5e, 9.0e10 on v5p, v6e, and 7x, double it for both directions at once. A chip's total egress is its per-link rate times its neighbor count, so a v6e chip can push 3.6e11 bytes per second outward if all four links run hot, and a collective is well designed exactly when they do.

That is enough to bound real operations from the armchair. An all-gather of D bytes around one ring axis of n chips must move each shard past n-1 positions; run the ring perfectly and the time floor is D(n-1)/n divided by the one-way link rate. Nothing in that formula is measured, and yet it called the shape of every ring-collective benchmark on this site's bench within the factor the rooflineThe floor model: latency is at least the larger of FLOPs over peak compute and bytes over bandwidth. Predict first, measure second.taught in /l/tpu → habit taught you to expect. When a measured collective misses the floor badly, the miss has a name (a congested axis, a slice without wraparound, DCN in the path), and finding which is the debugging.

A collective is well designed exactly when every link is busy.
§ 04

Past the slice

ICIThe inter-chip links (4.5e10 bytes per second each way per link on v5e); every collective resolves to hops over these.taught in /l/ici → ends at the slice boundary. Beyond it, traffic falls onto the data-center network, and the cliff is the point of this section: DCN carries 3.125e9 bytes per second per chip on v5e, 6.25e9 on v5p, 1.25e10 on v6e and 7x, between one and two orders of magnitude below the ICI numbers above. The host connection has the same flavor, PCIe at around 1.6e10 bytes per second per TPU generally and 3.2e10 on v6e. The hierarchy dictates strategy outright: the parallelism axes that communicate most stay inside the slice, and whatever crosses slices had better communicate rarely. When a later stage places data parallel across DCN and model parallel inside ICI, it is reading this table, not expressing a preference.

before you move on

Check yourself

01 What does wraparound buy, and when does a slice have it?

A factor of two on worst-case hops, since a ring's far chip is n/2 away instead of n-1. It exists when the slice spans full 4x4x4 cubes on 3D generations or a full axis of 16 on 2D ones.

02 Why do the chattiest parallelism axes stay inside the slice?

DCN carries one to two orders of magnitude less per chip than ICI, so whatever crosses slices has to communicate rarely.

assigned

Readings