ML@Scale
首发于ML@Scale
科普:分布式深度学习系统(二)

科普:分布式深度学习系统(二)

上一篇文章帮助大家梳理了一下分布式深度学习系统面临的问题以及相关进展。这一篇文章接着上一篇,详细介绍一下当把深度学习从CPU移植到GPU上后会碰到的几个明显问题,以及一些解决方法。


深度学习的计算模型

很多机器学习模型,包括绝大部分的神经网络(Neural Networks)、图模型(Graphical Models)、矩阵分解(MF)等等,它们的训练算法都可以被抽象成一个迭代收敛(iterative-convergent)的过程。虽然针对各个模型,开发者一般需要具体地写各种不同的计算函数,但是大体上,这些模型的代码逻辑一般都符合下面这个流程:

\begin{equation} \large \mbox{for } t = 1 \rightarrow T : \theta^{(t+1)} = \theta^{(t)} + \epsilon \Delta_{\mathcal{L}}(\theta^{(t)}, D^{(t)}) \end{equation}

这里 {\large \theta} 是这个模型的参数, {\large \Delta} 是一个更新函数: {\large \Delta} 函数每次拿当前状态下的模型参数 {\large \theta^{(t)}} 和训练数据 {\large D^{(t)}} 当做输入,然后计算并返回一组更新值(updates),然后把这组值直接加到当前参数 {\large \theta^{(t)}} 上,从而将参数更新到新的状态 {\large \theta^{(t+1)}} ;如此不断循环,直到达到提前给定的最大迭代次数 {\large T} ,或者满足某些循环终止条件(比如,参数更新值几乎趋近于0了)。举个具体的例子,比如在用随机梯度下降(SGD)训练神经网络时, {\large \theta} 就是这个神经网络要训练的参数,而 {\large \mathcal{L}} 就是你要优化的目标函数,{\large \Delta_{\mathcal{L}}} 等价于求梯度 {\large \nabla_{\mathcal{L}}} ,其对应的操作就是:在每次迭代中,取一个小batch的训练数据 {\large D^{(t)}} ,计算在这组数据上目标函数 {\large \mathcal{L}} 对于参数 {\large \theta} 的梯度;再对原始的梯度附加一系列操作,比如梯度下降取负,乘上learning rate(公式中的 {\large \epsilon} ),加上momentum,必要时再对梯度进行clip操作等等。随后,我们把更新值加上参数上,即做梯度下降,并得到一组更新后的参数,并检查当下神经网络是否已经收敛。

有了这个固定的计算流程,我们可以很容易的通过数据并行(见本专栏第一篇文章),把任何符合这种迭代收敛计算流程的模型训练程序,并行化到多个计算节点上进行加速。假设我们有 {\large P} 个计算节点,按照数据并行的思路,我们把数据均分成 {\large P} 份并且分发到这 {\large P} 个节点,然后让每个节点在被分到那份数据上单独执行 {\large \Delta_{\mathcal{L}}} 计算函数并计算一份更新值。然后我们再想办法把每个节点上算出来的更新值统一收集起来再更新模型参数,即可。

按照这个思路,我们可以很容易地把上面的计算流程重写成一个分布式版本:

\begin{equation} \large \mbox{for } t = 1 \rightarrow T : \theta^{(t+1)} = \theta^{(t)} + \epsilon \sum_{p=1}^P \Delta_{\mathcal{L}}(\theta^{(t)}, D^{(t)}_p) \end{equation}

和前面的单机版本比较,我们在这里给数据 {\large D^{(t)}} 加上了一个下标 {\large p} ,也就是说,这里我们有 {\large P} 个计算节点在不同的数据上同时执行计算函数 {\large \Delta_{\mathcal{L}}} 。当所有节点的计算都执行完毕后,它们的结果被汇总之后( \sum 符号)再用来更新参数 {\large \theta}

一旦我们把这种迭代收敛的计算模式抽象出来,按照这个构造,我可以很容易地把一个单机的机器学习程序改写成一个数据并行的版本,而不用考虑具体的模型细节。

但是,任何构造其实都隐含一些假设,我们在应用这些构造前需要适当反思一下当前应用场景是否遵从这些假设。那么,上面这一套用来并行迭代收敛算法的构造有什么假设呢?大体上,我们可以总结出下面这三点:

  • 我们希望上述训练程序的主要计算任务集中在 {\large \Delta} 函数上,因为真正被并行到多台机器上执行正是 {\large \Delta} 函数 --- 只有当主要的计算开销由多个机器共同承担时,分布式计算才能发挥其优势;
  • 在每个iteration内,每个节点上的 {\large \Delta} 运算应是相互独立的:第一台机器上 {\large \Delta} 函数的执行不依赖于第二台机器的结果,数据点跟数据点之间没有依赖性(dependency);
  • 在上面的构造中,模型参数 {\large \theta} 在分布式环境下是被所有节点共享的;也就是说,在每个iteration开始之前,每个节点都必须能够获得到当前(最新的)模型参数 {\large \theta^{(t)}} 。在 {\large \Delta} 函数执行完毕之后,每个节点上算出来的参数更新值 \Delta_{\mathcal{L}}(\theta^{(t)}, D^{(t)}_p) 可以很方便的被收集起来并用来更新模型参数( {\large \theta^{(t)} \rightarrow \theta^{(t+1)}})。

