星云Clustar首席科学家胡水海：GPU在联邦机器学习中的探索

近期，星云Clustar首席科学家胡水海，以“GPU在联邦机器学习中的探索”为题，全面详尽地讲解了目前解决联邦学习的性能与效率问题，以及解决思路。

在报告中胡水海提到，联邦学习的模型训练过程，很难绕开同态计算和密文传输，二者对算力和网络都有严苛的要求，星云Clustar也因此选择从GPU加速同态运算，以及高速网络助力密文传输效力的角度切入，来改善联邦学习的计算速度。

以下是胡水海的演讲全文：

目前在AI领域面临的一个很重大的问题，其实是数据孤岛问题。在企业层面，大部分公司在开发自己的AI模型的时候，其实并不缺少算法和应用场景，也不缺少优秀的人才，其所面临的最大问题是数据不足的问题。

每个企业都有一些自己的数据，但是这些数据彼此之间是相互割裂的，也没有一种方法将每个企业的数据高效地连通起来，所以一些小企业会面临数据不足以及大企业数据垄断问题。

另一方面，无论国内，还是国外，对数据隐私的保护都已经被重视了起来。其实，从2012年开始，国外欧盟已经在逐步起草一些法律法规来保护数据安全以及用户隐私，2018年5月份生效的GDPR更是将用户数据安全提升到了另一个高度。

另外，在今年1月份，美国加州也出台了相关法案，明确规定数据归用户所有。这些数据隐私保护的趋势都在表明：企业已经无法以明文的方式交换其拥有的数据。

而对于国内，从2009年开始也在逐步出台很多保护数据安全以及用户隐私的法案。总的来看，国内的数据法规政策有两大趋势，首先是对数据安全的保护事实上变得越来越严格，这直接体现为去年一些大数据公司在共享数据的时候，因为行为不当，受到了很严厉的法律惩罚。

另外一方面是对数据安全的保护变得越来越全面，在各个领域各个维度都出台了非常多的法律法规来保护数据隐私。所以，在上述背景下，解决数据孤岛问题其实就变得更加困难。但是联邦学习的出现为安全合规地连接数据孤岛，提供了一种非常有前景的方法。联邦学习是一项数据不出本地，就可完成机器学习多方协作建立模型的技术。换句话说这种数据不出本地的联合建模技术，正是解决国内企业数据孤岛现状的“良药”。

联邦学习与同态加密

联邦学习有很多的优点，首先能保证数据隔离，保证数据不会泄露到外部；其次联邦学习有无损的性质，保证联合建模的效果等同于直接用所有的数据进行建模的效果；再者，在联邦学习里所有数据参与方的地位都是对等的；最后，联邦学习能保证参与方共同获益，有助于打破数据巨头的垄断地位。联邦学习之所以能实现这些神奇的效果，其中有一项关键技术就是同态加密计算。

同态加密是一种特殊的非对称加密系统，一般加密后的密文是一段无法操作的二进制数，除非解密，不然不能对其进行计算或其他操作。而同态加密好的密文仍然能进行计算，得到仍然是加密的结果。

最重要的是对密文进行的计算，解密后跟对明文进行计算后的结果相同。这个特性可以让参与者进行数据运算，但无需知道参与计算的密文内容。同态计算的应用范围非常广，但是有一个显著的缺陷，就是性能太低，具体有多低后面会分析，不过有一种折中的方式勉强能被接受，那就是部分同态加密。部分同态加密分为加法同态和乘法同态。

比如Paillier算法是加法同态，只支持密文跟密文相加。而著名的RSA算法则是乘法同态，支持密文跟密文相乘。具体的原理就不详细展开了，大家可以参考相关论文。值得一提的是，同态加密虽然能够让联邦学习保护用户隐私，但它其实也为联邦学习带来了很大的技术挑战，这一点从与传统机器学习方法的对比中能够清晰看到。

首先，传统机器学习一般使用的是32-bit的基本运算，这些基本运算一般都有芯片指令的直接支持，而联邦学习中的Paillier/RSA算法依赖的是1024或2048-bit 甚至更长的大整数运算，且这些运算是模幂、模乘等复杂运算；其次，在分布式计算时，传统机器学习参数聚合使用内网传输，而联邦学习因为涉及不同的参与方，这些参与方可能位于不同的城市，所以联邦学习是使用广域网进行传输。

另一方面，从数据传输的角度来看，联邦学习对运算位数多要求1024或2048-bit ，所以传输密文数据体积比传统机器学习增加几十倍；因为联邦学习要求多次迭代，所以数据传输的次数也是传统机器学习的几倍。

总的算起来，如上图所示，联邦学习的部分同态计算的计算量是明文计算量上百倍，联邦学习的数据传输总量也比传统机器学习大100到1000倍。如果使用全同态的话，其计算量会是明文计算的上万倍。也正是基于这个原因，当前的联邦学习解决方案多采用部分同态加密。面临计算和传输方面的挑战，我们做了许多有价值的技术探索。

