Google TPU的发展历程与思考（二）

TPU v2 与 TPU v3

相较于 TPU v1 只能用于推理，TPU v2 致力于解决训练难题。

TPUv2 设计目标

训练与推理——仅仅是转变方向而已吗？

TPUv2 誓要解决更难的训练任务。事实上，训练与推理的难度相差比想象的要大。
1、训练并行化更难。推理阶段，每个推理都是独立的，因此 DSA 芯片集群可以横向扩展。而训练需要迭代运行数百万个示例，协调跨并行资源（因为训练到最后，必须生成一组一致的权重作为答案！）
2、计算更难。反向传播需要计算模型每一阶段变量的导数，包括更高精度的激活函数（甚至是超函数）和转置权重矩阵 $W^T$ 的乘法。
3、训练需要更多内存。权重更新数据来自前向和反向传播的临时中间值，当然中间值需要被保留，因而大大提高了存储要求；这些临时存储需求可能是权重存储空间的 10 倍。
4、它必须更具可编程性。训练算法和模型不断变化，因此在设计期间仅限于当前最佳实践算法的 DSA 可能会很快过时。
5、短整数可以用于推理，但要充分捕捉训练期间许多小权重更新的总和通常需要浮点运算。

TPUv2 架构

从推理到训练

论文提出，几十年来构建通用 CPU 的许多原则和经验在 TPU 中会变得不再适用。例如，以下是推理TPU v1 和 TPU v2 共享的特性，但在 CPU 中并不常见：

• 1-2 个大核 vs 32-64 个小核（PS：神威超算上用的申威CPU就是众核结构，有64个核）
• 由 2D 的脉动阵列组成的大量计算资源 vs 只有一些标量或SIMD计算单元
• 使用更窄的位宽来提高效率 vs 更长的位宽
• 没有缓存、分支跳转预测器等架构

总结完共性，再来看它们之间的差异，论文中有句话我特别喜欢：只要将 TPUv1 的架构设计稍作更改，就可以得到 TPUv2，而对此所做的更改，就是训练与推理的差别。TPUv1 中使用多个 SRAM（Weight Mem, Unified Buffer, Activation, Accumulator）来缓存数据，并且Weight RAM是只读的，但训练显然需要一个统一的 SRAM，以提高训练效率。因此第一个改动是，合并 Activation Storage 和 Accumulator，变成 Vector Mem（见下图）；用一个可编程的 Vector Unit 代替固定功能的 Activation Pipeline（这里面包含了池化与激活函数）；对 DNN 来说，因为 BF16 比 FP16 更好，因此 MXU 选择优先支持 BF16；由于需要更新 Weight，Weight SRAM 必须是可写的，且还另需一大块面积缓存上一层计算得到的中间结果；DRAM 是 Vector Memory 的次级内存，它们组成了一个编译器控制的内存层次结构；而使用 HBM
DRAM，带宽比 DDR3 增加了 20 倍，从而保证了计算资源的利用率；TPUv2 的 322位长的 VLIW 指令来自 local mem，相比 TPUv1，已经不再是 CPU 了。

架构总结

这里简单地列举并介绍 TPUv2 的架构情况：

• 高带宽内存 High Bandwidth Memory, HBM。因为 TPU v1 中的大多数神经网络都受到了内存的限制，故使用高带宽：它使用一个中间衬底，通过 32 条 128 位总线将 TPUv2 芯片连接到 4 个短 DRAM 芯片堆栈，从而提供了 20 倍于 TPUv1 的带宽
• 核间互联 Interconnect Router, ICI。为了实现 2D 环面连接以组成超算，芯片有四个自定义的核间互连(ICI)链路，每个链路运行在 TPUv2 中，每方向 496Gbits/s
• 核发射单元 Core Sequencer 。Core Sequencer 不再从 CPU 获取指令，而是从 Instruction Mem 取出 VLIW 指令，使用 4K 32-bit 的标量内存来执行标量运算，有 32 个 32 位的标量寄存器，而将向量指令送到 VPU，322 位宽的 VLIW 可以发送 8 个操作，2 个标量，2 个向量 ALU，向量 load，向量store，一对从矩阵乘来的排队数据
• 向量运算单元 Vector Processing Unit（VPU）进行向量计算，有 32K 128x32 位的16MiB 大小的片上内存Vmem，32 个 2D 的向量寄存器 Vregs，每个包含128x8 32位元素（4KiB）
• 为应对 batch normalization，增加了一个 SIMD 维度给vector unit
• 矩阵运算单元 MXU，输入 16 位的 FP 得到 32 位 FP 结果，一个芯片有多个 MXU，每个大小为 128 x 128
• 矩阵转置单元 Transpose Reduction Permute Unit，用于做 128x128 矩阵的转置

