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

【資料圖】

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

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

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

TPU v5e，為大模型而生

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

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

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

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

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

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

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

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

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

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

Jeff Dean：AI 硬件性能的下一個100倍將會更加困難

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

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

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

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

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

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

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

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

“與這個巨大的模型相比，稀疏模型的效率要高得多，”Dean解釋道?！八麄冎皇钦{(diào)用整個模型的正確部分——正確的部分也是在訓(xùn)練過程中學(xué)到的東西。然后模型的不同部分可以專門用于不同類型的輸入。最終的結(jié)果是，您最終會得到一些非常大的模型的正確 1% 或正確 10% 的結(jié)果，這會提高您的響應(yīng)能力和更高的準(zhǔn)確性，因?yàn)槟F(xiàn)在擁有比您大得多的模型容量可以進(jìn)行其他訓(xùn)練，然后可以調(diào)用正確的部分。”

根據(jù) Dean 的說法，稀疏性的另一個方面對于系統(tǒng)架構(gòu)師來說很重要，它與加速器中通常談?wù)摰募?xì)粒度稀疏性不同，在加速器中，單個向量或張量內(nèi)的稀疏性（通常每四個中的兩個）矩陣中的值設(shè)置為零，將其從密集型轉(zhuǎn)換為稀疏型），這也不同于粗粒度稀疏性，其中模型中的大型模塊要么被激活，要么不被激活。這種稀疏性看起來像這樣，我們將幾個 Dean 的圖表合并到一頁上，這樣您就可以將其全部理解：

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

“不過，您可能想做的一件事是，不是具有固定的計(jì)算成本，而是改變模型不同部分的計(jì)算成本。在每個示例上花費(fèi)相同的計(jì)算量是沒有意義的，因?yàn)橛行┦纠碾y度是原來的 100 倍。我們在真正困難的事情上花費(fèi)的計(jì)算量應(yīng)該是非常簡單的事情上的 100 倍。”Dean強(qiáng)調(diào)。

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

提高效率和計(jì)算能力的途徑

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

這就是稀疏性和自適應(yīng)計(jì)算。

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

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

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

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

規(guī)模對于訓(xùn)練和推理都很重要

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

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

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

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

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

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

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

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

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

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

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

使用混合精度和專門的數(shù)字表示來提高設(shè)備的實(shí)際吞吐量（Vahdat 稱之為“有效吞吐量”）。

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

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

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

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

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

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

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

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