Google今天正式發(fā)布了其第四代人工智能TPU v4 AI芯片，其速度達(dá)到了TPU v3的2.7被。Google實(shí)際上已經(jīng)于2020年就開始在自己的數(shù)據(jù)中心中使用了新的TPU v4。通過(guò)整合4096個(gè)TPU v4芯片成一個(gè)TPU v4 Pod，一個(gè)Pod性能就達(dá)到世界第一超算“富岳”的兩倍。這些算力可能在今年晚些時(shí)候向Google Cloud用戶開放此功能。且谷歌希望未來(lái)可能應(yīng)用于量子計(jì)算。

　　Google TPU系列研發(fā)

　　Google于2016年宣布了其首款內(nèi)部定制的定制AI芯片。

　　與傳統(tǒng)的CPU和GPU組合（用于訓(xùn)練和部署AI模型的最常見(jiàn)架構(gòu)）相比，第一款TPU ASIC提供了“數(shù)量級(jí)”更高的性能。

　　TPU V2于2018年上市，TPU V3于2019年上市，TPU v4于2020年推出。

　　借助TPU v3，這些服務(wù)器的機(jī)架需要大量的電源，以至于Google必須對(duì)其數(shù)據(jù)中心進(jìn)行改造以支持液體冷卻，這比傳統(tǒng)的風(fēng)冷系統(tǒng)具有更高的功率密度。

　　而今，在比TPU v3性能提高2.7倍的基礎(chǔ)上再整合4096個(gè)TPU v4為一個(gè)TPU v4 Pod，許多Pod再組成TPU v4 Pods，以提供超強(qiáng)的算力。

　　除了將這些系統(tǒng)用于自己的AI應(yīng)用程序（例如搜索建議，語(yǔ)言翻譯或語(yǔ)音助手）外，Google還將TPU基礎(chǔ)設(shè)施（包括整個(gè)TPU吊艙）以云的方式（付費(fèi)）開放給Google Cloud客戶。

　　1 TPU v4 Pod 整合算力 = 2 富岳

　　在其I/O開發(fā)者大會(huì)上，Google今天（美國(guó)時(shí)間5月18日）宣布了其下一代定制的張量處理單元（TPU）人工智能芯片。這是TPU芯片的第四代產(chǎn)品，Google稱其速度是上一版本的兩倍。正如Google首席執(zhí)行官桑達(dá)爾·皮查伊所指出的，這些芯片整合了4096個(gè)TPU v4，一個(gè)pod就可以提供超過(guò)一個(gè)exaflop（每秒百億億次浮點(diǎn)運(yùn)算）的AI計(jì)算能力，達(dá)到目前全球最快的超級(jí)計(jì)算機(jī)“富岳”的兩倍性能。

　　Google使用定制芯片為其自己的許多機(jī)器學(xué)習(xí)服務(wù)提供動(dòng)力，但與此同時(shí)它也將把這最新一代產(chǎn)品作為其Google云平臺(tái)的一部分提供給開發(fā)者。

　　“這是我們?cè)贕oogle部署過(guò)的最快的系統(tǒng)，對(duì)我們來(lái)說(shuō)是歷史性的里程碑，”Google CEO桑達(dá)爾·皮查伊說(shuō)?！耙郧耙氆@得1個(gè)exaflop的算力，通常需要建立一個(gè)定制的超級(jí)計(jì)算機(jī)，但我們今天已經(jīng)部署了許多這樣的計(jì)算機(jī)，很快就會(huì)在我們的數(shù)據(jù)中心有幾十個(gè)TPUv4 pods，其中許多將以90%或接近90%的無(wú)碳能源運(yùn)行。而我們的TPUv4 pods將在今年晚些時(shí)候提供給我們的云客戶”。

　　TPU是Google的第一批定制芯片之一，當(dāng)包括微軟在內(nèi)的其他公司決定為其機(jī)器學(xué)習(xí)服務(wù)采用更靈活的FPGA時(shí)，Google很早就在這些定制芯片上下了賭注。雖然它們的開發(fā)時(shí)間稍長(zhǎng)，而且隨著技術(shù)的變化很快就會(huì)過(guò)時(shí)，但用于特定場(chǎng)景的運(yùn)算時(shí)可以提供明顯更好的性能。

