Google TPU架構(gòu)演進(jìn)及十大經(jīng)驗(yàn)

David Patterson，Google杰出工程師、UC Berkeley榮譽(yù)退休教授、美國(guó)國(guó)家工程院、科學(xué)院院士、文理科學(xué)院“三院”院士。他是RISC（精簡(jiǎn)指令集計(jì)算機(jī)）、RAID（獨(dú)立磁盤(pán)冗余陣列）和NOW（工作站網(wǎng)絡(luò)）的締造者，他與John Hennessy的著作《計(jì)算機(jī)體系結(jié)構(gòu)：量化研究方法》在業(yè)內(nèi)久負(fù)盛名。

2017年，David Patterson加入Google TPU團(tuán)隊(duì)，2018年3月，他與John Hennessy共同獲得圖靈獎(jiǎng)，2008年獲ACM/IEEE Eckert-Mauchly 獎(jiǎng)（被譽(yù)為計(jì)算機(jī)體系結(jié)構(gòu)最高獎(jiǎng)），2000年獲得馮·諾依曼獎(jiǎng)?wù)隆?/p>

本文是他近期在加州大學(xué)伯克利分校的演講，他分享了Google TPU近十年的發(fā)展歷程以及心得體會(huì)，并闡述了提升機(jī)器學(xué)習(xí)硬件能效對(duì)碳足跡的影響。OneFlow社區(qū)對(duì)此進(jìn)行了編譯。

01一場(chǎng)由TPU引發(fā)的“地震”

2013年，Google AI負(fù)責(zé)人Jeff Dean經(jīng)過(guò)計(jì)算后發(fā)現(xiàn)，如果有1億安卓用戶每天使用手機(jī)語(yǔ)音轉(zhuǎn)文字服務(wù)3分鐘，消耗的算力就已是Google所有數(shù)據(jù)中心總算力的兩倍，何況全球安卓用戶遠(yuǎn)不止1億。

如果僅通過(guò)擴(kuò)大數(shù)據(jù)中心規(guī)模來(lái)滿足算力需求，不但耗時(shí)，而且成本高昂。因此，Google決定針對(duì)機(jī)器學(xué)習(xí)構(gòu)建特定領(lǐng)域計(jì)算架構(gòu)（Domain-specific Architecture），希望將深度神經(jīng)網(wǎng)絡(luò)推理的總體擁有成本（TCO）降低至原來(lái)的十分之一。

于是，Google在2014年開(kāi)始研發(fā)TPU，項(xiàng)目進(jìn)展神速，僅15個(gè)月后TPU就可在Google數(shù)據(jù)中心部署應(yīng)用，而且TPU的性能遠(yuǎn)超預(yù)期，它的每瓦性能是是GPU的30倍、CPU的80倍。（數(shù)據(jù)源自論文： https://arxiv.org/ftp/arxiv/papers/1704/1704.04760.pdf）

2016年，在Google I/O開(kāi)發(fā)者大會(huì)上，Google首席執(zhí)行官Sundar Pichai對(duì)外公布了TPU這一突破性成果，他介紹道：

“通過(guò)Google云平臺(tái)，用戶不但可以接觸到Google內(nèi)部使用的高性能軟件，還可以使用Google內(nèi)部開(kāi)發(fā)的專用硬件。機(jī)器學(xué)習(xí)的計(jì)算規(guī)模巨大，因此Google研發(fā)了機(jī)器學(xué)習(xí)專用硬件，也就是‘張量處理單元（TPU）’。TPU的每瓦性能比市面上所有GPU和FPGA都高出一個(gè)數(shù)量級(jí)。用戶可以通過(guò)Google云平臺(tái)體驗(yàn)TPU的優(yōu)異性能。DeepMind研發(fā)的AlphaGo在與韓國(guó)棋手李世石的對(duì)戰(zhàn)中使用的底層硬件就是TPU?！?/p>

