Chapter 8: Serving LLaMA 3 on TPUs

Scaling Book Exercises – Chapter 8

Chapter: 8 – Serving LLaMA 3 on TPUs
Source: handwritten Onyx Boox A4 PDF
Note: Transcribed from handwritten solutions; book markdown used only for exercise statements and notation.

Remark: The following questions don’t appear numbered in the Scaling Book. I will number them in order of appearance as 0.X, where $𝑋 \in ℕ$ and $𝑋 \geq 1$ .

Exercise 0.1 – LLaMA 3-70B KV cache size

Exercise statement

How large are LLaMA 3-70B’s KV caches per token? Assume we store them in int8. This determines how large our batch size can be on a given topology.

Solution

Scan pages: 1

Per-token KV cache size is:

\frac{2 𝑆 𝐿 𝐾 𝐻}{𝑆} = 2 𝐿 𝐾 𝐻

where $𝑆$ is sequence length. Hence, in int8:

2 \cdot 80 \cdot 8 \cdot 128 = 163840 Bytes \approx 164 KB

Exercise 0.2 – Memory use and smallest slice

Exercise statement

Let’s say we want to serve L3 70B at batch size 32 and 8192 sequence length with everything (params and KVs) in int8. How much total memory will this use? What’s the smallest slice we could serve this on?

Solution

Scan pages: 1–2

Let $𝑀$ be the number of parameters of the model, let $𝐵 = 32$ and $𝑆 = 8192$ . Assume int8.

\begin{aligned} Total Memory & = 𝑀 + 𝐵 𝑆 \cdot 2 𝐿 𝐾 𝐻 \\ = 70 𝑒 9 + 𝐵 𝑆 \cdot 164 𝑒 3 \\ \approx 113 GB \end{aligned}

Assuming we serve this model in TPU v5e, the smallest possible slice is a $4 \times 2$ since it gives us $128 GB$ of aggregated HBM.

Exercise 0.3 – Decode latency and throughput

Exercise statement

At this batch size and quantization on a TPU v5e 4x2, roughly what latency would we expect per decode step? What throughput (tokens / sec / chip)? What about a 4x4? Assume we perform our FLOPs in bfloat16 and everything is fully sharded.

Solution

Scan pages: 2–3

At $𝐵 = 32$ the MLP block is not compute bound. Therefore, using the general equation from chapter 6:

Step Time = \frac{𝐵 \times KV cache size}{BW \cdot 𝑁} + \frac{𝑀}{BW \cdot 𝑁}

where $𝑁$ is the total number of devices in the slice. Hence:

\begin{aligned} Step Time & = \frac{32 \cdot 8192 \cdot 164 𝑒 3}{8.2 𝑒 11 \cdot 8} + \frac{70 𝑒 9}{8.2 𝑒 11 \cdot 8} \\ \approx \frac{113 𝑒 9}{8.2 𝑒 11 \cdot 8} \\ \approx 0.017 s = 17 ms \end{aligned}

\begin{aligned} Tokens/s/chip & = 𝐵 \cdot \frac{1}{0.017} \cdot \frac{1}{𝑁} \\ = \frac{𝐵}{0.017 \cdot 𝑁} \\ = \frac{32}{0.017 \cdot 8} \\ = 235 tokens/s/chip \end{aligned}

If instead we use a $4 \times 4$ slice, we get:

Step Time = \frac{113 𝑒 9}{8.2 𝑒 11 \cdot 16} \approx 0.008 s = 8 ms

Tokens/s/chip = \frac{32}{0.008 \cdot 16} = 250 tokens/s/chip

Exercise 0.4 – Compute-bound batch sizes

Exercise statement

On TPU v5e, using bfloat16 weights and activations, how large do our batch sizes need to be for us to be compute-bound in our matmuls? What if we do int8 weights but perform our FLOPs in bfloat16? What about int8 weights with int8 FLOPs?

Solution

Scan pages: 3–4

For bfloat16 weights and activations:

Intensity(Matmul) = 𝐵 \geq \frac{1.97 𝑒 14}{8.2 𝑒 11} \approx 240

If int8 weights but FLOPs in bfloat16, note that:

\begin{aligned} Intensity(Matmul) & = \frac{2 𝐵 𝐷 𝐹}{𝐵 𝐷 + 𝐷 𝐹 + 𝐵 𝐹} \\ = 2 \cdot \frac{𝐵 𝐷 𝐹}{𝐵 𝐷 + 𝐷 𝐹 + 𝐵 𝐹} \\ \approx 2 𝐵 \geq 240 \Leftrightarrow 𝐵 \geq 120 \end{aligned}

If int8 weights and int8 FLOPs:

Intensity(Matmul) = 2 𝐵 \geq \frac{2 \cdot 1.97 𝑒 14}{8.2 𝑒 11} \Leftrightarrow 𝐵 \geq 240

