谷歌超硬核教科書(shū)來(lái)了！Jeff Dean帶貨揭Gemini訓(xùn)練秘籍：在TPU上scaling

作者：新智元 2025-02-21 13:45:00

谷歌團(tuán)隊(duì)發(fā)布LLM硬核技術(shù)教科書(shū)，從「系統(tǒng)視圖」揭秘LLM Scaling的神秘面紗。Jeff Dean強(qiáng)調(diào)書(shū)中藏著谷歌最強(qiáng)AI模型Gemini訓(xùn)練的更多信息。

由于深度學(xué)習(xí)的「黑箱」本性，從業(yè)者自我調(diào)侃道：

如果說(shuō)深度神經(jīng)網(wǎng)絡(luò)是現(xiàn)代版的「煉金術(shù)」，我們?cè)诠糯褪恰笩捊鹦g(shù)士」。

2018年5月3日， Science發(fā)表新聞，標(biāo)題直指「煉金術(shù)」，強(qiáng)調(diào)加強(qiáng)AI的科學(xué)基礎(chǔ)

但這次的谷歌的團(tuán)隊(duì)，卻有不一樣的看法：

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

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

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

谷歌最強(qiáng)的Gemini模型的訓(xùn)練，重度依賴JAX軟件棧+TPU硬件平臺(tái)。
如果你想了解更多詳情，來(lái)看看這本超棒的書(shū)：「How to Sacle Your Model」。

進(jìn)入教科書(shū)網(wǎng)站，可以看到大寫的標(biāo)題：「如何擴(kuò)大模型規(guī)模（How to Sacle Your Model）」。

正如小標(biāo)題所言，這本書(shū)關(guān)于在張量處理單元（TPU）上大語(yǔ)言模型的的系統(tǒng)觀點(diǎn)。

這是一本關(guān)于LLM底層硬核技術(shù)的教科書(shū)，簡(jiǎn)介如下：

訓(xùn)練大語(yǔ)言模型（LLMs）常常讓人感覺(jué)就像煉金術(shù)，但理解和優(yōu)化模型的性能其實(shí)并不復(fù)雜。
本書(shū)的目標(biāo)是揭開(kāi)在TPU上擴(kuò)展語(yǔ)言模型的科學(xué)謎團(tuán)：TPU是如何工作的，它們?nèi)绾蜗嗷ネㄐ牛琇LM在實(shí)際硬件上是如何運(yùn)行的，以及在訓(xùn)練和推理過(guò)程中如何對(duì)模型進(jìn)行并行化，以便在大規(guī)模運(yùn)行時(shí)實(shí)現(xiàn)高效性。
如果你想知道「訓(xùn)練這個(gè)LLM需要多貴的成本」、「要自己部署這個(gè)模型需要多少內(nèi)存」或者「什么是AllGather」這些問(wèn)題的答案，希望本書(shū)能對(duì)你有所幫助。

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

模型Scaling，無(wú)需恐懼

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

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

作者詳細(xì)解釋了為什么要模型擴(kuò)展及其目標(biāo)：

如果某種方法能在基準(zhǔn)測(cè)試中提升20%的性能，但同時(shí)使Roofline效率下降20%，那么這樣的優(yōu)化是沒(méi)有意義的。
許多有前景的模型架構(gòu)最終失敗，并不是因?yàn)樗鼈冊(cè)诶碚撋喜豢尚校且驗(yàn)樗鼈儫o(wú)法高效擴(kuò)展，或者沒(méi)有人投入足夠的精力去優(yōu)化它們的計(jì)算效率。
模型擴(kuò)展的目標(biāo)是在增加用于訓(xùn)練或推理的芯片數(shù)量時(shí)，實(shí)現(xiàn)吞吐量的線性增長(zhǎng)，這被稱為「強(qiáng)擴(kuò)展」（Strong Scaling）。
通常，增加芯片數(shù)量（即「并行計(jì)算」）可以減少計(jì)算時(shí)間，但同時(shí)也會(huì)帶來(lái)額外的芯片間通信開(kāi)銷。如果通信時(shí)間超過(guò)了計(jì)算時(shí)間，就會(huì)遇到「通信瓶頸」，導(dǎo)致無(wú)法實(shí)現(xiàn)理想的擴(kuò)展性能。如果對(duì)硬件足夠了解，能夠預(yù)測(cè)這些瓶頸的出現(xiàn)位置，就可以通過(guò)調(diào)整模型設(shè)計(jì)或重新配置系統(tǒng)來(lái)規(guī)避它們。

