Chapter 4: All the Transformer Math You Need to Know

Scaling Book Exercises – Chapter 4

Chapter: 4 – All the Transformer Math You Need to Know
Source: handwritten Onyx Boox A4 PDF
Note: Transcribed from handwritten solutions; book markdown used only for exercise statements and notation.

Exercise 1 – Parameter count, attention fraction, and KV cache

Exercise statement

How many parameters does a model with $𝐷 = 4096$ , $𝐹 = 4 \cdot 𝐷$ , $𝑉 = 32, 000$ , and $𝐿 = 64$ have? What fraction of these are attention parameters? How large are our KV caches per token? You can assume $𝑁 \cdot 𝐻 = 𝐷$ and multi-head attention with int8 KVs.

Solution

Scan pages: 1–2

The total parameters of a model are:

(3 𝐷 𝐹 + 4 𝐷 𝑁 𝐻 + 𝐷) 𝐿 + 𝐷 𝑉 .

If we assume $𝐹 = 4 𝐷$ , $𝑁 𝐻 = 𝐷$ , we get:

(16 𝐷^{2} + 𝐷) 𝐿 + 𝐷 𝑉 .

Let $𝐷 = 4096$ , $𝑉 = 32000$ , $𝐿 = 64$ . Hence:

(16 \cdot {(4096)}^{2} + 4096) \cdot 64 + 4096 \cdot 32000 \approx 17.3 \cdot 10^{9} = 17.3 Billion parameters .

Attention params are $4 𝐷^{2} 𝐿$ hence:

4 \cdot {(4096)}^{2} \cdot 64 \approx 4.3 Billion,

so the fraction is:

\frac{4.3}{17.3} \approx 0.24 .

The size of the KV cache is $2 \cdot 𝑆 \cdot 𝐿 \cdot 𝑁 \cdot 𝐻$ and since $𝑁 𝐻 = 𝐷$ , we get $2 \cdot 𝐿 \cdot 𝐷$ per token i.e.

2 \cdot 64 \cdot 4096 \approx 524 KB .

Exercise 2 – FLOPs under sharding

Exercise statement

How many total FLOPs are required to perform $𝐴 [𝐵_{𝑋}, 𝐷_{𝑌}] \cdot_{𝐷} 𝑊 [𝐷_{𝑌}, 𝐹]$ on {'X': 4, 'Y': 8, 'Z': 4}. How many FLOPs are performed by each TPU?

Solution

Scan pages: 2

Let $Mesh = {𝑋 : 4, 𝑌 : 8, 𝑍 : 4}$ .

Assume the intended operation is $𝐴 [𝐵_{𝑋}, 𝐷_{𝑌}] \cdot_{𝐷} 𝑊 [𝐷_{𝑌}, 𝐹]$ . Each device $(𝑖, 𝑗, 𝑘)$ performs

2 \cdot \frac{𝐵}{| 𝑋 |} \cdot \frac{𝐷}{| 𝑌 |} \cdot 𝐹

FLOPs. Hence total FLOPs is:

Total FLOPs = | 𝑋 | \cdot | 𝑌 | \cdot | 𝑍 | \cdot 2 \cdot \frac{𝐵}{| 𝑋 |} \cdot \frac{𝐷}{| 𝑌 |} \cdot 𝐹 = 8 𝐵 𝐷 𝐹 .

Exercise 3 – FLOPs for a tensor contraction

Exercise statement

How many FLOPs are involved in performing $𝐴 [𝐼, 𝐽, 𝐾, 𝐿] \cdot 𝐵 [𝐼, 𝐽, 𝑀, 𝑁, 𝑂] \to 𝐶 [𝐾, 𝐿, 𝑀, 𝑁, 𝑂]$ ?

Solution

Scan pages: 3

Assume the intended operation is:

𝐴 [𝐼, 𝐽, 𝐾, 𝐿] \cdot 𝐵 [𝐼, 𝐽, 𝑀, 𝑁, 𝑂] \to 𝐶 [𝐾, 𝐿, 𝑀, 𝑁, 𝑂]

by means of:

𝐶 [𝐾, 𝐿, 𝑀, 𝑁, 𝑂] = \sum_{𝑗 = 1}^{𝐽} \sum_{𝑖 = 1}^{𝐼} 𝐴 [𝑖, 𝑗, 𝐾, 𝐿] \cdot 𝐵 [𝑖, 𝑗, 𝑀, 𝑁, 𝑂] .

we do $𝐾 \cdot 𝐿 \cdot 𝑀 \cdot 𝑁 \cdot 𝑂 \cdot 2 \cdot 𝐼 \cdot 𝐽$ FLOPs.