回到分布式深度学习上,我们来逐个看看这三个假设是否都满足。

显然,第一个和第二个假设都很容易满足。对应到神经网络训练上, {\large \Delta} 函数即对应后向传播(backpropogation)求梯度的过程,显然,BP求梯度所涉及到的计算基本上就是神经网络训练的主要计算负载,所以第一个假设满足。同时,在绝大部分情况下,我们都会假设训练数据是i.i.d.的,用随机梯度下降训练时,每个计算节点只需要从它所负责的那块数据上取一个独立小batch算一个随机梯度(stochastic gradients)就可以了,这个计算不受限于其他节点上的计算,也没有数据之间的依赖,所以第二个假设也满足;再看第三个假设,在每个iteration开始之前,每个计算节点是否能够“轻易”获得当前的模型参数 {\large \theta^{(t)}}?在每个iteration的计算结束之后,我们能否同样“轻易”地把所有节点上产生的梯度都收集起来并更新当前参数?显然,满足这个假设需要一些额外的条件。

在单机上进行训练时,我们根本不需要考虑这两个步骤,因为参数就放在(GPU)内存上,我们只需要从内存上读或往内存上写就可以了。但是,当在多机集群上进行训练时,不管是在计算开始前读取模型参数,还是在计算结束后收集多个节点上的梯度,都会涉及到网络通讯,如何保证参数共享和梯度同步即为满足第三个条件要解决的核心问题。


参数服务器(Parameter Server)

说了这么多,其实就是为了引出参数服务器(Parameter Server)。Parameter server(PS) 就是为上面这种迭代收敛的计算模型而设计出的一套通讯接口和方法。顾名思义,PS的架构其实和课本上讲的client-server(CS)架构差不多。PS主要抽象出了两个主要概念:服务器(server, or master)和计算节点(client, or worker);服务器里面放了一些数据,而计算节点则可以向服务器发数据或者请求服务器回传数据。有了这两个概念,我们可以把分布式机器学习的计算流程和PS的服务器和计算节点这两个模块做如下的映射:PS的服务器端维护全局共享的模型参数 {\large \theta^{(t)}} (我们通常称为parameter state),而客户端则对应到执行计算任务 {\large \Delta} 的各个工作节点;同时,服务器端向客户端提供两个主要的API: push和pull。那么,PS架构如何满足解决上述的参数共享和梯度同步问题呢?在每个iteration开始的时候,所有的客户端先调用pull API向服务器发送一个请求,请求服务器回传最新的模型参数 {\large \theta^{(t)}} 。当每个计算节点收到 {\large \theta^{(t)}} 后,它就把这份最新的参数拷贝并覆盖到之前旧的参数( {\large \theta^{(t-1)}} )上(物理上通常这些参数存储在RAM或GPU内存上),然后执行{\large \Delta} 函数计算得到梯度更新值。换句话说,PS的pull API确保了每个计算节点在计算开始前都能获取一份最新参数的拷贝。

另一方面,梯度更新值计算完毕后,每个计算节点随后调用push API,把这组更新值发给服务器。服务器会收集所有计算节点发来的更新值,并且将它们“加”到当前维护的全局共享参数上( {\large \theta^{(t)}} \rightarrow {\large \theta^{(t+1)}})。在下一个iteration,当服务器再次收到客户端的pull请求时,它就会把更新过后的 {\large \theta^{(t+1)}} 发出去。因此,push API确保了梯度值的收集和模型参数的更新。

Parameter Server的架构示意图

这里我要重点说明一点:PS里面的“服务器”和“客户端”其实都是抽象的概念。虽然听起来像是一个中心化的架构,但是在真正的实现中,并不一定会有一台专门的机器被用来当做中心服务器;相反地,目前大部分的PS都把服务器端实现成一个分布式存储系统,以避免负载不均衡,并减少单机的通信瓶颈,这个我们后面再讨论。

