TPU Deep Dive
Course
Review
Slides
Exams
Sources
Theme: system
Sources
SB3
Sharded Matrices and How to Multiply Them
Notation
SB3
Sharded Matrices and How to Multiply Them
The four cases
SB3
Sharded Matrices and How to Multiply Them
Collective costs
SB3
Sharded Matrices and How to Multiply Them
AllToAll
SB3
Sharded Matrices and How to Multiply Them
Collective matmul
SB3
Sharded Matrices and How to Multiply Them
Problems
SB10
Programming TPUs in JAX
Collective matmul in shard_map
SB12
How to Think About GPUs
GPU collectives
J-sharding-api
`jax.sharding` API
J-shmap
`shard_map` tutorial
P-CollMM
Wang et al., "Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models"
X-flags
Flags guidance
windowed einsum, async collectives
X-async
Async HLO ops
P-Megatron
Shoeybi et al., "Megatron-LM"
§3
P-SeqPar
Korthikanti et al., "Reducing Activation Recomputation in Large Transformer Models"
PT-asynctp
"Introducing Async Tensor Parallelism in PyTorch"
P-Flux
Chang et al., "FLUX: Fast Software-based Communication Overlap on GPUs"
N-teoverlap
NeMo TP communication overlap
N-nccltests
NCCL-tests PERFORMANCE.md
D-7x
TPU7x (Ironwood)
Chapter 5. Day 3, first block, 55 minutes
Sharded matmuls and collectives on a torus
Read
Quiz
12
Flashcards
21
Teach-back
6
Question 1 of 12
Derivation: enter a number
An fp32[1024, 4096] array is sharded
A
[
I
X
Y
,
J
]
A[I_{XY}, J]
A
[
I
X
Y
,
J
]
on a mesh {X: 8, Y: 2}. How many bytes does each device hold, in MiB?
MiB
Check answer
Enter