探索 Google TPU 上机器学习训练的并发极限
机器学习
2021-03-17 v3 分布式、并行与集群计算
摘要
近期使用神经网络进行语言理解的结果需要前所未有的规模的训练硬件，数千颗芯片在单次训练运行中协同工作。本文提出在 Google TPU Multipod（一种具有 4096 颗 TPU-v3 芯片的网格）上扩展 ML 模型的技术。我们讨论了模型并行以克服数据并行中固定批大小带来的扩展限制、通信/集合优化、训练指标的分布式评估，以及主机输入处理扩展优化。这些技术在 TensorFlow 和 JAX 编程框架中均得到演示。我们还展示了 Google 近期提交至 MLPerf-v0.7 基准竞赛的性能结果，在 Google TPU-v3 Multipod 机器上于四个 MLPerf 模型中实现了 16 至 28 秒的创纪录训练时间。
引用
@article{arxiv.2011.03641,
title = {Exploring the limits of Concurrency in ML Training on Google TPUs},
author = {Sameer Kumar and James Bradbury and Cliff Young and Yu Emma Wang and Anselm Levskaya and Blake Hechtman and Dehao Chen and HyoukJoong Lee and Mehmet Deveci and Naveen Kumar and Pankaj Kanwar and Shibo Wang and Skye Wanderman-Milne and Steve Lacy and Tao Wang and Tayo Oguntebi and Yazhou Zu and Yuanzhong Xu and Andy Swing},
journal= {arXiv preprint arXiv:2011.03641},
year = {2021}
}