Elamur Research⌕
☰
home / papers / 2018 / Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

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

Paper · 2018 · Xianyan Jia, Shutao Song, Wei He et al.

Summary

Training large CNNs on ImageNet can take many hours even on several GPUs. Scaling out with synchronous data-parallel SGD runs into two obstacles: very large mini-batches tend to reduce test accuracy, and gradient exchange becomes a bottleneck as the GPU count grows. The authors, from Tencent and Hong Kong Baptist University, built Jizhi, a training system for dense GPU clusters, to address both. It runs forward and backward passes in half precision but applies the LARS layer-wise learning-rate rule in FP32. It also uses several regularization tweaks that keep accuracy intact at 64K mini-batches, plus a hybrid all-reduce that combines their own adaptive collective with NCCL's ring all-reduce. With this setup, ResNet-50 reaches 75.8% top-1 in 6.6 minutes on 2048 Tesla P40 GPUs (90 epochs), and AlexNet reaches 58.7% in about 4 minutes on 1024 P40s (95 epochs). Both results beat earlier GPU and KNL systems in speed while matching or exceeding their accuracy. The work shows that extreme data parallelism does not require giving up generalization.

Key points

  • Mini-batch size was pushed to 64K for both AlexNet and ResNet-50 on ImageNet without dropping below baseline accuracy. Prior 64K attempts had reached only 73.2% and 73.9% for ResNet-50.
  • An AlexNet ablation at 64K showed that three changes together raise top-1 accuracy from 55.8% to 58.8%: skipping weight decay on biases, skipping it on batch-norm parameters, and adding batch normalization after Pool5.
  • Mixed precision raised single-node ResNet-50 throughput from 172 to 218 images/second.
  • For ResNet-50 at 64 images per GPU, scaling efficiency on 1024 GPUs rose from 9.0% before optimization to 99.2%. At 32 per GPU it reached 87.9%, compared with 80.0% reported by Akiba et al.
  • For AlexNet on 512 GPUs, efficiency rose from 70% with the FP32 all-reduce baseline to 91.4% once tensor fusion, FP16 all-reduce and hybrid all-reduce were combined. Overall speedups over NCCL-based training reached up to 3x for AlexNet and 11x for ResNet-50.

Interesting ideas

  • I1. Compute sensitive optimizer statistics in higher precision. LARS computes norm ratios that can overflow or underflow FP16, which stalls training. The fix is to cast weights and gradients to FP32 only for the LARS step and then cast them back. This pattern of keeping bulk compute in low precision while protecting fragile scalar computations fits any low-precision pipeline with adaptive or normalized updates.
  • I2. Break throughput into per-device speed × device count × scaling efficiency. Framing system speed this way separates single-GPU work, such as FP16, from communication work, such as collectives and fusion. Teams can then target and measure each factor independently when planning training infrastructure.
  • I3. Exempt biases and normalization parameters from weight decay at large batch sizes. This is a cheap change that recovered noticeable accuracy at 64K. It is worth trying whenever large-batch training underperforms.
  • I4. Run cumulative ablations of training tricks. Turning strategies on one by one and reporting accuracy for each combination shows which ones actually matter. The same applies to the communication side, where FP32, fusion, FP16 and hybrid variants are measured separately. This evaluation style transfers well to any paper combining many small optimizations.
  • I5. Mix collective algorithms instead of relying on one library. Combining a custom adaptive all-reduce with NCCL's ring version gave large gains at scale. This suggests that communication strategies should be chosen according to tensor sizes and cluster topology.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksComputer Vision and Pattern Recognition · 3,486 worksHardware and Architecture · 167 works

Cite

Cite as Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, Tiegang Chen, Guangxiao Hu, Shaohuai Shi and Xiaowen Chu. 2018. Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes. arXiv:1807.11205.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.1807.11205 · arXiv 1807.11205

@misc{jia2018highly,
  title = {{Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes}},
  author = {Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, Tiegang Chen, Guangxiao Hu, Shaohuai Shi and Xiaowen Chu},
  year = {2018},
  doi = {10.48550/arxiv.1807.11205},
  eprint = {1807.11205},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.1807.11205}
}

References

