AI 時代的曙光就在這里，了解 AI 驅動軟件的成本結構與傳統軟件有很大差異至關重要。芯片微架構和系統架構在這些創新軟件的開發和可擴展性方面發揮著至關重要的作用。運行軟件的硬件基礎設施對資本支出和運營支出以及隨后的毛利率有明顯更大的影響，這與前幾代軟件相比，前幾代軟件的開發人員成本相對較高。因此，更加重要的是要投入大量精力來優化您的 AI 基礎設施，以便能夠部署 AI 軟件。在基礎設施方面具有優勢的公司在使用 AI 部署和擴展應用程序的能力方面也將具有優勢。

谷歌早在 2006 年就提出了構建 AI 專用基礎設施的想法，但這個問題在 2013 年達到了定點。他們意識到，如果他們想以任何規模部署 AI，就需要將現有數據中心的數量增加一倍。因此，他們開始為 2016 年投入生產的 TPU 芯片奠定基礎。將此與亞馬遜進行比較很有趣，亞馬遜在同一年意識到他們也需要構建定制芯片。

自 2016 年以來，谷歌現已構建了 6 種不同的 AI 芯片，TPU、TPUv2、TPUv3、TPUv4i、TPUv4 和 TPUv5。谷歌主要設計了這些芯片，并與博通進行了不同程度的中后端合作。這些芯片全部由臺積電代工。自 TPUv2 以來，這些芯片還使用了三星和 SK 海力士的 HBM 內存。雖然谷歌的芯片架構很有趣，我們將在本報告的后面深入探討，但還有一個更重要的話題在起作用。

谷歌擁有近乎無與倫比的能力，能夠以低成本和高性能可靠地大規模部署人工智能。話雖如此，讓我們為這個論點帶來一些合理性，因為谷歌也做出了與芯片級性能相關的虛偽聲明，需要糾正。我們認為，由于谷歌從微架構到系統架構的整體方法，與微軟和亞馬遜相比，谷歌在 AI 工作負載方面具有性能/總擁有成本 (perf/TCO) 優勢，而將生成人工智能商業化給企業和消費者的能力是一個不同的討論。

技術領域是一場永無休止的軍備競賽，人工智能是移動最快的戰場。隨著時間的推移，經過訓練和部署的模型架構發生了顯著變化。案例和重點是谷歌的內部數據。CNN 模型在 2016 年到 2019 年迅速上升，但隨后又下降了。與 DLRM、Transformers 和 RNN 相比，CNN 在計算、內存訪問、網絡等方面具有非常不同的配置文件。同樣的情況也發生在完全被 Transformer 取代的 RNN 上。

因此，硬件必須靈活地適應行業的發展并支持它們。底層硬件不能過度專注于任何特定的模型架構，否則它可能會隨著模型架構的變化而變得過時。芯片開發到大規模部署一般需要 4 年時間，因此，硬件可以被軟件想在其上做的事情拋在腦后。這已經可以從使用特定模型類型作為優化點的初創公司的某些 AI 加速器架構中看出。這是大多數 AI 硬件初創公司已經/將要失敗的眾多原因之一。

這一點在谷歌自己的 TPUv4i 芯片上尤為明顯，該芯片專為推理而設計，但無法在谷歌最好的模型（如 PaLM）上運行推理。上一代 Google TPUv4 和 Nvidia A100 不可能在設計時考慮到大型語言模型。同樣，最近部署的谷歌 TPUv5 和 Nvidia H100 不可能在設計時考慮到”AI墻”，也沒有為解決它而開發的新模型架構策略。這些策略是 GPT-4 模型架構的核心部分。

硬件架構師必須對機器學習在他們設計的芯片中的發展方向做出最好的猜測。這包括內存訪問模式、張量大小、數據重用結構、算術密度與網絡開銷等。

此外，芯片微架構只是人工智能基礎設施真實成本的一小部分。系統級架構和部署靈活性是更為重要的因素。今天，我們想深入探討 Google 的 TPU 微架構、系統架構、部署切片、可擴展性，以及他們在基礎設施方面與其他技術巨頭相比的巨大優勢。這包括我們在 TCO 模型中的想法，該模型將 Google 的 AI 基礎設施成本與 Microsoft、Amazon 和 Meta 的成本進行比較。

