[论文解读] Distributed TensorFlow with MPI

Abhinav Vishnu, Charles Siegel+1 more|arXiv (Cornell University)|Mar 7, 2016

Parallel Computing and Optimization Techniques参考文献 8被引用 34

一句话总结

该论文提出了一种轻量级扩展方法，通过消息传递接口（MPI）在大规模分布式集群上运行TensorFlow，实现了无需修改TensorFlow运行时即可在多个节点间高效并行训练。该方法在最多64个核心上实现了强可扩展性，在CIFAR-10-DNN上实现3.37倍加速，在Higgs数据集上实现2.6倍加速，证明了其在大规模机器学习工作负载中的高效性与可扩展性。

ABSTRACT

Machine Learning and Data Mining (MLDM) algorithms are becoming increasingly important in analyzing large volume of data generated by simulations, experiments and mobile devices. With increasing data volume, distributed memory systems (such as tightly connected supercomputers or cloud computing systems) are becoming important in designing in-memory and massively parallel MLDM algorithms. Yet, the majority of open source MLDM software is limited to sequential execution with a few supporting multi-core/many-core execution. In this paper, we extend recently proposed Google TensorFlow for execution on large scale clusters using Message Passing Interface (MPI). Our approach requires minimal changes to the TensorFlow runtime -- making the proposed implementation generic and readily usable to increasingly large users of TensorFlow. We evaluate our implementation using an InfiniBand cluster and several well knowndatasets. Our evaluation indicates the efficiency of our proposed implementation.

研究动机与目标

• 将TensorFlow的能力从单节点执行扩展至支持集群上的大规模分布式训练。
• 解决TensorFlow在处理来自模拟和实验的日益庞大的数据集时，因单节点限制而带来的局限性。
• 利用MPI作为分布式内存系统的通信骨干，确保节点间高效协调。
• 通过无需修改TensorFlow运行时来保持向后兼容性。
• 评估所提出的基于MPI的扩展在真实世界数据集上的可扩展性与性能。

提出的方法

• 在分布式TensorFlow架构中，将MPI作为工作节点之间的通信层，实现节点间同步与参数交换。
• 复用TensorFlow现有的计算图与会话模型，在保持其本机执行语义的同时，扩展其分布式操作能力。
• 利用MPI实现参数服务器与工作节点角色，实现高效的模型参数分发与梯度聚合。
• 通过抽象分布式协调机制来最小化对TensorFlow运行时的修改，避免深度集成到核心组件中。
• 利用控制依赖与通过MPI实现的变量同步，确保跨节点训练的一致性。
• 利用TensorFlow的C++后端与优化的线性代数库，最大化每个节点的性能。

实验结果

研究问题

• RQ1 能否通过轻量级、低修改量的方法有效将TensorFlow扩展至大规模集群？
• RQ2 使用MPI作为通信层是否能在保持TensorFlow性能特征的同时支持分布式训练？
• RQ3 所提出的基于MPI的扩展在真实世界机器学习数据集上，随着核心与节点数量增加，其可扩展性如何？
• RQ4 与单节点TensorFlow相比，该分布式实现的性能开销与加速比如何？
• RQ5 该扩展能否高效处理多样化的机器学习工作负载，包括DNN、CNN和表格数据，在MNIST、CIFAR-10、Higgs和Adult等数据集上表现如何？

主要发现

• 所提出的基于MPI的扩展在使用64个核心的CIFAR-10-DNN上实现了3.37倍加速，证明了其强可扩展性效率。
• 在Higgs数据集上，从20个核心扩展到80个核心时，实现2.6倍加速。
• 在MNIST上，16个核心时相对加速达到3.0倍，表明负载分布与通信开销控制良好。
• 该方法保持了高训练精度——在一小时内于CIFAR-10-DNN上达到约55%的准确率，证实了分布式处理下的模型保真度。
• 该实现无需修改TensorFlow运行时，因此具有高度可重用性与可部署性，适用于现有TensorFlow工作流。
• 在多个数据集上的评估证实了该方法的通用性与高效性，在各类监督学习任务中均表现出一致的性能提升。