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