三个月、零基础手搓一块TPU，能推理能训练，还是开源的

摘要：对于计算任务负载来说，越是专用，效率就越高，谷歌的 TPU 就是其中的一个典型例子。它自 2015 年开始在谷歌数据中心部署后，已经发展到了第 7 代。目前的最新产品不仅使用了最先进的制程工艺打造，也在架构上充分考虑了对于机器学习推理任务的优化。TPU 的出现

最近，大模型技术的发展，让人们再次重视起 AI 专用芯片。

对于计算任务负载来说，越是专用，效率就越高，谷歌的 TPU 就是其中的一个典型例子。它自 2015 年开始在谷歌数据中心部署后，已经发展到了第 7 代。目前的最新产品不仅使用了最先进的制程工艺打造，也在架构上充分考虑了对于机器学习推理任务的优化。TPU 的出现，促进了 Gemini 等大模型技术的进展。

这种芯片为何性能如此强大？或许最直接的了解方法就是尝试复刻它。近日，来自加拿大西安大略大学的工程师 Surya Sure 等人宣布已经利用暑假时间构建出了 TinyTPU：一种开源的 ML 推理、训练芯片。

项目 GitHub 地址：https://github.com/tiny-tpu-v2/tiny-tpu

有趣的是，他们并非芯片设计专业的学生，打造 TPU 是从理解多层感知机（MLP）这样的神经网络基本概念开始的。为此，他们还手工计算了网络推理和训练所需的数学运算。

让我们看看他是如何做到的。

手搓 TPU 的原因

选择造一块 TPU 的原因很简单：

构建一个用于机器学习工作负载的芯片看起来很酷；

之前还没有一个同时进行推理和训练的机器学习加速器的完整开源代码库。

我们都没有真正的硬件设计专业经验，这在某种程度上使得 TPU 更具吸引力，因为我们无法准确估计它的难度。在项目初期，我们确立了严格的设计理念：始终尝试「Hacky Way」（不靠谱的方法）。这意味着在咨询外部资源之前，先尝试那些我们首先想到的「愚蠢」想法。这种理念确保了我们没对 TPU 进行逆向工程，而是对其进行重新发明，因此我们自己推导出了 TPU 中使用的许多关键机制。

我们也想把这个项目当作一次不依赖人工智能代写代码的练习，因为我们感觉最近每当遇到小问题时，我们的第一反应就是求助于这些 AI 工具。我们希望培养一种特定的思维方式，以便我们能够将其付诸实践，并在未来的任何努力中运用它来解决难题。

在整个项目过程中，我们尽可能多地学习深度学习、硬件设计和算法创建方面的基础知识。我们发现，学习这些知识的最佳方式是将所有知识都画出来，并将其作为我们的第一反应。阅读这篇文章时，你会发现我们的解释是如何受到这种理念的启发的。

在继续之前，我们想先明确一下这篇文章涵盖的内容和不涵盖的内容。请注意，本文并非 TPU 的 1:1 复制品 —— 而是我们自己重新发明 TPU 的尝试。

什么是 TPU？

TPU 是谷歌设计的一款专用芯片 (ASIC)，旨在专门提高机器学习模型的推理和训练速度。GPU 既可以渲染图像，又可以运行机器学习工作负载，而 TPU 专用于执行数学运算，这使得它能够更高效率地完成其设计初衷。在芯片领域，尝试掌握单个任务比尝试掌握多个任务更容易，而且效果也更好，而 TPU 正是秉承了这一理念。

硬件设计快速入门：

在硬件中，我们处理的时间单位称为时钟周期。作为开发者，我们可以根据需要设置任意的时间段。通常，单个时钟周期的范围从 1 皮秒 (ps) 到 1 纳秒 (ns)，我们运行的任何操作都将在时钟周期之间执行。

时钟周期时序图展示了硬件中操作的同步方式。

我们用来描述硬件的语言叫做 Verilog。它是一种硬件描述语言，允许我们描述给定硬件模块的行为（类似于软件中的函数），但它不是以程序的形式执行，而是合成布尔逻辑门（AND 与、OR 或、NOT 非等），这些逻辑门可以组合起来，构建我们想要的任何芯片的数字逻辑。以下是 Verilog 中加法运算的一个简单示例：

在上面的例子中，信号 b 在下一个时钟周期的值被设置为信号 a 的当前值。你会发现，在大多数情况下，信号（变量）是按顺序的时钟周期更新的，而不是像软件设计中那样立即更新。

