上周，谷歌在論文預(yù)印本平臺(tái)arxiv上發(fā)表了其關(guān)于TPU v4的深入解讀論文《TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings》（TPU v4：通過(guò)光互聯(lián)可重配置的機(jī)器學(xué)習(xí)超級(jí)計(jì)算機(jī)，搭載硬件嵌入層加速）。

該論文將于今年六月在ISCA 2023（International Symposium on Computer Architecture，計(jì)算機(jī)架構(gòu)領(lǐng)域的頂級(jí)會(huì)議）上正式發(fā)表，而目前的預(yù)印本無(wú)疑為我們提供了可以一窺其全貌的機(jī)會(huì)。

為了TPU的可擴(kuò)展性設(shè)計(jì)專(zhuān)用光學(xué)芯片，谷歌也是拼了

從論文的標(biāo)題可以看到，谷歌TPU v4的一個(gè)主要亮點(diǎn)是通過(guò)光互連實(shí)現(xiàn)可重配置和高可擴(kuò)展性（也即標(biāo)題中的“optically reconfigurable”）。

而在論文的一開(kāi)始，谷歌開(kāi)門(mén)見(jiàn)山首先介紹的也并非傳統(tǒng)的MAC設(shè)計(jì)、片上內(nèi)存、HBM通道等AI芯片常見(jiàn)的參數(shù)，而是可配置的光學(xué)互聯(lián)開(kāi)關(guān)（reconfigurable optical switch）。

作為論文的重中之重，這里我們也詳細(xì)分析一下為什么光學(xué)互聯(lián)在TPU v4設(shè)計(jì)中占了這么重要的位置，以至于谷歌甚至為了它自研了一款光學(xué)芯片。

TPU v4從一開(kāi)始設(shè)計(jì)時(shí)，其目標(biāo)就是極高的可擴(kuò)展性，可以有數(shù)千個(gè)芯片同時(shí)加速，從而實(shí)現(xiàn)一個(gè)為了機(jī)器學(xué)習(xí)模型訓(xùn)練而設(shè)計(jì)的超級(jí)計(jì)算機(jī)。

在谷歌的設(shè)計(jì)中，超級(jí)計(jì)算機(jī)的拓?fù)浣Y(jié)構(gòu)為：將4x4x4（64）個(gè)TPU v4芯片互聯(lián)在一起形成一個(gè)立方體結(jié)構(gòu)（cube），然后再把4x4x4這樣的cube串聯(lián)在一起形成一個(gè)總共有4096個(gè)TPU v4的超級(jí)計(jì)算機(jī)。

TPU超級(jí)計(jì)算機(jī)（由4096個(gè)TPU v4組成）拓?fù)浣Y(jié)構(gòu)，圖/谷歌

在這樣的拓?fù)渲?，物理距離較近的TPU v4（即在同一個(gè)4x4x4 cube中的芯片）可以用常規(guī)的電互聯(lián)（例如銅絞線(xiàn)）方法連接，但是距離較遠(yuǎn)的TPU之間（例如在cube之間的互聯(lián)）就必須使用光互連，原因就在于在如此大規(guī)模的超級(jí)計(jì)算機(jī)中，芯片之間的數(shù)據(jù)互聯(lián)在很大程度上會(huì)決定整體計(jì)算的效率；

如果數(shù)據(jù)互聯(lián)效率不夠高的話(huà)，很多時(shí)候芯片都在等待來(lái)自其他芯片的數(shù)據(jù)到達(dá)以后開(kāi)始計(jì)算，這樣就形成了效率浪費(fèi)。為了避免這樣“芯片等數(shù)據(jù)”的情形出現(xiàn)，就必須確保芯片之間互聯(lián)能擁有高帶寬，低延遲。而光互連對(duì)于物理距離較遠(yuǎn)的芯片就成為了首選。

光互連在高性能計(jì)算中的使用也并非新聞，而谷歌在TPU v4中的主要突破是使用可重配置的光互連（即加入光路開(kāi)關(guān)，optical circuit switch OCS）來(lái)快速實(shí)現(xiàn)不同的芯片互聯(lián)拓?fù)洹?/p>

換句話(huà)說(shuō)，芯片之間的互聯(lián)并非一成不變的，而是可以現(xiàn)場(chǎng)可重配置的。這樣做可以帶來(lái)許多好處，其中最主要的就是可以根據(jù)具體機(jī)器學(xué)習(xí)模型來(lái)改變拓?fù)?，以及改善超?jí)計(jì)算機(jī)的可靠性。