Exercise 0.5 – Smallest serving topology by precision

Exercise statement

What is the smallest TPU v5e topology we could serve LLaMA 3-70B on using bfloat16, int8, and int4 (both KVs and parameters) with 8k context? You can think of KV caches as negligibly small for this one.

Solution

Scan pages: 4–5

Assume $𝐵 = 1$ .

For bfloat16:

\begin{aligned} Total Memory & = 2 𝑀 + 2 𝐵 𝑆 \cdot 2 𝐿 𝐾 𝐻 \\ = 2 \cdot 70 𝑒 9 + 2 \cdot 8192 \cdot 164 𝑒 3 \\ \approx 143 GB \end{aligned}

So we need at least a $4 \times 4$ TPU v5e slice.

For int8:

\begin{aligned} Total Memory & = 𝑀 + 𝐵 𝑆 \cdot 2 𝐿 𝐾 𝐻 \\ = 70 𝑒 9 + 8192 \cdot 164 𝑒 3 \\ \approx 71 GB \end{aligned}

So we need at least a $4 \times 2$ TPU v5e slice.

For int4:

Total Memory = \frac{𝑀}{2} + \frac{𝐵 𝑆 \cdot 2 𝐿 𝐾 𝐻}{2} \approx 35.5 GB

So we need at least a $2 \times 2$ TPU v5e slice.

Exercise 0.6 – Generate-step latency at maximum batch size

Exercise statement

Assume we use the largest batch size that fits on these topologies. What latency could we expect for each generate step?

Solution

Scan pages: 5–6

For bfloat16 in a $4 \times 4$ slice:

\frac{2 𝑀 + 2 𝐵 𝑆 \cdot 2 𝐿 𝐾 𝐻}{4 \times 4} \leq 16 GB

Note that $𝐵$ cannot be $𝐵 \geq 240$ , for otherwise our KV cache size will not fit in this topology. Hence we are memory bound, and choosing the max $𝐵$ that satisfies the conditions just makes us closer to our upper bound. Therefore:

bfloat16 Step Time = \frac{16 𝑒 9}{8.2 𝑒 11} \approx 0.019 s = 19 ms

For int8, $𝐵$ cannot be $𝐵 \geq 240$ , hence as in the bfloat16 case:

int8 step time with max B = 19 ms

For int4, the same argument as above holds, therefore:

int4 step time with max B = 19 ms

Exercise 0.7 – Throughput per chip

Exercise statement

For each of these, what throughput per chip does this give us (in terms of queries / chip)? Assume our median decode length is 512 tokens.

Solution

Scan pages: 6–7

For each case we can use the condition we derived in Exercise 2 of chapter 7 to obtain the max $𝐵$ that fits in each topology.

For int8 and a $2 \times 4$ slice:

𝐵 \leq \frac{| 𝑋 | \cdot HBM - 𝑀}{2 𝑆 𝐿 𝐾 𝐻} = \frac{8 \cdot 16 - 70 𝑒 9}{8192 \cdot 164 𝑒 3} \approx 42

For bfloat16 and a $4 \times 4$ slice:

𝐵 \leq \frac{| 𝑋 | \cdot HBM - 2 𝑀}{2 𝑆 \cdot 2 𝐿 𝐾 𝐻} \approx 43

For int4 and a $2 \times 2$ slice:

𝐵 \leq \frac{| 𝑋 | \cdot HBM - \frac{𝑀}{2}}{\frac{𝑆 \cdot 2 𝐿 𝐾 𝐻}{2}} \approx 43

Therefore:

Queries/chip bfloat16, 4x4 slice = \frac{1}{0.014 \cdot 512} \cdot 42 \cdot \frac{1}{16} \approx 0.26

Queries/chip int8, 4x2 slice = \frac{1}{0.014 \cdot 512} \cdot 42 \cdot \frac{1}{8} \approx 0.53

Queries/chip int4, 2x2 slice = \frac{1}{0.014 \cdot 512} \cdot 42 \cdot \frac{1}{4} \approx 1.07

Exercise 0.8 – Doubling the topology

Exercise statement

How would our peak throughput change if we doubled our topology for each of the above examples?

Solution

Scan pages: 8

If we double our topology, we can see that:

𝐵 \leq \frac{2 \cdot | 𝑋 | \cdot HBM - 𝑀}{𝑆 \cdot 2 𝐿 𝐾 𝐻} \approx 138

for all cases. Therefore, since:

\frac{138}{42} \approx 3.3

we $1.65 \times$ our peak throughput in each case.

Exercise 0.9 – Sharding on a TPU v5e 4x8

Exercise statement

Now let’s dig into the question of sharding. Let’s say we wanted to serve in bfloat16 on a TPU v5e 4x8. What sharding would we use for our model on a TPU v5e 4x8 during generation? Can we avoid being communication bound?