希臘神話中，特洛伊戰(zhàn)爭(zhēng)的起因是兩方爭(zhēng)奪世界上最美的女人——海倫，后世詩(shī)人將海倫的美貌“令成千戰(zhàn)艦為之起航”。我認(rèn)為T(mén)PU就像海倫，它的出現(xiàn)引起了“成千芯片與之競(jìng)逐”。

可以說(shuō)，TPU的問(wèn)世引發(fā)了硅谷的“地震”。TPU宣布誕生后，Intel耗資數(shù)十億美元收購(gòu)了多家芯片公司，阿里巴巴、Amazon等競(jìng)爭(zhēng)對(duì)手紛紛開(kāi)始研發(fā)類似產(chǎn)品。TPU重新喚起了人們對(duì)計(jì)算機(jī)架構(gòu)的關(guān)注，后來(lái)的幾年內(nèi)，出現(xiàn)了上百家相關(guān)初創(chuàng)企業(yè)，年均總?cè)谫Y額近20億美元，各種新奇的想法層出不窮。

五年后，Sundar Pichai又在2021年Google I/O開(kāi)發(fā)者大會(huì)公布TPU v4：

“AI技術(shù)的進(jìn)步有賴于計(jì)算基礎(chǔ)設(shè)施的支持，而TPU正是Google計(jì)算基礎(chǔ)設(shè)施的重要部分。新一代TPU v4芯片的速度是v3的兩倍多。Google用TPU集群構(gòu)建出Pod超級(jí)計(jì)算機(jī)，單臺(tái)TPU v4 Pod包含4096塊v4芯片，每臺(tái)Pod的芯片間互連帶寬是其他互連技術(shù)的10倍，因此，TPU v4 Pod的算力可達(dá)1 ExaFLOP，即每秒執(zhí)行10的18次方浮點(diǎn)運(yùn)算，相當(dāng)于1000萬(wàn)臺(tái)筆記本電腦的總算力?！?/p>

上圖展示了TPU的發(fā)展歷史。其中，Google尚未公布TPU v4i（TPU v4 lite）的相關(guān)細(xì)節(jié)。去年Google宣布TPU v4i已在云服務(wù)上可用，也發(fā)表了一篇關(guān)于TPU v4i的論文。

（https://www.gwern.net/docs/ai/scaling/hardware/2021-jouppi.pdf）

02 十年演進(jìn)，十大教訓(xùn)

過(guò)往十年，我們?cè)贛L計(jì)算架構(gòu)的發(fā)展中汲取了十大教訓(xùn)。

其中，前五個(gè)都和ML模型本身有關(guān)，后五個(gè)則關(guān)乎硬件和架構(gòu)。這些經(jīng)驗(yàn)對(duì)深度學(xué)習(xí)以外的領(lǐng)域也有借鑒意義。

教訓(xùn)一：DNN所需內(nèi)存空間和算力迅速增長(zhǎng)

我們閱讀近幾年的論文后發(fā)現(xiàn)，推理模型所需的內(nèi)存空間和算力平均每年增長(zhǎng)50%。由于芯片設(shè)計(jì)和部署至少各需要1年，投入實(shí)際使用并優(yōu)化需要3年。可見(jiàn)，從一款芯片開(kāi)始設(shè)計(jì)到生產(chǎn)周期結(jié)束的5年內(nèi)，模型所需的內(nèi)存空間和算力已增長(zhǎng)到大約8倍。因此，在芯片設(shè)計(jì)之初就要將這種增長(zhǎng)考慮在內(nèi)。

訓(xùn)練模型的長(zhǎng)速度比推理模型更快。根據(jù)OpenAI的統(tǒng)計(jì)，2012-2019年，SOTA訓(xùn)練模型的算力需求年均增長(zhǎng)10倍。備受關(guān)注的GPT-3模型的參數(shù)量更是從15億（GPT-2）增長(zhǎng)到1750億，提高了100倍。

教訓(xùn)二：DNN工作負(fù)載隨著DNN突破不斷演變

