参数服务器——分布式机器学习的新杀器

转自：微信公众号 数据极客
在大规模数据上跑机器学习任务是过去十多年内系统架构师面临的主要挑战之一，许多模型和抽象先后用于这一任务。从早期的MPI，到后来的Hadoop，乃至于目前使用较多的Spark，都曾被广泛应用于大规模机器学习处理任务。John Langford曾经在他的博客(机器学习领域最好的博客之一)上总结和对比了主流的分布式机器学习框架的抽象[1]：
MPI Gradient Aggregation：主要缺点是批任务求解器的速度不高，另外是MPI本身的问题无法支撑大规模数据。
MapReduce：解决了MPI无法支撑大数据的问题，但无法改进批处理求解器的训练性能，并且还引入了新的问题，包括迭代式计算的低效，节点之间通信低效。
基于图的抽象：由CMU推出的GraphLab是这方面的佼佼者，目前已经成立了Dato公司专门推广基于图的大规模机器学习系统。用图来做抽象可以解决许多机器学习问题，但仍然有许多问题无法很好高效求解，比如深度学习中的多层结构。
Parameter Server参数服务器：跟基于图的方法主要区别在于把模型参数存储和更新上升为主要组件，并且采用了异步机制提升处理能力，这是本文主要介绍的模型。
AllReduce：AllReduce本身就是MPI的原语，这其实是最显然和直接的分布式机器学习抽象，因为大部分算法的结构都是分布数据，在每个子集上面算出一些局部统计量，然后整合出全局统计量，并且再分配给各个节点去进行下一轮的迭代，这样一个过程就是AllReduce。AllReduce跟MapReduce有类似，但后者采用的是面向通用任务处理的多阶段执行任务的方式，而AllReduce则让一个程序在必要的时候占领一台机器，并且在所有迭代的时候一直跑到底，来防止重新分配资源的开销，这更加适合于机器学习的任务处理。AllReduce跟参数服务器都会是机器学习算法框架的重要抽象，DMLC组陈天奇开发的Rabit框架[2]是AllReduce模型的良好实现之一，其余的当然可以借助于vw这样的系统，它们都能直接运行在Hadoop上。基于同步的AllReduce模型并不是本文讨论的重点，我们只需要知道它很重要，很多算法离不了。

机器学习算法和计算机领域的其他算法相比，有自己的一些独特特点。例如：迭代性，模型的更新并非一次完成，需要循环迭代多次；容错性，即使在每个循环中产生一些错误，模型最终的收敛不受影响；参数收敛的非均匀性，模型中有些参数经过几个循环便不再改变，其他参数需要很长时间收敛。这些特点决定了机器学习系统的设计和其他计算系统的设计有很大不同，因此理想中的分布式机器学习任务，并不能随着机器的增加而能力线性提升，因为大量资源都会浪费在通讯，等待，协调，这些时间可能会占据大部分比例。参数服务器就是被提出来专门用于大规模最优化处理的框架，它特定用于这种需求：大规模的训练数据，比如TB甚至PB级别的；大规模的模型参数，在大规模的优化框架中，常常会有数十亿乃至千亿级别的参数需要估计。因此，在设计面临这种挑战的系统时，比如大规模深度学习系统，大规模Logistic Regression系统，大规模主题模型，大规模矩阵分解等等依赖于SGD或者L-BFGS最优化的算法，需要解决频繁访问修改模型参数时所需消耗的巨大带宽，以及如何提高并行度，减少同步等待造成的延迟，还有容错等挑战。

参数服务器的概念最早来自Alex Smola于2010年提出的并行LDA的框架[4]。它通过采用一个分布式的Memcached作为存放参数的存储，这样就提供了有效的机制用于在分布式系统不同的Worker节点之间同步模型参数，而每个Worker只需要保存它计算时所依赖的一小部分参数即可。当然，这里存放参数的存储跟做OLTP应用中的Key-Value抽象有所不同，因为以Key-Value为单元进行频繁的参数数据交互会导致过高的通信开销，因此参数服务器通常采用数学封装来进行参数同步，比如向量，张量，矩阵的行列等。

