Google TPU架構演進及十大經(jīng)驗

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

2017年，David Patterson加入Google TPU團隊，2018年3月，他與John Hennessy共同獲得圖靈獎，2008年獲ACM/IEEE Eckert-Mauchly 獎（被譽為計算機體系結構最高獎），2000年獲得馮·諾依曼獎章。

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

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

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

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

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

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

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

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

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

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

“AI技術的進步有賴于計算基礎設施的支持，而TPU正是Google計算基礎設施的重要部分。新一代TPU v4芯片的速度是v3的兩倍多。Google用TPU集群構建出Pod超級計算機，單臺TPU v4 Pod包含4096塊v4芯片，每臺Pod的芯片間互連帶寬是其他互連技術的10倍，因此，TPU v4 Pod的算力可達1 ExaFLOP，即每秒執(zhí)行10的18次方浮點運算，相當于1000萬臺筆記本電腦的總算力。”

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

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

02 十年演進，十大教訓

過往十年，我們在ML計算架構的發(fā)展中汲取了十大教訓。

其中，前五個都和ML模型本身有關，后五個則關乎硬件和架構。這些經(jīng)驗對深度學習以外的領域也有借鑒意義。

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

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

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

教訓二：DNN工作負載隨著DNN突破不斷演變

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

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

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

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

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

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

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

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

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

DNN需要使用多租戶技術（multi-tenancy）。不少深度學習論文的一個假設是同一時間只需運行一個模型，但在實際應用中，有不少情況都需要在不同模型中切換。

比如，機器翻譯涉及各種語言對，就需要不同的模型；傳統(tǒng)的軟件開發(fā)需要用到一個主模型和多個實驗模型；甚至有時因為對吞吐量和延遲有不同的側重要求，就需要不同的batch size，進而需要不同的模型。

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

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

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

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

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

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

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

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

現(xiàn)代芯片有多達10個金屬層，不存在這方面的問題，其最大難點是能耗，而脈動陣列的能效極高，使用脈動陣列可以使芯片容納更多乘法單元，從而分攤內(nèi)存訪問開銷。

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

作為一種針對特定領域的架構（DSA），TPU的難點在于既要進行針對性的優(yōu)化，同時還須保持一定的靈活性。Google在推出用于推理的TPU v1之后，決定攻克更難的問題——訓練。

訓練之所以比推理更加復雜，是因為訓練的計算量更大，包含反向傳播、轉置和求導等運算。而且訓練時需要將大量運算結果儲存起來用于反向傳播的計算，因此也需要更大的內(nèi)存空間。

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

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

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

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

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

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

然后，我們將DDR3改為HBM。因為從DDR3讀取參數(shù)速度太慢，影響性能，而HBM的讀寫速度快20倍。

第五步，我們在HBM和向量存儲區(qū)之間增加互連（Interconnect），用于TPU之間的連接，組成我們之前提到的Pod超級計算機。以上就是從TPU v1到TPU v2的改進。

教訓八：半導體技術的發(fā)展速度參差不齊

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

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

由于布線技術的進步相對滯后，如果仍像TPU v1一樣，每塊芯片只有一個Tensor Core，就會導致管道更為冗長，如果管道出了問題也會更加麻煩。因此，我們將兩個Tensor Core互相連接，這對編譯器而言也更友好。

Google做出TPU v2之后，希望再花一年時間完善v2，所以TPU v3沒有引進新技術，只是v2的改進版。

與v2相比，TPU v3有以下特點：

體積只大了不到10%；

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

時鐘頻率加快了30%，進一步加快計算速度；

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

容量翻倍，可使多種應用更加方便；

芯片間帶寬擴大30%；

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

上圖左上角即為TPU v1的主板。中間是v2，v2的散熱方式是風冷，所以圖中可見高高突起的風冷散熱器。右上角是v3，v3的運行溫度太高，所以只能采用液冷。左下角是TPU v2組成的Pod超級計算機，共有256張TPU，峰值性能為11 PFLOP/s；右側的TPU v3 Pod有1024張TPU，峰值性能可達100 PFLOP/s（1 PFLOP/s即每秒1015次浮點運算）。

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

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

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

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

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

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

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

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

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

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

此外，編譯器的后向ML兼容性非常重要。我的同事Luiz Barroso主管Google的一個與計算機架構無關的部門，他表示不希望在訓練中花太多時間，希望一晚上就可以訓練好模型，第二天可以直接部署。我們希望訓練和推理時結果一致，這就是我們說的后向ML兼容性，我們還希望它能在所有TPU上運行，而不是每次更改TPU時都要重新訓練。

為什么保持后向ML兼容性如此困難？因為浮點加法不符合結合律，所以運算順序可能會影響運算結果。而TPU的任務就是讓所有機器對編譯器而言都沒有區(qū)別，以便可以在重組代碼的同時獲得相同的高質(zhì)量結果，以實現(xiàn)后向ML兼容性。

教訓十：優(yōu)化的目標是Perf/TCO還是Perf/CapEx

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

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

之前提到，TPU v1有一個Tensor Core，v2和v3有兩個。到了v4時，基于對總體擁有成本的考慮，Google決定分開設計：用于訓練的TPU v4有兩個Tensor Core，用于推理的TPU v4i只有一個。這樣就大大提升了性能和總體擁有成本之間的比率。

上圖是TPU v4i和TPU v3的每瓦性能對比，紅色是TPU v4i，藍色是TPU v3，前者的每瓦性能是后者的兩倍以上。

03 提升機器學習能效，減少碳足跡

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

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

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

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

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

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

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

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

綜合上述四個因素，我們可將機器學習的能耗降低80倍，碳排放量降低700倍。這非常了不起。

OpenAI的GPT-3問世后引起了Google所有機器學習工程師的注意，他們都鉚足了勁想做得更好。18個月后，Google推出了GlaM模型，在相同的基準測試中它的表現(xiàn)比GPT-3更好。GlaM的參數(shù)是GPT-3的七倍，達到1.2萬億，但它的能耗卻并不大，因為它利用了稀疏性。GlaM是一個MoE模型（Mixture of Experts，專家混合模型），它平時只調(diào)用每個token中的8%的參數(shù)，而密集型模型會使用100%的參數(shù)。因此，GlaM中加速器的工作時長和能耗都降低了3倍。

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

審核編輯 ：李倩