具体来说，TPU 在执行矩阵乘法方面非常高效，它占 Transformer 计算操作的 80-90%（在超大型模型中高达 95%），以及 CNN 计算操作的 70-80%。每个矩阵乘法都代表 MLP（多层感知器）中单个层的计算，而在深度学习中，我们拥有多层级的 MLP，这使得 TPU 在处理大型模型时更加高效。

那么如何构建一块 TPU 呢？

当我们开始这个项目时，我们只知道方程 y = mx + b 是神经网络的基础构建块。然而，我们需要完全理解神经网络背后的数学原理，才能在 TPU 中构建其他模块。因此在开始编写任何代码之前，我们每个人都计算出了一个简单的 2 → 2 → 1 多层感知器 (MLP) 的数学原理。

用于解决 XOR 问题的 2→2→1 多层感知器的架构。

我们选择这个特定网络的原因是，我们的目标是针对 XOR 问题（神经网络的「Hello World」）进行推理和训练。XOR 问题是神经网络可以解决的最简单的问题之一。所有其他门（AND、OR 等）都可以仅使用一条线性线（一个神经元）来区分哪些输入对应于 0，哪些输入对应于 1，从而根据输入预测输出。但是，要对所有 XOR 进行分类，需要一个 MLP，因为它需要弯曲的决策边界，而这仅靠线性方程无法实现。

对于几何和第一性原理的论述，免费书籍《理解深度学习》（Understanding Deep Learning）非常值得一读：

现在，假设我们要进行连续推理（例如，自动驾驶汽车每秒进行多个预测）。这意味着我们需要同时发送多条数据。由于数据本质上是多维的，并且具有许多特征，因此我们需要非常大的矩阵维度。然而，XOR 问题简化了维度，因为只有两个特征（0 或 1）和 4 种可能的输入数据（0 和 1 的四种二进制组合）。这为我们提供了一个 4x2 矩阵，其中 4 是行数（批处理大小 batch size），2 是列数（特征大小）。

XOR 输入矩阵和目标输出：

每一行代表四个可能的异或输入之一，输出向量显示预期的异或结果。

我们对脉动阵列示例进行的另一项简化是，我们将使用 2x2 阵列，而不是 TPUv1 中使用的 256x256 阵列。然而，数学运算仍然忠实于原有的格式，因此实际上并没有简化，而是进行了规模缩小。

等式的第一步是将 m 与 x 相乘，以矩阵形式表示为 XW^T：

其中 X 是我们的输入矩阵，W 是我们的权重矩阵，b 是我们的偏差向量。

我们如何在硬件中执行矩阵乘法？嗯，我们可以使用一个叫做脉动阵列的单元！

TPU 的核心是一个叫做脉动阵列（systolic array）的单元。它由称为处理单元 (PE) 的独立构建块组成，这些构建块以网格状结构连接在一起。每个 PE 执行乘法 - 累加运算，这意味着它将传入的输入 X 与固定权重 W 相乘，并将其与传入的累加和相加，所有这些都在同一个时钟周期内完成。

处理元件（PE）架构显示乘法累加运算（无负载权重和启动标志）。

当这些 PE 连接在一起时，它们可以以脉动方式执行矩阵乘法，这意味着每个时钟周期都可以计算输出矩阵的多个元素。输入从左侧进入脉动阵列，然后每个时钟周期移动到右侧相邻的 PE。累加和从第一行 PE 的乘法输出开始，向下移动，并与每个连续 PE 的乘积相加，直到到达最后一行 PE，成为输出矩阵的一个元素。

脉动阵列架构展示了 PE 如何连接以执行矩阵乘法。

由于 TPU 只有一个单元（并且矩阵乘法在模型中占据了主要计算量），因此它可以非常轻松地推理和训练任何模型。

示例

现在，让我们来看看异或问题的示例。

我们的脉动阵列接受两个输入：输入矩阵和权重矩阵。对于我们的异或网络，使用以下权重和偏差进行初始化：

输入和权重调度

为了将输入批次输入到脉动阵列中，我们需要：

将 X 矩阵旋转 90 度

矩阵旋转 90 度，为脉动阵列输入做准备。

错开输入（将每行延迟 1 个时钟周期）

用于脉动阵列处理的输入矩阵交错模式。

要输入我们的权重矩阵，我们需要：

交错排列权重矩阵（与输入类似）

用于脉动阵列处理的权重矩阵交错模式

