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

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

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

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

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

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

　　云 TPU

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

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

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

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

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

　　雖然互連拓?fù)鋵⒂兄谏疃葘W(xué)習(xí)任務(wù)在 pod 中更好地?cái)U(kuò)展，但它不會(huì)對(duì)原始的理論性能帶來貢獻(xiàn)。

　　TPU 芯片

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

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

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

　　谷歌聲稱，隨著數(shù)據(jù)流通過 MXU 的收縮陣列，中間結(jié)果得到了「極大」的重復(fù)使用。

　　考慮到 MXU 的表現(xiàn)，我們相信谷歌不太可能在從 TPUv2 到 TPUv3 的轉(zhuǎn)變中將 MXU 做出改變。更有可能的結(jié)果是，谷歌將把 TPUv3 的 MXU 數(shù)量增加一倍。

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

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

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

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

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

　　市場(chǎng)回顧

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

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

　　然而，我們認(rèn)為 TPUv3 更多的應(yīng)該算 TPUv2.5，而不是新一代芯片。大多數(shù)新硬件開發(fā)似乎都是圍繞 TPUv3 芯片級(jí)別的系統(tǒng)展開的。