　　發(fā)布會(huì)上，Google還展示了其量子數(shù)據(jù)中心，盡管目前相關(guān)研究仍停留于早期階段，但谷歌顯然希望在未來(lái)有朝一日夢(mèng)想成真。基于目前的材料和實(shí)驗(yàn)場(chǎng)所限制，量子計(jì)算必須在極低的超導(dǎo)溫度下運(yùn)行。

　　或用于量子計(jì)算

　　發(fā)布會(huì)上，Google還展示了其量子數(shù)據(jù)中心，盡管目前相關(guān)研究仍停留于早期階段，但谷歌顯然希望在未來(lái)有朝一日夢(mèng)想成真?；谀壳暗牟牧虾蛯?shí)驗(yàn)場(chǎng)所限制，量子計(jì)算必須在極低的超導(dǎo)溫度下運(yùn)行。

　　Google TPU v4 性能詳解：打破MLPerf中的AI性能記錄

　　機(jī)器學(xué)習(xí)（ML）模型的快速訓(xùn)練對(duì)于提供以前無(wú)法實(shí)現(xiàn)的新產(chǎn)品，服務(wù)和研究突破的研究和工程團(tuán)隊(duì)至關(guān)重要。在Google，最近啟用ML的功能包括更有用的搜索結(jié)果和一個(gè)可以翻譯100種不同語(yǔ)言的ML模型。

　　行業(yè)標(biāo)準(zhǔn)MLPerf基準(zhǔn)測(cè)試對(duì)比的最新結(jié)果表明，谷歌已打造出世界上最快的ML訓(xùn)練超級(jí)計(jì)算機(jī)。Google使用此超級(jí)計(jì)算機(jī)以及最新的Tensor處理單元（TPU）芯片，在八項(xiàng)MLPerf基準(zhǔn)測(cè)試中的六項(xiàng)中創(chuàng)造了性能記錄。

　　圖1：與任何可用性類別中最快的非Google提交相比，Google最好的MLPerf Training v0.7研究提交速度有所提高。不論系統(tǒng)大小如何，總訓(xùn)練時(shí)間都將比較標(biāo)準(zhǔn)化，而系統(tǒng)大小介于8到4096個(gè)芯片之間。

　　使用TensorFlow，JAX和Lingvo中的ML模型實(shí)現(xiàn)實(shí)現(xiàn)了這些結(jié)果。在不到30秒的時(shí)間里，八個(gè)模型中的四個(gè)從零開始接受了訓(xùn)練。為了正確理解這一點(diǎn)，考慮到在2015年，花了超過(guò)三周的時(shí)間在可用的最先進(jìn)的硬件加速器上訓(xùn)練其中一種模型。谷歌最新的TPU超級(jí)計(jì)算機(jī)可以在五年后將相同模型的速度提高近五個(gè)數(shù)量級(jí)。

　　在此文中，我們將研究對(duì)比的一些細(xì)節(jié)，提交如何實(shí)現(xiàn)如此高的性能，以及它們對(duì)模型訓(xùn)練速度的所有意義。

　　MLPerf模型

　　選擇MLPerf模型來(lái)代表整個(gè)行業(yè)和學(xué)術(shù)界常見(jiàn)的尖端機(jī)器學(xué)習(xí)工作負(fù)載。上圖中，這是每個(gè)MLPerf模型的更多詳細(xì)信息：

　　DLRM代表了排名和推薦模型，這些模型是從媒體到旅行到電子商務(wù)等在線業(yè)務(wù)的核心

　　Transformer是自然語(yǔ)言處理（包括BERT）的最新發(fā)展浪潮的基礎(chǔ)

　　BERT使Google搜索獲得了“過(guò)去五年中最大的飛躍”

　　ResNet-50是廣泛用于圖像分類的模型

　　SSD是一種對(duì)象檢測(cè)模型，其重量很輕，可以在移動(dòng)設(shè)備上運(yùn)行

　　Mask R-CNN是廣泛使用的圖像分割模型，可用于自主導(dǎo)航，醫(yī)學(xué)成像和其他領(lǐng)域（您可以在Colab中進(jìn)行實(shí)驗(yàn)）

　　除了上述最大規(guī)模的行業(yè)領(lǐng)先結(jié)果之外，Google還使用Google Cloud Platform上的TensorFlow提供了MLPerf提交，供今天的企業(yè)使用。