深度學(xué)習(xí)是一個(gè)日新月異的領(lǐng)域。2016年，MLP（多層感知器）模型仍是主流，但到2020年，CNN、RNN和BERT等不同模型百花齊放。BERT是一種全新的Transformer模型，誕生于2018年，短短兩年后，四分之一以上的Google內(nèi)部應(yīng)用都在使用BERT模型，可見(jiàn)深度學(xué)習(xí)發(fā)展變化之快。因此，ML計(jì)算架構(gòu)需要能夠支持多種模型。

教訓(xùn)三：DNN模型可優(yōu)化

通常而言，計(jì)算機(jī)架構(gòu)師只需懂硬件、體系結(jié)構(gòu)、編譯器，如果還懂操作系統(tǒng)則更好，但他們不需要懂應(yīng)用。然而，構(gòu)建針對(duì)特定領(lǐng)域的架構(gòu)則需要軟硬件兼通。

對(duì)ML工程師而言，只要可以讓模型跑得更好，他們非常愿意根據(jù)硬件/編譯器改進(jìn)DNN模型。畢竟DNN模型不像GCC編譯器，后者已成為被廣泛采納的編譯器標(biāo)準(zhǔn)，不會(huì)輕易根據(jù)硬件改動(dòng)。

DNN模型之所以可以優(yōu)化，部分原因是這些程序本身不算龐大，大約只是成千上萬(wàn)行PyTorch或TensorFlow代碼，操作可行性較強(qiáng)。

Google的一篇論文介紹了一種模型優(yōu)化技術(shù)Platform-aware AutoML，AutoML使用的方法稱為“神經(jīng)架構(gòu)搜索（Neural Architecture Search）”，即機(jī)器自動(dòng)在搜索空間中尋找更優(yōu)的神經(jīng)網(wǎng)絡(luò)模型結(jié)構(gòu)。在上述論文的例子中，經(jīng)機(jī)器自動(dòng)優(yōu)化后的CNN1模型，在相同的硬件和編譯器上可實(shí)現(xiàn)相同的準(zhǔn)確率，而運(yùn)算性能為原模型的1.6倍。

教訓(xùn)四：影響推理體驗(yàn)的是延遲，而非批次規(guī)模

一些關(guān)于模型推理優(yōu)化的論文把重點(diǎn)放在數(shù)據(jù)批次規(guī)模（batch size）上，認(rèn)為要把batch size設(shè)置為1才能使延遲降到最低。然而，通過(guò)MLPerf基準(zhǔn)數(shù)據(jù)可見(jiàn)，Google的生產(chǎn)模型在batch size相當(dāng)大的情況下也能實(shí)現(xiàn)低延遲，這可能是因?yàn)檫@些模型是基于TPU開(kāi)發(fā)，因此更加高效。

教訓(xùn)五：生產(chǎn)端推理需要多租戶技術(shù)

DNN需要使用多租戶技術(shù)（multi-tenancy）。不少深度學(xué)習(xí)論文的一個(gè)假設(shè)是同一時(shí)間只需運(yùn)行一個(gè)模型，但在實(shí)際應(yīng)用中，有不少情況都需要在不同模型中切換。

比如，機(jī)器翻譯涉及各種語(yǔ)言對(duì)，就需要不同的模型；傳統(tǒng)的軟件開(kāi)發(fā)需要用到一個(gè)主模型和多個(gè)實(shí)驗(yàn)?zāi)Ｐ?；甚至有時(shí)因?yàn)閷?duì)吞吐量和延遲有不同的側(cè)重要求，就需要不同的batch size，進(jìn)而需要不同的模型。

如上圖所示，我們收集了8個(gè)模型的基準(zhǔn)數(shù)據(jù)，其中6個(gè)模型涉及多租戶。右方的柱狀圖展示了模型大?。ㄒ訫B計(jì)算）。紅色虛線表示單塊芯片的最大SRAM，可見(jiàn)不少模型需要的內(nèi)存遠(yuǎn)大于此，這意味著需要有存取速度極快的DRAM。部分芯片的設(shè)計(jì)思路是利用SRAM解決所有任務(wù)，但在多租戶應(yīng)用場(chǎng)景下，我們認(rèn)為這很難辦到。

