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.

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

References in corpus (7)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation
• Language Models are Few-Shot Learners
• SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks
• TensorFlow-Serving: Flexible, High-Performance ML Serving
• Scale MLPerf-0.6 models on Google TPU-v3 Pods
• Heterogeneity-Aware Cluster Scheduling Policies for Deep Learning Workloads