由于GPU在AI硬件市场的主导地位,NVIDIA的股价飙升。然而与此同时,谷歌众所周知的AI硬件TPU却并不对外出售,用户只能在谷歌云上租用虚拟机来使用它们。为什么谷歌没有加入销售AI硬件的游戏呢?
注释:TPU,全称为Tensor Processing Unit(张量处理单元),是由谷歌专门为加速机器学习工作负载而设计的一种应用特定集成电路(ASIC)芯片。与传统的CPU和GPU相比,TPU针对机器学习中的特定计算任务,如神经网络的训练和推理,进行了优化。
一个普遍的猜测是,谷歌想要吸引更多的客户使用其云服务。
不过,根据云服务客户的说法,这个理论并没有太多依据。没有企业级客户想要被锁定在一个特定的云服务提供商上,他们希望能够在需要的时候灵活地转移到另一个云。否则,如果供应商提高价格,他们将无能为力。如果他们因为使用TPU而被锁定在谷歌云上,他们宁愿不使用它。
所以,使用TPU吸引客户到谷歌云并不是不单独销售它们的理由。那么,真正的原因是什么呢?为了更好地回答这个问题,我们必须深入了解谷歌是如何启动TPU项目的。
谷歌为什么要启动TPU项目?
简单地说,TPU是为谷歌的特定需求设计的。早年间,GPU是无法满足深度学习的计算需求的。
我们尝试估计一下TPU项目是何时开始的。鉴于它是在2016年首次向公众发布,合理的猜测是它大约在2011年开始。如果是这样的话,谷歌相当有前瞻性,因为直到2012年AlexNet在ImageNet竞赛中取得了突破性的成绩,深度学习才在计算机视觉领域证明了其价值。
这说明谷歌很早就看到了这场AI革命,并希望为大规模计算提供更快的硬件,当时他们唯一的选择是为此构建一个新的解决方案。
这就是谷歌启动这个项目的原因,当然还有更多的问题。为什么当时的GPU不够好?谷歌看到了哪些潜在的改进,以至于有足够的理由启动他们的新硬件项目?
答案在于GPU和TPU的微架构,我们来对比一下GPU和TPU核心的设计思路。
GPU和TPU的设计思路对比
首先,让我们快速回顾一下CPU的背景知识。当一条指令到来时,它由指令解码器解码,并与寄存器中的数据一起输入到算术逻辑单元(ALU)。ALU完成所有的计算,并将结果返回到寄存器之一。如果CPU有多个核,它们可以并行工作。
而GPU呢?它是“图形处理单元”(Graphics Processing Unit)的缩写。它最初被设计用于图形计算,后来人们发现它也适用于机器学习任务。GPU中的大多数操作都是矩阵运算,这些运算可以并行执行。这也意味着与CPU相比,GPU不需要支持那么多种类的操作。
芯片执行任务的专业化程度越高,它在该任务上的运行速度就越快。
GPU的初始设计理念是创建一个功能简化版的CPU,它拥有更小但数量更多的核心,以便进行更快速的并行计算。与CPU相比,GPU支持的指令数量要少得多,这使得芯片上单个核心所占用的面积小得多。这样,它们就可以在芯片上集成更多的核,以实现大规模的并行计算。
为什么功能更少意味着芯片上占用的面积更小呢?在软件中,更多功能意味着更多代码。而在硬件中,所有功能都是使用逻辑电路而不是代码实现的。更多功能意味着电路更复杂,像CPU就必须在芯片上实现更多指令。
TPU进一步发展了为深度学习设计专用芯片的理念。TPU的标志性特征是其矩阵乘法单元(MXU)。由于矩阵乘法是深度学习中最频繁的操作,TPU为其构建了一个专用核心,即MXU。
这比GPU核心更加专业化,MXU只做一件事:矩阵乘法。
它的工作方式与传统的CPU/GPU核心非常不同。所有的动态性和通用性都被去除了。它有一个完全连接在一起的节点网格,每个节点仅以预定义的方式进行乘法和加法,结果直接被推送到下一个节点进行下一次乘法和加法。因此,一切都是预定义和固定的。
这样,我们就省去了指令解码的时间,因为它只是乘以和加上它接收到的任何内容。由于我们已经知道结果应该去哪里,并且没有必要将其存储用于接下来的任意操作,因此寄存器也省掉了。
除了MXU,TPU还为更好的可扩展性而设计。它有专门的端口用于高带宽的芯片间互连(ICI),可以以集群的形式使用。
TPU的速度和价格
从时间线来看,似乎NVIDIA在大约相同的时间有了类似的想法,它在2017年首次发布添加用于深度学习的张量计算核心(Tensor Cores)的V100 GPU,比谷歌的TPU公告晚了一年。
TPU今天是否仍然比GPU更快,目前并不清楚,我们找不到最新一代TPU和GPU的公开基准测试。不过,在谷歌云的定价模式下,如果用户使用的是相同的机器学习模型、数据集和训练周期下,那么使用TPU的成本会比使用其他硬件更低。
为什么TPU更便宜呢?一个重要原因是软件堆栈。
在软件堆栈上,谷歌的垂直集成度比GPU更好。谷歌有专门的工程团队为其构建整个软件堆栈,从模型实现(Vertex Model Garden)到深度学习框架(Keras、JAX和TensorFlow)再到为TPU优化的编译器(XLA)。
相比之下,GPU的生态则更为分散,它通常与多种不同的深度学习框架一起使用,如由Meta开发的PyTorch。这些框架和工具可能没有为特定类型的GPU进行同样程度的优化,导致性能可能不如专门为TPU设计的软件。
谷歌为什么不开始销售TPU?
TPU项目是出于专有使用目的而启动的,并且由于其较低的价格,在谷歌云上获得了相当好的用户基础。为什么谷歌不直接像NVIDIA一样向客户销售硬件呢?
简单地说, 可能主要有两个考量:
竞争格局:谷歌正在与OpenAI就生成性AI进行激烈的竞争,同时,为了降低成本,它正在持续裁员,不得不将有限的资源集中在最重要的项目上。如果谷歌开始销售TPU,它将同时与两个强大的对手竞争,NVIDIA和OpenAI,这可能不是目前最明智的策略。
成本和管理:直接销售硬件涉及高昂的开销和复杂的供应链管理。而通过云服务提供TPU,谷歌可以简化安装、部署和管理过程,减少不确定性和额外开销。这样的模式也简化了销售流程,无需额外建立硬件销售团队。
TPU方案的优势和缺点
这样的策略让谷歌获得了一些优势。
首先,他们可以拥有更具激进的TPU架构设计。TPU在连接芯片方面有独特的方式。与连接到同一板的多个GPU不同,TPU以立方体形式组织。他们将64个TPU排列在一个4x4x4的立方体中,以便它们之间进行更快的芯片间通信。在每个TPUv5pPod网络中有8960个芯片。它们可以很容易地一起使用。
其次,他们可以更快地迭代。由于他们只需要支持专有使用的一小部分用例,这大大减少了他们每一代芯片的研发周期。
从服务于其最重要的目的(在生成AI中竞争)的角度来看,这些优势使谷歌处于非常有利的位置。最重要的是,有了这种内部硬件解决方案,谷歌在进行AI研究和开发AI应用方面节省了大量的资金,因为它不需要以垄断价格从NVIDIA购买GPU。
到目前为止,我们已经讨论了谷歌TPU策略的许多优点,但是有没有缺点呢?当然,最大的问题是,谷歌变成了一个技术孤岛。
世界其他地区正在使用HuggingFace的模型和PyTorch进行创新,每个人都在快速地调整彼此的模型以开发更好的模型。然而,由于其基础设施主要围绕TensorFlow和JAX构建,谷歌无法轻松加入这个过程。当将外部模型投入生产时,它必须使用谷歌的框架重新实施一遍。
这个“技术孤岛”问题减慢了谷歌从外部世界获取良好解决方案的速度,并进一步将其与其他人隔离了。
AI硬件的未来是什么样的?
最后,让我们窥视一下AI硬件的未来。这里观察到两个趋势:
一方面,硬件将与应用程序进一步耦合。例如,支持更多精度格式以更好地服务大语言模型。NVIDIA今年发布了内嵌第二代Transformer引擎的Blackwell GPU,这使得在不改变用户代码的情况下,更容易地优化他们的硬件以适应Transformer模型,大量的协同设计正在发生。
另一方面,软件/算法很难轻易地跳出Transformer领域。否则由于缺乏硬件支持,它们会变慢。相反,为了提高性能,软件开发者会考虑硬件的特性来设计他们的模型。这意味着他们会优化算法,以便更好地利用硬件的特定功能和架构。
结果是硬件和软件彼此优化,它们都很难轻易跳脱当前的设计或算法。如果有一个与Transformer完全不同的新模型,它需要比Transformers好10倍才有可能被广泛采用。
总结
简单来说,谷歌启动TPU项目是因为当时GPU的计算能力跟不上需求,而TPU能提供更快、成本更低的计算,这对谷歌的AI研究和应用开发非常有帮助。但这也意味着谷歌在技术上有点自成一格,与外界不太兼容。未来,AI硬件可能会更加专门化,为某些AI模型,比如为Transformer专门设计。这会让硬件和模型结合得更加紧密,但也意味着它们很难跳出这个框架寻找新的可能性。