KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks

Abstract

Kronecker-factored Approximate Curvature (K-FAC) has recently been shown to converge faster in deep neural network (DNN) training than stochastic gradient descent (SGD); however, K-FAC's larger memory footprint hinders its applicability to large models. We present KAISA, a K-FAC-enabled, Adaptable, Improved, and ScAlable second-order optimizer framework that adapts the memory footprint, communication, and computation given specific models and hardware to improve performance and increase scalability. We quantify the tradeoffs between memory and communication cost and evaluate KAISA on large models, including ResNet-50, Mask R-CNN, U-Net, and BERT, on up to 128 NVIDIA A100 GPUs. Compared to the original optimizers, KAISA converges 18.1-36.3% faster across applications with the same global batch size. Under a fixed memory budget, KAISA converges 32.5% and 41.6% faster in ResNet-50 and BERT-Large, respectively. KAISA can balance memory and communication to achieve scaling efficiency equal to or better than the baseline optimizers. KAISA is open source and available at https://github.com/gpauloski/kfac_pytorch.

arXiv commentAccepted for publication at the International Conference for High Performance Computing, Networking, Storage and Analysis (SC21)

References in corpus (21)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• TensorFlow: A system for large-scale machine learning
• Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation
• Language Models are Few-Shot Learners
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• HOGWILD!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent
• Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training
• Mixed Precision Training
• Horovod: fast and easy distributed deep learning in TensorFlow
• Optimizing Neural Networks with Kronecker-factored Approximate Curvature
• SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks
• Large Batch Optimization for Deep Learning: Training BERT in 76 minutes
• Fast linear algebra is stable
• TensorFlow-Serving: Flexible, High-Performance ML Serving
• Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis
• Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks
• Practical Quasi-Newton Methods for Training Deep Neural Networks
• Scale MLPerf-0.6 models on Google TPU-v3 Pods
• A Kronecker-factored approximate Fisher matrix for convolution layers
• Heterogeneity-Aware Cluster Scheduling Policies for Deep Learning Workloads
• Eigenvalue Corrected Noisy Natural Gradient