将上述过程和单机上的神经网络训练过程作对比,我们发现一个主要的不同点:PS允许用梯度更新模型参数这一步发生在远端,而非当前计算节点上。作为一个工具,对于PS的用户来说,PS的优势主要体现在它的简单易用上:如果一个用户想把训练程序改成多机版,只需要在每个训练iteration开始前调一下PS的pull API,在计算结束后调用一下PS的push API,就可以让这个程序在多机上跑起来并且保证参数同步了,是不是很简单?其实这就是包括DistBelief等大部分CPU上的分布式机器学习系统的架构。但是,从设计者的角度来讲,想要实现这么一个系统并提供这两个有效的API,需要考虑的问题就复杂的多了。回到我第一篇专栏里面介绍关于分布式系统的几个问题,我们来一一讨论:

  • 架构:PS里面的服务器端具体应该怎么设计?物理上,这个服务器端应该放在哪?是一个单独的高性能服务器,还是一个专用的服务器集群,还是和计算节点直接共享硬件(前面讨论过)?计算节点怎么和服务器连接?
  • 存储:服务器端如何维护模型参数,用什么数据结构?怎么保证快速的索引和读写?客户端如何请求服务器端?客户端是否需要设计cache或者buffer来加速模型参数和梯度的获取?
  • API实现:push和pull API具体应该怎么实现,阻塞还是非阻塞?是否需要多线程?
  • 同步:服务器端收集梯度,以及客户端请求发回参数怎么保持同步性?完全同步是必须的吗?完全异步是否可行?同步和异步各有什么优点和缺点?
  • 网络通信:push和pull设计到参数和梯度通信。那么,如何保证通讯效率?如何有效利用网络带宽?通信量远大于可用带宽的时候应该怎么办?
  • 容错:如果服务器端宕机了,或者有某几个计算节点crash了,整个分布式系统应该怎么保证容错?
  • 如果计算节点是GPU,或者多GPU节点,问题又会怎么样?

上面我列举了很多问题,其中的某些问题在我下面列举的这几篇文章里面都有陆续讨论。如果你对这些问题比较感兴趣,我推荐按顺序读一下下面这几篇文章:

More Effective Distributed ML via a Stale Synchronous Parallel Parameter Server. Qirong Ho, James Cipar, Henggang Cui, Seunghak Lee, Jin Kyu Kim, Phillip B. Gibbons, Garth A. Gibson, Greg Ganger, Eric P. Xing. NIPS 2013.

Exploiting Bounded Staleness to Speed Up Big Data Analytics. Henggang Cui, James Cipar, Qirong Ho, Jin Kyu Kim, Seunghak Lee, Abhimanu Kumar, Jinliang Wei, Wei Dai, Gregory R. Ganger, Phillip B. Gibbons, Garth A. Gibson, and Eric P. Xing. ATC 2014.

Exploiting Iterative-ness for Parallel ML Computations. Henggang Cui, Alexey Tumanov, Jinliang Wei, Lianghong Xu, Wei Dai, Jesse Haber-Kucharsky, Qirong Ho, Gregory R. Ganger, Phillip B. Gibbons, Garth A. Gibson, and Eric P. Xing. SoCC 2014.

Managed Communication and Consistency for Fast Data-Parallel Iterative Analytics. Jinliang Wei, Wei Dai, Aurick Qiao, Qirong Ho, Henggang Cui, Gregory R. Ganger, Phillip B. Gibbons, Garth A. Gibson, and Eric P. Xing. SoCC 2015.

GeePS: Scalable Deep Learning on Distributed GPUs with a GPU-Specialized Parameter Server. Henggang Cui, Hao Zhang, Gregory R. Ganger, Phillip B. Gibbons, and Eric P. Xing. Eurosys 2016.

Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters. Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang Wei, Pengtao Xie, and Eric P. Xing. ATC 2017.

接下来我们详细探讨其中的某几个问题。


从CPU到GPU

经过上面的介绍,相信Parameter Server的架构大家心里有数了。下面我们实践一下:用PS把一个在单机上运行的神经网络训练程序并行化到 {\large P} 个机器上。按照我前面讲的思路,很容易将前面的迭代循环改写成如下这种形式:

\begin{equation} \large \begin{aligned} \mbox{fo} & \mbox{r } t = 1 \rightarrow T: \\ & \texttt{pull}(\large \theta^{(t)}) \\ & \nabla \theta^{(t)} = \Delta_{\mathcal{L}} (\theta^{(t)}, D_p^{(t)} )\\ & \texttt{push}(\nabla \theta^{(t)}) \end{aligned} \end{equation}

上面这段循环会运行在所有的 {\large P} 个计算节点上。

如果上述程序的主要计算函数 {\large \Delta_{\mathcal{L}}} 在CPU上完成,那么上述流程基本概括了绝大部分CPU上的深度学习(甚至更广泛的机器学习)系统,比如DistBelief,Project Adam的基本原理。

