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

文章插图
训练过程
基准测试
在基准测试阶段，Sky Computing 需要收集来自两个维度的数据：模型和设备。在模型维度，需要知道模型每一层所需的内存占用和计算量。通过结合模型的预计内存占用和设备的可用内存，可避免内存溢出；而所需计算量越大，同一设备完成该任务的时间就越久。在设备维度，需要知道设备的通讯延时、计算能力和可用内存等，受网络环境、当前运行负载等因素的影响。对于算力强、通信好但可用内存少的设备，应在内存不溢出的前提下，尽量多分配模型层（计算任务）。由于 Sky Computing 是一个负载均衡的联邦学习系统，因此我们在基准测试阶段只关心设备的机器学习的能力。通过在每个设备运行小型的机器学习测试任务，测探设备的AI计算能力。

文章插图
整体流程
分配
在决定任务分配方式时，经数学分析可知，分配方式本质上是一个NP-hard的混合整数线性规划问题。因此，在多项式时间内，我们无法得到一个最优解。而随着模型规模的不断增长，和设备数量的不断增多，计算最优解的成本显然是不可接受的。
因此，在实际情况中，我们不会直接计算求得最优解，而是尝试使用启发式算法得到近似解。在 Sky Computing 中，我们设计了一个两阶段的启发式算法：第一阶段为预分配，按照设备的实际可用内存大小进行模型的分配，并且计算每个设备实际的工作负载；第二阶段为分配调整，根据设备的负载量进行动态的调整，迭代降低整个系统的负载量。同时，为了验证 Sky Computing 的优越性，我们在实验中也设置了最优分配作为对比。

文章插图
【 Sky Computing：利用空间异构分布式计算特性加速联邦学习】实现架构
性能表现
我们在集群环境中，采用控制关键因素变量的方式，以联邦学习AI任务的forward 和 backward 的时间为指标，对Sky Computing的性能进行了验证。
实验结果
我们测试了三种分配方式（even：均匀分配，heuristic：启发式算法，optimal：最优分配）。在不同的计算资源数量规模和不同的模型大小下的表现，并记录了每次完成迭代所花费的时间。可以看到，随着设备数量的增多和模型深度的增加，我们的启发式算法的效果十分显著。在64个节点160层隐藏层的实验环境下，Sky Computing 比当前的均匀分配模型并行可加速55%。

文章插图
实验结果
其中，由于最优分配计算成本极高，在64节点时已难以计算，不适用于实际应用，仅作为小规模时的参考值。
更多特性
Sky Computing 是我们利用空间异构分布式计算特性加速联邦学习的一次成功尝试，获得了高达 55% 的性能提升。目前该项目仍处于开发阶段，未来我们将进行更加充分的实验，早日部署到实际应用中，并提供动态冗余等功能。
欢迎各位积极提issue和PR，共同为构建这一AI基础设施舔砖加瓦，解放AI生产力~
传送门
论文地址：
https://arxiv.org/abs/2202.11836
项目地址：
https://github.com/hpcaitech/SkyComputing
加入潞晨
潞晨科技的团队核心成员均来自美国加州大学伯克利分校，斯坦福大学，清华大学，北京大学，新加坡国立大学，新加坡南洋理工大学等国内外知名高校。目前，潞晨科技还在广纳英才，招聘全职/实习AI分布式系统、架构、编译器、网络、CUDA、SaaS、k8s等核心系统研发人员，开源社区运营、销售人员。
潞晨科技提供有竞争力的薪资回报，特别优秀的，还可以申请远程工作。也欢迎各位向潞晨科技引荐优秀人才，如果您推荐优秀人才成功签约潞晨科技，我们将为您提供数千元至数万元的推荐费。
工作地点：中国北京，新加坡，美国。(可相互转岗)
简历投递邮箱：hr@luchentech.com
参考链接：
https://arxiv.org/abs/2202.11836
https://36kr.com/p/1619922542065412?channel=wechat
https://pdfs.semanticscholar.org
— 完 —