有圖有真相：深度拆解谷歌TPU3.0，新一代AI協(xié)同處理器

　　通過(guò)將服務(wù)器主板集成到云 TPU 機架中，谷歌可以用相同的機架配置使機架數量增加一倍。在一個(gè)機架上標準化配置必然有助于降低硬件部署的成本和復雜性。

　　電腦架：TPUv2(左)和 TPUv3(右)

　　但是，為了實(shí)現更高的密度，谷歌必須從 4U 云 TPU 外形規格轉變?yōu)?2U 高密度外形規格。其數據中心溫度很高(公布的數據在 80°F 到 95°F 之間)，因此 TPUv2 風(fēng)冷散熱器必須很大。谷歌使用開(kāi)放式機架，所以利用風(fēng)來(lái)冷卻密集外形規格的熱插槽變得非常昂貴，使得水冷成為可行的替代方案。特別是對于像深度學(xué)習這樣的高價(jià)值服務(wù)。

　　將服務(wù)器主板轉移到 TPUv3 機架中還會(huì )縮短連接電纜，因此我們一般認為谷歌節省了大量電纜成本，并除去了 TPUv2 Pod 服務(wù)器機架中的閑置空間。

　　谷歌沒(méi)有展示主板與機架水互連的照片。

　　云 TPU

　　但是，它確實(shí)顯示了 TPUv3 云 TPU 的兩張照片。TPUv3 云 TPU 具有與 TPUv2 云 TPU 相似的布局。明顯的變化是水冷卻的增加。主板電源接頭的背面看起來(lái)相同。但是，主板前面還有四個(gè)附加連接器。照片正面(左)的兩個(gè)銀色大正方形是由四個(gè)連接器組成的集群。

　　TPUv3 主板(左上)，TPUv2 主板(左下)和 TPUv3 主板特寫(xiě)(右)

　　谷歌沒(méi)有提及其他連接器。我們認為最有可能的解釋是 Google 為環(huán)形超網(wǎng)格(hyper-mesh)添加了一個(gè)維度，也就是從 2D 環(huán)形超網(wǎng)格到 3D 環(huán)形超網(wǎng)格。

　　環(huán)形超網(wǎng)格互連圖：2D(左)和 3D(右)

　　去年，我們推測了幾種類(lèi)型的互連，并將其稱(chēng)為錯誤的互連——谷歌使用 32 條有線(xiàn) PCI-Express 3.0(每條鏈路 28GB / s)將服務(wù)器連接到云 TPU 上。我們認為，谷歌不太可能增加服務(wù)器主板和云 TPU 之間的帶寬，因為 PCI-Express 帶寬和延遲可能不是什么重要的性能限制因素。

　　雖然互連拓撲將有助于深度學(xué)習任務(wù)在 pod 中更好地擴展，但它不會(huì )對原始的理論性能帶來(lái)貢獻。

　　TPU 芯片

　　現在，我們要深入到芯片層面來(lái)回答以下問(wèn)題：「剩下的 2 倍性能改進(jìn)來(lái)自哪里?」谷歌概括其 TPUv2 核心為：

　　有兩個(gè)矩陣單元(MXU)每個(gè) MXU 都有 8GB 的專(zhuān)用高帶寬內存(HBM)每個(gè) MXU 的原始峰值吞吐量為 22.5 萬(wàn)億次但是 MXU 不使用標準浮點(diǎn)格式來(lái)實(shí)現其浮點(diǎn)吞吐量谷歌創(chuàng )造了自己的內部浮點(diǎn)格式，稱(chēng)為「bfloat」，意為「大腦浮點(diǎn)(brain floating point)」(在谷歌大腦之后)。Bfloat 格式使用 8 位指數和 7 位尾數，而不是 IEEE 標準 FP16 的 5 位指數和 10 位尾數。Bfloat 可以表示從~1e-38 到~3e38 的值，其動(dòng)態(tài)范圍比 IEEE 的 FP16 寬幾個(gè)數量級。谷歌之所以創(chuàng )造 bfloat 格式，是因為它發(fā)現在 IEEE 標準 FP16 的訓練中需要數據科學(xué)專(zhuān)家，以確保數據保持在 FP16 較為有限的范圍內。

