Chapter 5: How to Parallelize a Transformer for Training

Scaling Book Exercises – Chapter 5

Chapter: 5 – How to Parallelize a Transformer for Training
Source: handwritten Onyx Boox A4 PDF
Note: Transcribed from handwritten solutions; book markdown used only for exercise statements and notation.

Shared setup from the book

Let’s use LLaMA-2 13B as a basic model for this section. Here are the model details:

hyperparam	value
$𝐿$	$40$
$𝐷$	$5, 120$
$𝐹$	$13824$
$𝑁$	$40$
$𝐾$	$40$
$𝐻$	$128$
$𝑉$	$32, 000$

LLaMA-2 has separate embedding and output matrices and a gated MLP block.

Exercise 1 – LLaMA-2 13B parameter count

Exercise statement

How many parameters does LLaMA-2 13B have (I know that’s silly but do the math)? Note that, as in Transformer Math, LLaMA-3 has 3 big FFW matrices, two up-projection and one down-projection. We ignored the two “gating” einsum matrices in this section, but they behave the same as $𝑊_{in}$ in this section.

Solution

Scan pages: 1

From LLAMA2 we know that $𝑁 = 𝐾$ . Therefore:

\begin{aligned} MLP Params & = 3 𝐷 𝐹 per layer \\ Attention Params & = 4 𝐷 𝑁 𝐻 per layer \\ Embeddings & = 2 𝐷 𝑉 overall . \end{aligned}

In total we have $(3 𝐷 𝐹 + 4 𝐷 𝑁 𝐻) \cdot 𝐿 + 2 𝐷 𝑉$ parameters. Let $𝐿 = 40$ , $𝑁 𝐻 = 𝐷$ , $𝐷 = 5120$ , $𝐹 = 13824$ , and $𝑉 = 32, 000$ . Hence:

(3 \cdot 5120 \cdot 13824 + 4 {(5120)}^{2}) \cdot 40 + 2 \cdot 5120 \cdot 32 𝑒 3 = 13, 015, 449, 600 \approx 13 \cdot 10^{9} = 13 Billion .

Exercise 2 – Memory for BS=16M training

Exercise statement

Let’s assume we’re training with BS=16M tokens and using Adam. Ignoring parallelism for a moment, how much total memory is used by the model’s parameters, optimizer state, and activations? Assume we store the parameters in bf16 and the optimizer state in fp32 and checkpoint activations three times per layer (after the three big matmuls).

Solution

Scan pages: 1-2

Model Params $= 2 \cdot 13 𝑒 9 = 26 𝑒 9 Bytes = 26 GB .$

Optimizer State:

4 \cdot 13 𝑒 9 + 4 \cdot 13 𝑒 9 = 8 \cdot 13 𝑒 9 = 104 GB .

Activation Checkpointing:

\begin{aligned} 𝐿 \cdot (2 𝐵 𝑇 𝐹 + 2 𝐵 𝑇 𝐹 + 2 𝐵 𝑇 𝐷) & = 40 \cdot (4 𝐵 𝑇 𝐹 + 2 𝐵 𝑇 𝐷) \\ = 40 \cdot (𝐵 𝑇) \cdot (4 𝐹 + 2 𝐷) \\ = 40 \cdot (16 𝑒 6) \cdot (4 \cdot 13824 + 2 \cdot 5120) \\ \approx 42 TB . \end{aligned}

Exercise 3 – 32k training on TPU v5p 16x16x16

Exercise statement

Assume we want to train with 32k sequence length and a total batch size of 3M tokens on a TPUv5p 16x16x16 slice. Assume we want to use bfloat16 weights and a float32 optimizer, as above.

Can we use pure data parallelism? Why or why not?
Can we use pure FSDP? Why or why not? With pure FSDP, how much memory will be used per device (assume we do gradient checkpointing only after the 3 big FFW matrices).
Can we use mixed FSDP + tensor parallelism? Why or why not? If so, what should $𝑋$ and $𝑌$ be? How much memory will be stored per device? Using only roofline FLOPs estimates and ignoring attention, how long will each training step take at 40% MFU?

Solution

Scan pages: 2-11

Let:

\begin{aligned} HBM v5p & = 96 GB \\ bf16 FLOPs/s v5p & \approx 4.59 𝑒 14 \\ 𝑊_{ICI} v5p & \approx 1.8 𝑒 11 . \end{aligned}

3.1

