[Paper Review] Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

Xianyan Jia, Shutao Song+12 more|arXiv (Cornell University)|Jul 30, 2018

Advanced Neural Network Applications24 references313 citations

TL;DR

The paper presents a scalable training system using mixed-precision with LARS, enabling 64K mini-batches and optimized all-reduce to train ImageNet models (AlexNet and ResNet-50) in minutes rather than hours, outperforming prior systems.

ABSTRACT

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.

Motivation & Objective

• Address the generalization risk of large mini-batch training while achieving high throughput.
• Develop mixed-precision training with LARS to maintain accuracy at very large mini-batch sizes.
• Design optimized all-reduce algorithms for scalable communication on thousands of GPUs.
• Demonstrate state-of-the-art training speed on ImageNet for AlexNet and ResNet-50.
• Evaluate convergence and scalability across large GPU clusters with real hardware.

Proposed method

• Introduce mixed-precision training with LARS to enable large mini-batch sizes without accuracy loss.
• Use FP16 for forward/backward passes with FP32 master weights for stable updates.
• Eliminate weight decay on BN biases and BN parameters and add an extra BN layer for AlexNet to improve convergence.
• Develop tensor fusion and a hybrid all-reduce strategy combining hierarchical and ring-based approaches for scalable gradient aggregation.
• Employ RoCEv2 and GPUDirect RDMA to reduce communication latency and improve bandwidth in a 1024- and 2048-GPU setup.

Experimental results

Research questions

• RQ1 Can mixed-precision training with LARS maintain ImageNet accuracy at mini-batch sizes up to 64K?
• RQ2 What architectural and optimization adjustments are needed to preserve convergence with extremely large mini-batches?
• RQ3 How can all-reduce strategies be optimized for high scalability on large GPU clusters?
• RQ4 What is the impact of communication optimizations on overall training time for AlexNet and ResNet-50?

Key findings

• Mixed-precision training with LARS preserves top-1 accuracy for ResNet-50 at 64K mini-batch with 90 epochs (76.2% with LARS).
• AlexNet with 64K mini-batch and targeted architectural tweaks reaches 58.8% top-1 accuracy after 95 epochs.
• The system achieves 4 minutes (AlexNet, 95 epochs) and 6.6 minutes (ResNet-50, 90 epochs) training times on 1024 and 2048 Tesla P40 GPUs respectively.
• The approach yields up to 3x and 11x speedups over NCCL-based training for AlexNet and ResNet-50 on a 1024-GPU cluster.
• Overall, the ResNet-50 training reaches 75.8% top-1 accuracy in 6.6 minutes on 2048 GPUs, and 75.8% is achieved with competitive efficiency compared to prior work.
• Scaling efficiency improvements from 9.0% to 99.2% on 1024 GPUs when using optimized all-reduce and tensor fusion.