# Accelerating Distributed K-FAC with Smart Parallelism of Computing and Communication Tasks

arXiv: https://arxiv.org/abs/2107.06533 · HTML: https://mlsanity.com/paper/2107.06533/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 36 references that are in the map.
- First posted on arXiv 14 Jul 2021.

## Builds on (most cited first, 21 of 21)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning](https://mlsanity.com/paper/1602.07261/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [Practical Gauss-Newton Optimisation for Deep Learning](https://mlsanity.com/paper/1706.03662/) (2017)
- [A Kronecker-factored approximate Fisher matrix for convolution layers](https://mlsanity.com/paper/1602.01407/) (2016)
- [Practical Quasi-Newton Methods for Training Deep Neural Networks](https://mlsanity.com/paper/2006.08877/) (2020)
- [Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes](https://mlsanity.com/paper/1807.11205/) (2018)
- [Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters](https://mlsanity.com/paper/1706.03292/) (2017)
- [On the interplay between noise and curvature and its effect on optimization and generalization](https://mlsanity.com/paper/1906.07774/) (2019)
- [Continual Learning with Extended Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/2004.07507/) (2020)
- [Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks](https://mlsanity.com/paper/1811.12019/) (2018)
- [Towards Scalable Distributed Training of Deep Learning on Public Cloud Clusters](https://mlsanity.com/paper/2010.10458/) (2020)
- [Convolutional Neural Network Training with Distributed K-FAC](https://mlsanity.com/paper/2007.00784/) (2020)
- [Scalable and Practical Natural Gradient for Large-Scale Deep Learning](https://mlsanity.com/paper/2002.06015/) (2020)
- [Performance Modeling and Evaluation of Distributed Deep Learning Frameworks on GPUs](https://mlsanity.com/paper/1711.05979/) (2017)
- [MG-WFBP: Efficient Data Communication for Distributed Synchronous SGD Algorithms](https://mlsanity.com/paper/1811.11141/) (2018)
- [MG-WFBP: Merging Gradients Wisely for Efficient Communication in Distributed Deep Learning](https://mlsanity.com/paper/1912.09268/) (2019)

## Built on by (most cited first, 5 of 5)

- [On the Parameterization of Second-Order Optimization Effective Towards the Infinite Width](https://mlsanity.com/paper/2312.12226/) (2023)
- [Scalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning](https://mlsanity.com/paper/2206.15143/) (2022)
- [MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates](https://mlsanity.com/paper/2306.01685/) (2023)
- [Eva: A General Vectorized Approximation Framework for Second-order Optimization](https://mlsanity.com/paper/2308.02123/) (2023)
- [Jorge: Approximate Preconditioning for GPU-efficient Second-order Optimization](https://mlsanity.com/paper/2310.12298/) (2023)

## Abstract (arXiv)

Distributed training with synchronous stochastic gradient descent (SGD) on GPU clusters has been widely used to accelerate the training process of deep models. However, SGD only utilizes the first-order gradient in model parameter updates, which may take days or weeks. Recent studies have successfully exploited approximate second-order information to speed up the training process, in which the Kronecker-Factored Approximate Curvature (KFAC) emerges as one of the most efficient approximation algorithms for training deep models. Yet, when leveraging GPU clusters to train models with distributed KFAC (D-KFAC), it incurs extensive computation as well as introduces extra communications during each iteration. In this work, we propose D-KFAC (SPD-KFAC) with smart parallelism of computing and communication tasks to reduce the iteration time. Specifically, 1) we first characterize the performance bottlenecks of D-KFAC, 2) we design and implement a pipelining mechanism for Kronecker factors computation and communication with dynamic tensor fusion, and 3) we develop a load balancing placement for inverting multiple matrices on GPU clusters. We conduct real-world experiments on a 64-GPU cluster with 100Gb/s InfiniBand interconnect. Experimental results show that our proposed SPD-KFAC training scheme can achieve 10%-35% improvement over state-of-the-art algorithms.