昨夜晚間，谷歌突然對外披露公司新一代 AI 加速器Cloud TPU v5e。

【資料圖】

谷歌在博客中指出，我們正處于計算領(lǐng)域千載難逢的拐點(diǎn)。設計和構建計算基礎設施的傳統方法不再足以滿(mǎn)足生成式人工智能和LLM等工作負載呈指數級增長(cháng)的需求。事實(shí)上，過(guò)去五年里L(fēng)LM的參數數量每年增加 10 倍。因此，客戶(hù)需要具有成本效益且可擴展的人工智能優(yōu)化基礎設施。

Google Cloud 則通過(guò)提供領(lǐng)先的人工智能基礎設施技術(shù)、TPU 和 GPU，以滿(mǎn)足開(kāi)發(fā)者的需求，今天，谷歌宣布，公司在這兩個(gè)產(chǎn)品組合均得到重大增強。首先，我們正在通過(guò) Cloud TPU v5e 擴展我們的 AI 優(yōu)化基礎設施產(chǎn)品組合，這是迄今為止最具成本效益、多功能且可擴展的 Cloud TPU，現已提供預覽版。TPU v5e 提供與 Google Kubernetes Engine (GKE)、Vertex AI 以及 Pytorch、JAX 和 TensorFlow 等領(lǐng)先框架的集成，因此您可以通過(guò)易于使用、熟悉的界面開(kāi)始使用。

同時(shí)，谷歌還宣布，公司基于 NVIDIA H100 GPU的GPU 超級計算機 A3 VM將于下個(gè)月全面上市，為您的大規模 AI 模型提供支持。、

TPU v5e，為大模型而生

谷歌表示，Cloud TPU v5e是Google Cloud 最新一代 AI 加速器，專(zhuān)為提供中大規模訓練和推理所需的成本效益和性能而構建。雖然谷歌并沒(méi)披露更多關(guān)于這個(gè)芯片的工藝，但據猜測，這可能是基于5nm打造的。

了解谷歌TPU 芯片的讀者都知道，在之前的產(chǎn)品中，他們都有一個(gè)或兩個(gè) TensorCore 來(lái)運行矩陣乘法。與 v2 和 v3 Pod 類(lèi)似，v5e 每個(gè)芯片都有一個(gè) TensorCore。每個(gè) TensorCore 有 4 個(gè)矩陣乘法單元 (MXU)、一個(gè)向量單元和一個(gè)標量單元。下表顯示了 v5e 的關(guān)鍵規格及其值。相比之下，v4 Pod 每個(gè)芯片有 2 個(gè) TensorCore。

與 Cloud TPU v4 相比，TPU v5e 可為L(cháng)LM和新一代 AI 模型提供高達 2 倍的訓練性能和高達 2.5 倍的推理性能。TPU v5e 的成本不到 TPU v4 的一半，使更多組織能夠訓練和部署更大、更復雜的 AI 模型。

谷歌指出，v5e 每個(gè) Pod 的芯片占用空間較小，為 256 個(gè)，經(jīng)過(guò)優(yōu)化，成為transformer、文本到圖像和卷積神經(jīng)網(wǎng)絡(luò ) (CNN) 訓練、微調和服務(wù)的最高價(jià)值產(chǎn)品。

換而言之，通過(guò) TPU v5e Pod 平衡性能、靈活性和效率，允許最多 256 個(gè)芯片互連，總帶寬超過(guò) 400 Tb/s 和 100 petaOps INT8 性能。TPU v5e 的用途也非常廣泛，支持八種不同的虛擬機 (VM) 配置，范圍從一個(gè)芯片到單個(gè)片內的 250 多個(gè)芯片。這使得客戶(hù)能夠選擇正確的配置來(lái)服務(wù)于各種 LLM 和 gen AI 模型大小。

與此同時(shí)，Cloud TPU v5e 還為領(lǐng)先的 AI 框架（如 JAX、PyTorch 和 TensorFlow）以及流行的開(kāi)源工具（如 Hugging Face 的 Transformers 和 Accelerate、PyTorch Lightning 和 Ray）提供內置支持。谷歌透露，公司將通過(guò)即將發(fā)布的 PyTorch/XLA 2.1 版本進(jìn)一步加強對 Pytorch 的支持，其中包括 Cloud TPU v5e 支持，以及用于大規模模型訓練的模型和數據并行性等新功能。

據介紹，在公共預覽版，推理客戶(hù)可以使用 1 芯片 (v5litepod-1)、4 芯片 (v5litepod-4) 和 8 芯片 (v5litepod-8)?！盀榱藶榫哂懈嘈酒奶摂M機騰出空間，調度程序可能會(huì )搶占具有較少芯片的虛擬機。因此，8 芯片虛擬機很可能會(huì )搶占 1 芯片和 4 芯片虛擬機?！惫雀璞硎?。

