由于深度學習的「黑箱」本性，從業(yè)者自我調侃道：

如果說深度神經網絡是現(xiàn)代版的「煉金術」，我們在古代就是「煉金術士」。

2018年5月3日， Science發(fā)表新聞，標題直指「煉金術」，強調加強AI的科學基礎

但這次的谷歌的團隊，卻有不一樣的看法：

在許多方面，深度學習仍然帶有一定的「煉金術」色彩，但理解和優(yōu)化模型性能并不一定如此——即使是在大規(guī)模環(huán)境下！

近日，谷歌DeepMind科學家Jacob Austint在X上， 發(fā)布了基于JAX和TPU的大模型Scaling教科書《How to Sacle Your Model》。

Jeff Dean轉發(fā)原帖，并打起了廣告：

谷歌最強的Gemini模型的訓練，重度依賴JAX軟件棧+TPU硬件平臺。

如果你想了解更多詳情，來看看這本超棒的書：「How to Sacle Your Model」。

進入教科書網站，可以看到大寫的標題：「如何擴大模型規(guī)模（How to  Sacle Your Model）」。

正如小標題所言，這本書關于在張量處理單元（TPU）上大語言模型的的系統(tǒng)觀點。

這是一本關于LLM底層硬核技術的教科書，簡介如下：

訓練大語言模型（LLMs）常常讓人感覺就像煉金術，但理解和優(yōu)化模型的性能其實并不復雜。

本書的目標是揭開在TPU上擴展語言模型的科學謎團：TPU是如何工作的，它們如何相互通信，LLM在實際硬件上是如何運行的，以及在訓練和推理過程中如何對模型進行并行化，以便在大規(guī)模運行時實現(xiàn)高效性。

如果你想知道「訓練這個LLM需要多貴的成本」、「要自己部署這個模型需要多少內存」或者「什么是AllGather」這些問題的答案，希望本書能對你有所幫助。

教科書鏈接：https://jax-ml.github.io/scaling-book/

模型Scaling，無需恐懼

三四年前，大多數(shù)機器學習研究人員，可能并不需要了解模型擴展（model scaling）。

但如今，即便是「較小」的模型，也已經逼近硬件極限，因此研究要有真正的創(chuàng)新性，就必須考慮如何在大規(guī)模環(huán)境下提高效率。

作者詳細解釋了為什么要模型擴展及其目標：

如果某種方法能在基準測試中提升20%的性能，但同時使Roofline效率下降20%，那么這樣的優(yōu)化是沒有意義的。

許多有前景的模型架構最終失敗，并不是因為它們在理論上不可行，而是因為它們無法高效擴展，或者沒有人投入足夠的精力去優(yōu)化它們的計算效率。

模型擴展的目標是在增加用于訓練或推理的芯片數(shù)量時，實現(xiàn)吞吐量的線性增長，這被稱為 「強擴展」（Strong Scaling）。

通常，增加芯片數(shù)量（即「并行計算」）可以減少計算時間，但同時也會帶來額外的芯片間通信開銷。如果通信時間超過了計算時間，就會遇到 「通信瓶頸」，導致無法實現(xiàn)理想的擴展性能。如果對硬件足夠了解，能夠預測這些瓶頸的出現(xiàn)位置，就可以通過調整模型設計或重新配置系統(tǒng)來規(guī)避它們。

看不懂這些，也沒關系，畢竟這是谷歌最強模型Gemini同款的技術棧！

但作者誠意十足，表示：如果認真看完后，有晦澀的地方，請及時反饋，保證一定改。

作者保證：必有所得

從處理單個加速器到處理數(shù)萬個加速器，相對簡單的原則無處不在，了解這些原則可以讓你做很多有用的事情：

• 粗略評估模型的各個部分與理論最優(yōu)性能的接近程度。
• 在不同規(guī)模下，合理選擇并行計算方案（如何在多個設備間分配計算任務）。
• 估算訓練和運行大型Transformer模型所需的成本和時間。
• 設計能夠充分利用特定硬件特性的算法。
• 設計硬件時，基于對當前算法性能瓶頸的明確理解來進行優(yōu)化。

此書的目標是解釋TPU（以及 GPU）的工作原理，以及為了當前硬件上實現(xiàn)高效計算，Transformer架構如何不斷演化。

希望這些內容既能幫助研究人員設計新的模型架構，也能為工程師提供指導，以優(yōu)化當前一代的大語言模型（LLM）計算性能。

作者保證，讀完此書一定有所收獲：

在閱讀完本書后，應該能自信地為特定硬件平臺上的Transformer模型選擇最佳并行方案，并大致估算訓練和推理的耗時。

如果你仍然感到困惑，請告訴我們！我們希望知道如何讓這些內容更加清晰易懂。