我們還將從從業者的角度對大型模型研究、訓練和部署進行研究。我們還想深入研究 DLRM 模型，盡管目前是最大的大規模 AI 模型架構，但這些模型經常被低估。此外，我們將討論 DLRM 和 LLM 模型類型之間的基礎設施差異。最后，我們將討論谷歌利用 TPU 為外部云客戶取得成功的能力。同樣在最后，我們認為谷歌的 TPU 有一個異常的復活節彩蛋是一個錯誤。

谷歌的系統基礎設施優勢

谷歌在基礎設施方面的部分優勢在于，他們始終從系統級的角度設計 TPU。這意味著單個芯片很重要，但如何在現實世界的系統中一起使用它更為重要。因此，在我們的分析中，我們將逐層從系統架構到部署使用再到芯片級別。

雖然從系統的角度思考，但他們的系統規模比谷歌更小、更窄。此外，直到最近，Nvidia 還沒有云部署方面的經驗。谷歌在其 AI 基礎設施方面最大的創新之一是在 TPU、ICI 之間使用自定義網絡堆棧。相對于昂貴的以太網和 InfiniBand 部署，此鏈接具有低延遲和高性能。它更類似于 Nvidia 的 NVLink。

谷歌的 TPUv2 可以擴展到 256 個 TPU 芯片，與 Nvidia 當前一代 H100 GPU 的數量相同。他們使用 TPUv3 將這個數字增加到 1024，使用 TPUv4 增加到 4096。根據趨勢線，我們假設當前一代 TPUv5 可以擴展到 16,384 個芯片，而無需通過低效的以太網。雖然從大規模模型訓練的性能角度來看這很重要，但更重要的是他們將其劃分以供實際使用的能力。

谷歌的 TPUv4 系統每臺服務器有 8 個 TPUv4 芯片和 2 個 CPU。此配置與 Nvidia 的 GPU 相同，后者配備 8 個 A100 或 H100 服務器，每臺服務器 2 個 CPU。單個服務器通常是 GPU 部署的計算單元，但對于 TPU，部署單元是更大的“slice”，由 64 個 TPU 芯片和 16 個 CPU 組成。這 64 個芯片通過直接連接的銅纜在 4^3 立方體中與 ICI 網絡內部連接。

在這個 64 芯片單元之外，通信轉而轉移到光學領域。這些光收發器的成本是無源銅纜的 10 倍以上。

將此與 2023 Nvidia SuperPod 部署進行比較，后者使用 NVLink 最多配備 256 個 GPU，僅為 4096 的芯片的 2020 TPUv4 pod 的十六分之一。此外，基于 Nvidia 的第一方渲染和 DGX Superpod 系統，Nvidia 顯然不太關注密度和網絡成本。Nvidia 的部署通常是每個機架 4 個服務器。

除了 4 臺服務器總共 32 個 GPU 之外，通常，通信必須采用光學方式。因此，Nvidia 需要更多的光收發器來進行大規模部署。

谷歌OCS

谷歌部署了其定制光開關，它使用基于 mems 的微鏡陣列陣列在 64 個 TPU slice之間切換。簡短的總結是，谷歌聲稱他們的自定義網絡將吞吐量提高了 30%，使用的電力減少了 40%，資本支出減少了 30%，流程完成減少了 10%，并且在他們的網絡中減少了 50 倍的停機時間，

谷歌使用這些 OCS 來構建其數據中心主干。他們還使用它們將 TPU pod 互連和內部連接在一起。此 OCS 的一大優勢是信號僅保留在光域中，從任何 64 TPU slice到 4096 TPU Pod 內的任何其他 TPU slice。

將此與具有多個 Nvidia SuperPods 的 4,096 個 GPU 的 Nvidia GPU 部署進行比較。該系統需要在這些 GPU 之間進行多層切換，總共需要約 568 個 InfiniBand 交換機。谷歌只需要 48 個光開關來部署 4096 個 TPU。

