一个轻量级分布式机器学习系统的设计与实现

宋匡时,李 翀,张士波

(1.中国科学院 计算机网络信息中心,北京 100190; 2.中国科学院大学 计算机科学与技术学院,北京 100190)

0 概述

随着互联网信息处理技术的迅速发展,社会生活越来越多地依托于互联网平台服务,互联网因此积累了大量实体与行为数据。机器学习作为一种从数据中挖掘模式特征和模式关联的重要工具,在应用于精准推荐、文本相关性判别、语音识别、人脸识别等场景时能够创造巨大的社会价值与商业利益。

同时,互联网技术的高速发展也对机器学习模型训练的速度要求越来越高。一方面,在产业应用中,更快的模型训练意味着能提升模型时效性,更快捕捉到短期用户行为喜好的变化,也能解决超大模型训练问题;另一方面,在一些前沿研究领域中,神经网络架构搜索(Neural Network Architecture Search,NAS)需要随机采样大量不同超参数以及神经网络架构的训练效果作为NAS的训练样本。对于模型训练加速问题,目前主要有2种解决方式:一种是通过等价数学变换、稀疏量化、低精度训练方法降低计算的复杂度;另一种是组合若干个计算节点的资源,实现可扩展的模型存储与训练加速,最大化计算硬件的利用率,代表性的工作有参数服务器理论[1]与环拓扑All-Reduce[2]。在异步梯度下降方面,文献[3]证明了凸函数无锁并行随机梯度下降的收敛性,文献[4]证明了与SSP半同步的异步梯度下降的收敛性;在All-Reduce方面,研究者提出了层次化All-Reduce与大批量训练加速方法,并使用PCIe/NVLink[5]、多网卡、RDMA[6]提高物理通信效率。之后又出现大量大规模机器学习的研究[7]。

目前大规模机器学习成为机器学习算法与分布式系统的交叉研究领域,机器学习系统也正向细分领域高定制化的方向发展。现有的主流机器学习系统大多由国外一线互联网公司开发,如Google的Tensorflow[8]、Facebook的PyTorch[9]以及Amazon的MXNet[10],其侧重面向高性能设备的机器学习任务,对不同硬件指令做了大量的适配优化工作,不可避免地造成项目依赖复杂、调试困难、对内存资源消耗较高以及在多机扩展尤其是资源受限集群的分布式机器学习任务上稳定性不足的问题,在涉及高定制化、计算资源受限设备情况下,二次开发与定制非常复杂且不可控,新增、修改与定制化功能也十分困难,这催生了众多研究机构与公司针对自身需求开发各类高定制化机器学习与深度学习系统[11],设计满足基本生产环境要求的轻量级、低资源占用、容错措施灵活的高性能机器学习系统成为迫切需要。

本文设计一个轻量级高模块化并且可扩展的机器学习系统,移植并优化实现多种主流机器学习与深度学习算法,同时设计参数服务器与环拓扑All-Reduce分布式梯度同步方案,对算法模型进行并行训练加速,以提高数据并行训练的效率与稳定性。

1 模块化分层设计

为满足轻量级系统的设计要求,本文将系统划分为计算层、梯度同步层与通信层,层内为可独立拆分的功能模块与工具模块,功能封装为编程对象。图1展示了系统内的主要模块设计,其服务于算法移植与分布式模型训练。其中:计算层包括提高并行度的线程池与SIMD指令,用于提高算法模块计算性能;梯度同步层包括不同同步方式处理逻辑与共享内存哈希表、进度动态控制模块等;通信层包括网络I/O模块、多线程控制、消息队列等。这样的分层设计避免了循环依赖的耦合问题,一个功能调用仅需引入调用链上必要的模块,模块代码剥离方便清晰,可避免无关内存占用。

在精简外部依赖上,本文选取目前网络传输吞吐量较高、跨平台且轻量级的网络框架ZeroMQ作为系统的网络通信底层,上层则进一步抽象消息发送队列、重发线程、事情循环等基础网络功能,轻量且低资源消耗地为整个系统提供异步非阻塞的高吞吐量消息发送与接收处理能力。

本文系统在深度学习方面使用基于层与计算单元的模块化设计模式,优点在于能够集中控制内存的占用,功能层次简单。系统各模块按照训练阶段划分,提供对初始化、训练、梯度同步、预测等阶段统一的调用接口,方便对训练过程进行控制与干预。

2 算法移植与实现

