Google深度揭秘TPU:一文看懂内部原理,以及为何碾压GPU
搜索、街景、照片、翻譯,這些Google提供的服務,都使用了Google的TPU(張量處理器)來加速背后的神經網絡計算。
△?在PCB板上的Google首款TPU和部署了TPU的數據中心
去年Google推出TPU并在近期對這一芯片的性能和架構進行了詳細的研究。簡單的結論是:TPU與同期的CPU和GPU相比,可以提供15-30倍的性能提升,以及30-80倍的效率(性能/瓦特)提升。
這意味著,Google的服務既可以大規模運行于最先進的神經網絡,而且可以把成本控制在可接受的程度上。以下的內容,將深入解讀Google TPU中的技術,并且討論如何實現更為出色的性能。
通往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倍以上的性能提升。
RISC,CISC和TPU指令集
可編程性是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)運算。
TPU的核心:脈動陣列
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、GPU和TPU的每周期算術運算量進行一下對比:
| 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將成為快速、智能和價格實惠的重要芯片。【完】
原文發布于Google云
作者:
谷歌云 Kaz Sato, Staff Developer Advocate
谷歌大腦軟件工程師Cliff Young, Software Engineer
谷歌大腦杰出工程師David Patterson
總結
以上是生活随笔為你收集整理的Google深度揭秘TPU:一文看懂内部原理,以及为何碾压GPU的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 通过PageSpeed优化网站性能
- 下一篇: TensorFlow前向传播