超越GPU：TPU能成为接班人吗？

在计算机的世界里，硬件技术的发展一直在快速推进。今天，我们要谈论的就是一种特殊的处理器：TPU，全称是Tensor Processing Unit。在我们开始深入探讨TPU之前，先了解一下两个重要的芯片技术，FPGA和ASIC。

FPGA和ASIC

Altera是一家FPGA巨头

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 倍。

Google TPU V4

在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各有优势，选择使用哪种硬件取决于具体的应用场景和需求。

关注萤火架构，加速技术提升！