Exercise 4 – Self-attention arithmetic intensity and effective cost

Exercise statement

What is the arithmetic intensity of self-attention, ignoring the Q/K/V/O projections? Give the answer as a function of the Q and KV lengths $𝑇$ and $𝑆$ . At what context length is attention FLOPs-bound? Given the HBM bandwidth of our TPUs, plot the effective relative cost of attention to the FFW block as the context length grows.

Solution

Scan pages: 4–8

Assume our naive attention algorithm writes every intermediate result back to HBM and also that $𝐾 = 𝑁$ . Hence we do:

$𝐿 [𝐵, 𝑇, 𝑆, 𝑁] \leftarrow 𝑄 [𝐵, 𝑇, 𝑁, 𝐻] \cdot 𝐾 [𝐵, 𝑆, 𝑁, 𝐻]$ .
$𝑆 [𝐵, 𝑇, 𝑆, 𝑁] \leftarrow {SoftMax}_{𝑆} (𝐿 [𝐵, 𝑇, 𝑆, 𝑁])$ .
$𝑂 [𝐵, 𝑇, 𝑁, 𝐻] \leftarrow 𝑆 [𝐵, 𝑇, 𝑆, 𝑁] \cdot 𝑉 [𝐵, 𝑆, 𝑁, 𝐻]$ .

Assume that $𝑁 𝐻 = 𝐷$ .

Step 1:

\begin{matrix} Bytes Step 1 = 2 𝐵 𝑇 𝑁 𝐻 + 2 𝐵 𝑆 𝑁 𝐻 + 2 𝐵 𝑇 𝑆 𝑁 = 2 𝐵 𝑇 𝐷 + 2 𝐵 𝑆 𝐷 + 2 𝐵 𝑇 𝑆 𝑁, \\ FLOPs Step 1 = 2 𝐵 𝑁 𝑇 𝑆 𝐻 = 2 𝐵 𝑇 𝑆 𝐷 . \end{matrix}

Step 2:

\begin{matrix} Bytes Step 2 = 2 𝐵 𝑇 𝑆 𝑁 + 2 𝐵 𝑇 𝑆 𝑁 = 4 𝐵 𝑇 𝑆 𝑁, \\ FLOPs Step 2 = 𝑂 (𝐵 𝑇 𝑆 𝑁) . \end{matrix}

Step 3:

\begin{matrix} Bytes Step 3 = 2 𝐵 𝑇 𝑆 𝑁 + 2 𝐵 𝑆 𝑁 𝐻 + 2 𝐵 𝑇 𝑁 𝐻 = 2 𝐵 𝑇 𝑆 𝑁 + 2 𝐵 𝑆 𝐷 + 2 𝐵 𝑇 𝐷, \\ FLOPs Step 3 = 2 𝐵 𝑁 𝑇 𝐻 𝑆 = 2 𝐵 𝑇 𝑆 𝐷 . \end{matrix}

Assume Flash Attention, so $𝐿 [𝐵, 𝑇, 𝑆, 𝑁]$ does not need to be in HBM.

\begin{aligned} Bytes Total & = 4 (𝐵 𝑆 𝐷 + 𝐵 𝑇 𝐷) \\ FLOPs Total & = 2 𝐵 𝑇 𝑆 𝐷 + 2 𝐵 𝑇 𝑆 𝐷 = 4 𝐵 𝑇 𝑆 𝐷 \\ AI (𝑇, 𝑆) & = \frac{𝐵 𝑇 𝑆 𝐷}{𝐵 𝑆 𝐷 + 𝐵 𝑇 𝐷} = \frac{𝑇 𝑆 𝐷}{𝑆 𝐷 + 𝑇 𝐷} = \frac{𝑇 𝑆}{𝑇 + 𝑆} . \end{aligned}

Let $𝑆 = 𝑇$ :

AI = \frac{𝑇^{2}}{2 𝑇} = \frac{𝑇}{2} .

If $𝑇 \geq 2 \cdot AI (accelerator)$ , we are FLOPs bound. If we assume v5e, then $AI (v5e) \approx 240$ , hence $𝑇 \geq 480$ .

If by effective relative cost they mean:

Effective cost (𝑇) = \frac{𝑡_{attention} (𝑇)}{𝑡_{FFW} (𝑇)},

where $𝑇$ is context length, not time.

Assume $𝐵 > 240$ so that in the FFW block we become compute bound, hence:

\begin{matrix} 𝑡_{FFW} = \frac{6 𝐵 𝑇 𝐷 𝐹}{FLOPs/s} \\ 𝑡_{attention} = max (\frac{8 𝐵 𝑇 𝐷}{BW}, \frac{4 𝐵 𝑇^{2} 𝐷}{FLOPs/s}) \end{matrix}