　　我們相信谷歌已經(jīng)在 MXU 內部實(shí)現了硬件格式轉換，真正消除了轉換延遲和軟件開(kāi)發(fā)難題。從 FP16 到 bfloat 的格式轉換看起來(lái)像是直接把精度截斷到較小的尾數。將 FP16 轉換為 FP32，然后再將 FP32 轉換為 FP16 是已知的實(shí)踐;可以使用相同的技術(shù)把格式從 FP32 轉換成 bfloat，然后再從 bfloat 轉換成 FP16 或 FP32。

　　谷歌聲稱(chēng)，隨著(zhù)數據流通過(guò) MXU 的收縮陣列，中間結果得到了「極大」的重復使用。

　　考慮到 MXU 的表現，我們相信谷歌不太可能在從 TPUv2 到 TPUv3 的轉變中將 MXU 做出改變。更有可能的結果是，谷歌將把 TPUv3 的 MXU 數量增加一倍。

　　框圖：TPUv2(左)和 TPUv3(右)

　　距離上次芯片發(fā)布只有一年，因此芯片設計節奏非常短，沒(méi)有時(shí)間進(jìn)行重要的架構開(kāi)發(fā)。但是，足夠將現有的 MXU 核心壓縮為新的制造工藝、調整功耗和速度路徑，然后做一點(diǎn)額外的平面規劃工作，以在模具上沖壓更多的 MXU 核心。下表包含了我們所掌握的少量硬信息，以及我們對谷歌 TPUv3 芯片發(fā)展方向的最佳估計。

　　去年，我們估計 TPUv2 每個(gè)芯片需要消耗 200 瓦至 250 瓦?，F在我們知道，每個(gè)包中還包含 16GB 的 HBM，其 MXU 和 HBM 之間的帶寬為 2.4 TB /秒。

　　我們將堅持去年估計的 36 千瓦機架電源(一個(gè) TPUv3pod 總共需要 288 千瓦)。如果假設每臺雙插槽服務(wù)器的功率為 400 瓦，我們會(huì )把每臺 TPUv3 芯片的功率將向后降至 200 瓦左右，其中包括 32GB 的 HBM。如果這些芯片沒(méi)有密集地封裝在主板和機架上，或者被部署在較冷的數據中心，那它們可能不需要水冷。另一種選擇可能是谷歌正在他們的新 TPUv3 集群中部署單插槽服務(wù)器。將服務(wù)器功率降至 250 瓦以下可能會(huì )為 TPUv3 提供足夠的擴展空間，以達到 225 瓦。

　　假定最初 TPUv2 MXU 設計保守，隨后 TPUv3 工藝收縮、HBM 變寬且更快，并且路徑調整速度加快，則可以合理地期望每個(gè)核的性能在兩代之間保持不變，而無(wú)需徹底地對 MXU 進(jìn)行重新設計。

　　市場(chǎng)回顧

　　谷歌仍在部署 TPUv1 外接程序卡，用于推理任務(wù)，四個(gè) TPUv1 用于一臺服務(wù)器上。谷歌部署了 TPUv1 來(lái)加速網(wǎng)頁(yè)搜索和其他大規模推理任務(wù)——如果你最近使用了谷歌搜索引擎，你可能已經(jīng)使用了 TPUv1。

　　谷歌僅通過(guò)其測試云 TPU 實(shí)例提供 TPUv2 訪(fǎng)問(wèn)，并未預測何時(shí)通過(guò)服務(wù)級別協(xié)議供應可用產(chǎn)品。谷歌本周確實(shí)表示，它將在「今年年底」向客戶(hù)提供 TPUv2 pod 服務(wù)，但尚不清楚這是否為一項產(chǎn)品服務(wù)。我們最大的猜測是，谷歌將繼續等待，直到驗證和調試完 TPUv3 pod，以便在全球范圍內部署 TPU pod。谷歌內部正在使用 TPUv2 pod 進(jìn)行一些訓練任務(wù)。本周，谷歌沒(méi)有就何時(shí)部署基于 TPUv3 芯片的任何功能或服務(wù)發(fā)表任何聲明。我們認為，TPUv3 的發(fā)布旨在強調谷歌長(cháng)期致力于控制自己內部生態(tài)的承諾，以加速其 TensorFlow 深度學(xué)習框架。

　　然而，我們認為 TPUv3 更多的應該算 TPUv2.5，而不是新一代芯片。大多數新硬件開(kāi)發(fā)似乎都是圍繞 TPUv3 芯片級別的系統展開(kāi)的。