轻量级设计在降低资源消耗同时也应保持系统的高性能与实用性,本文系统移植并深度优化实现了主流机器学习与深度学习算法,可应用于推荐、点击率预估场景下的群体发现、行为序列分析、文本内容分析等任务。表1列出了目前已经移植实现的算法,通过SIMD向量指令、指令流水线、多核多线程、异步文件I/O多重加速最大化单机CPU计算能力,利用内存复用与预取节约内存并降低cache miss,用对凸优化友好的无锁多线程计算,大幅提高算法的运行效率。设计中考虑了模块化与可扩展性,对同类算法进行了统一抽象封装,很容易就可以将这些算法模块组合,实现将线性模型与深度模型组合的Neural factorization machines[12]、Wide & Deep[13],以及将变分贝叶斯与深度模型进行组合的Variational AutoEncoder[14]。同时也为数据并行训练提供了梯度同步接口,对梯度内存进行了统一管理,便于对梯度进行零拷贝分布式同步。

表1 移植实现的主流算法Table 1 Migrated and implemented mainstream algorithms

3 可扩展性系统设计

单机的训练不能满足大规模模型训练的参数存储与训练加速需求,需要机器学习系统具有较强的可扩展性。本节介绍针对数据并行模型训练的轻量级、高扩展性系统设计。

3.1 可扩展性设计的性能模型

数据并行的机器学习梯度同步方法按集群拓扑结构主要分为主从结构(又分一主多从、多主多从)、树结构、环结构。一主多从代表为Map-Reduce[15],多主多从代表为参数服务器,树结构代表为MPI All-Reduce,环结构代表为Baidu提出的通过Ring-base的All-Reduce。

定义1令N为节点个数,K为待同步的参数数量,H为同步梯度数据包大小,每个节点在一次同步中的通信次数为S,每字节传输的耗时为ξ,梯度下降优化参数的耗时为η,每次通信网络数据包传输延时θ,则同步梯度过程中的数据传输耗时由数据包传输耗时与合并参数耗时组成,表示为:

Φ=S(Hξ+θ)+Kη

根据定义1可知,不同网络拓扑组织方式直接影响H与S的取值,使用Adam[16]、FTRL[17]等不同的梯度优化方法则影响η的取值。根据定义1,不同拓扑结构下同步方案的数据传输耗时如表2所示。在一主多从情况下,每个节点只需与主节点一次通信传输全部参数,传输耗时以主节点单网卡串行队列瓶颈计算共收发KN个数据;树结构用分治的方法降低为lbN;而多主多从是每个节点都即承担计算也承担局部参数存储,既要向其他节点进行w-1次通信传输对应的局部参数,也要接收来自其他节点的局部参数,每个数据包大小为K/N;同样地,环结构由于基于邻居节点的迭代同步方法需要2(N-1)次通信,每个数据包大小为K/N。

表2 不同拓扑结构下的参数同步性能对比

Table 2 Comparison of parameter synchronization performance with different topologies

拓扑结构数据包大小通信次数数据传输耗时一主多从K1KNξ+θ+Kη多主多从KNN-1(N-1)KξN+(N-1)θ+Kη树结构Klb N(Kξ+θ+Kη)×lb N环结构KN2(N-1)2Kξ(N-1)N+2θ(N-1)+Kη

通过分析表2中传输耗时与节点个数N的关系可知,一主多从受限于单点瓶颈,传输耗时随节点个数呈线性关系,树结构则呈对数关系,多主多从与环结构呈更优的倒数关系。仅从同步参数的传输耗时上分析,多主多从在参数同步过程中通信次数是环结构的一半,但因为参数分布存储,前向预测过程同样需要N-1次获取参数的通信,环结构因为本身保存全量参数没有此消耗,所以两者在先前向后同步梯度的模型训练任务上总传输耗时相近。

在实际的多主多从同步过程中,引入高网络吞吐量的参数服务器作为只负责参数存储的节点,通信次数只与参数服务器数量有关,降低了传输耗时,且参数服务器支持异步梯度下降,减少了轮次间的训练依赖,并行度更高,因此,参数服务器方案被广泛用于Tensorflow、PyTorch等成熟框架和可扩展性机器学习训练任务。相比而言,环结构的Ring-AllReduce可以充分利用所有计算节点来追求更高的训练加速,由于梯度阻塞同步不会对训练收敛性产生影响,P2P的分布式拓扑避免了热点问题,因此常被用于计算机视觉等领域大批量训练数据的梯度聚合,在超大GPU集群上仅需分钟级别即可完成大规模模型训练[18]。2种同步方案均有其优势与适合的模型训练场景,因此,有必要分别设计参数服务器与环拓扑同步2种可扩展的模型训练方式。

3.2 参数服务器设计

参数服务器适合于参数总量超过单机内存的稀疏模型分布式训练场景,本文设计的参数服务器同步方案如图2所示,其中包括3个角色:一台主节点(Master)负责协调集群启动与广播集群拓扑,对所有节点进行心跳检测,出现节点异常时从拓扑中去除异常节点,自动重置拓扑结构;多台参数服务器通过一致性哈希分布式放置模型总量特征参数,满足训练过程中参数请求与梯度汇报的高吞吐处理需要;多台计算节点(Worker)负责读入训练数据并执行模型梯度计算。本文设计了适用于稀疏数据与稠密张量2种数据类型传输的Pull与Push协议,Worker在训练过程中无状态地向参数服务器Pull请求所需特征值后,向参数服务器Push汇报本轮计算的梯度值,由参数服务器将梯度信息进行筛选应用于随机梯度下降优化算法、并调控集群整体训练进度。

