Source-linked AI summary
Image Classification at Supercomputer Scale
Chris Ying, Sameer Kumar, Dehao Chen, Tao Wang, Youlong Cheng
TL;DR
Large-scale deep learning requires substantial computation, while distributed training must address accuracy and systems challenges to achieve end-to-end speedup. The paper combines distributed batch normalization, input-pipeline optimizations, and 2-D gradient summation, training ResNet-50 on ImageNet to 76.3% accuracy in 2.2 minutes on a 1024-chip TPU v3 Pod.
Problem
Large-scale synchronous distributed training must manage batch-size effects on accuracy, CPU input bottlenecks, and scalable gradient reduction to achieve end-to-end wall-clock speedup.
Source: §1 · p. 1, §3.1 · p. 3Method
The paper combines distributed batch normalization, input-pipeline optimizations, and 2-D gradient summation with large-batch training techniques.
Source: §1 · p. 1, §3.1 · p. 3, §3.2 · p. 3Results
76.3% accuracy was maintained with no accuracy drop while training ResNet-50 on 1024 TPU v3 chips in 2.2 minutes at over 1.05 million images/second.
Source: §5 · p. 6Takeaways & Limitations
The authors state that their techniques, evaluated on TPUs, are general and apply to other accelerators as well.
Source: §5 · p. 6The experiments explore scaling image classification using purely data parallelism, rather than a mix of data and model parallelism.
Source: §6 · p. 6
Abstract
from arXiv · showhide
Deep learning is extremely computationally intensive, and hardware vendors have responded by building faster accelerators in large clusters. Training deep learning models at petaFLOPS scale requires overcoming both algorithmic and systems software challenges. In this paper, we discuss three systems-related optimizations: (1) distributed batch normalization to control per-replica batch sizes, (2) input pipeline optimizations to sustain model throughput, and (3) 2-D torus all-reduce to speed up gradient summation. We combine these optimizations to train ResNet-50 on ImageNet to 76.3% accuracy in 2.2 minutes on a 1024-chip TPU v3 Pod with a training throughput of over 1.05 million images/second and no accuracy drop.
1 Introduction
The introduction frames large-scale image classification as a compute- and systems-challenging task, with model quality and end-to-end speed both requiring attention. It presents systems optimizations for scaling and describes the TPU hardware used as context.
1 Introduction
- ResNet-50 on ImageNet requires about 3.2 × 10^16 floating-point operations per epoch, and typical training can require up to 90 epochs.The paper uses this compute burden to motivate training at scale.Source: §1 · p. 1
- Scaling distributed synchronous SGD must address accuracy dependence on global and per-replica batch sizes, CPU input bottlenecks, and scalable all-reduce.The introduction notes that model-quality adaptations and systems optimizations are needed for real end-to-end speedup.Source: §1 · p. 1
- The paper presents systems optimizations for these challenges, evaluates each in controlled experiments, and combines them with machine-learning techniques for ResNet-50 training..Source: §1 · p. 1
1.1 Tensor Processing Units
- A TPU v2 Pod contains 256 chips, while the TPU v3 Pod contains 1024 chips and has a theoretical peak of 107.5 petaFLOPS mixed-precision throughput.Each TPU device has four chips and two cores per chip.Source: §1.1 · p. 2
2 Related Work
Prior work explored large-batch training and alternative ways to scale ResNet-50, including optimizer, normalization, precision, and iteration-count changes.
2 Related Work
- Keskar et al. [12] described a generalization gap between small- and large-batch training and suggested sharp minima as a possible explanation.The passage reports the proposed explanation as a possibility.Source: §2 · p. 2
- Goyal et al. [6] trained ResNet-50 on ImageNet with batch size 8192 in one hour, while You et al. [19] scaled it to batch size 32768 using layer-wise learning rates.Smith et al. [16] explored increasing batch size during training alongside learning-rate changes.Source: §2 · p. 2
- Akiba et al. [2] reported 15-minute ResNet-50 training, and Jia et al. [10] reduced training time to 8.7 minutes using mixed precision and optimized all-reduce.Smith and Topin [15] instead approached the problem by reducing the number of iterations.Source: §2 · p. 2
3 Methods
The methods combine established large-batch training techniques with distributed batch normalization, input-pipeline optimizations, and 2-D gradient summation. These systems methods target batch-size control, sustained accelerator throughput, and scalable gradient reduction.
3 Methods
- Convolutions and activations use bfloat16 while non-convolution operations use 32-bit values; training also uses linear learning-rate scaling, warm-up and decay, and LARS for batch sizes up to 32768.LARS enabled scaling beyond the batch size of 8192 reached with SGD and momentum.Source: §3 · pp. 2–3
3.1 Distributed Batch Normalization
- Distributed batch normalization reduces across a small peer group to compute mean and variance, decoupling its effective batch size from global batch size and replica count.This addresses the concern that small per-replica batch sizes can impair final validation accuracy.Source: §3.1 · p. 3
3.2 Input Pipeline Optimization
- The input pipeline combines dataset sharding, prefetching, fused JPEG decoding and cropping, and parallel parsing to sustain model throughput.The authors report that their fastest results would otherwise be input-bottlenecked.Source: §3.2 · pp. 3–4
3.3 2-D Gradient Summation
- The 2-D all-reduce performs concurrent reductions on tensor halves along two dimensions, reducing gradient-summation span from O(N^2) to O(N) network hops and doubling throughput over 1-D.On a torus, bidirectional ring exchange can further speed gradient summation.Source: §3.3 · p. 4
4 Analysis
Controlled experiments assess the individual effects of distributed batch normalization, input-pipeline optimizations, and gradient-summation algorithms.
Distributed Batch Normalization
- An effective BN batch size of 64 (group size 4) is selected, while larger groups increase training time.With group size 1, the effective BN batch size of 16 causes a large performance drop.Source: §4 · p. 5
Input Pipeline Optimizations
- Parallelization more than doubles base input throughput; fused JPEG decode and crop improve it by around 25% when added.Ablations show around a 10% drop for fused JPEG decode and crop and 6–13% drops for caching and prefetching; sharding ImageNet across a TPU Pod permits full in-memory caching.Source: §4 · p. 5, Figure 5 · p. 5
2-D Gradient Summation
- For ResNet-50 gradient reductions across TPU v2 chip counts, 2-D is more scalable than 1-D, with lower latency and higher throughput; on the 256-chip pod, torus links enable more efficient summation than baseline 1-D ring and 2-D mesh algorithms.Source: §4 · p. 5
5 Results
The combined techniques train ResNet-50 on ImageNet in 2.2 minutes on 1024 TPU v3 chips, reaching 76.3% accuracy with no accuracy drop. The paper also reports the training recipe and evaluation setup.
Results
- Training runs for 90 epochs for comparability with prior results; time is measured from TPU binary compilation completion to final checkpoint writing.Accuracy is measured on the validation set without removing blacklisted examples.Source: §5 · p. 5
- The model is the ResNet-50 variant popularized by Goyal et al. [6], which differs from canonical v1 [7], uses about 12% more compute, and achieves higher accuracy.Label smoothing is also used to improve generalization.Source: §5 · p. 5
- Table 1 summarizes large-scale ResNet-50 training results; Figure 6 varies effective BN size at fixed per-core batch size 16, while Figure 7 compares 1-D and 2-D reductions across chip counts.Figure 6 reports validation accuracy and total training time; torus links are available only on the full 256-chip TPU v2 Pod.Source: Table 1 · p. 6, Figure 6 · p. 6, Figure 7 · p. 6
- The training recipe combines bfloat16, learning-rate scaling, warmup and decay, LARS, distributed batch normalization, input-pipeline optimizations, and 2-D gradient summation with torus links.Training used TensorFlow on a 1024-chip TPU v3 Pod.Source: §5 · p. 6
- In 2.2 minutes, the 1024-chip TPU v3 Pod reaches 76.3% accuracy with over 1.05 million images/second throughput and no accuracy drop.The reported time includes constant start-up costs such as weight initialization; throughput was measured during training via profiling.Source: §5 · p. 6
6 Future Work
The paper explores scaling image classification using purely data parallelism.
6 Future Work
- The paper's scope is scaling image classification through purely data parallelism.Source: §6 · p. 6
- An alternative scaling approach combines data and model parallelism to shard computation.Source: §6 · p. 6
- That mixed-parallelism approach can use large clusters without necessarily requiring large global batch sizes.Source: §6 · p. 6