TPU 的线性扩展假设忽略硬件异质性与内存壁垒，并行技术需与编译器优化深度耦合，否则难以突破实际性能瓶颈。
《如何扩展模型：系统视角解析大型语言模型（LLMs）于 TPUs 架构中的应用》探讨了如何有效地将深度学习模型的运算规模从单一硬件扩展至数万台设备，同时保持运算性能的线性增长（强扩展）。该书目的是说明 TPU 与 GPU 硬件的特性，并解析 Transformer 架构如何在现有硬件上进行优化操作，对于设计新模型架构的研究人员及优化现有 LLM 性能的工程师均具实用性。书中以计算、内存、与通信限制为切入点，介绍效率瓶颈及相应解决方案。
内容分为三部分：第一部分说明基础概念，包括屋顶线分析（roofline analysis）的使用，TPU 的运作方式，以及碎型矩阵分片（sharded matrix）的计算方法；第二部分聚焦于 Transformer，深入探讨其数学运算细节（参数数量与运算需求），以及如何通过多种并行技术（如数据平行、张量平行、管道平行、专家平行）优化模型的培训与推论性能；第三部分则提供实务教学，展示如何以 JAX 编程 TPUs，并如何使用工具（如 TensorBoard profiler）侦测与修正实际问题。
这本书最终期望读者能具备在特定硬件运算平台上，选择并行技术与设置模型的能力，从而提升大型 Transformer 模型在当代硬件上的训练与推论效益。对 LLaMA-3 等热门开源模型的案例解析亦是亮点，提供具体实践指引，涵盖成本与性能考量。同时书中鼓励读者参与讨论并提供反馈，内容不断更新优化。
https://news.ycombinator.com/item?id=42936910