教訓(xùn)六：重要的是內(nèi)存，而非浮點(diǎn)運(yùn)算數(shù)

借用克林頓競(jìng)選總統(tǒng)時(shí)的口號(hào)——“重要的是經(jīng)濟(jì)，懂嗎？”（OneFlow譯注：當(dāng)時(shí)美國(guó)正值經(jīng)濟(jì)蕭條，克林頓將經(jīng)濟(jì)作為競(jìng)選演說(shuō)的重要話題，最終贏得選舉），在此，我想說(shuō)，“重要的是內(nèi)存，不是浮點(diǎn)運(yùn)算數(shù)（FLOPs），懂嗎？”

現(xiàn)代微處理器最大的瓶頸是能耗，而不是芯片集成度。Yahoo！創(chuàng)始人Mark Horowitz在十多年前就發(fā)現(xiàn)，訪問(wèn)片外DRAM的能耗是訪問(wèn)片上DRAM的100倍，是算術(shù)運(yùn)算的5000~10，000倍。因此，我們希望可以通過(guò)增加浮點(diǎn)運(yùn)算單元（FPU）來(lái)分?jǐn)們?nèi)存訪問(wèn)開(kāi)銷。基于Mark Horowitz的數(shù)據(jù)，芯片上的FPU數(shù)量被設(shè)置為10，000個(gè)左右。ML模型開(kāi)發(fā)人員常常試圖通過(guò)減少浮點(diǎn)運(yùn)算數(shù)來(lái)優(yōu)化模型，但其實(shí)減少內(nèi)存訪問(wèn)數(shù)才是更有效的辦法。

TPU v1有65，000多個(gè)乘法單元，比GPU、CPU等硬件高出許多倍。盡管它的時(shí)鐘頻率較低，僅為700MHz，但由于其乘法單元數(shù)量巨大，且每個(gè)乘法單元可進(jìn)行2個(gè)運(yùn)算操作，因此TPU v1每秒可執(zhí)行65，000×2×700M≈90 TeraOPS次操作。

上圖右側(cè)展示了運(yùn)算時(shí)的主要工作循環(huán)。65，000多個(gè)乘法器組成矩陣乘法單元（Matrix Multiply Unit）。計(jì)算時(shí)，首先啟動(dòng)累加器（Accumulator），然后通過(guò)激活函數(shù)管道（Activation Pipeline）進(jìn)行非線性函數(shù)運(yùn)算。累加器和激活函數(shù)輸出存儲(chǔ)（Activation Storage）是兩個(gè)主要功能單元之間的緩沖區(qū)。內(nèi)存（DDR3）向矩陣乘法單元輸入?yún)?shù)；最后，計(jì)算結(jié)果通過(guò)PCIe隊(duì)列返回服務(wù)器。

因此，TPU v1中主要的數(shù)據(jù)流動(dòng)如下圖紅色箭頭所示，此外的數(shù)據(jù)流動(dòng)還包括DDR3向其中輸入權(quán)重，以及計(jì)算輸出結(jié)果發(fā)送至主機(jī)。

TPU v1使用了脈動(dòng)陣列（systolic array），這一概念早在40年前就被提出，做法是以固定的時(shí)間間隔使數(shù)據(jù)從不同方向流入陣列中的處理單元（cell），最后將數(shù)據(jù)累積，以完成大型矩陣乘法運(yùn)算。由于70年代的芯片只有一個(gè)金屬層，不能很好地實(shí)現(xiàn)互連，所以Kung和Leiserson提出“脈動(dòng)陣列“以減少布線，簡(jiǎn)化連接。

現(xiàn)代芯片有多達(dá)10個(gè)金屬層，不存在這方面的問(wèn)題，其最大難點(diǎn)是能耗，而脈動(dòng)陣列的能效極高，使用脈動(dòng)陣列可以使芯片容納更多乘法單元，從而分?jǐn)們?nèi)存訪問(wèn)開(kāi)銷。

