最近學習了論文《TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings》。 谷歌的TPU v4（張量處理單元第四代）是專門為機器學習（ML）尤其是深度學習模型訓練設計的加速器，主要在大規模計算中應用，如大規模語言模型（LLM）、深度學習推薦模型（DLRM）以及其他需要高吞吐量、低延遲的計算任務。TPU v4針對前代TPU的不足，結合了新的硬件架構和先進的互聯技術，顯著提高了性能、可擴展性和可靠性。

1. TPU v4的三大核心創新

TPU v4引入了三大創新特性，以應對大規模機器學習計算中的關鍵挑戰：

光學電路交換（OCS）與光纖數據鏈路

TPU v4采用了光學電路交換（Optical Circuit Switches，OCS）技術，用光纖連接來實現超大規模系統的高效通信。通過這種光學交換方式，TPU v4能夠構建一個最大支持4000個TPU節點的超級計算機，同時保持良好的可擴展性和可靠性。與傳統電子包交換（如InfiniBand）相比，光學交換不僅能提高系統的吞吐量，還能減少功耗，且由于其高帶寬和低延遲的特性，能有效支撐深度神經網絡（DNN）訓練中的全對全通信模式（All-to-All communication），這對于處理大規模數據特別重要。

SparseCore硬件支持

SparseCore是TPU v4的一個新特性，專門優化了稀疏嵌入（embedding）操作，特別是對于深度學習推薦模型（DLRM）等需要大量稀疏矩陣計算的任務。SparseCore提供了一種數據流海量核心架構，允許將嵌入數據靈活地分配到TPU v4超級計算機的128TiB物理內存中，這種架構極大提升了稀疏計算的效率，能夠比傳統方法快5到7倍。

靈活的拓撲重配置

TPU v4在硬件拓撲上引入了光學交換的靈活性，能夠根據實際應用需求調整拓撲結構（如3D環形拓撲或扭曲環形拓撲），優化互連帶寬的使用。這種拓撲重配置能力使得TPU v4能根據不同計算任務的需求，選擇最適合的拓撲結構，從而大幅提高了性能并減少了通信瓶頸。

2. TPU v4的關鍵優勢 (1)可擴展性和可靠性

通過光學電路交換（OCS），TPU v4能夠支持超過4000個TPU節點的超級計算機規模，而系統能容忍1K個CPU主機不可用的情況。傳統的分布式系統通常在硬件故障時需要復雜的恢復機制，但TPU v4通過OCS提高了系統的可靠性，使得訓練過程更加穩定。

性能與能效的提升

TPU v4相較于TPU v3，提升了多個方面的性能，尤其在大規模語言模型（LLM）的訓練中，使用3K個TPU v4切片和3D環形拓撲顯著減少了訓練時間。OCS的光學交換技術相較于傳統的電子包交換，在減少功耗的同時，還優化了性能，因此TPU v4在能效方面優于同類的其他加速器，例如Nvidia的A100和Graphcore的MK2 IPU。

稀疏矩陣加速

SparseCore提供的稀疏計算加速功能是TPU v4的另一個亮點。這種硬件支持將稀疏矩陣操作效率提高了5到7倍，對于大規模推薦系統和其他稀疏計算密集型任務，具有極大的優勢。并且，這種加速僅占TPU v4總芯片面積和功耗的5%左右，成本效益非常高。

模塊化設計和拓撲靈活性

TPU v4的設計支持模塊化擴展，可以通過靈活配置拓撲來優化性能。用戶可以根據具體需求，選擇不同的拓撲結構（如3D環形、扭曲環形等），這樣不僅提升了互連帶寬，還能夠減少系統在不同應用場景下的性能瓶頸。

3. TPU v4與前代產品及其他競爭產品的比較

與TPU v3相比，TPU v4在多個方面的性能提升使得其成為訓練大規模模型（如PaLM模型）的理想選擇。PaLM是一個擁有540億參數的大規模語言模型，在TPU v4超級計算機上訓練時，能夠保持接近硬件浮點性能的57.8%。同時，TPU v4在功耗、性價比和計算密度方面，較Nvidia A100和Graphcore MK2 IPU等同類產品具備明顯優勢。

4. TPU v4的能效和環境影響

TPU v4的另一大優勢是其能效。在集成到高效能的倉庫級計算機（WSC）中時，TPU v4能夠顯著減少碳足跡。與傳統的深度學習加速器相比，TPU v4在相似規模的數據中心中，能夠消耗20倍更少的能源，且二氧化碳排放量降低約95%。這一點對于支持可持續的人工智能發展具有重要意義，尤其是在大規模機器學習模型訓練過程中。

5. 總結

TPU v4代表了谷歌在深度學習硬件加速領域的最新突破。通過光學電路交換技術、SparseCore加速和靈活的拓撲配置，TPU v4不僅提高了大規模機器學習模型訓練的性能，還在可靠性、能效和可擴展性方面做出了顯著改進。TPU v4的設計滿足了當今AI領域對于更強計算力、低功耗和高可靠性的需求，特別是在訓練大規模語言模型（LLM）和深度學習推薦模型（DLRM）等任務時，展現出無與倫比的優勢。