search
尋找貓咪~QQ 地點 桃園市桃園區 Taoyuan , Taoyuan

Google深度揭秘TPU:一文看懂內部原理,以及為何碾壓GPU

李林 舒石 編譯整理

量子位 報道 | 公眾號 QbitAI

原文發佈於Google雲

作者:

谷歌雲 Kaz Sato, Staff Developer Advocate

谷歌大腦軟體工程師Cliff Young, Software Engineer

谷歌大腦傑出工程師David Patterson

搜索、街景、照片、翻譯,這些Google提供的服務,都使用了Google的TPU(張量處理器)來加速背後的神經網路計算。

在PCB板上的Google首款TPU和部署了TPU的數據中心

去年Google推出TPU並在近期對這一晶元的性能和架構進行了詳細的研究。簡單的結論是:TPU與同期的CPU和GPU相比,可以提供15-30倍的性能提升,以及30-80倍的效率(性能/瓦特)提升。

這意味著,Google的服務既可以大規模運行於最先進的神經網路,而且可以把成本控制在可接受的程度上。以下的內容,將深入解讀Google TPU中的技術,並且討論如何實現更為出色的性能。

早在2006年,Google就在考慮為神經網路構建一個專用集成電路(ASIC)。2013年這個需求變得更加緊迫,當時Google意識到快速增長的計算需求,可能意味著數據中心的數量需要翻番才能滿足。

通常而言,ASIC的開發需要耗時數年。但具體到TPU而言,從設計到驗證、構建和部署到數據中心里,只需要15個月。

TPU ASIC採用了28nm工藝製造,主頻700MHz,功耗40W。為了儘快把TPU部署到現有的伺服器中,Google選擇把這個晶元打包成外部擴展加速器,然後插到SATA硬碟插槽里使用。所以TPU通過PCIe Gen3 x16匯流排與主機相連,也就是說12.5GB/s的有效帶寬。

要說明TPU的設計思路,需要先來簡介一下神經網路的計算。

這是一個TensorFlow Playground的例子。用以訓練一個神經網路,以標籤對數據進行分類,或者對缺失數據進行估計,或者推斷未來的數據。對於推斷來說,神經網路中的每個神經元都進行如下計算:

  • 輸入數據(x)乘以權重(w)以表示信號強度

  • 乘積加總,成為代表神經元狀態的唯一值

  • 應用激活函數(f),例如ReLU、Sigmoid等調節神經元

神經網路把輸入數據與權重矩陣相乘,並輸入激活函數

例如,對於有三個輸入數據和兩個全連接神經元的單層神經網路而言,需要把輸入和權重進行六次相乘,並得出兩組乘積之和。這個乘法和加法序列,可以寫成一個矩陣乘法,然後通過激活函數進一步處理矩陣的輸出。

在更複雜的神經網路架構中,乘法矩陣通常也是計算量最大的部分。

實際業務中需要多少次乘法運算?2016年7月,Google團隊調查了實際業務中,六個有代表性的神經網路應用,結果如下表所示:

如上表所示,每個神經網路中的權重數量從500萬到1億不等。每一個預測,都需要許多步的輸入數據和權重矩陣相乘,並輸入到激活函數中。

總而言之,計算量超大。作為優化的第一步,Google應用了一種稱為量化的技術進行整數運算,而不是在CPU或者GPU上對所有數學工作進行32位或者16位浮點運算。這能減少所需的內存容量和計算資源。

通常而言,神經網路的預測不需要32位或16浮點計算精度,通過一些方法,可以用8位整數對神經網路進行預測,並保持適當的準確度。

所謂量化,就是一種使用8位整數來近似預設的最小值和最大值之間任意數值的優化技術。

TensorFlow中的量化

量化是降低神經網路預測成本的利器,同時帶來的內存減少也很重要,特別是對於移動和嵌入式部署。舉個例子,在Inception中應用量化之後,這個圖像識別模型能從91MB壓縮到23MB,成功瘦身四分之三。

使用整數而不是浮點計算,大大減小了TPU的硬體尺寸和功耗。一個TPU鍾包含65,536個8位整數乘法器。雲環境中使用的主流GPU,通常包含數千個32位浮點乘法器。只要能用8位滿足精度需求,就能帶來25倍以上的性能提升。

可編程性是TPU的另一個重要設計目標。TPU不是設計用來運行某一種神經網路,而是要能加速許多不同類型的模型。

大多數當代CPU都採用了精簡指令集(RISC)。但Google選擇複雜指令集(CISC)作為TPU指令集的基礎,這一指令集側重於運行更複雜的任務。

我們來看看TPU的結構圖。

TPU包括以下計算資源:

  • 矩陣乘法單元(MUX):65,536個8位乘法和加法單元,運行矩陣計算

  • 統一緩衝(UB):作為寄存器工作的24MB容量SRAM

  • 激活單元(AU):硬體連接的激活函數