上图的sampler是并行LDA里的组件，可类比为通用参数服务器框架里的计算单元。Smola提出的模型是最早的参数服务器抽象，随后出现了许多改进，最出名的应当是Google的跨界高人Jeff Dean 2012年进一步提出了第一代Google大脑的解决方案DistBelief[5]，主要用于超大规模深度学习网络的训练。DistBelief将巨大的深度学习模型分布存储在全局的参数服务器中，计算节点通过参数服务器进行信息传递，很好地解决了SGD和L-BFGS算法的分布式训练问题。由于SGD和L-BFGS是机器学习的普遍性优化问题，因此尽管DistBelief是作为深度学习的系统框架而提出，但DistBelief的核心结构却可以应用到多种普通机器学习手段中。相比最早的参数服务器模型，DistBelief把该模型扩展成为更加通用和灵活的框架，豆瓣的Paracel[3]正是参考DistBelief的直接实现，先来看看Paracel和DistBelief模型：

图中是分布式异步SGD架构流程图，运行时，需要把训练数据分为多个子集，然后在每个子集上运行模型的多个副本，模型通过集中式的参数服务器通信，参数服务器存放了模型的全部参数和状态。异步体现在两方面：模型的副本独立运行；参数服务器的分片也各自独立运行。DistBelief没有过多谈论系统实现，从Paracel里我们可以看到具体的工程实现：总体上Paracel实现非常简单，参数服务器直接采用内存hashtable，并封装了对分网络，图，稀疏矩阵，稠密矩阵等数据格式用于参数同步。Paracel解决的另一问题是straggler问题：由于一些软硬件的原因，节点的计算能力往往不尽相同。对于迭代问题来说，每一轮结束时算得快的节点都需等待算得慢的节点算完，再进行下一轮迭代。这种等待在节点数增多时将变得尤为明显，从而拖慢整体的性能。Paracel放宽了“每个迭代步都等待”这个约束：当在一轮迭代结束时，算得快的节点可以继续下一轮迭代，但不能比最慢的节点领先参数s个迭代步。当领先超过s个迭代步，Paracel才会强制进行等待。这样异步的控制方式既从整体上省去了等待时间，也能间接地帮助慢的节点赶上。从优化问题的角度来看，虽然单迭代步收敛得慢了，然而每个迭代步的时间开销变少了，总体上收敛也就变快了。这种做法又叫Staleness Synchronous Parallel (SSP)，基本思想是允许各机器以不同步调对模型进行更新，但是加一个限制，使得最快的机器的进度和最慢机器的进度之差不要太大。这样做的好处是：既减轻慢的机器拖整个系统的后腿，又能保证模型的最终收敛。

SSP是相对于BSP(Bulk Synchronous Parallel)来说的，BSP是上世纪八十年代就提出的，它要求算法在每一次迭代结束后都要同步等待，因此会因为最慢的机器拖慢整个系统。BSP在绝大部分已有的分布式机器学习和数据挖掘框架框架中都在使用，例如Spark MLBase，Google Pregel，Apache Hama等。

SSP是由CMU Eric Xing的Petuum项目组提出的[6]，Paracel引入SSP使得豆瓣的参数服务器方案工程上更加成熟，在Paracel内部，SSP的等待通过调用MPI来实现。关于一致性收敛和Petuum，在下边还会有介绍。关于参数服务器，另一个重要的方面是容错设计。在几十台机器的集群上运行，这也许并不是一个问题，但是如果在有上千台机器的集群上运行任务，节点发生任务失败的概率就会大很多，如果缺乏容错设计，就会导致任务重启，从而浪费大量时间。不过，在Paracel的代码里并没有找到相关的处理逻辑，通常容错处理需要借助于Checkpoint来做快照，这样任务重启时无需从头进行，比如DistBelief就是这样处理。跟豆瓣的工程师咨询后已经确认，在开源版本的Paracel里确实还没有相关设计。

