TPU's linear scalability assumption ignores hardware heterogeneity and memory barriers. Parallel techniques need to be deeply coupled with compiler optimizations, otherwise, it is difficult to break through actual performance bottlenecks.
How to Expand Models: A Systematic Perspective on the Application of Large Language Models (LLMs) in TPU Architectures explores how to effectively scale the computational scale of deep learning models from single hardware to tens of thousands of devices while maintaining linear growth in computational performance (strong scaling). The book aims to explain the characteristics of TPU and GPU hardware, and to analyze how the Transformer architecture can be optimized on existing hardware, providing practical value for researchers designing new model architectures and engineers optimizing the performance of existing LLMs. The book introduces efficiency bottlenecks and corresponding solutions by focusing on computational, memory, and communication constraints.
The content is divided into three parts: Part One explains basic concepts, including the use of roofline analysis, how TPUs operate, and the calculation methods of sharded matrices; Part Two focuses on the Transformer, delving into its mathematical operation details (parameter count and computational requirements), and how to optimize model training and inference performance through various parallel techniques (such as data parallelism, tensor parallelism, pipeline parallelism, and expert parallelism); Part Three provides practical guidance, demonstrating how to program TPUs using JAX and how to use tools (such as TensorBoard profiler) to detect and resolve actual issues.
This book ultimately aims for readers to gain the ability to select parallel techniques and set models on specific hardware platforms, thereby improving the training and inference efficiency of large Transformer models on modern hardware. The case analysis of popular open-source models such as LLaMA-3 is also a highlight, providing specific practical guidelines that cover cost and performance considerations. Meanwhile, the book encourages readers to participate in discussions and provide feedback, with the content continuously updated and optimized.
https://news.ycombinator.com/item?id=42936910