但是,当把核心计算函数 {\large \Delta_{\mathcal{L}}} 挪到GPU上执行时,由于GPU有自己的内存,所以 {\large \Delta_{\mathcal{L}}} 所需要的输入以及输出也必须位于GPU内存上(以便GPU能快速读写)才能保证运行效率。在分布式环境下,{\large \theta^{(t)}} 作为 {\large \Delta_{\mathcal{L}}} 的一个输入,是从远端pull回来的。目前绝大部分的网络通讯协议只能从一台机器的RAM发到另一台机器的RAM(Nvidia其实有GPUDirect技术,能让两个GPU直接通信,不过这个对硬件有一定的要求,我们后面再谈),所以其实在pull执行之后和 {\large \Delta_{\mathcal{L}}} 函数执行之前,一般有一个隐藏的memcpy操作,把收到的参数 {\large \theta^{(t)}} 从RAM移动到GPU内存上;类似地,在 {\large \Delta_{\mathcal{L}}} 算完后,同样需要有一个memcpy操作把算出来的梯度从GPU内存移动到RAM上,才能调用push进行通信。

总结一下,有了PS的push和pull两个API,再加上两个memcpy,我们就设计出了一个可以在GPU集群上跑的深度学习系统了,听起来是不是特别简单?其实目前大部分的深度学习框架号称自己支持分布式GPU集群也大概就是做了上面这点简单的工作。

如果问题真的这么简单,那么其实这篇专栏文章写到这里就算结束了。理解了的同学,只要先到Github找一个PS的实现(比如这个这个),然后再按照上述思路就可以把任何深度学习框架,或者自己写的神经网络训练程序改成一个“能在GPU集群上并行跑的分布式深度学习系统了”。事实上,只要完成了上面这一步,市面上大部分的神经网络确实都可以在GPU集群上跑起来了。

然而,如果你真的动手去实现了这么一个系统,再去弄一个GPU集群,找几个流行的神经网络上测试一下自己的系统的性能,你会发现当你用很多台机器并行训练某个神经网络时,你并不能获得多大收益。比如,你用8台机器同时训练一个VGG19,可能只能获得2-3倍的加速;换句话说,你花8000刀买了8块显卡,最后发现你的系统实现让6000刀直接打了水漂。事实上,如果一个分布式系统用了8台机器实际只2-3台机器的收益,这个系统显然是不达标的。并且,如果你继续买更多的机器和GPU,你会发现加速比并不会继续提高,甚至有可能降低......


GPU快所产生的问题

接下来我们就来详细讨论一下上面这段程序在面对GPU集群的时候会出现什么样的问题(当然,如果你自己动手实现了,你能很容易观测到这些问题)。

如果把上述系统部署到GPU集群上,并且仔细profile一下整个训练过程,我们就很容易发现几个很明显的问题:

  • 内存移动:在GPU内存和RAM之间拷贝数据(memcpy)会产生一定的开销。这个开销如果跟CPU上的计算比可能可以忽略不计。但是在高性能计算领域(比如GPU上),完成一步计算和调用一个GPU内核程序(kernel launching)花的时间是差不多的,所以这个memcpy在很多情况下会和实际计算花差不多的时间。在这种接近于1:1的比例下,这个数据拷贝的开销就没法忽略不计了。
  • 通讯:机器跟机器之间的远程通讯需要一定的时间来完成(latency)。更进一步,我们可能还需要保证参数之间的同步性,也就是说,在每个iteration,服务器端必须等接收到所有计算节点的梯度更新值({\large \nabla \theta^{(t)}_p} ),然后把他们加到上个iteration的参数 {\large \theta^{(t)}} 上,才能响应计算节点的下一轮pull请求 --- 发出新的参数 {\large \theta^{(t+1)}} ;另一方面,客户端(每个计算节点)在完成一轮计算并把梯度push出去后,必须等待服务器回传新的参数 {\large \theta^{(t+1)}} ,才能开始下轮计算。考虑到网络状态一般都不稳定,网络带宽也不一定充裕;再考虑到每个机器的快慢不一(不同节点完成 {\large \Delta_{\mathcal{L}}} 所需时间有快有慢),那么想要保证上述同步性,可能就要花相当长的时间纯粹用来等待各个节点(包括服务器)上的通信完成,或者等待最慢的计算节点完成计算任务。
  • GPU内存:我们都知道,相比较于RAM,GPU内存通常是非常有限的(现在最大的Nvidia GPU也就只有大概12G内存)。但是训练神经网络的时会产生很多中间状态值(intermediate states),如各层的activation。在后向传播的时候,这些中间状态值会参与梯度的计算,所以必须保存下来,并且很轻易就把GPU内存占满了。假设我们训练的神经网络很大(或很深),GPU内存装不下了,就无法进行计算了。

