Google TPU v6e（Trillium）

是什么

Google TPU v6e 又称 Trillium，是 Google Cloud 提供的一种张量处理单元（tensor processing unit，TPU）。TPU 是专门用来执行机器学习计算的加速芯片，尤其擅长神经网络训练和推理中大量出现的矩阵运算与向量运算。Google Cloud 为 v6e 公布了规格表，列出每颗芯片的算力、内存和互连参数，因此既可以看到单颗芯片的配置，也可以看到由多颗芯片组成的 Pod 的规模。^[1]

由来与可用性

Google Cloud 表示，Trillium 于 2024 年早些时候发布，并在该文发布时已向 Google Cloud 客户普遍可用。^[3] Google 还表示，它使用 Trillium TPU 训练了新的 Gemini 2.0，并且单个 Jupiter 网络结构可以承载超过 100,000 颗 Trillium 芯片。^[3] 这两点把这款芯片放在 Google 自己的模型训练工作和数据中心网络设计中：Gemini 2.0 是在这款芯片上训练的模型，Jupiter 则是连接大量芯片的网络结构。

算力与内存

规格表列出每颗 TPU v6e 芯片的峰值 BF16 算力为 918 teraflops，峰值 FP8 算力也是 918 teraflops。^[1] BF16 和 FP8 分别是 16 位和 8 位浮点格式；位宽越低，数据占用的内存和传输量通常越小，硬件在单位时间内能完成的操作也可能更多，具体选择取决于模型和训练或推理阶段。规格表还列出每颗芯片有 32 GiB 的 HBM（high-bandwidth memory，高带宽内存）容量和 1,638 GB/s 的 HBM 带宽。^[1] HBM 是靠近加速芯片的高带宽内存，容量决定模型状态和中间结果能有多少留在片上，带宽决定这些数据能被多快地搬进搬出。

每颗 v6e 芯片按该表配置一个张量核心（TensorCore）和两个稀疏核心（SparseCore）。^[1] TensorCore 是执行矩阵乘加运算的单元；SparseCore 是面向稀疏计算的专用单元。Google 的架构文档说，v6e 的矩阵乘法单元为 256×256，并支持 8 位整数格式。^[2] 这描述了核心矩阵引擎的尺寸和数据类型支持；矩阵乘法是全连接层、注意力机制以及许多其他神经网络组件背后的基本运算。

互连与规模

规格表列出每颗芯片 800 GB/s 的双向芯片间互连带宽，以及 100 Gbps 的数据中心网络带宽。^[1] 前者描述 Pod 内芯片之间交换数据的速度，后者描述芯片连接到更广的数据中心网络的速度。该表列出每个 TPU v6e Pod 有 256 颗芯片。^[1] Google 的文档说，一个 v6e 切片（slice）可以按每台虚拟机 1 颗、4 颗或 8 颗 TPU 芯片来配置。^[2] 因此，Pod 是多颗芯片组成的大粒度物理分组，切片则是客户把芯片挂到虚拟机上时看到的分配单位。

在发展中的位置

v6e 是 Google TPU 产品线中通过 Google Cloud 提供的型号，公布了每颗芯片的参数，并采用 256 颗芯片组成的 Pod 设计。^[1] 它的配置体现了加速芯片设计的两个层面：一层是每颗芯片的算力和内存，负责本地计算；另一层是芯片间带宽和数据中心网络，负责跨多颗芯片的任务。Google 关于 Trillium 训练 Gemini 2.0 的说法，为这款芯片提供了一个具体的训练负载。^[3] 同一说法中，Jupiter 能承载超过 100,000 颗 Trillium 芯片，说明该设计可以扩展到单个 Pod 之外。^[3] 对使用者来说，实际看到的是每台虚拟机 1 颗、4 颗或 8 颗芯片的切片，位于 256 颗芯片的 Pod 内，并通过能连接到更大规模芯片数量的网络结构互连。^[2][1][3]