首頁 > 科技要聞 > 科技> 正文

芯片戰(zhàn)爭早已打響！谷歌15個月打造首個TPU，欲和老黃平起平坐

新智元 整合編輯：甘靈文 發(fā)布于：2024-04-01 17:14

在計算的歷史上，曾被丟棄或過時的想法仍然很有趣，有時甚至非常有用。

在這方面，我們這輩子能經歷的最重要的例子莫過于神經網絡了。

多數(shù)人可能都了解，在神經網絡近70年的歷史中，寒冬和泡沫交替出現(xiàn)，——事實上，藏在神經網絡背后的專用硬件加速器（ASIC）也是如此。

神經網絡和ASIC屬于是「先有雞還是先有蛋」的關系，大家都知道神經網絡在CPU上效率不高，但是，如果沒有證據證明神經網絡是有用的，公司憑什么要為神經網絡投資開發(fā)專門的硬件？

然而，正因為沒有合適的硬件，神經網絡也work不起來......

在歷史的周期演進中，構建專用硬件的項目一次又一次地被放棄，最終，作為外來者的GPU打破了僵局。

毫無疑問，我們現(xiàn)在正處于泡沫周期，在大模型的激勵下，從初創(chuàng)公司到超大規(guī)模企業(yè)，都在構建建立自己的硬件加速器。

他們之中的大多數(shù)都可能會失敗，但也肯定會有成功的設計成為未來計算領域的重要組成部分。

——而谷歌的張量處理單元（TPU）作為第一批大規(guī)模部署的加速器，肯定會成為幸存者之一。

當我們在談論TPU時，我們在談論什么

為什么谷歌的TPU既有趣又重要？因為這是谷歌，TPU被切實應用于谷歌龐大的服務（搜索、Android、Chrome、Gmail、地圖、Play 商店、YouTube、照片），用戶超過10億。此外，谷歌還擁有第三大公共云。

在英偉達的GPU主導了這個AI時代硬件供應的當下，谷歌的TPU是一個真正經歷了時間和規(guī)模考驗的競品。

以下的內容，小編分成兩部分：第一部分講故事，關于第一個谷歌TPU的前世今生；第二部分講技術，揭秘TPU的架構細節(jié)和性能。

起源

機器學習對谷歌來說是一件大事。畢竟谷歌的既定使命是「組織世界信息，使其普遍可用和有用（to organize the world's information and make it universally accessible and useful）」。

機器學習幫助谷歌發(fā)掘信息的價值，從圖像和語音識別到語言翻譯，再到大型語言模型，當然也少不了谷歌的「搖錢樹」——年入數(shù)十億美元的廣告業(yè)務。

2010年代初，谷歌的注意力開始轉向深度學習：

2011年：Jeff Dean、Greg Corrado和Andrew Ng發(fā)起了關于深度學習的研究項目——Google Brain。

2013年：繼AlexNet圖像識別項目取得成功后，谷歌以4400萬美元的價格收購了由Geoffrey Hinton、Alex Krizhevsky和Ilya Sutskever組成的初創(chuàng)公司。

2014年：谷歌收購了由Demis Hassabis、Shane Legg和Mustafa Suleyman創(chuàng)立的DeepMind，價格高達65000萬美元

2013年，當AlexNet的開發(fā)者Alex Krizhevsky來到谷歌時，他發(fā)現(xiàn)谷歌現(xiàn)有的模型都在CPU上運行。

于是，在公司工作的第一天，他出去從當?shù)氐囊患译娮由痰曩I了一臺GPU機器，把它插入網絡，然后開始在GPU上訓練神經網絡。

最終，大家意識到需要自己需要更多的GPU，2014年，谷歌決定以約13000萬美元的價格購買40,000個英偉達GPU。

深度學習所提供的能力可以大規(guī)模應用于谷歌的各種業(yè)務，然而，無論是在技術上還是在戰(zhàn)略上，依賴英偉達GPU都不一定是最佳解決方案。