應該注意的是，與第三方從 Nvidia 購買 Nvidia 的 InfiniBand 交換機相比，直接從 Google 的合同制造商處購買時， Google 的 OCS每個交換機的價格也高出 3.2 到 3.5 倍。不過，這不是一個公平的比較，因為它包括 Nvidia 約 75% 的數據中心毛利率。

如果我們只比較合同制造成本，谷歌的 IE 成本與 Nvidia 的成本；然后成本差異上升到 Nvidia InfiniBand 交換機的 12.8 到 14 倍。部署4096芯片所需的交換機數量為48 vs 568，IE為11.8x。Nvidia 的解決方案在交換機基礎上的制造成本更低。當包括額外的光收發器的成本時，這個等式趨于平衡或向有利于谷歌的方向移動。

每層交換之間的每個連接都是另一個需要更多布線的點。雖然其中一些可以通過直接連接的銅纜完成，但仍有多個點的信號也需要通過光纖傳輸。這些層中的每一層都會在每一層切換之間從電轉換為光再轉換為電。這將使大型電氣開關系統的功耗遠高于谷歌的OCS。

谷歌聲稱所有這些功率和成本的節省都非常大，以至于它們的網絡成本不到 TPU v4 超級計算機總資本成本的 5% 和總功率的不到 3%。這不僅僅是通過從電氣開關轉向內部光開關來實現的。

通過拓撲最小化網絡成本

雖然谷歌大力推動這一觀點，但重要的是要認識到 Nvidia 和 Nvidia 網絡的拓撲結構完全不同。Nvidia 系統部署了“non-blocking”的“Clos 網絡”。這意味著它們可以同時在所有輸入和輸出對之間建立全帶寬連接，而不會發生任何沖突或阻塞。此設計提供了一種可擴展的方法，用于連接數據中心中的許多設備、最大限度地減少延遲并增加冗余。

谷歌的 TPU 網絡放棄了這一點。他們使用 3D 環面拓撲連接三維網格狀結構中的節點。每個節點都連接到網格中的六個相鄰節點（上、下、左、右、前和后），在三個維度（X、Y 和 Z）中的每一個維度上形成一個閉環。這創建了一個高度互連的結構，其中節點在所有三個維度上形成一個連續的循環。

第一張圖比較合乎邏輯，但如果你想一想有點餓了，這個網絡拓撲簡直就是一個甜甜圈！

與 Nvidia 使用的 Clos 拓撲相比，torus 拓撲有幾個優點：

更低的延遲：3D 環面拓撲可以提供更低的延遲，因為它在相鄰節點之間有短而直接的鏈接。這在運行需要節點之間頻繁通信的緊密耦合的并行應用程序時特別有用，例如某些類型的 AI 模型。

更好的局部性：在 3D 環面網絡中，物理上彼此靠近的節點在邏輯上也很接近，這可以帶來更好的數據局部性并減少通信開銷。雖然延遲是一個方面，但功耗也是一個巨大的好處。

較低的網絡直徑：對于相同數量的節點，3D 環面拓撲的網絡直徑低于 Clos 網絡。由于相對于 Clos 網絡需要更少的交換機，因此可以節省大量成本。

另一方面，3D 環面網絡有很多缺點。

可預測的性能：Clos 網絡，尤其是在數據中心環境中，由于其非阻塞特性，可以提供可預測和一致的性能。它們確保所有輸入輸出對都可以在全帶寬下同時連接，而不會發生沖突或阻塞，而這在 3D 環面網絡中是無法保證的。

更易于擴展：在脊葉（spine-leaf ）架構中，向網絡添加新的葉交換機（例如，以容納更多服務器）相對簡單，不需要對現有基礎設施進行重大更改。相比之下，縮放 3D 環面網絡可能涉及重新配置整個拓撲，這可能更加復雜和耗時。

負載平衡：Clos 網絡在任意兩個節點之間提供更多路徑，從而實現更好的負載平衡和冗余。雖然 3D 環面網絡也提供多條路徑，但 Clos 網絡中的備選路徑數量可能更多，具體取決于網絡的配置。

總的來說，雖然 Clos 有優勢，但谷歌的 OCS 減輕了其中的許多優勢。OCS 支持在多個切片和多個 pod 之間進行簡單縮放。