他們進(jìn)一步指出，為了更輕松地擴展訓練作業(yè)，谷歌還在預覽版中引入了 Multislice 技術(shù)，該技術(shù)允許用戶(hù)輕松擴展 AI 模型，使其超出物理 TPU pod 的范圍，最多可容納數萬(wàn)個(gè) Cloud TPU v5e 或 TPU v4 芯片。

谷歌在博客中寫(xiě)道，到目前為止，使用 TPU 的訓練作業(yè)僅限于單個(gè) TPU 芯片slice，TPU v4 的最大slice大小為 3,072 個(gè)芯片。借助 Multislice，開(kāi)發(fā)人員可以在單個(gè) Pod 內通過(guò)芯片間互連 (ICI)：inter-chip interconnect 或通過(guò)數據中心網(wǎng)絡(luò ) (DCN：Data center network) 跨多個(gè) Pod 將工作負載擴展到數萬(wàn)個(gè)芯片。Multislice 技術(shù)同時(shí)還為谷歌最先進(jìn)的 PaLM 模型的創(chuàng )建提供了動(dòng)力。

“我們真正致力于使其成為一個(gè)可擴展的解決方案，”谷歌的Lohmeyer說(shuō)?！拔覀兛畿浖陀布M(jìn)行設計。在這種情況下，該軟件[和]硬件與Multislice等新軟件技術(shù)協(xié)同工作的魔力，使我們的客戶(hù)能夠輕松擴展其 AI 模型，超越單個(gè) TPU pod 或單個(gè) GPU 集群的物理邊界，”他解釋道?！皳Q句話(huà)說(shuō)，單個(gè)大型人工智能工作負載現在可以跨越多個(gè)物理 TPU 集群，擴展到數萬(wàn)個(gè)芯片，而且這樣做非常經(jīng)濟高效?！盠ohmeye強調。

Jeff Dean：AI 硬件性能的下一個(gè)100倍將會(huì )更加困難

十年前，Jeff Dean 在一張紙上做了一些數學(xué)計算，并發(fā)現：如果谷歌將人工智能功能添加到其同名搜索引擎中，它就必須將其數據中心占地面積擴大一倍，這就趨勢谷歌走上創(chuàng )建定制張量處理單元（TPU）矩陣數學(xué)引擎的道路。

十年過(guò)去了，人工智能變得更加復雜、計算密集，備受討論的 TPUv4 鐵雖然現在和未來(lái)很多年都很有用，但看起來(lái)有點(diǎn)過(guò)時(shí)了。全新推出的 TPUv5e將替代TPUv4 系統。

在日前開(kāi)幕的Hotchips演講中，Jeff Dean表示，谷歌專(zhuān)注于驅動(dòng)人工智能模型的三種不同方法——稀疏性、自適應計算和動(dòng)態(tài)神經(jīng)網(wǎng)絡(luò )——并且還試圖讓人工智能蛇吃掉它的尾巴而不是蠶食它。

在Jeff Dean看來(lái)，真正讓人工智能專(zhuān)家系統開(kāi)始設計人工智能處理器可加快整個(gè)芯片開(kāi)發(fā)周期，從而幫助不斷改進(jìn)的硬件進(jìn)入該領(lǐng)域，以滿(mǎn)足更快增長(cháng)的模型。

Dean 解釋說(shuō)，到目前為止創(chuàng )建的人工智能模型，整個(gè)模型的層數不斷增加，參數數量呈爆炸式增長(cháng)，由數十億、數百億、數千億的token snippets 數據驅動(dòng)，每次人工智能都會(huì )被激活。在新token上訓練的模型或針對已完成的模型提出token以進(jìn)行人工智能推理。

但是，有了像 Pathways 這樣的框架（它是 Google PaLM 系列模型的基礎），世界正在從專(zhuān)門(mén)針對不同任務(wù)的單獨 AI 模型轉向擁有單一基礎模型。

對于稀疏模型，AI 模型的各個(gè)部分會(huì )在需要時(shí)激活（activate ），并且僅激活這些部分。模型如何知道要激活哪些部分尚不清楚，這就是 Pathways 框架中的秘密武器，該框架已通過(guò) Gemini 模型得到完善，毫無(wú)疑問(wèn)使用了 Dean 所說(shuō)的技術(shù)。

值得注意的是，Pathways 框架不像 Google 創(chuàng )建的早期且可能更為初級的 TensorFlow 框架（該框架于 2015 年 11 月開(kāi)源）那樣是開(kāi)源的。因此，我們只能知道 Google 告訴我們的有關(guān) Pathways 和 Gemini 的信息。

