Chapter 7: All About Transformer Inference

Scaling Book Exercises – Chapter 7

Chapter: 7 – All About Transformer Inference
Source: handwritten Onyx Boox A4 PDF
Note: Transcribed from handwritten solutions; book markdown used only for exercise statements and notation.

Shared model setup for Exercises 1-7

The chapter uses an invented model based on LLaMA-2 13B with $𝐿 = 64$ (num_layers), $𝐷 = 4096$ (d_model), $𝐹 = 16384$ (ffw_dimension), $𝑁 = 32$ (num_heads), $𝐾 = 8$ (num_kv_heads), $𝐻 = 256$ (qkv_dim), and $𝑉 = 32128$ (num_embeddings).

Exercise 1 – Parameter count and int8 KV cache

Exercise statement

How many parameters does the above model have? How large are its KV caches per token in int8? You can assume we share the input and output projection matrices.

Solution

Scan pages: 1

From previous chapters we know that number of params is:

\begin{aligned} Params & = (3 𝐷 𝐹 + 2 𝐷 (𝑁 + 𝐾) \cdot 𝐻) \cdot 𝐿 + 𝐷 𝑉 \\ = (3 \cdot 4096 \cdot (4 \cdot 4096) + 2 \cdot 4096 \cdot 40 \cdot 256) \cdot 64 + 4096 \cdot 32128 \\ = 18 385 207 246 \approx 18 \cdot 10^{9} = 18 Billion . \end{aligned}

Per-token KV cache size is:

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

Hence in int8:

2 \cdot 64 \cdot 8 \cdot 256 = 262144 bytes \approx 262 KB .

Exercise 2 – Largest batch size at 128K context

Exercise statement

Say we want to serve this model on a TPUv5e 4x4 slice and can fully shard our KV cache over this topology. What’s the largest batch size we can fit, assuming we use int8 for everything and want to support 128k sequences? What if we dropped the number of KV heads to 1?

Solution

Scan pages: 2-3

Assume TPU v5e 4x4 slice where v5e HBM = $16 GB$ .

Note that:

Space per shard = \frac{𝑀 + 𝐵 2 𝑆 𝐿 𝐾 𝐻}{| 𝑋 |} .

\frac{𝑀}{| 𝑋 |} + \frac{𝐵 \cdot 2 \cdot 𝑆 \cdot 𝐿 \cdot 𝐾 \cdot 𝐻}{| 𝑋 |} \leq HBM .

where $𝑀$ is total params in int8. Therefore the largest batch size we can fit in is:

\begin{aligned} \frac{𝑀}{| 𝑋 |} + \frac{𝐵 \cdot 2 \cdot 𝑆 \cdot 𝐿 \cdot 𝐾 \cdot 𝐻}{| 𝑋 |} & \leq HBM \\ \Leftrightarrow \frac{𝑀}{| 𝑋 |} - HBM & \leq - \frac{𝐵 \cdot 2 \cdot 𝑆 \cdot 𝐿 \cdot 𝐾 \cdot 𝐻}{| 𝑋 |} \\ \Leftrightarrow \frac{| 𝑋 | \cdot HBM - 𝑀}{| 𝑋 |} & \geq \frac{𝐵 \cdot 2 \cdot 𝑆 \cdot 𝐿 \cdot 𝐾 \cdot 𝐻}{| 𝑋 |} \\ \Leftrightarrow \frac{| 𝑋 | \cdot HBM - 𝑀}{2 \cdot 𝑆 \cdot 𝐿 \cdot 𝐾 \cdot 𝐻} & \geq 𝐵 . \end{aligned}

For $𝑆 = 128 K$ we get:

𝐵 \leq \frac{16 \cdot (16 \cdot 10^{9}) - 18 \cdot 10^{9}}{2 \cdot 128 \cdot 10^{3} \cdot 64 \cdot 8 \cdot 256} = 7.04 \approx 7 .

And if we let $𝐾 = 1$ we get:

𝐵 \leq \frac{16 \cdot 16 \cdot 10^{9} - 18 \cdot 10^{9}}{2 \cdot 128 \cdot 10^{3} \cdot 64 \cdot 1 \cdot 256} \approx 56 .

Exercise 3 – Parameter-loading latency

Exercise statement

How long does it take to load all the parameters into the MXU from HBM assuming they’re fully sharded on a TPU v5e 4x4 slice? Assume int8 parameters. This is a good lower bound on the per-step latency.

Solution

Scan pages: 3

Assume TPU v5e 4x4 slice where v5e $BW = 8.2 \cdot 10^{11}$ . Therefore to load $\frac{𝑀}{| 𝑋 |}$ into the MXU it takes:

