arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2006.16668cs.CLcs.LGstat.ML
收藏

GShard:使用条件计算和自动分片扩展巨型模型

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

Google(谷歌)

Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen

摘要:

在许多具有大量训练数据和计算的真实世界机器学习应用中,神经网络缩放对于提高模型质量至关重要。虽然这种缩放趋势被肯定是一种确保更好的模型质量的方法,但在这条道路上存在着挑战,例如计算成本,编程的易用性以及在并行设备上的有效实现。GShard是一个由一组轻量级注释API和XLA编译器扩展组成的模块。它提供了一种优雅的方式来表达广泛的并行计算模式,只需对现有模型代码进行最小的更改。GShard使我们能够使用自动分片扩展多语言神经机器翻译Transformer模型,其中包含超过6000亿个参数的稀疏门控专家混合模型。我们证明,这样一个巨大的模型可以在4天内在2048个TPU v3加速器上有效地训练,与现有技术相比,从100种语言翻译成英语的质量要高得多。

英文摘要:

Neural network scaling has been critical for improving the model quality in many real-world machine learning applications with vast amounts of training data and compute. Although this trend of scaling is affirmed to be a sure-fire approach for better model quality, there are challenges on the path such as the computation cost, ease of programming, and efficient implementation on parallel devices. GShard is a module composed of a set of lightweight annotation APIs and an extension to the XLA compiler. It provides an elegant way to express a wide range of parallel computation patterns with minimal changes to the existing model code. GShard enabled us to scale up multilingual neural machine translation Transformer model with Sparsely-Gated Mixture-of-Experts beyond 600 billion parameters using automatic sharding. We demonstrate that such a giant model can efficiently be trained on 2048 TPU v3 accelerators in 4 days to achieve far superior quality for translation from 100 languages to English compared to the prior art.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