從拓?fù)浣Y(jié)構(gòu)來(lái)說(shuō)，不同的機(jī)器學(xué)習(xí)模型對(duì)于數(shù)據(jù)流的要求大致可以分為三類(lèi)，即數(shù)據(jù)并行（每塊芯片都加載整個(gè)模型，不同的芯片處理數(shù)據(jù)集中不同的數(shù)據(jù)），模型并行（模型中有些層特別大，因此每塊芯片只負(fù)責(zé)這樣很大的層中的一部分計(jì)算），以及流水線(xiàn)并行（把模型中的不同層交給不同的芯片計(jì)算），而不同的數(shù)據(jù)流就對(duì)應(yīng)了不同的TPU互聯(lián)拓?fù)洹?/p>

當(dāng)有了可重配置光互連之后，就可以根據(jù)具體模型數(shù)據(jù)流來(lái)調(diào)整TPU之間的互聯(lián)拓?fù)?，從而?shí)現(xiàn)最優(yōu)的性能，其提升可超過(guò)2倍。

使用可重配置光互連可以快速切換不同的芯片間互聯(lián)拓?fù)?，圖/谷歌

另一個(gè)優(yōu)勢(shì)就是可靠性。在這樣擁有海量芯片組成的超級(jí)計(jì)算機(jī)中，一個(gè)重要的考量就是，如果有一小部分的芯片不工作了，如何確保整體超級(jí)計(jì)算機(jī)仍然能維持較高的性能？如果使用常規(guī)的固定互聯(lián)架構(gòu)，那么一個(gè)芯片出故障可能會(huì)影響整個(gè)系統(tǒng)工作。

而在有了可重配置的光互連之后，需要做的只需要把出故障的芯片繞過(guò)，就不會(huì)影響整個(gè)系統(tǒng)的工作，最多會(huì)犧牲一點(diǎn)整體的性能。

谷歌在論文中給出了一個(gè)單芯片故障率和系統(tǒng)平均性能影響的曲線(xiàn)圖，在使用可配置光互連（以及光路開(kāi)關(guān)）時(shí)，假設(shè)芯片可靠率在99%的情況下，其整體系統(tǒng)的平均性能提升比不使用OCS可高達(dá)6倍，可見(jiàn)光互連開(kāi)關(guān)的重要性。

系統(tǒng)性能與芯片可靠性曲線(xiàn)圖——使用OCS可以大大提升系統(tǒng)性能，圖/谷歌

為了實(shí)現(xiàn)數(shù)據(jù)中心級(jí)的可配置光互連，需要光路開(kāi)關(guān)首先能高效擴(kuò)展到超高數(shù)量的互聯(lián)數(shù)（例如1000x1000），同時(shí)需要實(shí)現(xiàn)低開(kāi)關(guān)切換延遲，低成本，以及低信號(hào)損耗。

谷歌認(rèn)為現(xiàn)有的商用方案都不夠滿(mǎn)意，因此谷歌的做法是自研了一款光路開(kāi)關(guān)芯片Palomar，并且使用該芯片實(shí)現(xiàn)了全球首個(gè)數(shù)據(jù)中心級(jí)的可配置光互連，而TPU v4就是搭配了這款自研光路開(kāi)關(guān)芯片從架構(gòu)上實(shí)現(xiàn)了高性能。

谷歌自研的光路開(kāi)關(guān)芯片Palomar使用的是基于MEMS反射鏡陣列的技術(shù)，具體原理是使用一個(gè)2D MEMS反射鏡陣列，通過(guò)控制反射鏡的位置來(lái)調(diào)整光路，從而實(shí)現(xiàn)光路的切換。

使用MEMS的光路開(kāi)關(guān)芯片可以實(shí)現(xiàn)低損耗，低切換延遲（毫秒級(jí)別）以及低功耗。在經(jīng)過(guò)一系列優(yōu)化之后，光路系統(tǒng)的成本也控制得很低，在整個(gè)TPU v4超級(jí)計(jì)算機(jī)成本中占5%以下。

谷歌自研的Palomar MEMS光路開(kāi)關(guān)芯片，圖/谷歌

算法-芯片協(xié)同設(shè)計(jì)是TPU v4的靈魂

如果說(shuō)可重配置光互聯(lián)給TPU v4提供了良好的根基的話(huà)，那么算法-芯片協(xié)同設(shè)計(jì)就是TPU v4的靈魂。算法-芯片協(xié)同設(shè)計(jì)包括兩部分，一部分是如何根據(jù)算法優(yōu)化芯片，而另一方面是如何根據(jù)芯片去優(yōu)化算法，在TPU v4的架構(gòu)中，兩者都得到了仔細(xì)考慮。

我們首先分析TPU v4如何根據(jù)算法來(lái)優(yōu)化芯片。如果說(shuō)2017年TPU v1發(fā)表時(shí)候，其主要解決的還是卷積神經(jīng)網(wǎng)絡(luò)CNN的話(huà)，那么在2023年來(lái)看，CNN的加速問(wèn)題早已經(jīng)被解決得差不多了，更多的是如何處理目前如日中天的大模型的問(wèn)題。

