Accelerating Distributed K-FAC with Efficient Collective Communication and Scheduling

Kronecker-factored approximate curvature (K-FAC) has been shown to achieve faster convergence than SGD in training deep neural networks. However, existing distributed K-FAC (D-KFAC) relies on the all-reduce collective for communications and scheduling, which incurs excessive communications in each iteration. In this work, we propose a new form of D-KFAC with a reduce-based alternative to eliminate redundant communications. This poses new challenges and opportunities in that the reduce collective requires a root worker to collect the results, which considerably complicates the communication scheduling. To this end, we formulate an optimization problem that determines tensor fusion and tensor placement simultaneously aiming to minimize the training iteration time. We develop novel communication scheduling strategies and propose a placement-aware D-KFAC (PAD-KFAC) training algorithm, which further improves communication efficiency. Our experimental results on a 64-GPU cluster interconnected with 10Gb/s and 100Gb/s Ethernet show that our PAD-KFAC can achieve an average of 27% and 17% improvement over state-of-the-art D-KFAC methods, respectively.