谷歌TPU超算，大模型性能超英偉達，已部署數十臺：圖靈獎得主新作

與英偉達的GPU相比，谷歌TPU採用低精度計算，幾乎不影響深度學習處理效果的前提，比基於英偉達 A100 芯片的系統快 1.7 倍，能效高 1.9 倍。

我們還沒有看到能與 ChatGPT 相匹敵的 AI 大模型，但在算力基礎上，領先的可能並不是微軟和 OpenAI。

本週二，谷歌公佈了其訓練語言大模型的超級計算機的細節，基於 TPU 的超算系統已經可以比英偉達的同類更加快速、節能。

谷歌張量處理器（tensor processing unit，TPU）是該公司爲機器學習定製的專用芯片（ASIC），第一代發佈於 2016 年，成爲了 AlphaGo 背後的算力。與 GPU 相比，TPU採用低精度計算，在幾乎不影響深度學習處理效果的前提下大幅降低了功耗、加快運算速度。同時，TPU 使用了脈動陣列等設計來優化矩陣乘法與卷積運算。

當前，谷歌 90% 以上的人工智能訓練工作都在使用這些芯片，TPU 支撐了包括搜索的谷歌主要業務。作爲圖靈獎得主、計算機架構巨擘，大衛・帕特森（David Patterson）在 2016 年從 UC Berkeley 退休後，以傑出工程師的身份加入了谷歌大腦團隊，爲幾代 TPU 的研發做出了卓越貢獻。

如今 TPU 已經發展到了第四代，谷歌本週二由 Norman Jouppi、大衛・帕特森等人發表的論文《 TPU v4： An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings 》詳細介紹了自研的光通信器件是如何將 4000 多塊芯片並聯成爲超級計算機，以提升整體效率的。

TPU v4 的性能比 TPU v3 高 2.1 倍，性能功耗比提高 2.7 倍。基於 TPU v4 的超級計算機擁有 4096 塊芯片，整體速度提高了約 10 倍。對於類似大小的系統，谷歌能做到比 Graphcore IPU Bow 快 4.3-4.5 倍，比 Nvidia A100 快 1.2-1.7 倍，功耗低 1.3-1.9 倍。

除了芯片本身的算力，芯片間互聯已成爲構建 AI 超算的公司之間競爭的關鍵點，最近一段時間，谷歌的 Bard、OpenAI 的 ChatGPT 這樣的大語言模型（LLM）規模正在爆炸式增長，算力已經成爲明顯的瓶頸。

由於大模型動輒千億的參數量，它們必須由數千塊芯片共同分擔，並持續數週或更長時間進行訓練。谷歌的 PaLM 模型 —— 其迄今爲止最大的公開披露的語言模型 —— 在訓練時被拆分到了兩個擁有 4000 塊 TPU 芯片的超級計算機上，用時 50 天。

谷歌表示，通過光電路交換機（OCS），其超級計算機可以輕鬆地動態重新配置芯片之間的連接，有助於避免出現問題並實時調整以提高性能。

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

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

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

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

與 Infiniband 相比，OCS 的成本更低、功耗更低、速度更快，成本不到系統成本的 5%，功率不到系統功率的 3%。每個 TPU v4 都包含 SparseCores 數據流處理器，可將依賴嵌入的模型加速 5 至 7 倍，但僅使用 5% 的裸片面積和功耗。

‘這種切換機制使得繞過故障組件變得容易，’谷歌研究員 Norm Jouppi 和谷歌傑出工程師大衛・帕特森在一篇關於該系統的博客文章中寫道。‘這種靈活性甚至允許我們改變超級計算機互連的拓撲結構，以加速機器學習模型的性能。’

在新論文上，谷歌着重介紹了稀疏核（SparseCore，SC）的設計。在大模型的訓練階段，embedding 可以放在 TensorCore 或超級計算機的主機 CPU 上處理。TensorCore 具有寬 VPU 和矩陣單元，並針對密集操作進行了優化。由於小的聚集 / 分散內存訪問和可變長度數據交換，在 TensorCore 上放置嵌入其實並不是最佳選擇。在超級計算機的主機 CPU 上放置嵌入會在 CPU DRAM 接口上引發阿姆達爾定律瓶頸，並通過 4：1 TPU v4 與 CPU 主機比率放大。數據中心網絡的尾部延遲和帶寬限制將進一步限制訓練系統。

對此，谷歌認爲可以使用 TPU 超算的總 HBM 容量優化性能，加入專用 ICI 網絡，並提供快速收集 / 分散內存訪問支持。這導致了 SparseCore 的協同設計。

SC 是一種用於嵌入訓練的特定領域架構，從 TPU v2 開始，後來在 TPU v3 和 TPU v4 中得到改進。SC 相對划算，只有芯片面積的約 5% 和功率的 5% 左右。SC 結合超算規模的 HBM 和 ICI 來創建一個平坦的、全局可尋址的內存空間（TPU v4 中爲 128 TiB）。與密集訓練中大參數張量的全部歸約相比，較小嵌入向量的全部傳輸使用 HBM 和 ICI 以及更細粒度的分散 / 聚集訪問模式。

作爲獨立的核心，SC 允許跨密集計算、SC 和 ICI 通信進行並行化。下圖顯示了 SC 框圖，谷歌將其視爲‘數據流’架構（dataflow），因爲數據從內存流向各種直接連接的專用計算單元。

最通用的 SC 單元是 16 個計算塊（深藍色框）。每個 tile 都有一個關聯的 HBM 通道，並支持多個未完成的內存訪問。每個 tile 都有一個 Fetch Unit、一個可編程的 8-wide SIMD Vector Processing Unit 和一個 Flush Unit。獲取單元將 HBM 中的激活和參數讀取到 2.5 MiB 稀疏向量內存（Spmem）的圖塊切片中。scVPU 使用與 TC 的 VPU 相同的 ALU。Flush Unit 在向後傳遞期間將更新的參數寫入 HBM。此外，五個跨通道單元（金色框）執行特定的嵌入操作，正如它們的名稱所解釋的那樣。

與 TPU v1 一樣，這些單元執行類似 CISC 的指令並對可變長度輸入進行操作，其中每條指令的運行時間都取決於數據。

在特定芯片數量下，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 倍改進。