Solution

Scan pages: 8–9

The only sharding option we have is model parallelism. From chapter 7 we know that we will not be communication bound if:

\frac{2 𝐷 𝐹}{𝑁 𝑊_{HBM}} > \frac{2 𝐵 𝐷}{2 𝑊_{ICI}} \Leftrightarrow \frac{2 𝐹}{𝐵 𝛽} > 𝑁

where:

𝛽 = \frac{𝑊_{HBM}}{𝑊_{ICI}}

Therefore, for our setting:

𝑁 < \frac{2 𝐹}{138 \cdot 8} \approx 52

So we can avoid being communication bound since $𝑁 = 32$ in a $4 \times 8$ TPU v5e slice.

The KV cache will be sharded along the head dimension and among the batch dimension as well:

KV [𝐵_{𝑋}, 2, 𝑆, 𝐿, 𝐾_{𝑌}, 𝐻]

where Mesh = {X: 4, Y: 8}.

Exercise 0.10 – Prefill time

Exercise statement

Assume we achieve a 40% FLOPs utilization during prefill. How long will a prefill of length 8192 take on 16 TPU v5e chips?

Solution

Scan pages: 9

We can use the formula:

Total FLOPs = 2 𝑀 𝐵

Hence:

𝑡 = \frac{2 𝑀 𝐵}{16} \cdot \frac{1}{𝐶 \cdot 0.4} \approx 0.9

where we assumed bfloat16 FLOPs.

Exercise 0.11 – Decode completions and KV eviction

Exercise statement

Assume we have a median prefill length of 8192 tokens and a median decode length of 4096 tokens. Say we have a generate batch size of 32. On average how many sequences finish decoding per step? On average how many tokens are evicted from our KV cache each step?

Solution

Scan pages: 10

At batch size 32 we are memory bound during decoding, therefore:

Step Time = \frac{32 \cdot 8192 \cdot 164 𝑒 3 + 70 𝑒 9}{𝑁 𝑊_{HBM}} \approx 0.008 s = 8 ms

Sequence Time = Step Time \times 4096 \approx 35 s

Decodings Per Second = \frac{1}{StepTime \cdot 4096} \cdot 32 \approx 0.9

Prefills Per Second = \frac{1}{0.9} \approx 1.1

Per prefill step I do:

\frac{0.9 \cdot 32}{35} = 0.82 decodings

So per prefill step I evict:

0.82 \cdot 4096 \approx 3358 tokens

Exercise 0.12 – Prefill-to-generate server ratio

Exercise statement

Assume we do disaggregated serving with a median prefill length of 8192 and a median decode length of 512. Assume the prefill and generate latencies calculated above in bfloat16. What ratio of prefill:generate servers will you need to keep both fully saturated?

Solution

Scan pages: 11

Assuming each server produces:

Prefills/s = \frac{1}{0.9}

Decodings/s = \frac{32}{0.008 \cdot 512} = 7.81

Therefore we need to produce prefills at the same rate as the decoding servers:

𝑁 \cdot \frac{1}{0.9} = 7.81 \Leftrightarrow 𝑁 = 7.029

So we need 7 more prefill servers than decoding servers.

Exercise 1 – LLaMA 3-405B forward-pass bounds

Exercise statement

How many FLOPs does each forward pass for LLaMA 3-405B use per-token? Assuming we’re FLOPs bound, what is a lower bound on a single forward pass on $𝑁$ chips on TPU v5e? What if we’re comms bound? Ignore the fact that the model does not fit on a single chip.

Solution

Scan pages: 12

We can use the formula from this chapter:

Total FLOPs = 2 𝑀 𝐵

Hence, per token:

FLOPs/Token = 2 𝑀 = 2 \cdot 405 𝑒 9 = 810 𝑒 9

If we are FLOPs bound, a lower bound on $𝑁$ TPU v5e chips is:

𝑡 \geq \frac{2 𝑀}{𝑁 𝐶} = \frac{810 𝑒 9}{𝑁 \cdot 1.97 𝑒 14}

If we are communication bound, a lower bound is:

𝑡 \geq \frac{2 𝐵 𝐷}{𝑀_{𝑋} \cdot 𝑊_{ICI}} = \frac{2 \cdot 𝐵 \cdot 16384}{𝑀_{𝑋} \cdot 9.0 𝑒 10}

Exercise 2 – LLaMA 3-8B serving memory

Exercise statement

Assume we want to serve LLaMA 3-8B with BS240 using int8 weights and int8 KV caches. How many bytes are used by (a) model parameters, (b) KV caches, and (c) peak working activations, roughly? What’s the smallest topology we can run this on?

Solution

Scan pages: 13

