前言
大家好,我是jiantaoyab,上篇文章讲了FPGA和ASIC,讲解了 FPGA 如何实现通过“软件”来控制“硬件”,以及我们可以进一步把 FPGA 设计出来的电路变成一块 ASIC 芯片。今天我们来看看TPU。大家可以点击这篇文章TPU深入了解TPU。
TPU解决什么问题?
深度学习中计算量最大的过程其实是推断的过程。
推断部分,是指我们在完成深度学习训练之后,把训练完成的模型存储下来。这个存储下来的模型,是许许多多个向量组成的参数。然后,我们根据这些参数,去计算输入的数据,最终得到一个计算结果。
这个推断过程,可能是在互联网广告领域,去推测某一个用户是否会点击特定的广告;也可能是我们在经过高铁站的时候,扫一下身份证进行一次人脸识别,判断一下是不是你本人。
第一代的 TPU,首先优化的并不是深度学习的模型训练,而是深度学习的模型推断。
这个时候你可能要问了,那模型的训练和推断有什么不同呢?主要有三个点。
第一点,深度学习的推断工作更简单,对灵活性的要求也就更低。模型推断的过程,我们只需要去计算一些矩阵的乘法、加法,调用一些 Sigmoid 或者 RELU 这样的激活函数。这样的过程可能需要反复进行很多层,但是也只是这些计算过程的简单组合。
第二点,深度学习的推断的性能,首先要保障响应时间的指标。计算机关注的性能指标,有响应时间(Response Time)和吞吐率(Throughput)。我们在模型训练的时候,只需要考虑吞吐率问题就行了。因为一个模型训练少则好几分钟,多的话要几个月。而推断过程,像互联网广告的点击预测,我们往往希望能在几十毫秒乃至几毫秒之内就完成,而人脸识别也不希望会超过几秒钟。很显然,模型训练和推断对于性能的要求是截然不同的。
第三点,深度学习的推断工作,希望在功耗上尽可能少一些。深度学习的训练,对功耗没有那么敏感,只是希望训练速度能够尽可能快,多费点电就多费点儿了。这是因为,深度学习的推断,要 7×24h 地跑在数据中心里面。而且,对应的芯片,要大规模地部署在数据中心。一块芯片减少 5% 的功耗,就能节省大量的电费。而深度学习的训练工作,大部分情况下只是少部分算法工程师用少量的机器进行。很多时候,只是做小规模的实验,尽快得到结果,节约人力成本。少数几台机器多花的电费,比起算法工程师的工资来说,只能算九牛一毛了。
这三点的差别,也就带出了第一代 TPU 的设计目标。那就是,在保障响应时间的情况下,能够尽可能地提高能效比这个指标,也就是进行同样多数量的推断工作,花费的整体能源要显著低于 CPU 和 GPU。
FPU 的设计
专用电路和大量缓存,适应推断的工作流程
TPU并没有设计成一个独立的“CPU“,而是设计成一块像显卡一样,插在主板 PCI-E 接口上的板卡。
更进一步地,TPU 甚至没有像我们之前说的现代 GPU 一样,设计成自己有对应的取指令的电路,而是通过 CPU向 TPU 发送需要执行的指令。
在芯片模块图里面,有单独的矩阵乘法单元(Matrix Multiply Unit)、累加器(Accumulators)模块、激活函数(Activation)模块和归一化 / 池化(Normalization/Pool)模块。而且,这些模块是顺序串联在一起的。
这是因为,一个深度学习的推断过程,是由很多层的计算组成的。而每一个层(Layer)的计算过程,就是先进行矩阵乘法,再进行累加,接着调用激活函数,最后进行归一化和池化。
硬件的设计就是把整个流程变成一套固定的硬件电路。这也是一个 ASIC 的典型设计思路,其实就是把确定的程序指令流程,变成固定的硬件电路。
控制电路(Control)只占了 2%。这是因为,TPU 的计算过程基本上是一个固定的流程。不像CPU 那样,有各种复杂的控制功能,比如冒险、分支预测等等。
超过一半的 TPU 的面积,都被用来作为 Local Unified Buffer(本地统一缓冲区)(29%)和矩阵乘法单元(Matrix Mutliply Unit)了。
相比于矩阵乘法单元,累加器、实现激活函数和后续的归一 / 池化功能的激活管线(Activation Pipeline)也用得不多。
这是因为,在深度学习推断的过程中,矩阵乘法的计算量是最大的,计算也更复杂,所以比简单的累加器和激活函数要占用更多的晶体管。
而统一缓冲区(Unified Buffer),则由 SRAM 这样高速的存储设备组成。SRAM 一般被直接拿来作为 CPU 的寄存器或者高速缓存。
SRAM 比起内存使用的 DRAM 速度要快上很多,但是因为电路密度小,所以占用的空间要大很多。统一缓冲区之所以使用 SRAM,是因为在整个的推断过程中,它会高频反复地被矩阵乘法单元读写,来完成计算。
在性能上,TPU 比现在的 CPU、GPU 在深度学习的推断任务上,要快 15~30 倍。而在能耗比上,更是好出 30~80 倍。另一方面,Google 已经用 TPU 替换了自家数据中心里 95% 的推断任务。
30 倍。而在能耗比上,更是好出 30~80 倍。另一方面,Google 已经用 TPU 替换了自家数据中心里 95% 的推断任务。