基于 TPU v4的超級計算機性能解析

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

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

?

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

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

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

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

?

? 隨后，八個 64 芯片機架構成一臺 4096 芯片超算。 ?

?

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

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

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

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

? 最通用的 SC 單元是 16 個計算塊（深藍色框）。每個 tile 都有一個關聯(lián)的 HBM 通道，并支持多個未完成的內(nèi)存訪問。每個 tile 都有一個 Fetch Unit、一個可編程的 8-wide SIMD Vector Processing Unit 和一個 Flush Unit。獲取單元將 HBM 中的激活和參數(shù)讀取到 2.5 MiB 稀疏向量內(nèi)存 (Spmem) 的圖塊切片中。scVPU 使用與 TC 的 VPU 相同的 ALU。Flush Unit 在向后傳遞期間將更新的參數(shù)寫入 HBM。此外，五個跨通道單元（金色框）執(zhí)行特定的嵌入操作，正如它們的名稱所解釋的那樣。 ? 與 TPU v1 一樣，這些單元執(zhí)行類似 CISC 的指令并對可變長度輸入進行操作，其中每條指令的運行時間都取決于數(shù)據(jù)。 ? 在特定芯片數(shù)量下，TPU v3/v4 對分帶寬比高 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 大語言模型時的性能，展示了預訓練階段專家設計的 1.2 倍改進。 ?

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

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

編輯：黃飛

?