教訓(xùn)七：DSA既要專門(mén)優(yōu)化，也要靈活

作為一種針對(duì)特定領(lǐng)域的架構(gòu)（DSA），TPU的難點(diǎn)在于既要進(jìn)行針對(duì)性的優(yōu)化，同時(shí)還須保持一定的靈活性。Google在推出用于推理的TPU v1之后，決定攻克更難的問(wèn)題——訓(xùn)練。

訓(xùn)練之所以比推理更加復(fù)雜，是因?yàn)橛?xùn)練的計(jì)算量更大，包含反向傳播、轉(zhuǎn)置和求導(dǎo)等運(yùn)算。而且訓(xùn)練時(shí)需要將大量運(yùn)算結(jié)果儲(chǔ)存起來(lái)用于反向傳播的計(jì)算，因此也需要更大的內(nèi)存空間。

TPU v1只支持INT8計(jì)算，對(duì)訓(xùn)練而言動(dòng)態(tài)范圍不夠大，因此Google在TPU v2引入了一種的新的浮點(diǎn)格式BFloat16，用于機(jī)器學(xué)習(xí)計(jì)算。訓(xùn)練的并行化比推理的并行化更難。由于針對(duì)的是訓(xùn)練而非推理，所以TPU v2的可編程性也比TPU v1更高。

與TPU v1相比，TPU v2的改進(jìn)分為5步。第一步，TPU v1有兩個(gè)存儲(chǔ)區(qū)域：Accumulator和Activation Storage，前者負(fù)責(zé)儲(chǔ)存矩陣相乘結(jié)果，后者負(fù)責(zé)儲(chǔ)存激活函數(shù)輸出。

為了提升靈活性，TPU v2將上述兩個(gè)互相獨(dú)立的緩沖區(qū)調(diào)整位置后合并為向量存儲(chǔ)區(qū)（Vector Memory），從而提高可編程性，這也更類似傳統(tǒng)的內(nèi)存區(qū)。

第二步改進(jìn)針對(duì)的是激活函數(shù)管道（Activation Pipeline），TPU v1的管道內(nèi)包含一組負(fù)責(zé)非線性激活函數(shù)運(yùn)算的固定功能單元。TPU v2則將其改為可編程性更高的向量單元（Vector Unit），使其對(duì)編譯器和編程人員而言更易用。

第三步，將矩陣乘法單元直接與向量存儲(chǔ)區(qū)連接，如此一來(lái)，矩陣乘法單元就成為向量單元的協(xié)處理器。這種結(jié)構(gòu)對(duì)編譯器和編程人員而言更友好。

第四步，TPU v1使用DDR3內(nèi)存，因?yàn)樗槍?duì)的是推理，只需使用已有的權(quán)重，不需要生成權(quán)重。針對(duì)訓(xùn)練的TPU v2則不一樣，訓(xùn)練時(shí)既要讀取權(quán)重，也要寫(xiě)入權(quán)重，所以在v2中，我們將原本的DDR3改為與向量存儲(chǔ)區(qū)相連，這樣就既能向其讀取數(shù)據(jù)，又能向其寫(xiě)入數(shù)據(jù)。

然后，我們將DDR3改為HBM。因?yàn)閺腄DR3讀取參數(shù)速度太慢，影響性能，而HBM的讀寫(xiě)速度快20倍。

第五步，我們?cè)贖BM和向量存儲(chǔ)區(qū)之間增加互連（Interconnect），用于TPU之間的連接，組成我們之前提到的Pod超級(jí)計(jì)算機(jī)。以上就是從TPU v1到TPU v2的改進(jìn)。

教訓(xùn)八：半導(dǎo)體技術(shù)的發(fā)展速度參差不齊

回顧過(guò)去可以發(fā)現(xiàn)，各類技術(shù)的發(fā)展速度并不同步。計(jì)算邏輯的進(jìn)步速度很快，芯片布線的發(fā)展速度則較慢，而SRAM和HBM比DDR4和GDDR6的速度更快，能效更高。

