大规模训练是在现实世界中构建深度学习解决方案最具挑战性的方面之一。。从训练到优化，深度学习程序的生命周期需要健壮的基础设施构建块，以便能够并行化和扩展计算工作负载。尽管深度学习框架正在快速发展，但相应的基础架构模型仍处于初期阶段。在过去的几年中，技术巨头Google，Microsoft，Uber，DeepMind和其他公司定期发布了研究成果，以实现跨大型GPU基础架构的深度学习模型的并行化。
分布式和并行化计算的原理与深度学习程序生命周期的几乎所有阶段都息息相关。训练一个深度学习模型是一项非常昂贵的练习，它的执行也是如此。显而易见的答案是，可以利用大型GPU网络来分配深度学习程序的工作负载，但这绝非易事。众所周知，并发和并行编程是非常复杂的，在应用于大型神经网络时更是如此。今天，我想回顾一下在谷歌、DeepMind、Microsoft和Uber中用于并行化大规模深度学习模型训练的一些顶级架构。
谷歌的GPipe
GPipe专注于为深度学习项目扩展训练工作量。从基础架构的角度来看，训练过程的复杂性是深度学习模型中经常被忽视的一个方面。训练数据集越来越大，越来越复杂。例如，在卫生保健领域，遇到需要使用数百万高分辨率图像进行训练的模型并不罕见。结果，训练过程通常要花很长时间才能完成，由于内存和CPU消耗而导致的成本也非常高。
将深度学习模型的并行性分为数据并行性和模型并行性是研究深度学习模型并行性的有效方法。数据并行方法使用大量的机器集群来拆分输入数据。模型并行性尝试将模型移至具有专用硬件以加速模型训练的加速器，例如GPU或TPU。几乎所有的训练数据集都可以按照一定的逻辑并行化，但模型却不是这样。例如，一些深度学习模型是由可以独立训练的并行分支组成的。在这种情况下，典型的策略是将计算划分为分区，并将不同的分区分配给不同的分支。然而，这种策略在顺序堆叠层的深度学习模型中存在不足，这给高效地并行化计算带来了挑战。
GPipe利用一种称为管道的技术，将数据和模型并行性结合起来。从概念上讲，GPipe是一个分布式机器学习库，它使用同步随机梯度下降和管道并行性进行训练，适用于任何由多个连续层组成的DNN。GPipe将一个模型划分到不同的加速器上，并自动将一个小批量的训练实例分割成更小的微批量。该模型允许GPipe的加速器并行运行，以最大限度地提高训练过程的可伸缩性。
下图说明了神经网络的GPipe模型在四个加速器之间的划分。Fk是第k个划分的复合正向计算函数。Bk是对应的反向传播函数。Bk既依赖于上层的Bk+1，也依赖于Fk的中间激活。在模型上部中，我们可以看到网络的顺序特性是如何导致资源利用不足的。图片下部显示了GPipe方法，其中输入的mini-batch被分成更小的macro-batches，这些macro-batches可以由加速器同时处理。
Uber Horovod
Horovod是在社区中变得极为流行的Uber ML堆栈之一，并已被DeepMind或OpenAI等AI巨头的研究团队采用。从概念上讲，Horovod是用于大规模运行分布式深度学习训练工作的框架。
Horovod利用诸如OpenMPI之类的消息传递接口栈来使训练作业能够在高度并行且分布式的基础架构上运行，而无需进行任何修改。通过以下四个简单的步骤即可在Horovod中运行分布式TensorFlow训练工作：
- hvd.init（）初始化Horovod。
- config.gpu_options.visible_device_list = str（hvd.local_rank（））为每个TensorFlow进程分配一个GPU。
- opt = hvd.DistributedOptimizer（opt）使用Horovod优化器包装任何常规的TensorFlow优化器，该优化器使用ring-allreduce来平均梯度。
- hvd.BroadcastGlobalVariablesHook（0）将变量从第一个进程广播到所有其他进程，以确保一致的初始化。
这是基本TensorFlow训练作业的模板：
DeepMind的TF-Replicator
TF-Replicator专注于与TensorFlow程序如何利用Tensor处理单元（TPU）相关的可伸缩性的不同方面。TPU被认为是最先进的AI芯片之一，可为机器学习工作负载提供本地可扩展性。但是，在TensorFlow程序中使用TPU 需要专用的API，这会给不熟悉底层硬件模型的数据科学家带来可移植性问题和采用障碍。DeepMind的TF复制器通过提供一个更简单，友好的开发程序模型来利用TensorFlow程序中的TPU，从而解决了这一难题。
TF-Replicator关注的是可伸缩性的另一个方面，即TensorFlow程序如何利TPU。TPU被认为是最先进的人工智能芯片之一，它为机器学习工作负载提供了本机可伸缩性。然而，在TensorFlow程序中使用TPU需要专门的api，这会导致可移植性问题。DeepMind的TF-Replicator提供了一个更简单、对开发人员友好的编程模型，用于在TensorFlow程序中利用TPU。
TF-Replicator的魔力在于“in-graph replication”模式，其中每个设备的计算都复制到同一TensorFlow图中。设备之间的通信是通过连接设备相应子图中的节点来实现的。为了达到该级别的并行化，TF-Replicator利用TensorFlow的图重写模型在图中的设备之间插入本地通信。当提供一个TensorFlow图时，TF-Replicator首先独立地构建每个设备的计算，并在用户指定跨设备计算的地方留下占位符。一旦构建了所有设备的子图，TF-Replicator就会用实际的跨设备计算代替占位符来连接它们。
从编程模型的角度来看，使用TF-Replicator编写的代码看起来类似于为单个设备编写的本机TensorFlow代码。用户只需要定义（1）公开数据集的输入函数，（2）定义其模型逻辑的阶跃函数（例如，梯度下降的单个阶跃）。以下Python代码片段显示了一个简单的TF-Replicator程序：
为了优化不同设备之间的通信，TF-Replicator利用了最先进的MPI接口。在一些实验中，DeepMind能够在TPUv3 pod的512个核上，以2048的batches size训练著名的BigGAN模型。目前，TF-Replicator是DeepMind的主要TPU编程接口。
微软DeepSpeed
微软的DeepSpeed是一个新的开源框架，专注于优化大型深度学习模型的训练。当前版本包括ZeRO的第一个实现以及其他优化方法。从编程的角度来看，DeepSpeed是建立在PyTorch之上的，它提供了一个简单的API，让工程师们只需几行代码就可以利用并行化技术进行训练。DeepSpeed抽象了大规模训练的所有困难方面，例如并行化、混合精度、梯度累积和检查点，使开发人员可以专注于模型的构建。
DeepScale在四个关键方面都表现出色：
· 规模： DeepSpeed提供系统支持以运行多达1000亿个参数的模型，与其他训练优化框架相比，这代表了10倍的改进。
· 速度：在最初的测试中，DeepSpeed显示出的吞吐量是其他库的4到5倍。
· 成本：使用DeepSpeed训练模型的成本比其他方法低三倍。
· 可用性： DeepSpeed不需要重构PyTorch模型，仅需几行代码即可使用。
将深度学习模型的训练并行化是一项非常复杂的练习，超出了大多数机器学习团队的专业能力。利用谷歌、微软(Microsoft)、优步(Uber)或DeepMind等科技公司创建的框架和架构，肯定可以简化这些工作。
閱讀更多 不靠譜的貓 的文章
