Scalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning
The second-order optimization methods, notably the D-KFAC (Distributed
Kronecker Factored Approximate Curvature) algorithms, have gained traction on
accelerating deep neural network (DNN) training on GPU clusters. However,
existing D-KFAC algorithms require to compute and communicate a large volume of
second-order information, i.e., Kronecker factors (KFs), before preconditioning
gradients, resulting in large computation and communication overheads as well
as a high memory footprint. In this paper, we propose DP-KFAC, a novel
distributed preconditioning scheme that distributes the KF constructing tasks
at different DNN layers to different workers. DP-KFAC not only retains the
convergence property of the existing D-KFAC algorithms but also enables three
benefits: reduced computation overhead in constructing KFs, no communication of
KFs, and low memory footprint. Extensive experiments on a 64-GPU cluster show
that DP-KFAC reduces the computation overhead by 1.55x-1.65x, the communication
cost by 2.79x-3.15x, and the memory footprint by 1.14x-1.47x in each
second-order update compared to the state-of-the-art D-KFAC methods.