Chapter 2: How to Think About TPUs

Scaling Book Exercises — Chapter 2

Chapter: 2 — How to Think About TPUs
Source: handwritten Onyx Boox A4 PDF
Note: Transcribed from handwritten solutions; book markdown used only for exercise statements and notation.

This document contains only the handwritten solution content from the scan, with exercise statements copied from the Chapter 2 book markdown and kept separate from the solutions.

Exercise 1 — Bounding LLM latency

Exercise statement

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? Hint: use the numbers above.

Solution

Scan pages: 1-2

v4p $\Rightarrow$

HBM BW = $1.2 𝑒 12$
FLOPs/s bf16 = $2.75 𝑒 14$
HBM capacity = $32 GB$

$200 Billion = 200 \cdot 10^{9}$ (bf16). Split across 32 TPUs v4p $\Rightarrow$ Bytes per TPU:

\frac{200 \cdot 10^{9} \cdot 2}{32} = \frac{400 \cdot 10^{9}}{32} = 12.5 \cdot 10^{9} bytes = 12.5 GB

$12.5 GB \leq 32 GB .$ So it does fit.

Each TPU HBM has $12.5 GB$ of weights to send to the Tensor Cores:

𝑇_{comms} = \frac{12.5 \cdot 10^{9} bytes}{1.2 𝑒 12 bytes/s} \approx 10 \cdot 10^{- 3} s = 10 ms .

Exercise 2 — TPU details

Exercise statement

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 TPU v5p pod.

Solution

Scan pages: 2-5

TPU v5e pod

Pod size: $16 \times 16$
Host size: $4 \times 2$
HBM per chip: $16 GB$
FLOPs bf16 per chip = $1.97 𝑒 14$
Host size: $2 \times 4 = 8$

$16 \times 16$ : how many blocks of $2 \times 4$ fit:

One host (schematic):
+----+----+----+----+
|    |    |    |    |
+----+----+----+----+
|    |    |    |    |
+----+----+----+----+

4 hosts to cover one “row”, 8 hosts to cover one “column” $\Rightarrow 32 hosts .$ I could also just flatten: $16 \times 16 = 256$

[ ][ ][ ][ ][ ][ ][ ][ ] ...
<---------- 8 ---------> is a host

and then divide i.e $\frac{256}{8} = 32$ hosts.

TPU v5e has only one Tensor Core per chip $\Rightarrow 256 Tensor Cores .$ Also you can derive by cores per host and from the book you see using previous result: $32 \times 8 = 256 Tensor Cores$ where 32 is hosts in the pod and 8 is cores per host.

Total FLOPs/s = $256 \times 1.97 𝑒 14 \approx 500 𝑒 14$

Total HBM = $256 \times 16 \cdot 10^{9} = 4096 𝑒 9 = 4096 GB .$

TPU v5p pod

Pod Size: $16 \times 20 \times 28$
Host Size: $2 \times 2 \times 1$

Flatten $\Rightarrow$ $16 \times 20 \times 28 = 8960 chips .$

Total Hosts = $\frac{8960}{4} = 2240$

Tensor Cores = $8960 \times 2 = 17920$

Total HBM:

8960 \times 96 \cdot 10^{9} = 860160 \cdot 10^{9} \approx 860 \cdot 10^{3} \cdot 10^{9} = 860 TB .

Total FLOPs/s:

8960 \times 4.59 𝑒 14 \approx 41000 \times 10^{14} \approx 4.1 Exa FLOPs/s .

Exercise 3 — PCIe operational intensity

Exercise statement