\frac{𝑀}{| 𝑋 |} \cdot \frac{1}{BW} = \frac{18 \cdot 10^{9}}{16} \cdot \frac{1}{8.2 \cdot 10^{11}} \approx 0.0013 s = 1.3 ms .

Exercise 4 – Prefill and generation sharding

Exercise statement

Let’s say we want to serve this model on a TPUv5e 4x4 slice using int8 FLOPs and parameters/activations. How would we shard it for both prefill and decode? Hint: maybe answer these questions first:

What does ICI look like on a 4x4?
What’s the roofline bound on tensor parallelism?
How can we shard the KV caches?

For this sharding, what is the rough per-step latency for generation?

Solution

Scan pages: 4-6

Assume TPU v5e 4x4 slice. With this slice we don’t have wraparound links. To shard it both for prefill and generation, since during generation FSDP does not help, we could use tensor parallelism sharded over all the devices for both prefill and generation.

For prefill we are compute bound when:

𝑌 < 𝐹 \cdot \frac{𝑊_{ICI}}{𝐶} \cdot 2 .

Let $𝐹 = 4 \cdot 4096$ , $𝑊_{ICI} = 9 \cdot 10^{10}$ , and $𝐶 = 3.94 \cdot 10^{14}$ . Hence:

𝑌 < 16384 \cdot \frac{9 \cdot 10^{10}}{3.94 \cdot 10^{14}} \cdot 2 \approx 7.5 .

So we cannot do tensor parallelism with 16 devices and be compute bound.

We could instead do TP and CP with 4 way TP and 4 way CP, to be compute bound. So then a roofline bound will be:

𝑡 = \frac{4 \cdot 𝐵 𝐷 𝐹}{𝑋 𝑌 𝐶}, where 𝐵 = 7 \cdot 128 K .

During generation we will be memory bound instead of communication bound when using tensor parallelism when:

𝑌 < 𝐹 \cdot \frac{1}{𝐵} \cdot \frac{𝑊_{ICI}}{𝑊_{HBM}} .

where $\frac{𝑊_{ICI}}{𝑊_{HBM}} \approx \frac{1}{8}$ in a TPU v5e. Hence:

𝑌 < 4 \cdot 4096 \cdot \frac{1}{𝐵} \cdot \frac{1}{8} = \frac{2048}{𝐵} .

Since for $𝑆 = 128 K$ the maximum batch size is 7, then we get:

𝑌 < \frac{2048}{7} = 292 .

Therefore during generation we will be memory bound and a roofline bound will be:

𝑡 = \frac{2 𝐷 𝐹}{𝑋 \cdot 𝑊_{HBM}} .

We can shard the KV caches as:

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

where the mesh is defined as {X: 2, Y: 8}.

A rough per-step latency under this sharding is:

\begin{aligned} \frac{𝐵 \times KV cache size + Parameter size}{16 \cdot 𝑊_{HBM}} & = \frac{7 \times (2 \cdot 128 \cdot 10^{3} \cdot 64 \cdot 8 \cdot 256) + 18 \cdot 10^{9}}{16 \cdot 8.2 \cdot 10^{11}} \\ \approx 0.02 s . \end{aligned}

Exercise 5 – MoE parameter and FLOP accounting

Exercise statement

Let’s pretend the above model is actually an MoE. An MoE model is effectively a dense model with $𝐸$ copies of the FFW block. Each token passes through $𝑘$ of the FFW blocks and these $𝑘$ are averaged to produce the output. Let’s use $𝐸 = 16$ and $𝑘 = 2$ with the above settings.

How many total and activated parameters does it have? Activated means used by any given token.
What batch size is needed to become FLOPs bound on TPU v5e?
How large are its KV caches per token?
How many FLOPs are involved in a forward pass with $𝑇$ tokens?

Solution

Scan pages: 7-8

5.1

Let $𝐸 = 16$ and $𝑘 = 2$ , then:

Total Params = (3 𝐸 𝐷 𝐹 + 2 𝐷 (𝑁 + 𝐾) \cdot 𝐻) \cdot 𝐿 + 𝐷 𝑉 \approx 212 Billion .

Activated Params = (3 𝑘 𝐷 𝐹 + 2 𝐷 (𝑁 + 𝐾) \cdot 𝐻) \cdot 𝐿 + 𝐷 𝑉 \approx 31 Billion .

5.2

From the approximation of the MLP part they use in this chapter, to be FLOPs bound we need $𝑇_{math} > 𝑇_{comms}$ :