上面三个问题里面,第一个和第三个问题显然是GPU集群上特有的;第二个问题,也就是如何尽量降低多机同步的时间开销,几乎是分布式机器学习(甚至整个分布式系统)这个领域最精经典、最常见、也最棘手的问题,我前面列举的很多论文就是专门讨论这个问题的。为什么我要在GPU集群这个背景下把这个问题重新拿出来讨论呢?考虑下面这个问题:当我们衡量时间开销的时候,我们到底是在乎绝对时间还是相对时间?回到上面的例子,考虑下面这两种情况:

  • 情况A:假设每次循环中, {\large \Delta_{\mathcal{L}}} 里面的计算平均需要5分钟完成,而push和pull对应的参数同步大概需要10秒完成;
  • 情况B: {\large \Delta_{\mathcal{L}}} 需要0.4秒完成,push和pull对应的通信只需要0.1秒完成。

A和B两个系统哪个更差?答案显然是B,虽然从绝对时间上看,花0.1秒在通信上远比花10秒快,但是其实在评价这个系统的时候,我们并不在意绝对时间,而更看重的是通讯和计算的相对时间(通讯计算时间比)。如果我搭了一个系统有A这种性能,其实这个系统已经足够好,没必要再优化了。但是对于B,每个计算节点完成一次计算后,要额外等待1/4的计算时间才能开始下次运算,这就是说有1/5的时间计算资源是闲置的。

这个问题和GPU集群有什么关系呢?事实上,当把这类分布式系统问题拉到高性能计算(HPC)领域讨论的时候,问题是截然不同的(要知道,HPC领域的研究者们通常都是在毫秒和微秒级做优化)。GPU做深度学习的相关运算一般比CPU快了好几十倍,以前每个iteration {\large \Delta_{\mathcal{L}}} 可能在CPU要个好几秒才能完成,但是到了GPU上可能只需要几十毫秒。在分布式系统里面,网络通讯的快慢受很多因素影响,包括你要通讯的内容的大小,你的通讯程序写的好不好,你的Ethernet硬件怎么样、可用带宽高不高,等等问题。由于GPU上计算时间大幅减小,通信时间和计算时间的比例会越来越大;随着这个比例增大,对系统的要求也就越高。

下面我来举一个简单的例子,带大家估算一下GPU上的深度学习对网络通信的要求到底在什么水平。但提前声明一点,下面的例子已经被我简化过了,和实际情况可能有一些出入,但是在帮助理解通信量上是没有差别的。对更实际的情况感兴趣的读者,可以参考Poseidon论文

考虑在Nvidia Gefore Titan X (GM200)这块显卡上训练AlexNet这个网络,假设我们把batchsize设成256,大概每0.25秒这块GPU就可以完成一个iteration的计算。另一方面,AlexNet有大约61.5M个参数,也就是说,在每台机器上,每块显卡大概每0.25秒就会产生61.5M个梯度值。假设我们在8台机器上并行训练AlexNet,每个机器有且仅有一块显卡,并且我们把所有的数值都存成单浮点数(float)。这是背景,接下来我们来动手算算具体会产生多少的通信量。

  • 先考虑PS的客户端。为了保证参数的同步,每个计算节点在每个iteration计算结束后,要先把这61.5M个梯度值发出去,再从PS收回61.5M个值,作为更新后的参数 {\large \theta^{(t+1)}}
  • 再考虑服务器端:每个iteration内,服务器端得收61.5M * 8 = 492M个梯度回来,把这些梯度加到 {\large \theta^{(t)}} 上得到 {\large \theta^{(t+1)}} ,再把 {\large \theta^{(t+1)}} 分别发给每个计算节点,也就是再往外发送61.5M * 8 = 492M个参数出去。

我们先假设最理想的情况:我们有办法将通讯时间和计算时间完全重合起来吧,那么我们的网络带宽要到多少才能保证这个通讯任务迅速完成而不阻塞下一轮的运算呢?

我们来算算:对于每个计算节点,我们要在0.25s的计算时间内完成发61.5M和收61.5M这两步,共需要61.5M × 2 / 0.25s = 492M/s的网络带宽;而对于服务器端,我们要完成收492M的梯度再发492M的参数这两步操作,也就是需要492M * 2 / 0.25s = 3936 M/s的带宽。如果我们再进一步把这个M换成Ethernet领域常用的Gbps:

  • 对于每个计算节点,所需要的带宽大概是492M/s * 32 bit = 15744 Mbps = 15.7 Gbps;
  • 对于服务器端,所需带宽3936M/s * 32 bit = 125.9 Gbps;

