最近,随着大模型技术的飞速发展,AI专用芯片再次成为业界关注的焦点。在这些专用芯片中,谷歌的TPU以其卓越的性能和广泛的应用而闻名。自2015年首次在谷歌数据中心部署以来,TPU已经迭代至第七代,不仅在制程工艺上采用了最先进的技术,还在架构设计上对机器学习推理任务进行了深度优化。TPU的成功应用,极大地推动了Gemini等大模型技术的发展。
TPU为何能拥有如此强大的性能?加拿大西安大略大学的工程师Surya Sure等人选择了一条不同寻常的道路来探索这一问题——他们利用暑假时间,成功构建了一种开源的ML推理和训练芯片TinyTPU。这一项目不仅展示了他们的创新精神,也为理解TPU的工作原理提供了宝贵的实践经验。
Surya Sure和他的团队并非芯片设计专业的学生,他们从零开始,通过理解多层感知机(MLP)等神经网络基本概念,手工计算了网络推理和训练所需的数学运算,最终成功构建出了TinyTPU。他们的项目GitHub地址为:https://github.com/tiny-tpu-v2/tiny-tpu,这为对TPU感兴趣的研究者提供了宝贵的资源。
在构建TinyTPU的过程中,Surya Sure和他的团队采用了“Hacky Way”的设计理念,即先尝试那些看似“愚蠢”的想法,而不是直接依赖外部资源。这种理念不仅让他们在没有专业硬件设计经验的情况下成功构建了TPU,还让他们自己推导出了TPU中使用的许多关键机制。
TinyTPU的核心是一个脉动阵列(systolic array),它由多个处理单元(PE)组成,这些PE以网格状结构连接在一起,每个PE都能在同一个时钟周期内完成乘法-累加运算。这种设计使得TinyTPU在执行矩阵乘法时非常高效,而矩阵乘法正是Transformer和CNN等深度学习模型中计算量最大的部分。
为了构建TinyTPU,Surya Sure和他的团队首先计算了一个简单的2→2→1多层感知器(MLP)的数学原理,并选择了XOR问题作为他们的目标问题。XOR问题是神经网络可以解决的最简单的问题之一,它需要弯曲的决策边界,这仅靠线性方程无法实现。通过理解XOR问题的数学原理,他们成功地构建了能够处理这一问题的TinyTPU。
在TinyTPU的设计中,团队还采用了流水线技术和双倍缓冲机制来提高芯片的运行效率。流水线技术允许多个操作在TPU的不同阶段同时执行,而双倍缓冲机制则通过在计算过程中加载新的权重来减少总时钟周期数。这些设计使得TinyTPU能够持续高效地处理输入数据和权重矩阵。
除了推理功能外,TinyTPU还支持训练功能。在训练过程中,TinyTPU使用损失函数来评估模型的性能,并通过反向传播算法来更新权重。反向传播算法利用链式法则将损失函数的导数逐层传播回网络中,从而指导权重的更新方向。TinyTPU的设计使得前向传播和反向传播过程能够无缝衔接,从而实现了高效的模型训练。
Surya Sure和他的团队通过构建TinyTPU不仅深入理解了TPU的工作原理和优势所在,还为AI专用芯片的发展提供了新的思路和方法。他们的项目不仅为学术界和工业界提供了宝贵的资源和实践经验,也为推动AI技术的进一步发展做出了重要贡献。