\begin{aligned} \frac{2 \cdot 𝐵 \cdot 3 𝑘 𝐷 𝐹}{𝐶} & > \frac{2 \cdot 3 \cdot 𝐸 𝐷 𝐹}{BW} \\ \Leftrightarrow \frac{2 \cdot 𝐵 \cdot 3 𝑘 𝐷 𝐹}{2 \cdot 3 \cdot 𝐸 𝐷 𝐹} & > \frac{𝐶}{BW} \\ \Leftrightarrow 𝐵 \cdot \frac{𝑘}{𝐸} & > \frac{𝐶}{BW} \\ \Leftrightarrow 𝐵 & > \frac{𝐸}{𝑘} \cdot \frac{𝐶}{BW} . \end{aligned}

5.3

KV cache size per token does not change.

5.4

Using the formula from Chapter 4:

FLOPs = 6 \cdot 𝐵 𝑇 (3 𝑘 𝐷 𝐹 + 2 \cdot 𝐷 (𝑁 + 𝐾) \cdot 𝐻) \cdot 𝐿 .

Exercise 6 – Expert sharding and minimum slice

Exercise statement

With MoEs, we can do “expert sharding”, where we split our experts across one axis of our mesh. In our standard notation, our first FFW weight has shape [E, D, F] and we shard it as $[𝐸_{𝑍}, 𝐷_{𝑋}, 𝐹_{𝑌}]$ where $𝑋$ is only used during training as our FSDP dimension. Let’s say we want to do inference on a TPU v5e:

What’s the HBM weight loading time for the above model on a TPU v5e 8x16 slice with $𝑌 = 8$ , $𝑍 = 16$ ? How much free HBM is available per TPU?
What is the smallest slice we could fit our model on?

Solution

Scan pages: 8-10

6.1

Assuming we shard the MLPs as:

𝑊_{in1} [𝐸_{𝑍}, 𝐷, 𝐹_{𝑌}], 𝑊_{in2} [𝐸_{𝑍}, 𝐷, 𝐹_{𝑌}], 𝑊_{out} [𝐸_{𝑍}, 𝐹_{𝑌}, 𝐷] .

The attention params as:

𝑊_{𝑞 [𝐷, 𝑁_{𝑌}, 𝐻]}, 𝑊_{𝑘 [𝐷, 𝐾_{𝑌}, 𝐻]}, 𝑊_{𝑣 [𝐷, 𝐾_{𝑌}, 𝐻]}, 𝑊_{𝑜 [𝑁_{𝑌}, 𝐻, 𝐷]} .

and the shared input-output embedding as:

𝑉 [𝑉_{𝑌}, 𝐷] .

The total params per device is:

Params Per Device = (3 \frac{𝐸}{| 𝑍 |} 𝐷 \frac{𝐹}{| 𝑌 |} + 2 \frac{𝐷}{| 𝑌 |} (𝑁 + 𝐾) 𝐻) \cdot 𝐿 + \frac{𝐷 𝑉}{| 𝑌 |} \approx 2.3 Billion .

Hence the weight loading time in bf16:

\frac{2 \cdot 2.3 \cdot 10^{9}}{8.2 \cdot 10^{11}} \approx 0.005 s = 5 ms .

As free space we have approximately:

16 GB - 4.6 GB = 11.4 GB .

6.2

Fit $𝑍 = 16$ so that we can shard our expert copies, then we need to choose $𝑌$ so that:

\begin{aligned} 2 \cdot \frac{1}{| 𝑌 |} ((3 𝐷 𝐹 + 2 𝐷 (𝑁 + 𝐾) \cdot 𝐻) \cdot 𝐿 + 𝐷 𝑉) & \leq 16 \\ \Leftrightarrow \frac{2 \cdot 18 \cdot 10^{9}}{16} & \leq | 𝑌 | \\ \Leftrightarrow | 𝑌 | & \geq 2.25 . \end{aligned}

Here if we need a slice with axis sizes being multiple of 2, the smallest slice we could fit our model on is $4 \times 16$ .

Exercise 7 – 2D model sharding

Exercise statement

Here we’ll work through the math of what the ESTI paper calls 2D weight-stationary sharding. We describe this briefly in Appendix B, but try doing this problem first to see if you can work out the math. The basic idea of 2D weight stationary sharding is to shard our weights along both the $𝐷$ and $𝐹$ axes so that each chunk is roughly square. This reduces the comms load and allows us to scale slightly farther.

Here’s the algorithm for 2D weight stationary:

