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

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

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

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

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

與A100對打，速度快1.7倍

論文中，谷歌表示，對于規(guī)模相當?shù)南到y(tǒng)，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提交了結(jié)果。

如下展示了兩個系統(tǒng)在ResNet和BERT的結(jié)果，點之間的虛線是基于芯片數(shù)量的插值。

TPU v4和A100的MLPerf結(jié)果都擴展到比IPU更大的系統(tǒng)（4096個芯片對比256個芯片）。

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

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

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

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

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

如下如圖使用Roofline模型展示了峰值FLOPS/秒與內(nèi)存帶寬之間的關(guān)系。

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

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

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

TPU vs GPU

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

換句話說，這已經(jīng)不再是關(guān)于ChatGPT與Bard的對抗，而是TPU與GPU之間的對決，以及它們?nèi)绾斡行У剡M行矩陣乘法。

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

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

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

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

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

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

第四代TPU

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

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

這次的改進已經(jīng)成為構(gòu)建AI超算的公司之間競爭的關(guān)鍵點，因為像谷歌的Bard、或OpenAI的ChatGPT類似的大型語言模型已經(jīng)在參數(shù)規(guī)模上實現(xiàn)爆炸式增長。

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

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

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

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

谷歌研究員Norm Jouppi和谷歌杰出工程師David Patterson在關(guān)于該系統(tǒng)的博客文章中寫道，

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

盡管谷歌現(xiàn)在才發(fā)布有關(guān)其超級計算機的詳細信息，但自2020年以來，該超級計算機已在位于俄克拉荷馬州梅斯縣的數(shù)據(jù)中心內(nèi)上線。

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

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

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