GPU 在联邦机器学习中的探索

联邦学习是机器学习未来重要的发展方向之一，为了完成隐私保护下的机器学习，联邦学习使用了很多与传统分布式人工智能系统不一样的设计方法，同时也面临着更为复杂的学习环境。其中，联邦学习的模型训练过程很难绕开同态计算和密文传输，两者对算力和网络都提出了很高的要求。

本文根据星云 Clustar 首席科学家胡水海在 NVIDIA 2019 年 GTC 大会上的讲演《GPU 在联邦机器学习中的探索》整理而成，阐述了联邦学习密态计算和密文传输的性能与效率瓶颈，并提出相应的解决之道。

一、联邦学习背景

在人工智能领域，传统数据处理模式往往是一方收集数据，再转移到另一方进行处理、清洗并建模，最后把模型卖给第三方。但随着法规的完善和监控愈加严格，数据离开收集方或者用户不清楚模型的具体用途都可能会导致运营者触犯法律。为应对隐私泄漏风险，世界各国都采取了相应的措施。2018 年欧盟出台了首个关于数据隐私保护的法案《通用数据保护条例》。国内从 2009 年开始，出台了一系列隐私保护法案，数据隐私的保护正逐步严格化和全面化。

一般来说，人工智能需要通过大量的数据学习才能把数据后面的知识和价值挖掘出来。但现实的情况是一方面很多数据质量不好，缺乏标签；另一方面，数据完全分散在各个数据主体，企业的个案里面，是一个个数据孤岛。但使用传统的方法粗暴地将数据聚合是法律法规所禁止的。如何打破数据孤岛，打通人工智能应用的最后一公里，促进人工智能落地呢？联邦学习给我们提供了解决方案。

1.1. 联邦学习的核心原理

根据 2020 年 4 月最新发布的《联邦学习白皮书》，联邦学习的概念可以定义为：在进行机器学习的过程中，各参与方可以借助其他方数据进行联合建模。各方无需共享数据资源，即数据不出本地的情况下，进行数据联合训练，建立共享的机器学习模型。联邦学习可以分为三类：

1.横向联邦学习

在两个数据集的用户特征重叠较多而用户重叠较少的情况下，我们把数据集按照横向（即用户维度）切分，并取出双方用户特征相同而用户不完全相同的那部分数据进行训练。

2.纵向联邦学习

在两个数据集的用户重叠较多而用户特征重叠较少的情况下，我们把数据集按照纵向（即特征维度）切分，并取出双方用户相同而用户特征不完全相同的那部分数据进行训练。

3.联邦迁移学习

在两个数据集的用户与用户特征重叠都较少的情况下，我们不对数据进行切分，而可以利用迁移学习来克服数据或者标签不足的情况。

联邦学习实现了各个企业自有数据不出本地，而是通过加密机制下的参数交换方式共建模型。由于数据本身不移动，因此也不会涉及隐私泄漏和数据合规的问题。这样，建好的模型将在各自的区域仅为本地的目标服务。在这样的一个联邦机制下，各方可以在不披露底层数据和底层数据的加密（混淆）形态下共建模型，各个参与者的身份和地位相同，这就是为什么该体系叫做联邦学习。

1.2. 联邦学习中的技术挑战

联邦学习为了完成隐私保护下的机器学习，使用了很多与传统机器学习不一样的方法。首先，传统机器学习一般使用的是 32-bit 的基本运算，这些基本运算一般都有芯片指令的直接支持，而联邦学习中的 Paillier/RSA 算法依赖的是 1024-bit 或 2048-bit 甚至更长的大整数运算，而且大多是模幂，模乘这样的复杂算子。其次，传统的分布式机器学习的参数聚合使用内网传输，而联邦学习的参数聚合使用广域网进行传输，且密文数据体积要增加几十倍，传输的次数也是传统机器学习的几倍。

目前联邦学习同态加密采用的是 Paillier 加密，这是一种部分同态加密算法，它支持的计算类型有：

1.密文 + 密文

2.密文 + 明文

3.密文 * 明文

之所以没有采用更方便的全同态方案，那是因为全同态的计算量会是明文计算的 10W 倍，而部分同态计算的计算量是明文计算量的几千倍。

