Mesh-TensorFlow：超级计算机上的深度学习

批分割（数据并行）是分布式深度神经网络（DNN）训练的主要策略，这归因于其普遍适用性和对单程序多数据（SPMD）编程的支持。然而，批分割面临着一些问题，包括由于内存限制而无法训练非常大的模型、高延迟以及在小批量大小时的低效性。所有这些问题都可以通过更为通用的分布策略（模型并行）来解决。不幸的是，高效的模型并行算法往往难以发现、描述和实现，尤其是在大型集群上。我们介绍了Mesh-TensorFlow，这是一种用于指定一类通用分布式张量计算的语言。数据并行可以被视为沿“批”维度分割张量和操作，而在Mesh-TensorFlow中，用户可以指定任何张量维度以在处理器的多维网格的任意维度上进行分割。一个Mesh-TensorFlow图会编译成一个包含并行操作和集体通信原语（如Allreduce）的SPMD程序。我们使用Mesh-TensorFlow实现了Transformer序列到序列模型的一个高效的数据并行和模型并行版本。利用最多512个核心的TPU网格，我们训练了参数量高达50亿的Transformer模型，在WMT'14英法翻译任务和十亿词语言建模基准测试中超越了现有最佳结果。Mesh-TensorFlow可在https://github.com/tensorflow/mesh 获取。