← 全部论文 / 2018

Mesh-TensorFlow: Deep Learning for Supercomputers

2018 TPU、加速器与 Transformer 时代
下载 PDF

Mesh-TensorFlow: Deep Learning for Supercomputers - 中文验证版

英文原文卡片：mesh_tensorflow_2018.md

状态：已翻译。

元数据

  • Slug： mesh_tensorflow_2018
  • 年份： 2018
  • 会议： NeurIPS
  • 作者： Noam Shazeer et al.
  • 阅读状态： read complete
  • 计算范式： TPU、加速器与 Transformer 时代
  • 主要来源： PDF、抽取文本

计算设置

论文的目标明确是最多 512 个核心的 TPU 网格。它描述了一个 512 核心的 TPU 集群，配备 16x16x2 环形网络互联，并指出同一物理集群可以表示为逻辑网格，如 16x16x2、32x16 或 512。实验使用最多 16x32（即 512 个核心）的二维 TPUv2 网格。

报告的最大语言模型实验在 512 核心 TPUv2 集群上训练一个 4.9B 参数的 Transformer，耗时 13 小时。最大的翻译实验在 128 核心 TPUv2 集群上训练一个 2.9B 参数的模型，耗时 22 小时。论文还报告了最大模型超过 50% 的计算效率：在最大 11.5 PFLOP/s 中达到 6 PFLOP/s。

瓶颈

纯数据并行会遇到内存约束和同步代价。每个处理器持有参数的完整副本，而归约 batch 维度的操作需要对参数大小的张量进行 Allreduce。对于非常大的 Transformer，参数和激活值可能超过每核心内存，而同步复制参数所需的时间可能使纯 batch 拆分效率低下。

替代方案——模型并行——是必要的，但难以手工编程。Transformer 的维度（如词汇量、前馈隐藏层大小和注意力头）可以跨设备拆分，但正确的选择取决于张量形状、batch size、内存和物理网络拓扑。论文强调逻辑网格仅是一个命名抽象：物理拓扑仍然影响性能，尤其是 MPI Allreduce 等集合通信。

方法适配

Mesh TensorFlow 通过以下方式将 TensorFlow 式模型适配到加速器网格：

  • 给张量维度命名。
  • 给处理器网格维度命名。
  • 将逻辑张量维度映射到网格维度。
  • 使用 Allreduce 等集合操作编译单程序多数据程序。
  • 跨网格轴拆分 Transformer 维度，如词汇量、前馈宽度和注意力头。

这使得张量布局成为模型定义中的显式部分。关键的抽象是从命名的张量维度到命名的网格维度的映射。数据并行是一种布局，将 batch 映射到处理器网格。模型并行是另一种布局，将隐藏层或词汇量维度映射到网格。混合并行使用二维网格，例如将 batch 映射到行，同时将词汇量、前馈宽度和注意力头映射到列。

这种映射是一种计算设备适配，因为它让程序员能够表达内存和通信权衡，而无需为每个核心手工编写不同的程序。如果一个 einsum 归约出一个跨网格拆分的维度，编译器会插入集合操作。如果相邻张量的维度以相同方式拆分，则可以避免数据移动。如果一个张量维度未被映射，则它被复制。论文的示例展示了通信量如何根据 batch、隐藏层、输入/输出或收缩维度是否被拆分而变化。

对于 Transformer，论文同时扩展模型维度和集群维度。Batch 沿一个 TPU 网格维度拆分，而词汇量、前馈隐藏层维度和注意力头沿另一个维度拆分。如果 batch size 随数据并行轴扩展且拆分的模型维度随模型并行轴扩展，则此布局能保持性能。因此，该方法是一种在增长模型宽度和处理器数量的同时保持每核心内存和每核心算术量合理的方式。

证据

扩展证据是具体的。使用二维 TPUv2 布局，作者在最多 16x32 核心的网格上训练了前馈隐藏维度最高达 262144 和最高达 256 个注意力头的 Transformer 模型。他们报告在最大模型上超过 50% 的效率，具体为 11.5 PFLOP/s 中的 6 PFLOP/s。

对于 One Billion Word 语言模型基准，最大模型有 4.9B 参数，在 512 个 TPUv2 核心上训练 10 个 epoch 耗时 13 小时。所有模型的 batch size 为 256 个序列、每个序列 256 个 token。Batch 沿大小为 16 的网格维度拆分，模型维度沿大小为 32 的网格维度拆分。最大模型报告每词开发困惑度 24.0，当 logits 乘以 0.9 进行评估时降至 23.5。

对于 WMT14 英法翻译，最大模型有 2.9B 参数，在 128 个 TPUv2 核心上训练 3 个 epoch 耗时 22 小时。它使用 sacrebleu 达到 BLEU 43.9，论文称这是当时发表的最佳结果。表 3 显示模型大小随前馈宽度和注意力头增长：d_ff 2048 和 4 个头的 0.15B 参数达到 En-Fr BLEU 41.8，而 d_ff 65536 和 128 个头的 2.89B 参数达到 43.9。

系统证据是代价表和示例。论文按计算、通信和内存代价比较了布局，并解释了如果每处理器 batch 过小，性能会受损；如果隐藏层切分得过细，性能也会受损。Mesh TensorFlow 的贡献在于暴露这些布局决策，而非将它们隐藏在固定的数据并行训练背后。

历史影响

Mesh TensorFlow 使布局成为一等扩展抽象。它是从单机 Transformer 通向 TPU-Pod 语言模型训练的系统桥梁：模型定义、张量维度、处理器网格和集合通信成为一个设计面。后来的 TPU 规模工作（如 T5、GShard 和 Switch Transformer）继承了模型架构和网格布局必须协同设计的思想。

从历史上看，这是远离"仅增大 batch size"式扩展的重要一步。它表明当布局控制好内存和通信时，通过增加模型宽度和注意力头可以改善语言模型质量。它还将责任从手工编写的模型并行代码转移到了编译器生成的单程序多数据程序。

局限

  • 它假设处理器相同且可靠。
  • 张量维度必须能被网格维度整除。
  • 物理拓扑仍然影响集合通信。
  • 自动化布局搜索被留作未来工作。
  • 糟糕的布局可能复制计算或通信过多；该抽象暴露了布局，但不会自动解决每个布局选择。
  • 结果与 TPUv2 网格集合操作和 TensorFlow/XLA 时代的基础设施相关，因此确切的效率声明不应推广到每种加速器拓扑。

链接