$In [𝐵, 𝐷_{𝑋}] = {AllGather}_{𝑌 𝑍} (In [𝐵, 𝐷_{𝑋 𝑌 𝑍}])$
$Tmp [𝐵, 𝐹_{𝑌 𝑍}] {𝑈_{𝑋}} = In [𝐵, 𝐷_{𝑋}] \cdot_{𝐷} 𝑊_{in} [𝐷_{𝑋}, 𝐹_{𝑌 𝑍}]$
$Tmp [𝐵, 𝐹_{𝑌 𝑍}] = {AllReduce}_{𝑋 (Tmp [𝐵, 𝐹_{𝑌 𝑍}] {𝑈_{𝑋}})}$
$Out [𝐵, 𝐷_{𝑋}] {𝑈_{𝑌 𝑍}} = Tmp [𝐵, 𝐹_{𝑌 𝑍}] \cdot_{𝐹} 𝑊_{out} [𝐹_{𝑌 𝑍}, 𝐷_{𝑋}]$
$Out [𝐵, 𝐷_{𝑋 𝑌 𝑍}] = {ReduceScatter}_{𝑌 𝑍} (Out [𝐵, 𝐷_{𝑋}] {𝑈_{𝑌 𝑍}})$

Your goal is to work out $𝑇_{math}$ and $𝑇_{comms}$ for this algorithm and find when it will outperform traditional 3D model sharding?

Solution

Scan pages: 10-14

Let $𝑁 = | 𝑋 | \cdot | 𝑌 | \cdot | 𝑍 |$ be the total number of devices in the slice. From the algorithm we can see that:

\begin{aligned} 𝑇_{comms} & = \frac{2 𝐵 𝐷}{| 𝑋 |} \cdot \frac{1}{2 \cdot 𝑊_{ICI}} + \frac{2 \cdot 2 \cdot 𝐵 𝐹}{| 𝑌 | \cdot | 𝑍 |} \cdot \frac{1}{𝑊_{ICI}} + \frac{2 𝐵 𝐷}{| 𝑋 |} \cdot \frac{1}{2 \cdot 𝑊_{ICI}} \\ = \frac{2 𝐵 𝐷}{| 𝑋 | \cdot 𝑊_{ICI}} + \frac{4 𝐵 𝐹 \cdot | 𝑋 |}{𝑁 \cdot 𝑊_{ICI}} \\ = \frac{1}{𝑊_{ICI}} (\frac{2 𝐵 𝐷}{| 𝑋 |} + \frac{4 𝐵 𝐹 \cdot | 𝑋 |}{𝑁}) . \end{aligned}

\begin{aligned} 𝑇_{math} & = 2 \cdot 𝐵 \cdot \frac{𝐷}{| 𝑋 |} \cdot \frac{𝐹}{| 𝑌 | \cdot | 𝑍 |} \cdot \frac{1}{𝐶} + 2 \cdot 𝐵 \cdot \frac{𝐹}{| 𝑌 | \cdot | 𝑍 |} \cdot \frac{𝐷}{| 𝑋 |} \cdot \frac{1}{𝐶} \\ = \frac{4 𝐵 𝐷 𝐹}{| 𝑋 | \cdot | 𝑌 | \cdot | 𝑍 |} \cdot \frac{1}{𝐶} \\ = \frac{4 𝐵 𝐷 𝐹}{𝑁 \cdot 𝐶} . \end{aligned}

Note that given $𝑁$ devices, $𝑇_{comms}$ can change based in how we allocate devices for the different axes $𝑋$ , $𝑌$ , $𝑍$ . Therefore we need to find the setting where $𝑇_{comms}$ is minimum.

𝑇_{comms} (| 𝑋 |) = \frac{1}{𝑊_{ICI}} (\frac{2 𝐵 𝐷}{| 𝑋 |} + \frac{4 𝐵 𝐹 \cdot | 𝑋 |}{𝑁}) .

\frac{𝜕 𝑇_{comms}}{𝜕 | 𝑋 |} = \frac{1}{𝑊_{ICI}} (2 𝐵 𝐷 \cdot (- 1) \cdot \frac{1}{| 𝑋 |^{2}} + \frac{4 𝐵 𝐹}{𝑁}) .

Solving $\frac{𝜕 𝑇_{comms}}{𝜕 | 𝑋 |} = 0$ , we get a candidate for a stationary point:

\begin{aligned} \frac{4 𝐵 𝐹}{𝑁} & = \frac{2 𝐵 𝐷}{| 𝑋 |^{2}} \\ \Leftrightarrow \frac{𝑁}{4 𝐵 𝐹} \cdot 2 𝐵 𝐷 & = | 𝑋 |^{2} \\ \Leftrightarrow | 𝑋 |^{2} & = \frac{𝐷}{2 𝐹} \cdot 𝑁 . \end{aligned}