對(duì)于谷歌來(lái)說(shuō)，目前最關(guān)鍵的大模型就是決定了其公司主營(yíng)收入的推薦系統(tǒng)大模型，因此TPU v4的設(shè)計(jì)也針對(duì)推薦系統(tǒng)大模型做了相當(dāng)?shù)膬?yōu)化。

在推薦系統(tǒng)大模型中，目前的加速瓶頸是嵌入層（embedding layer）。嵌入層的目的是將高維度稀疏特征映射到低維度高密度特征，從而該高密度低維度特征可以被神經(jīng)網(wǎng)絡(luò)進(jìn)一步處理。

嵌入層的實(shí)現(xiàn)通常是一個(gè)查找表（look-up table），而這個(gè)查找表可以非常巨大至100GB的數(shù)量級(jí)。在一個(gè)推薦系統(tǒng)模型中可以有多個(gè)這樣的查找表，從而讓整個(gè)查找表的存儲(chǔ)量達(dá)到TB級(jí)別。

如此巨大的查找表會(huì)需要使用分布式計(jì)算，將每一個(gè)嵌入層的查找表都分布到多塊TPU v4芯片中進(jìn)行計(jì)算。

谷歌在論文中提到，在進(jìn)行這樣的嵌入層計(jì)算時(shí)，計(jì)算是以1D向量計(jì)算為主，而非2D矩陣或者3D張量計(jì)算；

其次，計(jì)算往往是稀疏的（因?yàn)檩斎胩卣魇窍∈璧?，因此并不是所有特征都?huì)有高密度計(jì)算）而且分布在不同的芯片上，因此需要能對(duì)于共享存儲(chǔ)進(jìn)行優(yōu)化，這樣不同的芯片可以進(jìn)行高效的數(shù)據(jù)交換。

為了實(shí)現(xiàn)對(duì)于嵌入層優(yōu)化，谷歌在TPU v4中專(zhuān)門(mén)設(shè)計(jì)了一種專(zhuān)用加速模塊，稱(chēng)為稀疏核（SparseCore，SC）。

每個(gè)SC都有自己的向量計(jì)算單元（scVPU），2.5 MB本地SRAM，以及可以訪(fǎng)問(wèn)高達(dá)128TB共享HBM的內(nèi)存訪(fǎng)問(wèn)接口。

除此之外，SC還有一些專(zhuān)門(mén)為嵌入層操作設(shè)計(jì)的專(zhuān)用加速邏輯，包括排序（Sort）、規(guī)約（Reduce）、拼接（Concat）等。

我們可以看到，其實(shí)每個(gè)SC的結(jié)構(gòu)都較為簡(jiǎn)單，因此在每個(gè)TPU v4中都部署了大量SC，但同時(shí)SC總體的面積開(kāi)銷(xiāo)和功耗開(kāi)銷(xiāo)都僅僅占TPU v4的5%左右。

谷歌在論文中比較了使用CPU運(yùn)行嵌入層（這也是常規(guī)運(yùn)行嵌入層的做法）以及使用TPU v4 SC運(yùn)行嵌入層，結(jié)果表明在運(yùn)行相同的推薦系統(tǒng)時(shí)，相對(duì)于把嵌入層在CPU上運(yùn)行，把嵌入層放在TPU v4的SC上可以把整體推薦系統(tǒng)的運(yùn)行速度提升6倍以上。

事實(shí)上，這也是領(lǐng)域?qū)Ｓ迷O(shè)計(jì)（domain-specific design）最吸引人的地方，即使用很小的芯片面積和功耗開(kāi)銷(xiāo)，可以得到非常大的性能提升。

而谷歌在TPU v4的設(shè)計(jì)中把這樣的領(lǐng)域?qū)Ｓ没O(shè)計(jì)放到了對(duì)于整個(gè)公司都最關(guān)鍵的地方（決定谷歌整體收入的推薦系統(tǒng)模型的核心瓶頸嵌入層），從而撬動(dòng)了非常大的收益。

TPU v4的稀疏核（SC）設(shè)計(jì)，圖/谷歌

除了在芯片層面進(jìn)行針對(duì)算法的優(yōu)化之外，TPU v4還在分布式計(jì)算拓?fù)鋵用鎸?shí)現(xiàn)了對(duì)于算法的優(yōu)化。

在前文中我們提到，TPU v4引入可重配置光互連可以針對(duì)不同的人工智能模型實(shí)現(xiàn)不同的TPU v4之間的互聯(lián)拓?fù)?，為此谷歌設(shè)計(jì)了一套機(jī)器學(xué)習(xí)算法來(lái)決定如何根據(jù)人工智能模型來(lái)配置光路開(kāi)關(guān)來(lái)提升性能。

