语言理解成新一代“太空竞赛”，谷歌用4096颗TPU-v3集群探索最优并行方法

图片出处：unsplash.com

深度学习革命正处于一场语言理解引领的 “太空竞赛” 之中。

梳理其中的发展进程，早期的代表性突破之一是谷歌的神经网络翻译系统 (GNMT, Google’s Neural Machine Translation System)，该系统使用了 LSTMs 和 Attention 机制，以实现显著的质量改进。在这之后出现了 Transformer。这是一种用于 seq2seq 任务的模型，采用避免循环 (recurrent) 的模型结构，完全依赖于注意力机制（self-attention）对输入输出的全局依赖关系进行建模，实现并行化处理，提高训练速度，在机器翻译等任务中取得了一定提升。

而大热的 BERT 就是从 Transformer 这个基本组件中衍生出来的预训练语言模型，它首先对一般语言知识进行 “预训练”，然后针对特定的语言任务进行微调。截止到目前，作为 OpenAI 开发的第三代语言模型，GPT-3 称得上有史以来训练复杂度最高的语言模型，是 AI 领域的重大突破。它共有 1,750 亿个参数，可见其规模的庞大，能直接用于新任务以及简短的演示，无需针对特定数据进行进一步的微调。

毋庸置疑，训练这些愈发大型的模型，所需要的计算设备也需要不断进化。

按照时间顺序来看，2012 年出现的突破性模型 AlexNet 使用两个 GPU 对模型执行并行训练。同年，谷歌利用其数据中心规模的 CPU 集群在 DistBelief 中开展异步训练。在这之后，深度学习革命引发了业界对 GPU 的 “追捧大潮”：2016 年年中到 2018 年年中，英伟达的收入每季度平均同比增长 50%。

2015 年，谷歌发展出了专门的神经网络加速器 —— 张量处理单元 (TPU,Tensor Processing Unit)，这是一款专用于机器学习算法的专用芯片，在性能功耗比、峰值性能和推理延迟等方面提高了 10 倍以上。

短短的两年内，谷歌的第二代 TPU 使用了 256 芯片的 pods 来训练一个具有近乎完美的并行扩展性的单一模型；第三代 TPU 增加到 1024 pods。其他厂商也在着手部署类似规模的集群，例如微软和 OpenAI 构建了一个 10,000 个 GPU 的大型集群。

很明显，当前语言理解领域的竞争就像一场现代太空竞赛，各组织竞相 “组装” 强劲算力硬件和超大规模模型，以寻求通用人工智能的突破。

那么，以公认占据领先地位的谷歌 TPU 为例，如何再突破并发计算的能力限制？

近日，一支团队便展示了他们在这个问题上所做的探索 —— 他们制作了 4096 个 TPU-v3 芯片构成的 TPU 集群 Multipods，并详细地介绍了如何对训练进行扩改进，以在六个 MLPerf 测试任务中实现最佳效果，其中 4 项测试任务的训练时间更是创纪录地被缩减到 16~23 秒。

研究发现，模型间的最佳并行化途径是存在差异的，比如说，MLPerf SSD、MaskRCNN 和 Transformer 模型并行化实现了最大规模的训练，而 e BERT 和 Resnet-50 则依赖数据并行化有更好表现。在技术层面，编译器对模型并行化的作用、在 multipod 网格上的快速梯度求和、分布式评估、输入 pipeline 优化和特定模型优化的组合有助于达到最强效果。

团队还演示了 TensorFlow 和 JAX 编程框架的性能。其中，JAX 的多客户端方法进一步降低了启动和编译成本。

此次论文。图片出处：Arxiv

Multipod

MLPerf 是业内广受欢迎和支持的首套衡量 ML 软硬件性能的通用基准，由图灵奖得主 David Patterson 联合企业和高校于 2018 年发起，为不同类别的 ML 加速器和框架在最新的 ML 任务上进行基准测试。

在最近一次的 MLPerf v0.7 训练榜单中，有来自英伟达、谷歌、阿里巴巴、富士通、英特尔的测试提交。

与之前的系统基准测试一样，MLPerf 基准测试套件也在推动性能的发展，随着谷歌 TPU-v3 和 TPU-v4 系统的 MLPerf 训练 v0.7 出炉，足以展示了 AI “跑分” 的实力，我们有能力实现大规模测试。

