TPU Deep Dive
Course
Review
Slides
Exams
Sources
Theme: system
Sources
SB5
How to Parallelize a Transformer for Training
Data parallelism
SB5
How to Parallelize a Transformer for Training
FSDP
SB5
How to Parallelize a Transformer for Training
Tensor parallelism
SB5
How to Parallelize a Transformer for Training
FSDP + TP
SB5
How to Parallelize a Transformer for Training
Pipelining
SB5
How to Parallelize a Transformer for Training
Across pods (DCN)
SB5
How to Parallelize a Transformer for Training
Problems
SB6
Training LLaMA 3 on TPUs
Sharding LLaMA 3-70B
SB6
Training LLaMA 3 on TPUs
Worked problems
SB12
How to Think About GPUs
GPU rooflines
P-Llama3
Llama Team, Meta, "The Llama 3 Herd of Models"
Tables 3–4, §3.3–3.4
R-maxtext
MaxText repo
base.yml mesh, sharding guide
P-GSPMD
Xu et al., "GSPMD: General and Scalable Parallelization for ML Computation Graphs"
§3.3 pipelining as sharding
D-7x
TPU7x (Ironwood)
G-8dd
Gupta & Mugazambi, "TPU 8t and TPU 8i technical deep dive"
Chapter 6. Day 3, second block, 55 minutes
Choosing a mesh for training
Read
Quiz
12
Flashcards
20
Teach-back
5
Question 1 of 12
Derivation: enter a number
Ironwood has 2,307 TFLOP/s BF16 and
W
ici
W_\text{ici}
W
ici
= 200 GB/s per torus axis. What is its ICI intensity
α
=
C
/
W
ici
\alpha = C/W_\text{ici}
α
=
C
/
W
ici
?
Check answer
Enter