上面讲述了不少参数服务器的背景和系统结构，那么为什么参数服务器能够具备更好的性能呢？仍以SGD为例说明：在传统同步SGD中，如果一台机器失效，整个训练过程将会延时；但是对于异步SGD来讲，如果某个模型副本的一台机器失效，其他模型副本仍然继续处理样本并更新参数服务器中的模型参数，因此异步SGD具备更好的鲁棒性。此外，多种异步处理形式给最优化过程带来进一步的随机性：模型实例最可能是使用一个稍微过时的参数来计算梯度，因为这时其他的副本可能已经更新了参数服务器上的参数。除此之外还有其他随机的来源，因为参数服务器组的每台机器是行为独立的，所以无法保证在给定时间点上，每个节点的参数被更新的次数相同，或者以同样的顺序被更新。更进一步，因为模型副本使用不同的线程来获取参数和推送梯度值，故在同一时间戳上，单个副本内的参数将有额外的稍微不一致的现象。尽管对于非凸问题的这些操作的安全性缺乏理论基础，但是在实践中，这种放松一致性要求的做法是相当有效的。传统同步SGD的最优化过程，每次迭代选取的方向是由全部训练数据决定，或者由随机选定的一小部分训练集指定(mini-batch)。而异步的做法由于上述更多的随机性则会同时在很多方向上由不同的mini-batch选定不同梯度方向，这就好比整个最优化过程是以一个区域为单位进行的，而区域内的点代表不同SGD的过程，因此这种并行化的工作会带来性能上的提升。

豆瓣的Paracel并不是唯一一种开源的参数服务器，这里继续介绍另外一个重要项目，来自Alex Smola的高徒——李沐设计的参数服务器[7]。这个项目在早期拥有一个独立域名http://parameterserver.org，后来因为李沐和陈天奇等国内英才成立的DMLC深度学习项目组，之前的项目也进行了重构因此转移到[7]所在的地址，而项目的背景介绍则在[8]和[9]。从架构上来说，ps-lite跟Paracel并没有什么不同，作为参数服务器，都需要一个全局分布式的key-value用来存储算法的模型或参数。当计算节点需要某个参数的时候，可以从参数服务器上读取。用户可定义不同的函数在参数服务器端对参数进行更新、过滤等操作。在大部分情况下，计算节点之间的通信都是通过参数服务器进行。图中W代表计算节点，X代表参数服务器节点。

根据作者的宣传，ps-lite应当属于第三代参数服务器，就是提供了更加通用的设计，在架构上包含一个Server Group和若干个Worker Group：
Server Group用来做参数服务器，每个Server Node存放一个参数分片，由Server Manager管理整个Server Group，维持整个Server Group的元数据的一致性视图，以及参数分片情况。
每个Worker Group运行一个应用，Worker Node只跟Server Node通信用来更新参数，Worker Node之间没有任何交互。每个Worker Group内有一个调度器，负责给Worker Nodes分配任务以及监控，如果有Worker Node挂掉或者新加入，调度器负责重新调度剩余的任务。

跟Paracel不同，ps-lite提供了多种数据一致性选择：
Sequential：这种情况下，所有任务顺序进行，因此数据严格一致，不会出现不同副本看到的数据有不同的情况，因此实际上跟前文介绍的BSP是等价的。
Eventual：这种情况下，所有任务并行执行，因此拥有最大的随机性。Eventual只适用于对于数据一致没有要求，非常健壮的算法，比如SGD。
Bounded Delay：每个任务需要设置最大超时时间，在该时间之前如果有任务未结束，那么新任务将会等待。Bounded Delay类似于上面的SSP，只不过这是用时间而SSP则是用迭代次数。
在容错设计方面，ps-lite通过给参数服务器引入多副本机制来提供：整个模型参数按照一致哈希分片存储，在默认情况下，采用链式复制确保参数多副本，如下图所示：

