flâneur — a map of the web's best reading

How to Think About GPUs | How To Scale Your Model

jax-ml.github.io · 13,968 words · saved by 4 readers

We love TPUs at Google, but GPUs are great too. This chapter takes a deep dive into the world of NVIDIA GPUs – how each chip works, how they’re networked together, and what that means for LLMs, especially compared to TPUs. This section builds on Chapter 2 and Chapter 5, so you are encouraged to read them first. Jacob Austin† †Google DeepMind Swapnil Patil† Adam Paszke† Reiner Pope* *MatX Aug. 18, 2025 A modern ML GPU (e.g. H100, B200) is basically a bunch of compute cores that specialize in matrix multiplication (called Streaming Multiprocessors or SMs) connected to a stick of fast memory (called HBM). Here’s a diagram: Each SM, like a TPU’s Tensor Core, has a dedicated matrix multiplication core (unfortunately also called a Tensor Core), a vector arithmetic unit (called a Warp Scheduler), and a fast on-chip cache (called SMEM). Unlike a TPU, which has at most 2 independent “Tensor Cores”, a modern GPU has more than 100 SMs (132 on an H100). Each of these SMs is much less powerful than

How to Think About GPUs | How To Scale Your Model How to Think About GPUs Part 12 of How To Scale Your Model ( Part 11: Conclusion | The End ) We love TPUs at Google, but GPUs are great too. This chapter takes a deep dive into the world of GPUs – how each chip works, how they're networked together, and what that means for LLMs, especially compared to TPUs. While there are a multitude of GPU architectures from NVIDIA, AMD, Intel, and others, here we will focus on NVIDIA GPUs. This section builds on Chapter 2 and Chapter 5 , so you are encouraged to read them first. Authors Affiliation Jacob Aus

Explore this link on the map →

saved by

related reading