为了应对密态环境下联邦学习模型训练中的挑战，星云 Clustar 给出了使用 GPU 加速联邦学习计算，以及使用 RDMA 加速联邦学习通信的一些探索。

二、GPU 联邦学习加速

2.1. 密态计算加速

联邦学习密态环境下的计算存在以下四种特点：

1.数据加解密及密态计算，不同数据的计算互不影响（计算高度并行）

2.计算公式本身不复杂，但重复执行次数巨大（重复的轻量级运算）

3.联邦学习中，数据 IO 时间不到计算时间的 0.1%（计算密集型任务）

4.数据以批量形式产生，并且数据量巨大（批量大数据）

虽然在联邦学习的密态计算上，GPU 拥有四大优势：加速高度并行的计算任务、加速重复的轻量级计算任务、加速密集型计算任务及加速海量数据批量计算任务。然而联邦学习计算需进一步处理 2048-bit 的大整数运算、大量的模幂运算以及缓存大量中间计算结果，此时 GPU 表现就面临一些挑战：

1.2048-bit 大整数运算。例如采用 1024-bit 的 Paillier 同态加密，传统类型数据的体积会膨胀到 2048-bit，而 GPU 流处理器是不支持大整数运算的。

2.联邦学习计算涉及到大量的模幂计算，而 GPU 做模幂运算的代价是非常大的。

3.联邦学习需要缓存大量中间计算结果。但由于成本和功耗的限制，GPU 的显存是非常有限的。

为了应对这些挑战，星云 Clustar 设计了一些方案来最大化 GPU 对联邦学习的加速。

2.1.1. 基于分治思想做元素间并行

以计算大整数乘法 ab 为例来介绍如何基于分治的思想来做元素级的并行。 首先将 N 比特位长的大整数 a 和 b 分解成高位和低位两部分，那么 ab 就可以表示为下图包含 4 个互不关联子项的一个表达式。

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

2.1.2. 平方乘算法+蒙哥马利算法

这一部分想要解决的问题是如何高效地进行模幂运算。比较朴素的方法是先计算，再将计算的结果对 c 取模。这样存在的问题是计算复杂度高达，并且中间的乘积结果很大。

这里优化的思路为平方乘+蒙哥马利算法。平方乘算法主要基于这样一个观察：要计算的值，并不一定需要将 a 自乘 k 次，而是可以先计算出的值，然后求平方。以此类推，我们只需要 次的乘法运算就可以得到的值。根据这个思想，可以将 b 表示为 2 进制数，然后通过次乘法以及取模运算得到计算结果。平方乘算法的优点是将复杂度降低到并且中间计算结果的大小不超过 c。缺点是需要做 2N 次取模运算，但对 GPU 来说做取模运算的时间代价很高。为了解决这个问题，这里引入了蒙哥马利模乘算法来高效完成上图第 3 步中的模乘计算。蒙哥马利算法的优点是计算模乘的过程中不需要进行取模运算，从而大大加快了运算速度。

2.1.3. 中国剩余定理减小中间计算结果

下图是联邦学习解密运算的计算公式。可以观察到这个计算公式的最终计算结果长度为 N 比特。但是中间计算结果长度为 2N 比特，因此就需要 2 倍的显存进行存储。由于 GPU 上高速显存是非常稀缺的，这一点对计算性能影响很大。

在介绍如何使用中国剩余定理减小计算结果之前，先简单介绍一下中国剩余定理。中国剩余定理是数论领域的一个著名定理，说的是给定一组两两互质的整数 n1, n2, n3……nk 和任意一组整数 a1, a2, a3……ak, 那么通过这两组数构造下面这个同余方程组一定有解，并且解一定同余于 N。

下面简要介绍下如何通过中国剩余定理分解解密计算，从而减小中间计算结果。首先，定义跟这两个子项，并依据这两个子项构造一个满足中国剩余定理的同余方程组。这里用 CRT(, )来表示这个同余方程组的解。可以证明解密计算公式等价于同余方程组的解 ，所以可以通过计算这个新的表达式来求解 m 的值。

对上图这三个计算表达式，这里有两个观察：首先，三部分的中间计算结果都不超过 N 比特，因此减小了中间计算结果。此外，计算公式从 2N 比特数的模幂运算简化成 N 比特数的模幂运算，计算量大幅减小。

