比A100系統(tǒng)快1.7倍，從Google TPU v4看AI芯片的未來
相比較使用英偉達(dá) A100 所創(chuàng)建的超級計(jì)算機(jī)，使用谷歌 TPUv4 構(gòu)建的超級計(jì)算機(jī)速度快 1.2-1.7 倍，功耗降低 1.3-1.9 倍。
目前在人工智能計(jì)算領(lǐng)域，英偉達(dá)占據(jù)主導(dǎo)地位，超過 90% 的人工智能項(xiàng)目開發(fā)都使用其芯片。而谷歌試圖在該市場上和英偉達(dá)競爭，TPU 的全稱為 Tensor Processing Units，是谷歌自研的芯片。
為了TPU的可擴(kuò)展性設(shè)計(jì)專用光學(xué)芯片，谷歌也是拼了
從論文的標(biāo)題可以看到，谷歌TPU v4的一個(gè)主要亮點(diǎn)是通過光互連實(shí)現(xiàn)可重配置和高可擴(kuò)展性（也即標(biāo)題中的“optically reconfigurable”）。而在論文的一開始，谷歌開門見山首先介紹的也并非傳統(tǒng)的MAC設(shè)計(jì)、片上內(nèi)存、HBM通道等AI芯片常見的參數(shù)，而是可配置的光學(xué)互聯(lián)開關(guān)（reconfigurable optical switch）。作為論文的重中之重，這里我們也詳細(xì)分析一下為什么光學(xué)互聯(lián)在TPU v4設(shè)計(jì)中占了這么重要的位置，以至于谷歌甚至為了它自研了一款光學(xué)芯片。
TPU v4從一開始設(shè)計(jì)時(shí)，其目標(biāo)就是極高的可擴(kuò)展性，可以有數(shù)千個(gè)芯片同時(shí)加速，從而實(shí)現(xiàn)一個(gè)為了機(jī)器學(xué)習(xí)模型訓(xùn)練而設(shè)計(jì)的超級計(jì)算機(jī)。在谷歌的設(shè)計(jì)中，超級計(jì)算機(jī)的拓?fù)浣Y(jié)構(gòu)為：將4x4x4（64）個(gè)TPU v4芯片互聯(lián)在一起形成一個(gè)立方體結(jié)構(gòu)（cube），然后再把4x4x4這樣的cube用連在一起形成一個(gè)總共有4096個(gè)TPU v4的超級計(jì)算機(jī)。
在這樣的拓?fù)渲校锢砭嚯x較近的TPU v4（即在同一個(gè)4x4x4 cube中的芯片）可以用常規(guī)的電互聯(lián)（例如銅絞線）方法連接，但是距離較遠(yuǎn)的TPU之間（例如在cube之間的互聯(lián)）就必須使用光互連，原因就在于在如此大規(guī)模的超級計(jì)算機(jī)中，芯片之間的數(shù)據(jù)互聯(lián)在很大程度上會決定整體計(jì)算的效率；如果數(shù)據(jù)互聯(lián)效率不夠高的話，很多時(shí)候芯片都在等待來自其他芯片數(shù)據(jù)到達(dá)以開始計(jì)算，這樣就形成了效率浪費(fèi)。為了避免這樣“芯片等數(shù)據(jù)”的情形出現(xiàn)，就必須確保芯片之間互聯(lián)能擁有高帶寬，低延遲。而光互連對于物理距離較遠(yuǎn)的芯片就成為了首選。
光互連在高性能計(jì)算中的使用也并非新聞，而谷歌在TPU v4中的主要突破是使用可重配置的光互連（即加入光路開關(guān)，optical circuit switch OCS）來快速實(shí)現(xiàn)不同的芯片互聯(lián)拓?fù)?。換句話說，芯片之間的互聯(lián)并非一成不變的，而是可以現(xiàn)場可重配置的。這樣做可以帶來許多好處，其中最主要的就是可以根據(jù)具體機(jī)器學(xué)習(xí)模型來改變拓?fù)?，以及改善超級?jì)算機(jī)的可靠性。
從拓?fù)浣Y(jié)構(gòu)來說，不同的機(jī)器學(xué)習(xí)模型對于數(shù)據(jù)流的要求大致可以分為三類，即數(shù)據(jù)平行（每塊芯片都加載整個(gè)模型，不同的芯片處理數(shù)據(jù)集中不同的數(shù)據(jù)），模型并行（模型中有些層特別大，因此每塊芯片只負(fù)責(zé)這樣很大的層中的一部分計(jì)算），以及流水線并行（把模型中的不同層交給不同的芯片計(jì)算），而不同的數(shù)據(jù)流就對應(yīng)了不同的TPU互聯(lián)拓?fù)洹．?dāng)有了可重配置光互連之后，就可以根據(jù)具體模型數(shù)據(jù)流來調(diào)整TPU之間的互聯(lián)拓?fù)?，從而?shí)現(xiàn)最優(yōu)的性能，其提升可超過2倍。
另一個(gè)優(yōu)勢就是可靠性。在這樣擁有海量芯片組成的超級計(jì)算機(jī)中，一個(gè)重要的考量就是，如果有一小部分的芯片不工作了，如何確保整體超級計(jì)算機(jī)仍然能維持較高的性能？如果使用常規(guī)的固定互聯(lián)架構(gòu)，那么一個(gè)芯片出故障可能會影響整個(gè)系統(tǒng)工作。而在有了可重配置的光互連之后，需要做的只需要把出故障的芯片繞過，就不會影響整個(gè)系統(tǒng)的工作，最多會犧牲一點(diǎn)整體的性能。谷歌在論文中給出了一個(gè)單芯片故障率和系統(tǒng)平均性能影響的曲線圖，在使用可配置光互連（以及光路開關(guān)）時(shí)，假設(shè)芯片可靠率在99%的情況下，其整體系統(tǒng)的平均性能提升比不使OCS可高達(dá)6倍，可見光互連開關(guān)的重要性。
為了實(shí)現(xiàn)數(shù)據(jù)中心級的可配置光互連，需要光路開關(guān)首先能高效擴(kuò)展到超高數(shù)量的互聯(lián)數(shù)（例如1000x1000），同時(shí)需要實(shí)現(xiàn)低開關(guān)切換延遲，低成本，以及低信號損耗。谷歌認(rèn)為現(xiàn)有的商用方案都不夠滿意，因此谷歌的做法是自研了一款光路開關(guān)芯片Palomar，并且使用該芯片實(shí)現(xiàn)了全球首個(gè)數(shù)據(jù)中心級的可配置光互連，而TPU v4就是搭配了這款自研光路開關(guān)芯片從架構(gòu)上實(shí)現(xiàn)了高性能。谷歌自研的光路開關(guān)芯片Palomar使用的是基于MEMS反射鏡陣列的技術(shù)，具體原理是使用一個(gè)2D MEMS反射鏡陣列，通過控制反射鏡的位置來調(diào)整光路，從而實(shí)現(xiàn)光路的切換。使用MEMS的光路開關(guān)芯片可以實(shí)現(xiàn)低損耗，低切換延遲（毫秒級別）以及低功耗。在經(jīng)過一系列優(yōu)化之后，光路系統(tǒng)的成本也控制得很低，在整個(gè)TPU v4超級計(jì)算機(jī)成本中占5%以下。
Google自研TPU，五年更新四代
Google最早于2016年宣布首款內(nèi)部定制的AI芯片，區(qū)別于訓(xùn)練和部署AI模型的最常見的組合架構(gòu)，即CPU和GPU組合，第一代TPU在那場世界著名的人機(jī)圍棋大戰(zhàn)助力AlphaGo打敗李世石“一戰(zhàn)成名”，宣告并不是只有GPU才能做訓(xùn)練和推理。
Google第一代TPU采用28nm工藝制程，功耗大約40W，僅適用于深度學(xué)習(xí)推理，除了AlphaGo，也用在Google搜索、翻譯等機(jī)器學(xué)習(xí)模型中。
2017年5月，Google發(fā)布了能夠?qū)崿F(xiàn)機(jī)器學(xué)習(xí)模型訓(xùn)練和推理的TPU v2，達(dá)到180TFLOPs浮點(diǎn)運(yùn)算能力，同時(shí)內(nèi)存帶寬也得以提升，比同期推出的CPU AI工作負(fù)載提升30倍，比GPU AI工作負(fù)載提升15倍，被基于4塊TPU v2的AlphaGo擊敗的世界圍棋冠軍柯潔最直觀地感受了這一切。
2018年5月，Google又發(fā)布第三代TPU v3，性能是上一代TPU的兩倍，實(shí)現(xiàn)420TFLOPs浮點(diǎn)運(yùn)算，以及128GB的高帶寬內(nèi)存。
按照一年一次迭代更新的節(jié)奏，Google理應(yīng)在2019年推出第四代TPU，不過這一年的I/O大會上，Google推出的是第二代和第三代TPU Pod，可以配置超過1000顆TPU，大大縮短了在進(jìn)行復(fù)雜的模型訓(xùn)練時(shí)所需耗費(fèi)的時(shí)間。
在AI芯片發(fā)展史上，無論是從片上內(nèi)存上，還是從可編程能力來看，Google TPU都是不可多得的技術(shù)創(chuàng)新，打破GPU的“壟斷”地位，且打開云端AI芯片的新競爭格局。
發(fā)展五年的Google TPU在今天依然保持著強(qiáng)勁的競爭力，未來的世界是什么樣的？Google TPU已經(jīng)告訴了我們一小部分答案。
TPU v4與人工智能芯片的未來
從TPU v4的設(shè)計(jì)中，我們可以看到人工智能芯片未來的一些方向，而這些方向是我們在Nvidia的GPU等其他主流人工智能芯片的設(shè)計(jì)中也看到的：
首先就是對于高效互聯(lián)和規(guī)?；闹С?。隨著人工智能模型越來越大，對于這類模型的支持主要依賴人工智能芯片的可擴(kuò)展性（即如何讓多芯片可以高效并可靠地一起分工合作來加速這樣的大模型），而不是一味提高單芯片的能力來支持大模型，因?yàn)槟Ｐ偷难葸M(jìn)總是要比芯片的設(shè)計(jì)迭代更快。在這個(gè)領(lǐng)域，不同的芯片公司會有不同的側(cè)重，例如AMD側(cè)重較為微觀層面的使用chiplet來實(shí)現(xiàn)封裝級別的可擴(kuò)展性，Nvidia有NvLink等芯片技術(shù)來實(shí)現(xiàn)單機(jī)多卡之間的可擴(kuò)展性和性能提升，谷歌則直接為了海量TPU互聯(lián)設(shè)計(jì)了一款光路開關(guān)芯片；但是這些公司之間的共性，即對于人工智能芯片可擴(kuò)展性的支持以滿足大模型的需求，卻是相當(dāng)一致的。從這個(gè)角度來看，未來可擴(kuò)展性（例如數(shù)據(jù)互聯(lián)帶寬）有可能會成為與峰值算力一樣的人工智能芯片主要指標(biāo)，而這也讓人工智能芯片設(shè)計(jì)更加跨界：即不僅僅是需要對于數(shù)字邏輯和計(jì)算機(jī)架構(gòu)方面的資源，同時(shí)也需要在封裝、數(shù)據(jù)互聯(lián)等領(lǐng)域都有積累。
此外，人工智能芯片與算法之間的結(jié)合繼續(xù)保持緊密關(guān)系，算法-芯片協(xié)同設(shè)計(jì)仍然將是未來人工智能芯片繼續(xù)提升性能的主要手段之一。我們目前已經(jīng)看到了谷歌、Nvidia等在算法-芯片協(xié)同設(shè)計(jì)中的大量成果：包括對于新的數(shù)制（Nvidia的FP16、FP8，谷歌的BF16等）的支持，對于計(jì)算特性的支持（Nvidia對于稀疏計(jì)算的支持），以及對于模型關(guān)鍵算法的直接專用加速器的部署（Nvidia的transformer acclerator，谷歌的SC等）。隨著摩爾定律未來越來越接近物理極限，預(yù)計(jì)未來人工智能芯片性能進(jìn)一步提升會越來越倚賴算法-芯片協(xié)同設(shè)計(jì)，而另一方面，由于有算法-芯片協(xié)同設(shè)計(jì)，我們預(yù)計(jì)未來人工智能芯片的性能仍然將保持類似摩爾定律的接近指數(shù)級提升，因此人工智能芯片仍然將會是半導(dǎo)體行業(yè)未來幾年最為熱門的方向之一，也將會成為半導(dǎo)體行業(yè)未來繼續(xù)發(fā)展的重要引擎。