第一个探索是使用GPU来加速联邦学习计算。如上图，我们首先进行四个观察方向的可行性分析，第一个观察数据加解密及密态计算，不同数据的计算其实并不存在很大的关联性，因此计算是高度并行的。而GPU正好适合加速高度并行的计算任务。

第二个观察是联邦学习很多计算公式其实本身并不复杂，但重复执行次数巨大。举例而言，联邦学习需要经常运行 A的B次方这种幂计算，而A和B往往是1024比特甚至更长的数字。所以，即使是简单的公式，但是重复运算的次数非常多，而GPU正好适合加速这种重复的轻量级计算。

第三个观察是在联邦学习里，数据IO时间占比非常少，可能不到计算时间的0.1%，这说明联邦学习符合计算密集型的任务，而GPU适合加速计算密集型任务。第四个观察是联邦学习里训练模型的数据通常是以批量形式的产生为主，符合大数据的特征，而GPU正好适合加速海量数据的批量计算。 

GPU加速联邦学习计算的挑战和解决方案

上述四个观察虽然确定了GPU能够加速联邦计算的方向，但同时也提出了三个挑战。第一个挑战是联邦学习计算需要做2048-bit大整数运算，而GPU流处理器不直接支持大整数运算；第二个挑战是联邦学习计算涉及大量的模幂运算，而GPU做除法或者模幂运算的代价非常大；第三个挑战是联邦学习计算需要缓存大量中间计算结果，而由于成本和能耗的限制，GPU显存非常有限。

针对三个挑战，我们提出了三个解决方案。第一个方案是基于分治思想做元素级并行。如图所示，以计算大整数乘法a*b为例，首先我们将N比特位长的大整数a和b分解成高位和低位两部分，分解之后其a和b以及a*b的表达式如图。仔细观察a*b的表达式，发现四个子项的计算不存在数据关联性，可以并行计算。

基于这个思想，我们可以通过递归的方式将大整数乘法分解成很多可并行计算的小整数乘法，这样GPU就能发挥并行计算的优势完成大整数乘法的快速计算。不仅如此，对于联邦学习涉及的其他大整数运算，也可以做类似的元素级并行。

第二个解决方案是用平方乘算法+蒙哥马利算法解决GPU做模幂运算代价大的问题。其核心是如何高效计算模幂运算ab mod c ，其中a,b,c均为N比特大整数。对于这个问题，最容易想到的朴素算法是先计算ab的值，然后将计算结果对c取模。但这样会使问题计算复杂度高达O(2^N)，并且中间的乘积结果很大。我们采用的方法是通过平方乘算法进行优化。平方乘算法主要基于的观察是：我们要计算a^K，并不一定需要将a自乘k次，而是可以先计算出a^k/2的值，然后求平方。

以此类推，我们只需要 logk 次的乘法运算就可以得到ak的值。根据这个思想，我们可以将b表示为2进制数，然后通过O(N)次乘法以及取模运算得到计算结果。这类方法的优点是将复杂度降低到O(N)并且中间计算结果的大小不超过c。

缺点是需要做2N次取模运算，对GPU来说，做取模运算的时间代价很高。为了解决这个问题，我们引入了蒙哥马利模乘算法来高效完成第3步中的模乘计算。蒙哥马利算法的优点能够让复杂度下降到O(N)，中间结果大小不超过c，完全避免了取模/除法运算，从而大大加快了运算速度。

对于第三个挑战，如何减少中间计算结果，我们给出的解决方案是通过中国剩余定理。中国剩余定理是数论领域的一个著名定理，说的是给定一组两两互质的整数n1,n2,…,nk和任意一组整数a1,a2,…,ak，那么通过这两组数构造的下面这个同余方程组一定有解，并且解一定同余于N。

 那么怎么使用呢？首先定义mp跟mq这两个子项，并依据这两个子项构造一个满足中国剩余定理的同余方程组。如上图所示，并用CRT(mp,mq)来表示这个同余方程组的解。可以证明解密计算公式等价于同余方程组的解mod pq，所以可以通过计算这个新的表达式来求解m的值。根据上面三个计算表达式，会有两个观察结论。首先，三部分的中间计算结果都不超过N比特，因此减小了中间计算结果。此外，计算公式从2N比特数的模幂运算简化成N比特数的模幂运算，计算量大幅减小。

最后看一下GPU加速联邦学习的初步评测结果。我们主要评测了四种运算：同态加密、同态解密、密态乘法和密态加法在三种优化下的加速比。对比的baseline是14核2.2Ghz的服务器级CPU，而使用的CPU代码是高度优化的。

结果如上图：对于比较复杂的同态加密和解密，在经过三种优化手段后，GPU为联邦学习带来了差不多6倍的加速比。对于计算相对简单的密态乘法和密态加法，GPU为联邦学习分别带来了30倍以上和400倍以上的加速比。

