原文来源:硅星人
图片来源:由无界 AI生成
谷歌的大语言模型Gemini 1.0,一经问世就技惊四座。首先在性能方面,无论是从对文本、图像、音频的理解,还是对57种领域文本以及数学问题的推理,几乎都超过了自然语言界的霸主GPT-4。而且在谷歌Gemini的官网上还有这一段话,内容是“Gemini是第一个在MMLU(海量多任务语言理解)上超过人类专家的模型”。
为什么谷歌Gemini这么厉害呢?其实最核心的原因来自于训练谷歌Gemini所使用的两种服务器端推理芯片TPU v4和TPU v5e。虽然从结果来看,TPU跟GPU都是用来训练和推理,但是两个人在“解题思路”上,相距甚远。
谷歌在TPU v4上采用了一种名为光学回路开关(optical circuit switch,以下简称OCS)。顾名思义,就是服务器组用光来传递信息。上方是注入模块,使用850纳米激光二极管,发射850纳米的显示光。中间是用来拆分或者合并光的二向色分光器;下方是用来反射的2D MEMS阵列;左右两端是用来投射O波段信号光的2D透镜阵列和光纤准直器阵列。
简易版OCS原理
传统服务器组用的是铜线导电,通过规定时间内铜线上电子传输的个数来传递信息。但是任何介质,哪怕是不计成本的采用金或者银,依然有电阻的存在,必然会减慢电子在这之间传递的速度。没有任何的光电转换环节,全都是由光来运送信息,没有中间商赚差价。
我们不妨对比一下市面上最常见,也是Open AI目前的解决方案,英伟达H100所使用的Infiniband技术。这个技术使用交换结构拓扑。所有传输都在通道适配器处开始或结束。每个处理器都包含一个主机通道适配器(HCA ),每个外设都有一个目标通道适配器(TCA)。
说得直白一点,H100就和过去的驿站差不多。驿站会代替权力中心接收消息,这样管理者就可以直接去驿站来查看自己未来的任务。但是TPU v4则是直接烧狼烟,不仅快,而且跳过中间所有的消息转达设施,一步到位。根据谷歌的说法,OCS系统和底层光学组件成本价格不到TPU v4系统成本的5%,其功耗不到整个系统的5%,又便宜又好用。
说到这,不得不提一个谷歌创造的概念,叫做MFU。全称为模型FLOPs利用率(Model FLOPs Utilization),这是一个不同于其他衡量计算利用率的方法,不包含任何反向传递期间的作业,这意味着MFU衡量的效率直接转化为端到端的训练速度。那也就是说,MFU最大上限一定是100%,同时数字越高代表训练速度越快。
谷歌想通过MFU这个概念,来评估TPU v4的负载能力以及运行效率,以此得出到底可以给谷歌Gemini上多少强度。对一系列纯解码器变压器语言模型(下图以GPT为例)配置进行了深入的基准测试,这些配置的参数从数十亿到数万亿不等。测试规定,模型大小(参数单位十亿)是芯片数的二分之一。
TPU v4在不同场景下的MFU
在谷歌的基准测试活动中,TPU v4实现了44-56%的MFU。根据下图的对比可以明显看到,尤其是多个TPU v4串联的情况。而这也正是OCS的神奇之处,同时谷歌也希望借助OCS向世人们传达一个道理:“要相信光。”
不同芯片的MFU对比
既然TPU v4都这么厉害了,那TPU v5e作为迭代产品,效果只会更好。事实上谷歌之所以将这个模型命名为Gemini,也是因为这两个TPU。Gemini的本意是双子座,代表的是希腊神话中的天神宙斯与斯巴达王后勒达所生的双胞胎卡斯托耳和波鲁克斯。TPU v4和TPU v5e某种意义上也是双胞胎,因此得名Gemini。当然,另外有一种可能是开发者特别喜欢《圣斗士星矢》,里面黄金十二宫篇大BOSS就是双子座的撒加,绝招是银河星爆。不过我觉得后者的可能性低一点。
还有一点,服务器最大的成本来自运营和维护。但是在这种规模下,使用传统的定期对持久集群存储进行权重检查点操作的方法是不可能保持高吞吐量的,毕竟要周期地关闭重启某个机组。所以对于Gemini,谷歌使用模型状态的冗余内存副本,并且在任何计划外的硬件故障时,可以直接从完整的模型副本中快速恢复。与PaLM和PaLM-2比,尽管使用了大量的培训资源,但这大大加快了恢复时间。结果,最大规模的培训工作的总体产出从85%增加到97%。
谷歌Gemini背后的服务器组也是第一款支持嵌入的硬件的超级计算机。嵌入是一种比较贴近谷歌业务的算法。嵌入本身是一种相对低维的空间,模型可以将高维向量转换为这种低维空间。借助嵌入,可以更轻松地对表示字词的稀疏向量等大型数据输入进行机器学习。理想情况下,嵌入会将语义上相似的输入置于嵌入空间中彼此靠近的位置,以捕获输入的一些语义。
嵌入是广告、搜索排名、YouTube和Google Play中使用的深度学习推荐模型 (DLRM) 的关键组成部分。每个TPU v4都包含第三代稀疏核心数据流处理器,可将依赖嵌入的模型加速5至7倍,但仅使用5%的芯片面积和功耗。
嵌入处理需要大量的端到端通信,因为嵌入分布在在模型上协同工作的TPU芯片周围。这种模式强调共享内存互连的带宽。这就是TPU v4使用3D环面互连的原因(相对于TPU v2和v3使用 2D 环面)。TPU v4的3D环面提供了更高的二等分带宽,即跨越互连中间从一半芯片到另一半的带宽,以帮助支持更多数量的芯片和更好地展现稀疏核心性能。
TPU v4
毫无疑问,谷歌是一家软件公司,谷歌Gemini是一款软件产品,但是谷歌胜就胜在硬件上。谷歌强调Gemini是一个功能强大的多模态模型(Multimodal Model)。多模态指的是结合了多种感官输入形式来做出更明智的决策,相当于是模型可以用复杂、多样和非结构化的数据来求解。
一个短视频(360p到1080p)的数据量大约几十mb到几百mb,一条语音的数据量约为几百kb,一行文字的数据量约为几b。对于传统的大语言模型,如果只处理文字信息,那么对服务器的载荷要求不会很高,毕竟数据量小。可是一旦将视频、图片等一并处理,那么服务器的载荷就会成指数倍增长。事实上不是其他模型不能弄多模态,是其他服务器扛不住这么大数据量,硬件拖了软件的后腿。谷歌之所以敢弄,更多的原因是TPU v4和TPU v5e在大规模串联的场景下,高载荷高MFU无疑是谷歌的绝招。
然而谷歌并不能高兴太早,首先英伟达的张量计算GPU H200再过些日子就要发售了。像Open AI这样的大客户估计是第一时间就能拿到手。到时候GPT无论是推理还是训练,都可能会来到一个非常夸张的量级。
其次,微软也出手了。2019年时,微软启动了一个名为雅典娜的项目。通过设计和开发定制人工智能芯片,以满足训练大型语言模型和驱动人工智能应用的独特需求。而且雅典娜和TPU相似,也是一个内部项目,能够减少对英伟达等第三方硬件提供商的依赖。雅典娜是一个极其神秘的项目,它的性能是多少,它的外观是怎样,外界一概不知。唯一能够了解的,是有些幸运的Open AI员工已经开始对雅典娜测试了。
虽然不清楚H200和雅典娜具体能为Open AI带来哪些变化,不过可以确定的是,谷歌以及谷歌Gemini压力都挺大的。
TPU对谷歌意味着什么?
说到谷歌的TPU还有一个小故事,TPU的全称是Tensor Processing Unit,中文是张量处理单元,是 Google 开发的专用集成电路(ASIC),专门用于加速机器学习。虽然在公众视野中,谷歌是2016年5月的I/O大会上宣布使用TPU,可实际在2015年的时候TPU就已经问世。
神经网络翻译技术从2014年9月提出,百度2015年5月上线首个互联网NMT系统(神经网络机器翻译系统)。NMT技术克服了传统方法将句子分割为不同片段进行翻译的缺点,而是充分利用上下文信息,对句子进行整体的编码和解码,从而产生更为流畅的译文。
但是以当时的技术,神经翻译对服务器造成的负载压力非常大,主要当时硬件没那么强大,也没有H100、A100这样的产品。从发表的论文和谷歌Blog来看,谷歌实际上已经在当时具备了这样的技术,没有使用的核心原因就是硬件遭不住。经过大约半年的测试,TPU已经可以完美适配到服务器中,作为算法加速器使用。那到了I/O大会那天,谷歌一咬牙一跺脚,“我们也有NMT了!”
谷歌的TPU最早是不对外开放的,直到2018年才开始有云TPU业务。就是用户可以在云上购买TPU,来快速解决需要张量计算的业务。2022年的时候,谷歌云上更新了TPU v4业务,意味着现在用户也可以购买并使用上文提到的各种技术来享受更高的推理和训练能力。
谷歌和OpenAI最大的不同在于,前者有很强的业务要求,谷歌Gemini不会是一个赚钱的工具,它对于绝大多数人来说更像是个新鲜的高科技玩具。但是云TPU v4就不一样了,截至3月31日的季度,谷歌云业务的收入为74亿美元,运营利润为1.91亿美元,利润率为2.5%。谷歌云业务实现首次盈利,但是利润率还是太低,尤其是对比亚马逊云28%的利润。所以云TPU就成为了谷歌业务的一个大型增长点,一旦谷歌Gemini的效果反响良好,那么对谷歌云TPU业务的推广来说,无疑是王炸。