Deep Neural Network Training With Distributed K-FAC

doi:10.1109/TPDS.2022.3161187

Deep Neural Network Training With Distributed K-FAC

2022-12-01

2

PRE

AI

DOI：10.1109/TPDS.2022.3161187

Original

Original request for help

Save

Abstract

En 中文

Scaling deep neural network training to more processors and larger batch sizes is key to reducing end-to-end training time; yet, maintaining comparable convergence and hardware utilization at larger scales is challenging. Increases in training scales have enabled natural gradient optimization methods as a reasonable alternative to stochastic gradient descent and variants thereof. Kronecker-factored Approximate Curvature (K-FAC), a natural gradient method, preconditions gradients with an efficient approximation of the Fisher Information Matrix to improve per-iteration progress when optimizing an objective function. Here we propose a scalable K-FAC algorithm and investigate K-FAC's applicability in large-scale deep neural network training. Specifically, we explore layer-wise distribution strategies, inverse-free second-order gradient evaluation, and dynamic K-FAC update decoupling, with the goal of preserving convergence while minimizing training time. We evaluate the convergence and scaling properties of our K-FAC gradient preconditioner, for image classification, object detection, and language modeling applications. In all applications, our implementation converges to baseline performance targets in 9-25% less time than the standard first-order optimizers on GPU clusters across a variety of scales.

Keywords:

Training

Parallel processing

Program processors

Convergence

Computational modeling

Data models

Deep learning

Optimization methods

neural networks

scalability

high-performance computing

Deep Neural Network Training With Distributed K-FAC

Abstract

Journal

Organization