KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks
This work addresses scalability issues for researchers and practitioners training large deep neural networks, offering an incremental improvement over existing K-FAC optimizers.
The paper tackles the problem of K-FAC's large memory footprint hindering its applicability to large deep neural network models by presenting KAISA, an adaptive second-order optimizer framework that adapts memory, communication, and computation for specific models and hardware, resulting in 18.1-36.3% faster convergence across applications and up to 41.6% faster convergence under fixed memory budgets.
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.