AI 术语表 · 基础 · 效率与基础设施
TPU
TPU（Tensor Processing Unit，张量处理单元）是 Google 为加速深度学习计算（尤其是矩阵乘法）而自主设计的 AI 专用芯片。
工作原理
深度学习计算的大部分工作是大矩阵之间的乘法。TPU 把处理矩阵乘法的电路放在芯片的核心位置并做得很大，使其在相同功耗下能完成更多矩阵运算。大型模型以“Pod”为单位训练，即用专用互连网络把多个 TPU 芯片连在一起。程序由名为 XLA 的编译器转换成适合 TPU 的形式后运行。
在哪里会看到
TPU 不是单独购买的产品，主要是在 Google Cloud 上租用。在 Google Cloud 的价格表和实例列表、Google Colab 的运行时选择菜单中可以看到它。代码方面，JAX 与它最为契合，TensorFlow 和 PyTorch（通过 XLA 对接）也可以使用。研究论文描述实验环境时，有时也会出现“在 TPU 上训练”的说法。
为什么重要
它是 NVIDIA GPU 之外，可用于大规模 AI 训练和服务的代表性选择。据悉，Google 使用 TPU 训练和部署自家模型。它也是科技巨头自研 AI 芯片这一趋势的代表案例。在 GPU 供应紧张时，它会被当作替代方案来考虑。
注意事项
用 CUDA 编写的代码或 CUDA 专用内核无法直接在 TPU 上运行。开源推理引擎和最新技术往往先推出 GPU 版本，因此在 TPU 上需要另外确认是否支持。比较性能时，与其比较单颗芯片，不如看相同成本下能处理多少相同的工作，这样更贴近实际。使用场景实际上被绑定在 Google Cloud 上，这一点也需要考虑。
GPU 与 TPU
GPU(NVIDIA)通用并行芯片
- CUDA 生态广泛
- 可在多种云上使用或直接购买
- 新技术通常率先支持
TPU（Google）AI 专用芯片
- 专为矩阵乘法优化
- 主要在 Google Cloud 上使用
- 与 JAX、XLA 契合度高
공통 两者都用于大规模模型训练和推理
示例
一家初创公司在选择模型训练基础设施时，如果一直用 JAX 编写代码，就会把 Google Cloud 的 TPU 列入候选。相反，依赖 PyTorch 和 CUDA 专用推理引擎的服务，可能会权衡代码修改的负担，选择继续使用 GPU 云。