朱传家 刘 鑫* 方佳瑞

1(江南计算技术研究所 江苏 无锡 214083)2(清华大学计算机科学与技术系 北京 100084)

0 引 言

随着机器学习技术的快速发展，深度学习已经在图像识别、语音识别、自然语言处理等领域进行了成功的应用[1]。随着训练数据规模成倍的增长和深度学习模型越来越复杂，训练这些网络非常耗时，训练时间通常几天甚至数周，如何提高训练效率成为研究的热点。分布式扩展技术是一种有效的方法。分布式扩展方式分类的方式有多种，按网络通信的内容是参数还是数据划分可以分为数据并行和模型并行；按网络通信的时序划分可以分为同步方式和异步方式。常用的分布式扩展方式有参数服务器分布式扩展方式和去中心化的分布式扩展方式两种。为了更好地使用深度学习方法，工业界推出了多款开源的深度学习框架，影响比较广泛的有伯克利大学推出的Caffe[2]、Google开发的Tensorflow[3]、微软开发的CNTK[4]、Torch[5]等框架。其中，Caffe完全由C++编写，具有代码结构清晰、模块化好、运算效率高、可移植性好的特点。

本文基于国产超级计算平台“神威太湖之光”对开源深度学习框架Caffe进行了分布式扩展研究，对比了同步方式下参数服务器分布式扩展方式和去中心化分布式扩展方式在大规模并行时的性能。本文虽基于国产超算平台，但对大规模分布式扩展有借鉴意义。

1 背景介绍

1.1 “神威太湖之光”简介

“神威太湖之光”是世界上首台峰值运算速度超过十亿亿次量级的超级计算机，也是中国第一台全部采用自主技术构建的世界第一的超级计算机[6]。“神威太湖之光”计算机系统采用基于高密度弹性超节点和高流量复合网络的高效能体系结构，由高速计算系统、辅助计算系统、高速计算互联网络、辅助计算互联网络、高速计算存储系统、辅助计算存储系统和相应的软件系统等组成。系统总体架构如图1所示。

图1 “神威太湖之光”系统架构图

系统峰值运算性能125.436 PFLOPS，实测LINPACK持续运算性能为93.015 PFLOPS，LINPACK系统效率达到74.153%，内存总容量1 024 TB，访存总带宽4 473.16 TB/s，互联网络的网络链路带宽14 GB/s，I/O聚合带宽341 GB/s，网络对分带宽70 TB/s。系统功耗15.371 MW，性能功耗比6 051.131 MFLOPS/W。系统采用40 960个SW26010异构众核处理器[7]。SW26010处理器采用片上计算阵列集群和分布式共享存储相结合的异构众核体系结构，使用64位自主神威指令集，全芯片260核心，芯片标准工作频率1.5 GHz，峰值运算速度3.168 TFLOPS。SW26010处理器的架构如图2所示。

图2 SW26010异构众核处理器架构

1.2 深度神经网络

深度神经网络是一种能够从高维数据中提取信息的神经网络，卷积神经网络(CNN)是经典的深度学习算法。

CNN由输入层、输出层和卷积层、池化层等多个隐藏层组成。输入层一般为一个二维向量。卷积层是CNN的核心，用于从输入层或较低级别的特征图提取更高层次的特征。池化层的作用是简化卷积层的输出。输出层完成对结果的预测，一般为一个Softmax函数。CNN组成如图3所示。

图3 CNN组成图

CNN需要为给定的训练数据调整和更新其卷积核参数或权重。反向传播是计算神经网络中权重梯度优化所需梯度的有效方法[8]。反向传播算法包括正向传播和反向传播。正向传播过程中顺着网络结构依次计算，最后得到网络的输出结果；反向传播过程遵从链式法则，计算损失函数相对于权重的梯度。最后采用优化算法对权重进行更新。

随机梯度下降算法(SGD)[9]具有结构简单、收敛速度快、效果好的优点，因此在深度神经网络算法中得到了广泛应用。在大数据背景下，深度神经网络的数据并行更多的是通过分布式随机梯度下降算法。对于该算法中参数更新方式的选择，目前主要有同步SGD和异步SGD两种机制。同步SGD需要利用所有节点上的参数信息，而慢节点所带来的同步等待使得数据并行时的加速比并不理想。异步SGD虽然单次训练速度快，但是其固有的随机性使得网络在训练过程中达到相同收敛点耗费的时间更长，且在训练后期可能会出现震荡现象。