3D 環面拓撲面臨的最大問題是錯誤可能是一個更大的問題。錯誤可能會突然出現并發生。即使主機可用性為 99%，2,048 個 TPU 的幻燈片也將具有接近 0 的正常工作能力。即使在 99.9% 的情況下，使用 2,000 個 TPU 運行的訓練在沒有 Google 的 OCS 的情況下也有 50% 的有效輸出。

OCS 的美妙之處在于它支持動態重新配置路由。

盡管有一些節點出現故障，但仍需要備件以允許調度作業。操作員無法在不冒失敗風險的情況下從 4k 節點 pod 實際調度兩個 2k 節點切片?；?Nvidia 的訓練運行通常需要過多的開銷，專門用于檢查點、拉出故障節點并重新啟動它們。谷歌通過繞過故障節點路由在某種程度上簡化了這一點。

OCS 的另一個好處是切片可以在部署后立即使用，而不是等待整個網絡。

部署基礎設施——用戶的視角

從成本和功耗的角度來看，基礎設施效率很高，允許谷歌每美元部署更多的 TPU，而不是其他公司可以部署的 GPU，但這意味著沒有任何用處。谷歌內部用戶體驗到的最大優勢之一是他們可以根據自己的模型定制基礎設施需求。

沒有任何芯片或系統能夠匹配所有用戶想要的內存、網絡和計算配置文件類型。芯片必須通用化，但與此同時，用戶需要這種靈活性，他們不想要一種放之四海而皆準的解決方案。Nvidia 通過提供許多不同的 SKU 變體來解決這個問題。此外，它們還提供一些不同的內存容量層級以及更緊密的集成選項，例如 Grace + Hopper 和為SuperPods準備的 NVLink Network。

谷歌負擔不起這種奢侈。每個額外的 SKU 意味著每個單獨 SKU 的總部署量較低。這反過來又降低了他們整個基礎設施的利用率。更多的 SKU 也意味著用戶更難在需要時獲得他們想要的計算類型，因為某些選項將不可避免地被超額訂閱。然后，這些用戶將被迫使用次優配置。

因此，谷歌面臨著一個棘手的問題，即向研究人員提供他們想要的確切產品，同時還要最大限度地減少 SKU 差異。與 Nvidia 必須支持其更大、更多樣化的客戶群的數百種不同規模的部署和 SKU 相比，谷歌恰好有 1 個 TPUv4 部署配置，即 4,096 個 TPU。盡管如此，谷歌仍然能夠以一種獨特的方式對其進行slice和切塊，使內部用戶能夠擁有他們想要的基礎設施的靈活性。

Google 的 OCS 還支持創建自定義網絡拓撲，例如扭曲環面網絡。這些是 3d 環面網絡，其中某些維度是扭曲的，這意味著網絡邊緣的節點以非平凡、非線性的方式連接，從而在節點之間創建額外的捷徑。這進一步提高了網絡直徑、負載平衡和性能。

谷歌的團隊充分利用這一點來協助某些模型架構。以下是 2022 年 11 月僅 1 天的各種 TPU 配置的流行情況快照（按芯片數量和網絡拓撲）。有 30 多種不同的配置，盡管許多配置在系統中具有相同數量的芯片，以適應正在開發的各種模型架構。這是谷歌對他們使用 TPU 和靈活性的巨大深刻見解。此外，它們還有許多甚至未被描繪的較少使用的拓撲。

為充分利用可用帶寬，用戶沿 3D 環面的一個維度映射數據并行性，并在其他維度上映射兩個模型并行參數。谷歌聲稱最佳拓撲選擇可以使性能提高 1.2 到 2.3 倍。

規模最大的 AI 模型架構：DLRM

如果不討論深度學習推薦模型 (DLRM)，任何關于 AI 基礎設施的討論都是不完整的。這些 DLRM 是百度、Meta、字節跳動、Netflix 和谷歌等公司的支柱。它是廣告、搜索排名、社交媒體訂閱等領域年收入超過 1 萬億美元的引擎。這些模型包含數十億個權重，對超過一萬億個示例進行訓練，