分析表明，對(duì)于目前最熱門(mén)的大語(yǔ)言模型（包括GPT-3）的訓(xùn)練，使用機(jī)器學(xué)習(xí)算法查找到的最優(yōu)TPU光互聯(lián)拓?fù)渑渲每梢蕴嵘?.2-2.3倍的性能。

圖/谷歌

最后，谷歌還為了TPU v4專(zhuān)門(mén)設(shè)計(jì)了一套神經(jīng)網(wǎng)絡(luò)架構(gòu)搜索（NAS）算法，可以根據(jù)TPU v4的特性來(lái)優(yōu)化人工智能模型，從而確保經(jīng)過(guò)優(yōu)化的模型可以最高效地運(yùn)行在TPU v4上，并且充分利用TPU v4的資源。

與人工優(yōu)化相比，使用該NAS可以實(shí)現(xiàn)推薦系統(tǒng)高達(dá)10%的運(yùn)行時(shí)間優(yōu)化，這大約相當(dāng)于每年節(jié)省數(shù)千萬(wàn)美元的成本。

TPU v4與人工智能芯片的未來(lái)

從TPU v4的設(shè)計(jì)中，我們可以看到人工智能芯片未來(lái)的一些方向，而這些方向是我們?cè)贜vidia的GPU等其他主流人工智能芯片的設(shè)計(jì)中也看到的：

首先就是對(duì)于高效互聯(lián)和規(guī)?；闹С帧ｋS著人工智能模型越來(lái)越大，對(duì)于這類(lèi)模型的支持主要依賴(lài)人工智能芯片的可擴(kuò)展性（即如何讓多芯片可以高效并可靠地一起分工合作來(lái)加速這樣的大模型），而不是一味提高單芯片的能力來(lái)支持大模型，因?yàn)槟Ｐ偷难葸M(jìn)總是要比芯片的設(shè)計(jì)迭代更快。

在這個(gè)領(lǐng)域，不同的芯片公司會(huì)有不同的側(cè)重，例如AMD側(cè)重較為微觀層面的使用chiplet來(lái)實(shí)現(xiàn)封裝級(jí)別的可擴(kuò)展性，Nvidia有NvLink等芯片技術(shù)來(lái)實(shí)現(xiàn)單機(jī)多卡之間的可擴(kuò)展性和性能提升，谷歌則直接為了海量TPU互聯(lián)設(shè)計(jì)了一款光路開(kāi)關(guān)芯片；

但是這些公司之間的共性，即對(duì)于人工智能芯片可擴(kuò)展性的支持以滿(mǎn)足大模型的需求，卻是相當(dāng)一致的。從這個(gè)角度來(lái)看，未來(lái)可擴(kuò)展性（例如數(shù)據(jù)互聯(lián)帶寬）有可能會(huì)成為與峰值算力一樣的人工智能芯片主要指標(biāo)，而這也讓人工智能芯片設(shè)計(jì)更加跨界：即不僅僅是需要對(duì)于數(shù)字邏輯和計(jì)算機(jī)架構(gòu)方面的資源，同時(shí)也需要在封裝、數(shù)據(jù)互聯(lián)等領(lǐng)域都有積累。

此外，人工智能芯片與算法之間的結(jié)合繼續(xù)保持緊密關(guān)系，算法-芯片協(xié)同設(shè)計(jì)仍然將是未來(lái)人工智能芯片繼續(xù)提升性能的主要手段之一。

我們目前已經(jīng)看到了谷歌、Nvidia等在算法-芯片協(xié)同設(shè)計(jì)中的大量成果：包括對(duì)于新的數(shù)制（Nvidia的FP16、FP8，谷歌的BF16等）的支持，對(duì)于計(jì)算特性的支持（Nvidia對(duì)于稀疏計(jì)算的支持），以及對(duì)于模型關(guān)鍵算法的直接專(zhuān)用加速器的部署（Nvidia的transformer acclerator，谷歌的SC等）。

隨著摩爾定律未來(lái)越來(lái)越接近物理極限，預(yù)計(jì)未來(lái)人工智能芯片性能進(jìn)一步提升會(huì)越來(lái)越倚賴(lài)算法-芯片協(xié)同設(shè)計(jì)，而另一方面，由于有算法-芯片協(xié)同設(shè)計(jì)，我們預(yù)計(jì)未來(lái)人工智能芯片的性能仍然將保持類(lèi)似摩爾定律的接近指數(shù)級(jí)提升，因此人工智能芯片仍然將會(huì)是半導(dǎo)體行業(yè)未來(lái)幾年最為熱門(mén)的方向之一，也將會(huì)成為半導(dǎo)體行業(yè)未來(lái)繼續(xù)發(fā)展的重要引擎。??

圖/谷歌