How to Scale Models

Scaling LLMs requires looking at compute, memory, and communication bottlenecks together. Understanding the limits of TPU and GPU lets you predict where a model will bottleneck and estimate training and inference costs more accurately.

The key is maintaining a state where throughput scales near-linearly as the number of chips increases. When communication becomes slower than compute, you fall into a communication bound, and even within a single chip, memory bandwidth and total capacity become bottlenecks. In particular, operations with high FLOPs/byte like matrix multiplication align well with the TPU's systolic array, and the GPU's Tensor Cores are evolving in the same direction.

• Use Roofline analysis to see whether compute, communication, or memory is creating the limit.
• Explain TPU architecture and inter-chip links, and cover sharding matrix multiplication and data reshuffling.
• Use Transformer math to calculate parameter counts, FLOPs, and KV cache size.
• For Training and inference, compare FSDP, Megatron sharding, pipeline parallelism, expert parallelism, and disaggregated serving.
• Reduce memory requirements with Rematerialisation, ZeRO, host offload, and gradient accumulation.

The latter half continues with a LLaMA 3 application example, a JAX implementation, profiling and debugging, and a new GPU chapter. Ultimately, the goal is to determine which combination of parallelization is most efficient for given hardware, and to be able to estimate approximate training time and inference cost.