上圖虛線框內(nèi)展示了單個(gè)Tensor Core運(yùn)算單元。TPU v2中有兩個(gè)互連的Tensor Core。

由于布線技術(shù)的進(jìn)步相對(duì)滯后，如果仍像TPU v1一樣，每塊芯片只有一個(gè)Tensor Core，就會(huì)導(dǎo)致管道更為冗長(zhǎng)，如果管道出了問(wèn)題也會(huì)更加麻煩。因此，我們將兩個(gè)Tensor Core互相連接，這對(duì)編譯器而言也更友好。

Google做出TPU v2之后，希望再花一年時(shí)間完善v2，所以TPU v3沒(méi)有引進(jìn)新技術(shù)，只是v2的改進(jìn)版。

與v2相比，TPU v3有以下特點(diǎn)：

體積只大了不到10%；

矩陣乘法單元（MXU）的數(shù)量翻倍，因此峰值性能也翻倍；

時(shí)鐘頻率加快了30%，進(jìn)一步加快計(jì)算速度；

內(nèi)存帶寬擴(kuò)大了30%；

容量翻倍，可使多種應(yīng)用更加方便；

芯片間帶寬擴(kuò)大30%；

可連接的節(jié)點(diǎn)數(shù)是之前的4倍。

上圖左上角即為T(mén)PU v1的主板。中間是v2，v2的散熱方式是風(fēng)冷，所以圖中可見(jiàn)高高突起的風(fēng)冷散熱器。右上角是v3，v3的運(yùn)行溫度太高，所以只能采用液冷。左下角是TPU v2組成的Pod超級(jí)計(jì)算機(jī)，共有256張TPU，峰值性能為11 PFLOP/s；右側(cè)的TPU v3 Pod有1024張TPU，峰值性能可達(dá)100 PFLOP/s（1 PFLOP/s即每秒1015次浮點(diǎn)運(yùn)算）。

從TPU v3到TPU v4i，矩陣乘法單元的數(shù)量再次翻倍，但芯片面積卻沒(méi)有擴(kuò)大。如前所述，計(jì)算邏輯的發(fā)展速度是最快的。

如果想了解TPU v4i，可以閱讀論文《Ten Lessons From Three Generations Shaped Google’s TPUv4i》。TPU v4i中，單個(gè)Tensor Core有4個(gè)矩陣乘法單元，是v3的兩倍，且v4i的片上內(nèi)存更大。此外，TPU v4i的時(shí)鐘頻率加快了10%，矩陣乘法單元中使用四位加法器（4-input adder），可以大幅節(jié)省芯片面積和功耗。

性能計(jì)數(shù)器（Performance counter）的重要性不言而喻，Google在v4i的很多地方都放置了性能計(jì)數(shù)器，可以更好地協(xié)助編譯器，并能更清楚地掌握運(yùn)行情況。

性能計(jì)數(shù)器沒(méi)有緩存，它們都在編譯器控制的4D DMA （直接存儲(chǔ)器訪問(wèn)）之下，并且可以進(jìn)行自定義互連。最后，為了控制更多的MPU（微處理器）和CMEM，VLIW（超長(zhǎng)指令字）指令拓寬到將近400位。

教訓(xùn)九：編譯器優(yōu)化和ML兼容性十分重要

XLA（加速線性代數(shù)）編譯器可對(duì)全程序進(jìn)行分析和優(yōu)化，優(yōu)化分為與機(jī)器無(wú)關(guān)的高級(jí)操作和與機(jī)器相關(guān)的低級(jí)操作，高級(jí)優(yōu)化操作將影響TPU和GPU，所以通常我們不會(huì)改動(dòng)高級(jí)優(yōu)化操作，以免導(dǎo)致失靈，但我們可以改動(dòng)低級(jí)優(yōu)化操作。

