谷歌長文總結四代TPU打造經驗：里程碑式的TPUv4是怎樣煉成的？

机器之心發表於2022-02-13
谷歌
作為圖靈獎得主、計算機架構巨擘，David Patterson 在 2016 年從伯克利退休後，以傑出工程師的身份加入了谷歌大腦團隊，為幾代 TPU 的研發做出了卓越貢獻。

幾個月前谷歌強勢推出了 TPUv4，並撰寫論文講述了研發團隊的設計思路和從前幾代 TPU 中吸取的經驗。

圖片


論文連結：https://www.gwern.net/docs/ai/2021-jouppi.pdf

當前，深度神經網路在商用領域的特定架構（DSA）已初步建立。該論文回顧了四代 DSA 的演變歷程。生產經驗推動了新設計的誕生，不再受限於 CPU 的緩慢發展和摩爾定律的衰退。

圖片

圖 1：從 2015 年開始部署在谷歌資料中心的 3 個 TPU 的框圖。

下表展示了幾代 TPU 的關鍵特性，其中 TPUv4i 中 i 代表的是「推理」。

圖片


TPUv1 是谷歌第一代 DNN DSA（圖 1 左），能夠處理推理任務。近年來，模型訓練的規模逐漸增大，所以一種新的改進是新增一個片到片的定製互連結構 ICI，使得搭載 TPUv2 的超級計算機晶片數量可多達 256 個。

與 TPUv1 不同，TPUv2 每個晶片有兩個 TensorCore。晶片上的全域性線不會隨著特徵尺寸的縮小而縮小 ，因此相對延遲會增加。每個晶片有兩個較小的核，以避免單個大型全晶片核的過度延遲。谷歌並沒有繼續增加核的數量，因為他們相信兩個強大的核比許多不夠強的核更易於高效地編譯程式。

TPUv3 則在 TPUv2 的基礎上微調了設計，採用相同的技術，擁有 2 倍的 MXU 和 HBM 容量，並將時脈頻率、記憶體頻寬和 ICI 頻寬提高至 1.3 倍。一臺 TPUv3 超級計算機可以擴充套件到 1024 個晶片。TPUv3 在使用 16 位浮點（bfloat16 vs IEEE fp16 時與 Volta GPU 相當。然而，Volta 需要使用 IEEE fp32 來訓練谷歌的生產工作負載，因此 TPUv3 快了 5 倍以上。一些擴充套件到 1024 個晶片的應用程式可以獲得 97%-99% 的完美線性加速。

2015 年以來的十大經驗

1. 邏輯、線路、SRAM 和 DRAM 的發展不平衡

Horowitz 對操作功耗的見解啟發了許多 DSA 的設計。表 2 顯示 7 nm 相對於 45 nm 的平均增益是 2.6 倍，但並不均勻：

圖片

  • SRAM 訪問只提高了 1.3 - 2.4 倍，部分原因是 SRAM 的密度擴充套件速度比之前慢。65 nm 與 7 nm 相比， SRAM 每平方毫米的容量密度僅約為理想擴充套件的 1/5；

  • 得益於封裝的革新，DRAM 訪問量提高了 6.3 倍。HBM 將短 DRAM dies 棧放置在寬匯流排上靠近 DSA 的位置；

  • 但每單位長度線路的功耗提高不到 2 倍（表 2 未顯示這一點）。較差的線路延遲縮放導致 TPUv2/v3 相比於 TPUv1，從使用 1 個較大核變為使用 2 個較小核。


邏輯方面的改進速度比線路和 SRAM 快得多。HBM 比 GDDR6 和 DDR DRAM 更節能，且 HBM 的每 GB/s 的頻寬成本最低。

2. 利用預先的編譯器最佳化 

自 20 世紀 80 年代以來，新架構的命運與編譯器的質量密切相關。事實上，編譯器問題可能會破壞安騰（Itanium）的 VLIW 架構，但許多 DSA 依賴於 VLIW，包括 TPU。

架構師希望在模擬器上開發出出色的編譯器，然而大部分進展發生在硬體可用之後，因為這樣編譯器編寫者能夠測量程式碼所花費的實際時間。因此，如果能夠利用預先的編譯器最佳化而不是從頭開始重新編寫，那麼快速實現架構的全部潛力會容易得多。

3. 效能 / TCO 和效能 / CapEx 的設計

CapEx 指的是一個專案的價格。OpEx 指的是運營成本，包括電力消耗和電力供應。標準的會計處理方法是在 3-5 年內攤銷計算機 CapEx，因此 3 年的 TCO = CapEx + 3 ✕ OpEx。