MLPerf v0.7 是第三个 AI 训练跑分榜单，其规则中还增加了新的模型：BERT，大型预训练语言模型；DLRM，深度学习推荐系统；更大版本、经强化的 MiniGo。MLPerf 训练基准包括在特定数据集上按照预定义的收敛测试度训练模型，同时遵循特定的参数、优化和时序方法。

图 1

在此次研究中，团队成员组装了一个带有 4096 个芯片的 TPU-v3 “multipod”，如图 1 所示，以探讨其在 MLPerf 模型中并发性的限制。需要注意的是，此规模为 MLPerf v0.6 训练基准测试时使用的 TPU-v3 的四倍大，且 multipod 采用了 128x32 的二维网格拓扑结构 (在 Y 边缘有 pod 内环形链接)。

图 2

从图 2 可以清晰看到，在此设置中，两个 pod 通过跨 pod 的链路沿网格的 x 维方向进行连接。这些 links 比标准的 TPU-v3pod 内的 links 更长。

团队最终选择用 MLPerf 模型的一个子集，在 multipod 上进行基准测试，其中包括 BERT、 ResNet-50、 Transformer 和 SSD。

在 BERT 和 ResNet-50 中，团队使用 batch parallelism 进行扩展，而在 Transformer 和 SSD 中，结合使用模型并行化和 batch parallelism。

而在 Mask-RCNN 和 DLRM 模型，前者可用的 batch parallelism 非常有限，后者的可扩展性受限于全局的 batch size，通信成本很快就盖过了横向扩展的优势。

测试表现

作者在论文中对上述所讨论的所有框架、基础设施和模型优化能够实现的性能给出了阐释。在此，主要挑选几个重点简述一下：

表 1

首先，表 1 给出了在谷歌 TPU multipods 上进行的 6 个 MLPerf v-0.7 基准测试的端到端时间，以及谷歌 TPU 提交到 MLPerf v0.6 基准的加速过程。对 DLRM 而言，1.21 分钟的最佳结果是在 TPU-v4 机器上实现的，但 TPUv4 不作为本文研究的重点，所以给出 TPU-v3 的结果是 2.4 分钟。

表 2

其次，表 2 比较了 TensorFlow 和 JAX 两个框架在包含有 4096 和 2048 TPU 芯片的大型系统上的初始化时间。鉴于两者采用了不同的分布式编程模型，其中，JAX 使用多客户端分布式系统，每个客户端都有一个代码副本编译自己的图，因此对于不同大小的系统会生成恒定的初始化时间；而 TensorFlow 为 Python 进程提供了全局可见性和对系统的整体控制，如果系统超级大，随之时间花费会更大。

这就刚好解释了两个框架在初始化时间上的差异性：TensorFlow 从 498 秒到 1040 秒不等，而 JAX 的初始化时间显然要低得多，从 122 秒到 294 秒不等，因此也就说明了在性能和可用性之间作出权衡的必要性。

图 5

上图 5 则是随着 TPU 芯片数量的增加，ResNet-50 的端到端和吞吐率的加速情况，从图中的曲线走势发现，吞吐量提速比端到端提速更接近理想的扩展性也就不足为奇。

图 6

为了进一步地研究阿姆达尔定律（Amdahl's Law）在大型系统中的瓶颈，图 6 分解了 step time 内的计算和 allreduce（通信）算法成本，其中蓝色区域显示计算时间，红色区域表示通信时间，两者之和代表设备上的步长时间。可以清晰看到，随着系统规模的扩大，计算时间不断减少，而通信时间 (即 allreduce 时间) 几乎保持不变。

图 7

图 7 显示了改变 TPU 芯片数量的 BERT 加速情况。BERT 展现了从包含 16 到包含 4096 个芯片的系统上的最高可扩展性。

最后，作者又从模型并行化在基准测试的加速情况、针对模型 / 系统额外优化和可扩展性的关系、不同芯片数量对不同系统的测试时间的影响以及 TPU 和 GPU 的可扩展性比较等等。综合分析得出结论：本文提出的技术使 TPU 实现更低的端到端时间和更高的加速，在 BERT、ResNet-50、Transformer、MarkRCNN、DLRM 和 SSD 上的性能都优于其他。