“與這個(gè)巨大的模型相比，稀疏模型的效率要高得多，”Dean解釋道?！八麄冎皇钦{用整個(gè)模型的正確部分——正確的部分也是在訓練過(guò)程中學(xué)到的東西。然后模型的不同部分可以專(zhuān)門(mén)用于不同類(lèi)型的輸入。最終的結果是，您最終會(huì )得到一些非常大的模型的正確 1% 或正確 10% 的結果，這會(huì )提高您的響應能力和更高的準確性，因為您現在擁有比您大得多的模型容量可以進(jìn)行其他訓練，然后可以調用正確的部分?！?/p>

根據 Dean 的說(shuō)法，稀疏性的另一個(gè)方面對于系統架構師來(lái)說(shuō)很重要，它與加速器中通常談?wù)摰募毩６认∈栊圆煌?，在加速器中，單個(gè)向量或張量?jì)鹊南∈栊裕ㄍǔＣ克膫€(gè)中的兩個(gè)）矩陣中的值設置為零，將其從密集型轉換為稀疏型），這也不同于粗粒度稀疏性，其中模型中的大型模塊要么被激活，要么不被激活。這種稀疏性看起來(lái)像這樣，我們將幾個(gè) Dean 的圖表合并到一頁(yè)上，這樣您就可以將其全部理解：

“當今大多數稀疏工作對每個(gè)專(zhuān)家都使用相同的大小和結構，”Dean說(shuō)?！八赃@里有一些綠色專(zhuān)家為他們服務(wù)。這里有一些學(xué)習路由功能，可以了解哪個(gè)專(zhuān)家擅長(cháng)哪種事情，然后將一些示例發(fā)送給適當的專(zhuān)家。計算平衡通常是通過(guò)每個(gè)專(zhuān)家的計算量相等以及每個(gè)專(zhuān)家的示例數量相等的流量來(lái)實(shí)現的。對于計算機架構師來(lái)說(shuō)，這意味著(zhù)跨加速器的全面混洗性能非常重要。對于所有稀疏模型來(lái)說(shuō)都是如此——您希望能夠以正確的方式將事物從模型的一個(gè)部分快速路由到另一部分?！?/p>

“不過(guò)，您可能想做的一件事是，不是具有固定的計算成本，而是改變模型不同部分的計算成本。在每個(gè)示例上花費相同的計算量是沒(méi)有意義的，因為有些示例的難度是原來(lái)的 100 倍。我們在真正困難的事情上花費的計算量應該是非常簡(jiǎn)單的事情上的 100 倍?！盌ean強調。

事實(shí)證明，一些小型專(zhuān)家可能只需要少量計算，并且可以用于生產(chǎn)中使用的模型中大約 90% 的提示。專(zhuān)家們?yōu)榱颂幚砀鼜碗s的事情而變得更大，具有不同的計算結構和可能更多的層，并且它們的計算量更大，因此運行成本更高。如果您正在運行人工智能服務(wù)，您將希望能夠將成本歸因于所提供的專(zhuān)家答案的價(jià)值，以便您可以適當收費。

提高效率和計算能力的途徑

當然，這不是谷歌的獨家理論——該公司之所以談?wù)撍且驗?Pathways 框架這樣做：

這就是稀疏性和自適應計算。

Dean 說(shuō)，上圖中提到的最后一件事是動(dòng)態(tài)神經(jīng)網(wǎng)絡(luò )，這意味著(zhù)可以在正在運行的系統中添加或刪除容量——這是我們幾十年來(lái)通用服務(wù)器所擁有的（雖然不是在 X86 平臺上，但奇怪的是，這就是 Arm 和 RISC-V 可能能夠趕上大型機和 RISC/Unix 系統的地方）。CPU 及其工作負載的情況（虛擬機管理程序級別肯定存在動(dòng)態(tài)分配）對于 GPU、TPU 和其他 AI 計算引擎來(lái)說(shuō)也是如此。您希望能夠在運行推理或訓練時(shí)動(dòng)態(tài)地為任何給定模型添加或減少核心池的容量。

來(lái)自 Google 的具有 5000 億個(gè)參數的 PaLM 模型在 Pathways 上進(jìn)行了訓練，并通過(guò)在具有 6,144 個(gè) TPUv4 引擎的一對 pod 之間動(dòng)態(tài)分配資源來(lái)實(shí)現這一目標，但 TPUv4 引擎實(shí)際上分布在總共 24,576 個(gè)引擎的 6 個(gè) pod 中，所有這些都鏈接在一起通過(guò)高速數據中心網(wǎng)絡(luò )。像這樣：

以下是 Dean 希望給系統架構師留下深刻印象的關(guān)鍵要點(diǎn)：

加速器的連接性（帶寬和延遲）很重要

規模對于訓練和推理都很重要

稀疏模型給內存容量和高效路由帶來(lái)壓力

機器學(xué)習軟件必須能夠輕松表達有趣的模型 - 就像上圖中所示的函數稀疏性

