KAISA：一种用于深度神经网络的自适应二阶优化器框架
机器学习
2021-09-21 v2 分布式、并行与集群计算
摘要
Kronecker 因子近似曲率（K-FAC）近来被证明在深度神经网络（DNN）训练中比随机梯度下降（SGD）收敛更快；然而，K-FAC 较大的内存占用阻碍了其在大模型上的适用性。我们提出 KAISA，一个支持 K-FAC 的、自适应、改进且可扩展的二阶优化器框架，它针对特定模型和硬件调整内存占用、通信与计算，以提升性能并增强可扩展性。我们量化了内存与通信成本之间的权衡，并在大模型上评估 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 能平衡内存与通信，实现等于或优于基线优化器的扩展效率。KAISA 是开源的，可在 https://github.com/gpauloski/kfac_pytorch 获取。
引用
@article{arxiv.2107.01739,
title = {KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks},
author = {J. Gregory Pauloski and Qi Huang and Lei Huang and Shivaram Venkataraman and Kyle Chard and Ian Foster and Zhao Zhang},
journal= {arXiv preprint arXiv:2107.01739},
year = {2021}
}
备注
Accepted for publication at the International Conference for High Performance Computing, Networking, Storage and Analysis (SC21)