TPU v2 的设计遵循了 TPU v1 中的设计原则：使用大型二维矩阵乘法单元（Matrix Multiply Unit）脉动阵列来减少面积和能量，以及一个软件控制的大容量片上存储器，而不是 cache。

但不同的是，TPU v2 芯片上有两个核，而单个核内 MXU 相比 TPU v1 小了四分之一，这是因为有些卷积规模不到 256x256，单用一个大核，会导致 MXU 的利用率不是很高，但要是使用 16 个 64x64 MXUs 的就需要更多的面积，因此选择了两个 128x128。这里有一个很简单的判断依据，MXU 得到的输入和输出结果所需的带宽与其周长成正比，但它提供的计算量却与面积成正比，与周长的平方成正比。因此越大的 MXU 其受限于带宽瓶颈的问题越明显。另一个理由是，单个大核下，脉动阵列过大会产生较大延迟，且芯片上核数过多也会增加编译程序的复杂度。

TPUv3

TPU v3 实际上就是 TPU v2 的增强版。TPU v3 相比 TPU v2 有约 1.35 倍的时钟频率、ICI带宽和内存带宽，加上两倍的 MXU 数量，因此峰值性能提高 2.7 倍。而采用水冷方式，能使功率提升 1.6 倍。Google 将1024 个 TPU v3 组成超级计算机，在同等技术条件下，TPUv3 的模具尺寸只比 TPU v2 大 6%，尽管每个 MXU 的数量是 TPU v2 的两倍，其中的原因很简单，因为工程师在布局布线 TPU v2 时发现了一种更好的布局方法，因此在布局 TPU v3 时获得了更好的平面布局设计。

下表列出了这三代芯片与同时代 GPU 相比的数据，

TPU组成的超级计算机

模型训练常需要大规模数据才能完成，TPUv2 的一个重要改进是增加了核间互连结构(ICI)，使最多 256 个 TPUv2 组成一个超级计算机。Google 基于 TPU v2 开发的超级计算机，会使用多个 TPUv2 组成集群对深度学习网络进行训练（见下图），超级计算机中，大多数流量都是来自所有节点的权重更新时的 all-reduce 流程。
通过交换机提供的虚拟电路和无死锁路由功能，再加上 TPUv2 本身存在的核间互连(ICI)结构，便可以构建出上图这样的 TPUv2 集群：TPU v2 集群的二维网孔结构，提供了 15.9T/s 的平分带宽，相比传统的集群策略，省去了 Infiniband 网卡、Infiniband 交换机的成本，以及与集群 CPU 的通信延迟。

而有意思的是，使用专用的硬件并在各个节点上小批量训练数据，几乎没有尾部延迟，同步并行训练成为可能。最近的研究显示，同步训练在同等资源的条件下可以击败异步 SGD 训练，是因为同步训练不需要参数服务器，直接允许节点间的通信，使用 all-reduce 模式就可以获得一致的权重，解决了异步训练时参数服务器与各工作节点带宽不足的问题。

与GPU的比较

TPU 架构与 GPU 相比，有哪些不同或相同点？

• 多节点并行方面。TPU 通过 ICI 连接并行，并通过 TPU 上层的编译器 XLA，和 TensorFlow 支持 all-reduce 操作。多 GPU 则通过NVIDIA NVLink 和主机控制的 InfiniBand 网络或交换机连接在一起。
• 运算数据类型。TPU 提供了 BF16 数据类型，Volta GPU 则使用 FP16 数据类型。
• 架构。TPUv2 采用了双核结构，顺序执行的架构，计算部分由脉动阵列(systolic array)完成，编译器会控制 scratchad mem 优化访存行为，并尽量让计算与访存行为交叠。GPU 有非常多小的计算核，且对延迟容忍度较高，配合有大量线程、CUDA 和寄存器文件等，完成大量计算。

参考文献

[1] Jouppi, Norman P., et al. “Ten lessons from three generations shaped google’s tpuv4i: Industrial product.” 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA). IEEE, 2021.
[2] Jouppi, Norman P., et al. “In-datacenter performance analysis of a tensor processing unit.” Proceedings of the 44th annual international symposium on computer architecture. 2017.
[3] Jouppi, Norman P., et al. “A domain-specific supercomputer for training deep neural networks.” Communications of the ACM 63.7 (2020): 67-78.