具有分布式预处理的深度神经网络可扩展 K-FAC 训练
机器学习
2022-07-01 v1 分布式、并行与集群计算
摘要
二阶优化方法，尤其是 D-KFAC（分布式 Kronecker 因子近似曲率）算法，在加速 GPU 集群上深度神经网络（DNN）训练方面已获得关注。然而，现有 D-KFAC 算法需要在预处理梯度之前计算并通信大量二阶信息，即 Kronecker 因子（KF），导致巨大的计算与通信开销以及高内存占用。在本文中，我们提出 DP-KFAC，一种新颖的分布式预处理方案，将不同 DNN 层的 KF 构建任务分配给不同工作节点。DP-KFAC 不仅保留了现有 D-KFAC 算法的收敛性质，还带来三点益处：降低 KF 构建的计算开销、无需通信 KF、以及低内存占用。在 64-GPU 集群上的大量实验表明，与最先进的 D-KFAC 方法相比，DP-KFAC 在每次二阶更新中将计算开销降低 1.55x–1.65x，通信成本降低 2.79x–3.15x，内存占用降低 1.14x–1.47x。
引用
@article{arxiv.2206.15143,
title = {Scalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning},
author = {Lin Zhang and Shaohuai Shi and Wei Wang and Bo Li},
journal= {arXiv preprint arXiv:2206.15143},
year = {2022}
}
备注
13 pages