KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks
KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks
复制标题
DOI：
10.1145/3458817.3476152
复制
发表时间：
2021-07
期刊：
影响因子：
--
通讯作者：
J. G. Pauloski;Qi Huang;Lei Huang;S. Venkataraman;K. Chard;Ian T. Foster;Zhao Zhang
中科院分区：
文献类型：
--
作者：
J. G. Pauloski;Qi Huang;Lei Huang;S. Venkataraman;K. Chard;Ian T. Foster;Zhao Zhang
kronecker因子近似曲率（K-FAC）最近被证明在深度神经网络（DNN）训练中比随机梯度下降（SGD）收敛得更快；然而，K-FAC较大的内存占用阻碍了它对大型模型的适用性。我们提出了KAISA，一个支持k - facc的、可适应的、改进的和可扩展的二阶优化器框架，它可以适应给定的特定模型和硬件的内存占用、通信和计算，以提高性能和增加可扩展性。我们量化了内存和通信成本之间的权衡，并在大型模型上评估了KAISA，包括ResNet-50， Mask R-CNN， U-Net和BERT，在多达128个NVIDIA A100 gpu上。与最初的优化器相比，KAISA在具有相同全局批处理大小的应用程序之间的收敛速度提高了18.1-36.3%。在固定内存预算下，KAISA在ResNet-50和BERT-Large上的收敛速度分别提高了32.5%和41.6%。KAISA可以平衡内存和通信，以实现等于或优于基线优化器的扩展效率。
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.