2.1.4. 测评结果

接下来，可以看一下 GPU 加速联邦学习的初步评测结果。目前主要评测了四种运算：同态加密、同态解密、密态乘法和密态加法。对比的 baseline 是 14 核 2.2GHz 的服务器级 CPU，使用的 CPU 代码也是高度优化的。

可以看到，对于比较复杂的同态加密和解密，在经过三种优化手段后，GPU 为联邦学习带来了差不多 6 倍的加速比。对于计算相对简单的密态乘法和密态加法，GPU 为联邦学习分别带来了 30 倍以上和 400 倍以上的加速比。这里展示的只是初步优化后的加速结果，相信在使用 GPU 加速联邦学习计算上，还有更大的提升空间等着我们大家一起来探索 。

2.2. 密态网络通信加速

联邦学习有以下两大通信场景：

1.数据中心内部不同机构间通信

2.不同机构的数据中心跨区域通信

数据中心内通信面临的主要挑战是高速网络时代如何加速联邦学习通信。跨区域通信场景主要的挑战是如何在高延迟、高丢包率网络环境下加速联邦学习通信。

2.2.1. GDR 技术

对于第一个挑战，我们提出了通过 RDMA 网络技术优化两点间通信，然后通过动态参数聚合模型优化多点间通信。

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

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

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

为了评估 GDR 能为联邦学习带来多大的性能提升，对于 AlexNet 和 VGG16 两种模型，我们分别测试了它们在 TCP 和 GDR 两种网络下的训练效率：

上图的测试结果表示，使用 GDR 为两种模型的训练分别带来了超过 60%和超过 50%的训练性能的提升。

2.2.2. 动态参数聚合模型

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

在联邦学习场景下，多方之间网络状况更为复杂，因此前面提到的两种静态模型都无法很好优化多方间通信。我们的核心想法是根据网络当前状态，包括物理拓扑、实时流量负载等，生成最合适的参数聚合方案。

具体来说，我们动态生成的参数聚合模型会将训练节点分成多个组，每个组构成一个比较小的树结构。然后不同树结构再通过一个环连接起来。对比 Parameter Server，我们模型的多对一通信更轻量，因此不容易发生拥塞。而对比 Ring Allreduce，我们模型的通信依赖链要短很多。

2.2.3. MLT 传输协议

对于第二个挑战，这里星云 Clustar 目前在探索设计一种机器学习专用的网络传输协议 MLT。

我们对联邦学习跨区域通信有以下几点观察：

1.随着物理距离增加，跨区域通信时间在联邦学习中的时间占比越来越大

2.跨区域主干网具有高延迟、高丢包率等特征，丢包侦测与丢包恢复代价很大

3.机器学习模型训练可以容忍一定的丢包率

4.当丢包率小于 15%时，即使不做丢包恢复，模型收敛所需要的轮次并不会变多

5.当丢包率低于 15%时，不做丢包重传能显著减少模型训练时间

我们做了一些初步的实验来评测 MLT 的性能，对比 TCP 我们发现 MLT 通过减少不必要的丢包重传，能够大幅度缩短联邦学习模型训练的时间。

三、小结

联邦学习可以看作是一种安全计算和人工智能的结合体，它的重要价值在于打破数据孤岛，通过鼓励具有相同数据结构或不同数据结构的各方共同参与训练，提高模型的整体效果。作为一种具有前景的新型领域，联邦学习还存在很多值得去研究和优化的点等待各位去探索。

作者介绍：

胡水海，星云 Clustar 首席科学家，星云 Clustar 新技术研发负责人，香港科技大学博士，数据中心网络、云计算、RDMA 网络、分布式 AI 系统方向专家，在顶级网络学术会议及期刊（SIGCOMM、NSDI、CoNEXT、ToN 等）上发表过多篇论文。

目前负责星云 Clustar 产品新技术研究及创新等工作，带领团队进行星云深度安全 AI 处理器产品的研发；长期聚焦于人工智能在制造、金融、保险、银行等多个领域的技术研究，致力于将联邦学习技术推向 AI 商用落地。

本文第二作者为杨林彬。