Imagine we’re forced to store a big weight matrix $𝐴$ of type bf16[D, F], and a batch of activations $𝑥$ of type bf16[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 $𝐵 ≪ 𝐷$ , and $𝐹 = 4 𝐷$ . What is the smallest batch size $𝐵$ we need to remain FLOPs bound over PCIe? Assume PCIe bandwidth of 1.6e10 bytes / second.

Solution

Scan pages: 6-7

[Host] -- PCIe --> [TPU]

Assume:

TPU v6e.
$bf16 [𝐵, 𝐷] \cdot_{𝐷} bf16 [𝐷, 𝐹] \to bf16 [𝐵, 𝐹]$
$𝐵 ≪ 𝐷$ , $𝐹 = 4 𝐷$ .
PCIe BW = $1.6 𝑒 10 bytes/s$ .
FLOPs/s = $9.20 𝑒 14$ .

Assuming I can overlap communication with compute:

AI ≔ \frac{2 \cdot 𝐵 𝐷 𝐹}{2 𝐵 𝐷 + 2 𝐷 𝐹 + 2 𝐵 𝐹} \geq \frac{9.20 \cdot 10^{14}}{1.6 \cdot 10^{10}}

Using $𝐹 = 4 𝐷$ and $𝐵 ≪ 𝐷$ :

AI \approx \frac{2 \cdot 𝐵 \cdot 𝐷 \cdot 4 𝐷}{2 \cdot 𝐷 \cdot 4 𝐷} = \frac{𝐵 \cdot 𝐷^{2}}{𝐷^{2}}

$\Rightarrow 𝐵 \geq 5.75 \cdot 10^{4} = 57.5 \cdot 10^{3}$

Like 57 sequences of 1024 seq len.

Exercise 4 — General matmul latency

Exercise statement

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

How long will this multiplication take as a function of $𝐵$ ? Hint: it may help to calculate how long it will take to load the arrays from HBM and how long the multiplication will actually take. Which is bottlenecking you?
What if we wanted to run this operation out of VMEM? How long would it take as a function of $𝐵$ ?

Solution

Scan pages: 7-13

v5e:

HBM BW = $8.2 𝑒 11$
FLOPs/s int8 = $3.94 𝑒 14$

VMEM BW $\approx$ HBM BW $\times 22$ $\approx 8.2 \cdot 10^{11} \times 22 = 180.4 \times 10^{11} bytes/s .$

$int8 [𝐵, 4096] \cdot int8 [4096, 16384] \to int8 [𝐵, 16384]$

$𝑇 = max (𝑇_{comms}, 𝑇_{math})$

Note: $16384 = 4 \cdot 4096$ , so let $𝐷 = 4096$ and $𝐹 = 4096 \cdot 4 \Rightarrow 𝐹 = 4 \cdot 𝐷$ as before.

$[𝐵, 𝐷] \cdot [𝐷, 4 𝐷] \to [𝐵, 4 𝐷]$

𝑇_{comms} = \frac{𝐵 \cdot 𝐷 + 4 \cdot 𝐷^{2} + 𝐵 \cdot 4 \cdot 𝐷}{8.2 𝑒 11} = \frac{𝐵 \cdot (𝐷 + 4 \cdot 𝐷) + 4 \cdot 𝐷^{2}}{8.2 𝑒 11} = \frac{𝐵 \cdot (5 𝐷) + 4 \cdot 𝐷^{2}}{8.2 𝑒 11}

𝑇_{math} = \frac{2 \cdot 𝐵 \cdot 𝐷 \cdot 𝐹}{3.94 𝑒 14} = \frac{8 \cdot 𝐵 \cdot 𝐷^{2}}{3.94 𝑒 14}

𝑇_{operation} = max (𝑇_{comms} (𝐵), 𝑇_{math} (𝐵)) = max (\frac{𝐵 \cdot 5 \cdot 𝐷 + 4 \cdot 𝐷^{2}}{8.2 𝑒 11}, \frac{8 \cdot 𝐵 \cdot 𝐷^{2}}{3.94 𝑒 14})

\Rightarrow 𝑇_{operation} (𝐵, BW) = max (\frac{𝐵 \cdot 5 \cdot 𝐷 + 4 \cdot 𝐷^{2}}{BW}, \frac{8 \cdot 𝐵 \cdot 𝐷^{2}}{3.94 𝑒 14})

Since: VMEM = $128 MiB = 128 \cdot 10^{6}$ .

$4096 \times (4 \cdot 4096) \approx 68 \cdot 10^{6}$ . At least the weights do fit in VMEM.

Bottleneck: compute bound when?

\begin{matrix} 𝑇_{comms} (𝐵, BW) \leq 𝑇_{math} (𝐵) \Leftrightarrow \frac{𝐵 \cdot 5 \cdot 𝐷 + 4 \cdot 𝐷^{2}}{BW} \leq \frac{8 \cdot 𝐵 \cdot 𝐷^{2}}{3.94 𝑒 14 FLOPs/s} \\ \Leftrightarrow \frac{𝐵 \cdot 5 \cdot 𝐷 + 4 \cdot 𝐷^{2}}{8 \cdot 𝐵 \cdot 𝐷^{2}} \leq \frac{BW}{3.94 𝑒 14} = {AI}^{- 1} \Leftrightarrow \frac{𝐵 \cdot 5 \cdot 𝐷}{8 \cdot 𝐵 \cdot 𝐷^{2}} + \frac{4 \cdot 𝐷^{2}}{8 \cdot 𝐵 \cdot 𝐷^{2}} \leq {AI}^{- 1} \\ \Leftrightarrow 0 \leq \frac{5}{8 \cdot 𝐷} + \frac{4}{8 \cdot 𝐵} \leq {AI}^{- 1} \Leftrightarrow \frac{4}{8 \cdot 𝐵} \leq {AI}^{- 1} - \frac{5}{8 \cdot 𝐷} \\ \Leftrightarrow \frac{1}{2 \cdot 𝐵} \leq {AI}^{- 1} - \frac{5}{8 \cdot 𝐷} \Leftrightarrow \frac{1}{𝐵} \leq 2 \cdot {AI}^{- 1} - \frac{10}{8 \cdot 𝐷} \\ \Leftrightarrow 𝐵 \geq \frac{1}{2 \cdot (\frac{1}{AI}) - \frac{10}{8 \cdot 𝐷}} \Leftrightarrow 𝐵 \geq \frac{1}{2 \cdot (\frac{BW}{3.94 𝑒 14}) - 3 \cdot 10^{- 4}} \end{matrix}

${BW}_{HBM} = 8.2 𝑒 11$

𝐵 \geq \frac{1}{2 \cdot \frac{8.2 𝑒 11}{3.94 𝑒 14} - 3 \cdot 10^{- 4}} \approx \frac{1}{4 \cdot 10^{- 3}} = {(4 \cdot 10^{- 3})}^{- 1} = \frac{1}{4} \cdot 10^{3} = 250

${BW}_{VMEM} = 180.4 𝑒 11$

𝐵 \geq \frac{1}{2 \cdot \frac{180.4 𝑒 11}{3.94 𝑒 14} - 3 \cdot 10^{- 4}} \approx 11 .

Exercise 5 — ICI bandwidth

Exercise statement

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

How soon will the first byte arrive at its destination?
How long will the total transfer take?

Solution

Scan pages: 14-17

4x4 slice, no wraparound.
Source = [00] = TPU{0,0}; receiver = [33] = TPU{3,3}.

P1: across the top row, then down the right column.
P2: down the left column, then across the bottom row.

[00]--P1-->[01]--P1-->[02]--P1-->[03]
  |                                |
  P2                               P1
  v                                v
[10]       [11]       [12]       [13]
  |                                |
  P2                               P1
  v                                v
[20]       [21]       [22]       [23]
  |                                |
  P2                               P1
  v                                v
[30]--P2-->[31]--P2-->[32]--P2-->[33]

v5e $\Rightarrow (4, 4)$ slice has no wraparound.

Assume $𝑇_{comms, hop} = 𝛼 + \frac{𝐷}{BW}$ for one hop. Then, assuming pipelining:

𝑇_{comms} = #Hops \cdot 𝛼 + \frac{𝐷}{BW} where 𝛼 = 1 𝜇 𝑠

$#Hops = 6 \Rightarrow$ first byte arrives after $6 𝜇 𝑠$ .

Different algos, but a lower bound on time is:

𝑇 \geq \frac{𝐷}{2 \cdot BW}

since it needs to receive $𝐷$ bytes at $2 * BW$ :

           |
           |
           v
----> [TPU[3, 3]]

Algo 1

Data Split = $\frac{2 \times 8 \times 128 \times 8192}{2}$

path1 = P1 arrows: [00] -> [01] -> [02] -> [03] -> [13] -> [23] -> [33]
path2 = P2 arrows: [00] -> [10] -> [20] -> [30] -> [31] -> [32] -> [33]

\Rightarrow max (\frac{𝐷 (bytes)}{2 \cdot {BW}_{path1}}, \frac{𝐷 (bytes)}{2 \cdot {BW}_{path2}}) = \frac{𝐷}{2 \cdot BW} \geq lower bound

Algo 2

Everything through path 1:

max (\frac{𝐷}{{BW}_{path1}}, 0) = \frac{𝐷}{{BW}_{path1}} > lower bound

BW = $4.5 𝑒 10$ . Therefore:

\frac{𝐷}{2} = \frac{2 \times 8 \times 128 \times 8192}{2} = 2^{3} \cdot 2^{7} \cdot 2^{13} = 2^{23}

\Rightarrow \frac{2^{23}}{4.5 \cdot 10^{10}} \approx 0.00018 = 0.18 \cdot 10^{- 3} s

$𝑇_{comms} = 6 𝜇 𝑠 + 0.18 ms .$

Exercise 6 — Pulling it all together

Exercise statement

Imagine you have a big matrix $𝐴$ : 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? Hint: use the numbers above.

Solution

Scan pages: 18-20

Question 6 setup, schematic:

Host 0 (4x2)             Host 1 (4x2)
+----+----+              +----+----+
|    |----|--------------|    |    |
+----+----+              +----+----+
|    |    |              |    |    |
+----+----+              +----+----+
|    |    |              |    |    |
+----+----+              +----+----+
|    |    |              |    |    |
+----+----+              +----+----+

No wraparound links.

$int8 𝐴 [128 \cdot 1024, 128 \cdot 1024]$ . Assume each host DRAM has an even part of the array $𝐴 \Rightarrow \frac{128^{2} \cdot 1024^{2}}{2} \approx 8.6 GB$ per host, so it fits in v5e HBM, which is $16 GB$ .

Path

[Host_B] --PCIe--> [TPU] --ICI--> [TPU[0,0]]

[Host_A] --PCIe--> [TPU[0,0]]

Assuming pipelining:

𝑇_{comms} (𝐴) ≔ \frac{8.6 GB}{{BW}_{PCIe}} = \frac{8.6 \cdot 10^{9}}{1.6 \cdot 10^{10}} \approx 5.3 \cdot 10^{- 1} s = 0.53 s

$Bytes in bf16 [8, 128 \times 1024] = 8 \times 128 \times 1024 \times 2 \approx 2 GB$ . So let’s assume it’s in HBM already. Assuming intended operation is:

[128 \times 1024, 128 \times 1024] @ [128 \times 1024, 8]

\Rightarrow 𝑇 = max (𝑇_{comms}, 𝑇_{math}) = max (0.53 s, \frac{{(128 \cdot 1024)}^{2} \cdot 8 \cdot 2}{1.97 𝑒 14}) \approx max (0.53 s, 0.001 s) = 0.53 s .