Scalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning

TL;DR

This paper introduces DP-KFAC, a distributed preconditioning method for second-order DNN training that reduces computation, communication, and memory costs while maintaining convergence, demonstrated on a 64-GPU cluster.

Contribution

DP-KFAC distributes Kronecker factor construction across workers, significantly improving efficiency over existing D-KFAC algorithms.

Findings

Reduces computation overhead by up to 1.65x.

Decreases communication costs by up to 3.15x.

Lowers memory footprint in second-order updates.

Abstract

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.