深度｜傳說中既能訓練又能推理、擁有超高性能的谷歌第二代TPU到底有多牛？

雲計算 Google 機器學習 GPU 將門創投 2017-05-25

來源：The Next Platform

Google在今年的Google I/O開發者大會上，推出了極具震撼的第二代TPU芯片，風頭大大蓋過了前幾天Nvidia剛剛推出的Volta架構GPU。通過四個TPU芯片的組合板可以實現180Tera次每秒用於深度學習訓練和推理的浮點運算，如果再將64個這樣的TPU系統板通過網絡連接起來，浮點性能可以達到驚人的11.5PFLOPS，遠遠超過世界上前五百超級計算機的算力！

谷歌大腦研究團隊主管Jeff Dean表示：“我們的TPU擁有世界一流的ASIC設計團隊，第二代TPU芯片採用了與第一代完全不同的設計架構,第一代芯片由於只進行推理運算所以體積較小，而在第二代芯片中我們加入了能滿足訓練需求的硬件結構，同時還考慮瞭如何將這些計算單元組合成大型系統的架構方法。”

接下來我們將緊隨Google的腳步為大家解讀這一性能趕超傳統超級計算機的新型TPU。雖然Google沒有提供關於TPU2和母版的詳細資料，但是我們根據收集到的資料進行分析總結，盡力呈現出這一芯片的架構和原理。

乍一看新型的TPU2架構與CrayXT或XC比較類似，我們猜測其通過母版核心實現芯片的內部互聯，而與外部的通信端口則位於板子外部。TPU2板擁有4個TPU2單元，每一個能實現高達45TFLOPs的浮點性能，而通過系統機的板上互聯則可以實現高達180 TFLOPs的浮點性能。

我們可以看到在作用兩邊有八個互聯端口，另外還有兩個額外借口在板子左邊。如果TPU2板像AMD的最新Vega架構一樣直接與flash存儲連接，這樣的架構將會十分有趣。左邊連個額外的接口可以用於直接連接存儲器，或者可以用於上行連接與其他TPU互聯的上級系統。

我們猜測每一個TPU板擁有兩個與外部通信的接口，左側兩個接口用於與本地存儲通信並用於與其他節點進行連接（如果這一架構足夠快足夠靈活的話Google可能利用弱耦合或者緊內存的共享協議來實現這樣的互聯。）

下圖是TPU2的Pod板集群示意圖，表示這一裝置在進行機器學習時擁有11.5Peta次每秒的浮點計算能力。

讓我們來看看這個被稱為TPU Pod的裝置，這是完全開放式的結構。這個二維裝置總共有8行TPU單元，每一行有四個TPU板水平排列。我們還不清楚機架是完整的TPU板或者是半個開發板，只可以看到每一個TPU節點伸出了6個端口，同時其中的兩個端口在上方將他們的機箱互聯起來。

在最上一排TPU機箱上面，這兩個接入機箱的端口並沒有出現在TPU單元中。我們猜測這是鏈接用於高速存儲本地數據的flash機箱。最後我們還看到每一個Pod中至少有32個TPU母版（4*8)，意味著它擁有128個TPU芯片在同時運行，稍微計算一下我們發現整套系統的計算能力大概在11.5千萬億次。

如此強大的計算能力意味著生產力的大幅提升，Google最新的大規模翻譯模型在目前市面上最好的GPU上需要1整天的訓練時間，而利用1/8個TPU pod的算力在一下午就可以完成。但我們要記住一點TPU是專門為TensorFlow量身定做的，而商用GPU則是通用性較強的計算裝置，但Google感興趣的一直是實現高性能的專用集成電路ASIC。

需要補充的是，回想一下第一代的TPU芯片，雖然用於高速的高效的推斷能力，但是需要首先在GPU上訓練模型而後移植到TPU上，而實驗新模型的訓練時間和訓練週期會耗費開發人員大量的時間。這就是為什麼對於深度學習來說最好的硬件就是在一個設備上同時實現訓練和推理。現在我們終於等到了可以有多種選擇的這一天，比如不久後Intel的Knights Mill 和NVIDIA剛發佈的Volta GPU。

NVIDIA的Volta GPU，利用tensor core的機器學習訓練加速處理元素，可以在一個設備上達到120 萬億次的計算能力，比去年出的Pascal GPU在計算力上提高很多。雖然這個很厲害，但還是Google的發佈會更抓人眼球，即便離走入生活還有一定的距離，但是GPU越來越快的計算能力依舊令人印象深刻。

Dean 提到，Volta的結構很有趣，因為NVIDIA意識到了矩陣乘法對於應用加速的重要性。Google的第一代TPU也是用了同樣的想法，提高了矩陣乘法推理的速度。他說，“加快線性代數運算總是很好用的。”

從用戶的角度看，Google將於這些TPU放到Google 雲平臺上，開放給大家使用。Dean說到，他們不想侷限用戶的選擇，在提供TPU的同時，也會繼續提供現有的Volta GPU 和 Skylake Xeon CPU 組合。Google 會提供多個選擇讓開發者建造和運行他們的模型，並且會在雲裡為合格的研究團體提供1000臺TPU，做開源的機器學習工作。

Dean 還說，在Google內部，雖然有些機器學習任務還會在GPU和CPU上完成，但是很多會轉到這個專門為TensorFlow設計，並且兼顧訓練推理的TPU上完成。雖然我們還不清楚新TPU在訓練和推理的耗電參數，我們可以猜測它還是會比Volta更優化。畢竟Volta是一個很大的芯片，並且兼容很多種任務。NVIDIA選用了16位浮點，但是我們認為TPU結構選用了16位或8位。我們希望在以後可以和工程師確認。

TPU 2和第一代不同的是，第一代選用了量化整數運算，而第二代用的是浮點運算。這樣一來你就不用把訓練好的模型轉換使用量化整數做推理了。你直接用訓練時用的浮點做推理就好，這樣實施起來容易了很多。