随机梯度下降算法的并行算法[10]如下：

Wt+1=Wt-ηE(ΔWt)

(1)

式中：Wt为第t次迭代的权重值；η为学习率；ΔWt为第t次迭代W关于损失函数的梯度；E(ΔWt)为第t次迭代各节点W关于损失函数的梯度的均值。

1.3 Caffe深度学习框架

Caffe是伯克利视觉和学习中心(BVLC)推出的一款深度学习框架[2]。它具有以下优点：(1) 完全基于C++代码，且结构清晰，层次分明，具有非常好的效率和可移植性；(2) 采用了模块化的架构，可以实现新的数据格式、网络层、损失函数等的方便扩展；(3) 具有非常完善的开发交流社区等。

Caffe模型[11]包括Solver、Nets、Layers和Blobs。Solver是求解器，负责协调模型的优化；Nets、Layers和Blobs组成一个具体的模型。Nets是网络，通过逐层定义的方式定义，从数据输入层到损失层自上而下的定义整个模型。一次迭代可以看作Net的一次正向和反向传播。Layer是层，也是计算的基本单元，可以进行很多运算，比如卷积、池化、内积等。blobs是Caffe的标准数组结构，它提供了统一的内存接口。

BVLC版Caffe对多GPU并行有很好的支持，但不支持CPU的分布式扩展。国家超级计算无锡中心已经成功地将Caffe移植到“神威太湖之光”超级计算平台上，并针对申威异构众核处理器开发出针对卷积、矩阵乘等深度学习核心计算模块的算法库swDNN[12]。swDNN的性能相比K40M图形处理器上的cuDNN算法库具有1.91～9.75倍的双精度浮点数性能优势[6]。

2 相关工作

随着数据规模的急剧增大和模型的日趋复杂，采用分布式扩展技术降低深度神经网络的训练时间成为研究热点。Google的DistBelief[13]和微软的Adams[14]项目都是利用数据并行和模型并行训练大规模模型的分布式框架。文献[15]提出了基于GPU集群的模型并行系统，文献[16]提出了基于异步交互的参数服务器方式用于解决分布式机器学习问题。Disha Shrivastava等提出了一种基于Apache Spark平台的数据并行和模型并行混合的分布式方案。

上述的分布式扩展方案在速度和精度方面展示了比较好的性能，然而，由于它们对通信的要求很高且大部分研究基于GPU平台，因此，我们不能在国产超算平台上照搬上述的研究。我们提出了两种基于同步方式的数据并行方式进行Caffe的分布式扩展，分别是参数服务器分布式扩展方式和去中心化分布式扩展方式。

3 参数服务器分布式扩展设计

参数服务器分布式扩展方式由一个或多个参数服务器节点和多个计算节点组成。参数服务器节点(Parameter Server)存储和更新全部的网络参数W。计算节点进行网络的前向特征计算和反向梯度计算。反向传播时，每个网络层从下到上依次计算损失函数相对权重的梯度，并将所在层的梯度传递给参数服务器节点。参数服务器接收到所有层梯度信息后进行参数更新操作，并将更新后的参数传递给计算节点。参数服务器分布式扩展方式的架构示意如图4所示。

图4 参数服务器方式架构示意图

我们将一个SW26010处理器作为一个节点，采用基于MPI的并行编程语言。选用一个节点(通常为0号进程)作为参数服务器节点，其他节点(非0号进程)作为计算节点。计算节点的数量根据分布式扩展规模确定。

在一次迭代中，Caffe中算法实现流程为：

(1) 0号进程向其他进程广播网络参数。

(2) 非0号进程进行正向计算，从上向下依次计算每层的特征。

(3) 非0号进程进行反向计算。从下向上依次计算每层的网络参数相对于损失函数的梯度。对层j，计算出该层网络参数的梯度信息后，0号进程对该层梯度信息进行聚集操作。每层依次进行，直至反向传播完成。

(4) 所有层的梯度信息聚集操作完成后，0号进程利用梯度信息进行网络参数的更新。

