性能和效率都超越英偉達(dá) A100，這樣的超算我有不止十臺(tái)。

我們還沒(méi)有看到能與 ChatGPT 相匹敵的 AI 大模型，但在算力基礎(chǔ)上，領(lǐng)先的可能并不是微軟和 OpenAI。 ? 本周二，谷歌公布了其訓(xùn)練語(yǔ)言大模型的超級(jí)計(jì)算機(jī)的細(xì)節(jié)，基于 TPU 的超算系統(tǒng)已經(jīng)可以比英偉達(dá)的同類(lèi)更加快速、節(jié)能。 ? 谷歌張量處理器（tensor processing unit，TPU）是該公司為機(jī)器學(xué)習(xí)定制的專(zhuān)用芯片（ASIC），第一代發(fā)布于 2016 年，成為了 AlphaGo 背后的算力。與 GPU 相比，TPU 采用低精度計(jì)算，在幾乎不影響深度學(xué)習(xí)處理效果的前提下大幅降低了功耗、加快運(yùn)算速度。同時(shí)，TPU 使用了脈動(dòng)陣列等設(shè)計(jì)來(lái)優(yōu)化矩陣乘法與卷積運(yùn)算。 ? 當(dāng)前，谷歌 90% 以上的人工智能訓(xùn)練工作都在使用這些芯片，TPU 支撐了包括搜索的谷歌主要業(yè)務(wù)。作為圖靈獎(jiǎng)得主、計(jì)算機(jī)架構(gòu)巨擘，大衛(wèi)?帕特森（David Patterson）在 2016 年從 UC Berkeley 退休后，以杰出工程師的身份加入了谷歌大腦團(tuán)隊(duì)，為幾代 TPU 的研發(fā)做出了卓越貢獻(xiàn)。 ? ?

?

? 如今 TPU 已經(jīng)發(fā)展到了第四代，谷歌本周二由 Norman Jouppi、大衛(wèi)?帕特森等人發(fā)表的論文《 TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings 》詳細(xì)介紹了自研的光通信器件是如何將 4000 多塊芯片并聯(lián)成為超級(jí)計(jì)算機(jī)，以提升整體效率的。 ?

?TPU v4 的性能比 TPU v3 高 2.1 倍，性能功耗比提高 2.7 倍?；?TPU v4 的超級(jí)計(jì)算機(jī)擁有 4096 塊芯片，整體速度提高了約 10 倍。對(duì)于類(lèi)似大小的系統(tǒng)，谷歌能做到比 Graphcore IPU Bow 快 4.3-4.5 倍，比 Nvidia A100 快 1.2-1.7 倍，功耗低 1.3-1.9 倍。 ? 除了芯片本身的算力，芯片間互聯(lián)已成為構(gòu)建 AI 超算的公司之間競(jìng)爭(zhēng)的關(guān)鍵點(diǎn)，最近一段時(shí)間，谷歌的 Bard、OpenAI 的 ChatGPT 這樣的大語(yǔ)言模型（LLM）規(guī)模正在爆炸式增長(zhǎng)，算力已經(jīng)成為明顯的瓶頸。 ? ? 由于大模型動(dòng)輒千億的參數(shù)量，它們必須由數(shù)千塊芯片共同分擔(dān)，并持續(xù)數(shù)周或更長(zhǎng)時(shí)間進(jìn)行訓(xùn)練。谷歌的 PaLM 模型 —— 其迄今為止最大的公開(kāi)披露的語(yǔ)言模型 —— 在訓(xùn)練時(shí)被拆分到了兩個(gè)擁有 4000 塊 TPU 芯片的超級(jí)計(jì)算機(jī)上，用時(shí) 50 天。 ? 谷歌表示，通過(guò)光電路交換機(jī)（OCS），其超級(jí)計(jì)算機(jī)可以輕松地動(dòng)態(tài)重新配置芯片之間的連接，有助于避免出現(xiàn)問(wèn)題并實(shí)時(shí)調(diào)整以提高性能。 ?

? 下圖展示了 TPU v4 4×3 方式 6 個(gè)「面」的鏈接。每個(gè)面有 16 條鏈路，每個(gè)塊總共有 96 條光鏈路連接到 OCS 上。要提供 3D 環(huán)面的環(huán)繞鏈接，相對(duì)側(cè)的鏈接必須連接到相同的 OCS。因此，每個(gè) 4×3 塊 TPU 連接到 6 × 16 ÷ 2 = 48 個(gè) OCS 上。Palomar OCS 為 136×136（128 個(gè)端口加上 8 個(gè)用于鏈路測(cè)試和修復(fù)的備用端口），因此 48 個(gè) OCS 連接來(lái)自 64 個(gè) 4×3 塊（每個(gè) 64 個(gè)芯片）的 48 對(duì)電纜，總共并聯(lián) 4096 個(gè) TPU v4 芯片。 ?

