Scaling Book Part 2. TPU Questions

TPU Architecture Practice Questions

Question 1: LLM Latency Bounds

Say you want to sample from a 200B parameter model in bf16 that’s split across 32 TPU v4p. How long would it take to load all the parameters from HBM into the systolic array?

Model size = $200\ast 10^9\ast 2\text{bytes}=400GB=4e11\text{bytes}$

Model size per chip= $4e11/32=1.25e10$

TPU v4p HBM BW = $1.23e12\text{bytes/s}$

Time = $1.25e10/1.23e12=0.01s=10ms$

Question 2: TPU Pod Specifications

Consider a full TPU v5e pod. How many total CPU hosts are there? How many TPU TensorCores? What is the total FLOPs/s for the whole pod? What is the total HBM?

Do the same exercise for a TPU v5p pod.

Chip Type	Pod Size	Host Size	Total Hosts	Total FLOPs/s	Total HBM
TPU v5e	16×16	4×2	8	$9.2e14\ast 16\ast 16=2.35e17$	$16\ast 16\ast 16GB=4.096TB$
TPU v5p	16x20x28	2x2x1	4	$4.59e14\ast 16\ast 20\ast 28=4.11e18$	$16\ast 20\ast 28\ast 96GB=0.86PB$

Question 3: PCIe Operational Intensity

Imagine we’re forced to store a big weight matrix $A$ of type $\text{bfloat16}[D,F]$, and a batch of activations $x$ of type $\text{bfloat16}[B,D]$ in host DRAM and want to do a matrix multiplication on them. This is running on a single host, and we’re using a single TPU v6e chip attached to it.

You can assume $B\ll D$, and $F=4D$ (we’ll see in future chapters why these are reasonable assumptions). What is the smallest batch size $B$ we need to remain FLOPs bound over PCIe? Assume PCIe bandwidth of 1.5e10 bytes/second.

Intensity(matmul) = $BDF/(BD+DF+BF)$

Assuming $B<<D\text{and}F=4D$ :

Intensity(matmul) = $BD\ast 4D/D\ast 4D=B$

For flop-bound, $B>9.20e14/1.5e10=>B>6.13e4$

Correct answer and approach, incorrect thought process

Since the data needs to be first loaded into DRAM, I should not have thought of the data read operation ‘directly’ in these calculations above.

operations time = $2BDF/9.2e14$ seconds weight read+write time = $2(BD+DF+BF)/1.5e10$ seconds

This is assuming we can overload operations with weight loading (either overlap reading/writing data from DRAM while load from host happens, or overlap compute while read/write from DRAM happens)

Question 4: General MatMul Latency

Let’s say we want to multiply a weight matrix int8[16384, 4096] by an activation matrix of size int8[B, 4096] where B is some unknown batch size. Let’s say we’re on 1 TPUv5e to start.

a) How long will this multiplication take as a function of B?

b) What if we wanted to run this operation out of VMEM? How long would it take as a function of B?

Data size:

• W = $16384\ast 4096\text{bytes}=6.7e7\text{bytes}$
• X = $B\ast 4096\text{bytes}=B\ast .4e5\text{bytes}$
• X*W = $B\ast 16384\text{bytes}=B\ast 1.6e5\text{bytes}$

Ops = $2\ast 16384\ast 4096\ast B=B\ast 1.3e8$

For v5e chip:

Time (ops) = $B\ast 1.3e8/3.94e14=B\ast 3.3e-7$

Out of HBM:

Time (comm) = $(6.7e7+2Be4)/8.1e11=(827+0.24\ast B)e-7$

For FLOP-bound:

$3.3B>(827+0.246B)=>3.05B>827=>B>270$

Out of VMEM:

Time (comm) = 1/22 * ${\text{Time(comm)}}_{\text{HBM}}$

For FLOP-bound:

$3.3B\ast 22>(827+0.246B)=>72.35B>827=>B>11.42$

Question 5: ICI Bandwidth

Let’s say we have a TPU v5e 4x4 slice. Let’s say we want to send an array of type bfloat16[8, 128, 8192] from TPU{0,0} to TPU{3, 3}. Let’s say the per-hop latency for TPU v5e is $1\mu s$.

a) How soon will the first byte arrive at its destination?

b) How long will the total transfer take?

For v5e no wrap-around

Can send data down + right

Total bytes = $2\ast 8\ast 128\ast 8192=1.7e7$

Bytes per second = $2\ast 4.5e10=9e10$

First byte = 6 hops = $6\mu s$

Total time = $1.7e7/9e10=188\mu s$

Question 6: Multi-Component Performance (Hard)

Imagine you have a big matrix A: int8[128 * 1024, 128 * 1024] sharded evenly across a TPU v5e 4x4 slice but offloaded to host DRAM on each chip. Let’s say you want to copy the entire array to TPU{0, 0} and multiply it by a vector bf16[8, 128 * 1024]. How long will this take?

Steps:

1. Calculate size of array on each slice
2. Calculate data transfer time from each slice to T{0, 0}
3. Calculate time to read from T{0, 0} DRAM (both matrix and vector)
4. Calculate time to perform multiplication
5. Calculate time to write output to DRAM