SC21 Proceedings

The International Conference for High Performance Computing, Networking, Storage, and Analysis

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

Authors: J. Gregory Pauloski (University of Chicago); Qi Huang (University of Texas); Lei Huang (Texas Advanced Computing Center (TACC)); Shivaram Venkataraman (University of Wisconsin, Madison); Kyle Chard and Ian Foster (University of Chicago, Argonne National Laboratory (ANL)); and Zhao Zhang (Texas Advanced Computing Center (TACC))

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.

Back to Technical Papers Archive Listing