XLA編譯器可以處理多達(dá)4096個(gè)芯片的多核并行，2D向量和矩陣功能單元的數(shù)據(jù)級(jí)并行性，以及322~400位VLIW指令集的指令級(jí)并行。由于向量寄存器和計(jì)算單元是2D，這就要求功能單元和內(nèi)存中有良好的數(shù)據(jù)布局。此外，由于沒(méi)有緩存，所以編譯器需要管理所有的內(nèi)存?zhèn)鬏敗?/p>

最后的問(wèn)題是，與CPU相比，DSA的軟件棧還不夠成熟。那么編譯器優(yōu)化最終能夠提速多少？

實(shí)際上的提速相當(dāng)可觀。其中重要的優(yōu)化之一稱為算子融合（Operator Fusion），如將矩陣乘法與激活函數(shù)進(jìn)行融合，省略將中間結(jié)果寫(xiě)入HBM再讀取出來(lái)的步驟。上圖是我們的MLPerf基準(zhǔn)測(cè)試結(jié)果，可見(jiàn)，使用算子融合平均可以帶來(lái)超兩倍的性能提升。

上圖顯示了編譯器優(yōu)化的提速效果。藍(lán)色表示使用GPU，紅色表示使用TPU，經(jīng)過(guò)短短十幾個(gè)月的優(yōu)化后，不少模型的性能都提升到了兩倍。要知道，對(duì)C++編譯器而言，如果能在一年內(nèi)把性能提升5%-10%就已經(jīng)非常了不起了。

此外，編譯器的后向ML兼容性非常重要。我的同事Luiz Barroso主管Google的一個(gè)與計(jì)算機(jī)架構(gòu)無(wú)關(guān)的部門(mén)，他表示不希望在訓(xùn)練中花太多時(shí)間，希望一晚上就可以訓(xùn)練好模型，第二天可以直接部署。我們希望訓(xùn)練和推理時(shí)結(jié)果一致，這就是我們說(shuō)的后向ML兼容性，我們還希望它能在所有TPU上運(yùn)行，而不是每次更改TPU時(shí)都要重新訓(xùn)練。

為什么保持后向ML兼容性如此困難？因?yàn)楦↑c(diǎn)加法不符合結(jié)合律，所以運(yùn)算順序可能會(huì)影響運(yùn)算結(jié)果。而TPU的任務(wù)就是讓所有機(jī)器對(duì)編譯器而言都沒(méi)有區(qū)別，以便可以在重組代碼的同時(shí)獲得相同的高質(zhì)量結(jié)果，以實(shí)現(xiàn)后向ML兼容性。

教訓(xùn)十：優(yōu)化的目標(biāo)是Perf/TCO還是Perf/CapEx

在將研究成果應(yīng)用到實(shí)際生產(chǎn)時(shí)，我們優(yōu)化的目標(biāo)是什么？Google構(gòu)建硬件是為了用在自己的數(shù)據(jù)中心，所以我們所要控制的成本是指總體擁有成本（TCO），包括資本成本（采購(gòu)成本）和運(yùn)行成本（電力、冷卻、空間成本）。資金成本是一次性的，而運(yùn)行成本需要持續(xù)支出3~4年。

因此，芯片和主板生產(chǎn)商只需要考慮產(chǎn)品性能/資本成本的比率；而Google卻要考慮整個(gè)硬件生命周期的成本，關(guān)注性能/總體擁有成本之間的比率。如上面的餅狀圖所示，電力可占總體擁有成本的一半。所以，如果把眼光擴(kuò)大到總體擁有成本上，在系統(tǒng)設(shè)計(jì)時(shí)就很可能會(huì)做出不同的取舍。

之前提到，TPU v1有一個(gè)Tensor Core，v2和v3有兩個(gè)。到了v4時(shí)，基于對(duì)總體擁有成本的考慮，Google決定分開(kāi)設(shè)計(jì)：用于訓(xùn)練的TPU v4有兩個(gè)Tensor Core，用于推理的TPU v4i只有一個(gè)。這樣就大大提升了性能和總體擁有成本之間的比率。