For LLaMA 3-8B we have $𝐿 = 32$ , $𝐷 = 4096$ , $𝐹 = 14336$ , $𝑁 = 32$ , $𝐾 = 8$ , $𝐻 = 128$ , and $𝑉 = 128256$ .

Assuming int8 for everything, we have:

𝑀 \approx 8 𝑒 9 Bytes

\frac{KV caches Size}{𝑆} = \frac{𝐵 𝑆 \cdot 2 𝐿 𝐾 𝐻}{𝑆} = 240 \cdot 2 \cdot 32 \cdot 8 \cdot 128 \approx 15.7 MB

For peak working activations we will assume:

In [𝐵, 𝑆, 𝐷] \cdot_{𝐷} 𝑊_{in} [𝐷, 𝐹] \to Out [𝐵, 𝑆, 𝐹]

Hence per token:

𝐵 𝐹 = 240 \cdot 14336 Bytes \approx 3.4 MB

are used. The smallest topology is a single TPU v5e.

Exercise 3 – LLaMA 3-405B serving under a 15 ms limit

Exercise statement

How would you serve LLaMA 3-405B on TPU v5e? Assume int8 weights and bfloat16 FLOPs. Let’s say we have a firm limit of 15ms / token. What’s the highest-throughput configuration we could achieve? What is the theoretical minimum step time?

Solution

Scan pages: 14–17

For LLaMA 3-405B we have $𝐿 = 126$ , $𝐷 = 16384$ , $𝐹 = 53248$ , $𝑁 = 128$ , $𝐾 = 8$ , $𝐻 = 128$ , and $𝑉 = 128256$ . Assume int8 weights and bfloat16 FLOPs. Then we know that:

KV cache size/token = 2 𝐿 𝐾 𝐻 = 2 \cdot 126 \cdot 8 \cdot 128 \approx 258 𝑒 3 Bytes = 258 KB

and:

\begin{aligned} Total Memory & = 𝑀 + 𝐵 𝑆 \cdot 2 𝐿 𝐾 𝐻 \\ = 405 𝑒 9 + 𝐵 𝑆 \cdot 258 𝑒 3 \end{aligned}

Hence we need a TPU v5e slice topology with at least $408 GB$ of aggregated HBM, assuming $𝐵 = 1$ . Since HBM of TPU v5e is $16 GB$ , we need:

| 𝑋 | \cdot 16 GB \geq 408 GB \Leftrightarrow | 𝑋 | \geq 25.5

Note that a lower bound in latency is:

\frac{405 𝑒 9}{| 𝑋 | \cdot 𝑊_{HBM}} \leq Step Time

Hence, if we need to have a latency of $\leq 15 ms$ , we have to satisfy:

\frac{405 𝑒 9}{0.015 \cdot 8.2 𝑒 11} \leq | 𝑋 | \Leftrightarrow | 𝑋 | \geq 32.9

So we need more than 32 devices. We can consider an $8 \times 8$ TPU v5e slice, and use model parallelism. During prefill we will be memory bound since:

\frac{𝐹}{2188} = 24 ≱ | 𝑋 | = 64

However, we can improve our latency during generation.

Note that for a choice of $𝐵$ we are memory bound if:

𝑇_{HBM comms} (int8) = \frac{𝐷 𝐹}{| 𝑋 | \cdot 𝑊_{HBM}}

𝑇_{ICI comms} (bfloat16) = \frac{2 𝐵 𝐷}{𝑊_{ICI}}

\frac{𝐷 𝐹}{| 𝑋 | \cdot 𝑊_{HBM}} > \frac{2 𝐵 𝐷}{𝑊_{ICI}} \Leftrightarrow \frac{𝐹 \cdot 𝑊_{ICI}}{2 \cdot | 𝑋 | \cdot 𝑊_{HBM}} > 𝐵

So then, in order to be memory bound, we need:

𝐵 \leq 52.0

At this batch size we’re memory bound in the MLP. We can solve now for our max $𝑆$ length that still fits:

𝑆 \leq \frac{| 𝑋 | \cdot HBM - 𝑀}{𝐵 \cdot 2 𝐿 𝐾 𝐻} \approx 46130

However, we need $𝑆$ to be smaller, for otherwise we will not hit the $15 ms$ requirement.

Step time = \frac{52 \cdot 𝑆 \cdot 258 𝑒 3 + 405 𝑒 9}{64 \cdot 8.2 𝑒 11} \leq 0.015

𝑆 \leq \frac{64 \cdot 8.2 𝑒 11 \cdot 0.015 - 405 𝑒 9}{52 \cdot 258 𝑒 3} \approx 28488

Therefore, with $𝐵 = 52$ and $𝑆 = 28488$ we have:

Step Time = 15 ms

and a throughput of:

Tokens/s = \frac{𝐵}{0.015} = \frac{52}{0.015} = 3466