链式复制会导致网络带宽占用增加数倍，而ps-lite还提供了另外一种容错设计：先聚合再复制。聚合的意思是在机器学习算法中，参数在很多时候是可以累加的，比如梯度。采用先聚合再复制的方式，可以降低网络带宽占用。

根据作者在Logistic Regression上的测试，ps-lite可以比传统实现所占用时间缩短1到两个数量级。

ps-lite目前在DMLC项目中处于核心基础地位，因为大部分的分布式机器学习算法都会基于它来进行，包括DMLC最新推出的热门深度学习框架MXNet。ps-lite的代码整体非常简单，便于修改和移植，而且DMLC项目组目前也给它增加了资源管理器的集成，使得Yarn能够来管理参数服务器的资源分配。
下面再来介绍由CMU机器学习系领头人Eric Xing带领的小组推出的参数服务器Petuum。Eric Xing和Smola和李沐同来自CMU，但却做出了两份独立的工作，个中缘由不在本文八卦之行列，毕竟跟双方都不熟悉。因此这里只谈论技术。事实上，Petuum是最早开源的参数服务器，其目的都是在DistBelief之后期望推出通用的参数服务器设计。跟Paracel和ps-lite一样，Petuum也采用C++开发，Eric Xing据此给出的解释是目前Petuum仍然处于原型阶段，是学术产品，所以没有考虑通用的语言如Java等。Petuum目前分成几个子项目，分别包含了参数服务器Bosen，以及基于Bosen和Caffe的分布式深度学习系统Poseidon，后者不是本文介绍的范围。Bosen的系统设计建立于机器学习的特征之上，目前包含两个主要模块：Key Value Store和Scheduler，一致性协议是上文介绍过的SSP。通过调节SSP的staleness参数，SSP可以转化成数据流系统常用的BSP（Bulk Synchronous Parallel) 协议或者早期机器学习系统（如Yahoo LDA）使用的ASP(Asynchronous Parallel)。Scheduler用于模型并行，它提供的编程接口主要包含三个操作：
Schedule: 调度节点根据模型参数的相互依赖性和收敛的不均匀性，自动选择一个待更新的参数子集。
Push: 调度节点令计算节点并行地为选好的参数计算Update。
Pull：调度节点从计算节点收集Update，并更新参数。

Petuum/Bosen的架构如上图所示。跟其他的参数服务器并没有大的差别，但模块化设计更加良好，比如一致性模型，调度这些重要功能都放入单独组件。比较遗憾的是Petuum/Bosen也没有在容错设计上有所考虑，这跟Eric Xing宣称的原型系统也相吻合，因此跟Paracel类似，Bosen目前主要适用于几十台机器的集群，在更大集群上处理有风险。

从代码上来看，Bosen的结构相比Paracel和ps-lite都要复杂不少，主要原因是Bosen在所有的组件，包括存储，调度，还有Worker上面都是多线程实现，调度器的设计更为复杂，因为对机器学习模型参数的更新进行细粒度的调度，能根据参数的优先级自动调整更新次序，并根据参数的相关性防止不安全的并行。我们来进一步分析功能：

如上图所示，在Bosen中，SSP跟参数服务器封装在一起，称为SSPTable，供多个Worker节点以类似分布式共享内存的方式访问，让Worker节点操作SSPTable跟操作本地内存一样。Worker通过SSPTable更新参数服务器的同时，SSPTable同时更新SSP一致性控制器内部的设置。SSPTable是参数服务器节点运行的主体结构，Bosen在每个参数服务器节点上可以启动多个SSPTable，这是它显著区别于ps-lite和Paracel的地方之一。之所以这样设计，是因为Petuum项目需要根据不同的算法来配置并行任务。例如，Petuum的示例程序中，运行深度学习DNN算法，就配置了神经网络层数两倍的SSPTable数量，而其他一些算法比如随机森林，矩阵分解，只配置了2个SSPTable。究竟应当如何配置并没有从Petuum的文档和代码中找到说明，因为本文只是介绍性文字，故而这里不去深究，在使用中需要注意这一点。SSPTable内部运行若干Background线程，来运行SSP控制逻辑，线程的数目跟SSPTable数保持一致。Bosen实现了多种SSP模型，默认的SSP实现需要根据客户端(运行在参数服务器节点)汇报的时钟计数来决策，时钟在这里代表机器学习算法的计算单元，比如一次迭代。其他还有SSP Push和SSP Aggregation，前者只是在SSP基础之上提供额外的接口用于异步推送SSPTable的指定行数据，后者则修改了SSP的实现原语，如何使用这两个模型目前还没找到正式说明，使用SSP应当基本可以满足全部需求。

