Source-linked AI summary
Extremely Large Minibatch SGD: Training ResNet-50 on ImageNet in 15 Minutes
Takuya Akiba, Shuji Suzuki, Keisuke Fukuda
TL;DR
Training ResNet-50 on ImageNet with very large minibatches requires methods that address accuracy and systems challenges. The paper combines optimizer and normalization changes with a distributed hardware-software system, achieving 90-epoch training in 15 minutes at comparable accuracy.
Problem
The paper addresses how to scale deep-network training to large minibatches while preventing accuracy loss and maintaining stable, practical distributed systems.
Method
The approach combines RMSprop warm-up, slow-start learning rates, batch normalization without moving averages, and Chainer-based synchronous data-parallel training.
Results
15 minutes achieved for 90-epoch ResNet-50 training on ImageNet with 1024 Tesla P100 GPUs and comparable 74.9% accuracy.
Takeaways & Limitations
ResNet-50 can be trained on ImageNet with a 32k minibatch without severely degrading accuracy.
Abstract
from arXiv · showhide
We demonstrate that training ResNet-50 on ImageNet for 90 epochs can be achieved in 15 minutes with 1024 Tesla P100 GPUs. This was made possible by using a large minibatch size of 32k. To maintain accuracy with this large minibatch size, we employed several techniques such as RMSprop warm-up, batch normalization without moving averages, and a slow-start learning rate schedule. This paper also describes the details of the hardware and software of the system used to achieve the above performance.
1 Introduction
The paper addresses the algorithmic and systems challenges of scaling ResNet-50 training to very large minibatches without losing accuracy. On ImageNet, it reports 15-minute training with 1024 GPUs and 74.9% accuracy.
- Distributed computing is needed to scale deep-network training to larger datasets and more complex models while shortening training time.
- 15 minutes achieved for 90-epoch ResNet-50 training on ImageNet with comparable 74.9% accuracy.
- The central challenge combines preventing accuracy loss at large minibatch sizes with designing stable, practical hardware-software systems.
- The reported system used 1024 Tesla P100 GPUs, Chainer, and a total minibatch size of 32,768.
2 Training Procedure for Large Minibatches
The training procedure combines optimizer warm-up, a slow-start learning-rate schedule, and revised batch normalization to address optimization and estimation problems caused by large minibatches.
- The procedure builds on an earlier training procedure and changes selected settings while referring detailed specifications to Appendix A.
- RMSprop Warm-up: RMSprop warm-up starts training with RMSprop and gradually transitions to SGD to address initial optimization difficulty.
- Slow-Start Learning Rate Schedule: The slow-start schedule uses a longer initial phase and lower initial learning rate to further address initial optimization difficulty.
- Batch Normalization without Moving Averages: Batch normalization uses only the last minibatch and all-reduce statistics across workers because moving averages became inaccurate at larger minibatch sizes.
3 Software and Hardware Systems
The system combines Chainer-based synchronous data-parallel software with optimized communication and an in-house 1024-GPU cluster interconnected by InfiniBand.
- Software: ChainerMN enables multi-node synchronous data-parallel training, using NCCL and Open MPI for communication.
- Software: Half-precision floats reduce all-reduce communication overhead, with preliminary experiments finding relatively small effects on final accuracy.
- Hardware: MN-1 contains 128 nodes, each with eight NVIDIA Tesla P100 GPUs, two eight-core Xeon processors, and 256 GB memory.
- Hardware: The nodes are interconnected by Mellanox InfiniBand FDR.
4 Experimental Results
With 1024 GPUs and a total minibatch size of 32k, the system completed 90-epoch training quickly while preserving validation accuracy comparable to prior ResNet-50 results.
- The experiment used a per-worker minibatch size of 32 and a total minibatch size of 32k across 1024 workers.
- Training Time: 897.9 ± 3.3 seconds was required for 90 epochs on 1024 GPUs, including validation after every epoch.
- Training Time: 70% scaling efficiency was achieved against a single-GPU baseline, and 80% against a single-node eight-GPU baseline.
- Figure 1 compares average all-reduce communication time with whole-iteration time across different GPU counts.
- Accuracy: 74.94% ± 0.09 top-1 single-crop validation accuracy was obtained after 90 epochs with 1024 GPUs.
A.1 RMSprop Warm-up
The update rule combines momentum SGD and RMSprop, beginning near RMSprop and smoothly transitioning toward SGD through an ELU-like schedule.
- A.1 RMSprop Warm-up: The update rule combines momentum SGD and RMSprop, with αSGD and αRMSprop controlling their balance.Setting αRMSprop = 0 gives standard momentum SGD, while setting αSGD = 0 gives RMSprop.
- A.1 RMSprop Warm-up: Training starts with RMSprop (αSGD ≈ 0) and smoothly switches to SGD (αSGD = 1).
- A.1 RMSprop Warm-up: The transition uses a function similar to an ELU activation, with αSGD increasing exponentially before reaching 1 at the βcenter-th epoch.
2. After that, it increases linearly until the βcenter + 1
The RMSprop-to-SGD transition reaches αSGD = 1 by the βcenter + 1/2βperiod-th epoch and then remains at 1; the method uses separate learning rates.
- 2. After that, it increases linearly until the βcenter + 1: αSGD becomes 1 by the 2βperiod-th epoch and remains 1 for the rest of training.The supplied passage states βcenter = 10 and βperiod = 5 throughout the experiments.
- 2. After that, it increases linearly until the βcenter + 1: The optimizer uses ηRMSprop = 0.0003 and combines it with SGD through αRMSprop = (1 − αSGD)ηRMSprop/ηSGD.
- 2. After that, it increases linearly until the βcenter + 1: The authors report that sudden optimizer switching severely impacts training, motivating a smooth transition from RMSprop to SGD.They examined linear and sigmoid transition functions in preliminary experiments.
A.2 Slow-Start Learning Rate Schedule
The paper uses a modified slow-start learning-rate schedule to address initial optimization difficulty, with a lower initial rate and a longer initial phase than the linear-rule schedule.
- A.2 Slow-Start Learning Rate Schedule: ηbase = 12.8 for 1024 workers with local batch size 32 under the linear rule.The total batch size is 1024 × 32, and ηbase follows 0.1 · btotal/256.
- A.2 Slow-Start Learning Rate Schedule: Compared with the linear-rule schedule, the modified schedule has a longer initial phase and a lower initial learning rate.
- A.2 Slow-Start Learning Rate Schedule: The slow-start SGD schedule uses 0.5 · ηbase for 40 epochs, followed by 0.075 · ηbase for 30, 0.01 · ηbase for 15, and 0.001 · ηbase for 5.