看到这个数据,你大概会惊呆了。你工作用的笔记本或者台式机的标准带宽大概是1Gbps,最高档的AWS GPU instance大概可以提供给你20Gbps的带宽,收费$16每小时;甚至可能市面上压根儿买不到这么高带宽的Ethernet。更何况我们才用了8台机器,随着机器的增多,通信量显然会线性增长。而且我前面给的数据是2年前的显卡的计算速度 ,现在最新版的Volta可能已经是这个的2 - 4倍了,也就是说计算时间会更小 -- 上面算术的分母会更小……

不过好消息是有一个很简单的工程方案可以极大的缓解上面这个问题。我前面说过,PS里面的“服务器”或“客户端”只是一个虚拟的概念,我们在实际工程实现中,可以让集群中的任何一个物理节点既充当计算节点(worker)又是一个服务器(server)。这样我们就可以把服务器端实现成一个分布式的存储系统 -- 我们要服务器要维护的模型参数平均分布到多台物理机器,从而可以充分利用每台机器的通讯带宽了。考虑这种情况,同样是一个8台机器的集群,让每台机器既充当服务器又充当计算节点,每台机器作为服务器端只需要负责1/8的参数通信。那么每台机器每秒需要多少带宽才能完成任务呢?相似的:每台物理节点,作为服务器需要收一次发一次,而同时作为计算节点需要发一次收一次,所以总体上大概需要在每秒传输4 × 61.5M / 0.25s × 7/8 = 840 M/s(为什么这里要乘以7/8?读者可以自己想想),对应的Ethernet带宽大约是26Gbps。可以看到,这个方法虽然大大的降低了服务器端的通讯总量,但是26Gbps还是一个太大的值,而且考虑到我们之前还做了大量的假设:我们假设了通信时间和计算时间是完全重叠的,还假设了通信的稳定性(不会有可用带宽不稳定的情况),而在真实环境下这些假设几乎都不可能满足。因此,相比较于CPU环境下,分布式GPU对通信有极其高的要求。

另外,以上对于通信量的估算也详细解释了我在知乎上对Facebook那篇训练ResNet论文的评论:ResNet的计算时间和其他网络比是偏长的(层数多,计算多),而参数量却很少(都是卷积层,没多少参数),因此通信计算时间比并不大,所以把ResNet分布式并行起来并获得不错的加速相比较于其他网络结构,更简单。当然,这也从另一个侧面反映了ResNet这个网络的优点。

回到这一节刚开始提的第三个问题:GPU内存空间不够。一个可行的解决方案就是使用模型并行(model parallelism),把内存消耗过大的模型分割成很多份放到不同的计算节点(GPU)上。由于这一系列的专栏主要讨论数据并行,在这里我们就不展开讨论模型并行了。而另一个可行的方法就是GeePS这篇文章里面提出的方法。

接下来我要介绍的我自己的两个工作,Poseidon和GeePS:

Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters. Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang Wei, Pengtao Xie, Eric P. Xing. ATC 2017.

GeePS: Scalable Deep Learning on Distributed GPUs with a GPU-Specialized Parameter Server. Henggang Cui, Hao Zhang, Gregory R. Ganger, Phillip B. Gibbons, Eric P. Xing. Eurosys 2016.

主要就是为了解决上面讨论的这三个问题的。


Poseidon,海神,深度学习

Poseidon这个系统提出了几个方案来解决上面提到的前两个问题。

我们在前面的讨论中明确了一点:网络通信无论如何都是有时间开销的,Ethernet的带宽高低或者latency大小只会影响这个时间的长短,但并不能把这个时间降到零。以这一点为前提,设想,如果按照前面讲的 {\large \texttt{pull} \rightarrow \Delta_{\mathcal{L}} \rightarrow \texttt{push}} 三段式流程串行的进行通信和计算,无论这个通信是快是慢,这个时间开销都会导致在分布式环境下每个iteration的时间比单机版要长,进而导致整个系统无论如何都无法达到线性加速(为什么?)。所以,把通信和计算重叠(overlap)起来以便“掩盖”通信时间几乎是一个必须的步骤。

回到神经网络的训练过程上,怎么设计系统来重叠计算和通信?大家回顾一下后向传播的细节,注意两点:首先,神经网络的计算是一层接着一层完成的,不管是前向还是后向传播,算完本层才能算下一层;另一方面,在后向传播的过程中,一旦后一层拿到前一层的输入,这一层的计算就不再依赖于前一层了。Poseidon利用了这种层与层之间的依赖性设计了一套pipeline,称为Wait-free Backpropagation(WFBP),用多线程的方法把计算和通信尽量重叠了起来。具体来说,WFBP试图并行执行一部分相互独立的通信和计算,以隐藏参数同步的时间开销:在后向传播的时候,当第 {\large i} 层的计算完成后,第 {\large i} 层参数的同步(push和pull)和它前面所有 {\large i-1} 层的梯度计算是独立的,也就是说二者可以并行;同时,由于参数的更新也是相互独立的,没有必要在所有参数都更新完成后再统一pull,取而代之,我们可以在某一层参数更新完成后立刻pull回本地,把上行通信和下行通信也并行起来。这么做其实还有一个额外的好处:很多常见的神经网络都是前面几层计算量大,而后面几层的参数多,比如很多常用的CNN中前面的卷积层计算比较费时,而后面的全链接层则参数量大,这样WFBP正好可以把整个计算流程中计算最耗时和通信最耗时的两大块重叠起来,大大隐藏了开销。

