上周，Google在論文預印本平台arxiv上發表了其關於TPU v4的深入解讀論文《TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings》（TPU v4：通過光互聯可重配置的機器學習超級計算機，搭載硬體嵌入層加速）。
該論文將於今年六月在ISCA 2023（International Symposium on Computer Architecture，計算機架構領域的頂級會議）上正式發表，而目前的預印本無疑為我們提供了可以一窺其全貌的機會。
為了TPU的可擴展性設計專用光學晶片，Google也是拼了
從論文的標題可以看到，GoogleTPU v4的一個主要亮點是通過光互連實現可重配置和高可擴展性（也即標題中的“optically reconfigurable”）。
而在論文的一開始，Google開門見山首先介紹的也並非傳統的MAC設計、片上記憶體、HBM通道等AI晶片常見的參數，而是可配置的光學互聯開關（reconfigurable optical switch）。
作為論文的重中之重，這里我們也詳細分析一下為什麼光學互聯在TPU v4設計中占了這麼重要的位置，以至於Google甚至為了它自研了一款光學晶片。
TPU v4從一開始設計時，其目標就是極高的可擴展性，可以有數千個晶片同時加速，從而實現一個為了機器學習模型訓練而設計的超級計算機。
在Google的設計中，超級計算機的拓撲結構為：將4x4x4（64）個TPU v4晶片互聯在一起形成一個立方體結構（cube），然後再把4x4x4這樣的cube串聯在一起形成一個總共有4096個TPU v4的超級計算機。
TPU超級計算機（由4096個TPU v4組成）拓撲結構，圖/Google
在這樣的拓撲中，物理距離較近的TPU v4（即在同一個4x4x4 cube中的晶片）可以用常規的電互聯（例如銅絞線）方法連接，但是距離較遠的TPU之間（例如在cube之間的互聯）就必須使用光互連，原因就在於在如此大規模的超級計算機中，晶片之間的數據互聯在很大程度上會決定整體計算的效率；
如果數據互聯效率不夠高的話，很多時候晶片都在等待來自其他晶片的數據到達以後開始計算，這樣就形成了效率浪費。為了避免這樣“晶片等數據”的情形出現，就必須確保晶片之間互聯能擁有高帶寬，低延遲。而光互連對於物理距離較遠的晶片就成為了首選。
光互連在高性能計算中的使用也並非新聞，而Google在TPU v4中的主要突破是使用可重配置的光互連（即加入光路開關，optical circuit switch OCS）來快速實現不同的晶片互聯拓撲。
換句話說，晶片之間的互聯並非一成不變的，而是可以現場可重配置的。這樣做可以帶來許多好處，其中最主要的就是可以根據具體機器學習模型來改變拓撲，以及改善超級計算機的可靠性。
從拓撲結構來說，不同的機器學習模型對於數據流的要求大致可以分為三類，即數據並行（每塊晶片都加載整個模型，不同的晶片處理數據集中不同的數據），模型並行（模型中有些層特別大，因此每塊晶片只負責這樣很大的層中的一部分計算），以及流水線並行（把模型中的不同層交給不同的晶片計算），而不同的數據流就對應了不同的TPU互聯拓撲。
當有了可重配置光互連之後，就可以根據具體模型數據流來調整TPU之間的互聯拓撲，從而實現最優的性能，其提升可超過2倍。
使用可重配置光互連可以快速切換不同的晶片間互聯拓撲，圖/Google
另一個優勢就是可靠性。在這樣擁有海量晶片組成的超級計算機中，一個重要的考量就是，如果有一小部分的晶片不工作了，如何確保整體超級計算機仍然能維持較高的性能？如果使用常規的固定互聯架構，那麼一個晶片出故障可能會影響整個系統工作。
而在有了可重配置的光互連之後，需要做的只需要把出故障的晶片繞過，就不會影響整個系統的工作，最多會犧牲一點整體的性能。
Google在論文中給出了一個單晶片故障率和系統平均性能影響的曲線圖，在使用可配置光互連（以及光路開關）時，假設晶片可靠率在99%的情況下，其整體系統的平均性能提升比不使用OCS可高達6倍，可見光互連開關的重要性。
系統性能與晶片可靠性曲線圖——使用OCS可以大大提升系統性能，圖/Google
為了實現數據中心級的可配置光互連，需要光路開關首先能高效擴展到超高數量的互聯數（例如1000×1000），同時需要實現低開關切換延遲，低成本，以及低信號損耗。
Google認為現有的商用方案都不夠滿意，因此Google的做法是自研了一款光路開關晶片Palomar，並且使用該晶片實現了全球首個數據中心級的可配置光互連，而TPU v4就是搭配了這款自研光路開關晶片從架構上實現了高性能。
Google自研的光路開關晶片Palomar使用的是基於MEMS反射鏡陣列的技術，具體原理是使用一個2D MEMS反射鏡陣列，通過控制反射鏡的位置來調整光路，從而實現光路的切換。
使用MEMS的光路開關晶片可以實現低損耗，低切換延遲（毫秒級別）以及低功耗。在經過一系列優化之後，光路系統的成本也控製得很低，在整個TPU v4超級計算機成本中占5%以下。
Google自研的Palomar MEMS光路開關晶片，圖/Google
算法-晶片協同設計是TPU v4的靈魂
如果說可重配置光互聯給TPU v4提供了良好的根基的話，那麼算法-晶片協同設計就是TPU v4的靈魂。算法-晶片協同設計包括兩部分，一部分是如何根據算法優化晶片，而另一方面是如何根據晶片去優化算法，在TPU v4的架構中，兩者都得到了仔細考慮。
我們首先分析TPU v4如何根據算法來優化晶片。如果說2017年TPU v1發表時候，其主要解決的還是卷積神經網絡CNN的話，那麼在2023年來看，CNN的加速問題早已經被解決得差不多了，更多的是如何處理目前如日中天的大模型的問題。
對於Google來說，目前最關鍵的大模型就是決定了其公司主營收入的推薦系統大模型，因此TPU v4的設計也針對推薦系統大模型做了相當的優化。
在推薦系統大模型中，目前的加速瓶頸是嵌入層（embedding layer）。嵌入層的目的是將高維度稀疏特徵映射到低維度高密度特徵，從而該高密度低維度特徵可以被神經網絡進一步處理。
嵌入層的實現通常是一個查找表（look-up table），而這個查找表可以非常巨大至100GB的數量級。在一個推薦系統模型中可以有多個這樣的查找表，從而讓整個查找表的存儲量達到TB級別。
如此巨大的查找表會需要使用分布式計算，將每一個嵌入層的查找表都分布到多塊TPU v4晶片中進行計算。
Google在論文中提到，在進行這樣的嵌入層計算時，計算是以1D向量計算為主，而非2D矩陣或者3D張量計算；
其次，計算往往是稀疏的（因為輸入特徵是稀疏的，因此並不是所有特徵都會有高密度計算）而且分布在不同的晶片上，因此需要能對於共享存儲進行優化，這樣不同的晶片可以進行高效的數據交換。
為了實現對於嵌入層優化，Google在TPU v4中專門設計了一種專用加速模塊，稱為稀疏核（SparseCore，SC）。
每個SC都有自己的向量計算單元（scVPU），2.5 MB本地SRAM，以及可以訪問高達128TB共享HBM的記憶體訪問接口。
除此之外，SC還有一些專門為嵌入層操作設計的專用加速邏輯，包括排序（Sort）、規約（Reduce）、拼接（Concat）等。
我們可以看到，其實每個SC的結構都較為簡單，因此在每個TPU v4中都部署了大量SC，但同時SC總體的面積開銷和功耗開銷都僅僅占TPU v4的5%左右。
Google在論文中比較了使用CPU運行嵌入層（這也是常規運行嵌入層的做法）以及使用TPU v4 SC運行嵌入層，結果表明在運行相同的推薦系統時，相對於把嵌入層在CPU上運行，把嵌入層放在TPU v4的SC上可以把整體推薦系統的運行速度提升6倍以上。
事實上，這也是領域專用設計（domain-specific design）最吸引人的地方，即使用很小的晶片面積和功耗開銷，可以得到非常大的性能提升。
而Google在TPU v4的設計中把這樣的領域專用化設計放到了對於整個公司都最關鍵的地方（決定Google整體收入的推薦系統模型的核心瓶頸嵌入層），從而撬動了非常大的收益。
TPU v4的稀疏核（SC）設計，圖/Google
除了在晶片層面進行針對算法的優化之外，TPU v4還在分布式計算拓撲層面實現了對於算法的優化。
在前文中我們提到，TPU v4引入可重配置光互連可以針對不同的人工智慧模型實現不同的TPU v4之間的互聯拓撲，為此Google設計了一套機器學習算法來決定如何根據人工智慧模型來配置光路開關來提升性能。
分析表明，對於目前最熱門的大語言模型（包括GPT-3）的訓練，使用機器學習算法查找到的最優TPU光互聯拓撲配置可以提升1.2-2.3倍的性能。
圖/Google
最後，Google還為了TPU v4專門設計了一套神經網絡架構搜索（NAS）算法，可以根據TPU v4的特性來優化人工智慧模型，從而確保經過優化的模型可以最高效地運行在TPU v4上，並且充分利用TPU v4的資源。
與人工優化相比，使用該NAS可以實現推薦系統高達10%的運行時間優化，這大約相當於每年節省數千萬美元的成本。
TPU v4與人工智慧晶片的未來
從TPU v4的設計中，我們可以看到人工智慧晶片未來的一些方向，而這些方向是我們在Nvidia的GPU等其他主流人工智慧晶片的設計中也看到的：
首先就是對於高效互聯和規模化的支持。隨著人工智慧模型越來越大，對於這類模型的支持主要依賴人工智慧晶片的可擴展性（即如何讓多晶片可以高效並可靠地一起分工合作來加速這樣的大模型），而不是一味提高單晶片的能力來支持大模型，因為模型的演進總是要比晶片的設計疊代更快。
在這個領域，不同的晶片公司會有不同的側重，例如AMD側重較為微觀層面的使用chiplet來實現封裝級別的可擴展性，Nvidia有NvLink等晶片技術來實現單機多卡之間的可擴展性和性能提升，Google則直接為了海量TPU互聯設計了一款光路開關晶片；
但是這些公司之間的共性，即對於人工智慧晶片可擴展性的支持以滿足大模型的需求，卻是相當一致的。從這個角度來看，未來可擴展性（例如數據互聯帶寬）有可能會成為與峰值算力一樣的人工智慧晶片主要指標，而這也讓人工智慧晶片設計更加跨界：即不僅僅是需要對於數字邏輯和計算機架構方面的資源，同時也需要在封裝、數據互聯等領域都有積累。
此外，人工智慧晶片與算法之間的結合繼續保持緊密關系，算法-晶片協同設計仍然將是未來人工智慧晶片繼續提升性能的主要手段之一。
我們目前已經看到了Google、Nvidia等在算法-晶片協同設計中的大量成果：包括對於新的數制（Nvidia的FP16、FP8，Google的BF16等）的支持，對於計算特性的支持（Nvidia對於稀疏計算的支持），以及對於模型關鍵算法的直接專用加速器的部署（Nvidia的transformer acclerator，Google的SC等）。
隨著摩爾定律未來越來越接近物理極限，預計未來人工智慧晶片性能進一步提升會越來越倚賴算法-晶片協同設計，而另一方面，由於有算法-晶片協同設計，我們預計未來人工智慧晶片的性能仍然將保持類似摩爾定律的接近指數級提升，因此人工智慧晶片仍然將會是半導體行業未來幾年最為熱門的方向之一，也將會成為半導體行業未來繼續發展的重要引擎。
圖/Google
來源：快科技