GPU不是ASIC，它不是為神經網絡量身打造的，對于需要大規(guī)模應用的谷歌來說，相當于要付出很大的額外代價；同時，依賴單一供應商顯然也是戰(zhàn)略上的重大風險。

谷歌當時有兩種選擇：現(xiàn)場可編程門陣列（FPGA）和專用集成電路（ASIC）。

當時的實驗證明，F(xiàn)PGA的性能打不過GPU，而定制ASIC可能在推理方面產生10倍的性價比優(yōu)勢。

快速交付

開發(fā)ASIC的時間成本是一個重要的考量，如果落地周期太長，那么硬件本身也就失去了意義。為此谷歌動用了一切可以快速訪問的資源和現(xiàn)有知識。

谷歌迅速招募了一支才華橫溢、經驗豐富的團隊，其中包括David Patterson，——伯克利RISC原始設計的開發(fā)者，也是RISC-V指令集架構開發(fā)的關鍵人物。

為了趕時間，谷歌沒有去從頭開發(fā)新的架構。幸運的是，在35年前，就有人為他們準備好了合適的架構。

1978年，卡內基梅隆大學的H.T.Kung和Charles E. Leiserson發(fā)表了《Systolic Arrays (for VLSI)》，提出了「systolic system」。

論文地址：https://www.eecs.harvard.edu/htk/static/files/1978-cmu-cs-report-kung-leiserson.pdf

A systolic system is a network of processors which rhythmically compute and pass data through the system….In a systolic computer system, the function of a processor is analogous to that of the heart. Every processor regularly pumps data in and out, each time performing some short computation so that a regular flow of data is kept up in the network.

systolic system是一個處理器網絡，它有節(jié)奏地計算并通過系統(tǒng)傳遞數(shù)據......處理器的功能類似于心臟，每個處理器都會定期將數(shù)據泵入和泵出，每次都執(zhí)行一些簡短的計算，以便在網絡中保持常規(guī)的數(shù)據流。

作者同時提出了架構的一種應用：矩陣計算。

Many basic matrix computations can be pipelined elegantly and efficiently on systolic networks having an array structure. As an example, hexagonally connected processors can optimally perform matrix computation......These systolic arrays enjoy simple and regular communication paths, and almost all processors used in the network are identical. As a result, special purpose hardware devices based on systolic arrays can be built inexpensively using the VLSI technology.

許多基本的矩陣計算可以在具有數(shù)組結構的脈動網絡上優(yōu)雅而有效地執(zhí)行流水線。例如，六邊形連接的處理器可以最佳地執(zhí)行矩陣計算......這些脈動陣列享有簡單而規(guī)則的通信路徑，并且網絡中使用的幾乎所有處理器都是相同的。因此，使用VLSI技術可以廉價地構建基于脈動陣列的專用硬件設備。

——這不就來了嗎！簡單、便宜、還適合算矩陣！真是老天爺喂飯吃。

有了工程師和架構之后，谷歌還與當時的LSI公司（現(xiàn)在是Broadcom的一部分）合作，幫助芯片落地。

另外，全新的芯片意味著需要全新的指令集架構、全新的編譯器以及上層軟件，這是一個龐大的工程。

2015年初，第一批TPU正式部署在谷歌的數(shù)據中心，此時，距離項目啟動僅僅過去了15個月，看一下這個龐大的研發(fā)團隊：

論文地址：https://arxiv.org/ftp/arxiv/papers/1704/1704.04760.pdf

為了這個進度，谷歌也不得不做了很多取舍，包括使用較老的工藝（28nm），以及較低的時鐘頻率（700MHz）.

2016年3月，使用了TPU的AlphaGo Zero擊敗了當時的世界圍棋冠軍李世石，舉世震驚。

TPU架構

在神經網絡的運算中，矩陣乘法是關鍵，下面是一個簡單的例子：

那么，TPU是如何計算矩陣乘法的呢？

在上面的故事中，我們提到了TPU的原理：systolic arrays（脈動陣列），假設有如下的2*2矩陣乘法：

計算結果如下圖所示。每個角上的方塊代表一個乘法/累加單位 （MAC），可以執(zhí)行乘法和加法運算。