Therefore:

Effective cost (𝑇) = \frac{max (\frac{8 𝐵 𝑇 𝐷}{BW}, \frac{4 𝐵 𝑇^{2} 𝐷}{FLOPs/s})}{\frac{6 𝐵 𝑇 𝐷 𝐹}{FLOPs/s}} = max (\frac{8 𝐵 𝑇 𝐷}{BW}, \frac{4 𝐵 𝑇^{2} 𝐷}{FLOPs/s}) \cdot \frac{FLOPs/s}{6 𝐵 𝑇 𝐷 𝐹} .

If $𝑇 \geq 480$ , then:

Effective cost (𝑇) = \frac{4 𝐵 𝑇^{2} 𝐷}{6 𝐵 𝑇 𝐷 𝐹} = \frac{4 𝑇}{6 𝐹} = \frac{4 𝑇}{24 𝐷} = \frac{𝑇}{6 𝐷} .

If $𝑇 < 480$ , then:

Effective cost (𝑇) = \frac{8 𝐵 𝑇 𝐷}{BW} \cdot \frac{FLOPs/s}{6 𝐵 𝑇 𝐷 𝐹} = \frac{8 𝐵 𝑇 𝐷}{6 𝐵 𝑇 𝐷 𝐹} \cdot \frac{1}{240} = \frac{8}{24 𝐷} \cdot \frac{1}{240} = \frac{8}{24 \cdot 240 \cdot 𝐷} = \frac{8}{5760 \cdot 𝐷} \approx \frac{13}{10^{- 4}} \cdot \frac{1}{𝐷} .

Assume $𝐷 \approx 8 𝑘$ :

Exercise 5 – Attention FLOPs vs. QKVO projection FLOPs

Exercise statement

At what sequence length are self-attention FLOPs equal to the QKVO projection FLOPs?

Solution

Scan pages: 9

Assuming FLOPs without training, $𝐷 = 𝑁 𝐻$ , $𝑆 = 𝑇$ , and $𝑁 = 𝐾$ . Then we have to solve for:

4 𝐵 𝑇^{2} 𝐷 = 8 𝐵 𝑇 𝐷 𝑁 𝐻 \Rightarrow 4 𝐵 𝑇^{2} 𝐷 = 8 𝐵 𝑇 𝐷^{2} \Rightarrow 𝑇 = \frac{8}{4} 𝐷 \Rightarrow 𝑇 = 2 𝐷 .

So at $𝑇 = 2 𝐷$ sequence length, self-attention FLOPs equal the QKVO projection FLOPs.

Exercise 6 – Rematerialization FLOPs

Exercise statement

Say we only save the output of each of the 7 main matmuls in a Transformer layer during our forward pass, namely Q, K, V, O + the three FFW matrices. How many extra FLOPs do we need to “rematerialize” during the backwards pass?

Solution

Scan pages: 9–12

Define a general layer as:

𝑂 = 𝑓 (𝐼, 𝑊) .

In a computational graph, during the backward pass, it will receive $\frac{𝑑 𝐿}{𝑑 𝑂}$ , and it will compute $\frac{𝑑 𝐿}{𝑑 𝐼}$ , $\frac{𝑑 𝐿}{𝑑 𝑊}$ and send $\frac{𝑑 𝐿}{𝑑 𝐼}$ down the graph. Note that in general:

\frac{𝑑 𝐿}{𝑑 𝑊} = \frac{𝑑 𝐿}{𝑑 𝑂} \cdot \frac{𝑑 𝑂}{𝑑 𝑊} (𝑊, 𝐼)

so it depends on $𝐼$ as well. With this in mind, we can see that if we draw the computational graph we will need to rematerialize:

$𝐼 = gelu (𝑋 𝑊_{in1}) ⊙ (𝑋 𝑊_{in2})$ for $\frac{𝑑 𝐿}{𝑑 𝑊_{out}}$ .
$𝐴 = softmax (𝑄 𝐾^{𝑇}) 𝑉$ for $\frac{𝑑 𝐿}{𝑑 𝑊_{𝑜}}$ .
$softmax (𝑄 𝐾^{𝑇})$ for $\frac{𝑑 𝐿}{𝑑 𝑉}$ .

where we have ignored the layernorm operations. Therefore we need:

FLOPs = 4 𝐵 𝑇 𝑆 𝑁 𝐻 + 𝑂 (𝐵 𝑇 𝐹)

to rematerialize during the backward pass.

Exercise 7 – DeepSeek V3 utilization

Exercise statement