We cannot use data parallelism since a lower bound in memory per chip we need is $130 GB$ . (In fact we need more if we consider the activation checkpointing.)

3.2

Assume FSDP shards the whole architecture as follows. For the attention params:

𝑊_{𝑞} [𝐷_{𝑋}, 𝑁, 𝐻] 𝑊_{𝑘} [𝐷_{𝑋}, 𝑁, 𝐻] 𝑊_{𝑣} [𝐷_{𝑋}, 𝑁, 𝐻] 𝑊_{𝑜} [𝑁, 𝐻, 𝐷_{𝑋}]

and for the MLP params as usual:

𝑊_{in1} [𝐷_{𝑋}, 𝐹] 𝑊_{in2} [𝐷_{𝑋}, 𝐹] 𝑊_{out} [𝐹, 𝐷_{𝑋}]

and for input output embeddings as:

𝑊_{embed} [𝑉, 𝐷_{𝑋}] 𝑊_{outembed} [𝐷_{𝑋}, 𝑉]

Therefore the total params per device is:

\begin{aligned} (3 \frac{𝐷}{| 𝑋 |} 𝐹 + 4 \frac{𝐷}{| 𝑋 |} 𝐻 𝑁) \cdot 𝐿 + 2 \frac{𝐷}{| 𝑋 |} 𝑉 & = \frac{1}{| 𝑋 |} \cdot ((3 𝐷 𝐹 + 4 𝐷 𝐻 𝑁) \cdot 𝐿 + 2 𝐷 𝑉) \\ \approx \frac{Model Params}{| 𝑋 |} . \end{aligned}

For LLAMA2 13B therefore:

Device Params = \frac{13 \cdot 10^{9}}{16^{3}} \approx 3.2 \cdot 10^{6} = 3.2 Million .

Assuming optimizer states also sharded, we get that each device needs at least:

10 \cdot (3.2 \cdot 10^{6}) = 32 \cdot 10^{6} = 32 MB .

If we also take into account the activation checkpointings which are of shapes:

[𝐵_{𝑋}, 𝑇, 𝐹] [𝐵_{𝑋}, 𝑇, 𝐹] [𝐵_{𝑋}, 𝑇, 𝐷]

we get:

\begin{aligned} Checkpoints per device & = 𝐿 \cdot (2 \frac{𝐵}{| 𝑋 |} \cdot 𝑇 𝐷 + 4 \frac{𝐵}{| 𝑋 |} 𝑇 𝐹) Bytes \\ = \frac{1}{| 𝑋 |} \cdot (𝐿 \cdot (𝐵 𝑇) \cdot (4 𝐹 + 2 𝐷)) Bytes . \end{aligned}

Let $𝐵 𝑇 = 3 \cdot 10^{6}$ , then:

\frac{1}{16^{3}} \cdot (40 \cdot (3 \cdot 10^{6}) \cdot (4 \cdot 13824 + 2 \cdot 5120)) \approx 1.9 GB .

So memory used per device is $32 MB + 1.96 GB$ .

We could use FSDP since with FSDP sharding everything fits in memory, but we would be communication bound since we are compute bound iff:

\frac{𝐵}{| 𝑋 |} > \frac{2550}{𝑀_{𝑋}} \Leftrightarrow \frac{𝑀_{𝑋} \cdot 𝐵}{2550} > | 𝑋 | .

Let $𝐵 = 3 \cdot 10^{6}$ and $𝑀_{𝑋} = 3$ , therefore:

\frac{9 \cdot 10^{6}}{2550} \approx 3530 > | 𝑋 |,

and in our slice $| 𝑋 | = 4046$ .

3.3

Assume we shard the attention matrices as:

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

and for the MLP params as:

𝑊_{in1} [𝐷_{𝑋}, 𝐹_{𝑌}] 𝑊_{in2} [𝐷_{𝑋}, 𝐹_{𝑌}] 𝑊_{out} [𝐹_{𝑌}, 𝐷_{𝑋}]

and for input and output embeddings as:

𝑉_{embed} [𝑉_{𝑌}, 𝐷_{𝑋}] 𝑉_{outembed} [𝐷_{𝑋}, 𝑉_{𝑌}]

Therefore total params per device is:

