--- area: tech-engineering category: ai-ml companies_orgs: [] date: '2025-08-28' draft: true guest: '' insight: '' layout: post.njk project: [] series: '' source: https://www.youtube.com/watch?v=hYrtS0cFaZA speaker: Best Partners TV status: evergreen summary: 加拿大工程师团队在三个月内，从零基础出发，成功设计并开源了一款能同时支持推理和训练的**TPU**原型芯片。他们通过“第一性原理”和“不靠谱方法”的独特理念，深入理解并重新发明了**MLP**、脉动阵列等核心机制，展现了独立解决复杂问题的能力。 tags: - ai-chip-design - principle - technology title: 从零手搓TPU：三个月实现AI推训一体芯片设计的奇迹 --- ### 项目缘起：零基础挑战AI芯片设计 这是一个听起来几乎不可能完成的项目：仅用三个月时间，一群非芯片设计专业的学生，从零开始，成功打造出了一块能够运行的**TPU**（Tensor Processing Unit: 谷歌设计的专用AI芯片，用于加速机器学习模型的推理和训练）原型。这块芯片不仅能进行推理，还能支持训练，并且已完全开源。尽管谷歌的**TPU**已迭代至第七代，背后有顶尖工程师团队和先进工艺支持，但来自加拿大西安大略大学的工程师们，利用一个暑假的时间，完成了这项看似不可能的挑战，并将他们的项目命名为TinyTPU。 项目的发起者们最初对**MLP**（Multi-Layer Perceptron: 一种前馈人工神经网络，包含至少一个隐藏层）这样的神经网络基本概念都需从头理解。为了彻底搞清楚网络推理和训练的数学原理，他们甚至手工计算了所有必要的运算。这种从最基础开始的钻研精神，为整个项目奠定了独特的基调。 ### 独特的设计理念与动力 团队成员表示，构建一个机器学习的专用芯片听起来非常酷，同时，当时市面上还没有一个能够同时支持推理和训练的、完整的机器学习加速器开源代码库。更重要的是，他们没有任何硬件设计的相关经验，但这反而成为一种动力，因为无法准确估计难度，从而减少了心理负担。 他们还定下了一个特别的设计理念：始终尝试那些“不靠谱的方法”（Hacky Way），即在咨询外部资源之前，先尝试自己想到的“愚蠢”想法。在这种理念驱动下，他们不是在逆向工程谷歌的**TPU**，而是在真正地重新发明它，许多关键机制都是他们自己推导出来的。 此外，他们希望将此项目作为不依赖AI代写代码的一个练习。面对小问题，许多人习惯性求助AI工具，而他们则希望培养独立解决难题的思维方式。在整个过程中，他们尽可能多地学习了深度学习、硬件设计和算法创建的基础知识，并发现最好的学习方式就是将所有知识画出来，这在他们的技术文档中得到了淋漓尽致的体现。 ### TPU基础概念与高效秘诀 需要明确的是，TinyTPU并非谷歌**TPU**的1:1复制品，而是基于团队自身理解重新发明的一个尝试，因此接下来将展示一条全新的探索路径。 首先，了解**TPU**的本质至关重要。**TPU**是谷歌设计的**ASIC**（Application-Specific Integrated Circuit: 专用集成电路，为特定用途而设计）芯片，专门用于提高机器学习模型的推理和训练速度。与**GPU**（Graphics Processing Unit: 图形处理器）既能渲染图像又能运行机器学习任务不同，**TPU**完全专注于数学运算，这也是其效率如此之高的重要原因，因为在芯片领域，专注单一任务往往比兼顾多项任务更容易做得好。 理解硬件设计需要两个基础概念： 1. **时钟周期**: 这是硬件处理操作的时间单位，开发者可以设置任意时间段，通常在1皮秒到1纳秒之间，所有操作都在这个时钟周期内执行。 2. **硬件描述语言Verilog**（硬件描述语言: 一种用于描述数字电路的语言）: 它与软件编程语言不同，不是以程序形式执行的，而是通过合成布尔逻辑门（如与、或、非等）来构建芯片的数字逻辑。在Verilog中，一个简单的加法运算，需要将信号`b`在下一个时钟周期的值设置为信号`a`的当前值，这与软件中变量立即更新的方式截然不同。 **TPU**之所以高效的另一个原因在于其在执行矩阵乘法时的优势。在**Transformer模型**（Transformer Model: 一种主要用于自然语言处理的深度学习模型架构）中，矩阵乘法占计算操作的80%-90%，超大型模型中甚至可高达95%；即便在**CNN**（Convolutional Neural Network: 卷积神经网络，一种常用于图像处理的深度学习模型）中，也能占70%-80%。每个矩阵乘法都代表了**MLP**中单个层的计算，而深度学习模型往往包含多层**MLP**，这使得**TPU**在处理大型模型时的效率倍增。 ### 从零开始构建：异或问题与数学原理 这群零基础的工程师如何从零开始构建**TPU**呢？他们的起点非常基础，只知道`y = mx + b`这个方程是构建神经网络的基础模块。然而，要在硬件中实现神经网络，必须完全理解其背后的数学原理。因此，在编写任何代码之前，团队每个人都手工计算了一个简单的、用于解决**异或**（XOR: 一种逻辑运算，当两个输入不同时输出真，相同时输出假）问题的2→2→1的**MLP**网络。 选择**异或**问题是因为它被称为神经网络的“Hello World”，是最简单的需要用神经网络解决的问题之一。像与（AND）、或（OR）这些逻辑门，用一条线（即一个神经元）就能区分输入对应的输出，但**异或**问题需要弯曲的决策边界，这只有**MLP**才能实现。 假设要进行连续推理，如自动驾驶汽车每秒进行多个预测，就需要同时处理多条数据。这种数据通常是多维的，具有许多特征，需要大维度的矩阵。但**异或**问题简化了维度，只有两个特征（0或1）和四种可能的输入，从而形成一个4x2的矩阵（4是批处理大小，2是特征大小），这为他们的实验提供了一个合适的简化模型。 **异或**的输入矩阵和目标输出都很简单：四行分别代表四种二进制组合（0,0）、（0,1）、（1,0）、（1,1），对应的目标输出是0、1、1、0。 矩阵乘法是神经网络计算的核心，在数学上表示为`XW^T + b`，其中`X`是输入矩阵，`W`是权重矩阵，`b`是偏差向量。 ### 硬件推理实现：脉动阵列与流水线技术 为了在硬件中执行矩阵乘法，他们用到了一个关键结构——**脉动阵列**（Systolic Array: 一种并行计算架构，数据以“脉动”方式流经处理单元）。他们对这个**脉动阵列**做了简化，用2x2的阵列代替了**TPUv1**中的256x256阵列，但数学运算是完全类似的，只是规模缩小了。 **TPU**的核心就是**脉动阵列**，它由**PE**（Processing Element: 处理单元，脉动阵列中的基本计算节点）组成。这些**PE**以网格状结构连接，每个**PE**都会执行乘法累加运算，即将传入的输入`X`与固定权重`W`相乘，再与传入的累加和相加。所有操作在一个时钟周期内完成。当这些**PE**连接起来，就能以脉动方式执行矩阵乘法，每个时钟周期可以计算输出矩阵的多个元素。输入从左侧进入**脉动阵列**，每个时钟周期向右移动到相邻的**PE**；累加和从第一行**PE**的乘法输出开始，向下移动，与每个连续**PE**的乘积相加，直到到达最后一行**PE**，成为输出矩阵的一个元素。正因为**TPU**专注于这个核心单元，而矩阵乘法又占模型计算量的绝大部分，所以它能够高效地处理各种模型的推理和训练。 具体到**异或**问题的处理过程，他们的**脉动阵列**可以接受输入矩阵和权重矩阵。对于这个**异或**网络，初始化的权重和偏差有特定的值。将输入批次传入**脉动阵列**需要两个步骤：首先要将`X`矩阵旋转90度，再错开输入，即每行延迟1个时钟周期；输入权重矩阵时也要进行类似的交错排列，并且需要转置。这里需要说明的是，旋转和交错没有数学意义，只是为了让**脉动阵列**正常工作，而转置则是为了数学对齐，确保矩阵运算的正确。 为了实现交错，他们在**脉动阵列**的上方和左侧设计了几乎相同的累加器。由于激活是逐个输入的，所以适合采用**FIFO**（First-In, First-Out: 先进先出队列，一种数据存储和访问方式）的数据存储方案。但他们的累加器有两个输入端口，一个用于手动写入权重，另一个用于将上一层的输出写回作为当前层的输入。权重**FIFO**的逻辑类似，但没有第二个端口。 完成矩阵乘法后，下一步是添加偏差。他们在**脉动阵列**的每一列下创建了一个偏差模块。当总和从最后一行流出时，直接会进入偏差模块计算预激活的值，用变量`Z`表示。偏差向量会被添加到每个`Z`行，这一步就像线性方程一样，只是扩展到了多维形式，每一列都代表一个特征。 之后需要应用**激活函数**（Activation Function），他们选择了**Leaky ReLU**（Leaky Rectified Linear Unit: 一种激活函数，允许小的负值通过，解决ReLU的“死亡神经元”问题）。这是一个逐个元素的操作，所以在每个偏差模块下都有一个激活模块，偏差模块的输出会直接流入到激活模块，得到的结果用`H`表示。**Leaky ReLU**的计算规则是：输入为正时输出本身，为负时输出输入乘以0.5（泄漏因子β=0.5）。以**异或**网络的第一层为例，**脉动阵列**会先计算矩阵乘法，然后添加偏差`b1`，再应用**Leaky ReLU**得到`H1`，负值乘以0.5，正值保持不变。 这里有个关键设计思路是**流水线技术**（Pipelining: 一种将操作分解为多个阶段，并行处理以提高效率的技术）。为什么不把偏差和激活合并到一个时钟周期里呢？因为**流水线**允许不同阶段的操作同时进行，就像装配线一样。当激活模块处理一个数据时，偏差模块已经在处理下一个了，这样能让所有模块保持忙碌，避免闲置。如果某个模块在一个周期内执行多个操作，就会成为瓶颈，所以将操作分解到多个时钟周期是更高效的方法。 为了进一步提高效率，他们设计了一个称为“travelling chip enable”的信号，用紫色圆点表示。由于所有组件都是交错排列的，只需在第一个累加器上断言一个时钟周期的启动信号，就能准确传播到相邻模块，依次激活**脉动阵列**、偏差和激活模块，确保每个模块只在需要的时候工作，不浪费电量。 接下来的挑战是如何处理多层网络。当开始新的层时，需要使用新的权重矩阵，而**脉动阵列**是权重平稳的。如何更换权重呢？他们借鉴了电子游戏中的**双缓冲**（Double Buffering: 一种通过使用两个缓冲区交替工作来提高效率或避免显示问题的技术）概念，通过添加第二个“影子”缓冲区，在当前层计算时加载下一层的权重，从而让总时钟周期减少了一半。 为了实现**双缓冲**，他们还添加了两个信号：用蓝点表示的“切换”信号，它会将影子缓冲区的权重复制到活动缓冲区，从左上角传播到右下角；以及一个用绿点表示的“接受”标志，表示权重要向下移动一行，新权重进入顶行，每行依次下移。通过这两个信号的配合，就能让**脉动阵列**执行持续的推理，通过不断输入新的权重和数据，计算任意层级的前向传播，最大化**PE**的利用率。例如，在第二层中，第一层的输出`H1`成为输入，经过同样的矩阵乘法、加偏差、激活操作，得到最终的预测结果。对于**异或**问题来说，所有值都是正数，所以**Leaky ReLU**的输出保持不变。 ### 训练功能实现：反向传播与梯度下降 推理功能完成后，他们面临的下一个挑战是训练。有趣的是，用于推理的架构可以直接用来训练，因为训练本质上也是矩阵乘法，只是多了一些步骤。假设推理得到的预测结果是`[0.8, 0.3, 0.1, 0.9]`，而目标是`[1, 0, 0, 1]`，显然模型表现不佳，需要通过训练来改进。 训练的核心是损失函数，他们选择了**均方误差**（MSE - Mean Squared Error: 一种常用的回归损失函数，衡量预测值与真实值之间差异的平方平均值），可以理解为预测与目标之间的“距离”，用`L`表示。但训练并不需要损失值本身，而是需要它的导数，这个导数指示了应该向哪个方向调整权重才能减小损失，就像一个指向“更好性能”的指南针一样。 这里还用到了微积分中的**链式法则**（Chain Rule: 微积分中的一个规则，用于计算复合函数的导数），它能将复杂的梯度计算分解成更小的部分，逐层计算梯度并向后传播。整个过程可以分为几步：首先计算损失相对于最终激活的变化率，然后通过激活函数的导数（即**Leaky ReLU**）来计算损失相对于预激活的梯度，以此类推。 等绘制完完整的计算图后，他们发现了一个惊人的对称性：反向传播中的最长链与前向传播非常相似。在前向传播中，激活矩阵是与转置权重矩阵相乘的；而反向传播中，梯度矩阵是与未转置的权重矩阵相乘的，就像是照镜子一样。这种对称性让他们能够复用许多前向传播的设计，例如将梯度传播到隐藏层，再通过第一层的激活，只需要根据`Z1`中正负值的混合，应用相应的**Leaky ReLU**梯度即可（例如正值为1，负值为0.01）。 在计算激活导数时，还需要用到前向传播中得到的激活值`H`，这就需要存储每一层的输出。为此他们创建了一个**统一缓冲区**（UB - Unified Buffer: 一种设计，用于存储和管理数据，提供多个读写端口以减少数据争用）。在前向传播计算`H`值后，立即存储。**UB**还设计了两个读写端口，因为需要同时访问两个值，包括每行/列的2个输入或者权重，这样就最大限度地减少了数据争用。读取时，只需提供起始地址和数量，**UB**就会在后台运行，每个时钟周期读取2个值，比每个周期用指令加载更加高效。 由于位于所有**脉动阵列**下方的模块都需要处理逐个输出的列向量，这促使他们将这些模块统一为**向量处理单元**（VPU - Vector Processing Unit: 一种专门处理向量运算的处理器单元）。这不仅让设计变得更加优雅，也更具有可扩展性，例如当**脉动阵列**超过2x2时，`N`个这样的模块就可以统一管理。 他们还对激活导数模块做了一个小小的优化。由于`H2`的值只用来计算了一次梯度，所以他们在**VPU**内部创建了一个小型缓存（H-cache），而把其他的`H`值存储在**UB**中，因为需要用来进行多个导数计算。 **VPU**通过整合控制信号（即**VPU**的通路位，pathway bits），可以选择性地启用或者跳过特定的操作，同时支持推理和训练。例如，正向传播时会应用偏差和激活，跳过损失和导数计算；反向传播时则启用所有的模块，但只在反向链中计算激活函数的导数。由于**流水线技术**，流经**VPU**的值会经过所有四个模块，只不过没有使用的模块只会充当寄存器，不执行计算。 接下来的几个导数计算可以直接利用**脉动阵列**的矩阵乘法能力，因为有三个关键的恒等式：如果`Z = XW^T`，那么对权重求导和对输入`X`求导可以得到特定形式，而偏差的导数就是1。这意味着可以用前面得到的梯度与`X`、`W^T`和1相乘，得到权重和输入的梯度，再乘以学习率来更新参数。 有了这些设计之后，反向传播就能顺畅运行。首先从**UB**获取梯度和`H`矩阵，通过**脉动阵列**计算权重梯度，同时直接输入到梯度下降模块，再用当前权重和梯度更新参数。整个过程就会像流水一样不间断地执行。 随着功能的完善，他们的指令集也从最初的24位扩展到了94位，但确保了每一位都是必需的，从而在不影响速度和效率的情况下实现了所有的控制功能。最终，通过不断重复前向传播、反向传播、权重更新的循环，网络性能逐步提升。在**GTKWave**（一种开源的波形查看器，用于查看和分析数字电路模拟结果）中的波形模拟显示，一个周期后内存中的权重和偏差已经得到了有效更新。 ### 总结与展望 这个项目虽然只是一个简单的原型，真正的**TPU**远比这复杂得多，但整个过程中最令人惊叹的，不只是他们在三个月内从零开始构建出了能推理和训练的**TPU**，更在于他们采用的“第一性原理”的方法。他们从最基础的数学和硬件原理出发，重新发明而不是复制了现有的设计。这种方法让他们在没有专业背景的情况下，不仅完成了技术挑战，也更为深入地理解了每一个环节背后的原理。 对于想要深入了解这个项目的观众，可以访问他们的GitHub地址以及项目的官网。团队成员苏里亚·苏雷（Surya Sure）也在X平台进行了分享，所有链接都可在视频简介中找到。