GShard：利用条件计算与自动分片扩展巨型模型
计算与语言
2020-07-01 v1 机器学习
机器学习
摘要
神经网络扩展对于在拥有海量训练数据和算力的许多真实机器学习应用中提升模型质量至关重要。尽管这一扩展趋势被确认为提升模型质量的可靠途径，但道路上仍存在诸如计算成本、编程便利性以及在并行设备上的高效实现等挑战。GShard 是一个由一组轻量级标注 API 和对 XLA 编译器的扩展组成的模块。它以极小的对现有模型代码的改动，优雅地表达了广泛的并行计算模式。GShard 使我们能够利用自动分片将多语言神经机器翻译 Transformer 模型与稀疏门控混合专家（Sparsely-Gated Mixture-of-Experts）扩展至超过 6000 亿参数。我们展示了这样一个巨型模型可以在 2048 个 TPU v3 加速器上于 4 天内高效训练，从而相较于先前最佳方法，在从 100 种语言到英语的翻译上取得远优的质量。
引用
@article{arxiv.2006.16668,
title = {GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding},
author = {Dmitry Lepikhin and HyoukJoong Lee and Yuanzhong Xu and Dehao Chen and Orhan Firat and Yanping Huang and Maxim Krikun and Noam Shazeer and Zhifeng Chen},
journal= {arXiv preprint arXiv:2006.16668},
year = {2020}
}