WFBP示意图:图中灰色代表计算,绿色代表向PS发送梯度更新(push),浅蓝色代表向PS请求最新参数(pull)

我们在普通的Parameter Server架构上加入WFBP后进行了一些实验,结果证明WFBP的优化效果是十分显著的。对于Inception-v3和ResNet-152等参数分布比较均匀,且没有巨大的全连接层的网络,只是使用WFBP就能在32个Titan X节点上达到接近30倍的加速。对WFBP Pipeline具体实现感兴趣的同学,可以参考Poseidon论文

WFBP的主要思想是通过多线程、多stream技术把通信和计算尽量重叠起来,但是仅有这一步仍是远远不够的。我们又在一个稍微大一点、通讯计算时间比更大的网络VGG19上做了一组实验:仅仅依赖WFBP,在32个节点我们大概最多能达到20x的加速比,虽然相比较于TensorFlow原版的负加速比(是的,TensorFlow自带的分布式训练在某些神经网络上可能是负加速,分布式训练比单机还要慢)已经很不错了,但是这个结果并不令人满意:如果你花了10万刀买了机器,大概有3万多刀打了水漂。

另外值得强调的是,我们上面的几组实验用了40Gbps的Ethernet,这个已经算是奢侈品了。考虑到大部分的云计算平台(AWS等),以及大多数的机器学习实验室其实很难有这种硬件配置(毕竟不是人人都是Google或Facebook,感兴趣的同学可以去看看Facebook的论文里面用了多少带宽的Ethernet),这个对硬件配置的要求显然并不现实。实际上,我们又把上面测加速比的实验重新搬到10Gbps和1Gbps的Ethernet上重新跑了几次,看到结果后心态直接崩了:由于可用带宽瞬间减小了很多,上述的好几个实验在32台机器上大概只剩个位数的加速比了(具体参见Poseidon论文)。我们用的GPU还是2015年的Titan X,而且这几个网络里面最大的VGG也仅有120M个参数而已。随着GPU的飞速发展,以及各种大模型被开发出来,这个问题会变得原来越严重。

问题出在哪?与单机上高速的PCIe总线相比,多机间的网络传输速度是很慢的,而深度神经网络的一大特点就是参数多,通信量大。设想,如果总通信时间超过了计算时间,就算你能够设计出完美的重叠通信和计算的方法,也无法掩盖通信时间 -- 你永远无法达到线性加速,因为在这种环境下,每个iteration所需的时间是由通信时间决定的,而并非计算时间了。其次,突发数据传输(burst)也是一大问题,如果网络通信匀速发生,那么只要可用的带宽大于所需要的通信速度,就能得到比较好的传输性能;如果通讯负载不均匀,在大部分时间网络带宽处于空闲状态,但是在某几个时间点突发很多信息需要传输,这就会导致在这一时刻对Ethernet带宽的要求剧增,并超过上限带宽。结果就是,在空闲时刻,Ethernet的带宽利用不充分;在忙碌的时候又负载失衡,在峰值传输点达到时网络通信会阻塞。

回到训练VGG19这个具体问题上。通过一些分析和实验,我们发现主要问题在于VGG19有两个巨大的全连接层(FC),同步这两层参数消耗了大量的时间,使得通信时间超过了计算时间,拖慢了速度。为了解决这个问题,Poseidon采取了一套叫做混合通信(hybrid communication)的机制,试图直接减小通信量(但不损失精度)。如果通信总量减少了,通信时间自然就变短了。什么是hybrid communication?大体思想是,对神经网络不同类型的层,使用不同类型的通信模式。举个例子,在Poseidon里面,对于一个卷积神经网络里面的卷积层,Poseidon直接使用PS进行参数同步,因为卷积层有参数共享,参数个数一般都很少;而针对全连接层这种参数特别多的层,Poseidon会根据情况,在必要时选择另一个叫做Sufficient Factor Broadcasting(SFB)的模式来进行参数同步。