优化方法

接下来，我们再简单了解下在 4096 芯片 TPU 的 multipod 上扩展 TensorFlow 和 JAX 框架可用的优化技术。

模型并行化：在数据并行性受限的模型中，采用并行化技术在 multipod 上实现更高的并发性。借助 XLA 的 SPMD 分区器基于轻量级注释自动划分模型图。在 SSD 和 MaskRCNN 分割模型中，通过标注输入图像来实现空间分割。SPMD 分区器可以自动地沿空间维度并行化计算。这些模型一般具有较大的空间维度 (MaskRCNN 为 8000x1333, SSD 为 300x300)。两个模型都支持沿着 8 个核进行空间分区，以实现最高的并发级别。在诸如 MLPerf transformer 基准测试这样的语言模型中，可以采用 Shazeer 等人所描述的划分特征维度的方法，以实现后续任务中的并行化计算。

weight update 切分：在传统的数据并行中，优化器在每个训练步骤结束时复制和更新模型权值。然而，当每个核的 mini batch size 很小时，这种计算可能会变得很重要。以 MLPerf BERT 模型举例，在 512 个 TPU-v3 芯片上，LAMB 优化器的权重更新时间大约是 step time 的 18%。weight update 更新切分技术首先执行一个全局的 reduce-scatter，然后每个加速器都有一个由累加梯度组成的切分来分配计算量。在接下来的一步中，将对已更新权重的片进行全局更新所有。为了实现更好的加速，在数据和模型并行性方面启用了权重更新切分。在权值被复制的分割模型中，weight update 切分方案类似于数据并行。

优化全局求和：梯度求和步骤对于 MLPerf 基准实现强扩展至关重要。为了优化大型 TPU-v3 multipod 上的梯度求和，利用沿 Y 维的环面缠绕连接。双向环用于执行沿 Y 维的 reduce-scatter 操作，其输出是沿 Y 环的梯度求和的分片。然后，沿着 X 维方向执行 reduce-scatter 操作。接下来进行权重更新计算，以梯度分片作为输入。更新后的权重首先沿 X 传播，然后沿 Y 方向传播，分两步进行。在 MLPerf transformer 基准测试中，通过将模型权重分片到高达 4 个相邻的 TPU 核上，执行分布式矩阵乘法操作。这些核心沿着 x 维上的一条线放置。在 ML 训练的前向传输过程中，所有的 reduce 调用都沿着 x 近邻的短环执行；与数据并行性一样，Y 维上的梯度和保持不变。

评估指标的分布式计算：训练和评估计算是在 TPU 加速器上紧密循环执行的。训练循环的结果更新了每个 TPU 加速器上 HBM 存储器中的模型权重。然后，使用更新后的权重来评估 MLPerf 规则中指定的 epoch 数的输出指标。在评测中，当评估批处理的大小大于评估数据集中的样本数时，评估数据集中填充虚拟样本。在 TensorFlow 实现中，评估输出张量用于计算 TPU 主控主机上的评估指标 (例如 Resnet-50 基准测试中的 top-1 精度)。然而，在 JAX 实现中，评估质量指标的计算是通过全局求和调用实现完全分配的。

输入 pipeline 优化：扩展 ResNet-50 模型的挑战之一是主机输入 pipeline 中的负载不平衡。由于 multipod 的规模很大，一些主机输入 pipeline 在解压缩大的 JPEG 图像时成本很大。对于此问题，团队的解决方案是将未压缩的图像存储在主机内存中，以便主机输入 pipeline 仅执行以下操作：i）随机裁剪，ii）随机翻转和 iii）图像归一化，按照 MLPerf 参考中指定的固定均值和方差。这显示大大增加了主机输入 pipeline 的吞吐量，从而允许它创建一个大型预取缓冲区。因此，当主机 Pipeline 对大的输入图像进行预处理时，它可以向 TPUs 提供预取缓冲区中的图像，从而消除 multipod 系统上的输入 pipeline 负载不平衡的情况。

根据团队在论文中的描述，上述技术是通用的，可以应用于其他模型、框架和 ML 加速器架构。

Refrence：

[1]https://arxiv.org/pdf/2011.03641.pdf