转置它！

权重矩阵转置以实现正确的数学对齐。

请注意，旋转和交错没有任何数学意义 —— 它们只是为了使脉动阵列正常工作。转置也只是为了进行数学记账 —— 它是使矩阵数学运算正常工作所必需的，因为我们在神经网络图中设置权重指针的方式。

交错和先进先出 (FIFO)

为了执行交错，我们为权重和输入设计了几乎相同的累加器，分别位于脉动阵列的上方和左侧。

由于激活是逐个输入到脉动阵列中的，我们认为先进先出队列 (FIFO) 是最佳的数据存储方案。然而，传统的 FIFO 和我们构建的累加器之间略有不同。我们的累加器有两个输入端口 —— 一个用于手动将权重写入 FIFO，另一个用于将激活模块的上一层输出写回到输入 FIFO（上一层的输出是当前层的输入）。

我们还需要以类似的方式为每一层加载权重，因此我们复制了权重 FIFO 的逻辑，但没有第二个端口。

脉动阵列矩阵乘法：

偏差和激活

等式的下一步是添加偏差。为了在硬件中实现这一点，我们需要在脉动阵列的每一列下创建一个偏差模块。我们可以看到，当总和移出脉动阵列的最后一行时，我们可以立即将它们输入到偏差模块中，以计算预激活。我们将用变量 Z 表示这些值。

偏差向量 b 会在矩阵的所有行上转发 —— 这意味着它会被添加到每个 Z 行。

现在，我们的方程看起来很像我们在高中学到的 —— 只不过是多维形式，其中从脉动阵列流出的每一列都代表其自身的特征！

接下来，我们必须应用激活函数，为此我们选择了 Leaky ReLU。[5] 这也是一个逐元素的操作，类似于偏差函数，这意味着我们需要在每个偏差模块下（以及通过代理在脉动阵列的每一列下）都有一个激活模块，并且我们可以将偏差模块的输出立即流式传输到激活模块中。我们将用 H 表示这些激活后的值。

Leaky ReLU 函数逐元素应用：

其中 α=0.5 是我们的泄漏因子。对于矩阵，这适用于每个元素。

以我们的异或运算示例为例，我们来看看第 1 层如何处理数据。首先，脉动阵列计算

然后添加偏差：

最后，LeakyReLU 逐个元素应用：

负值乘以 0.5，正值保持不变。

带偏置和 LeakyReLU 的脉动阵列：

流水线

现在你可能会问 —— 为什么我们不把偏置项和激活项合并在一个时钟周期内？这是因为流水线技术！流水线允许多个操作在 TPU 的不同阶段同时执行 —— 无需等待一个完整的操作完成后再开始下一个操作，而是将工作分解成可以重叠的阶段。

可以把它想象成一条装配线：当一个工作模块（激活模块）处理一个零件时，前一个工作模块（偏置模块）已经在处理下一个零件了。这使得所有模块都保持忙碌状态，而不是让它们闲置等待上一个阶段完成。这也会影响 TPU 的运行速度 —— 如果某个模块试图在一个周期内执行多个操作，那么该模块将成为 TPU 的瓶颈，因为其他模块的运行速度最终只能与该模块相同。因此，尽可能将操作分解成多个时钟周期是高效且最佳实践。

流水线阶段显示操作如何跨时钟周期重叠。

为了尽可能高效地运行芯片，我们采用了另一种机制，即传播「启动」信号，我们称之为「移动芯片使能」（用紫色圆点表示）。由于我们设计中的所有组件都是交错排列的，我们意识到可以非常优雅地在第一个累加器上断言一个时钟周期的启动信号，并在相邻模块需要开启时将其准确传播到它们。

这将延伸到脉动阵列，最终延伸到偏置和激活模块，其中相邻的 PE 和模块（从左上角到右下角）在连续的时钟周期内开启。这确保了每个模块仅在需要时执行计算，并且不会在后台浪费电量。

双倍缓冲

现在，我们知道开始一个新的层意味着我们必须使用新的权重矩阵计算相同的。如果我们的脉动阵列是权重平稳的，我们该如何做到这一点？我们该如何改变权重？

在思考这个问题时，我们偶然发现了双倍缓冲的概念，它源自电子游戏。双缓冲存在的原因是为了防止显示器上出现所谓「画面撕裂」的情况。归根结底，像素加载需要时间，我们希望以某种方式「隐藏」这段时间。如果你仔细观察，就会发现这与我们目前在脉动阵列中面临的问题完全相同。幸运的是，游戏设计师已经想出了解决这个问题的方案。通过添加第二个「影子」缓冲区（在当前层计算时保存下一层的权重），我们可以在计算过程中加载新的权重，从而将总时钟周期数减少一半。