27 references, 6 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, and others. 2016. Tensorflow: a system for large-scale machine learning.. In OSDI, Vol. 16. 265–283.
  2. Takuya Akiba, Shuji Suzuki, and Keisuke Fukuda. 2017. Extremely large minibatch sgd: Training resnet-50 on imagenet in 15 minutes. arXiv preprint arXiv:1711.04325 (2017). find
  3. M. Barnett, L. Shuler, R. van de Geijn, S. Gupta, D. G. Payne, and J. Watts. 1994. Interprocessor collective communication library (InterCom). In Proceedings of IEEE Scalable High Performance Computing Conference. 357–364. https://doi.org/10.1109/SHPCC.1994.296665 — [doi:10.1109/shpcc.1994.296665](https://doi.org/10.1109/shpcc.1994.296665) ↗
  4. Minsik Cho, Ulrich Finkler, Sameer Kumar, David Kung, Vaibhav Saxena, and Dheeraj Sreedhar. 2017. PowerAI DDL. arXiv preprint arXiv:1708.02188 (2017). find
  5. Valeriu Codreanu, Damian Podareanu, and Vikram Saletore. 2017. Scale out for large minibatch SGD: Residual network training on ImageNet-1K with improved accuracy and reduced time to train. arXiv preprint arXiv:1711.04291 (2017). — [arXiv:1711.04291](https://arxiv.org/abs/1711.04291) find
  6. Matthieu Courbariaux, Y Bengio, and Jean-Pierre David. 2015. Training deep neural networks with low precision multiplications. (12 2015). find
  7. Christopher De Sa, Megan Leszczynski, Jian Zhang, Alana Marzoev, Christopher R Aberger, Kunle Olukotun, and Christopher Ré. 2018. High-Accuracy Low-Precision Training. arXiv preprint arXiv:1803.03383 (2018). — [arXiv:1803.03383](https://arxiv.org/abs/1803.03383) find
  8. Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009. Imagenet: A large-scale hierarchical image database. In Computer Vision and Pattern Recognition, 2009. CVPR 2009. IEEE Conference on. Ieee, 248–255. find
  9. Aditya Devarakonda, Maxim Naumov, and Michael Garland. 2017. AdaBatch: Adaptive Batch Sizes for Training Deep Neural Networks. arXiv preprint arXiv:1712.02029 (2017). find
  10. Andrew Gibiansky. 2017. Bringing HPC techniques to deep learning. http://research.baidu.com/bringing-hpc-techniques-deep-learning ↗
  11. Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. 2016. Deep learning. Vol. 1. MIT press Cambridge. find
  12. Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. 2017. Accurate, large minibatch SGD: training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 (2017). find
  13. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. 770–778.
  14. Sergey Ioffe and Christian Szegedy. 2015. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167 (2015). find
  15. Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. 2016. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04836 (2016). find
  16. Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. 2012. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems. 1097–1105.
  17. Anders Krogh and John A. Hertz. 1992. A Simple Weight Decay Can Improve Generalization. In Advances in Neural Information Processing Systems 4, J. E. Moody, S. J. Hanson, and R. P. Lippmann (Eds.). Morgan-Kaufmann, 950–957. http://papers.nips.cc/paper/563-a-simple-weight-decay-can-improve-generaliza… ↗
  18. Jianmin Chen Zhifeng Chen Andy Davis Jeffrey Dean Matthieu Devin Sanjay Ghemawat Geoffrey Irving Michael Isard Manjunath Kudlur Josh Levenberg Rajat Monga Sherry Moore Derek G. Murray Benoit Steiner Paul Tucker Vijay Vasudevan Pete Warden Martin Wicke Yuan Yu Xiaoqiang Zheng Martin Abadi, Paul Barha… — TensorFlow: A system for large-scale machine learning
  19. Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaev, Ganesh Venkatesh, and others. 2017. Mixed precision training. arXiv preprint arXiv:1710.03740 (2017). find
  20. Alexander Sergeev and Mike Del Balso. 2018. Horovod: fast and easy distributed deep learning in TensorFlow. arXiv preprint arXiv:1802.05799 (2018).
  21. Shaohuai Shi and Xiaowen Chu. 2017. Performance Modeling and Evaluation of Distributed Deep Learning Frameworks on GPUs. arXiv preprint arXiv:1711.05979 (2017). find
  22. Shaohuai Shi, Qiang Wang, Xiaowen Chu, and Bo Li. 2018. Modeling and Evaluation of Synchronous Stochastic Gradient Descent in Distributed Deep Learning on Multiple GPUs. arXiv preprint arXiv:1805.03812 (2018). find
  23. Samuel L Smith, Pieter-Jan Kindermans, and Quoc V Le. 2017. Don’t Decay the Learning Rate, Increase the Batch Size. arXiv preprint arXiv:1711.00489 (2017). find
  24. Rajeev Thakur, Rolf Rabenseifner, and William Gropp. 2005. Optimization of collective communication operations in MPICH. The International Journal of High Performance Computing Applications 19, 1 (2005), 49–66. find
  25. Pijika Watcharapichat, Victoria Lopez Morales, Raul Castro Fernandez, and Peter Pietzuch. 2016. Ako: Decentralised deep learning with partial gradient exchange. In Proceedings of the Seventh ACM Symposium on Cloud Computing. ACM, 84–97. find
  26. Yang You, Igor Gitman, and Boris Ginsburg. 2017. Scaling SGD Batch Size to 32K for ImageNet Training. CoRR abs/1708.03888 (2017). arXiv:1708.03888 http://arxiv.org/abs/1708.03888 — [arXiv:1708.03888](https://arxiv.org/abs/1708.03888) ↗
  27. Yang You, Zhao Zhang, C Hsieh, James Demmel, and Kurt Keutzer. 2017. ImageNet training in minutes. CoRR, abs/1709.05011 (2017).

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/jia-2018-highly-scalable-deep-learning-training/