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

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

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

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

在這方面,我們這輩子能經(jīng)歷的最重要的例子莫過于神經(jīng)網(wǎng)絡(luò)了。

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

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

然而,正因?yàn)闆]有合適的硬件,神經(jīng)網(wǎng)絡(luò)也work不起來......

在歷史的周期演進(jìn)中,構(gòu)建專用硬件的項(xiàng)目一次又一次地被放棄,最終,作為外來者的GPU打破了僵局。

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

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

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

當(dāng)我們在談?wù)揟PU時,我們在談?wù)撌裁?/p>

為什么谷歌的TPU既有趣又重要?因?yàn)檫@是谷歌,TPU被切實(shí)應(yīng)用于谷歌龐大的服務(wù)(搜索、Android、Chrome、Gmail、地圖、Play 商店、YouTube、照片),用戶超過10億。此外,谷歌還擁有第三大公共云。

在英偉達(dá)的GPU主導(dǎo)了這個AI時代硬件供應(yīng)的當(dāng)下,谷歌的TPU是一個真正經(jīng)歷了時間和規(guī)模考驗(yàn)的競品。

以下的內(nèi)容,小編分成兩部分:第一部分講故事,關(guān)于第一個谷歌TPU的前世今生;第二部分講技術(shù),揭秘TPU的架構(gòu)細(xì)節(jié)和性能。

起源

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

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

2010年代初,谷歌的注意力開始轉(zhuǎn)向深度學(xué)習(xí):

2011年:Jeff Dean、Greg Corrado和Andrew Ng發(fā)起了關(guān)于深度學(xué)習(xí)的研究項(xiàng)目——Google Brain。

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

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

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

于是,在公司工作的第一天,他出去從當(dāng)?shù)氐囊患译娮由痰曩I了一臺GPU機(jī)器,把它插入網(wǎng)絡(luò),然后開始在GPU上訓(xùn)練神經(jīng)網(wǎng)絡(luò)。

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

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

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

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

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

快速交付

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

谷歌迅速招募了一支才華橫溢、經(jīng)驗(yàn)豐富的團(tuán)隊(duì),其中包括David Patterson,——伯克利RISC原始設(shè)計的開發(fā)者,也是RISC-V指令集架構(gòu)開發(fā)的關(guān)鍵人物。

為了趕時間,谷歌沒有去從頭開發(fā)新的架構(gòu)。幸運(yùn)的是,在35年前,就有人為他們準(zhǔn)備好了合適的架構(gòu)。

1978年,卡內(nèi)基梅隆大學(xué)的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是一個處理器網(wǎng)絡(luò),它有節(jié)奏地計算并通過系統(tǒng)傳遞數(shù)據(jù)......處理器的功能類似于心臟,每個處理器都會定期將數(shù)據(jù)泵入和泵出,每次都執(zhí)行一些簡短的計算,以便在網(wǎng)絡(luò)中保持常規(guī)的數(shù)據(jù)流。

作者同時提出了架構(gòu)的一種應(yīng)用:矩陣計算。

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ù)組結(jié)構(gòu)的脈動網(wǎng)絡(luò)上優(yōu)雅而有效地執(zhí)行流水線。例如,六邊形連接的處理器可以最佳地執(zhí)行矩陣計算......這些脈動陣列享有簡單而規(guī)則的通信路徑,并且網(wǎng)絡(luò)中使用的幾乎所有處理器都是相同的。因此,使用VLSI技術(shù)可以廉價地構(gòu)建基于脈動陣列的專用硬件設(shè)備。

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

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

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

2015年初,第一批TPU正式部署在谷歌的數(shù)據(jù)中心,此時,距離項(xiàng)目啟動僅僅過去了15個月,看一下這個龐大的研發(fā)團(tuán)隊(duì):

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

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

2016年3月,使用了TPU的AlphaGo Zero擊敗了當(dāng)時的世界圍棋冠軍李世石,舉世震驚。

TPU架構(gòu)

在神經(jīng)網(wǎng)絡(luò)的運(yùn)算中,矩陣乘法是關(guān)鍵,下面是一個簡單的例子:

那么,TPU是如何計算矩陣乘法的呢?

在上面的故事中,我們提到了TPU的原理:systolic arrays(脈動陣列),假設(shè)有如下的2*2矩陣乘法:

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

在此圖中,黃色值是從頂部和左側(cè)輸入矩陣的輸入。淺藍(lán)色值是存儲的部分總和。深藍(lán)色值是最終結(jié)果。

第一步,a11和b11加載到左上角的MAC中,相乘并存儲結(jié)果。

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

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

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

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

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

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

更詳細(xì)的設(shè)計可以擴(kuò)展成下圖這樣子:

DDR3 DRAM / Weight FIFO:權(quán)重存儲在通過DDR3-2133接口連接到TPU v1的DDR3 RAM芯片中。權(quán)重通過PCIe從主機(jī)的內(nèi)存預(yù)加載,然后可以傳輸?shù)綑?quán)重FIFO存儲器中,供矩陣乘法單元使用。

Matrix Multiply Unit:256 x 256大小的矩陣乘法單元,頂部輸入256個權(quán)重值,左側(cè)是256個input值。

Accumulators:運(yùn)算結(jié)果從脈動陣列的底部匯總到累加器中(內(nèi)存)。

Activation:激活函數(shù)。

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

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

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

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

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

Read_Host_MemoryRead_WeightsLoop_Start    Matrix_Multiply    ActivateLoop_EndWrite_Host_Memory

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

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

2013年,TPU v1與英特爾的Haswell CPU和英偉達(dá)的K80 GPU進(jìn)行了比較:

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

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

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

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

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

文章來源:新智元

 

新智元

網(wǎng)友評論

聚超值•精選

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