加速联邦学习跨机构跨区域通信的探索

上面讲的是如何应对联邦学习计算方面的挑战，那么在传输方面，即在加速联邦学习跨机构跨区域通信方面，主要考虑联邦学习通信的两大场景：场景一是数据中心内部不同机构间通信（主要是云服务器），场景二是不同机构的数据中心跨区域通信（地理位置不同）。

其中，数据中心内通信场景的主要挑战是高速网络时代如何加速联邦学习通信；而跨区域通信场景的主要挑战是如何在高延迟、高丢包率网络环境下加速联邦学习通信。针对场景一带来的挑战，我们采用的解决方案是通过RDMA网络技术优化两点间通信，然后通过动态参数聚合模型优化多点间通信来解决。

在这里也提一下数据传输的背景，现在正处在数据中心高速网络时代，如上图所示，数据中心网络带宽近年来高速增长，100G，200G网络对于大规模商用数据中心来说，已经非常普遍。当然，网络带宽的高速增长也对通信带来了巨大挑战！10-100倍的带宽增长带来了三个问题，第一，收发两端相同时间需要处理10-100x的网络数据包，第二，网络突发流量现象变得更加严重，第三，网络流完成时间大大减少意味着拥塞控制需要更快响应。

传统的TCP网络由于存在CPU负载高、端处理延迟大以及吞吐量瓶颈等几个问题，不太适用于高速网络。所以在高速网络下，RDMA取代TCP已经成为了一个趋势。具体表现在：通过内核旁路以及将传输层卸载到网卡硬件上，RDMA能实现高吞吐、低时延、低CPU负载的两点间通信，非常适合用于加速联邦学习数据中心内的通信。

但是要将RDMA应用于联邦学习数据中心内通信，我们还需要解决GPU跟RDMA网卡之间高效协作的问题。我们注意到GPU与RDMA网卡之间的通信存在从GPU到内存以及从内存到网卡的多次数据拷贝。这会增大传输延迟, 降低吞吐量和浪费CPU。

为了解决这一问题，我们在联邦学习通信中引入了英伟达的GPU-Direct-RDMA 技术，实现了GPU和RDMA网卡之间的直接数据拷贝。一方面通信吞吐量从20G提升到了100G，另一方面也将传输延迟最多降低了1000倍。

最后我们评估了GRDMA能为联邦学习带来性能提升的程度，对于AlexNet和VGG16两种模型，分别测试了他们在TCP和GRDMA两种网络下的训练效率。初步的测试结果如上图显示，使用GRDMA分别带来了超过60%和超过50%的训练性能提升。

关于优化联邦学习多点间通信，Parameter Server和Ring Allreduce是目前使用最广泛的两种参数聚合模型。但他们都分别有一些缺点。ParameterServer的问题是存在多个worker节点给单个server节点发送参数的多对一通信方式。在超售网络下，这种通信方式的性能会因为链路拥塞而大幅度下降。Ring Allreduce的问题是存在一个很长的通信依赖链。一旦某一跳发生阻塞，RingAllreduce 的长依赖链会使整个聚合任务停滞。

对于跨区域通信场景问题，首先有以下几点观察，第一，随着物理距离增加，跨区域通信时间在联邦学习中的时间占比越来越大；第二，跨区域主干网具有高延迟、高丢包率等特征，丢包侦测与丢包恢复代价很大；第三，机器学习模型训练可以容忍一定的丢包率，即我们通过实验发现，当丢包率小于15%时，即使不做丢包恢复，模型收敛所需要的轮次并不会变多。另外我们还发现，当丢包率低于15%时，不做丢包重传能显著减少模型训练时间。

那么为什么机器学习模型训练可以容忍部分丢包呢？原因是目前模型训练大多采用随机梯度下降(SGD)方式通过多轮迭代进行，丢失一部分数据不影响训练算法找到模型收敛点。如图所示，蓝线是不丢包的情况下模型训练的收敛路径，而在有丢包的情况下，随机梯度下降能让模型训练选择另外一条路径达到收敛点。

基于上述观察，我们设计了一个机器学习专用的网络传输协议 --- MLT。核心思想是：在不影响模型收敛的前提下，允许一定的丢包，不做重传，从而降低跨区域通信时间。将MLT跟传统的TCP以及UDP进行对比可以发现，TCP可以看作是做百分百丢包重传的可靠传输，UDP可以看作是百分百丢包不重传的不可靠传输，而MLT位于两者之间，是根据机器学习训练的特点，选择重传一部分丢失的数据包，使丢包率控制在不影响模型收敛的范围内，并通过避免不必要的丢包重传来降低联邦学习的通信时间。

具体到实验评测如上图，MLT可以通过减少不必要的丢包重传，能够大幅缩短联邦学习模型训练的时间。

雷锋网 (公众号：雷锋网) 、雷锋网、雷锋网

。