Sky Computing：利用空间异构分布式计算特性加速联邦学习

在AI浪潮中，无论是企业还是国家，对算力的需求都日益高涨。近期启动的“东数西算”项目，更是从宏观层面大力打造AI基础设施。但位于不同地理位置的计算机之间通信延迟较高，如何统筹兼顾、高效利用不同地区的计算能力，是当下亟待解决的重大议题。
与此同时，在大数据时代的背景下，如何保护隐私数据也成为社会热点，国家出台了数据安全、隐私保护的一系列法规。
针对以上难点，Sky Computing成功利用空间异构分布式计算特性，在保证用户数据隐私的前提下，可对联邦学习加速达55%。
开源地址：
https://github.com/hpcaitech/SkyComputing

文章插图
“东数西算”项目布局
空间异构分布式计算
随着深度学习的不断发展，模型的尺寸日益增长，目前的主流模型，例如 BERT 和 GPT-3 都有着数以亿计的参数。尽管这些模型在预测精度和性能提升方面有了长足的进步，但同样也给存储和运算等带来了极大的压力。为了加速AI模型训练的速度，分布式机器学习得以应运而生，它通常使用大量高速互联的同类型处理器，如超级计算机。

文章插图
超级计算机
而空间异构分布式计算则进一步将拥有不同计算能力、通讯能力的计算资源组合在一起，作为一个大的集群完成大型计算任务。其中参与计算的硬件资源可以是大型专业计算服务器，也可以是小型的智能设备。目前，空间异构分布式计算作为一种新形式的异构计算，正在得到越来越多的关注。以我国为例，随着「东数西算」工作的推行，越来越多的计算资源将广泛地分布到西部各个地区，如何协调这类混合计算集群联合高效工作，也将成为高性能计算应用的研究热点。
近年来，云服务的规模、范围和对象都被不断扩展，越来越多的企业选择将自己的数据存储和数据计算相关业务部署在云端。然而，将所有服务依托于云端环境的缺点在于数据的迁移成本极高；同时，数据的隐私性和可靠性也难以保证；此外，分布在不同地区的云算力之间高昂的通信成本，也使得他们难以有效联合完成高算力任务。

文章插图
云计算
联邦学习
为保护数据的隐私性，Google 于 2016 年提出联邦学习，这是一种加密的分布式机器学习技术。顾名思义，它通过搭建一个虚拟的「联邦」，将大大小小的数据孤岛联合到一起。每一个数据孤岛都像是这个「联邦」中的一个州，既保持一定的独立自主（比如商业机密，用户隐私），又能在数据不被对外共享的前提下共同建模，提升AI模型效果。目前，联邦学习广泛被运用在智能终端的模型训练中，如各个语音助手例如Siri, Alex等等。

文章插图
联邦学习
在现有的联邦学习模型并行中，模型被均匀分配给各个训练设备。然而，如前文所述，由于联邦学习的训练设备往往是用户的智能终端，性能差异较大，使用均匀分配，往往会造成通信时间瓶颈。
正如我们都知道木桶效应：木桶的盛水量由最短的那块木板决定。而在传统的联邦学习中，存在类似现象：训练速度由最慢的那个设备决定。
例如，对于处于使用模型并行的同一个联邦学习任务中的智能手机和树莓派，它们会被分配相同的任务量。但由于智能手机的运算能力远超树莓派，智能手机被迫闲置等待树莓派的任务完成。

文章插图
木桶效应
Sky Computing
Sky Computing 针对以上痛点，通过负载均衡，将不同规模和能力的云服务器智能互联，达到大规模计算的算力需求，同时通过联邦学习的方式，仅在云服务器内部访问用户数据，避免数据迁移和隐私泄露。
负载均衡
要解决负载均衡的问题，首先要了解什么是「负载」。在计算机中，无论进行哪种操作，究其本质，负载都可以理解为「完成任务所需的时间」。由于在联邦学习中，训练模型的计算总量是固定的，因此如果我们能通过自适应的方式智能分配计算任务，便能够使得每个设备完成计算任务的耗时相同，确保整体训练的时间最优。而为了得到一个好的分配方式，我们需要首先得到模型和设备相关信息，然后再进行实际的适当分配操作。因此，对于训练模型，我们需要分为两个阶段：基准测试和分配。