GoogleTPU超算,大模型性能超英偉達,已部署數十臺:圖靈獎得主新作

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

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

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

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

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

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

論文連結:https://arxiv.org/ftp/arxiv/papers/2304/2304.01433.pdf

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

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

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

Google表示,通過光電路交換機(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% 的裸片面積和功耗。

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

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

對此,Google認為可以使用 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 框圖,Google將其視為「資料流」架構(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 倍。

下圖展示了Google自用的推薦模型(DLRM0)在不同晶片上的效率。TPU v3 比 CPU 快 9.8 倍。TPU v4 比 TPU v3 高 3.1 倍,比 CPU 高 30.1 倍。

Google探索了 TPU v4 超算用於 GPT-3 大語言模型時的性能,展示了預訓練階段專家設計的 1.2 倍改進。

雖然Google直到現在才公佈有關其超級計算機的詳細資訊,但自 2020 年以來,基於 TPU 的 AI 超算一直在位於俄克拉荷馬州的資料中心發揮作用。Google表示,Midjourney 一直在使用該系統訓練其模型,最近一段時間,後者已經成為 AI 畫圖領域最熱門的平臺。

Google在論文中表示,對於同等大小的系統,其晶片比基於英偉達 A100 晶片的系統快 1.7 倍,能效高 1.9 倍,後者與第四代 TPU 同時上市,並被用於 GPT-4 的訓練。

對此,英偉達發言人拒絕置評。

當前英偉達的 AI 晶片已經進入 Hopper 架構的時代。Google表示,未對第四代 TPU 與英偉達目前的旗艦 H100 晶片進行比較,因為 H100 在Google晶片之後上市,並且採用了更先進的製程。

但同樣在此,Google暗示了下一代 TPU 的計劃,其沒有提供更多細節。Jouppi 告訴路透社,Google擁有開發「未來晶片的健康管道」。

TPU v4 比當代 DSA 晶片速度更快、功耗更低,如果考慮到互連技術,功率邊緣可能會更大。通過使用具有 3D 環面拓撲的 3K TPU v4 切片,與 TPU v3 相比,Google的超算也能讓 LLM 的訓練時間大大減少。

性能、可擴展性和可用性使 TPU v4 超級計算機成為 LaMDA、MUM 和 PaLM 等大型語言模型 (LLM) 的主要算力。這些功能使 5400 億參數的 PaLM 模型在 TPU v4 超算上進行訓練時,能夠在 50 天內維持 57.8% 的峰值硬體浮點性能。

Google表示,其已經部署了數十臺 TPU v4 超級計算機,供內部使用和外部通過Google雲使用。

參考內容:https://www.reuters.com/technology/google-says-its-ai-supercomputer-is-faster-greener-than-nvidia-2023-04-05/

相關文章

當創業者開始用 ChatGPT 裁員

當創業者開始用 ChatGPT 裁員

他說,從沒想到科幻片的場景有一天會發生在自己身上。 作者 | 黎詩韻編輯| 衛詩婕 幾個月來,對話機器人 ChatGPT 展現出的能力令世界...

ChatGPT 背後的經濟賬

ChatGPT 背後的經濟賬

拿投資和賺錢之前,要先去了解下成本。 ChatGPT 能否取代 Google、百度這樣的傳統搜尋引擎?為什麼中國不能很快做出 ChatGPT...

ChatGPT 的「神功」,是如何煉成的?

ChatGPT 的「神功」,是如何煉成的?

最強對話式 AI,站在巨人的肩膀上。 AIGC 在這幾個月成了一個大熱的話題。 頗有些風水輪流轉的感覺,如同年初大火特火的 web3 一樣,...