张量处理单元 TPU 是什么？谷歌为神经网络定制的芯片

张量处理单元（Tensor Processing Unit，TPU）是谷歌公司自主研发的一种专门用来加速神经网络计算的芯片。它属于专用集成电路，是为深度学习中最核心的运算，也就是大规模的矩阵乘法，量身定制的。谷歌在 2016 年的开发者大会上首次公开了它，并透露它已经在谷歌的数据中心中运行了一年多。

为什么谷歌要自己造芯片

大约在 2013 年，谷歌的工程师估算：如果每个用户每天只用语音搜索三分钟，按照当时的计算方式，运行深度神经网络所需的计算量，就需要把谷歌的数据中心规模扩大一倍。这促使谷歌决定开发一种专门用于神经网络计算的芯片，以大幅提高效率、降低成本。第一代 TPU 从立项到在数据中心中部署只用了大约 15 个月。2017 年，谷歌发表论文介绍了第一代 TPU 的设计，报告称在执行神经网络推理时，它的性能和每瓦性能都比同时期的 CPU 和 GPU 高出一个数量级左右。

核心设计：脉动阵列

神经网络的计算主要由大量的矩阵乘法构成。TPU 的核心是一个由大量乘加单元组成的二维阵列，第一代 TPU 中有 256 乘 256，也就是六万多个乘加单元，采用“脉动阵列”的结构：数据像心脏泵血一样，有节奏地在阵列中的计算单元之间流动，每个单元从相邻单元接收数据，做一次乘法和加法，再把结果传给下一个单元。这样，数据从内存中读取一次之后，可以在阵列中被反复地使用很多次，大大减少了数据在芯片和内存之间来回搬运的次数，而数据搬运恰恰是功耗和速度的主要瓶颈。TPU 还去掉了通用处理器中用于复杂控制、分支预测、缓存管理等的大量电路，把芯片面积尽可能地用在计算上。第一代 TPU 使用 8 位整数进行计算，专门用于推理；后来的版本支持 bfloat16 这种谷歌提出的浮点格式，它保留了与 32 位浮点数相同的数值范围、但精度较低，非常适合神经网络的训练，如今已被众多芯片厂商采用。

发展历程

第一代 TPU 只用于推理。2017 年发布的第二代 TPU 开始支持训练，并且可以通过高速的专用网络把许多块芯片连接起来，组成被称为“Pod”的超级计算集群。此后谷歌几乎每隔一两年就推出新一代的 TPU，性能和规模不断提升，后续的版本中使用了液冷技术，并引入光学电路交换等技术来灵活地连接成千上万块芯片。谷歌从 2018 年起通过谷歌云对外提供 TPU 的租用服务。AlphaGo、AlphaFold、BERT、PaLM 以及 Gemini 系列大模型，都主要在 TPU 上训练。

软件和生态

TPU 主要通过谷歌开发的 XLA 编译器来使用：开发者用 TensorFlow、JAX 或 PyTorch 编写模型，XLA 把计算图编译成能在 TPU 上高效运行的指令。JAX 与 TPU 的配合尤其紧密，被谷歌内部的大量研究所采用。与英伟达的 GPU 相比，TPU 在大规模的矩阵运算和大集群的训练上效率很高，但通用性较弱，软件生态也不如 CUDA 那样成熟和广泛。TPU 的成功推动了整个行业开发 AI 专用芯片的热潮，亚马逊、微软、Meta，以及国内的华为、寒武纪等公司都推出了各自的 AI 芯片。