文章摘要
Tiny-TPU项目由一群硬件设计新手发起,旨在通过挑战自我来理解并重新发明TPU的工作原理。他们采用“先尝试简单方法”的设计哲学,避免直接参考外部资料,从而培养独立思考和解决问题的能力。该项目还强调不依赖AI工具,旨在锻炼一种能够应用于未来复杂问题的思维方式。
文章总结
Tiny-TPU:从零开始的TPU设计与实现
TPU(张量处理单元)是谷歌专为加速机器学习模型的推理和训练而设计的专用芯片。与GPU不同,TPU专注于数学运算,尤其是矩阵乘法,这使得它在处理深度学习模型时更加高效。本文详细介绍了我们如何从零开始设计和实现一个简化的TPU,尽管我们并没有硬件设计的专业背景,但通过不断尝试和探索,我们成功“重新发明”了TPU的核心机制。
设计初衷与挑战
我们选择构建TPU的原因很简单:硬件设计对我们来说是一个全新的领域,这使得TPU更具挑战性。我们希望通过这个项目证明,只要有决心,我们能够完成任何看似不可能的任务。在设计过程中,我们遵循了一个严格的原则:先尝试“笨办法”。这意味着在查阅外部资料之前,我们先尝试自己想到的简单方法。这种思维方式确保我们不是在逆向工程TPU,而是在重新发明它,从而帮助我们独立推导出TPU的关键机制。
此外,我们希望通过这个项目培养一种独立思考的能力,而不是依赖AI工具来解决每一个小问题。我们希望在未来的项目中,能够运用这种思维方式来解决更复杂的问题。
TPU的核心:矩阵乘法与脉动阵列
TPU的核心任务是执行矩阵乘法,尤其是深度学习模型中占据80-90%计算量的矩阵运算。为了实现这一点,我们使用了脉动阵列(Systolic Array),它由多个处理单元(PE)组成,每个PE在每个时钟周期内执行乘加运算。脉动阵列的设计使得多个输出矩阵元素可以在每个时钟周期内同时计算,从而极大地提高了计算效率。
从XOR问题入手
我们选择从XOR问题(神经网络的“Hello World”)开始,因为它是最简单的神经网络问题之一。XOR问题需要多层感知机(MLP)来解决,因为它需要非线性的决策边界。通过解决XOR问题,我们逐步理解了神经网络背后的数学原理,并将其应用到TPU的设计中。
硬件设计基础:时钟周期与Verilog
在硬件设计中,时间的基本单位是时钟周期,开发者可以根据需求设置时钟周期的长度。我们使用Verilog语言来描述硬件模块的行为,Verilog代码最终会被合成为布尔逻辑门(如与、或、非等),从而构建出我们所需的数字逻辑电路。
脉动阵列与双缓冲机制
脉动阵列的核心是处理单元(PE),它们以网格结构连接在一起,每个PE执行乘加运算。为了进一步提高效率,我们引入了双缓冲机制,它允许我们在计算当前层的同时,加载下一层的权重,从而减少时钟周期的浪费。这种机制使得TPU能够进行连续推理,即在不中断的情况下处理多个输入和权重。
前向传播与反向传播
TPU不仅能够进行推理,还能够进行训练。训练过程的核心是反向传播,它通过计算损失函数的导数来调整模型的权重。我们发现,反向传播的计算过程与前向传播非常相似,都是通过矩阵乘法来实现的。通过脉动阵列,我们可以高效地计算这些梯度,并更新模型的参数。
控制单元与指令集
为了实现推理和训练,我们设计了一个控制单元,并使用自定义指令集(ISA)来控制所有的操作。ISA的长度为94位,确保我们能够高效地执行所有必要的操作。通过将所有的模块整合在一起,我们最终实现了完整的TPU架构,能够同时支持推理和训练。
总结
通过这个项目,我们不仅成功设计和实现了一个简化的TPU,还深入理解了神经网络和硬件设计的核心原理。尽管我们最初对硬件设计一无所知,但通过不断尝试和探索,我们最终完成了这个极具挑战性的任务。这个项目不仅证明了我们的能力,也为我们未来的技术探索奠定了坚实的基础。
评论总结
评论内容总结:
对项目实现方式的疑问
- skybrian 询问项目是硬件实现还是模拟,并希望了解硬件的具体类型和制作过程。
引用: "Did you build real hardware or is it simulated somehow? If it's hardware, what kind and how did you make it?" - UncleOxidant 建议尝试在FPGA上运行项目。
引用: "Have you tried it out in an FPGA?"
- skybrian 询问项目是硬件实现还是模拟,并希望了解硬件的具体类型和制作过程。
对项目工具和技术的建议
- utopcell 提到Google团队使用Chisel而非SystemVerilog,建议考虑切换工具。
引用: "The Google team used Chisel instead of SystemVerilog. You could consider switching to that if it makes sense for your project." - zoobab 建议使用LiteX构建原型。
引用: "Maybe try to build a proto with LiteX?"
- utopcell 提到Google团队使用Chisel而非SystemVerilog,建议考虑切换工具。
对项目成果的赞赏与目标探讨
- jacquesm 赞扬项目在未知领域的探索,认为这种项目最具满足感。
引用: "Sometimes it is the projects where you don't know that you really don't know what you are doing that are the most satisfying." - ganiszulfa 赞赏项目的详细描述和动画,并询问项目的最终目标是否为将TPU应用于消费级或边缘设备。
引用: "Amazing project, and amazing write-up, I especially like the animations. What's the end goal here? Putting these TPUs in the consumer hands or edge devices?"
- jacquesm 赞扬项目在未知领域的探索,认为这种项目最具满足感。