谷歌等大多數公司更關心生產應用程式的效能 / TCO (perf/TCO)，而不是基準的原始效能或效能 / CapEx (perf/CapEx) 。雖然 TCO 是谷歌在產品設計過程中最佳化的主要變數，但資本支出仍然會影響一些業務決策。

4. 支援向後 ML 相容

一些 DNN 具有上市時間限制，因為及時性可能會帶來經濟價值。在這種情況下，向後 ML 相容性成為一種原則。新 CPU 的目標是相同的，即應該獲得完全相同的結果，包括相同的異常行為和相關效能，從 TPUv2 開始，訓練和圍繞TPU生成的服務都是按照這種情況設計的。

5. 推理 DSA 需要全域性範圍的風冷

75W TPUv1 和 280W TPUv2 採用風冷，但 450W 功率的 TPUv3 使用了液冷。液冷需要將 TPU 放置在幾個相鄰的機架中，以分攤冷卻基礎設施。這種放置限制對於訓練已經由幾個相鄰機架組成的超級計算機來說不是問題。此外，將訓練限制在擁有更多空間的幾個資料中心的缺點是很小的，因為沒有必要進行廣泛部署。

6. 一些推理應用程式需要浮點運算

DNN 量化旨在維持使用整數下的實時推理模型質量，即使所有訓練都是在浮點下完成的。量化計算可以節省面積和功耗，但這需要和質量降低、延遲部署進行權衡，並且一些應用程式在量化時不能很好地工作。

7. 生產推理通常需要多租戶模式

與 CPU 一樣，DSA 應該支援多租戶（multi-tenancy）模式。如果應用程式使用許多模型，共享則可以降低成本並減少延遲。例如，翻譯 DNN 需要許多語言對，而語音 DNN 必須處理多種方言。多租戶模式還支援多種批大小，以權衡吞吐量和延遲。另一方面，採用良好的軟體工程實踐方式也是很重要的，例如面向一小部分客戶嘗試新功能，或緩慢部署新版本以減少出現問題的機會。表 3 顯示了超過 80% 的生產推理工作負載需要多租戶。

圖片


8. DNN 的記憶體和計算需求每年約增長 1.5 倍 

與基準測試不同，程式設計師不斷改進生產應用程式，這通常會增加記憶體大小和計算需求。表 4 跟蹤了仍然在 TPUv1/v2/v3 上執行的四個原始生產推理應用程式的記憶體大小和計算量的年均增長。生產 DNN 每年約增長 1.5 倍，增長速度與摩爾定律一樣快，就像早期的 PC 軟體一樣。這個比率表明架構師應該提供空間，以便 DSA 可以在其整個生命週期內保持有用。

圖片


9. DNN 工作負載隨著 DNN 的突破而發展 

表 3 中的 DNN 約佔谷歌 2020 年推理工作負載的 100%。2016 年的 MLP 和 CNN 仍然很受歡迎，儘管一些應用程式從 MLP 切換到 BERT DNN (28%)。BERT 出現於 2018 年，但它已經佔了 28% 的工作負載。為了提高質量，transformer編碼器加上 LSTM 解碼器 (RNN0) 和 Wave RNN (RNN1) 取代了 LSTM (29%)。這個經驗說明了可程式設計性和靈活性對於推理DSA的重要性。

10. 推理 SLO 的限制是 P99 延遲，而不是批大小

近來 DSA 的論文根據批大小重新定義了延遲限制，通常設定為 1。表 5 給出了生產應用程式的 P99 時間 SLO 和 MLPerf 推理基準測試，並顯示了近來符合 SLO 的 TPU 的最大批。顯然，資料中心應用程式限制了延遲，而不是批大小。未來的 DSA 應該利用更大的批大小。

圖片


這十條經驗如何塑造了現在的 TPUv4i？
鑑於利用預先的編譯器最佳化和向後 ML 相容性的重要性，再加上重用早期硬體設計的好處，TPUv4i 沿襲了 TPUv3 的一些設計。該論文集中討論了 TPUv3 和下一代的不同之處，谷歌重新考慮了構建單個晶片的策略，使晶片既能最佳化訓練，又能用於推理。

編譯相容，而不是二進位制相容。鑑於 TPUv2 和 TPUv3 共享 322 位 VLIW 指令包長度，傳統架構思想是 TPUv4i 和 TPUv4 嘗試保持向後的二進位制相容性，但谷歌選擇了編譯相容。

在片上SRAM儲存中增加了普通記憶體CMEM。DSA 在編譯之後首先考慮的是記憶體系統。限制記憶體開銷可以提高 perf/CapEx，但是可能會影響 perf/TCO。除了推理之外，多租戶、快速增長的 DNN 和 HBM 的高能效，都促使 TPUv4i 繼續使用類似於 TPUv3 的 HBM。