基礎知識

要閱讀此書，作者提醒讀者：

對LLM（大語言模型）和Transformer架構有基本的了解，但不一定熟悉它們在大規(guī)模計算中的運作方式。

應該了解LLM訓練的基礎知識，并且最好對JAX有一定的了解。

下面的背景資料，有助于了解所需的基礎知識：

博客鏈接：https://jalammar.github.io/illustrated-transformer/

JAX講義:https://github.com/rwitten/HighPerfLLMs2024

整體結構

在本書中，將解答以下問題：

• 矩陣乘法的計算時間如何估算？在多大規(guī)模下，它的計算受限于計算能力、內存帶寬還是通信帶寬？
• TPU是如何連接在一起組成訓練集群的？系統(tǒng)的各個部分分別具備多少帶寬？
• 在多個TPU之間進行數(shù)據(jù)收集（gather）、分發(fā)（scatter）或重新分布（re-distribute）需要多少時間？
• 如何高效地計算跨設備分布的矩陣乘法？

這些內容能幫助讀者，深入理解LLM在現(xiàn)代硬件上的運行機制，并學會如何優(yōu)化訓練和推理的效率。

《第1章》介紹屋頂線分析（Roofline Analysis），并探討限制模型擴展的關鍵因素，包括通信、計算和內存。

《第2章》和《第3章》詳細講解TPU和現(xiàn)代GPU的工作原理，既包括作為獨立芯片的運行機制，也涵蓋了更關鍵的內容——它們如何通過芯片間互連（inter-chip links）形成一個計算集群，并受到帶寬和延遲的限制。

五年前，機器學習領域的架構還十分多樣化包——括卷積神經網絡、長短時記憶網絡、多層感知機和Transformer等。如今，Transformer架構一家獨大。

Transformer結構的每一個細節(jié)，都非常值得深入理解，包括：矩陣的具體尺寸、歸一化（Normalization）發(fā)生的位置、各部分包含多少參數(shù)和FLOPs（浮點運算次數(shù)）。

《第4章》將詳細解析Transformer的數(shù)學計算，幫助你掌握如何計算訓練和推理過程中的參數(shù)量和FLOPs。

這些計算將揭示：

• 模型的內存占用有多大？
• 計算和通信的時間消耗分布如何？
• 注意力機制（Attention）和前饋網絡（Feed-Forward Blocks）何時成為計算的瓶頸？

通過這些分析，將能夠更精確地優(yōu)化Transformer訓練和推理的效率，并更深入地理解其計算特性。

圖示2：標準Transformer層，每個矩陣乘法（matmul）以圓圈中的點表示。所有參數(shù)（不包括歸一化層）以紫色顯示。

《第5章：訓練》和《第7章：推理》是本書的核心內容，在這兩章中將討論一個根本問題：

給定一個大小和一定數(shù)量芯片的模型，如何將模型并行化，以保持在「強擴展」（strong scaling）范疇內？

這個看似簡單的問題，其實有著令人意外的復雜答案。

從高層次來看，主要有四種并行化技術用于將模型分布到多個芯片上：數(shù)據(jù)并行（Data Parallelism）、張量并行（Tensor Parallelism）、流水線并行（Pipeline Parallelism）以及專家并行（Expert Parallelism）。

圖3：純數(shù)據(jù)并行（前向傳播）示意圖。激活（Activations）（左側）完全按照批次維度（batch dimension） 進行分片。

這種方法通過將批次分配到多個 TPU 上，實現(xiàn)了數(shù)據(jù)并行，從而在沒有額外通信負擔的情況下，加速模型的計算。

此外，還有多種技術可以減少內存需求，比如重新計算（Rematerialization）、優(yōu)化器/模型分片（Optimizer/Model Sharding，也稱為ZeRO）、主機卸載（Host Offload）、梯度累積（Gradient Accumulation）。

在這兩章中將討論這些技術，并幫助理解如何在新的架構或設置中選擇最適合的并行化策略。

《第6章》和《第8章》是實際操作教程，應用這些概念于LLaMA-3，更直觀地理解如何在實際應用中進行操作。

最后，《第9章》和《第10章》將討論如何在JAX中實現(xiàn)這些想法，并介紹當代碼出現(xiàn)問題時如何進行性能分析和調試。

在《第11章》中，會給出進一步閱讀清單和更深入的參考文獻。

在整個過程中，會給出一些需要自己動手解決的問題。

作者溫馨提示：

請不要覺得有壓力要按順序閱讀所有章節(jié)，也不一定要全部閱讀完。

我們鼓勵你留下反饋意見。

目前這是草稿版本，未來會繼續(xù)修訂和改進。

當前的目錄，翻譯如下：