雖然谷歌早在2020年，就在自家的數據中心上部署了當時最強的AI芯片——TPU v4。

但直到今年的4月4日，谷歌才首次公布了這臺AI超算的技術細節。

論文地址：https://arxiv.org/abs/2304.01433

相比于TPU v3，TPU v4的性能要高出2.1倍，而在整合4096個芯片之后，超算的性能更是提升了10倍。

另外，谷歌還聲稱，自家芯片要比英偉達A100更快、更節能。

與A100對打，速度快1.7倍

論文中，谷歌表示，對于規模相當的系統，TPU v4可以提供比英偉達A100強1.7倍的性能，同時在能效上也能提高1.9倍。

另外，谷歌超算速度還要比Graphcore IPU Bow快約4.3倍至4.5倍。

谷歌展示了TPU v4的封裝，以及4個安裝在電路板上的封裝。

與TPU v3一樣，每個TPU v4包含兩個TensorCore（TC）。每個TC包含四個128x128矩陣乘法單元（MXU），一個具有128個通道（每個通道16個ALU），以及16 MiB向量存儲器（VMEM）的向量處理單元（VPU）。

兩個TC共享一個128 MiB的公共存儲器（CMEM）。

值得注意的是，A100芯片與谷歌第四代TPU同時上市，那么其具體性能對比如何？

谷歌分別展示了在5個MLPerf基準測試中每個DSA的最快性能。其中包括BERT、ResNET、DLRM、RetinaNet、MaskRCNN。

其中，Graphcore IPU在BERT和ResNET提交了結果。

如下展示了兩個系統在ResNet和BERT的結果，點之間的虛線是基于芯片數量的插值。

TPU v4和A100的MLPerf結果都擴展到比IPU更大的系統（4096個芯片對比256個芯片）。

對于相似規模的系統，TPU v4在BERT上比A100快1.15倍，比IPU快大約4.3倍。對于ResNet，TPU v4分別快1.67倍和大約4.5倍。

對于在MLPerf基準測試上的功耗使用情況，A100平均上使用了1.3倍至1.9倍的功率。

峰值每秒浮點運算次數是否能預測實際性能？許多機器學習領域的人認為峰值每秒浮點運算次數是一個很好的性能代理指標，但實際上并非如此。

例如，盡管在峰值每秒浮點運算次數上僅具有1.10倍的優勢，TPU v4在兩個MLPerf基準測試上比IPU Bow在相同規模的系統上快4.3倍至4.5倍。

另一個例子是，A100的峰值每秒浮點運算次數是TPU v4的1.13倍，但對于相同數量的芯片，TPU v4卻快1.15倍至1.67倍。

如下如圖使用Roofline模型展示了峰值FLOPS/秒與內存帶寬之間的關系。

那么，問題來了，谷歌為什么不和英偉達最新的H100比較？

谷歌表示，由于H100是在谷歌芯片推出后使用更新技術制造的，所以沒有將其第四代產品與英偉達當前的旗艦H100芯片進行比較。

不過谷歌暗示，它正在研發一款與Nvidia H100競爭的新TPU，但沒有提供詳細信息。谷歌研究員Jouppi在接受路透社采訪時表示，谷歌擁有「未來芯片的生產線」。

TPU vs GPU

在ChatGPT和Bard「決一死戰」的同時，兩個龐然大物也在幕后努力運行，以保持它們的運行——英偉達CUDA支持的GPU（圖形處理單元）和谷歌定制的TPU（張量處理單元）。

換句話說，這已經不再是關于ChatGPT與Bard的對抗，而是TPU與GPU之間的對決，以及它們如何有效地進行矩陣乘法。

由于在硬件架構方面的出色設計，英偉達的GPU非常適合矩陣乘法任務——能有效地在多個CUDA核心之間實現并行處理。

因此從2012年開始，在GPU上訓練模型便成為了深度學習領域的共識，至今都未曾改變。

而隨著NVIDIA DGX的推出，英偉達能夠為幾乎所有的AI任務提供一站式硬件和軟件解決方案，這是競爭對手由于缺乏知識產權而無法提供的。

相比之下，谷歌則在2016年推出了第一代張量處理單元（TPU），其中不僅包含了專門為張量計算優化的定制ASIC（專用集成電路），并且還針對自家的TensorFlow框架進行了優化。而這也讓TPU在矩陣乘法之外的其他AI計算任務中具有優勢，甚至還可以加速微調和推理任務。

此外，谷歌DeepMind的研究人員還找到了一種能夠創造出更好矩陣乘法算法的方法——AlphaTensor。

然而，即便谷歌通過自研的技術和新興的AI計算優化方法取得了良好的成果，但微軟與英偉達長久以來的深度合作，則通過利用各自在行業上的積累，同時擴大了雙方的競爭優勢。

第四代TPU

時間回到21年的谷歌I/O大會上，劈柴首次公布了谷歌最新一代AI芯片TPU v4。

「這是我們在谷歌上部署的最快的系統，對我們來說是一個具有歷史意義的里程碑。」

這次的改進已經成為構建AI超算的公司之間競爭的關鍵點，因為像谷歌的Bard、或OpenAI的ChatGPT類似的大型語言模型已經在參數規模上實現爆炸式增長。

這意味著它們遠遠大于單個芯片所能存儲的容量，對算力需求是一個巨大的「黑洞」。

因此這些大模型必須分布在數千個芯片上，然后這些芯片必須協同工作數周，甚至更長時間來訓練模型。

目前，谷歌迄今為止公開披露的最大的語言模型PaLM，有5400億參數，便是在50天內將其分割到兩臺4000芯片的超級計算機上進行訓練的。

谷歌表示，自家的超級計算機能夠輕松地重新配置芯片之間的連接，能夠避免問題，并進行性能調優。

谷歌研究員Norm Jouppi和谷歌杰出工程師David Patterson在關于該系統的博客文章中寫道，

「電路交換使得繞過失效組件變得容易。這種靈活性甚至允許我們改變超算互連的拓撲結構，以加速機器學習模型的性能。」

盡管谷歌現在才發布有關其超級計算機的詳細信息，但自2020年以來，該超級計算機已在位于俄克拉荷馬州梅斯縣的數據中心內上線。

谷歌表示，Midjourney使用該系統訓練了其模型，最新版的V5讓所有人見識到圖像生成的驚艷。

最近，劈柴在接受紐約時報采訪稱，Bard將從LaMDA轉到PaLM上。

現在有了TPU v4超算的加持，Bard只會變得更強。