算法语言如算法1所示。

算法1参数服务器方式分布式扩展算法

for(i=0; i

{

0号进程向其他进程广播网络参数W;

for(j=0; j

{

非0号进程 forward_cpu();

}

for(j=Net_depth; j > 0; j--)

{

非0号进程计算层j梯度ΔW

0号进程聚集ΔW

}

0号进程进行梯度平均和参数更新;

}

4 去中心化方式的分布式扩展设计

通过理论分析，我们发现第3节中参数服务器方式的分布式扩展方式具有实现简单的特点，但是也存在一些不足，具体表现在两方面：一是参数服务器节点和多个计算节点之间存在一对多或多对一的通信，容易产生通信瓶颈；二是反向传播过程中，每一层立刻对梯度信息进行聚集操作，通信延迟比较大。针对这两个问题，我们提出了去中心化方式的分布式扩展框架。

去中心化方式的分布式扩展中所有的节点都是计算节点，且节点之间不分主次。计算节点拥有完全相同的完整的初始网络参数。在网络优化时，每个计算节点进行独立的网络计算，即正向计算和反向计算。在完整的网络计算完成后，计算节点对所有网络层的梯度信息进行聚集平均操作。计算节点自主进行参数更新。

与参数服务器分布式扩展算法相同，去中心化方式的分布式扩展算法实现时，也将一个SW26010处理器作为一个节点，采用基于MPI的并行编程语言。具体的算法实现可以分为两部分，即网络参数初始化和网络参数优化。网络参数初始化完成所有节点初始网络参数的同步，整个训练过程只需要一次网络参数初始化操作；网络参数优化通过数万次的迭代优化实现。

初始化时，任意节点(通常为0号进程)进行网络参数广播，完成所有节点的网络参数同步；后面为上万次的网络参数优化迭代。在一次迭代中，去中心化方式的分布式扩展的操作为：

(1) 进程从上向下对网络中每层进行特征计算；

(2) 进程从下向上对网络中每层计算梯度ΔW；

(3) 进程计算出所有网络层ΔW后，所有进程进行梯度信息聚集平均操作，聚集平均操作完成后，所有进程拥有相同的梯度信息；

(4) 所有进程自主进行网络参数更新。

算法语言如算法2所示。

算法2去中心化分布式扩展设计

0号进程向其他进程广播网络参数W

for(i=0; i

{

for(j=0; j

{

所有进程 forward_cpu();

}

for(j=Net_depth; j > 0; j--)

{

所有进程计算层j梯度ΔW

所有进程聚集平均ΔW

}

所有进程进行参数更新操作;

}

4.1 Caffe梯度内存存放位置的重构

由于BVLC版Caffe中每个网络层梯度信息在内存中的存放地址是不连续的，因此，去中心化方式的进程梯度信息聚集操作具体实现时需要先将网络层的梯度信息拷贝到连续的地址空间再进行聚集操作。聚集操作完成后，再将梯度信息从连续的地址空间拷贝回各网络层的梯度存放地址空间。这会浪费大量的时间，特别是在迭代次数非常大时。因此，我们从Caffe框架层面对Caffe的梯度信息存放方式进行了重构，将各层分散的梯度信息内存分布方式改进为连续的内存存放方式，如图5所示。

图5 Caffe梯度信息内存重分配

4.2 梯度信息聚集算法的优化

传统的MPI_Allreduce算法是为低延迟的网络设计的，在高性能计算中，MPI_Allreduce的效率比较低。我们将一种环状的聚集算法引入到国产超级计算平台上，如图6所示。

图6 SW26010 ringAllreduce

其具体步骤为：(1) scatter-reduce；(2) allgather。在第1步中，CPU将交换数据，使得每个CPU最终都有一个最终结果的数据块。在第2步中，CPU将交换那些块，使得所有CPU最终得到完整的最后结果。

RingAllreduce和MPI_Allreduce在神威太湖之光的性能对比如图7和图8所示。可以看到，在常用的深度学习模型下，512个节点上RingAllreduce相比MPI_Allreduce的性能提升可达2.3～4.1倍的加速效果。

图7 RingAllreduce和MPI_Allreduce性能 对比图(通信数据量为60.97 MB)

图8 RingAllreduce和MPI_Allreduce性能对比图

5 实 验

我们进行了两组实验以验证分布式扩展算法的正确性(精度)和性能。

5.1 正确性验证

正确性测试主要测试分布式扩展后的训练精度是否和单节点的训练精度一致。我们选取Cifar10模型对Cifar10数据集进行训练，分别比较单节点、参数服务器分布式扩展方式、去中心化分布式扩展方式的训练精度。训练迭代次数为40 000次。具体参数及结果如表1所示。

表1 Caffe单节点及分布式训练精度表

5.2 性能测试

参数服务器分布式扩展方式和去中心化分布式扩展方式中，网络的前向计算和后行计算过程是相同的，不同的部分在于前者在每次迭代优化中需要进行一次网络参数广播操作和多次梯度信息聚集操作，后者只需要进行一次梯度信息聚集操作。因此节点间的通信时间是衡量两种分布式扩展方式性能优劣的评价标准。

我们在“神威太湖之光”超算平台上进行实验，一个SW26010处理器作为一个节点，处理器选用大共享模式。选择三个经典的深度学习模型Alexnet、GoogleNet和ResNet进行实验。数据集采用Imagenet数据集。分布式扩展规模最大为512个计算节点。测试结果如图9-图11所示。

图9 Alexnet模型下两种分布式扩展通信性能对比图

图10 Googlenet模型下两种分布式扩展通信性能对比图

图11 Resnet模型下两种分布式扩展通信性能对比图

5.3 结果分析

5.3.1正确性测试分析

参数服务器分布式扩展方式和去中心化的分布式扩展方式训练精度与单节点的训练精度相比，仅下降了0.1%和0.2%，在正常波动范围内。这说明参数服务器分布式扩展方式和去中心化的分布式扩展方式的实现是正确的，得到的结果也是正确的。

5.3.2性能测试分析

对Alexnet网络，当并行规模小于16时，参数服务器方式和去中心化方式通信时间差别不大；当并行规模大于16时，去中心化方式通信时间随着并行规模的扩大增长缓慢，参数服务器方式的聚集时间却随着并行规模的增加而急剧增长；当并行规模为256时，参数服务器方式所需通信时间是去中心化方式的10倍。

对GoogleNet网络，当并行规模小于8时，参数服务器方式和去中心化方式梯度聚集时间差别不大；当并行规模大于8时，去中心化方式梯度聚集时间随着并行规模的扩大增长缓慢，而参数服务器方式的通信时间随着并行规模的增加而急剧增长；当并行规模为32时，参数服务器方式所需通信时间是去中心化方式的77倍。

对ResNet网络，当并行规模小于8时，参数服务器方式和去中心化方式通信时间差别不大；当并行规模大于8时，去中心化方式通信时间随着并行规模的扩大增长缓慢，参数服务器方式的聚集时间却随着并行规模的增加而急剧增长；当并行规模为32时，参数服务器方式所需通信时间是去中心化方式的98倍。

去中心化的分布式扩展方式相比参数服务器方式的分布式扩展方式在节点间通信指标上可以加速最多达98倍。这个结果是合理的，因为去中心化的分布式扩展与参数服务器分布式扩展相比，在每次迭代中少了网络参数的广播操作，并将多次的梯度信息聚集操作变为一次梯度信息聚集操作，这减少了网络间通信的次数以及因为同步而引起的时延。

基于上述的设计和实验结果可得出结论，去中心化的分布式扩展方式明显优于参数服务器分布式扩展方式。因此，我们将在基于“神威太湖之光”平台的Caffe上采用去中心化的分布式扩展方式。

6 结 语

本文在“神威太湖之光”平台上对深度学习框架Caffe进行了基于同步梯度更新方式的分布式扩展方式研究。首先实现了基于同步梯度更新方式的参数服务器分布式扩展方式；其次实现了去中心化的分布式扩展方式，对Caffe梯度信息存放内存、梯度信息聚集算法等进行了改进。实验结果表明，两种分布式扩展方式均能很好地对Caffe框架进行扩展，但是去中心化的分布式扩展方式在节点间的通信效率方面更有优势，是更好的分布式扩展算法。