看不懂這些，也沒(méi)關(guān)系，畢竟這是谷歌最強(qiáng)模型Gemini同款的技術(shù)棧！

但作者誠(chéng)意十足，表示：如果認(rèn)真看完后，有晦澀的地方，請(qǐng)及時(shí)反饋，保證一定改。

作者保證：必有所得

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

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

此書(shū)的目標(biāo)是解釋TPU（以及 GPU）的工作原理，以及為了當(dāng)前硬件上實(shí)現(xiàn)高效計(jì)算，Transformer架構(gòu)如何不斷演化。

希望這些內(nèi)容既能幫助研究人員設(shè)計(jì)新的模型架構(gòu)，也能為工程師提供指導(dǎo)，以優(yōu)化當(dāng)前一代的大語(yǔ)言模型（LLM）計(jì)算性能。

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

在閱讀完本書(shū)后，應(yīng)該能自信地為特定硬件平臺(tái)上的Transformer模型選擇最佳并行方案，并大致估算訓(xùn)練和推理的耗時(shí)。
如果你仍然感到困惑，請(qǐng)告訴我們！我們希望知道如何讓這些內(nèi)容更加清晰易懂。

基礎(chǔ)知識(shí)

要閱讀此書(shū)，作者提醒讀者：

對(duì)LLM（大語(yǔ)言模型）和Transformer架構(gòu)有基本的了解，但不一定熟悉它們?cè)诖笠?guī)模計(jì)算中的運(yùn)作方式。
應(yīng)該了解LLM訓(xùn)練的基礎(chǔ)知識(shí)，并且最好對(duì)JAX有一定的了解。

下面的背景資料，有助于了解所需的基礎(chǔ)知識(shí)：

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

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

整體結(jié)構(gòu)

在本書(shū)中，將解答以下問(wèn)題：

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

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

《第1章》介紹屋頂線分析（Roofline Analysis），并探討限制模型擴(kuò)展的關(guān)鍵因素，包括通信、計(jì)算和內(nèi)存。

《第2章》和《第3章》詳細(xì)講解TPU和現(xiàn)代GPU的工作原理，既包括作為獨(dú)立芯片的運(yùn)行機(jī)制，也涵蓋了更關(guān)鍵的內(nèi)容——它們?nèi)绾瓮ㄟ^(guò)芯片間互連（inter-chip links）形成一個(gè)計(jì)算集群，并受到帶寬和延遲的限制。

五年前，機(jī)器學(xué)習(xí)領(lǐng)域的架構(gòu)還十分多樣化包——括卷積神經(jīng)網(wǎng)絡(luò)、長(zhǎng)短時(shí)記憶網(wǎng)絡(luò)、多層感知機(jī)和Transformer等。如今，Transformer架構(gòu)一家獨(dú)大。

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

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

這些計(jì)算將揭示：

模型的內(nèi)存占用有多大？
計(jì)算和通信的時(shí)間消耗分布如何？
注意力機(jī)制（Attention）和前饋網(wǎng)絡(luò)（Feed-Forward Blocks）何時(shí)成為計(jì)算的瓶頸？

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

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

《第5章：訓(xùn)練》和《第7章：推理》是本書(shū)的核心內(nèi)容，在這兩章中將討論一個(gè)根本問(wèn)題：

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

這個(gè)看似簡(jiǎn)單的問(wèn)題，其實(shí)有著令人意外的復(fù)雜答案。

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

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

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

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

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

《第6章》和《第8章》是實(shí)際操作教程，應(yīng)用這些概念于LLaMA-3，更直觀地理解如何在實(shí)際應(yīng)用中進(jìn)行操作。