2020-04-08 15:08:38

编者按：Google、Uber、DeepMind和Microsoft这四大科技公司是当前将深度学习研究广泛应用于自身业务的典型代表，跻身全球深度学习研究水平最高的科技公司之列。GPipe、Horovod、TF Replicator和DeepSpeed分别是这四家公司开发应用的深度学习框架，它们结合了深度学习研究及其基础设施的前沿技术，以提高深度学习模型的训练效率。这四个深度学习框架各有什么特点和优势呢？本文将对此做一个基本介绍。

AI 四巨头 Google、DeepMind、Microsoft、Uber 深度学习框架大比拼_深度学习

作者 | Jesus Rodriguez

译者 | 苏本如，责编 | 夕颜

封图 | CSDN下载自东方IC

出品 | CSDN（ID：CSDNnews）

以下为译文：

大规模深度学习训练是在现实世界中构建深度学习解决方案最具挑战性的方面之一。 正如俗语所说，你最大的优点可以成为你最大的缺点，这一点自然也适用于深度学习模型。整个深度学习空间在一定程度上取决于深度神经网络（DNN）跨越GPU拓扑的能力。然而，同样的扩展能力导致了计算密集型程序的创建，这给大多数组织带来了操作上的挑战。从训练到优化，深度学习项目的生命周期需要健壮的基础设施构建块，以便能够并行化和扩展计算工作负载。

尽管深度学习框架正在快速发展，但相应的基础架构模型仍处于初期阶段。在过去的几年里，技术巨头谷歌、微软、优步（Uber）、DeepMind和其他公司定期发布了各自的研究成果，以便在大型GPU等基础设施上实现深度学习模型的并行化。

分布式和并行化计算的原理与深度学习程序生命周期的几乎所有阶段都息息相关。训练一个深度学习模型可能是一个非常昂贵的任务，运行也是如此。显而易见的答案是，可以利用大型GPU网络来分配深度学习程序的工作负载，但这绝非易事。众所周知，并发和并行编程是非常复杂的，尤其在应用于大型神经网络时更是如此。大型科技公司每天都在面临这些挑战，因为它们必须为关键业务应用运行极其复杂的深层神经网络。

今天，我想回顾一下谷歌、微软、DeepMind、和优步（Uber）等科技巨头用于并行化大规模深度学习模型训练的一些顶级框架。具体来说，我们将讨论以下四个框架：

  • 谷歌的 GPipe

  • 优步的Horovod

  • DeepMind的TF-Replicator

  • 微软的 DeepSpeed

 

谷歌的GPipe

GPipe专注于为深度学习项目提高训练工作量。从基础设施的角度来看，训练过程的复杂性是深度学习模型中经常被忽视的一个方面。训练数据集越来越大，也越来越复杂。例如，在医疗保健领域，遇到需要使用数百万高分辨率图像进行训练的模型已不罕见。因此，训练过程通常需要很长的时间才能完成，并且由于内存和CPU消耗而导致的成本也非常高。

将深度学习模型的并行性分为数据并行性和模型并行性，是一种研究深度学习模型并行性的有效方法。数据并行方法使用大量的机器集群来拆分输入数据。模型并行性尝试将模型移至加速器上，如GPU或TPU，它们具有加速模型训练的特殊硬件。在较高的层次上，几乎所有的训练数据集都可以按照一定的逻辑进行并行化，但在模型上却不是这样。例如，一些深度学习模型是由可以独立训练的并行分支组成的。在这种情况下，经典的策略是将计算划分为多个分区，并将不同的分区分配给不同的分支。然而，这种策略在连续堆叠层的深度学习模型中存在不足，这给高效并行计算带来了挑战。

GPipe采用一种称为管道的技术，将数据和模型并行性结合起来。从概念上讲，GPipe是一个分布式机器学习库，它使用同步随机梯度下降和管道并行性进行训练，适用于任何由多个连续层组成的深度神经网络（DNN）。GPipe将一个模型划分到不同的加速器上，并自动将一个小批量的训练实例分割成更小的微批量。该模型允许GPipe的加速器并行运行，以最大限度地提高了训练过程的可伸缩性。

下图说明了具有多个连续层的神经网络的GPipe模型在四个加速器之间的划分。Fk是第k个划分的复合正向计算函数。Bk是对应的反向传播函数。Bk既依赖于上层的Bk+1，也依赖于Fk的中间激活。在图片的上部中，我们可以看到网络的顺序特性如何导致资源利用不足的。图片的下部显示了GPipe方法，其中输入的小批量（mini-batch）被划分为更小的宏批量（macro-batch），这些宏批量（macro-batch）可以由加速器同时处理。

AI 四巨头 Google、DeepMind、Microsoft、Uber 深度学习框架大比拼_深度学习_02

谷歌开源了GPipe的一种实现，作为TensorFlow项目的一部分。

开源地址：https://github.com/tensorflow/lingvo/blob/master/lingvo/core/gpipe.py

 

优步的 Horovod

Horovod是Uber的机器学习（ML）堆栈之一，它已经在社区中非常流行，并且已经被DeepMind和OpenAI等人工智能巨头的研究团队采用。从概念上讲，Horovod是一个用于大规模运行分布式深度学习训练工作的框架。

Horovod利用诸如OpenMPI之类的消息传递接口栈，来使训练作业能够在高度并行和分布式的基础设施上运行，而无需进行任何修改。通过以下四个简单的步骤即可在Horovod中运行分布式TensorFlow训练工作：

  1. hvd.init初始化Horovod。

  2. config.gpu_options.visible_device_list = str(hvd.local_rank)为每个TensorFlow进程分配一个GPU。

  3. opt=hvd.DistributedOptimizer(opt)使用Horovod优化器包装任何常规的TensorFlow优化器，该优化器使用ring-allreduce来处理平均梯度。

  4. hvd.BroadcastGlobalVariablesHook(0)将变量从第一个进程广播到所有其他进程，以确保一致的初始化。

下面这个代码示例是一个基本的TensorFlow训练作业的模板，你可以从中看到上面的四个步骤：

 1import tensorflow as tf
2import horovod.tensorflow as hvd# Initialize Horovod
3hvd.init# Pin GPU to be used to process local rank (one GPU per process)
4config = tf.ConfigProto
5config.gpu_options.visible_device_list = str(hvd.local_rank)# Build model…
6loss = …
7opt = tf.train.AdagradOptimizer(0.01)# Add Horovod Distributed Optimizer
8opt = hvd.DistributedOptimizer(opt)# Add hook to broadcast variables from rank 0 to all other processes during
9# initialization.
10hooks = [hvd.BroadcastGlobalVariablesHook(0)]# Make training operation
11train_op = opt.minimize(loss)# The MonitoredTrainingSession takes care of session initialization,
12# restoring from a checkpoint, saving to a checkpoint, and closing when done
13# or an error occurs.
14with tf.train.MonitoredTrainingSession(checkpoint_dir=“/tmp/train_logs”,
15 config=config,
16 hooks=hooks) as mon_sess:
17 while not mon_sess.should_stop:
18 # Perform synchronous training.
19 mon_sess.run(train_op)

 

DeepMind的TF-Replicator

TF-Replicator专注于TensorFlow程序如何利用Tensor处理单元（TPU）有关的可伸缩性的另一个方面。TPU被认为是最先进的人工智能芯片之一，它为机器学习工作负载提供了本机可扩展性。然而，在TensorFlow程序中使用TPU需要专门的API，这会给不熟悉底层硬件模型的数据科学家们带来可移植性问题和采用障碍。DeepMind的TF Replicator通过提供一个更简单、对开发人员更友好的编程模型来利用TensorFlow程序中的TPU，从而解决了这一难题。

TF-Replicator的魔力依赖于一种“图内复制（in-graph replication）”模型，其中每个设备的计算被复制到同一张TensorFlow图中。设备之间的通信是通过连接设备对应子图中的节点来实现的。为了达到这种级别的并行化，TF-Replicator利用TensorFlow的图重写模型在图中的设备之间插入本机通信。当呈现一个TensorFlow图时，TF Replicator首先独立地为每个设备构建计算，并在用户指定跨设备计算的地方留下占位符。一旦构建了所有设备的子图，TF Replicator就会用实际的跨设备计算替换占位符来连接它们。

AI 四巨头 Google、DeepMind、Microsoft、Uber 深度学习框架大比拼_深度学习_03

从编程模型的角度来看，使用TF-Replicator编写的代码看起来类似于为单个设备编写的本机TensorFlow代码。用户只需定义：（1）一个公开数据集的输入函数，和（2）一个定义其模型逻辑的阶跃函数（例如，梯度下降的单个阶跃）。下面的代码片段展示了一个简单的TF-Replicator程序：

 1# Deploying a model with TpuReplicator.
2repl = tf_replicator.TpuReplicator(
3num_workers=1, num_tpu_cores_per_worker=8
4)
5with repl.context:
6model = resnet_model
7base_optimizer = tf.train.AdamOptimizer
8optimizer = repl.wrap_optimizer(base_optimizer)# ... code to define replica input_fn and step_fn.per_replica_loss = repl.run(step_fn, input_fn)
9train_op = tf.reduce_mean(per_replica_loss)with tf.train.MonitoredSession as session:
10repl.init(session)
11for i in xrange(num_train_steps):
12session.run(train_op)
13repl.shutdown(session)

为了优化不同设备之间的通信，TF-Replicator利用了最先进的MPI接口。在一些实验中，DeepMind能够在一个TPUv3 pod的512个核心上，以2048的batch size批量训练著名的BigGAN模型。目前，TF-Replicator是DeepMind公司的TPU的主要编程接口。

 

微软的DeepSpeed

微软的DeepSpeed是一个新的开源框架，专注于优化大型深度学习模型的训练。当前版本包含了ZeRO的第一个实现以及其他优化方法。从编程的角度来看，DeepSpeed是在PyTorch之上构建的，它提供了一个简单的API，允许工程师只需几行代码就可以利用并行化技术来进行训练。DeepSpeed抽象了大规模训练的所有困难方面，例如并行化、混合精度、梯度累积和检查点，使得开发人员可以专注于模型的构建。

从功能角度来看，DeepSpeed在以下四个关键方面表现出色：

规模：DeepSpeed可以为运行高达1,000亿个参数的模型提供系统支持，这和其他训练优化框架相比，提高了10倍。

速度：在最初的测试中，DeepSpeed表现出的吞吐量比其他库高出4-5倍。

成本：可以使用DeepSpeed进行训练的模型，其成本比其他替代方案少三倍。

可用性：DeepSpeed不需要重构PyTorch模型，仅需几行代码即可使用。

AI 四巨头 Google、DeepMind、Microsoft、Uber 深度学习框架大比拼_深度学习_04

将深度学习模型的训练并行化是一项非常复杂的工作，超出了大多数机器学习团队的专业知识。利用谷歌、微软、优步或DeepMind等科技公司创建的框架和架构，肯定可以简化这些工作。

在不久的将来，我们希望看到这些框架的一些版本包含在主流深度学习堆栈中，以加速核心深度学习社区准入的民主化。