Training LLaMA 3 on TPUs | How To Scale Your Model

[jacobaustin123]

Training LLAMA on TPUs!

[Serega6678]

Thank you for this section!
There is a type in the "Question: LLaMA 3-70B was pretrained with a batch size of about 4M tokens..."
In the table in the answer, the Total is incorrect, after the table Total seems correct

Thanks for this wonderful book again!

[jacobaustin123]

Fixed!

[dhaya]

In the answer to the 4th question:

44 days * 8960 / 171

The denominator should be 225,

[dhaya]

Also the same question mentions "gradients 4 times per layer", but the answer assumes 3.

[jacobaustin123]

Fixed. Thank you!

[tengyifei]

What's the meaning of "number of gradient checkpoints" -- is it the factor of the total size of residuals compared to the size of the layer input?

[jacobaustin123]

https://jax-ml.github.io/scaling-book/transformers/#gradient-checkpointing should explain it.

[kerrywang]

When computing the T_math, should we, in practice, also take into consideration of TPU utilization?

eg: in the LLaMA 70B, we assumed 40% FLOPs utilization, should we instead compute $$T_{math} = \frac {B \cdot D \cdot F} {N \cdot C \cdot \text{FLOPS UTILIZATION}}$$

[jacobaustin123]

Generally no, while LLAMA 70B might have an aggregate FLOPs utilization of 40%, the matmuls themselves will have very good FLOPs utilization (upwards of 90% often). The low overall utilization (esp during decoding) comes from things like attention, layer-norms, and imperfect ICI overlapping.

In other words, you should include that discount factor when saying total llama training time = total llama FLOPs / (total llama FLOPs/s * 0.4) but not when talking about a specific matmul.

[henryhmko]

Thanks for making an amazing resource! I just noticed two small typos:

1. Question 4

1752 days to train. That's 6 and a half years.

To "4.8 years"

2. How to shard Llama3-70B intro

with 4M token batch size (1024 sequences of length 8192 per batch)

To "1024 sequences of length 4096 per batch" from the Llama3 paper

[jacobaustin123]

Fixed. Thanks for flagging these.