DEVOURED
Promising model architectures routinely fail because no one makes them run efficiently at scale, so Google DeepMind researchers published a free comprehensive book on LLM hardware optimization.
What: Google DeepMind researchers (Austin, Douglas, Frostig, Levskaya, and others) released a free online book covering TPU/GPU architecture, roofline analysis, Transformer mathematics (parameters, FLOPs, memory), parallelization techniques (data, tensor, pipeline, expert), and practical examples training and serving LLaMA 3 in JAX across thousands of chips.
Why it matters: Signals that LLM optimization knowledge has become essential for competitive ML research - scaling laws push models to hardware frontiers where a 20% benchmark improvement is irrelevant if it costs 20% roofline efficiency, and promising architectures routinely fail because no one makes them run efficiently.
Takeaway: Free online at jax-ml.github.io/scaling-book - use it to estimate training costs, choose parallelism strategies for your model size and chip count, or understand why your LLM is running slower than expected.
Deep dive
- Google DeepMind researchers (Austin, Douglas, Frostig, Levskaya, Chen, Vikram, Lebron, Choy, Ramasesh, Webson, Pope) published a comprehensive free online book demystifying how to scale language models on TPU and GPU hardware
- Covers three main parts: hardware fundamentals (roofline analysis, TPU/GPU architecture, inter-chip communication), Transformer specifics (parameter counting, FLOP calculations, parallelization for training and inference), and practical tutorials (JAX programming, profiling, debugging)
- Explains how to choose among 4 primary parallelism techniques (data, tensor, pipeline, expert) plus memory optimization techniques (rematerialization, optimizer sharding/ZeRO, host offload, gradient accumulation) based on model size and available chips
- Includes worked examples for training and serving LLaMA 3 on TPUs, with cost and latency estimates
- Motivation: ML research has shifted from being purely algorithmic to requiring deep systems knowledge because models now run so close to hardware limits that efficiency directly impacts research viability
- Key insight from authors: "A 20% win on benchmarks is irrelevant if it comes at a 20% cost to roofline efficiency" - promising architectures routinely fail because researchers don't make them run efficiently at scale
- Book teaches how to estimate: training cost and time for large Transformers, memory needed to serve models, whether code is compute-bound or memory-bound, best parallelism scheme for given hardware topology
- Covers both TPUs (Google's ML-optimized chips with systolic arrays) and GPUs (with Tensor Cores), including detailed networking and bandwidth specifications for clusters
- Explains Transformer mathematics in detail: exact matrix sizes, parameter counts, FLOPs for forward/backward passes, KV cache sizes, when attention becomes important vs feed-forward layers
- Includes sections on disaggregated serving architecture for inference, profiling TPU code with JAX/TensorBoard, and understanding the JAX + XLA compilation stack
- Emphasizes "strong scaling" as the goal: adding more chips should give proportional throughput increase, but communication overhead and memory bottlenecks can break this
- Discusses hardware-software co-design challenge: TPU designers had to bet on what algorithms would look like 2-3 years in advance, with TPUs optimized for matrix multiplication (N FLOPs per byte) proving highly successful
- Free online at jax-ml.github.io/scaling-book with 12 sections, worked problems, and practical examples
Decoder
- Roofline analysis: Framework for identifying whether code is limited by compute, memory bandwidth, or communication
- Strong scaling: Linear speedup from adding chips - double the chips, double the throughput
- Systolic array: TPU's specialized matrix multiplication architecture that pumps data rhythmically through a processor grid
- FLOPs: Floating point operations (total adds/multiplies), not FLOPS (operations per second)
Original article
This book discusses the science of scaling language models. It covers how TPUs and GPUs work, how they communicate with each other, how LLMs run on real hardware, and how to parallelize models during training and inference so they run efficiently at massive scale. The book answers questions about how expensive training a model should be, how much memory is needed to serve models, and more.
Explore all 66 articles in today's digest
Fresh dev news, AI-digested daily. Never miss a trend at devoured.fyi