Source-linked AI summary
Yet Another Accelerated SGD: ResNet-50 Training on ImageNet in 74.7 seconds
Masafumi Yamazaki, Akihiko Kasagi, Akihiro Tabuchi, Takumi Honda, Masahiro Miwa, Naoto Fukumoto, Tsuguchika Tabaru, Atsushi Ike, Kohta Nakashima
TL;DR
Large-mini-batch distributed training offers a path to faster deep learning but makes scalability and validation accuracy difficult to maintain. The paper combines accuracy, framework, and communication optimizations to address these challenges. On ResNet-50/ImageNet, it reports 74.7-second training with 75.08% top-1 validation accuracy using 2,048 GPUs.
Problem
Large-mini-batch distributed deep learning must scale across large clusters without compromising validation accuracy.
Source: §I · p. 1Method
The paper combines optimization techniques for accuracy, framework throughput, and communication in an MXNet-based distributed training system.
Source: §III · p. 1, A. Accuracy Improvement · p. 1, C. Communication Optimizations · p. 2Results
74.7 seconds and 75.08% top-1 validation accuracy were achieved for ResNet-50 on ImageNet using 2,048 GPUs and an 81,920 mini-batch.
Source: §I · p. 1, §IV · p. 3Takeaways & Limitations
The reported techniques enable large-mini-batch ResNet-50 training with high throughput and validation accuracy on a large GPU cluster.
Source: §IV · p. 3, §V · p. 3
Abstract
from arXiv · showhide
There has been a strong demand for algorithms that can execute machine learning as faster as possible and the speed of deep learning has accelerated by 30 times only in the past two years. Distributed deep learning using the large mini-batch is a key technology to address the demand and is a great challenge as it is difficult to achieve high scalability on large clusters without compromising accuracy. In this paper, we introduce optimization methods which we applied to this challenge. We achieved the training time of 74.7 seconds using 2,048 GPUs on ABCI cluster applying these methods. The training throughput is over 1.73 million images/sec and the top-1 validation accuracy is 75.08%.
I. INTRODUCTION
The paper addresses the communication and accuracy challenges of scaling distributed deep learning with large mini-batches. It combines optimization techniques and reports fast, accurate ResNet-50 training on ImageNet.
I. INTRODUCTION
- Distributed data-parallel training reduces computation time but incurs increasing allreduce communication overhead on large clusters.Each process computes different mini-batches, combines gradients, and updates shared weights.Source: §I · p. 1
- Large mini-batches reduce the number of SGD updates and generally worsen validation accuracy, motivating techniques that preserve accuracy at larger scales.The paper frames large-mini-batch training as a central challenge for scalable distributed learning.Source: §I · p. 1, §II · p. 1
- 75.08% validation accuracy was achieved for ResNet-50 on ImageNet with an 81,920 mini-batch in 74.7 seconds using 2,048 GPUs.The experiment used the ABCI cluster and a self-optimized MXNet framework.Source: §I · p. 1
II. RELATED WORKS
- Prior work applied warmup, variable mini-batches, and LARS to improve large-mini-batch training, while earlier systems reported faster training with thousands of processors.Reported prior results include 15 minutes with 1,024 GPUs, 6.6 minutes with 2,048 GPUs, and 1.05 million images/sec with 1,024 TPU v3 processors.Source: §II · p. 1
III. OUR APPROACH
- The paper introduces techniques for improving both validation accuracy and training throughput in large-scale distributed learning.The approach is implemented in an MXNet-based deep learning framework.Source: §III · p. 1
A. Accuracy Improvement
The accuracy-improvement approach stabilizes optimization and adapts training behavior for very large mini-batches. It combines learning-rate control with label smoothing and batch-normalization tuning.
1) Learning Rate Control:
- Warmup gradually raises the learning rate, while LARS adjusts each layer's rate according to its weight and gradient norms.These techniques address instability caused by high learning rates and unequal layer behavior during large-mini-batch training.Source: 1) Learning Rate Control · pp. 1–2
- Optimized learning-rate decay patterns were selected from step, polynomial, linear, and other schedules through repeated trials.Source: 1) Learning Rate Control · p. 2
2) Other techniques:
- Label smoothing improved accuracy for training with an 81,920 mini-batch.The passage connects this technique to previously reported gains with 32,768 mini-batches.Source: 2) Other techniques · p. 2
- Batch-normalization moving averages were tuned because independently computed per-process statistics become inaccurate during large-mini-batch training.Weights are synchronized across processes, but the moving averages are computed independently.Source: 2) Other techniques · p. 2
B. Framework optimizations
The framework optimizations reduce startup, norm-computation, and communication overheads that become important on large GPU clusters. They parallelize initialization and norm calculations while aggregating gradients into larger allreduce operations.
B. Framework optimizations
- The framework optimizations were motivated by bottlenecks that become visible in massive clusters despite occupying little time on smaller systems.CPU and GPU profilers were used to identify these bottlenecks in MXNet.Source: B. Framework optimizations · p. 2
1) Parallel DNN model initialization:
- Identical random seeds let every process initialize model weights in parallel, avoiding the broadcast operation from a root process.This synchronizes initial weights without incurring broadcast costs that increase with process count.Source: 1) Parallel DNN model initialization · p. 2
2) Batched norm computations on the GPU:
- A specialized GPU kernel computes layer norms in parallel by launching enough threads to occupy CUDA cores.The kernel addresses poor utilization when individual ResNet-50 layers contain too few weights for the available GPU cores.Source: 2) Batched norm computations on the GPU · p. 2
C. Communication Optimizations
- Allreduce communication becomes a significant bottleneck on large clusters because communication time grows while per-GPU calculation time shrinks.Distributed training requires exchanging each layer's gradients among all processes.Source: C. Communication Optimizations · p. 2
- Gradients from multiple layers were gathered into several-megabyte allreduce operations to reduce frequent communication calls and improve network-bandwidth use.Per-layer allreduce is especially inefficient when gradient tensors are small.Source: 1) Adjusting data size of communication · p. 2
2) Optimal scheduling of communications:
The framework overlaps grouped gradient communication with backward processing and evaluates large-scale ResNet-50 training on ABCI. It achieves 74.7-second training with 75.08% validation accuracy, 1.7 million images per second, and 77.0% scalability on 2,048 GPUs.
2) Optimal scheduling of communications:
- Allreduce starts for completed layer groups once gradient data exceed a threshold, overlapping communication with backward processing.Statistical grouping avoids the overhead of using allgather to identify commonly completed layers.Source: 2) Optimal scheduling of communications · p. 2
IV. ENVIRONMENT AND EXPERIMENTAL RESULT
- Mixed precision, a fine-control learning-rate optimizer, warmup, and LARS are used to stabilize and enable large-mini-batch training.Computation and communication use half precision, while weight updates use single precision.Source: §IV · p. 3
- 74.7 seconds completes ResNet-50 ImageNet training with 75.08% validation accuracy on the optimized framework.The measurement includes initialization and memory allocation time under the MLPerf v0.5.0 rule.Source: §IV · p. 3
- 1.7 million images per second and 77.0% scalability are achieved with 2,048 GPUs.The reported scalability remains quite good until 2,048 GPUs; Fig. 2 compares the result with an ideal throughput curve.Source: §IV · p. 3, Figure 2 · p. 3
- 81,920-image mini-batches reach over 75% validation accuracy despite the challenge of too few updates per epoch.At this batch size, ImageNet provides only 16 updates per epoch and 1,440 total updates in the reported setup.Source: §IV · p. 3
- The training and validation accuracy comparison indicates no overfitting with batch normalization and label smoothing.Fig. 4 provides the comparison between training and validation accuracy.Source: §IV · p. 3
V. CONCLUSION
The paper develops techniques for large mini-batch training on large GPU clusters without loss of validation accuracy, achieving fast ResNet-50 training. Future work will formulate these techniques and apply them to general DNN training.
V. CONCLUSION
- 75.08% validation accuracy was achieved for ResNet-50 with an 81,920 mini-batch in 74.7 seconds.The result used the MXNet-based framework and is presented as the paper's concluding benchmark.Source: §V · p. 3
- The developed techniques support large mini-batch training on large-scale GPU clusters without loss of validation accuracy.Source: §V · p. 3
VI. FUTURE WORK
- Future work will formulate the techniques and apply them to general DNN training.The authors also describe the resulting ResNet-50 training time as approximately a minute.Source: §VI · p. 3
APPENDIX
The appendix records a ResNet-50 training run, including model configuration, training epochs, evaluation checkpoints, and final evaluation accuracy.
APPENDIX
- The logged run includes batch normalization with momentum 0.9 and epsilon 1e-05.Source: APPENDIX · p. 5
- The appendix records training epochs 0 through 4 and includes an eval_start event.Source: APPENDIX · p. 5
- 0.7343 evaluation accuracy was recorded at epoch 85 during the logged training run.Source: APPENDIX · p. 5
- 0.75082 evaluation accuracy was recorded at epoch 89.Source: APPENDIX · p. 5