SFB其实是个轻量级的分布式机器学习通信框架,想深究这一块的同学可以读读这篇文章。这套通信模式可以用到很多具备一定性质的机器学习模型上,而很多神经网络正具备这个性质。原理上,我们发现全连接层的参数实际上是个大矩阵,而且在后向传播时候产生的的梯度矩阵其实是两个向量的乘积(秩为1):具体来说,一个参数矩阵大小为 {N \times M} 的FC层,在算梯度的过程中,假设batchsize是 B ,那么它在这个batch上的梯度可以分解为 B 个长度为NM 的向量的乘积,这些向量被称为Sufficient Factor(SF)。在实际中, B 往往是远小于 NM 的,因此我们可以选择发送SF来减小通信量,而不用再发送原来那个很大的梯度矩阵,在通信结束后再用SF重构原始梯度矩阵即可。

Microsoft发表在OSDI 2014上的Project Adam系统也使用了一个类似的技术进行优化,不过Project Adam只优化了上行通信(push),下行拉取最新参数(pull)时,仍需发送整个大参数矩阵。这一方面没有解决下行通信量大的问题;另一方面,这种上行发SF下行发参数矩阵的设计导致只能有一个物理节点作为服务器来接受所有计算节点的SF,再由这个节点把更新过后的参数矩阵返还给所有其他节点。如果你理解了我前面讲的PS的设计,你大概会发现这会导致一个严重的问题:我们在实现PS的时候通常使用分布式存储,以便尽量分散利用每个节点的网络传输带宽;但是在这种设计下,某一个物理节点需要承担很大的通信任务,这个节点很有可能会成为一个通信瓶颈。不同的是,在SFB里面,每个节点发送自己产生的SF给其他所有节点(broadcast),每个节点收到其他节点发来的SF后在本地重构出梯度,再更新本地参数。每个节点承担相同的通信量,大大减小了负载不均衡的可能性。但是SFB也有自身的问题:P2P通信的开销是与集群规模的平方成正比的,当集群规模太大的时候,使用P2P的通信方式不仅不能使通信开销减小,反而可能增大。

所以这里Poseidon的hybrid communication就发挥其优势了。Poseidon并不是在任何情况下对任何FC层矩阵都使用SFB进行通信,而是用一个自适应的方法:用一个不等式去判断当前情况SFB和PS哪个更适合。不等式很简单,简单的推导就能得到,感兴趣的同学可以自己动手算一下。有了SFB之后,Poseidon处理大网络的能力得到了增强。对于VGG19,在32个节点以及很小的可用网络带宽上,也能轻松达到30x的加速;对于更加丧心病狂的VGG19-22k(把输出层从1000改为了21841,ImageNet 22K extreme classification的一个使用场景)也得到了28.5x的加速。相比较一些其它减少通信的方法,Poseidon也具备一定的优势,具体的可以参加Poseidon论文中的细节。


本篇总结和下篇预告

谈完了这些技术细节,我们可以再谈谈在设计Poseidon这个系统的一些思考。与Caffe, TensorFlow,Torch这种深度学习工具包不同,我们希望Poseidon更多的成为一个平台,而不是框架,或是工具包。在这个平台上,你可以把任何语言写的深度学习程序 -- 不管是Python写的TensorFlow或PyTorch程序,还是按照Caffe Prototxt定义的神经网络,抑或是Lua写torch程序 -- 都轻松的部署到Poseidon上面,甚至不需要改一行代码,就可以在一个GPU集群上跑起来,并且或得满意的加速。所以Poseidon尽量保存了原有框架的编程语言接口,所以当你把一个TensorFlow程序部署到Poseidon上的时候,几乎不用改任何代码。另一方面,我们把Poseidon的底层代码写的非常轻量级,用Poseidon把一个现有的非常复杂的深度学习框架(比如TensorFlow)改成分布式版本大概只需要多写一百行代码就可以了。在另一方面,我们在部署上花了很大的功夫,让Poseidon非常易于安装。当然,这方面会更多的涉及到关于深度学习的产品设计,就不再深入讨论了。

这篇文章想要讲的内容基本已经涵盖到了。那么,Poseidon系统的名字是什么意思呢?Poseidon,海神,深度学习......这个就留给各位读者自己联想了......

我会在第三篇文章,也就是这一系列专栏的最后一篇中谈谈GeePS这个工作,顺便聊聊当下分布式深度学习领域其他比较热点的话题,比如异步训练、动态神经网络、增强学习,等等。


当然,文章的最后不忘打个广告:Petuum Inc.上个月完成了9300万美元的B轮融资,欢迎对机器学习和分布式系统感兴趣的各位投简历。Petuum新一轮融资过后会在湾区建立office,在湾区和匹兹堡都有headcount,欢迎各位咨询!



* 这篇文章的封面图来自:facebook.com/convolutio.

发布于 2017-11-16

文章被以下专栏收录