谷歌 TPU v4

TPU v4 是什么

TPU v4 是谷歌为机器学习打造的芯片，以及围绕它搭起来的超级计算机。它从 2020 年起投入使用；谷歌在 ISCA 2023 上发表的论文（作者包括 Norman Jouppi、David Patterson 等人）称它是公司的第五个专用架构（domain-specific architecture），也是第三台面向机器学习模型的超级计算机^[1]。所谓「专用」，指芯片围绕一类计算——神经网络的矩阵运算——来设计，而不是去做通用处理器。

芯片本身

每颗 TPU v4 芯片有两个 TensorCore，每个 TensorCore 里有 4 个矩阵乘法单元，负责神经网络层里最核心的乘加运算^[2]。峰值算力为 275 万亿次浮点运算每秒（teraflops，bf16 或 int8 精度）；bf16 是 16 位浮点格式，int8 是 8 位整数格式，都比完整的 32 位浮点更省^[2]。每颗芯片配 32 GiB 的 HBM2 内存，带宽 1,200 GB/s；实测功耗最小 90 W、平均 170 W、最大 192 W^[2]。它采用 7 纳米工艺，而 TPU v3 用的是 16 纳米^[2]。

从芯片到整机

芯片以 4,096 颗为单位组成一个 pod（整机集群），合计 1.1 exaflops（bf16 或 int8）^[2]。论文给出的数字是：整机规模是 TPU v3 的 4 倍，整体快近 10 倍；单颗芯片比 TPU v3 快 2.1 倍，每瓦性能好 2.7 倍^[2]。

更有意思的是芯片之间的连接方式。光路交换（optical circuit switch）可以动态重连 4×4×4 芯片立方体之间的链路，用户可以自行选择拓扑，比如扭转的三维环面（twisted 3D torus），而不是接受一张固定不变的网络^[1]。谷歌称光交换占系统成本不到 5%、占系统功耗不到 3%^[1]。2023 年 12 月的 Gemini 1.0 报告里，谷歌描述每个 4,096 颗芯片的「超级 pod」都配有一台专用光交换机，能在大约 10 秒内把 4×4×4 立方体重构成三维环面^[1]。

给嵌入准备的 SparseCore

TPU v4 还加了 SparseCore，这是一种数据流处理器，专门加速推荐系统这类嵌入（embedding）密集的模型。谷歌报告说，它让这类模型快 5 到 7 倍，而只占用约 5% 的芯片面积和功耗^[1]。嵌入本质上是一张大表，把标识符映射成向量，瓶颈在访存而不是密集算术，所以单开一块硬件来处理更划算。

与其他芯片的对比

对规模相近的系统，谷歌报告 TPU v4 比英伟达 A100 快 1.2 到 1.7 倍，功耗低 1.3 到 1.9 倍；这是谷歌自己的对比，不是独立评测^[1]。

谁跑在上面

2022 年 4 月的 PaLM 用 6,144 颗 TPU v4 芯片、跨两个 pod、通过 Pathways 系统训练；TPU v4 论文称这是到当时为止描述过的最大 TPU 配置^[3]。2023 年 12 月的 Gemini 1.0 报告说，Gemini 系列模型在 TPUv5e 和 TPUv4 上训练，其中 Gemini Ultra 用了分布在多个数据中心的、数量很大的 TPUv4 加速器集群^[1]。

意义

TPU v4 体现了机器学习硬件的两个方向。一是分工：把密集矩阵运算和稀疏嵌入查询交给不同的硬件块，而不是让一个通用单元全包。二是可重构网络：如果互连能在软件层面改变形状，同一台机器就能服务通信模式不同的任务。Gemini 1.0 报告里对光交换的描述，说明这套设计被沿用到了后来的系统里^[1]。