微软为 ChatGPT 打造专用超算,砸下几亿美元,用了上万张 A100。现在,谷歌首次公布了自家 AI 超算的细节 —— 性能相较上代 v3 提升 10 倍,比 A100 强 1.7 倍。此外,据说能和 H100 对打的芯片已经在研发了。
虽然谷歌早在 2020 年,就在自家的数据中心上部署了当时最强的 AI 芯片 ——TPU v4。
但直到今年的 4 月 4 日,谷歌才首次公布了这台 AI 超算的技术细节。
相比于 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 包含四个 128x128 矩阵乘法单元(MXU),一个具有 128 个通道(每个通道 16 个 ALU),以及 16 MiB 向量存储器(VMEM)的向量处理单元(VPU)。
两个 TC 共享一个 128 MiB 的公共存储器。
值得注意的是,A100 芯片与谷歌第四代 TPU 同时上市,那么其具体性能对比如何?
谷歌分别展示了在 5 个 MLPerf 基准测试中每个 DSA 的最快性能。其中包括 BERT、ResNET、DLRM、RetinaNet、MaskRCNN。
其中,Graphcore IPU 在 BERT 和 ResNET 提交了结果。
如下展示了两个系统在 ResNet 和 BERT 的结果,点之间的虚线是基于芯片数量的插值。
TPU v4 和 A100 的 MLPerf 结果都扩展到比 IPU 更大的系统。
对于相似规模的系统,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 年推出了第一代张量处理单元,其中不仅包含了专门为张量计算优化的定制 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 只会变得更强。
参考资料: