[翻译]大型企业如何进行分布式机器学习?

  • by

原文链接:https://thesequence.substack.com/p/edge74

大规模培训是在现实世界中构建深度学习解决方案最具挑战性的方面之一。正如一句古老的谚语所说,你最大的优点可以成为你最大的缺点,这当然适用于深度学习模式。整个深度学习空间之所以成为可能,部分原因在于深度神经网络能够跨GPU拓扑进行扩展。然而,同样的扩展能力必然需要计算密集型程序的开发,这对大多数组织来说是操作上的挑战。从训练到优化,深度学习项目的生命周期需要强大的基础设施构建块,以便能够并行化和扩展计算工作负荷。虽然深度学习框架的发展速度很快,但相应的基础设施模型仍然相对较新。在过去的几年里,科技巨头谷歌、微软、Uber、DeepMind和其他公司定期公布各自的努力,以便在大型GPU基础设施上实现深度学习模型的并行化。

分布式计算和并行计算的原理几乎贯穿于深度学习程序生命周期的任何阶段。 训练一个深度学习模型可能是一个非常昂贵的事情,它的执行也是如此。显而易见的方案是利用大型GPU网络来分配深度学习项目的工作量,但这远不是一件容易的事。并发和并行编程是出了名的复杂,当应用于大型神经网络时更是如此。大型科技公司每天都要面对这些挑战,因为它们为任务关键型应用程序操作极其复杂的深层神经网络。本文回顾一下Google、DeepMind、Microsoft和Uber使用的一些顶级体系结构,这些体系结构用于并行化大规模深度学习模型的培训:

  • Google GPipe 
  • Uber Horovod 
  • DeepMind TF-Replicator 
  • Microsoft DeepSpeed 

Google’s GPipe 

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

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

GPipe通过一种称为流水线的技术将数据和模型并行性结合起来。从概念上讲, GPipe是一个分布式机器学习库,它使用同步随机梯度下降和流水线并行进行训练,适用于任何由多个连续层组成的DNN。 GPipe将一个模型划分到不同的加速器上,并自动将一个小批量的训练示例分割成更小的小批量。此模型允许GPipe的加速器并行运行,最大限度地提高培训过程的可伸缩性。

下图说明了GPipe模型,其中一个具有连续层的神经网络被划分为四个加速器。Fk是第k个分区的复合正向计算函数。Bk是相应的反向传播函数。Bk依赖于上层的Bk+1和Fk的中间激活函数。在顶层模型中,我们可以看到网络的连续性如何导致资源利用不足。下图显示了GPipe方法,其中输入的batch被划分为更小的batch,这些batch可以由加速器同时处理。 

Image for post
Image credit: The original paper
Google open-sourced implementation of GPipe as part of the TensorFlow project. 

Uber Horovod 

Horovod 是在社区中非常流行的uberml栈之一,并已被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训练的模板: 

DeepMind TF-Replicator 

TF Replicator关注于TensorFlow程序如何利用张量处理单元(TPU)实现可伸缩性。TPU被认为是最先进的人工智能芯片之一,它为机器学习工作负载提供了本地可伸缩性。然而,在TensorFlow程序中使用TPU 需要专门的API, 这就给不熟悉底层硬件模型的数据科学家带来了可移植性问题和应用障碍。DeepMind的TF Replicator通过提供一个更简单、对开发人员友好的编程模型来解决这一挑战,该模型可以在TensorFlow程序中利用tpu。

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

Image credit: DeepMind

从编程模型的角度来看,使用TF Replicator编写的代码看起来类似于为单个设备编写的本机TensorFlow代码。用户只需定义(1)一个公开数据集的输入函数,和(2)一个定义其模型逻辑的步长函数(例如,梯度下降的单个步长)。以下代码段显示了一个简单的TF Replicator程序:

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

Microsoft DeepSpeed 

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

从功能角度来看,DeepScale在四个关键方面表现出色:

  • 规模: DeepSpeed为运行高达1000亿个参数的模型提供了系统支持,与其他训练优化框架相比,这意味着提高了10倍
  • 速度: 在最初的测试中,DeepSpeed显示出比其他库高4-5倍的吞吐量。
  • 成本: 使用DeepSpeed训练模型的成本比其他方法低三倍。
  • 可用性: DeepSpeed不需要重构PyTorch模型,只需几行代码即可使用。
Image for post
Image credit: Microsoft Research Blog

结论 

并行化 深度学习模型的训练是一项非常复杂的工作,超出了大多数机器学习团队的专业知识范围。利用Google、Microsoft、Uber或DeepMind等科技公司创建的框架和架构,当然可以简化这些工作。在不久的将来,我们希望看到这些框架的一些版本包含在主流深度学习技术栈中,以使核心深度学习社区更加平民化。


Leave a Reply

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.