Google TPU v5e

是什么

Google TPU v5e 是 Google 张量处理单元（tensor processing unit，TPU）系列中的一款机器学习加速芯片。TPU 是 Google 自己设计的芯片，专门做神经网络里占主导地位的大规模矩阵运算。维基百科把 v5e 列为 2023 年推出的产品^[2]；Google 对它的定位同时覆盖训练和推理（serving）^[1]，也就是既用来训练模型，也用来把模型跑给用户。

一颗芯片里有什么

每颗 v5e 芯片以一个 TensorCore 为核心。这个 TensorCore 里有 4 个矩阵乘法单元，负责神经网络层里最密集的线性代数运算；此外还有 1 个向量单元，处理逐元素以及其他非矩阵的运算，以及 1 个标量单元，处理周边的控制和地址计算^[1]。这样分工的意义在于让矩阵单元尽可能一直忙着做乘加，不像矩阵乘法的杂活交给另外两个单元。

算力与内存

按 Google 给出的数字，单颗芯片的 bfloat16 峰值算力为 197 TFLOP，INT8 峰值算力为 393 TOPs^[1]。bfloat16 是一种 16 位浮点格式，指数范围与 32 位浮点数相同，因此在训练中很常用；INT8 是 8 位整数格式，多用于推理。INT8 的数字大约是 bfloat16 的两倍，这在同一套乘法硬件用更低精度一次做更多次运算时是常见现象。

每颗芯片配 16 GB 的高带宽内存（high bandwidth memory，HBM），Google 文档给出的带宽是 800 GiBps^[1]；维基百科列出的容量同样是 16 GB，带宽写作 819 GB/s^[2]。容量决定一个模型、或一个批次能不能塞进一颗芯片，带宽决定权重和激活值能以多快的速度喂给运算单元。

主机与芯片间互联

芯片不是单独工作的。一台主机配有 512 GiB 的 DRAM，最多搭载 8 颗 v5e 芯片^[1]。同一台主机上的芯片之间、跨主机的芯片之间，都通过芯片间互联（inter-chip interconnect，ICI）通信：每颗芯片有 4 个 ICI 端口，双向芯片间带宽为 400 GBps^[1]。由于一层模型往往要切分到很多颗芯片上，这条链路承载的是分布式训练中的梯度和激活值，它经常决定一个任务能扩展到什么程度。

Pod 的规模

Google 把 v5e 芯片组成 256 颗的 Pod，用二维环面（two-dimensional torus）互联：每颗芯片与相邻芯片相连，网格边缘绕回，因此每颗芯片的链路数相同^[1]。一个完整的 Pod 的 bfloat16 峰值算力为 50.63 PFLOPs^[1]。Pod 内的全归约（all-reduce）带宽为 51.2 TB/s，二分带宽（bisection bandwidth）为 1.6 TB/s^[1]。全归约是让每颗芯片最终都拿到所有芯片梯度之和的集合通信操作，也就是数据并行训练里的同步步骤；二分带宽衡量的是同一时刻能有多少流量穿过网络中部，当很多芯片需要同时交换数据时，这个数字很关键。

训练与推理两种用途

Google 把 v5e 同时定位在训练和推理上，训练任务以吞吐为优化目标，推理任务以延迟为优化目标^[1]。区别很实际：训练任务希望单位时间内塞进尽可能多的样本，在意的是每秒完成的总工作量，所以大批次、高利用率是好事；推理任务一次回答一个请求，在意的是单个回答多快返回，因此小批次、短响应时间更重要。同一款芯片要能兼顾这两种用法，容量、带宽和互联这三组参数才会被放在一起，作为判断它适合什么负载的依据。