接下来看看Bosen的调度器设计，这是Petuum项目区别于其他参数服务器的主要组件之一。它的主要思想是在并行的同时避免模型出错，这个概念叫做Structure-Aware Parallelization(SAP)[11]，在Petuum项目里叫做Strads。

Strads系统由若干Master节点，若干Worker节点和一个Cordinator节点构成。调度流程如下：Master执行Schedule接口，作用是选出可以并行的参数子集，这个过程可能需要Strads从参数服务器读取参数数据。用来更新这些参数的任务被Cordinator通过Push接口下发到Worker节点进行计算，参数服务器通过Pull接口从Worker节点接收参数然后更新存储。为了更有效地执行任务，Strads把调度流程流水线化，Master无需等待Worker的结果就提前把任务准备好，Cordinator依次把待执行的任务下发到Worker。在决定如何并行从而避免模型出错这方面，并不存在一个通用的做法，因为这跟不同的模型有很大关系，因此Strads把这个工作留给不同算法来进行，例如对主题模型，Lasso，矩阵分解，都提供了相应的实现。

Eric Xing把他眼中的若干分布式机器学习模型的适用场景做了个概要图，可以看到，Petuum/Bosen的定位在于在较小集群上运行，但同时需要大量参数(百亿，千亿级)的场景，而DistBelief和ps-lite这些工作是运行在大规模集群上的参数服务器方案。至于为何是这样的结论（更大量参数），从架构上还没有得出很明显的结论，只能说Petuum项目相比其他参数服务器，对于更广泛的机器学习算法上考虑更多，然而由于缺乏容错机制，所以最好还是运行在独立中小集群之上(几十台服务器规模)。

上面介绍了3种最知名的开源参数服务器，随着这种模型为更多人所接受，一些公司也推出了相关的框架和产品，例如微软研究院[12]和英特尔[13]。下面分别简要介绍一下：

微软研究院的参数服务器是跟它在15年底开源的机器学习工具包DMTK一起发布的，DMTK最知名的项目就是老师木主导的LightLDA，这是一个超大规模主题模型(百万主题级别)，在最初的版本中，LightLDA正是基于Petuum来开发的。DMTK内部的参数服务器项目叫做Multiverso，在架构上比较简单，由于是来自研究院的项目，因此对大集群，容错等方面考虑得并不多。一致性模型方面，Multiverso包含BSP，SSP，以及ASP(Asynchronous Parallel)，ASP就是指全异步，所有的任务相互没有等待。在一致性模型的选择上，Multiverso跟ps-lite类似。

英特尔的DistML是对Spark的通用机器学习库MLLib的一个补充，它跟Spark和Hadoop的关系如下图所展示。这是本文谈论到目前为止第一个JVM上的实现。DistML利用Spark RDD任务来运行Worker，在Spark之外新增了两个组件，一个是Databus用于参数传送，另一个就是参数服务器本身，基于Actor实现。尽管Spark RDD任务具备容错功能，但参数服务器本身并没有类似ps-lite那样的多副本机制。由于Spark RDD的限制，DistML的参数更新只能在一次迭代完成后进行，因此这并不是严格意义上的参数服务器方案，所以也就谈不上采用何种一致性模型，但DistML第一次在Spark集群上增加了类似参数服务器的模型抽象，在网上看到有人用它尝试过4亿参数和200亿维度的Logistic Regression，从这个角度来说，DistML也是满足了许多依赖Spark进行算法开发的人的需求。