上圖是TPU v4i和TPU v3的每瓦性能對(duì)比，紅色是TPU v4i，藍(lán)色是TPU v3，前者的每瓦性能是后者的兩倍以上。

03 提升機(jī)器學(xué)習(xí)能效，減少碳足跡

2021年10月的IEEE Spectrum雜志有一篇文章提到，訓(xùn)練某一模型需要數(shù)年時(shí)間，花費(fèi)1000億美元，總碳排放量相當(dāng)于紐約一個(gè)月的排碳量，如果還要進(jìn)一步提升模型準(zhǔn)確度，這些數(shù)字還會(huì)更夸張。

2022年1月，又有文章表示，根據(jù)當(dāng)前的算力需求增長(zhǎng)曲線預(yù)計(jì)，到2026年，訓(xùn)練最大AI模型的成本將相當(dāng)于美國(guó)的GDP，大概是20萬(wàn)億美元。

Google研究了ML硬件的能源消耗。不同于全生命周期消耗的能源（包括從芯片制造到數(shù)據(jù)中心構(gòu)建的所有間接碳排放），我們只關(guān)注硬件運(yùn)行時(shí)的能源消耗。

展望未來(lái)，我們有辦法讓機(jī)器學(xué)習(xí)的能耗降低到原來(lái)的100倍，碳排放量降低1000倍。我們可以從四方面協(xié)同著手，極大地促進(jìn)機(jī)器學(xué)習(xí)在更多領(lǐng)域的可持續(xù)發(fā)展：

第一個(gè)因素是模型。Google在2017年公布Transformer模型，四年后，又開(kāi)發(fā)了Primer模型，其計(jì)算質(zhì)量相同，但能效更高。Primer的能耗和碳排放量相比Transformer降低了4倍。

第二個(gè)因素是硬件。2017年所使用的P100 GPU和當(dāng)前最新TPU的性能相差了14倍。所以，前兩個(gè)因素結(jié)合，可以將能耗和碳排放量降低60倍。

第三個(gè)因素是數(shù)據(jù)中心的能效。Google的PUE大約是其他數(shù)據(jù)中心的1.4倍。所以，前三個(gè)因素累積，可以將能耗和碳排放量降低80倍。

第四個(gè)因素是數(shù)據(jù)中心的地理位置。即使在同一個(gè)國(guó)家，不同地區(qū)使用無(wú)碳能源的比例也可能大不相同。在Google所有數(shù)據(jù)中心所在地中，俄克拉荷馬州的無(wú)碳能源占比最高，Primer模型就是在此處訓(xùn)練的，這可將碳排放量在上述基礎(chǔ)上降低9倍。

綜合上述四個(gè)因素，我們可將機(jī)器學(xué)習(xí)的能耗降低80倍，碳排放量降低700倍。這非常了不起。

OpenAI的GPT-3問(wèn)世后引起了Google所有機(jī)器學(xué)習(xí)工程師的注意，他們都鉚足了勁想做得更好。18個(gè)月后，Google推出了GlaM模型，在相同的基準(zhǔn)測(cè)試中它的表現(xiàn)比GPT-3更好。GlaM的參數(shù)是GPT-3的七倍，達(dá)到1.2萬(wàn)億，但它的能耗卻并不大，因?yàn)樗昧讼∈栊?。GlaM是一個(gè)MoE模型（Mixture of Experts，專家混合模型），它平時(shí)只調(diào)用每個(gè)token中的8%的參數(shù)，而密集型模型會(huì)使用100%的參數(shù)。因此，GlaM中加速器的工作時(shí)長(zhǎng)和能耗都降低了3倍。

最后，與GPT-3不同的是，GlaM在俄克拉荷馬州使用清潔能源進(jìn)行訓(xùn)練，因此累計(jì)下來(lái)，其碳排放量降低了14倍。所以GlaM的例子表明，相比V100 GPU，使用TPU v4既減少了碳排放量，而且計(jì)算質(zhì)量更好。

審核編輯 ：李倩