四維張量 DMA。記憶體系統架構對於任何 DNN 加速器都是至關重要的，因此應被設計為滿足常用情況工作負載的最大效能，同時對於其他模型足夠靈活。TPUv4i 包含張量 DMA 引擎，分佈在整個晶片的非核心部分，以減輕互連延遲和線路擴充套件的影響。

自定義片上互連(OCI)。快速增長和不斷髮展的 DNN 工作負載促使 TPU 的非核心具備更大的靈活性。過去 TPU 設計的每個元件都是點對點連線的(如圖 1 所示)，隨著記憶體頻寬的增加和元件數量的增加，點對點方法變得過於昂貴，需要大量路由資源和 die area。它還需要預先選擇支援哪種通訊模式，這限制了未來軟體使用晶片的方式。

因此 TPUv4i 增加了一個共享的片上互連 OCI，它連線了 die 上所有的部件，並且可以基於當前的部件來擴充套件它的拓撲。OCI 對於增加 CMEM 來說特別重要，在 HBM、CMEM 和 VEM 之間進行分配和資料傳輸的方式會持續進化。

運算改進。另一個重要變化是運算單元。需要量化的風險和 ML 向後相容性的重要性意味著，儘管要進行推理，仍然要保留 TPUv3 中的 bf16 和 fp32。谷歌還希望對 TPUv1 的應用程式進行量化，以便輕鬆地將其移植到 TPUv4i，TPUv4i 也支援 int8。

時鐘速率和 TDP。用於推理的風冷和減少 TCO 導致時脈頻率變為 1.05 GHz，晶片 TDP 為 175W，再次接近 TPUv1(75W)而不是 TPUv3(450W)。

ICI 擴充套件。為了給未來的 DNN 增長提供空間，TPUv4i 有 2 個 ICI 鏈路，這樣每塊板上的 4 個晶片可以透過模型分割槽快速訪問附近的晶片記憶體(TPUv3 使用 4 個 ICI 鏈路)。隨著軟體的成熟和 DNN 的增長，應用程式可能會用到這些擴充套件。

工作負載分析特徵。基於 TPUv1 的經驗，TPUv4i 包含了大量的追蹤和效能計數器硬體特徵，特別是在非核心部分。它們被軟體堆疊用於測量和分析使用者工作負載中的系統級瓶頸，並指導持續的編譯級和應用級最佳化，如下圖2所示。這些特徵增加了設計時間和麵積，但這種嘗試是值得的。隨著 DNN 工作負載的增加和發展，這些特徵可以顯著提升系統級效能，並提高產品在整個生命週期中的生產力(如表 4 所示)。

圖片


TPUv4i 效能分析

下圖比較了TPUv1、TPUv3和TPUv4i在生產推理應用上相對於TPUv2的效能和perf/TDP。TPUv3 和 TPUv4i 都比 TPUv2 快 1.9 倍以上，TPUv1 的速度大概是 TPUv2 的 70% 左右。更大、更熱的 TPUv2/v3 晶片有兩個核心，而較小的 TPUv4i 晶片有一個核心，這使 TPUv4i 晶片在 perf/TCO 和部署方面更勝一籌。

圖片

下圖使用 MLPerf 推理 0.5-0.7 版本基準測試比較了 TPUv3 和 TPUv4i 相比於 NVIDIA T4 的效能和 perf/(system)TDP。

圖片

更深度的效能表現：CMEM

下圖展示了 TPUv4i 增加 CMEM 給 MLPerf 推理 0.5-0.7 基準測試帶來的收益。由於延遲和延遲受限服務效能之間的關係在某些效用點上是非線性的，因此離線的平均收益只有 1.1 倍，而服務的平均收益是 1.3 倍。

圖片

為了理解 CMEM 對生產應用的影響，下圖比較了啟用 CMEM 的雙倍 HBM 頻寬與啟用 CMEM 的標準 HBM 頻寬的效能。透過編譯標誌來禁用 CMEM，同時在 TPUv4 執行時禁用一個核心，使得一個核心的 HBM 頻寬增加了一倍。

圖片

下圖是藉助 Roofline 模型解釋內在原理。它將應用程式劃分為計算受限或記憶體受限，使用操作強度來確定應用程式相對於其 Roofline 的位置。

圖片

最後，下圖展示了小於 128 MB 的 CMEM 對應用程式和 MLPerf 服務基準測試的影響。

圖片

更多研究細節，參見論文原文。


相關文章