探索ML中的并发性限制在Google TPU上进行培训
Exploring the limits of Concurrency in ML Training on Google TPUs
摘要:
最近的研究结果表明,使用神经网络的语言理解需要规模空前的训练硬件,在一次训练中需要数千个芯片的协作。本文介绍了在Google TPU Multipod上缩放ML模型的技术,这是一个带有4096 TPU-v3芯片的网格。我们讨论了模型并行来克服数据并行中固定批量大小的扩展限制,通信/集体优化,训练指标的分布式评估,以及主机输入处理扩展优化。这些技术在TensorFlow和JAX编程框架中都有演示。我们还展示了最近Google提交给MLPerf-v0.7基准测试竞赛的性能结果,在Google TPU-v3 Multipod机器上的四个MLPerf模型中,实现了16到28秒的创纪录训练时间。
英文摘要:
Recent results in language understanding using neural networks have required training hardware of unprecedentedscale, with thousands of chips cooperating on a single training run. This paper presents techniques to scaleML models on the Google TPU Multipod, a mesh with 4096 TPU-v3 chips. We discuss model parallelism toovercome 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 aredemonstrated in both the TensorFlow and JAX programming frameworks. We also present performance resultsfrom the recent Google submission to the MLPerf-v0.7 benchmark contest, achieving record training times from16 to 28 seconds in four MLPerf models on the Google TPU-v3 Multipod machine.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。