近年来，机器学习数据集和模型的规模不断增加，从而可以在各种任务中获得更好的结果。与此同时，硬件加速（GPU，TPU）也在不断改进，但速度明显放慢。模型增长和硬件改进之间的差距增加了并行性的重要性，即在多个硬件设备上训练单个机器学习模型。一些ML体系结构，尤其是小型模型，有助于并行性，并且可以在硬件设备之间轻松划分，但在大型模型中，同步成本会导致性能下降，从而无法使用它们。

来自Google Brain，GPipe的新论文（https://arxiv.org/pdf/1811.06965.pdf）提出了一种模型并行性的新技术，该技术允许在多个硬件设备上训练大型模型，性能几乎提高1：1（论文显示在4x硬件上的处理能力是3.5x）。GPipe库将是开源的，它自动分析TensorFlow神经网络模型的结构，并将训练数据和模型委托给多个硬件设备，同时应用独特的反向传播优化技术。

GPipe有助于ML模型包含更多参数，从而在训练中获得更好的结果。为了证明该技术的有效性，Google团队创建了更大版本的AmoebaNet，2017 ImageNet获胜者，使用更大的图像（480×480）作为输入，并在ImageNet、CIFAR-10、CIFAR-100和其他计算机视觉指标上实现了最先进的(SOTA)结果。

背景

机器学习中的并行性通常分为两类：

• 模型并行性——在训练中使用模型并行性时，机器学习模型被划分为K个硬件设备，每个设备都持有模型的一部分。一种简单的并行建模方法是通过在每个设备上托管N/K层(“stages”)将N层神经网络划分为K个设备。更复杂的方法是通过分析每一层的计算复杂度来保证每个设备处理相似的计算复杂度。标准模型并行性允许训练更大的神经网络，但由于设备之间一直在等待，而且在给定的时间内只有一个设备可以执行计算，因此在性能上受到了很大的冲击。
• 数据并行 - 在数据并行中，机器学习模型在K个硬件设备上复制，一小批训练样本分为K个微批。每个设备执行微批次的forward和backward 传递，当它完成该过程时，它将更新的模型权重与其他设备同步，然后计算整个小批量的更新权重。在每个小批量计算结束时，K模型的权重都是同步的（相同的值）。

虽然微批次允许并行，但当使用模型并行时，每个阶段自然仍然必须等待前一阶段的结果，导致空闲的“Bubble”，如图所示。在图中，F0,i-F3,i是跨越四个阶段的单个微批次的成员，并且Fi,0-Fi,3是在单个阶段中执行的四个微批量计算。

模型+数据并行（ GPipe：Huang et al.2018）

尽管存在Bubble问题，数据并行性仍然有效，但是存在额外的问题 - 通信开销。随着模型的增长，在设备之间同步整个模型的要求成为训练过程中的瓶颈，大大减慢了它的速度。附图说明了在大型神经网络中，通信开销如何构成大部分训练时间。通信开销现象鼓励创建非常大的小批量，但这些通常是训练网络的错误选择，并且可能在生产中呈现较差的结果。

PipeDream：快速高效的管道并行DNN训练

GPipe的工作原理

GPipe使用模型和数据并行，这种组合通常称为“pipelining”。它为以前的流水线技术提供了两个关键贡献 - 自动并行和设备内存优化。

自动并行

该软件接收神经网络的体系结构，mini-batch size以及可用于计算的硬件设备的数量作为输入。然后，它会自动将网络层分为几个阶段，将小批量分成微批，将它们分布在各个设备上。为了将模型划分为K个阶段，GPipe根据其激活函数和训练数据的内容估算每个层的成本。虽然本文没有详细说明如何在GPipe中完成，但常用的技术是通过神经网络运行数据样本，测量每个层的计算时间，并相应地进行划分。GPipe为每一层提供了一个可选的成本估算函数作为输入，允许更复杂的技术改进其内部机制。

设备内存优化

当计算神经网络中的backward pass时，需要网络的forward pass 激活来执行计算。通常这意味着在神经网络中具有micro-batch of size N 和L layers，在 forward pass之后O（N×L）激活被保持在设备存储器中，以准备backward pass。

GPipe使用不同的方法，应用计算内存权衡 - 不是将NxL激活保留在内存中，它只将N个激活保留在阶段的最后一层（stage = group of layers）。在这种情况下，每次backward pass开始（从最后一层），forward pass激活被重新计算并保存在内存中。当单个样本的backward pass结束时，存储器中的激活被丢弃并重新计算以用于下一个样本的backward pass。使用这种方法，设备存储器一次只保留一组激活，从而获得宝贵的存储器。由于一般硬件趋势是设备速度增长快于设备存储器，因此这种权衡通常是有用的。

GPipe应用程序

在该论文中，研究人员将AmoebaNet从1.553亿个参数扩展到5.57亿个参数，并输入480×480 ImageNet图像，而不是标准AmoebaNet模型使用的下采样331×331图像。结果是ImageNet Top-1 Accuracy（84.3％vs 83.5％）和Top-5 Accuracy（97.0％vs 96.5％）的改进，这标志着一种新的先进技术。

使用扩展的AmoebaNet模型进行迁移学习，研究人员在CIFAR-10，CIFAR-100，Oxford-IIIT Pets，Stanford Cars和Food-101数据集上获得了最新的（SOTA）结果。该模型的结果不如两个数据集中的最新结果 - FGVC Aircraft和Birdsnap，这可以用这些数据集中最先进的模型来解释，除了ImageNet数据，这些数据集还利用了来自谷歌图像搜索的980万张经过预处理的图像。

实施细节

GPipe是用TensorFlow编写的，并且是开源的。该技术不是TensorFlow独有的，也可以在其他平台上实现。

结论

作为一个开源库，GPipe将允许机器学习从业者以相对较低的成本训练更大的模型。可以肯定的是，结果将会是一种新的大型机器学习模型，它将获得优于现有模型的结果，同时也会增加对完整尺寸的图像数据的使用。

这种数据并行性的突破虽然对每个人都有用，但对于拥有大量计算和数据的大型组织来说，自然会提供一种特殊的优势。