? 根據(jù)這樣的排布，TPU v4（中間的 ASIC 加上 4 個(gè) HBM 堆棧）和帶有 4 個(gè)液冷封裝的印刷電路板 (PCB)。該板的前面板有 4 個(gè)頂部 PCIe 連接器和 16 個(gè)底部 OSFP 連接器，用于托盤(pán)間 ICI 鏈接。 ?

?

? 隨后，八個(gè) 64 芯片機(jī)架構(gòu)成一臺(tái) 4096 芯片超算。 ?

?

? 與超級(jí)計(jì)算機(jī)一樣，工作負(fù)載由不同規(guī)模的算力承擔(dān)，稱(chēng)為切片：64 芯片、128 芯片、256 芯片等。下圖顯示了當(dāng)主機(jī)可用性從 99.0% 到 99.9% 不等有，及沒(méi)有 OCS 時(shí)切片大小的「有效輸出」。如果沒(méi)有 OCS，主機(jī)可用性必須達(dá)到 99.9% 才能提供合理的切片吞吐量。對(duì)于大多數(shù)切片大小，OCS 也有 99.0% 和 99.5% 的良好輸出。 ?

? 與 Infiniband 相比，OCS 的成本更低、功耗更低、速度更快，成本不到系統(tǒng)成本的 5%，功率不到系統(tǒng)功率的 3%。每個(gè) TPU v4 都包含 SparseCores 數(shù)據(jù)流處理器，可將依賴(lài)嵌入的模型加速 5 至 7 倍，但僅使用 5% 的裸片面積和功耗。 ? 「這種切換機(jī)制使得繞過(guò)故障組件變得容易，」谷歌研究員 Norm Jouppi 和谷歌杰出工程師大衛(wèi)?帕特森在一篇關(guān)于該系統(tǒng)的博客文章中寫(xiě)道?！高@種靈活性甚至允許我們改變超級(jí)計(jì)算機(jī)互連的拓?fù)浣Y(jié)構(gòu)，以加速機(jī)器學(xué)習(xí)模型的性能?！??

? 在新論文上，谷歌著重介紹了稀疏核（SparseCore，SC）的設(shè)計(jì)。在大模型的訓(xùn)練階段，embedding 可以放在 TensorCore 或超級(jí)計(jì)算機(jī)的主機(jī) CPU 上處理。TensorCore 具有寬 VPU 和矩陣單元，并針對(duì)密集操作進(jìn)行了優(yōu)化。由于小的聚集 / 分散內(nèi)存訪(fǎng)問(wèn)和可變長(zhǎng)度數(shù)據(jù)交換，在 TensorCore 上放置嵌入其實(shí)并不是最佳選擇。在超級(jí)計(jì)算機(jī)的主機(jī) CPU 上放置嵌入會(huì)在 CPU DRAM 接口上引發(fā)阿姆達(dá)爾定律瓶頸，并通過(guò) 4:1 TPU v4 與 CPU 主機(jī)比率放大。數(shù)據(jù)中心網(wǎng)絡(luò)的尾部延遲和帶寬限制將進(jìn)一步限制訓(xùn)練系統(tǒng)。 ? 對(duì)此，谷歌認(rèn)為可以使用 TPU 超算的總 HBM 容量?jī)?yōu)化性能，加入專(zhuān)用 ICI 網(wǎng)絡(luò)，并提供快速收集 / 分散內(nèi)存訪(fǎng)問(wèn)支持。這導(dǎo)致了 SparseCore 的協(xié)同設(shè)計(jì)。 ?

? SC 是一種用于嵌入訓(xùn)練的特定領(lǐng)域架構(gòu)，從 TPU v2 開(kāi)始，后來(lái)在 TPU v3 和 TPU v4 中得到改進(jìn)。SC 相對(duì)劃算，只有芯片面積的約 5% 和功率的 5% 左右。SC 結(jié)合超算規(guī)模的 HBM 和 ICI 來(lái)創(chuàng)建一個(gè)平坦的、全局可尋址的內(nèi)存空間（TPU v4 中為 128 TiB）。與密集訓(xùn)練中大參數(shù)張量的全部歸約相比，較小嵌入向量的全部傳輸使用 HBM 和 ICI 以及更細(xì)粒度的分散 / 聚集訪(fǎng)問(wèn)模式。 ? 作為獨(dú)立的核心，SC 允許跨密集計(jì)算、SC 和 ICI 通信進(jìn)行并行化。下圖顯示了 SC 框圖，谷歌將其視為「數(shù)據(jù)流」架構(gòu)（dataflow），因?yàn)閿?shù)據(jù)從內(nèi)存流向各種直接連接的專(zhuān)用計(jì)算單元。 ?

