Source-linked AI summary
Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes
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, Xiaowen Chu
TL;DR
Large-batch distributed training must balance scalability with model accuracy and communication overhead. The paper combines mixed-precision training, large-batch optimization, and optimized all-reduce algorithms, achieving fast ImageNet training with reported accuracy. Its experiments focus on AlexNet and ResNet-50.
Problem
Large mini-batches can improve distributed-training scalability but may reduce generalization, while gradient communication can bottleneck large GPU clusters.
Method
The system combines mixed-precision training, LARS-based large-batch optimization, pipelined input, and hybrid all-reduce communication.
Results
58.7% top-1 test accuracy is achieved for AlexNet in 4 minutes, and 75.8% top-1 test accuracy for ResNet-50 in 6.6 minutes.
Takeaways & Limitations
The system trains AlexNet and ResNet-50 with 64K mini-batches without loss of accuracy while improving ImageNet training speed.
Takeaways & Limitations
The experiments use AlexNet and ResNet-50 as representative CNN models, with reported baseline top-1 accuracies of 58.8% and 75.3%.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Large-batch distributed training improves throughput and scalability but creates accuracy and communication challenges. The paper addresses these challenges with optimization strategies, mixed precision, and optimized all-reduce methods.
- Motivation: Large datasets and deep networks can require days or weeks to train, motivating distributed synchronous SGD across clusters.Training ResNet-50 takes 29 hours using 8 Tesla P100 GPUs.
- Motivation: Larger mini-batches improve weak scaling by increasing system throughput and reducing the number of model updates.
- Challenges: Larger mini-batches can lower test accuracy because they create a generalization gap.With batch size 64K, ResNet-50 accuracy drops from 75.4% to 73.2% in one reported result.
- Challenges: Gradient aggregation can bottleneck distributed training as GPU counts grow, so both single-GPU throughput and scaling efficiency must improve.The system throughput is expressed as T = S · N · e.
- Contributions: The paper scales AlexNet and ResNet-50 to a 64K mini-batch without loss of accuracy using mixed-precision training with LARS and related normalization and regularization changes.
- Contributions: The system improves single-GPU performance with half-precision training and scaling efficiency with hybrid optimized all-reduce.The hybrid strategy combines adaptive all-reduce with ring-based all-reduce in NCCL.
2 RELATED WORK
Prior work advances large-batch, low-precision, and distributed training, but existing approaches face scalability, accuracy, or applicability limitations. The paper builds on these techniques for scalable ImageNet training.
- Research Landscape: Related work spans large-batch training, low-precision training, and distributed training on heterogeneous clusters.
- Large-batch Training: Large-batch methods progressively scale ResNet-50 from 8K to 32K mini-batches using learning-rate warmup, linear scaling, and LARS.Reported systems finish ResNet-50 in one hour, 50 minutes, 20 minutes, and 15 minutes under different hardware configurations.
- Large-batch Training: Dynamic mini-batch adaptation has been tested with piece-wise constant learning-rate schedules but is not easily applied to polynomial decay.
- Low-precision Training: Low-precision arithmetic can reduce time and energy costs, but round-off and quantization errors may affect SGD convergence and accuracy.
- Distributed Training: Ring all-reduce reduces communication load as node counts increase, while small tensors can reduce its bandwidth utilization.
- Distributed Training: The paper combines useful elements from prior work with additional optimizations to obtain high scalability on ImageNet for AlexNet and ResNet-50.
3 SYSTEM OVERVIEW
The distributed training system consists of input, training, and communication modules. These modules pipeline data, optimize computation and model updates, and improve scaling efficiency.
- System Modules: The system overview contains input pipeline, training, and communication modules.
- Input Pipeline: The input pipeline delivers the next step’s data before the current step finishes to reduce CPU and GPU idle time.
- Training Module: The training module handles model construction and variable management while using mixed-precision computation and LARS for model updates.
- Communication Module: The communication module uses tensor fusion and hybrid all-reduce to improve scaling efficiency.
4 SYSTEM IMPLEMENTATION AND OPTIMIZATIONS
The system combines mixed-precision computation with LARS, model-architecture changes, and communication optimizations to support large-batch ImageNet training and efficient distributed execution.
- Mixed-Precision Training with LARS: FP16 forward and backward propagation is combined with FP32 weight and gradient casts before and after applying LARS.This avoids FP16 dynamic-range problems that can make gradients vanish and stall training.
- Mixed-Precision Training with LARS: 76.2% top-1 accuracy is maintained for ResNet-50 with a 64K mini-batch using LARS with mixed-precision training.
- Improvements on Model Architecture: Leaving bias and batch-normalization parameters unregularized improves AlexNet convergence, with about 1.3% higher accuracy at the same epoch count.The affected parameters comprise only 0.02% of AlexNet parameters, and reduced L2-regularization computations slightly shorten runtime.
- Improvements on Model Architecture: AlexNet requires an additional batch-normalization layer after Pool5 because the Pool5 feature-map distribution develops larger variance and maximum values during training.The changing feature scaling makes training difficult at a 64K mini-batch, motivating the Pool5 normalization change.
- Improvements on Communication Strategies: Tensor fusion packs small gradient tensors into buffers before all-reduce, improving network bandwidth use and reducing latency from many small transfers.Fusion is triggered when the accumulated tensor size exceeds threshold θ.
- Improvements on Communication Strategies: The hierarchical all-reduce performs intra-group reduction, inter-group ring all-reduce, and intra-group broadcast, reducing steps from 2(p − 1) to 4(k − 1) + 2(p/k − 1).The group-size parameter k is tunable; the highest performance in the 1024-GPU cluster occurred at k = 16.
5 EXPERIMENTAL RESULTS
Experiments on ImageNet evaluate accuracy, convergence, single-GPU speed, and distributed scalability for AlexNet and ResNet-50. The system maintains baseline-level accuracy with 64K mini-batches while substantially reducing training time and improving scaling efficiency.
- Overall experimental results: 4 minutes with 1,024 Tesla P40 GPUs completes AlexNet ImageNet training, improving on the previous 11-minute result with a 32K mini-batch.The reported AlexNet result uses the system’s large-scale training configuration.
- Convergence Analysis: 76.2% top-1 accuracy for mixed-precision ResNet-50 training versus 76.3% for single-precision training at 90 epochs.The method keeps FP32 master weights while using FP16 tensors in forward and backward passes.
- Convergence Analysis: 71.9% top-1 accuracy with LARS improves to 76.2% after eliminating weight decay on batch-normalization parameters for ResNet-50.The combined changes meet the baseline test accuracy.
- Convergence Analysis: 58.8% top-1 accuracy at 64K mini-batch size is reached for AlexNet after omitting selected regularization and adding batch normalization after Pool5.The added optimization strategies also improve convergence speed relative to using LARS alone.
- Training Speed and Scalability: 87.9% scaling efficiency at batch size 32 per GPU exceeds the reported 80.0% efficiency for 1,024 GPUs.For AlexNet with batch size 128 per GPU, tensor fusion, FP16 all-reduce, and hybrid all-reduce raise scaling efficiency to 91.4%.
6 CONCLUSION
The system addresses communication bottlenecks and the generalization challenges of large-batch distributed training. It achieves high ImageNet accuracy for AlexNet and ResNet-50 in minutes using optimized training methods and all-reduce algorithms.
- Data communication can bottleneck distributed training, while larger mini-batches make it harder to preserve generalization ability.
- 58.7% top-1 test accuracy is achieved for AlexNet in 4 minutes using 1024 Tesla P40 GPUs.The training uses 95 epochs on ImageNet.
- 75.8% top-1 test accuracy is achieved for ResNet-50 in 6.6 minutes using 2048 Tesla P40 GPUs.The training uses 90 epochs on ImageNet and outperforms existing systems.