← 全部论文 / 2018

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

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

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism - 中文验证版

英文原文卡片：gpipe_2018.md

状态：已翻译。

元数据

  • Slug： gpipe_2018
  • 年份： 2018
  • 会议： arXiv
  • 作者： Yanping Huang et al.
  • 阅读状态： read complete
  • 计算范式： TPU/GPU 模型并行训练
  • 主要来源： PDF、抽取文本

计算设置

论文明确给出了加速器设置。AmoebaNet 的规模测试在 Cloud TPU v2 上进行，每个加速器 8GB 内存，使用 224×224 输入和 minibatch 大小 128。Transformer 的规模测试在 Cloud TPU v3 上进行，每个加速器核心 16GB 内存，词汇量 32K，序列长度 1024，batch size 32。通信开销实验在一台配备多个 NVIDIA P100 GPU 但没有 NVLink 的单机上运行，强制走较慢的 PCIe 设备—主机—设备传输路径。

内存表指出每个模型参数需要 12 字节，因为训练期间使用了 RMSProp。没有 GPipe 时，单个 8GB 加速器只能容纳一个 82M 参数的 AmoebaNet-D。在单个加速器上使用重计算和 batch 拆分可容纳 318M 参数；使用 8 个加速器上的 8 个流水并行分区可容纳 1.8B 参数。对于 TPU v3 上的 Transformer，直接的单加速器可容纳 282.2M 参数，而 128 个分区可容纳一个 83.9B 参数的 Transformer。

容纳能力不仅取决于参数字节，还取决于激活值。对于 82M AmoebaNet 基线，每个参数 12 字节约 0.98 GB，但峰值激活值增加 6.26 GB，使得在 8GB 加速器上总计约 7.24 GB。使用重计算时，318M 单加速器案例使用约 3.82 GB 的参数状态和 3.46 GB 的峰值激活内存，再次接近 7.28 GB。因此，GPipe 报告的极限是真实的内存受限容纳，而非宽松的参数计数估算。

瓶颈

GPipe 针对两个耦合的瓶颈：设备内存容量和加速器利用率。巨型顺序网络无法放入单个 GPU/TPU，因为权重、优化器状态和激活值超过了内存。直接的逐层模型并行可以将层放置在不同设备上，但顺序网络会使大多数设备在一个层分区运行时处于空闲状态。论文将这种空闲时间称为"气泡开销"。

第二个瓶颈是通信模式。张量拆分的单程序多数据方法可以扩展单个矩阵乘法，但需要频繁的类 all-reduce 通信和高速互联。GPipe 则仅在分区边界进行通信。无 NVLink 的 P100 实验测试了在互联较弱时流水并行分区是否仍然有帮助。

方法适配

GPipe 将网络表示为一个层序列，并将连续的层组划分为流水线"单元"。每个单元放置在一个单独的加速器上，在分区边界处插入通信原语。用户指定分区数 K、微批次数 M 和层序列。一个代价估算器尝试平衡分区计算，使流水线阶段耗时相近。

关键的调度适配是 batch 拆分。一个大小为 N 的 minibatch 被拆分为 M 个相等的微批次，流经 K 个分区。每个微批次的梯度使用与前向传播相同的参数计算。在整个 minibatch 结束时，梯度被累积并应用一次同步更新。这避免了 PipeDream 式的权重过时和多个版本参数副本。

关键的内存适配是重计算。在前向传播期间，每个加速器仅存储分区边界处的激活值。在反向传播期间，它按需重新计算分区内部的函数。这以额外的计算换取大幅缩小的激活值占用，这在加速器内存而非原始 FLOPs 阻碍模型增长时恰好是正确的权衡。

流水线适配由气泡开销参数化。论文推导了在微步数上摊销的气泡时间，并报告当 M >= 4K 时开销可忽略。该规则将 batch size 变成了一个利用率旋钮：足够的微批次使多个设备保持忙碌，即使模型本身是顺序的。

证据

内存证据是核心。在 8GB 加速器上，直接的 AmoebaNet 训练可容纳 82M 参数，参数内存 1.05GB，峰值激活内存 6.26GB。在单个加速器上，GPipe 可容纳 318M 参数，并将峰值激活内存降至 3.46GB。使用 8 个流水并行分区，AmoebaNet 达到 1.8B 参数，约为无 GPipe 基线的 25 倍。扩展并非完全线性，因为 AmoebaNet 的各层是不平衡的。

对于 TPU v3 上的 Transformer 内存扩展，直接的单加速器可容纳一个 3 层、282.2M 参数的 Transformer。使用重计算的单流水线可容纳 13 层、785.8M 参数。128 个流水线可容纳 1663 层、83.9B 参数，提高了 298 倍。最大规模扩展更接近线性，因为各层具有相似的参数数量和输入大小。

吞吐量证据支持微批次规则。在 TPU 上，当 M 远大于 K 时，Transformer 训练几乎线性扩展：在 M=32 时，归一化吞吐量从 K=2 时的 1.8 上升到 K=4 时的 3.4 和 K=8 时的 6.3。当 M=1 时，实际上没有流水并行。在没有 NVLink 的 P100 主机上，将分区数从 2 增加到 8 对 AmoebaNet-D(18,128) 实现了 2.7 倍加速，对 24 层 Transformer 实现了 3.3 倍加速。

最终任务证据表明更大的模型不仅仅是内存容纳演示。一个 557M 参数的 AmoebaNet-B(18,512)，在 ImageNet 2012 上以 480x480 输入和 4 个分区训练，达到单裁剪 84.4% top-1 和 97% top-5 验证准确率。对于多语言 NMT，单个 128 层 6B 参数的 Transformer 在 103 种语言上训练，在 100 个语言对上优于单独训练的 350M 参数双语 Transformer Big 模型。论文还报告了将有效 batch size 从 260K 增加到 1M 再到 4M token 时，德英 BLEU 从 30.92 提升到 31.86 再到 32.71。

历史影响

GPipe 使流水并行和重计算成为超出单个加速器内存的模型的实用方案。它将模型扩展重新表述为放置、调度和激活内存问题。后来的大型模型系统重用了相同的思路：跨设备拆分层、将 batch 拆分为微批次、重叠流水线阶段、重新计算激活值，以及控制梯度更新语义。

论文连接了视觉时代的大型 CNN 和 Transformer 时代的规模：它展示了一个 557M 参数的 AmoebaNet 和一个 6B 多语言 Transformer，而内存表则构建了更大的分区 Transformer。这使流水并行成为一种通用的加速器原语，而非一次性架构技巧。

局限

  • GPipe 假设模型可以表示为层序列，且单层可放入一个加速器。
  • 当微批次数相对于分区数过少时，气泡开销仍然显著。
  • 不平衡的层会降低扩展性，如 AmoebaNet 相对于更均匀的 Transformer 所见。
  • BatchNorm 和其他依赖 batch 的层需要特殊处理，因为统计量是在微批次上计算的。
  • 重计算花费额外的计算来节省内存，因此当内存是限制资源时该权衡最优。

链接