以每秒超過 300,000 個查詢的速度處理推理。這些模型的大小 (10TB+) 甚至遠遠超過了最大的transformer模型，例如 GPT4，大約為 1TB+（模型架構差異）。

上述所有公司的共同點是，它們依靠不斷更新的 DLRM 來推動其在電子商務、搜索、社交媒體和流媒體服務等各個行業中的個性化內容、產品或服務業務。這些模型的成本是巨大的，必須針對它共同優化硬件。DLRM 并不是一成不變的，而是隨著時間的推移不斷改進，但在繼續之前讓我們先解釋一下通用模型架構。我們將盡量保持簡單。

DLRM 旨在通過對分類和數值特征進行建模來學習用戶-項目交互的有意義的表示。該架構由兩個主要組件組成：嵌入組件（處理分類特征）和多層感知器 (MLP) 組件（處理數字特征）。

用最簡單的術語來說， the 多層感知器組件是密集的。這些特征被饋送到一系列完全連接的層中。這類似于舊的 GPT 4 之前的transformer架構，它們也是密集的，密集層可以很好地映射到硬件上的大規模矩陣多單元。

嵌入組件對于 DLRM 來說是非常獨特的，也是使其計算配置文件如此獨特的組件。DLRM 輸入是表示為離散、稀疏向量的分類特征。一個簡單的谷歌搜索只包含整個語言中的幾個詞。這些稀疏輸入不能很好地映射到硬件中的大量矩陣乘法單元，因為它們從根本上更類似于哈希表，而不是張量。由于神經網絡通常在密集向量上表現更好，因此使用嵌入將分類特征轉換為密集向量。

稀疏輸入：[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]

密集向量：[0.3261477, 0.4263801, 0.5121493]

嵌入函數將分類空間（英語單詞、社交媒體帖子的參與度、對某種帖子的行為）映射到更小的密集空間（100 個向量代表每個單詞）。這些功能是使用查找表實現的，查找表是 DLRM 的重要組成部分，通常構成 DLRM 模型的第一層。嵌入表的大小可以有很大的不同，從幾十兆字節到幾百千兆字節甚至 TB 不等。

Meta 推出 2 年的 DLRM 參數超過 12 萬億個，需要 128 個 GPU 來運行推理。如今，最大的生產 DLRM 模型至少大了好幾倍，并且僅僅為了保存模型嵌入就消耗了超過 30TB 的內存。預計明年嵌入量會增加到 70TB 以上！因此，這些表需要在許多芯片的內存中進行分區。共有三種主要的分區方法：列分片（column sharding）、行分片（row sharding）和表分片（table sharding）。

DLRM 的性能在很大程度上取決于內存帶寬、內存容量、矢量處理性能以及芯片之間的網絡/互連。嵌入查找操作主要由小的收集或分散內存訪問組成，這些訪問具有低算術強度（FLOPS 根本無關緊要）。對嵌入表的訪問基本上是非結構化的稀疏性。每個查詢都必須從 30TB 以上的嵌入中提取數據，這些嵌入分布在數百或數千個芯片上。這會導致用于 DLRM 推理的超級計算機的計算、內存和通信負載不平衡。

這對于 MLP 和類似 GPT-3 的轉換器中的密集操作有很大不同。 芯片 FLOPS/秒仍然是主要性能驅動因素之一當然，除了 FLOPs 之外beyond FLOPs還有多種因素阻礙性能，但 仍然可以在 Chinchilla 風格的 LLM 中實現超過 71% 的硬件觸發器利用率。

谷歌的 TPU 架構

谷歌的 TPU 在架構中引入了一些關鍵創新，使其有別于其他處理器。與傳統處理器不同，TPU v4 沒有專用的指令緩存。相反，它采用類似于 Cell 處理器的直接內存訪問 (DMA) 機制。TPU v4 中的矢量緩存不是標準緩存層次結構的一部分，而是用作暫存器。便簽本與標準緩存的不同之處在于它們需要手動寫入，而標準緩存會自動處理數據。由于不需要服務于大型通用計算市場，谷歌可以利用這種更高效的基礎設施。這確實會在一定程度上影響編程模型，盡管 Google 工程師認為 XLA 編譯器堆?？梢院芎玫靥幚磉@個問題。對于外部用戶則不能這樣說。