除了以上这些参数服务器，Spark社区本身也有支持参数服务器的计划[14]，从issue的建议来看，Spark计划采用可选的BSP和SSP作为一致性模型，容错设计采用Checkpoint定期刷盘。由于是真正的参数服务器模型，跟Spark RDD本质上的BSP有冲突，因此这需要对Spark核心的修改，目前的实现在[15]，也许在Spark上使用参数服务器应当不会很久了。

在商业化方面，国内的百度，阿里，今日头条是已知大规模采用参数服务器的公司。阿里的材料可以从QCon 2015上海的演讲中获得：

从上图显示，阿里的方案采用HBase作为参数服务器，Worker则采用内部的数据流式计算引擎，并且只实现了DistBelief里的异步SGD—Downpour SGD，因此，这个架构主要是服务在线实时计算，而不是通用的参数服务器。另据称今日头条的参数服务器集群已达到4000台规模，百度也是拿C++自行研发专用集群，这些公司并没有采用开源方案，主要原因应当还是在于超大规模集群上的容错处理，因此，这两家公司目前在大规模集群应用参数服务器的经验属于领先(阿里的上述方案目前只用于特定应用的在线实时计算，在容错架构上不需更多复杂的设计)。

这篇文章是本人在2014年就有想法撰写的，一直拖拉到了2016年初才完成初步轮廓。短短的一年多内，参数服务器发生了巨大的变化：DMLC项目组的横空出世，Petuum的重构，Paracel的开源，以及其他众多公司的工作，并且这些工作全部是国人或者国人主导的项目。究其原因，是因为在最热门的机器学习算法中，包括深度学习，推荐引擎，计算广告等领域，参数服务器相对于其他抽象都具有无可比拟的优势，再加上其架构实现简单，因此出现这么多的可选就并不是一件意外的事情了。在为自己的项目引入参数服务器的同时，需要深刻了解每种方案背后的应用场景。在当前来说，所有的开源方案，都还不具备在大规模商业集群上使用的能力，因此这也是为何会有公司自己攒相关轮子的缘故。如果你想打造自己的轮子，DMLC的ps-lite是一个好的起点；如果你对机器学习算法很精通并且愿意探索更优秀的性能，你值得在Petuum基础上进一步研究；如果你只想等着在已有中等规模集群上使用，那么可以等待Spark社区。可以预计，在接下来的一两年中，相关的工作会更快地推进，尤其是在可运维性和容错方面有更加周到的考虑，这也是基础架构领域值得大干特干的若干领域之一。

[1] http://hunch.net/?p=151364
[2] https://github.com/dmlc/rabit
[3] http://paracel.io/
[4] Alex Smola, An Architecture for Parallel Topic Models. VLDB, 2010
[5] Jeff Dean. Large scale distributed deep networks. In NIPS, 2012
[6] Solving the stragglerproblem with bounded staleness. In HotOS (2013).
[7] https://github.com/dmlc/ps-lite
[8] Mu Li, Dave Andersen, Alex Smola, Scaling Distributed Machine Learning with the Parameter Server, In OSDI, 2014
[9] Mu Li, Zhou Li , Alex Smola, Parameter server for distributed machine learning, In NIPS, 2013
[10] Eric P Xing, Petuum: A New Platform for Distributed Machine Learning on Big Data, In SIGKDD 2015
[11] Seunghak Lee, Jin Kyu Kim, Xun Zheng, Qirong Ho, Garth A. Gibson, and Eric P. Xing, On Model Parallelization and Scheduling Strategies for Distributed Machine Learning, In NIPS 2014
[12] https://github.com/Microsoft/DMTK
[13] https://github.com/intel-machine-learning/DistML
[14] https://issues.apache.org/jira/browse/SPARK-6932
[15] https://github.com/chouqin/spark/tree/ps-on-spark-1.3