图2 优化传输量与收敛性的参数服务器架构

Fig.2 Parameter server architecture for optimizing transmission quantity and convergence

参数服务器同步方案的瓶颈是其在网络传输高延迟与丢包情况下的无效与重复梯度问题,本文通过传输压缩与约束梯度的方式对此加以改进。

3.2.1 网络传输压缩

考虑到此场景下为稀疏数据传输,本文采用稀疏编码,由特征哈希的编号与浮点数值的键值对表示稀疏参数与梯度。通过变长整数压缩特征编号,同时采用半精度浮点型压缩传输的梯度值。该设计考虑到模型梯度值精度与训练较为鲁棒,且混合精度训练可提升部分任务的效果,在提高传输数据压缩率的同时,使模型可以在混合精度下训练。通信底层借助高效的网络通信传输模块,满足节点高吞吐量的要求。

3.2.2 无效与重复梯度问题

当Worker汇报给参数服务器的梯度落后于总体训练轮次产生的无效梯度时,向多个Worker同时汇报相近的梯度方向,重复累计步长偏离最优解下降方向,导致参数服务器的异步梯度下降很不稳定。为此,通过Worker汇报给参数服务器轮次信息,设计了轮次间隔约束,其约束最快与最慢训练轮次的计算节点训练轮次间隔不大于一个阈值,当触发轮次间隔约束时阻塞最新轮次Worker的参数拉取,控制了异步计算进度差异产生的无效梯度数量。此外,使用一阶导的外积无偏估计二阶导矩阵的方法DC-ASGD[19],补偿延迟梯度。算法1与算法2分别描述了计算节点与参数服务器上的算法过程。

算法1计算节点

repeat

从参数服务器拉取批数据所需参数wt与step

repeat

计算参数梯度gm=fm(wt)

if 触发轮次间隔约束 then

sleep若干毫秒,等待最慢节点

until 获得全部所需参数

until forever

算法2参数服务器

输入学习率η,方差调控参数λt,Step最大间隔阈值E

初始化t=0,w0随机初始化,暂存参数wbak(m)=w0,计算节点编号m∈{1,2,…,M}

repeat

if 接收并解压到gmthen

if step触发轮次间隔约束 then

continue

wt+1←wt-η(gm+λtgm⊙gm⊙(wt-wbak(m)))

t ← t + 1

else if 拉取参数请求 then

根据step更新轮次间隔约束条件

wbak(m)←wt

压缩wt发送给计算节点m

end if

until forever

3.2.3 稳定性提升

除主节点负责检测去除异常节点并重置拓扑结构的容错处理外,参数服务器因在集群中负责参数存储、梯度下降优化与进度控制而成为重点,其稳定性至关重要。本文设计了共享内存哈希表,避免因内存紧张等原因进程崩溃后丢失本机存储的参数数据;为进一步防止参数服务器节点掉电导致整个集群的训练失败,通过一致性哈希的虚拟节点路由,将2个副本备份在2台不同的参数服务器节点上,其中一台若被检测离线,则向备份虚拟节点地址发送请求,采用定时主备参数一致性同步的方式,避免频繁同步对训练过程造成影响。

3.3 环拓扑设计

图3 Ring-AllReduce梯度同步方案Fig.3 Ring-AllReduce gradient synchronization scheme

目前环拓扑All-Reduce一般采用基于轮次阻塞的方案,如Uber开源的horovod[20]中,为了控制在多台机器上独立进程的迭代保持步调一致,会引入一个根节点作为协调者,当根节点收到所有进程完成上轮迭代的信息后,广播开始新一轮迭代。这种方法在控制同步过程的同时增加了网络通信开销、屏障阻塞中的资源浪费。

3.3.1 动态Ring-AllReduce同步方案

为解决上述问题,本文提出一种不引入根节点,通过P2P通信对分布式训练进度进行动态平衡的动态Ring-AllReduce梯度同步方案。首先通过超时重试机制保证可达性,直到收到回执(ACK),确保网络抖动丢包发生时消息也能不丢消息;其次将T轮同步迭代过程用更小粒度的2(N-1)×T轮子迭代表示,松弛约束条件使得收到的消息不必在本轮子迭代内处理,借助一个消息队列对收到的子迭代轮次领先的消息进行缓存,在开启新子迭代时优先检查队列,若可从队列里读到消息则跳过网络请求直接从消息队列中读取,分析时序依赖关系可知最快节点不会比最慢节点领先超过N-1个子迭代,实现了动态同步过程。此方法在没有轮次阻塞、不破坏同步迭代时序性的情况下,提升了环拓扑All-Reduce在同步过程中的资源利用率与鲁棒性。

