芯東西（公眾號：aichip001）
作者 | ZeR0
編輯 | 漠影
芯東西8月30日報道，昨晚，在時隔四年首次回歸線下的谷歌雲年度大會Next23上，谷歌雲發布最新雲端AI晶片TPU v5e，並罕見放出一段2分39秒的視頻，披露了其雲計算TPU數據中心內部實景，包括TPU、數據中心網絡、光電路交換、水冷系統、生物識別安全驗證等。
谷歌TPU晶片為加速AI計算而生，是AI晶片中專用集成電路（ASIC）的代表之作，近年一直被視作全球AI晶片霸主英偉達GPU的勁敵。因此谷歌AI晶片的一舉一動都備受業界關注。
據半導體研究和諮詢公司SemiAnalysis的分析師Dylan Patel和Daniel Nishball爆料，谷歌擁有的算力資源比OpenAI、Meta、亞馬遜、甲骨文和CoreWeave加起來還要多，其下一代大模型Gemini已經開始在新的TPUv5 Pod上進行訓練，算力達到GPT-4的5倍，基於其目前的基礎設施建設情況，到明年年底可能達到20倍。
值得一提的是，谷歌相信AI工作負載沒有唯一的標準，因此不只用自研AI晶片，還宣布跟英偉達擴大合作夥伴關係，為全球谷歌雲客戶提供基於英偉達GPU和生成式AI技術的雲服務。英偉達創始人兼CEO黃仁勛也特意來到現場為谷歌雲站台，與谷歌雲CEO托馬斯·庫里安（Thomas Kurian）進行爐邊對話。
谷歌雲CEO托馬斯·庫里安（左）與英偉達創始人兼CEO黃仁勛（右）同台（圖源：英偉達）
谷歌雲將於下個月全面上市的A3虛擬機，配備8張英偉達H100 GPU，能將訓練AI模型的效率提高3倍，將網絡帶寬分配效率提高10倍。此外，谷歌雲不僅是全球首家提供英偉達L4 GPU的雲供應商，還是全球首批獲得英偉達DGX GH200 AI超級計算機訪問權限的公司之一。
一、TPU v5e：單位成本訓練性能翻倍，跑大模型和生成式AI能效高
上次谷歌發布旗艦雲端AI晶片TPU v4，還是在2021年5月的谷歌IO開發者大會上，開發者直到去年才真正用上TPU v4。
業界對第五代旗艦AI晶片TPU v5的出場翹首已久。沒想到谷歌不按套路出牌，這次先亮出了為大規模、中等規模AI訓練和推理而打造的優化版晶片TPU v5e。
TPU v5e的系統架構
這是一款訓推一體產品，每顆TPU v5e晶片包含1個張量核心，每個張量核心有4個矩陣乘法單元（MXU）、1個向量單元和1個標量單元。
谷歌雲將它稱作是迄今最具成本效益、多功能和可擴展的TPU，能確保用戶「將其TPU集群擴展到以前無法達到的水平」。
與Cloud TPU v4相比，TPU v5e每美元可提供高達2倍的訓練性能，對於大型語言模型和生成式AI模型，每美元可提供2.5倍的推理性能，而成本不到TPU v4的一半。
谷歌通過TPU v5e Pod平衡性能、靈活性和效率。TPU v5e支持多達256個晶片互連，總帶寬超過400Tb/s，INT8性能達到100petaOps，從而解決更複雜的計算任務。
TPU v5e的關鍵規格
TPU v5e現已推出預覽版，將支持8種不同的虛擬機配置，從一個TPU晶片到單個切片中超過250個TPU晶片。谷歌雲客戶可按需選擇配置，以服務於不同規模的大模型和生成式AI模型。
TPU v5e支持的2D切片情況
TPU切片中，每個TPU虛擬機包含1、4或8個晶片，4晶片或者更少晶片將具有相同的非均勻內存訪問（NUMA）節點。對於8晶片的TPU虛擬機，CPU-TPU通信在NUMA分區內將更加高效。例如在下圖中，CPU0-Chip0的通信比CPU0-Chip4的通信快。
AssemblyAI技術副總裁Domenic Donato說，在他們的生產ASR模型上運行推理時，Cloud TPU v5e每美元的性能始終比市場上同類解決方案高出4倍。
Gridspace機器學習主管Wonkyum Lee亦談道，其速度基準顯示，當在谷歌Cloud TPU v5e上訓練和運行時，AI模型的速度提高了5倍。推理指標規模也有顯著改進，他們現在可以在1秒內實時處理1000秒的內部語音到文本和情感預測模型，處理速度提高了6倍。
二、引入Multislice優化橫向擴展，提高AI模型開發效率
從歷史來看，通過橫向擴展基礎設施協調大規模AI工作負載需要手動處理故障、日誌記錄、監控和其他基本操作。
谷歌正著力簡化TPU的操作，TPU v5e提供了與Google Kubernetes Engine（GKE）、Vertex AI、PyTorch、JAX、TensorFlow等主流框架的集成，並對各種主流開源工具提供內置支持，方便開發者使用熟悉的介面。
以知名生成式AI獨角獸企業Anthropic為例，據Anthropic聯合創始人Tom Brown分享，GKE使其能大規模運行和優化GPU和TPU基礎設施，Vertex AI使其能通過Vertex AI模型商店向客戶分發模型。此外，他提到由A3和帶有Multislice的TPU v5e驅動的谷歌下一代AI基礎設施，將為Anthropic的工作負載帶來性價比優勢。
為了更容易地擴展訓練工作，谷歌在預覽版中引入了Multislice技術，這是一種將AI模型分配給數萬個TPU晶片的方法，已經為谷歌最先進的PaLM模型的構建提供了動力。
該技術能夠超越物理TPU Pod的邊界，通過單個Pod內的晶片間互連（ICI）或數據中心網絡（DCN）上的多個Pod，將工作負載擴展到多達數萬個TPU v5e或TPU v4晶片。
截至目前，使用TPU的訓練任務僅限於單個TPU晶片切片，TPU v4的最大切片大小限制為3072個晶片。
三、為什麼谷歌要在機群中部署TPU v5e？
恰逢國際晶片頂會Hot Chips 2023正在舉行，兩位谷歌Fellow——谷歌首席科學家Jeff Dean和谷歌網絡技術主管Amin Vahdat——共同發表了主題為《機器學習模型的激動人心的方向和對計算硬體的影響》的演講。
Vahdat分享了為什麼谷歌在其機群中部署了TPU v5e引擎，以及為什麼可能帶有i或e字母版本的TPU v6正在開發中、很快就會部署。
過去五年，AI模型變得越來越複雜，大型語言模型的參數數量以每年10倍的速度增長。而GPU和TPU的性能每年最多增長2~3倍，單晶片性能增長根本撐不住，所以必須通過集群擴展和稀疏化模型來應對飆漲的算力需求。但Vahdat認為這兩種方法很快就會遇到瓶頸。
正如谷歌在博客文章中所寫，我們正處於計算機行業千載難逢的拐點。傳統的設計和構建計算基礎設施的方式已無法滿足生成式AI和大型語言模型等工作負載的指數級增長需求。客戶需要既具有成本效益又具有可擴展性的AI優化基礎設施，同時還不能犧牲掉性能或靈活性。
為了將每TCO性能提高100倍，谷歌必須做很多事情：
1、創建專門的硬體TPU，用於密集矩陣乘法。矩陣乘法單元等功能可優化大型矩陣運算。
2、採用HBM內存，將這些矩陣數學引擎的內存帶寬提高10倍。
3、創建專門的硬體加速器，用於稀疏矩陣中的scatter/gather操作，這被稱作Sparsecore，它嵌入在TPU v4i、TPU v4和可能的TPU v5e引擎中。
4、採用液冷可最大限度地提高系統電源效率，從而提高經濟效益。
谷歌輸送給超算的垂直電源，由其最新一代水冷系統冷卻，該系統設計考慮了可持續性
5、使用混合精度和專門的數字表示，來提高設備的實際有效吞吐量。
6、具有用於參數分配的同步、高帶寬互連，這是一種光電路交換（OCS），支持高帶寬內存（HBM），當系統上的任務發生變化時幾乎能即時重新配置網絡，還提高了機器的容錯能力。對於一個擁有成千上萬個計算引擎且工作負載需運行數月的系統來說，這是一件大事。
谷歌光電路交換技術使超算能進行重新配置，並為工作負載創建高效的機器形態和規格，按需動態擴展或縮小，無需人工參與
為了應對這些挑戰，構建的計算基礎設施類型必須改變。Vahdat說：「我們在過去五六十年間所形成的傳統智慧實際上已被拋到了九霄雲外。」
谷歌專注於優化硬體和軟體以跨系統集群動態管理工作負載和功耗：
對於內存受限的任務，電壓和電流強度可能會有很大差異，正如Vahdat所說，試圖管理數千到數萬個計算引擎集群的功耗是「介於困難和不可能之間的事情」。
通過不在集群中創建大量熱點，可延長設備的使用壽命並減少中斷，這對於AI訓練等同步工作非常具有破壞性。與其回滾到檢查點並從那裡開始，不如從一開始就避免中斷。
以下是如何利用核心頻率和電壓來平衡一些事情。在集群上查看任務放置的控制平面與這些任務的電源參數之間存在持續的干擾，並且這些任務在運行時的放置和移動是一個持續的過程。
四、用AI輔助設計AI晶片，三個方法加速AI計算
Dean在演講中分享的內容，跟TPU v5e相關度較少，但有助於為加速AI模型提供啟發。
他重點分享了谷歌專注於稀疏性、自適應計算、動態神經網絡這3種方法來驅動AI模型，並試圖真正讓AI專家系統開始設計AI晶片，以加快整個晶片開發周期、幫助優化AI硬體。
為了更快推出新的TPU晶片，谷歌已經使用其AI增強型EDA工具來幫助設計TPU v4i和TPU v4晶片的邏輯塊，很可能還包括TPU v5e。
據Dean分享，目前要花大約三年才能推出晶片，6~12個月用於設計和探索，1年用於實施設計，6個月用於將其放到晶圓廠進行流片上，12個月用於投入生產、測試和提升。顯然，硬體設計越接近新興的AI模型越好。
Dean希望讓系統架構師記住這些關鍵要點：
1、加速器的連接性（帶寬和延遲）很重要。
2、擴展對於訓練和推理都很重要。
3、稀疏模型給內存容量和高效路由帶來壓力。
4、機器學習軟體必須能夠輕鬆表達有趣的模型，例如函數稀疏性。
5、功率、可持續性和可靠性確實很重要。
他談道，AI模型的層數不斷增加，參數數量呈爆炸式增長，由數十億、數百億、數千億個tokens數據片段驅動。每次AI模型在一個新token上訓練或在一個已完成的模型呈現一個新token來進行AI推理時，整個模型都會被激活。
有了像Pathways這樣的框架，世界正在從針對不同任務的單獨AI模型轉向擁有能解決數百萬任務的單一基礎模型。
擁有5000億個參數的谷歌PaLM模型，便是在Pathways上進行訓練的，通過在有6144個TPU v4引擎的2個Pod之間動態分配資源實現，TPU v4實際上分布在總共24576個引擎的6個Pod中。PaLM花了56天進行訓練，下圖是第5.71天的快照。
谷歌Pathways框架採用稀疏性和自適應計算，在運行推理或訓練時動態地為給定模型添加或減少核心池的容量。該框架的一個秘密武器，是知道應該激活稀疏模型的哪些部分，它已經通過谷歌正在秘密籌備的「殺手級」大模型項目Gemini得到完善。不過谷歌沒將Pathways框架開源。
以稀疏方式激活模型，能夠大大降低進一步訓練和生產推理所需的計算成本。
這裡談到的稀疏性，跟通常說的單個向量或張量內稀疏性的加速器細粒度稀疏性不是一回事。它也不同於粗粒度稀疏性，模型中的大型模塊要麼被激活，要麼不被激活。為了直觀反映Dean說的稀疏性，外媒The Next Platform將Dean的幾張圖拼到一張圖中：
據Dean解釋，當前大多數稀疏化工作對每個專家都用相同的尺寸和結構，所以有一些綠色的專家，通過學習一些路由函數，可以了解哪個專家擅長做哪種事，然後將一些例子發送給合適的專家。計算平衡通常是通過每個專家的計算量相等以及每個專家的示例數量相等來實現的。
對於計算機架構師來說，這意味著跨加速器的全對全shuffle性能非常重要，對於所有稀疏模型都是成立的。不過你可能想能夠做的，不是用固定的計算成本，而是改變模型不同部分的計算成本。「我們應該在真正困難的事情上花費的計算量，應該是非常簡單的事情上花費計算量的100倍。」Dean說。
結語：拉大AI基礎設施性價比優勢， 在雲計算競賽中越來越關鍵
二十年來，從構建為如今大模型盛世奠定基礎的谷歌Transformer架構，到自研TPU晶片、構建為YouTube、Gmail、谷歌地圖、谷歌Play和Android等數十億用戶提供服務的AI優化基礎設施，谷歌已建立一系列行業領先的AI能力。
谷歌雲人工智慧業務和解決方案副總裁Phil Moyer稱，Anthropic、Character AI、Typeface等生成式AI獨角獸創企中，70%以上都是谷歌雲的客戶；獲得總投資的AI初創公司中，有1/4正在採用谷歌雲。
但這顯然遠遠不夠，如今生成式AI和大模型已經成為吸引雲客戶的一大關鍵因素，要應對亞馬遜雲科技、微軟等強勁雲計算對手的猛烈攻勢，谷歌不僅要繼續完善支持模型訓練、調優和全球範圍內服務的整個生命周期的端到端軟體和服務，而且需不斷提高其AI基礎設施的性價比優勢。