TPU v4 擁有用于暫存器的 160MB SRAM 以及 2 個 TensorCore，每個 TensorCore 都有 1 個向量單元和 4 個矩陣乘法單元 (MXU) 和 16MB 向量內存 (VMEM)。兩個 TensorCore 共享 128MB 內存。它們支持 BF16 的 275 TFLOPS，還支持 INT8 數據類型。TPU v4 的內存帶寬為 1200GB/s。芯片間互連 (ICI) 通過六個 50GB/s 鏈路提供 300GB/s 的數據傳輸速率。

TPU v4 中包含一個 322b 超長指令字 (VLIW) 標量計算單元。在 VLIW 架構中，指令被組合成一個單一的長指令字，然后被分派到處理器執行。這些分組指令，也稱為束，在程序編譯期間由編譯器顯式定義。VLIW 包包含多達 2 條標量指令、2 條矢量 ALU 指令、1 條矢量加載和 1 條矢量存儲指令，以及 2 個用于將數據傳入和傳出 MXU 的插槽。

Vector Processing Unit (VPU) 配備了 32 個 2D 寄存器，包含 128x 8 個 32b 元素，使其成為一個 2D 矢量 ALU。矩陣乘法單元 (MXU) 在 v2、v3 和 v4 上為 128x128，v1 版本采用 256x256 配置。發生這種變化的原因是谷歌模擬了四個 128x128 MXU 的利用率比一個 256x256 MXU 高 60%，但四個 128x128 MXU 占用的面積與 256x256 MXU 相同。MXU 輸入使用 16b 浮點 (FP) 輸入并使用 32b 浮點 (FP) 進行累加。

這些更大的單元允許更有效的數據重用以突破內存墻。

谷歌 DLRM 優化

谷歌是最早開始在其搜索產品中大規模使用 DLRM 的公司之一。這種獨特的需求導致了一個非常獨特的解決方案。上述架構的主要缺陷在于它無法有效處理 DLRM 的嵌入。Google 的主要 TensorCore 非常大，與這些嵌入的計算配置文件不匹配。谷歌必須在他們的 TPU 中開發一種全新類型的“SparseCore”，它不同于上面描述的用于密集層的“TensorCore”

SparseCore (SC) 為 Google 的 TPU 中的嵌入提供硬件支持。早在 TPU v2 中，這些特定領域的處理器就有直接綁定到每個 HBM 通道/子通道的塊。它們加速了訓練深度學習推薦模型 (DLRM) 中內存帶寬最密集的部分，同時僅占用大約 5% 的芯片面積和功率。通過在每個 TPU v4 芯片而非 CPU 上使用快速的 HBM2 進行嵌入，與將嵌入留在主機 CPU 的主內存上相比，谷歌展示了其內部生產 DLRM 的 7 倍加速（TPU v4 SparseCore vs TPU v4 Embeddings on Skylake-SP )。

SparseCore 支持從 HBM 進行快速內存訪問，使用專用的獲取、處理和刷新單元將數據移動到稀疏向量內存 (Spmem) 的組，并由可編程的 8 寬 SIMD 向量處理單元 (scVPU) 更新。這些單元的 16 個計算塊進入一個 SparseCore。

額外的跨通道單元執行特定的嵌入操作（DMA、排序、稀疏減少、分叉、連接）。每個 TPU v4 芯片有 4 個 SparseCore，每個有 2.5MB 的 Smem。展望未來，我們推測由于 HBM3 上子通道數量的增加，TPUv5 的 SparseCores 數量將繼續增加到 6，tiles數量將增加到 32。

雖然遷移到 HBM 帶來的性能提升是巨大的，但性能擴展仍然受到互連對分帶寬的影響。TPU v4 中 ICI 的新 3D 環面有助于進一步擴展嵌入查找性能。然而，當擴展到 1024 個芯片時，由于 SparseCore 開銷成為瓶頸，改進會下降。

如果谷歌認為他們的 DLRM 需要增加超過約 512 個芯片的大小和容量，這個瓶頸可能會導致每個圖塊的 Smem 也隨著 TPUv5 增加。

編輯：黃飛