為了控制MUX、UB和AU進行計算,Google定義了十幾個專門為神經網路推理而設計的高級指令。以下是五個例子。

簡而言之,TPU設計封裝了神經網路計算的本質,可以針對各種神經網路模型進行編程。為了編程,Google還創建了一個編譯器和軟體棧,將來自TensorFlow圖的API調用,轉化成TPU指令。

從TensorFlow到TPU:軟體堆棧

典型的RISC處理器提供簡單計算的指令,例如乘法或加法。這些事所謂的標量(Scalar)處理器,因為它們每個指令處理單一運算,即標量運算。

即使主頻千兆赫茲的CPU,仍然需要很長時間才能通過一系列標量運算來完成大型矩陣的計算。改進的方法就是矢量(Vector)運算,同時針對多個數據元素執行相同的操作。

GPU的流處理器(SM)就是一種高效的向量處理器,贊單個時鐘周期內,可以處理數百到數千次運算。

至於TPU,Google為其設計了MXU作為矩陣處理器,可以在單個時鐘周期內處理數十萬次運算,也就是矩陣(Matrix)運算。

MXU有著與傳統CPU、GPU截然不同的架構,稱為脈動陣列(systolic array)。之所以叫「脈動」,是因為在這種結構中,數據一波一波地流過晶元,與心臟跳動供血的方式類似。

如圖所示,CPU和GPU在每次運算中都需要從多個寄存器(register)中進行存取;而TPU的脈動陣列將多個運算邏輯單元(ALU)串聯在一起,復用從一個寄存器中讀取的結果。

MXU中的權值陣列專門為矩陣乘法運算進行了優化,並不適用於通用計算。

脈動陣列中,輸入向量與權值矩陣相乘

脈動陣列中,輸入矩陣與權值矩陣相乘

MXU的脈動陣列包含256 × 256 = 65,536個ALU,也就是說TPU每個周期可以處理65,536次8位整數的乘法和加法。

TPU以700兆赫茲的功率運行,也就是說,它每秒可以運行65,536 × 700,000,000 = 46 × 1012次乘法和加法運算,或每秒92萬億(92 × 1012)次矩陣單元中的運算。

TPU中的MXU


每周期運算量
CPU數個
CPU (向量擴展)數十
GPU數萬
TPU數十萬

這種基於複雜指令集計算(CISC)的矩陣運算設計,實現了出色的性能功耗比:TPU的性能功耗比,比同時期的CPU強83倍,比同時期的GPU強29倍。

極簡這一點,在Google之前發布的TPU論文第8頁提到過。與CPU和GPU相比,單用途的TPU就是一個單線程晶元,不需要考慮緩存、分支預測、多道處理等問題。

TPU的設計之簡潔,從沖模平面圖上就能看出來:

黃色代表運算單元;藍色是數據單元;綠色是I/O,紅色是控制邏輯單元。

與CPU和GPU相比,TPU的控制單元更小,更容易設計,面積只佔了整個沖模的2%,給片上存儲器和運算單元留下了更大的空間。而且,TPU的大小隻有其他晶元的一半。矽片越小,成本越低,良品率也越高。

確定性,是單用途帶來的另一個優勢。CPU和GPU需要考慮各種任務上的性能優化,因此會有越來越複雜的機制,帶來的副作用就是這些處理器的行為非常難以預測。

而用TPU,我們能輕易預測運行一個神經網路、得出預測,需要多長時間,這樣,我們能讓晶元以吞吐量接近峰值的狀態運行,同時嚴格控制延遲。

以上面提到的MLP0為例,在同樣將延遲控制在7毫秒之內的情況下,TPU的吞吐量是CPU和GPU的15到30倍。

各種處理器上每秒可運行的MLP0預測

下面,是TPU、CPU、GPU在六種神經網路上的性能對比。在CNN1上,TPU性能最為驚人,達到了CPU的71倍。

如上文所述,TPU性能強勁的秘訣,是因為它專註於神經網路推斷。這使得量化選擇、CISC指令集、矩陣處理器和最小設計都成為可能。

神經網路正推動計算模式的轉變,Google預計未來幾年中,TPU將成為快速、智能和價格實惠的重要晶元。【完】

招聘

量子位正在招募編輯記者、運營、產品等崗位,工作地點在北京中關村。相關細節,請在公眾號對話界面,回復:「招聘」。

One More Thing…

今天AI界還有哪些事值得關注?在量子位(QbitAI)公眾號對話界面回復「今天」,看我們全網搜羅的AI行業和研究動態。筆芯~

另外,,如果你研究或者從事AI領域,小助手會把你帶入量子位的交流群里。

掃碼強行關注『量子位』



熱門推薦

本文由 yidianzixun 提供 原文連結

寵物協尋 相信 終究能找到回家的路
寫了7763篇文章,獲得2次喜歡
留言回覆
回覆
精彩推薦