Since $| 𝑋 |$ can only be positive in our case, we get:

𝑋^{*} = \sqrt{\frac{𝐷}{2 𝐹} \cdot 𝑁} .

Let $𝐹 = 4 𝐷$ to get:

𝑋^{*} = \sqrt{\frac{𝑁}{8}} .

Since $\frac{𝜕^{2} 𝑇_{comms}}{𝜕 | 𝑋 |^{2}} > 0$ at $| 𝑋 | = 𝑋^{*}$ , we know that $𝑋^{*}$ is a global minimum and $𝑇_{comms}$ becomes:

\begin{aligned} 𝑇_{comms} & = \frac{1}{𝑊_{ICI}} (\frac{2 𝐵 𝐷 \sqrt{8}}{\sqrt{𝑁}} + \frac{4 𝐵 𝐹 \sqrt{𝑁}}{𝑁 \sqrt{8}}) \\ = \frac{1}{𝑊_{ICI}} (\frac{2 𝐵 𝐷 \sqrt{8}}{\sqrt{𝑁}} + \frac{4 𝐵 𝐹}{\sqrt{𝑁} \sqrt{8}}) \\ = \frac{1}{𝑊_{ICI}} (\frac{16 𝐵 𝐷}{\sqrt{8} \sqrt{𝑁}} + \frac{4 𝐵 𝐹}{\sqrt{𝑁} \sqrt{8}}) \\ = \frac{1}{\sqrt{8} \cdot 𝑊_{ICI}} (\frac{4 𝐵 \cdot (4 𝐷 + 𝐹)}{\sqrt{𝑁}}) \\ = \frac{1}{\sqrt{8} \cdot 𝑊_{ICI}} (\frac{4 𝐵 \cdot 8 𝐷}{\sqrt{𝑁}}) \\ = \frac{1}{\sqrt{8} \cdot 𝑊_{ICI}} \cdot \frac{32 𝐵 𝐷}{\sqrt{𝑁}} . \end{aligned}

\begin{aligned} \frac{𝑇_{math}}{𝑇_{comms}} > 1 & \Leftrightarrow \frac{4 𝐵 𝐷 𝐹}{𝑁 𝐶} \cdot \frac{\sqrt{8} 𝑊_{ICI} \sqrt{𝑁}}{32 𝐵 𝐷} > 1 \\ \Leftrightarrow \frac{4 \sqrt{8} 𝑊_{ICI} \sqrt{𝑁} 𝐹}{32 𝑁 𝐶} > 1 \\ \Leftrightarrow \frac{4 \sqrt{8}}{32} \cdot \frac{1}{\sqrt{𝑁}} \cdot \frac{𝑊_{ICI} 𝐹}{𝐶} > 1 \\ \Leftrightarrow \frac{16 \cdot 8}{32^{2}} \cdot \frac{1}{𝑁} \cdot \frac{𝑊_{ICI}^{2} 𝐹^{2}}{𝐶^{2}} > 1 \\ \Leftrightarrow \frac{16 \cdot 8}{32^{2}} \cdot \frac{𝐹^{2}}{𝛼^{2}} > 𝑁 \\ \Leftrightarrow \frac{1}{8} {(\frac{𝐹}{𝛼})}^{2} > 𝑁 . \end{aligned}

If we compare to traditional 3D model sharding where the condition to be compute bound is:

𝑀_{𝑦} \cdot \frac{𝐹}{𝛼} > 𝑁 .

We see that with 2D weight stationary:

\frac{\frac{1}{8} {(\frac{𝐹}{𝛼})}^{2}}{𝑀_{𝑥 𝑦 𝑧} (\frac{𝐹}{𝛼})} = \frac{1}{8 𝑀_{𝑥 𝑦 𝑧}} \cdot \frac{𝐹}{𝛼} .

During inference however if we are communication bound in the MLP because of a small batch size, we see that:

\begin{aligned} \frac{𝑇_{comms-2D}}{𝑇_{comms-3D}} & = \frac{32}{\sqrt{8}} \cdot \frac{𝐵 𝐷}{𝑊_{ICI} \sqrt{𝑁}} \cdot \frac{𝑊_{ICI} 𝑀_{𝑥 𝑦 𝑧}}{4 𝐵 𝐷} \\ = \frac{32}{\sqrt{8} \cdot 4} \cdot \frac{𝑀_{𝑥 𝑦 𝑧}}{\sqrt{𝑁}} \\ = 2 \cdot \sqrt{2} \cdot \frac{𝑀_{𝑥 𝑦 𝑧}}{\sqrt{𝑁}} . \end{aligned}