? 最通用的 SC 單元是 16 個(gè)計(jì)算塊（深藍(lán)色框）。每個(gè) tile 都有一個(gè)關(guān)聯(lián)的 HBM 通道，并支持多個(gè)未完成的內(nèi)存訪(fǎng)問(wèn)。每個(gè) tile 都有一個(gè) Fetch Unit、一個(gè)可編程的 8-wide SIMD Vector Processing Unit 和一個(gè) Flush Unit。獲取單元將 HBM 中的激活和參數(shù)讀取到 2.5 MiB 稀疏向量?jī)?nèi)存 (Spmem) 的圖塊切片中。scVPU 使用與 TC 的 VPU 相同的 ALU。Flush Unit 在向后傳遞期間將更新的參數(shù)寫(xiě)入 HBM。此外，五個(gè)跨通道單元（金色框）執(zhí)行特定的嵌入操作，正如它們的名稱(chēng)所解釋的那樣。 ? 與 TPU v1 一樣，這些單元執(zhí)行類(lèi)似 CISC 的指令并對(duì)可變長(zhǎng)度輸入進(jìn)行操作，其中每條指令的運(yùn)行時(shí)間都取決于數(shù)據(jù)。 ? 在特定芯片數(shù)量下，TPU v3/v4 對(duì)分帶寬比高 2-4 倍，嵌入速度可以提高 1.1-2.0 倍。 ?

? 下圖展示了谷歌自用的推薦模型（DLRM0）在不同芯片上的效率。TPU v3 比 CPU 快 9.8 倍。TPU v4 比 TPU v3 高 3.1 倍，比 CPU 高 30.1 倍。 ? ?

? 谷歌探索了 TPU v4 超算用于 GPT-3 大語(yǔ)言模型時(shí)的性能，展示了預(yù)訓(xùn)練階段專(zhuān)家設(shè)計(jì)的 1.2 倍改進(jìn)。 ?

? 雖然谷歌直到現(xiàn)在才公布有關(guān)其超級(jí)計(jì)算機(jī)的詳細(xì)信息，但自 2020 年以來(lái)，基于 TPU 的 AI 超算一直在位于俄克拉荷馬州的數(shù)據(jù)中心發(fā)揮作用。谷歌表示，Midjourney 一直在使用該系統(tǒng)訓(xùn)練其模型，最近一段時(shí)間，后者已經(jīng)成為 AI 畫(huà)圖領(lǐng)域最熱門(mén)的平臺(tái)。 ? ? ? 谷歌在論文中表示，對(duì)于同等大小的系統(tǒng)，其芯片比基于英偉達(dá) A100 芯片的系統(tǒng)快 1.7 倍，能效高 1.9 倍，后者與第四代 TPU 同時(shí)上市，并被用于 GPT-4 的訓(xùn)練。 ? 對(duì)此，英偉達(dá)發(fā)言人拒絕置評(píng)。 ? ? 當(dāng)前英偉達(dá)的 AI 芯片已經(jīng)進(jìn)入 Hopper 架構(gòu)的時(shí)代。谷歌表示，未對(duì)第四代 TPU 與英偉達(dá)目前的旗艦 H100 芯片進(jìn)行比較，因?yàn)?H100 在谷歌芯片之后上市，并且采用了更先進(jìn)的制程。 ? 但同樣在此，谷歌暗示了下一代 TPU 的計(jì)劃，其沒(méi)有提供更多細(xì)節(jié)。Jouppi 告訴路透社，谷歌擁有開(kāi)發(fā)「未來(lái)芯片的健康管道」。 ?

? ? TPU v4 比當(dāng)代 DSA 芯片速度更快、功耗更低，如果考慮到互連技術(shù)，功率邊緣可能會(huì)更大。通過(guò)使用具有 3D 環(huán)面拓?fù)涞?3K TPU v4 切片，與 TPU v3 相比，谷歌的超算也能讓 LLM 的訓(xùn)練時(shí)間大大減少。 ? 性能、可擴(kuò)展性和可用性使 TPU v4 超級(jí)計(jì)算機(jī)成為 LaMDA、MUM 和 PaLM 等大型語(yǔ)言模型 (LLM) 的主要算力。這些功能使 5400 億參數(shù)的 PaLM 模型在 TPU v4 超算上進(jìn)行訓(xùn)練時(shí)，能夠在 50 天內(nèi)維持 57.8% 的峰值硬件浮點(diǎn)性能。 ? 谷歌表示，其已經(jīng)部署了數(shù)十臺(tái) TPU v4 超級(jí)計(jì)算機(jī)，供內(nèi)部使用和外部通過(guò)谷歌云使用。

編輯：黃飛

?