Source-linked AI summary
Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, Kaiming He
TL;DR
Large minibatches can make distributed SGD efficient but create optimization challenges and had not been shown to preserve accuracy at very large scales. The paper combines linear learning-rate scaling with warmup and trains ResNet-50 on ImageNet in one hour on 256 GPUs while matching small-minibatch accuracy.
Problem
As models and datasets grow, training takes longer, while evidence was lacking that minibatches as large as 8192 could preserve generalization accuracy.
Method
The paper uses a hyper-parameter-free linear learning-rate scaling rule with a warmup phase to address early optimization difficulties in large-minibatch SGD.
Results
Training ResNet-50 on ImageNet with minibatches up to 8192 images matches small-minibatch accuracy, including at 256 workers in one hour.
Takeaways & Limitations
The approach supports efficient distributed training of visual recognition models and simplifies scaling from single- to multi-GPU implementations without hyper-parameter search.
Takeaways & Limitations
Minibatch scaling cannot continue indefinitely: beyond a certain point, accuracy degrades rapidly, with the threshold occurring near 8k images in ImageNet experiments.
Abstract
from arXiv · showhide
Deep learning thrives with large neural networks and large datasets. However, larger networks and larger datasets result in longer training times that impede research and development progress. Distributed synchronous SGD offers a potential solution to this problem by dividing SGD minibatches over a pool of parallel workers. Yet to make this scheme efficient, the per-worker workload must be large, which implies nontrivial growth in the SGD minibatch size. In this paper, we empirically show that on the ImageNet dataset large minibatches cause optimization difficulties, but when these are addressed the trained networks exhibit good generalization. Specifically, we show no loss of accuracy when training with large minibatch sizes up to 8192 images. To achieve this result, we adopt a hyper-parameter-free linear scaling rule for adjusting learning rates as a function of minibatch size and develop a new warmup scheme that overcomes optimization challenges early in training. With these simple techniques, our Caffe2-based system trains ResNet-50 with a minibatch size of 8192 on 256 GPUs in one hour, while matching small minibatch accuracy. Using commodity hardware, our implementation achieves ~90% scaling efficiency when moving from 8 to 256 GPUs. Our findings enable training visual recognition models on internet-scale data with high efficiency.
1. Introduction
As model and dataset scale improve accuracy, they also increase training time, motivating practical distributed synchronous SGD. The report shows that linear learning-rate scaling plus warmup enables accurate ImageNet training with minibatches up to 8192 and extends to more complex vision tasks.
- Motivation: Larger datasets and neural network architectures improve accuracy, but their growth also increases training time and makes large-scale deep learning harder to manage.The introduction frames training-time reduction as necessary for exploring large-scale deep learning.
- Contribution: 8192 images: ResNet-50 training reaches the 256-image minibatch baseline’s accuracy in 1 hour on 256 GPUs.The original 256-image setup used 8 Tesla P100 GPUs and required 29 hours.
- Method: A hyper-parameter-free linear scaling rule adjusts learning rates with minibatch size, while warmup uses lower initial learning rates to overcome early optimization difficulties.The techniques keep other hyper-parameters unchanged in the described scaling approach.
- Findings: Optimization difficulty, rather than poor generalization, is identified as the main problem caused by large minibatches on ImageNet.The finding is based on the report’s comprehensive experiments.
- Findings: The linear scaling rule and warmup also generalize to object detection and instance segmentation through Mask R-CNN.The introduction specifically lists object detection and instance segmentation as more complex tasks where the guideline applies.
2. Large Minibatch SGD
This section motivates large-minibatch SGD for distributed learning and presents linear learning-rate scaling with gradual warmup to preserve accuracy despite early optimization difficulties. It also identifies limits to scaling and explains why Batch Normalization requires preserving per-worker minibatch size.
- Motivation: Large minibatches enable distributed data parallelism without reducing per-worker workload or sacrificing model accuracy.The goal is to replace small minibatches while maintaining training and generalization accuracy.
- Linear Scaling Rule: When minibatch size is multiplied by k, multiply the learning rate by k while keeping other hyperparameters unchanged.This linear scaling rule is reported as effective across a broad range of minibatch sizes and helps match training curves and accuracy.
- Linear Scaling Rule: With linear scaling and warmup, small- and large-minibatch SGD achieve the same final accuracy and closely matching training curves.The empirical results suggest the approximation underlying linear scaling can hold on large-scale, real-world data.
- Limits: ∼8k images marks a point beyond which minibatch scaling causes rapid accuracy degradation in ImageNet experiments.The linear scaling condition also fails early in training, when network parameters change rapidly.
- Warmup: A gradual warmup ramps the learning rate from η to ˆη = kη over 5 epochs, avoiding a sudden increase that can cause training-error spikes.Constant warmup at η for the first 5 epochs was insufficient for large k, whereas gradual ramping supports healthy convergence at the start.
- Batch Normalization: Changing per-worker minibatch size changes the Batch Normalization loss function because its statistics exhibit different random variation at different n.BN computes statistics across samples, so the minibatch dimension is part of the optimized loss definition.
3. Subtleties and Pitfalls of Distributed SGD
Distributed SGD contains implementation subtleties that can silently alter hyper-parameter definitions and degrade accuracy. Correct handling requires care with weight decay, momentum updates, gradient aggregation, and per-epoch data shuffling.
- General implementation subtleties: Implementation errors can change hyper-parameter definitions, producing models that train but have unexpectedly higher error.These issues can be difficult to discover and should be considered explicitly when implementing the underlying solver.
- Weight decay: Scaling the cross-entropy loss is not equivalent to scaling the learning rate when weight decay is present.Weight decay corresponds to the gradient of an L2-regularization term and is added separately to aggregated sample-dependent gradients.
- Momentum correction: Apply momentum correction after changing the learning rate when using the learning-rate-absorbed momentum variant.Without correction, increasing η_t+1 far above η_t makes the history term too small and can destabilize training.
- Gradient aggregation: Normalize the per-worker loss by total minibatch size kn, not per-worker size n, when aggregating gradients across k workers.Allreduce sums gradients, so the missing 1/k factor must be incorporated into the loss; canceling k can also produce incorrect weight decay.
- Data shuffling: Use a single random shuffling of the training data per epoch and divide it among all k workers.This preserves fair comparisons with baselines that use shuffled training data.
4. Communication
The communication design overlaps layer-wise gradient aggregation with backpropagation and uses allreduce across GPUs and servers. Its implementation combines hierarchical reductions, bandwidth-aware interserver algorithms, and a 50Gbit network shown sufficient for ResNet-50.
- Overlapping communication and computation: Layer-wise gradient aggregation runs as soon as each gradient is computed, while backpropagation continues on the next layer.This overlap is enabled because gradients across layers have no data dependency.
- Allreduce design: Allreduce transforms each GPU’s locally computed gradients into the sum of all k gradients available on every GPU.Aggregation becomes harder to hide as parameter counts and GPU compute performance increase.
- Allreduce design: The hierarchical allreduce reduces gradients within each server, sums the server buffers across servers, and broadcasts the results.NCCL handles local reduction and broadcast for buffers of size 256 KB or more; smaller buffers use GPU-to-host copies and CPU reduction.
- Interserver algorithms: For interserver allreduce, halving/doubling generally has lower latency than the ring algorithm and performed much better for buffers up to a million elements.Halving/doubling uses 2 log2(p) communication steps, versus 2(p −1) for the ring algorithm.
- Network requirements: 50Gbit network bandwidth is sufficient for distributed synchronous SGD on ResNet-50, whose peak bandwidth requirement is 12.8 Gbit/s.The estimate uses approximately 25 million parameters, a 100MB parameter size, and 120 ms of backpropagation time on one Tesla P100 GPU.
- Network requirements: During the forward pass, the network can support less latency-sensitive tasks such as reading data or saving network snapshots.The peak aggregation bandwidth requirement occurs during backpropagation.
5. Main Results and Analysis
With linear learning-rate scaling and gradual warmup, minibatches up to 8k match small-minibatch optimization and validation performance, enabling ResNet-50 training on 256 workers in one hour. These results extend to ResNet-101 and transfer learning, while performance deteriorates near or beyond the useful minibatch regime.
- Main result: 256 workers train ResNet-50 on ImageNet in one hour while matching small-minibatch accuracy, without additional hyper-parameter tuning up to 8k images.Linear scaling and warmup enable seamless scaling between small and large minibatches.
- Optimization: Gradual warmup makes 8k-minibatch training error closely match the 256-minibatch baseline after about 20 epochs, whereas no warmup remains inferior and constant warmup causes an error spike.The 8k run uses η = 3.2 from linear scaling, starting at η = 0.1 and increasing gradually.
- Minibatch-size limits: 8k minibatches preserve stable validation error across a broad range of sizes, but 16k degrades ImageNet validation error and COCO transfer performance.For 256-to-8k pretraining, Mask R-CNN box and mask AP are nearly identical; 16k worsens both outcomes.
- ResNet-101: 0.28% is ResNet-101’s error increase with an 8k minibatch and η = 3.2, reaching 22.36% versus 22.08% for the 256-minibatch baseline.Training ResNet-101 takes 92.5 minutes on 256 Tesla P100 GPUs.
- Transfer learning: Nearly identical box and mask AP across 1-to-8-GPU configurations demonstrates that linear learning-rate scaling generalizes beyond classification to Mask R-CNN.Large-minibatch ImageNet pretraining also transfers across datasets and tasks without observed generalization issues.