为了实现这一点，我们还需要添加一些信号来移动数据。首先，我们需要一个信号来指示何时切换影子缓冲区和活动缓冲区中的权重。我们将这个信号称为「切换」信号（用蓝点表示），它将影子缓冲区中的值复制到活动缓冲区。它从脉动阵列的左上角传播到右下角（与移动芯片使能的路径相同，但仅在脉动阵列内传播）。然后，我们需要另一个信号来指示何时需要将权重向下移动一行，我们将其称为「接受」标志（用绿点表示），因为每一行都表示接受一组新的权重。这会将新的权重移动到脉动阵列的顶行，并将每一行权重向下移动到脉动阵列的下一行。这两个控制标志协同工作，使我们的双缓冲机制正常工作。

如果你还没注意到，这让脉动阵列能够执行一项强大的功能…… 持续推理！我们可以持续输入新的权重和输入，并计算任意层级的前向传播。这触及了脉动阵列的核心设计理念：我们希望最大化 PE 的利用率。我们希望始终保持脉动阵列的馈送！

对于第 2 层，第 1 层的输出 (H1）现在成为我们的输入：

添加偏差并应用激活：

所有值均为正数，因此它们会保持不变。这些就是我们对异或问题的最终预测！

正向传递演练（使用双倍缓冲）

控制单元和指令集（ISA）

我们推理的最后一步是创建一个控制单元，使其使用自定义指令集 (ISA) 来断言所有控制标志并通过数据总线加载数据。包括数据总线在内，我们的 ISA 长度为 24 位，这使得我们的测试平台更加优雅，因为我们可以每个时钟周期传递一串位，而无需单独设置多个标志。

然后，我们将所有东西整合在一起，推理功能完全正常运行！这是一个重要的里程碑。

反向传播与训练

好了，我们已经解决了推理问题 —— 那么训练呢？有趣的是，我们可以将用于推理的架构用于训练！为什么？因为训练只是矩阵乘法，只是多了一些步骤。

接下来才是真正令人兴奋的地方。假设我们刚刚对异或问题进行了推理，得到的预测结果类似于 [0.8, 0.3, 0.1, 0.9]，而我们实际想要的是 [1, 0, 0, 1]。我们的模型表现很差，需要改进它。这时训练就派上用场了。我们将使用一个叫做损失函数的东西来准确地告诉模型它的表现有多差。为了简单起见，我们选择了均方误差 (MSE)—— 可以把它想象成测量预测结果与实际期望结果之间的「距离」，就像测量篮球投篮偏离目标的距离一样。我们用 L 表示损失。

因此，在完成最后一层的激活函数（我们称之为 H2）的计算后，立即将它们输入到损失模块中，以计算我们的预测有多糟糕。这些损失模块位于激活模块的正下方，并且只有在到达最后一层时才会使用它们。但关键在于：训练时实际上并不需要计算损失值本身。你只需要它的导数。因为这个导数告诉我们应该向哪个方向调整权重以减小损失。这就像拥有一个指向「更好性能」的指南针。

链式法则的魔力

这就是微积分发挥作用的地方。为了改进我们的模型，我们需要弄清楚每个权重的变化如何影响损失。链式法则让我们把这个庞大的计算分解成更小、更易于管理的部分。

梯度的链式法则：

这使我们能够逐层计算梯度，并通过网络向后传播它们。

让我们一步步回顾一下发生的事情。

1、计算

损失相对于我们的最终激活有多少变化。

2、计算通过对激活函数（在我们的例子中是 LeakyReLU）求导。

3、计算

由于 z2 的所有元素均为正，因此 LeakyReLU 梯度为：

前向传播和后向传播的美妙对称性

绘制出完整的计算图后，我们发现了一个惊人的现象：反向传播中的最长链与前向传播非常相似！在前向传播中，我们将激活矩阵与转置后的权重矩阵相乘。在后向传播中，我们将梯度矩阵与未转置的权重矩阵相乘。这就像照镜子一样！

将梯度传播到隐藏层：

并通过第一层的激活：

当 Z1 中正负值混合时，梯度为：

一旦我们有了所有这些单独的导数，我们就可以将它们相乘，以找到关于损失的任何导数。

之后，我们必须计算激活导数。

公式为：

这也是一种逐元素计算，这意味着我们可以像损失模块（以及偏差和激活模块）一样构建它，但它会执行不同的计算。然而，关于这个模块，需要注意的是，它需要我们在前向传播过程中计算的激活值。

现在你可能想知道 —— 我们实际上是如何在硬件中计算导数的？让我们以 Leaky ReLU 为例，它非常简单，但却演示了关键原理。请记住，Leaky ReLU 会根据输入是正还是负应用不同的操作。导数遵循相同的模式：对于正输入，它输出 1，对于负输入，它输出一个小的常数（我们使用了 0.01）。

Leaky ReLU 梯度：

硬件中的 Leaky ReLU 导数实现展示了条件逻辑。

它的好处在于它只是一个简单的比较 —— 无需复杂的算法。硬件可以在一个时钟周期内计算出这个导数，从而保持我们的流水线顺畅运行。同样的原则也适用于其他激活函数：它们的导数通常简化为硬件可以高效执行的基本操作。这一洞见促使我们首先计算长链 —— 获取所有

梯度，就像我们在前向传播中计算激活一样。我们可以缓存这些梯度并重复使用，遵循我们已经掌握的高效模式。

你会注意到一个很棒的模式正在浮现：所有这些位于脉动阵列下方的模块都会处理逐个输出的列向量。这促使我们想到将它们统一起来，形成一个我们称之为向量处理单元 (VPU) 的东西 —— 因为它们正是这样做的，逐元素地处理向量！[6]

这不仅使用起来更加优雅，而且当我们将 TPU 扩展到 2x2 脉动阵列以上时也非常有用，因为我们将拥有 N 个这样的模块（N 是脉动阵列的大小），每个模块都必须单独与它们交互。将这些模块统一到一个父模块下，使我们的设计更具可扩展性和优雅性！

矢量处理单元 (VPU) 架构展示统一的元素级操作。

此外，通过为每个模块整合控制信号（称之为 VPU 通路位），我们可以选择性地启用或跳过特定操作。这使得 VPU 足够灵活，能够同时支持推理和训练。例如，在正向传递过程中，我们希望应用偏差和激活函数，但跳过计算损失或激活函数导数。当转换到反向传递时，所有模块都会参与，但在反向链中，我们只需要计算激活函数导数。由于流水线技术，所有流经 VPU 的值都会经过四个模块，任何未使用的模块都只是充当寄存器，将其输入转发到输出，而不执行计算。

接下来的几个导数很有趣，因为我们实际上可以使用矩阵乘法（和脉动阵列）借助以下三个恒等式来计算导数：

1、如果我们有

并对权重求导，我们得到：

2、如果我们有

并对输入 X 取导数，我们得到：

3、对于偏差项，导数就是 1。

这意味着我们可以将前面的与 X、W^T 和 1 相乘，得到

，得到损失函数关于所有第二层参数的梯度。由于所有梯度实际上都是梯度矩阵，我们可以使用收缩阵列。

，然后我们可以将所有这些乘以

和

现在需要注意的是，激活函数导数

都需要我们在前向传播过程中计算后激活函数 (H)。这意味着我们需要将每一层的输出存储在某种内存中才能进行训练。为此，我们创建了一个新的暂存器内存模块，我们称之为统一缓冲区 (UB)。这样，就可以在前向传播过程中计算 H 值后立即存储它们。

以及权重导数

我们意识到，通过使用 UB 来存储输入和权重累加器，以及手动将偏差和泄漏因子加载到各自的模块中，我们也可以省去这些累加器。相比每个时钟周期使用指令集加载新数据，这也是一种更好的做法。由于我们想要同时访问两个值（脉动阵列每行 / 列的 2 个输入或 2 个权重），我们添加了两个读写端口。我们对每个数据原语（输入、权重、偏差、泄漏因子、后激活）都进行了这样的操作，以最大限度地减少数据争用，因为我们拥有许多不同类型的数据。

要读取值，我们需要提供起始地址和值的数量，以及我们希望 UB 读取的起始地址和位置数量，这样它每个时钟周期就会读取 2 个值。写入是一种类似的机制，我们需要指定要写入两个输入端口的值。读取机制的优点在于，一旦我们提供起始地址，它就会在后台运行，直到读取给定位置的数量，这意味着我们只需要每隔几个时钟周期提供一条指令。