在计算机的世界里,硬件技术的发展一直在快速推进。今天,我们要谈论的就是一种特殊的处理器:TPU,全称是Tensor Processing Unit。在我们开始深入探讨TPU之前,先了解一下两个重要的芯片技术,FPGA和ASIC。
FPGA和ASIC
FPGA,全称为现场可编程门阵列,这是一种特殊的处理器,我们可以通过程序代码来控制其中的电路连线,从而实现某种特定的算法单元。简单来说,它就像一块空白的画板,可以根据需要进行编程和配置,实现各种硬件功能。使用FPGA,我们可以快速低成本地验证硬件设计,或者少量定制芯片的发行。
而ASIC,全称为专用集成电路,是为专门用途设计的芯片,如相机芯片、音频芯片、网卡芯片、挖矿芯片等等。可以说,对FPGA进行编程,就是把FPGA电路变成了一个ASIC。相比之下,FPGA需要的晶体管更多、能耗更大、单片成本也高。然而,FPGA没有硬件研发成本,但ASIC需要仿真、验证、流片等流程,如果不能大量发行,研发成本就不能被摊薄,反而整体成本更高。
了解了这两种技术之后,我们再来来看看TPU是什么?
什么是TPU?
TPU全称为Tensor Processing Unit,是一种张量处理器,属于一种ASIC。
张量在物理学和数学上有着广泛的应用,它经常被用来表达几何性质和物理规律。对于初次接触张量这个概念的同学,应该完全不知其所云,我这里简单介绍下。
大多数同学应该都听说过或接触过向量或者矩阵,在编程中向量可以使用一维数组表达,矩阵可以使用二维数组来表达,那么张量呢?它是多维数组。张量处理器就是专门用来进行多维数组运算的专用处理器,当然用它来搞搞矩阵、向量和单个数字的运算也完全没有问题,因为向量、矩阵都是张量的子集。
理解了什么是TPU,我们还要多问个为什么,为什么需要TPU?
为什么需要TPU?
其实在上面介绍ASIC时,我已经提到过了一些,包括为专门的用途设计和更低的成本,这里再展开详细说明下。
TPU主要是为了解决机器学习的计算问题,特别是深度学习。Google最初设计TPU是用来解决模型推断计算的问题,模型推断就是使用训练好的模型进行计算,输入参数、获取结果。比如我们使用ChatGPT聊天、使用Stable Diffusion出图等等都是模型推断。模型推断有以下特点:
- 训练后的模型一般会被广泛应用,相比训练,后续的计算量会很大;
- 推断工作的灵活性要求低,需要的电路相对简单;
- 推断工作要保证响应时间,尽可能地低,专门的硬件会更容易处理;
- 推断工作要尽可能低的功耗,还是因为大规模计算问题,累积的能耗不可忽视。
根据这些特点,使用专用的硬件设计就很有必要了,TPU就产生了,它内置了大量的矩阵运算单元,可以并行处理大量的矩阵运算,从而相比CPU和GPU大大提高了计算效率;同时由于它的专用硬件设计,其能耗也更低,这对于大规模计算降低成本特别重要。在Google论文里,它可以比CPU、GPU 快15~30 倍,能耗比可以高出 30~80 倍。
在TPU后来的迭代演化中,Google已经发布了四代TPU,它不仅能用于模型推断,也开始用于模型训练,每一代的性能都有显著的提升,广泛应用于语音识别、图片识别等深度学习任务中,表现出了很高的效率和性能。除了Google,其他的一些厂商,如阿里巴巴、华为等,也研发了自己的定制化AI芯片,这些芯片在设计理念上与TPU类似,都是为了提高深度学习任务的计算效率。
TPU的实现
我们继续了解下TPU是怎么做到这么优秀的。
它的设计思路非常简洁且实用。首先,它可以像显卡一样插在主板的PCI-E接口上,通过CPU向TPU发送指令,这样就可以向前兼容,设计发行时间也大大缩短。
其次,TPU的主要组成部分包括一个矩阵计算单元(固定硬件电路)和一个本地统一缓存区。这两个部分占据了TPU超过一半的空间。这是因为深度学习的计算主要是矩阵运算,矩阵计算单元的存在大大提高了运算速度。而缓存区则使用了SRAM,比常用的DRAM更快,尽管密度小,占用空间大,但是提高了数据的读写速度。
根据深度学习的特点,操作数都能控制在 -3 到 3,精度放到 2^6,一般就够了,因此在Google的TPU中,使用了8bit来表示浮点数,而且这样做的好处有很多。相比于32位或64位的浮点数,8位可以大大降低计算的复杂性,提高计算效率;还可以减少存储需求,降低内存带宽的压力。
TPU的局限性
如果你经常浏览AI资讯,提到硬件计算,似乎没什么人提及TPU,大多数谈的还是GPU,比如美帝最近又限制了Nvidia高端显卡的出口,甚至传言消费级的4090都要限制了,各大厂商都还在使用GPU搞深度学习、搞大模型,怎么没人用TPU呢?这是怎么回事呢?
首先虽然TPU的设计理念非常先进,性能和效率都比传统的CPU和GPU要高。然而,由于TPU是专门为深度学习设计的,所以在其他类型的计算任务中,它的性能可能并不出色。此外,TPU还需要特定的软件支持,这也可能限制了它的广泛应用。
其次成本,虽然使用TPU的总体成本比较低,但是由于TPU的生产技术相对复杂,生产成本相对较高,专业的芯片厂商可能不懂也不太关注后期的运行成本,生产积极性不高,所以一直是Google、华为、阿里巴巴这种云计算企业在引领这个区域,这是普通用户不容易接触到的。
最后,关于为什么很多厂商都在使用GPU做训练和推理,这主要是因为GPU的通用性和成熟的软件生态。GPU不仅可以用于深度学习,还可以用于图形处理、科学计算等多种任务。此外,GPU得到了TensorFlow、PyTorch等主流深度学习框架的良好支持,这使得开发者可以很方便地使用GPU进行深度学习的开发。相比之下,虽然TPU在深度学习任务中的性能更高,但由于其专用性和软件支持的限制,目前还无法完全替代GPU。
你觉得未来TPU能接GPU的班吗?
更多学习资源
我注册了一个微/信/公/众/号:萤火架构,后续会分享很多架构方面的真实经验和认知,欢迎关注,以免错过精彩内容。
总的来说,TPU是一种强大的处理器,它在处理深度学习任务时表现出色,能大大提高计算效率,降低能耗,总体成本较低。但是显卡和TPU各有优势,选择使用哪种硬件取决于具体的应用场景和需求。