Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

By Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuan Yang, Liwei Yu and 4 more. 2018

423 citations recorded by the source corpus.

Synchronized stochastic gradient descent (SGD) optimizers with data parallelism are widely used in training large-scale deep neural networks. Although using larger mini-batch sizes can improve the system scalability by reducing the communication-to-computation ratio, it may hurt the generalization ability of the models. To this end, we build a highly scalable deep learning training system for dense GPU clusters with three main contributions: (1) We propose a mixed-precision training method that significantly improves the training throughput of a single GPU without losing accuracy. (2) We propose an optimization approach for extremely large mini-batch size (up to 64k) that can train CNN models on the ImageNet dataset without losing accuracy. (3) We propose highly optimized all-reduce algorithms that achieve up to 3x and 11x speedup on AlexNet and ResNet-50 respectively than NCCL-based training on a cluster with 1024 Tesla P40 GPUs. On training ResNet-50 with 90 epochs, the state-of-the-art GPU-based system with 1024 Tesla P100 GPUs spent 15 minutes and achieved 74.9\% top-1 test accuracy, and another KNL-based system with 2048 Intel KNLs spent 20 minutes and achieved 75.4\% accuracy. Our training system can achieve 75.8\% top-1 test accuracy in only 6.6 minutes using 2048 Tesla P40 GPUs. When training AlexNet with 95 epochs, our system can achieve 58.7\% top-1 test accuracy within 4 minutes, which also outperforms all other existing systems.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Health Informatics, Efficiency & Compression, ML Systems & Hardware

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Health Informatics

In Efficiency & Compression

In ML Systems & Hardware

Citations in the corpus

Measuring the Algorithmic Efficiency of Neural Networks (2020, neutral): “6 [Huang et al., 2016] Jia et al., 2018] Jia, X., Song, S., He, W., Wang, Y., Rong, H., Zhou, F., Xie, L., Guo, Z., Yang, Y., Yu, L., et al. (2018 6, 10, 11, 16 [van den Oord et al., 2017] van den Oord, A., Li, Y., Babuschkin, I., Simonyan, K., Vinyals, O., Kavukcuoglu, K., van den Driessche, G.,…”
Large Batch Optimization for Deep Learning: Training BERT in 76 minutes (2019, methodology): “All the successful implementations are based on momentum SGD (He et al., 2016; Goyal et al., 2017) or L ARS optimizer (Ying et al., 2018; Jia et al., 2018; Mikami et al., 2018; You et al., 2018; Yamazaki et al., 2019).”
Flower: A Friendly Federated Learning Research Framework (2020, background): “The optimization of distributed training with and without federated concepts has been covered from many angles (Dean et al., 2012; Jia et al., 2018; Chahal et al., 2018; Sergeev & Balso, 2018; Dryden et al., 2016).”
Lookahead Optimizer: k steps forward, 1 step back (2019, methodology): “Other approaches for improving convergence on ImageNet can require hundreds of GPUs, or tricks such as ramping up the learning rate and adaptive batch-sizes [10, 16].”
Deep Leakage from Gradients (2019, background): “There are many works that aim to improve the scalability of distributed training, both at the algorithm level [7, 11, 16, 23, 32] and at the framework level [1, 2, 6, 29, 33].”
A Baseline for Few-Shot Image Classification (2019, background): “Batchnormalization [44] is used, but its parameters are excluded from weight decay [45].”
Demystifying Parallel and Distributed Deep Learning (2018, methodology): “The use of large-batch methods and reduced floating-point precision has sparked a “race” toward accelerating the training process of specific model/dataset combinations in image processing [84, 123, 176, 191, 256] and language modeling [192, 203, 257].”
On the Computational Inefficiency of Large Batch Sizes for Stochastic Gradient Descent (2018, background): “We observe markedly worse performance for these techniques in domains other than image classiﬁcation, where large batch optimization has received the most attention [15, 33].”
DAPPLE: a pipelined data parallel approach for training large models (2020, background): “Some prior studies [2, 3, 10, 28, 43, 51] focus on reducing the comm overheads for data parallelism.”
MLPerf Training Benchmark (2019, methodology): “ResNet-50 is a residual network (He et al., 2016a;b); such networks and their derivatives remain the state of the art in image classiﬁcation, and system studies commonly use them (Goyal et al., 2017; Jia et al., 2018; Mikami et al., 2018; Ying et al., 2018; Sun et al., 2019).”
Heterogeneity-Aware Asynchronous Decentralized Training (2019, methodology): “More recently, All-Reduce [41], an alternative distributed solution utilizing the advanced Ring All-Reduce algorithm [15], is shown to provide superior performance than PS [26, 31, 45, 52].”
Augment Your Batch: Improving Generalization Through Instance Repetition (2020, background): “Increasing batch size while mitigating accuracy degradation is actively researched in the ML and systems communities [8, 19, 13, 23, 25, 40].”
Communication-Efficient Decentralized Learning with Sparsification and Adaptive Peer Selection (2020, background): “, all-reduce) [13]–[15] as MPI has a long history in the HPC community for providing efficient communication primitives [16].”
Fault Tolerance in Iterative-Convergent Machine Learning (2018, background): “Reduced precision (Example 2.1) and bounded staleness consistency (Example 2.2) have already been the focus of much attention in both the ML and systems communities (Zhang et al., 2017a; Jia et al., 2018; Wei et al., 2015; Dai et al., 2015).”
BlueConnect: Decomposing all-reduce for deep learning on heterogeneous network hierarchy (2019, methodology): “Therefore, parallelization by distributing the deep learning training to many (upwards of hundreds) GPUs over a cluster or on a cloud environment is critical to cut the training time from weeks to hours and minutes (Goyal et al., 2017; Iandola et al., 2015; Jia et al., 2018; You et al., 2017a;b).”