TPU及其他ASIC简介

尽管CPU提供强大的灵活性，GPU为多种任务提供大规模并行处理能力，但还有一类处理器，专为单一任务实现最高效率而打造：即专用集成电路（ASIC）。对于某些大型AI计算需求，这些专用芯片能够提供通用硬件无法比拟的性能和能效。

什么是ASIC？

专用集成电路（ASIC）是为特定目的而设计制造的芯片。不同于CPU被设计来运行完整的操作系统和无数种软件，ASIC的逻辑门是物理上布局好的，用于执行非常有限的操作集。可以将其理解为可编程的厨房电器（能混合、切碎、搅拌）与一台简单高效的咖啡研磨机之间的区别。研磨机只做一件事，但比通用电器更快、能耗更低。

在机器学习的背景下，ASIC旨在加速神经网络的数学基本运算，主要是矩阵乘法和卷积。通过去除所有不必要的组件，例如CPU复杂的条件分支预测或GPU的图形渲染管线，AI ASIC可以将其所有的芯片资源和功耗预算用于这些核心计算。

谷歌张量处理器：一个典型范例

AI ASIC中最知名的例子是谷歌的张量处理器（TPU）。最初在内部开发，用于加速谷歌搜索和照片等服务的推理。现在，TPU通过谷歌云平台向公众提供。它们从零开始设计，以极高的速度和效率执行TensorFlow和JAX等机器学习框架中定义的操作。

TPU架构的核心创新是脉动阵列。

脉动阵列架构

脉动阵列是由简单、相同的处理单元（PE）组成的网格，这些单元连接到最近的邻居。数据以有节奏的波浪状模式流经此网格，类似于血液通过循环系统泵送的方式，这也是“脉动”一词的由来。

矩阵乘法$C = A \cdot B$的工作方式如下：

1. 载入权重： 一个矩阵的值（例如模型权重）被预先载入到处理单元阵列中。
2. 数据流传输： 另一个矩阵的值（例如输入激活值）从一侧流式传输到阵列中。
3. 计算与传递： 当输入数据流过时，每个PE执行乘累加（MAC）操作。它将传入的激活值与其存储的权重相乘，并将结果加到从其邻居传递过来的值上。然后它将自己的输入数据和累加结果传递给序列中的下一个PE。

这种设计非常高效，因为数据不断移动并被计算。权重矩阵的值被多次重复使用，无需反复从内存中获取，这显著减少了内存带宽瓶颈和功耗。

脉动阵列示意图。激活值从顶部流入，而权重从侧面载入。每个处理单元（PE）执行计算并将数据传递给其邻居，从而实现高度并行和高效的计算流程。

权衡：性能与灵活性

ASIC的专用性带来了明确的权衡。对于其设计目的的任务，它们的性能和能效可以比高端GPU好一个数量级。这通常以每秒万亿次操作（TOPS）衡量，更重要的是，以每瓦特TOPS衡量。

针对典型矩阵密集型AI计算需求，不同处理器类型的归一化每瓦性能。TPU的专用特性使其能够以明显更低的功耗执行目标操作。

然而，这种性能是以牺牲灵活性为代价的。TPU不是通用处理器。它不能运行任意Python代码或渲染用户界面。它针对特定的操作集和数据类型（如bfloat16或INT8）进行了优化。如果您的模型使用新的、不受支持的操作，则必须修改模型或等待硬件和软件支持加入后才能在TPU上运行。

这形成了硬件选择的多样性：

• CPU： 灵活性最高，并行性能最低。
• GPU： 灵活性与高并行性能之间有良好平衡。
• ASIC： 灵活性最低，针对特定任务的并行性能最高。

除了TPU：一个不断发展的生态体系

尽管谷歌的TPU是知名度最高的AI ASIC，但它绝非唯一。对高效AI计算的需求促使了其他专用芯片的出现：

• AWS Trainium和Inferentia： 亚马逊的定制ASIC，专门为AWS生态体系内的经济高效模型训练（Trainium）和推理（Inferentia）而设计。
• 专业初创公司： Cerebras Systems、SambaNova Systems和Groq等公司已开发出新颖的架构，以应对大规模AI计算，通常在内存和数据流方面采用独特的方法。

这些多样化解决方案的存在表明，对于达到一定规模的组织，设计或使用定制芯片是优化性能和成本的可行策略。对于大多数工程师和数据科学家而言，眼前的选择将是在通过云服务提供商使用这些ASIC，或者坚持使用更传统的基于GPU的基础设施之间。您的决定将取决于您计算需求的规模、预算，以及您的模型与ASIC专用能力的契合程度。