DeepSeek v3 says it was trained for 2.79M H800 hours on 14.8T tokens. Given that it has 37B activated parameters, roughly what hardware utilization did they achieve? Hint: note that they used FP8 FLOPs without structured sparsity.

Solution

Scan pages: 12

FLOPs for V3 = $6 \cdot 37 𝑒 9 \cdot 14.8 𝑒 12$ .

$H800 SXM FP8 tensor core FLOPs/s = 1474 TFLOPs/s$ .

Available FLOPs in $2.79 M$ H800 hours, if using FP8 = $2.79 𝑒 6 \cdot 3600 \cdot 1474 𝑒 12$ .

Hardware Utilization:

\frac{6 \cdot 37 \cdot 14.8 \cdot 10^{21}}{2.79 \cdot 3600 \cdot 1474 \cdot 10^{18}} = \frac{6 \cdot 37 \cdot 14.8}{2.79 \cdot 3600 \cdot 1474} \cdot 10^{3} \approx 0.00016 \cdot 10^{3} = 0.16 .

Exercise 8 – MoE compute-bound batch size

Exercise statement

Mixture of Experts (MoE) models have $𝐸$ copies of a standard dense MLP block, and each token activates $𝑘$ of these experts. What batch size in tokens is required to be compute-bound for an MoE with weights in int8 on TPU v5e? For DeepSeek, which has 256 routed experts and $𝑘 = 8$ , what is this number?

Solution

Scan pages: 13–15

Assume training FLOPs. $𝐸$ copies of standard MLP blocks, and each token activates $𝑘$ of these experts. Consider the following model of MoE:

ℎ_{𝑡} = \sum_{𝑖 = 1}^{𝐸} 𝑔_{𝑖, 𝑡} {FFN}_{𝑖} (𝑥_{𝑡}),

where ${FFN}_{𝑖} (𝑥_{𝑡}) = 𝑊_{𝑖} 𝑥_{𝑡}$ and $𝑊_{𝑖} \in ℝ^{𝑑 \times 𝑑}$ .

To compute the MoE layer, we will group the tokens belonging to an expert and then matmul. So in total we will do $𝐸$ matmuls. Assume that every token randomly chooses $𝑘$ numbers from ${1, \dots, 𝐸}$ .

Define:

𝑆_{𝑒} = \sum_{𝑖 = 1}^{𝑇} 𝑋_{𝑖, 𝑒},

where $𝑋_{𝑖, 𝑒} = 1$ if token $𝑖$ has sampled $𝑒 \in {1, \dots, 𝐸}$ .

The expected value of $𝑆_{𝑒}$ gives us the expected value of the batch dimension of the matmul of expert $𝑒$ :

𝐸 [𝑆_{𝑒}] = \sum_{𝑖 = 1}^{𝑇} 𝑃 (𝑋_{𝑖, 𝑒}) = \sum_{𝑖 = 1}^{𝑇} \frac{𝑘}{𝐸} = \frac{𝑇 \cdot 𝑘}{𝐸} .

So then we do:

For e in {1,...,E}
  1: [T*k/E, D] @ int8[D, D]

Therefore:

\begin{matrix} FLOPs = 𝐸 \cdot (2 \cdot \frac{𝑇 \cdot 𝑘}{𝐸} \cdot 𝐷^{2}) = 2 𝑇 𝑘 𝐷^{2} \\ Bytes = 𝐸 \cdot (\frac{2 𝑇 𝑘}{𝐸} \cdot 𝐷 + 𝐷^{2}) = 2 𝑇 𝑘 𝐷 + 𝐸 𝐷^{2} . \end{matrix}

AI = \frac{2 \cdot 𝑇 \cdot 𝑘 \cdot 𝐷^{2}}{2 \cdot 𝑇 \cdot 𝑘 \cdot 𝐷 + 𝐸 \cdot 𝐷^{2}} \approx \frac{2 \cdot 𝑇 \cdot 𝑘 \cdot 𝐷^{2}}{𝐸 \cdot 𝐷^{2}} = \frac{2 𝑇 𝑘}{𝐸} .

Therefore for being compute bound:

\begin{matrix} \frac{2 𝑇 𝑘}{𝐸} \geq 240 \\ \Rightarrow 𝑇 \geq \frac{𝐸}{𝑘} \cdot \frac{240}{2} = \frac{𝐸}{𝑘} \cdot 120 . \end{matrix}

So for DeepSeek V3 this number is $3840$ . (Note: While transcribing this execirse from the handwritten notes, I realized that I made an error and assumed that the operations were done in int8, but they are done in FP16, so I’ve corrected the of by 2 error in this transccribed notes, in the handwritten notes is not corrected).