功率、可持續性和可靠性確實(shí)很重要

Google Fellow以及工程副總裁Amin Vahdat在隨后的演講中展示了人工智能行業(yè)面臨的模型大小增長(cháng)的指數曲線(xiàn)：

據觀(guān)察，我們絕對沒(méi)有理由相信模型的復雜性會(huì )因此而降低，計算能力的需求將會(huì )降低。但根據我們的估計，模型每年增長(cháng) 10 倍，GPU 和 TPU 的性能最多每年增長(cháng) 2 到 3 倍。公司必須通過(guò)擴展來(lái)彌補這一點(diǎn)，這很困難，并改進(jìn)他們的模型，這也很困難。

這就是為什么谷歌已經(jīng)在其機群中部署了 TPUv5e 引擎。在我們看來(lái)，谷歌的 TPUv6 可能正在開(kāi)發(fā)中并準備就緒很快就會(huì )部署，以幫助支持Gemini模型的商業(yè)化。

為了使每 TCO 的性能提高 100 倍，Vahdat進(jìn)行了一個(gè)深入的講座，介紹了如何衡量人工智能或通用計算平臺的相對價(jià)值，在人工智能出現之前，我們一直同意這一點(diǎn)系統——谷歌必須做很多事情，當中包括：

創(chuàng )建專(zhuān)門(mén)的硬件 - TPU - 用于密集矩陣乘法。

使用 HBM 內存將這些矩陣數學(xué)引擎的內存帶寬提高 10 倍。

創(chuàng )建專(zhuān)門(mén)的硬件加速器，用于稀疏矩陣中的分散/聚集操作 - 我們現在稱(chēng)之為 Sparsecore，它嵌入在 TPUv4i、TPUv4 和可能的 TPUv5e 引擎中。

采用液體冷卻可最大限度地提高系統電源效率，從而提高經(jīng)濟效益。

使用混合精度和專(zhuān)門(mén)的數字表示來(lái)提高設備的實(shí)際吞吐量（Vahdat 稱(chēng)之為“有效吞吐量”）。

并具有用于參數分配的同步、高帶寬互連，事實(shí)證明，這是一種光路開(kāi)關(guān)，可以在系統上的作業(yè)發(fā)生變化時(shí)幾乎即時(shí)重新配置網(wǎng)絡(luò )，并且還提高了機器的容錯能力。對于一個(gè)擁有數萬(wàn)個(gè)計算引擎且工作負載需要數月才能運行的系統來(lái)說(shuō)，這是一件大事，世界各地的 HPC 中心都非常清楚這一點(diǎn)。

“我們?yōu)閼獙@一挑戰而必須構建的計算基礎設施類(lèi)型必須改變，”Vahdat在主題演講中說(shuō)道?！拔艺J為，值得注意的是，如果我們試圖在通用計算上做到這一點(diǎn)，我們就不會(huì )達到今天的水平。換句話(huà)說(shuō)，我們在過(guò)去 50 到 60 年間所形成的傳統智慧實(shí)際上已經(jīng)被拋到了九霄云外。我認為可以公平地說(shuō)，在谷歌，但更重要的是，在整個(gè)社區，機器學(xué)習周期將占據我們想做的事情中越來(lái)越多的部分?！?/p>

谷歌專(zhuān)注于優(yōu)化硬件和軟件以跨系統集群動(dòng)態(tài)管理工作負載和功耗的一件事：

對于受內存限制的部分，電壓和電流強度可能會(huì )有很大差異，正如 Vahdat 所說(shuō)，試圖管理數千到數萬(wàn)個(gè)計算引擎集群的功耗是“介于困難和不可能之間的事情”。通過(guò)不在集群中創(chuàng )建大量熱點(diǎn)（這可能發(fā)生在 Google 訓練 PaLM 模型時(shí)），可以延長(cháng)設備的使用壽命并減少中斷，這對于 AI 訓練等同步工作非常具有破壞性，就像 HPC 模擬一樣和建模。與其回滾到檢查點(diǎn)并從那里開(kāi)始，不如從一開(kāi)始就避免中斷。

以下是如何利用核心頻率和電壓來(lái)平衡一些事情。

當然，快速更新TPU芯片是谷歌工作最重要的一部分。

Dean表示，目前芯片的問(wèn)世大約需要三年時(shí)間。這需要六到十二個(gè)月的時(shí)間進(jìn)行設計和探索，一年的時(shí)間來(lái)實(shí)施設計，六個(gè)月的時(shí)間與晶圓進(jìn)行流片，十二個(gè)月的時(shí)間用于將其投入生產(chǎn)、測試和提升。

目前尚不清楚人工智能可以在多大程度上縮短芯片開(kāi)發(fā)周期或可以減少多少人力，Dean也沒(méi)有提供任何估計。但顯然，硬件設計越接近新興人工智能模型就越好。