李林 舒石 編譯整理
量子位 報道 | 公眾號 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領域,小助手會把你帶入量子位的交流群里。
△ 掃碼強行關注『量子位』