提升12倍!香港浸会大学与MassGrid发布低带宽高效AI训练新算法

随着训练数据量的增加和深度神经网络(DNN)的日益复杂,分布式计算环境(如GPU集群)被广泛采用,以加速DNN的训练。分布式计算网络在机器学习方面的瓶颈在于节点之间的数据传输效率,那如何在这一网络下高效地进行AI训练?2018年,香港浸会大学异构计算实验室与MassGrid合作,通过研究提出一种可用于低带宽网络的全局Top-k稀疏化的分布式同步SGD算法,并通过实验论证出在低带宽网络下也能高效进行AI训练。目前实验结果及相关论文已被ICDCS workshop收录。数据并行的分布式同步随机梯度下降(S-SGD)方法是训练大规模神经网络常用的优化器之一。与单节点的SGD相比,S-SGD将工作负载分配给多个计算节点以加速训练,但它也引入了在每次迭代中交换模型参数或梯度的通信开销。举例说明假设有P个节点用S-SGD训练DNN模型。在每次迭代中,所有计算节点都会采用不同的小批量(mini-batch)数据来并行计算模型的梯度。然后,对每个节点的梯度进行平均后来更新模型,这便引入较大的通信开销。由于计算节点的加速器(如GPU和TPU)的计算能力比网络速度的增长快得多,网络通信性能通常成为训练的性能瓶颈,