Timezone: »

 
Oral
EXPLORING THE LIMITS OF CONCURRENCY IN ML TRAINING ON GOOGLE TPUS
Sameer Kumar · Yu Wang · Cliff Young · James Bradbury · Naveen Kumar · Dehao Chen · Andy Swing

Tue Apr 06 04:20 PM -- 04:40 PM (PDT) @

Recent results in language understanding using neural networks have required training hardware of unprecedented scale, with thousands of chips cooperating on a single training run. This paper presents techniques to scale ML models on the Google TPU Multipod, a mesh with 4096 TPU-v3 chips. We discuss model parallelism to overcome scaling limitations from the fixed batch size in data parallelism, communication/collective optimizations, distributed evaluation of training metrics, and host input processing scaling optimizations. These techniques are demonstrated in both the TensorFlow and JAX programming frameworks. We also present performance results from the recent Google submission to the MLPerf-v0.7 benchmark contest, achieving record training times from 16 to 28 seconds in four MLPerf models on the Google TPU-v3 Multipod machine.

Author Information

Sameer Kumar (Google)
Yu Wang (Google)
Cliff Young (google.com)
James Bradbury (Google)
Naveen Kumar (Google)
Dehao Chen (Google)
Andy Swing (Google)

Related Events (a corresponding poster, oral, or spotlight)

More from the Same Authors

  • 2021 : Panel Discussion »
    Luis Ceze · Cliff Young · Chris Lattner
  • 2021 : Q&A for Cliff Young »
    Cliff Young
  • 2021 : TPU »
    Cliff Young
  • 2020 Oral: MLPerf Training Benchmark »
    Peter Mattson · Christine Cheng · Gregory Diamos · Cody Coleman · Paulius Micikevicius · David Patterson · Hanlin Tang · Gu-Yeon Wei · Peter Bailis · Victor Bittorf · David Brooks · Dehao Chen · Debo Dutta · Udit Gupta · Kim Hazelwood · Andy Hock · Xinyuan Huang · Daniel Kang · David Kanter · Naveen Kumar · Jeffery Liao · Deepak Narayanan · Tayo Oguntebi · Gennady Pekhimenko · Lillian Pentecost · Vijay Janapa Reddi · Taylor Robie · Tom St John · Carole-Jean Wu · Lingjie Xu · Cliff Young · Matei Zaharia
  • 2020 Poster: MLPerf Training Benchmark »
    Peter Mattson · Christine Cheng · Gregory Diamos · Cody Coleman · Paulius Micikevicius · David Patterson · Hanlin Tang · Gu-Yeon Wei · Peter Bailis · Victor Bittorf · David Brooks · Dehao Chen · Debo Dutta · Udit Gupta · Kim Hazelwood · Andy Hock · Xinyuan Huang · Daniel Kang · David Kanter · Naveen Kumar · Jeffery Liao · Deepak Narayanan · Tayo Oguntebi · Gennady Pekhimenko · Lillian Pentecost · Vijay Janapa Reddi · Taylor Robie · Tom St John · Carole-Jean Wu · Lingjie Xu · Cliff Young · Matei Zaharia