在此圖中，黃色值是從頂部和左側輸入矩陣的輸入。淺藍色值是存儲的部分總和。深藍色值是最終結果。

第一步，a11和b11加載到左上角的MAC中，相乘并存儲結果。

第二步，a12和b21加載到左上角的MAC中，相乘并添加到先前計算的結果中。這一步得到了結果矩陣的左上角值。同時，b11被傳輸?shù)接疑辖堑腗AC，乘以新加載的a21，并存儲結果；a11被傳輸?shù)阶笙陆堑腗AC，乘以新加載的b12，并存儲結果；

第三步，b21被傳輸?shù)接疑辖堑腗AC，乘以新加載的值a22，結果被添加到以前存儲的結果中；a12被傳輸?shù)阶笙陆堑腗AC，乘以新加載的b22，并將結果添加到先前存儲的結果中。此時得到了結果矩陣的右上角和左下角值。同時，a12和b21被傳輸?shù)接蚁陆堑腗AC，相乘并存儲結果。

第四步，將a22和b22傳輸?shù)接蚁陆堑腗AC，相乘并將結果添加到先前存儲的值中，從而得到結果矩陣的右下角值。

至此，2*2矩陣乘法完成。

完整計算系統(tǒng)的最簡單表示如下：

首先要注意的是，TPUv1依賴于通過PCIe（高速串行總線）接口與主機進行通信。它還可以直接訪問自己的DDR3存儲，

更詳細的設計可以擴展成下圖這樣子：

DDR3 DRAM / Weight FIFO：權重存儲在通過DDR3-2133接口連接到TPU v1的DDR3 RAM芯片中。權重通過PCIe從主機的內存預加載，然后可以傳輸?shù)綑嘀谾IFO存儲器中，供矩陣乘法單元使用。

Matrix Multiply Unit：256 x 256大小的矩陣乘法單元，頂部輸入256個權重值，左側是256個input值。

Accumulators：運算結果從脈動陣列的底部匯總到累加器中（內存）。

Activation：激活函數(shù)。

Unified Buffer / Systolic Data Setup：應用激活函數(shù)的結果存儲在統(tǒng)一緩沖區(qū)存儲器中，然后可以作為輸入反饋到矩陣乘法單元，以計算下一層所需的值。

TPU v1執(zhí)行8位整數(shù)乘法，利用量化來避免消耗更大的浮點計算。

TPU v1使用CISC（復雜指令集）設計，只有大約20條指令。這些指令是由主機通過PCIe接口發(fā)送給TPU的，而不是從內存中獲取的。

矩陣運算采用可變大小的B*256輸入，將其乘以256x256恒定權重輸入，生成B*256輸出，所以需要B次流水線循環(huán)才能完成。

TPU的整個執(zhí)行過程大概像這樣子：

Read_Host_MemoryRead_WeightsLoop_Start    Matrix_Multiply    ActivateLoop_EndWrite_Host_Memory

由于TPU v1的指令集非常簡單，可以保證芯片用于解碼和相關活動的開銷非常低，只有2%的芯片區(qū)域專門用于控制。

而24%的芯片面積專用于矩陣乘法單元，29%用于存儲輸入和中間結果的統(tǒng)一緩沖區(qū)存儲器。

2013年，TPU v1與英特爾的Haswell CPU和英偉達的K80 GPU進行了比較：

TPU v1的MAC數(shù)量是K80 GPU的25倍，片上內存是K80 GPU的3.5倍。

TPU v1的推理速度比K80 GPU和Haswell CPU快15到30倍。

TPU v1的相對計算效率是GPU的25到29倍。

時至今日，這個「臨時趕工」的ASIC，已經不斷完善和壯大，在英偉達的統(tǒng)治之下，開辟了另一條道路。

參考資料：https://thechipletter.substack.com/p/googles-first-tpu-architecture

文章來源：新智元

 
芯片   英偉達   谷歌
新智元

網友評論

聚超值•精選

推薦 手機 筆記本 影像 硬件 家居 商用 企業(yè) 出行 未來
二維碼 回到頂部