3.3.2 梯度融合

本文设计了梯度融合器BufferFusion将多块梯度内存逻辑融合为一块伪连续内存,使用时将多个梯度内存指针注册到BufferFusion上后,梯度在同步过程中会被作为一块连续内存操作,批量进行压缩与网络传输,借助SIMD向量化并行累计梯度并计算均值,将多个Tensor的同步融合为一个操作,进一步降低网络传输成本并提升计算效率。

4 实验与结果分析

本文通过设计4组实验,分别测试系统在稀疏数据场景和稠密模型训练场景下单节点与多节点模型的训练性能与扩展性,以及网络抖动情况下的稳定性。设置实验环境为多台8核、8 GB内存、千兆网卡的运算设备,基于高速以太网进行本地网络连接。

4.1 单机性能对比实验

为验证轻量级模块化设计的计算性能指标,本文对比了目前主流的稀疏模型LibFM与稠密深度模型Tensorflow计算框架在单机CPU下的计算性能。图4和图5所示的实验结果表明,本文系统相比LibFM在二阶交叉隐向量维度为64时达到了15倍的性能加速,CPU上的深度模型训练相比Tensorflow也有提高,证明轻量级设计与性能优化可达预期效果。

图4 与LibFM系统在计算性能上的对比Fig.4 Comparison with LibFM system in computational performance

图5 与Tensorflow-CPU系统在计算性能上的对比

Fig.5Comparison with Tensorflow-CPU system in computational performance

4.2 Wide&Deep模型训练

Wide & Deep是用于推荐与CTR预估的算法模型,是由逻辑回归与DNN组合的二分类模型,其DNN输入为离散特征的Embedding向量,融合了高维稀疏特征与稠密类别特征。借助上文设计的参数服务器,稀疏参数与类别稠密特征都通过Pull/Push完成参数获取与梯度汇报。实验采用Avazu数据集学习点击率,设置4个参数服务器节点、4个计算节点,将训练数据拆分4个输入计算节点。单机训练全量数据与4个节点数据并行训练的损失函数与准确率对比如图6所示,可以看出,本文系统在4个计算节点上尽管梯度下降有较强抖动,但都最终达到了与单机相近的准确率与收敛效果,且训练数据吞吐量大幅提高,发挥了分布式并行的优势。

图6 参数服务器训练下的损失函数与准确率曲线

Fig.6Loss function and accuracy curves under parameter server training

4.3 CNN模型训练

卷积神经网络CNN是稠密张量组成的深度模型。实验任务设置为MNIST手写数字识别,分别采用4个与8个计算节点用Ring-AllReduce同步方式训练一个十分类CNN神经网络。模型由3个卷积层与池化层、2个全连接层组成串联而成,输出十分类概率分布。实验结果如图7和图8所示,可以看出曲线平滑且快速收敛,多节点训练加速效果明显,在4节点上达到了3.4倍的加速,而在8节点模型上达到了6倍的加速。

图7 Ring-AllReduce训练下的损失函数与准确率曲线

Fig.7Loss function and accuracy curves under Ring-AllReduce training

图8 训练加速效果

4.4 网络抖动下的系统表现

在网络抖动下的实验中,用随机自旋模拟更频繁的网络抖动,并分别对消息重发同步与Ring-AllReduce同步2种方案的同步耗时进行比较,结果如表3所示。实验结果表明,随着同步轮次与待同步的浮点数个数的增加,Ring-AllReduce同步拥有最高的同步效率,且随着同步规模增大优势更加明显。其中消息重发同步在210个浮点数时性能最差,分析原因是在实验模拟的网络抖动环境下,小规模浮点数计算耗时短,导致消息重发同步中数据包重发变得非常频繁,无效的网络通信与I/O占用成为很大的负担,而动态Ring-AllReduce同步受网络抖动与运算耗时差异影响小,在动态平衡中形成了对最慢节点的异步等待。

表3 网络抖动下的同步耗时对比Table 3 Time comparison of synchronization under network jitter

5 结束语

本文设计一种轻量级模块化的分布式机器学习系统,以满足大规模特征分布式训练高定制化、低耦合与低资源消耗的需求。实验结果表明,该系统对于稀疏或稠密模型训练均有较好的加速性能、扩展性与稳定性。笔者已证明SIMD向量指令对模型计算效率具有优化效果,并且已将研究内容以LightCTR项目的形式开源在Github上[21],后续将设计深度学习算法在GPU上的SIMT异构计算,以及模块化的多机多卡系统,进一步提升本文系统的训练加速性能。