[Paper Review] Scale MLPerf-0.6 models on Google TPU-v3 Pods

Sameer Kumar, Victor Bitorff+10 more|arXiv (Cornell University)|Sep 21, 2019

Advanced Data Storage TechnologiesComputer Science11 references35 citations

TL;DR

The paper demonstrates scaling MLPerf-0.6 training benchmarks to Google TPU-v3 Pods (1024 chips) using optimizations like distributed evaluation, gradient summation acceleration, model parallelism, and weight update sharding, achieving record performance on ResNet-50, SSD, and Transformer.

ABSTRACT

The recent submission of Google TPU-v3 Pods to the industry wide MLPerf v0.6 training benchmark demonstrates the scalability of a suite of industry relevant ML models. MLPerf defines a suite of models, datasets and rules to follow when benchmarking to ensure results are comparable across hardware, frameworks and companies. Using this suite of models, we discuss the optimizations and techniques including choice of optimizer, spatial partitioning and weight update sharding necessary to scale to 1024 TPU chips. Furthermore, we identify properties of models that make scaling them challenging, such as limited data parallelism and unscaled weights. These optimizations contribute to record performance in transformer, Resnet-50 and SSD in the Google MLPerf-0.6 submission.

Motivation & Objective

• Demonstrate scalable MLPerf-0.6 training on TPU-v3 Pods (1024 chips) and identify bottlenecks at large scale.
• Develop and validate optimization techniques to improve throughput and maintain target accuracy.
• Showcase model-parallelism and compiler-assisted optimizations to enable scaling of diverse ML workloads.

Proposed method

• Lower TensorFlow graphs with the XLA compiler to optimize TPU-v3 execution.
• Use mixed precision (bfloat16) for compute while keeping critical non-convolutional ops in 32-bit to maintain accuracy.
• Apply a suite of optimizations: distributed evaluation, 2-D gradient summation, and weight update sharding to reduce bottlenecks on large pods.
• Employ model-parallelism techniques (spatial partitioning and halo exchanges) to increase concurrency across cores.
• Use data-parallelism and in-memory evaluation for Transformer and GNMT to mitigate infrastructure overheads.
• Reduce evaluation overheads and optimize input pipelines to sustain high throughput.

Experimental results

Research questions

• RQ1 How can MLPerf-0.6 benchmarks scale when trained on Google TPU-v3 Pods (up to 1024 chips)?
• RQ2 What optimizations are necessary to achieve peak throughput while preserving MLPerf accuracy targets at scale?
• RQ3 What model-specific parallelism and optimizer strategies best enable large-scale scaling for ResNet-50, SSD, Mask-RCNN, Transformer, and GNMT?

Key findings

• Achieved large-scale scaling on TPU-v3 Pods up to 1024 chips, with record performance for several benchmarks.
• Distributing evaluation and tightly coupling train/eval loops reduces bottlenecks and improves scalability.
• 2-D gradient summation with pipelined data transfers yields substantial speedups (over 1.5x in ResNet-50).
• Model parallelism techniques (spatial partitioning and weight update sharding) enable scaling for smaller per-core workloads and for transformers with large parameter counts.
• SSD saw a 1.6x speedup on 4 cores enabling scaling to 2048 TPU cores, and Mask-RCNN scaling benefits from multi-core distribution in first and second stages.
• Transformer and GNMT scaling required dedicated optimizations including in-memory evaluation, reduced maximum sequence length, and memory/compute optimizations to reach target accuracy on large batches.