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

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

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

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

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

與A100對打，速度快1.7倍

論文中，谷歌表示，對于規(guī)模相當(dāng)?shù)南到y(tǒng)，TPU v4可以提供比英偉達(dá)A100強(qiáng)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基準(zhǔn)測試中每個DSA的最快性能。其中包括BERT、ResNET、DLRM、RetinaNet、MaskRCNN。

其中，Graphcore IPU在BERT和ResNET提交了結(jié)果。

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

TPU v4和A100的MLPerf結(jié)果都擴(kuò)展到比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基準(zhǔn)測試上的功耗使用情況，A100平均上使用了1.3倍至1.9倍的功率。

峰值每秒浮點(diǎn)運(yùn)算次數(shù)是否能預(yù)測實(shí)際性能？許多機(jī)器學(xué)習(xí)領(lǐng)域的人認(rèn)為峰值每秒浮點(diǎn)運(yùn)算次數(shù)是一個很好的性能代理指標(biāo)，但實(shí)際上并非如此。

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

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

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

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

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

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

TPU vs GPU

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

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

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

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

而隨著NVIDIA DGX的推出，英偉達(dá)能夠?yàn)閹缀跛械腁I任務(wù)提供一站式硬件和軟件解決方案，這是競爭對手由于缺乏知識產(chǎn)權(quán)而無法提供的。

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

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

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

第四代TPU

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

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

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

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

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

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

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

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

「電路交換使得繞過失效組件變得容易。這種靈活性甚至允許我們改變超算互連的拓?fù)浣Y(jié)構(gòu)，以加速機(jī)器學(xué)習(xí)模型的性能?！?/span>

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

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

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

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