FLASH: 面向高性能的联邦学习硬件加速架构

第20届USENIX网络系统设计与实现研讨会USENIX Symposium on Network System Design and Implementation (NSDI) 将于2023年4月17日至19日在美国马萨诸塞州波士顿市召开。NSDI是网络系统领域最高水平学术会议之一，本次秋季投稿共288篇，接收46篇（接收率16%）。在本次会议中，香港科技大学智能网络与系统实验室iSING Lab和星云 Clustar 合作的论文《FLASH: Towards a High-performance Hardware Acceleration Architecture for Cross-silo Federated Learning》被收录，今天给大家带来这篇论文的解读。因为隐私保护技术包含高复杂度的运算，计算开销大，对联邦学习性能方面带来较大影响，该文章针对该问题，设计了一种专用的硬件架构——FLASH，对联邦学习中复杂的密码学运算提供充分的加速。

近年来，人们对数据安全以及隐私保护的关注度逐渐提升，世界各国也相继推出法律法规，比如欧盟出台的《通用数据保护条例》（GDPR）以及我国施行的《个人信息保护法》，这些法律法规限制了数据在不同持有方之间的流通。在这样的背景下，联邦学习应运而生，联邦学习推动数据可用不可见的思想，应用多种隐私保护技术，在保护数据信息的前提下实现运算，利用数据价值。但是这些隐私保护技术包含高复杂度的运算，计算开销大，对联邦学习的性能带来了较大的影响。我们设计了FLASH，一种专用的硬件架构，对联邦学习中复杂的密码学运算提供充分的加速。

论文题目：

FLASH: Towards a High-performance Hardware Acceleration Architecture for Cross-silo Federated Learning

作者：

Junxue Zhang,  Xiaodian Cheng,  Wei Wang,  Liu Yang,  Jinbin Hu,  Kai Chen

对联邦学习的常见应用场景进行分析后，我们整理出九种在联邦学习中高频出现的密码学运算，主要包括Paillier半同态加密算法、RSA加密算法等。这些运算正是我们加速的重点。然而，这些运算有着不同的运算流程，这对设计一套统一的硬件加速电路带来了较大的挑战。对于FPGA和ASIC这样基于硬件设计的加速设备，一旦完成电路逻辑的烧录或芯片的固化，电路的功能很难进一步更改。因此，为了支持九种不同的运算，一种简单的设计思路是，将系统划分为九个独立的部分，在每个部分分别构建其中一种运算的计算电路。但是由于芯片上资源有限，对资源进行分割，会导致每个部分的性能均严重不足。

因此，我们对这些密码学运算进行进一步的分析，发现（一）这九种密码学操作都有共同的算子部分（Operator）—— 模幂运算以及模乘运算，即这九种密码学的操作均可以用模幂运算和（或）模乘运算，配合一些简单的运算组合而成；（二）这九种密码学操作超过95%的运算时间均消耗在模幂运算和（或）模乘运算上。

基于这样的发现，我们提出了 FLASH，一种高效的联邦学习加速硬件体系结构。FLASH的核心思想如下：（一）将主要的硬件资源用于实现加速模幂、模乘两个核心算子的加速引擎，并通过流水线并行、提供充足的片上内存来提升引擎的性能；（二）设计数据流调度控制器，实现动态、按需地将引擎组合成上层软件需要的密码学操作。

FLASH的系统框架中包含与上位机通信的DMA模块、存储原始数据和计算结果的DDR、执行运算管理的调度模块以及执行运算的计算引擎。

图一 FLASH整体架构

计算引擎是FLASH的核心功能模块，为了在硬件上实现高效的模乘运算，我们选择了蒙哥马利算法，该算法通过将高复杂度的模运算转换为对硬件友好的移位、截位操作，从而大幅降低模乘的运算复杂度，算法的计算流程如下。

图二 蒙哥马利模乘算法

基于蒙哥马利算法，我们设计了流水的蒙哥马利乘法电路。为了充分发挥硬件电路高并行度的优势，我们在硬件上设计了300个并行的蒙哥马利模乘模块，可以同时独立地处理多次模乘运算。下图显示了多引擎之间任务流水下发、处理的流程。

图三 计算引擎间的流水和并行

显然，模幂运算可以通过反复调用模乘运算实现，因此，我们可以根据当前密码学运算的需求，通过配置模乘运算的循环次数，将计算引擎的功能在模乘和模幂之间进行切换。

此外，我们还设计了数据流的动态调度，通过调整引擎之间数据流流向，并结合引擎功能的切换，即可在硬件层面上支持不同的运算电路。如下图所示，动态调度功能的实现基于使能和关闭不同引擎之间的电路连接，实现不同的功能布局，在逻辑上将300个引擎划分到三个engine slot中，同一engine slot中的所有引擎均被配置为相同的功能（即模乘和模幂），这样就构成了三个执行模乘或者模幂运算的engine slots。通过engine slots之间的不同功能组合，我们就可以实现多种运算，比如下图(b)中的Paillier加密流程和(c)中的Paillier解密流程。engine slot中的具体引擎数量，可以根据当前密码学运算中模幂和模乘的运算量进行调整，以达到engine slots之间吞吐的平衡。

图四 FLASH的数据流调度实例

首先，我们展示实现在U250 FPGA加速卡上的FLASH原型机性能。FLASH的加速性能随着数据量的变化有所波动。对于九种基本的密码学运算，FLASH原型机可以分别最高达到8核CPU和GPU的14倍和3.4倍；对于联邦学习应用，相较于CPU和GPU，FLASH则可以最高达到6.8倍和2.0倍的加速。

图五 FLASH原型机在联邦学习应用中相较于CPU和GPU的加速

接下来，根据仿真结果，我们得到基于28nm ASIC工艺实现的FLASH，可以进一步达到FPGA原型机的7.11倍性能；基于12nm ASIC工艺实现的FLASH，则可以达到原型机的23.64倍性能。

在本文中，我们定位了联邦学习中严重降低性能的九种密码学操作，并提出FLASH，一种高效的硬件体系结构来加速该些密码学操作。我们的实验展现了FLASH在FPGA原型机上的性能优势，而基于ASIC实现的FLASH则可以达到更高的加速效果。

本文来源：FATE开源社区

往期推荐