\begin{aligned} (3 \frac{𝐷}{| 𝑋 | | 𝑌 |} 𝐹 + 4 \frac{𝐷}{| 𝑋 |} \cdot \frac{𝐻 𝑁}{| 𝑌 |}) \cdot 𝐿 + 2 \frac{𝐷}{| 𝑋 | | 𝑌 |} 𝑉 & = \frac{1}{| 𝑋 | | 𝑌 |} \cdot (𝐿 \cdot (3 𝐷 𝐹 + 4 𝐷 𝐻 𝑁) + 2 𝐷 𝑉) \\ = \frac{Model Params}{| 𝑋 | | 𝑌 |} . \end{aligned}

For LLAMA2 13B we get:

Params per device = \frac{13 \cdot 10^{9}}{{(16)}^{3}} \approx 3.2 Million .

If we take into account activation state then we get:

10 \cdot (3.2 \cdot 10^{6}) = 32 MB .

Activation checkpoints will be sharded with shapes:

[𝐵_{𝑋}, 𝑇, 𝐹_{𝑌}] [𝐵_{𝑋}, 𝑇, 𝐹_{𝑌}] [𝐵_{𝑋}, 𝑇, 𝐷_{𝑌}]

so we get:

\begin{aligned} Checkpoints per device & = 𝐿 \cdot (2 \frac{𝐵}{| 𝑋 |} \cdot 𝑇 \frac{𝐷}{| 𝑌 |} + 4 \frac{𝐵}{| 𝑋 |} 𝑇 \frac{𝐹}{| 𝑌 |}) Bytes \\ = \frac{1}{| 𝑋 | | 𝑌 |} \cdot (𝐿 \cdot (𝐵 𝑇) \cdot (4 𝐹 + 2 𝐷)) Bytes . \end{aligned}

Let $𝐵 𝑇 = 3 \cdot 10^{6}$ , then:

\frac{1}{16^{3}} \cdot (40 \cdot (3 \cdot 10^{6}) \cdot (4 \cdot 13824 + 2 \cdot 5120)) \approx 1.9 GB .

Therefore for memory used per device we get $\approx 1.9 GB$ .

If we use FSDP and TP to minimize communications we need to set

𝑋_{opt} = \sqrt{\frac{𝐵}{𝐹} \cdot \frac{𝑀_{𝑋}}{𝑀_{𝑌}} \cdot 𝑁}

Let $𝑁 = 16^{3}$ , $𝑀_{𝑋} = 2$ , $𝑀_{𝑌} = 1$ , $𝐵 = 3 \cdot 10^{6}$ , and $𝐹 = 13824$ . Therefore:

𝑋_{opt} = \sqrt{\frac{3 \cdot 10^{6}}{13824} \cdot 2 \cdot 16^{3}} \approx 1333.3 \approx 1333 .

Since $𝑁 = 𝑋_{opt} \cdot 𝑌 \Rightarrow \frac{16^{3}}{1333} = 𝑌 = 3.07 .$ To get integer values therefore we choose $𝑋_{opt} = 1024$ and $𝑌 = 4$ , which is closer to our optimum. Having chosen this, we will be compute bound iff:

\frac{𝐵}{𝑁} > \frac{𝛼^{2}}{𝑀_{𝑋} 𝑀_{𝑌} 𝐹}

where $𝛼 = \frac{𝐶}{𝑊_{ICI}} = 2550$ for v5p.

Let $𝐹 = 13824$ , $𝑁 = 16^{3}$ , $𝑀_{𝑋} = 2$ , $𝑀_{𝑌} = 1$ , and $𝐵 = 3 \cdot 10^{6}$ :

\begin{aligned} \frac{𝐵}{𝑁} & = \frac{3 \cdot 10^{6}}{16^{3}} \approx 733 \\ \frac{𝛼^{2}}{𝑀_{𝑋} 𝑀_{𝑌} \cdot 𝐹} & = \frac{{(2550)}^{2}}{2 \cdot 13824} = 235 . \end{aligned}

therefore we are compute bound.

Assuming we are compute bound, using only roofline FLOPs estimates and ignoring attention we get that:

Time per Step = \frac{3 𝑒 6 \cdot 13 𝑒 9 \cdot 6}{16^{3} \cdot (0.4 \cdot 4.59 𝑒 14)} = 0.312 s = 312 ms .

Transcription uncertainties

The final subpart is labeled 3.2 again in the scan; I preserved the repeated handwritten label.
On scan page 6, the slice size appears to be written as $|X| = 4046$ ; I marked it as uncertain in the transcription.