　　世界上最快的ML訓(xùn)練超級(jí)計(jì)算機(jī)

　　Google在本次MLPerf訓(xùn)練回合中使用的超級(jí)計(jì)算機(jī)是Cloud TPU v3 Pod的四倍，后者在之前的比賽中創(chuàng)下了三項(xiàng)記錄。該系統(tǒng)包括4096個(gè)TPU v3芯片和數(shù)百個(gè)CPU主機(jī)，所有這些都通過(guò)超快速，超大規(guī)模的定制互連進(jìn)行連接。總體而言，該系統(tǒng)可提供430多個(gè)PFLOP峰值性能。

　　表1：所有這些MLPerf提交的內(nèi)容都是在Google新型ML超級(jí)計(jì)算機(jī)上以33秒或更短的時(shí)間從零開始進(jìn)行訓(xùn)練的。2個(gè)

　　使用TensorFlow，JAX，Lingvo和XLA進(jìn)行大規(guī)模訓(xùn)練

　　使用數(shù)千個(gè)TPU芯片訓(xùn)練復(fù)雜的ML模型需要在TensorFlow，JAX，Lingvo和XLA中結(jié)合算法技術(shù)和優(yōu)化。

　　XLA是支持所有Google MLPerf提交的基礎(chǔ)編譯器技術(shù)，TensorFlow是Google的端到端開源機(jī)器學(xué)習(xí)框架，Lingvo是使用TensorFlow構(gòu)建的序列模型的高級(jí)框架，而JAX是一個(gè)基于可組合功能轉(zhuǎn)換的以研究為中心的新框架。

　　上面的記錄設(shè)置規(guī)模依賴于模型并行性，按比例縮放的批次歸一化，有效的計(jì)算圖啟動(dòng)以及基于樹的權(quán)重初始化。

　　上表中的所有TensorFlow，JAX和Lingvo提交（ResNet-50，BERT，SSD和Transformer的實(shí)現(xiàn)）都在2048或4096 TPU芯片上進(jìn)行了訓(xùn)練，每個(gè)不到33秒。

　　TPU v4 vs TPU v3，性能提高2.7倍

　　Google的第四代TPU ASIC提供了TPU v3的矩陣乘法TFLOP的兩倍以上，顯著提高了內(nèi)存帶寬，并且在互連技術(shù)方面取得了進(jìn)步。Google的TPU v4 MLPerf提交利用了這些新的硬件功能以及互補(bǔ)的編譯器和建模優(yōu)勢(shì)。結(jié)果表明，在上一屆MLPerf培訓(xùn)比賽中，TPU v4比TPU v3的性能平均提高了2.7倍，達(dá)到了類似的規(guī)模。

　　圖2：Google的MLPerf培訓(xùn)v0.7中的TPU v4結(jié)果研究提交的結(jié)果，與Google的MLPerf培訓(xùn)v0.6可用提交中的TPU v3結(jié)果相比，平均提高了2.7倍，在相同規(guī)模的64核心芯片條件下。此性能改進(jìn)歸功于TPU v4中的硬件創(chuàng)新以及軟件的改進(jìn)。

　　沒(méi)有有關(guān)節(jié)點(diǎn)、體系結(jié)構(gòu)或設(shè)計(jì)的信息，但可以假定它們位于7nm類節(jié)點(diǎn)上。同樣，F(xiàn)LOP中只有2倍，而性能則是2.7倍，這意味著它們?cè)谔崛≡糉lop方面更加有效。

　　Google已經(jīng)在使用這項(xiàng)新技術(shù)，并將很快擴(kuò)大其使用范圍

　　Google實(shí)際上已經(jīng)在使用新的TPU v4。該公司表示已經(jīng)在自己的數(shù)據(jù)中心中使用了該技術(shù)。至于何時(shí)可以訪問(wèn)Google Cloud客戶，這家搜索巨頭尚未提供確切的日期。但它確實(shí)說(shuō)，它們將在“今年晚些時(shí)候”上市。

　　該公司表示，目前，它將很快在其數(shù)據(jù)中心運(yùn)行數(shù)十個(gè)最新的Pod。并指出，其中許多將“以或接近” 90％的無(wú)碳能源運(yùn)行。至少就Google自身的運(yùn)營(yíng)而言，使新系統(tǒng)也成為最高效的系統(tǒng)之一。